setup.py 15.7 KB
Newer Older
1
import importlib.util
2
import io
3
import logging
4
5
import os
import re
6
import subprocess
bnellnm's avatar
bnellnm committed
7
import sys
8
import warnings
9
from shutil import which
10
from typing import Dict, List
11

Woosuk Kwon's avatar
Woosuk Kwon committed
12
import torch
13
14
15
from packaging.version import Version, parse
from setuptools import Extension, find_packages, setup
from setuptools.command.build_ext import build_ext
bnellnm's avatar
bnellnm committed
16
from torch.utils.cpp_extension import CUDA_HOME
17

18
19
20
21
22
23
24
25
26

def load_module_from_path(module_name, path):
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module
    spec.loader.exec_module(module)
    return module


27
ROOT_DIR = os.path.dirname(__file__)
28
logger = logging.getLogger(__name__)
29

30
31
32

def embed_commit_hash():
    try:
youkaichao's avatar
youkaichao committed
33
34
35
36
37
38
39
        if "BUILDKITE_COMMIT" in os.environ:
            # ci build
            commit_id = os.environ["BUILDKITE_COMMIT"]
        else:
            commit_id = subprocess.check_output(["git", "rev-parse", "HEAD"],
                                                encoding="utf-8").strip()

40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
        commit_contents = f'__commit__ = "{commit_id}"\n'

        version_file = os.path.join(ROOT_DIR, "vllm", "commit_id.py")
        with open(version_file, "w", encoding="utf-8") as f:
            f.write(commit_contents)

    except subprocess.CalledProcessError as e:
        warnings.warn(f"Failed to get commit hash:\n{e}",
                      RuntimeWarning,
                      stacklevel=2)
    except Exception as e:
        warnings.warn(f"Failed to embed commit hash:\n{e}",
                      RuntimeWarning,
                      stacklevel=2)


embed_commit_hash()

58
59
60
61
62
# cannot import envs directly because it depends on vllm,
#  which is not installed yet
envs = load_module_from_path('envs', os.path.join(ROOT_DIR, 'vllm', 'envs.py'))

VLLM_TARGET_DEVICE = envs.VLLM_TARGET_DEVICE
63

64
65
66
67
# vLLM only supports Linux platform
assert sys.platform.startswith(
    "linux"), "vLLM only supports Linux platform (including WSL)."

68
69
MAIN_CUDA_VERSION = "12.1"

bnellnm's avatar
bnellnm committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

def is_sccache_available() -> bool:
    return which("sccache") is not None


def is_ccache_available() -> bool:
    return which("ccache") is not None


def is_ninja_available() -> bool:
    return which("ninja") is not None


def remove_prefix(text, prefix):
    if text.startswith(prefix):
        return text[len(prefix):]
    return text


class CMakeExtension(Extension):

    def __init__(self, name: str, cmake_lists_dir: str = '.', **kwa) -> None:
92
        super().__init__(name, sources=[], py_limited_api=True, **kwa)
bnellnm's avatar
bnellnm committed
93
94
95
96
97
        self.cmake_lists_dir = os.path.abspath(cmake_lists_dir)


class cmake_build_ext(build_ext):
    # A dict of extension directories that have been configured.
98
    did_config: Dict[str, bool] = {}
bnellnm's avatar
bnellnm committed
99
100
101
102
103

    #
    # Determine number of compilation jobs and optionally nvcc compile threads.
    #
    def compute_num_jobs(self):
104
105
        # `num_jobs` is either the value of the MAX_JOBS environment variable
        # (if defined) or the number of CPUs available.
106
        num_jobs = envs.MAX_JOBS
107
108
        if num_jobs is not None:
            num_jobs = int(num_jobs)
109
            logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs)
110
111
112
113
114
115
116
        else:
            try:
                # os.sched_getaffinity() isn't universally available, so fall
                #  back to os.cpu_count() if we get an error here.
                num_jobs = len(os.sched_getaffinity(0))
            except AttributeError:
                num_jobs = os.cpu_count()
bnellnm's avatar
bnellnm committed
117

118
        nvcc_threads = None
119
120
121
122
123
        if _is_cuda() and get_nvcc_cuda_version() >= Version("11.2"):
            # `nvcc_threads` is either the value of the NVCC_THREADS
            # environment variable (if defined) or 1.
            # when it is set, we reduce `num_jobs` to avoid
            # overloading the system.
124
            nvcc_threads = envs.NVCC_THREADS
125
126
            if nvcc_threads is not None:
                nvcc_threads = int(nvcc_threads)
127
128
129
                logger.info(
                    "Using NVCC_THREADS=%d as the number of nvcc threads.",
                    nvcc_threads)
