setup.py 15.9 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
def _build_core_ext() -> bool:
    return not _is_neuron() and not _is_tpu()


278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
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
298

299

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

    # 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
319
def get_nvcc_cuda_version() -> Version:
320
321
322
323
    """Get the CUDA version from nvcc.

    Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
    """
324
    assert CUDA_HOME is not None, "CUDA_HOME is not set"
bnellnm's avatar
bnellnm committed
325
    nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"],
326
327
328
329
330
331
332
                                          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


333
334
335
336
def get_path(*filepath) -> str:
    return os.path.join(ROOT_DIR, *filepath)


337
def find_version(filepath: str) -> str:
338
339
340
341
342
    """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
343
344
        version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
                                  fp.read(), re.M)
345
346
347
348
349
        if version_match:
            return version_match.group(1)
        raise RuntimeError("Unable to find version string.")


350
def get_vllm_version() -> str:
351
    version = find_version(get_path("vllm", "version.py"))
352

353
    if _is_cuda():
bnellnm's avatar
bnellnm committed
354
        cuda_version = str(get_nvcc_cuda_version())
355
356
357
358
        if cuda_version != MAIN_CUDA_VERSION:
            cuda_version_str = cuda_version.replace(".", "")[:3]
            version += f"+cu{cuda_version_str}"
    elif _is_hip():
359
360
361
362
363
        # 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}"
364
365
    elif _is_neuron():
        # Get the Neuron version
bnellnm's avatar
bnellnm committed
366
        neuron_version = str(get_neuronxcc_version())
367
368
369
        if neuron_version != MAIN_CUDA_VERSION:
            neuron_version_str = neuron_version.replace(".", "")[:3]
            version += f"+neuron{neuron_version_str}"
370
371
    elif _is_openvino():
        version += "+openvino"
372
373
    elif _is_tpu():
        version += "+tpu"
374
375
    elif _is_cpu():
        version += "+cpu"
376
377
    elif _is_xpu():
        version += "+xpu"
378
    else:
379
        raise RuntimeError("Unknown runtime environment")
380

381
382
383
    return version


384
def read_readme() -> str:
Stephen Krider's avatar
Stephen Krider committed
385
386
387
388
389
390
    """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 ""
391
392


393
394
def get_requirements() -> List[str]:
    """Get Python package dependencies from requirements.txt."""
395
396
397

    def _read_requirements(filename: str) -> List[str]:
        with open(get_path(filename)) as f:
398
            requirements = f.read().strip().split("\n")
399
400
401
402
403
404
405
406
407
408
        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")
409
        cuda_major, cuda_minor = torch.version.cuda.split(".")
410
411
        modified_requirements = []
        for req in requirements:
412
413
            if ("vllm-flash-attn" in req
                    and not (cuda_major == "12" and cuda_minor == "1")):
414
415
416
417
                # vllm-flash-attn is built only for CUDA 12.1.
                # Skip for other versions.
                continue
            modified_requirements.append(req)
418
        requirements = modified_requirements
419
    elif _is_hip():
420
        requirements = _read_requirements("requirements-rocm.txt")
421
    elif _is_neuron():
422
        requirements = _read_requirements("requirements-neuron.txt")
423
424
    elif _is_openvino():
        requirements = _read_requirements("requirements-openvino.txt")
425
426
    elif _is_tpu():
        requirements = _read_requirements("requirements-tpu.txt")
427
    elif _is_cpu():
428
        requirements = _read_requirements("requirements-cpu.txt")
429
430
    elif _is_xpu():
        requirements = _read_requirements("requirements-xpu.txt")
431
432
    else:
        raise ValueError(
433
434
            "Unsupported platform, please use CUDA, ROCm, Neuron, "
            "OpenVINO, or CPU.")
435
436
437
    return requirements


bnellnm's avatar
bnellnm committed
438
439
ext_modules = []

440
441
442
if _build_core_ext():
    ext_modules.append(CMakeExtension(name="vllm._core_C"))

443
if _is_cuda() or _is_hip():
bnellnm's avatar
bnellnm committed
444
445
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))

446
if _build_custom_ops():
bnellnm's avatar
bnellnm committed
447
448
    ext_modules.append(CMakeExtension(name="vllm._C"))

449
450
451
package_data = {
    "vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]
}
452
if envs.VLLM_USE_PRECOMPILED:
453
    ext_modules = []
Simon Mo's avatar
Simon Mo committed
454
455
    package_data["vllm"].append("*.so")

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