130
131
132
            else:
                nvcc_threads = 1
            num_jobs = max(1, num_jobs // nvcc_threads)
bnellnm's avatar
bnellnm committed
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

        return num_jobs, nvcc_threads

    #
    # Perform cmake configuration for a single extension.
    #
    def configure(self, ext: CMakeExtension) -> None:
        # If we've already configured using the CMakeLists.txt for
        # this extension, exit early.
        if ext.cmake_lists_dir in cmake_build_ext.did_config:
            return

        cmake_build_ext.did_config[ext.cmake_lists_dir] = True

        # Select the build type.
        # Note: optimization level + debug info are set by the build type
        default_cfg = "Debug" if self.debug else "RelWithDebInfo"
150
        cfg = envs.CMAKE_BUILD_TYPE or default_cfg
bnellnm's avatar
bnellnm committed
151
152
153
154
155
156
157
158
159
160

        # where .so files will be written, should be the same for all extensions
        # that use the same CMakeLists.txt.
        outdir = os.path.abspath(
            os.path.dirname(self.get_ext_fullpath(ext.name)))

        cmake_args = [
            '-DCMAKE_BUILD_TYPE={}'.format(cfg),
            '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={}'.format(outdir),
            '-DCMAKE_ARCHIVE_OUTPUT_DIRECTORY={}'.format(self.build_temp),
161
            '-DVLLM_TARGET_DEVICE={}'.format(VLLM_TARGET_DEVICE),
bnellnm's avatar
bnellnm committed
162
163
        ]

164
        verbose = envs.VERBOSE
bnellnm's avatar
bnellnm committed
165
166
167
168
169
170
171
        if verbose:
            cmake_args += ['-DCMAKE_VERBOSE_MAKEFILE=ON']

        if is_sccache_available():
            cmake_args += [
                '-DCMAKE_CXX_COMPILER_LAUNCHER=sccache',
                '-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache',
172
                '-DCMAKE_C_COMPILER_LAUNCHER=sccache',
bnellnm's avatar
bnellnm committed
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
            ]
        elif is_ccache_available():
            cmake_args += [
                '-DCMAKE_CXX_COMPILER_LAUNCHER=ccache',
                '-DCMAKE_CUDA_COMPILER_LAUNCHER=ccache',
            ]

        # Pass the python executable to cmake so it can find an exact
        # match.
        cmake_args += ['-DVLLM_PYTHON_EXECUTABLE={}'.format(sys.executable)]

        #
        # Setup parallelism and build tool
        #
        num_jobs, nvcc_threads = self.compute_num_jobs()

        if nvcc_threads:
            cmake_args += ['-DNVCC_THREADS={}'.format(nvcc_threads)]

        if is_ninja_available():
            build_tool = ['-G', 'Ninja']
            cmake_args += [
                '-DCMAKE_JOB_POOL_COMPILE:STRING=compile',
                '-DCMAKE_JOB_POOLS:STRING=compile={}'.format(num_jobs),
            ]
        else:
            # Default build tool to whatever cmake picks.
            build_tool = []
        subprocess.check_call(
            ['cmake', ext.cmake_lists_dir, *build_tool, *cmake_args],
            cwd=self.build_temp)

    def build_extensions(self) -> None:
        # Ensure that CMake is present and working
        try:
            subprocess.check_output(['cmake', '--version'])
        except OSError as e:
            raise RuntimeError('Cannot find CMake executable') from e

        # Create build directory if it does not exist.
        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)

216
        targets = []
bnellnm's avatar
bnellnm committed
217
218
219
        # Build all the extensions
        for ext in self.extensions:
            self.configure(ext)
220
            targets.append(remove_prefix(ext.name, "vllm."))
bnellnm's avatar
bnellnm committed
221

222
        num_jobs, _ = self.compute_num_jobs()
bnellnm's avatar
bnellnm committed
223

224
225
226
227
228
229
        build_args = [
            "--build",
            ".",
            f"-j={num_jobs}",
            *[f"--target={name}" for name in targets],
        ]
bnellnm's avatar
bnellnm committed
230

231
        subprocess.check_call(["cmake", *build_args], cwd=self.build_temp)
232
233


234
def _is_cuda() -> bool:
235
236
237
    has_cuda = torch.version.cuda is not None
    return (VLLM_TARGET_DEVICE == "cuda" and has_cuda
            and not (_is_neuron() or _is_tpu()))
238
239


240
def _is_hip() -> bool:
241
242
    return (VLLM_TARGET_DEVICE == "cuda"
            or VLLM_TARGET_DEVICE == "rocm") and torch.version.hip is not None
243
244


245
246
247
248
def _is_neuron() -> bool:
    torch_neuronx_installed = True
    try:
        subprocess.run(["neuron-ls"], capture_output=True, check=True)
249
    except (FileNotFoundError, PermissionError, subprocess.CalledProcessError):
250
        torch_neuronx_installed = False
251
    return torch_neuronx_installed or VLLM_TARGET_DEVICE == "neuron"
252
253


254
255
256
257
def _is_tpu() -> bool:
    return VLLM_TARGET_DEVICE == "tpu"


258
259
260
261
def _is_cpu() -> bool:
    return VLLM_TARGET_DEVICE == "cpu"


262
263
264
265
def _is_openvino() -> bool:
    return VLLM_TARGET_DEVICE == "openvino"


266
267
268
269
def _is_xpu() -> bool:
    return VLLM_TARGET_DEVICE == "xpu"


270
271
272
273
def _build_custom_ops() -> bool:
    return _is_cuda() or _is_hip() or _is_cpu()


274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
def get_hipcc_rocm_version():
    # Run the hipcc --version command
    result = subprocess.run(['hipcc', '--version'],
                            stdout=subprocess.PIPE,
                            stderr=subprocess.STDOUT,
                            text=True)

    # Check if the command was executed successfully
    if result.returncode != 0:
        print("Error running 'hipcc --version'")
        return None

    # Extract the version using a regular expression
    match = re.search(r'HIP version: (\S+)', result.stdout)
    if match:
        # Return the version string
        return match.group(1)
    else:
        print("Could not find HIP version in the output")
        return None
Woosuk Kwon's avatar
Woosuk Kwon committed
294

295

296
297
298
def get_neuronxcc_version():
    import sysconfig
    site_dir = sysconfig.get_paths()["purelib"]
299
300
    version_file = os.path.join(site_dir, "neuronxcc", "version",
                                "__init__.py")
301
302
303
304
305
306
307
308
309
310
311
312
313
314

    # Check if the command was executed successfully
    with open(version_file, "rt") as fp:
        content = fp.read()

    # Extract the version using a regular expression
    match = re.search(r"__version__ = '(\S+)'", content)
    if match:
        # Return the version string
        return match.group(1)
    else:
        raise RuntimeError("Could not find HIP version in the output")


bnellnm's avatar
bnellnm committed
315
def get_nvcc_cuda_version() -> Version:
316
317
318
319
    """Get the CUDA version from nvcc.

    Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
    """
320
    assert CUDA_HOME is not None, "CUDA_HOME is not set"
bnellnm's avatar
bnellnm committed
321
    nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"],
322
323
324
325
326
327
328
                                          universal_newlines=True)
    output = nvcc_output.split()
    release_idx = output.index("release") + 1
    nvcc_cuda_version = parse(output[release_idx].split(",")[0])
    return nvcc_cuda_version


329
330
331
332
def get_path(*filepath) -> str:
    return os.path.join(ROOT_DIR, *filepath)


333
def find_version(filepath: str) -> str:
334
335
336
337
338
    """Extract version information from the given filepath.

    Adapted from https://github.com/ray-project/ray/blob/0b190ee1160eeca9796bc091e07eaebf4c85b511/python/setup.py
    """
    with open(filepath) as fp:
Woosuk Kwon's avatar
Woosuk Kwon committed
339
340
        version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
                                  fp.read(), re.M)
341
342
343
344
345
        if version_match:
            return version_match.group(1)
        raise RuntimeError("Unable to find version string.")


346
def get_vllm_version() -> str:
347
    version = find_version(get_path("vllm", "version.py"))
348

349
    if _is_cuda():
bnellnm's avatar
bnellnm committed
350
        cuda_version = str(get_nvcc_cuda_version())
351
352
353
354
        if cuda_version != MAIN_CUDA_VERSION:
            cuda_version_str = cuda_version.replace(".", "")[:3]
            version += f"+cu{cuda_version_str}"
    elif _is_hip():
355
356
357
358
359
        # Get the HIP version
        hipcc_version = get_hipcc_rocm_version()
        if hipcc_version != MAIN_CUDA_VERSION:
            rocm_version_str = hipcc_version.replace(".", "")[:3]
            version += f"+rocm{rocm_version_str}"
360
361
    elif _is_neuron():
        # Get the Neuron version
bnellnm's avatar
bnellnm committed
362
        neuron_version = str(get_neuronxcc_version())
363
364
365
        if neuron_version != MAIN_CUDA_VERSION:
            neuron_version_str = neuron_version.replace(".", "")[:3]
            version += f"+neuron{neuron_version_str}"
366
367
    elif _is_openvino():
        version += "+openvino"
368
369
    elif _is_tpu():
        version += "+tpu"
370
371
    elif _is_cpu():
        version += "+cpu"
372
373
    elif _is_xpu():
        version += "+xpu"
374
    else:
375
        raise RuntimeError("Unknown runtime environment")
376

377
378
379
    return version


380
def read_readme() -> str:
Stephen Krider's avatar
Stephen Krider committed
381
382
383
384
385
386
    """Read the README file if present."""
    p = get_path("README.md")
    if os.path.isfile(p):
        return io.open(get_path("README.md"), "r", encoding="utf-8").read()
    else:
        return ""
387
388


389
390
def get_requirements() -> List[str]:
    """Get Python package dependencies from requirements.txt."""
391
392
393

    def _read_requirements(filename: str) -> List[str]:
        with open(get_path(filename)) as f:
394
            requirements = f.read().strip().split("\n")
395
396
397
398
399
400
401
402
403
404
        resolved_requirements = []
        for line in requirements:
            if line.startswith("-r "):
                resolved_requirements += _read_requirements(line.split()[1])
            else:
                resolved_requirements.append(line)
        return resolved_requirements

    if _is_cuda():
        requirements = _read_requirements("requirements-cuda.txt")
405
        cuda_major, cuda_minor = torch.version.cuda.split(".")
406
407
        modified_requirements = []
        for req in requirements:
408
409
            if ("vllm-flash-attn" in req
                    and not (cuda_major == "12" and cuda_minor == "1")):
410
411
412
413
                # vllm-flash-attn is built only for CUDA 12.1.
                # Skip for other versions.
                continue
            modified_requirements.append(req)
414
        requirements = modified_requirements
415
    elif _is_hip():
416
        requirements = _read_requirements("requirements-rocm.txt")
417
    elif _is_neuron():
418
        requirements = _read_requirements("requirements-neuron.txt")
419
420
    elif _is_openvino():
        requirements = _read_requirements("requirements-openvino.txt")
421
422
    elif _is_tpu():
        requirements = _read_requirements("requirements-tpu.txt")
423
    elif _is_cpu():
424
        requirements = _read_requirements("requirements-cpu.txt")
425
426
    elif _is_xpu():
        requirements = _read_requirements("requirements-xpu.txt")
427
428
    else:
        raise ValueError(
429
430
            "Unsupported platform, please use CUDA, ROCm, Neuron, "
            "OpenVINO, or CPU.")
431
432
433
    return requirements


bnellnm's avatar
bnellnm committed
434
435
ext_modules = []

436
if _is_cuda() or _is_hip():
bnellnm's avatar
bnellnm committed
437
438
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))

439
if _build_custom_ops():
bnellnm's avatar
bnellnm committed
440
441
    ext_modules.append(CMakeExtension(name="vllm._C"))

442
443
444
package_data = {
    "vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]
}
445
if envs.VLLM_USE_PRECOMPILED:
446
    ext_modules = []
Simon Mo's avatar
Simon Mo committed
447
448
    package_data["vllm"].append("*.so")

bnellnm's avatar
bnellnm committed
449
setup(
Woosuk Kwon's avatar
Woosuk Kwon committed
450
    name="vllm",
451
    version=get_vllm_version(),
Woosuk Kwon's avatar
Woosuk Kwon committed
452
    author="vLLM Team",
453
    license="Apache 2.0",
Woosuk Kwon's avatar
Woosuk Kwon committed
454
455
    description=("A high-throughput and memory-efficient inference and "
                 "serving engine for LLMs"),
456
457
    long_description=read_readme(),
    long_description_content_type="text/markdown",
458
    url="https://github.com/vllm-project/vllm",
459
    project_urls={
460
461
        "Homepage": "https://github.com/vllm-project/vllm",
        "Documentation": "https://vllm.readthedocs.io/en/latest/",
462
463
464
465
466
    },
    classifiers=[
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
Woosuk Kwon's avatar
Woosuk Kwon committed
467
        "Programming Language :: Python :: 3.11",
468
        "Programming Language :: Python :: 3.12",
469
470
471
        "License :: OSI Approved :: Apache Software License",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
bnellnm's avatar
bnellnm committed
472
    packages=find_packages(exclude=("benchmarks", "csrc", "docs", "examples",
473
                                    "tests*")),
474
475
    python_requires=">=3.8",
    install_requires=get_requirements(),
Woosuk Kwon's avatar
Woosuk Kwon committed
476
    ext_modules=ext_modules,
477
    extras_require={
478
        "tensorizer": ["tensorizer>=2.9.0"],
479
    },
480
    cmdclass={"build_ext": cmake_build_ext} if _build_custom_ops() else {},
Simon Mo's avatar
Simon Mo committed
481
    package_data=package_data,
Ethan Xu's avatar
Ethan Xu committed
482
483
484
485
486
    entry_points={
        "console_scripts": [
            "vllm=vllm.scripts:main",
        ],
    },
Woosuk Kwon's avatar
Woosuk Kwon committed
487
)