setup.py 17.8 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

zhuwenwen's avatar
zhuwenwen committed
18
19
20
21
from typing import Optional, Union
import subprocess
from pathlib import Path

gaoqiong's avatar
gaoqiong committed
22
add_git_version = False
23
24
if int(os.environ.get('ADD_GIT_VERSION', '0')) == 1:
    add_git_version = True
25
26
27
28
29
30
31
32
33

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


34
ROOT_DIR = os.path.dirname(__file__)
35
logger = logging.getLogger(__name__)
36

37
38
39

def embed_commit_hash():
    try:
youkaichao's avatar
youkaichao committed
40
41
42
43
44
45
46
        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()

47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
        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()

65
66
67
68
69
# 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
70

71
72
73
# vLLM only supports Linux platform
assert sys.platform.startswith(
    "linux"), "vLLM only supports Linux platform (including WSL)."
74

75
76
MAIN_CUDA_VERSION = "12.1"

bnellnm's avatar
bnellnm committed
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

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:
99
        super().__init__(name, sources=[], py_limited_api=True, **kwa)
bnellnm's avatar
bnellnm committed
100
101
102
103
104
        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.
105
    did_config: Dict[str, bool] = {}
bnellnm's avatar
bnellnm committed
106
107
108
109
110

    #
    # Determine number of compilation jobs and optionally nvcc compile threads.
    #
    def compute_num_jobs(self):
111
112
        # `num_jobs` is either the value of the MAX_JOBS environment variable
        # (if defined) or the number of CPUs available.
113
        num_jobs = envs.MAX_JOBS
114
115
        if num_jobs is not None:
            num_jobs = int(num_jobs)
116
            logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs)
117
118
119
120
121
122
123
        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
124

125
        nvcc_threads = None
126
127
128
129
130
        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.
131
            nvcc_threads = envs.NVCC_THREADS
132
133
            if nvcc_threads is not None:
                nvcc_threads = int(nvcc_threads)
134
135
136
                logger.info(
                    "Using NVCC_THREADS=%d as the number of nvcc threads.",
                    nvcc_threads)
137
138
139
            else:
                nvcc_threads = 1
            num_jobs = max(1, num_jobs // nvcc_threads)
bnellnm's avatar
bnellnm committed
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

        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"
157
        cfg = envs.CMAKE_BUILD_TYPE or default_cfg
bnellnm's avatar
bnellnm committed
158
159
160
161
162
163
164
165
166
167

        # 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),
168
            '-DVLLM_TARGET_DEVICE={}'.format(VLLM_TARGET_DEVICE),
bnellnm's avatar
bnellnm committed
169
170
        ]

171
        verbose = envs.VERBOSE
bnellnm's avatar
bnellnm committed
172
173
174
175
176
177
178
        if verbose:
            cmake_args += ['-DCMAKE_VERBOSE_MAKEFILE=ON']

        if is_sccache_available():
            cmake_args += [
                '-DCMAKE_CXX_COMPILER_LAUNCHER=sccache',
                '-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache',
179
                '-DCMAKE_C_COMPILER_LAUNCHER=sccache',
bnellnm's avatar
bnellnm committed
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
216
217
218
219
220
221
222
            ]
        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)

223
        targets = []
bnellnm's avatar
bnellnm committed
224
225
226
        # Build all the extensions
        for ext in self.extensions:
            self.configure(ext)
227
            targets.append(remove_prefix(ext.name, "vllm."))
bnellnm's avatar
bnellnm committed
228

229
        num_jobs, _ = self.compute_num_jobs()
bnellnm's avatar
bnellnm committed
230

231
232
233
234
235
236
        build_args = [
            "--build",
            ".",
            f"-j={num_jobs}",
            *[f"--target={name}" for name in targets],
        ]
bnellnm's avatar
bnellnm committed
237

238
        subprocess.check_call(["cmake", *build_args], cwd=self.build_temp)
bnellnm's avatar
bnellnm committed
239

240

241
def _is_cuda() -> bool:
242
243
244
    has_cuda = torch.version.cuda is not None
    return (VLLM_TARGET_DEVICE == "cuda" and has_cuda
            and not (_is_neuron() or _is_tpu()))
245
246
247


def _is_hip() -> bool:
248
249
    return (VLLM_TARGET_DEVICE == "cuda"
            or VLLM_TARGET_DEVICE == "rocm") and torch.version.hip is not None
250
251


252
253
254
255
def _is_neuron() -> bool:
    torch_neuronx_installed = True
    try:
        subprocess.run(["neuron-ls"], capture_output=True, check=True)
256
    except (FileNotFoundError, PermissionError, subprocess.CalledProcessError):
257
        torch_neuronx_installed = False
258
    return torch_neuronx_installed or VLLM_TARGET_DEVICE == "neuron"
259
260


261
262
263
264
def _is_tpu() -> bool:
    return VLLM_TARGET_DEVICE == "tpu"


265
266
def _is_cpu() -> bool:
    return VLLM_TARGET_DEVICE == "cpu"
267
268


269
270
271
272
def _is_openvino() -> bool:
    return VLLM_TARGET_DEVICE == "openvino"


273
274
275
276
def _is_xpu() -> bool:
    return VLLM_TARGET_DEVICE == "xpu"


277
278
279
280
def _build_custom_ops() -> bool:
    return _is_cuda() or _is_hip() or _is_cpu()


281
282
def _build_core_ext() -> bool:
    return not _is_neuron() and not _is_tpu()
283

284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304

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
305

306

307
308
309
def get_neuronxcc_version():
    import sysconfig
    site_dir = sysconfig.get_paths()["purelib"]
310
311
    version_file = os.path.join(site_dir, "neuronxcc", "version",
                                "__init__.py")
312
313
314
315
316
317
318
319
320
321
322
323
324
325

    # 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
326
def get_nvcc_cuda_version() -> Version:
327
328
329
330
    """Get the CUDA version from nvcc.

    Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
    """
331
    assert CUDA_HOME is not None, "CUDA_HOME is not set"
bnellnm's avatar
bnellnm committed
332
    nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"],
333
334
335
336
337
338
339
                                          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


340
341
342
343
def get_path(*filepath) -> str:
    return os.path.join(ROOT_DIR, *filepath)


344
def find_version(filepath: str) -> str:
345
346
347
348
349
    """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
350
351
        version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
                                  fp.read(), re.M)
352
353
354
355
356
        if version_match:
            return version_match.group(1)
        raise RuntimeError("Unable to find version string.")


zhuwenwen's avatar
zhuwenwen committed
357
358
359
360
361
362
def get_sha(root: Union[str, Path]) -> str:
    try:
        return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=root).decode('ascii').strip()
    except Exception:
        return 'Unknown'

363

zhuwenwen's avatar
zhuwenwen committed
364
365
366
def get_version_add(sha: Optional[str] = None) -> str:
    vllm_root = os.path.dirname(os.path.abspath(__file__))
    add_version_path = os.path.join(os.path.join(vllm_root, "vllm"), "version.py")
367
368
369
370
371
372
373
    if add_git_version:
        if sha != 'Unknown':
            if sha is None:
                sha = get_sha(vllm_root)
            version = 'das.opt1.' + sha[:7]
    else:
        version = 'das.opt1'
zhuwenwen's avatar
zhuwenwen committed
374
375
376
377
378
379
380
381
382
383
384

    # dtk version
    if os.getenv("ROCM_PATH"):
        rocm_path = os.getenv('ROCM_PATH', "")
        rocm_version_path = os.path.join(rocm_path, '.info', "rocm_version")
        with open(rocm_version_path, 'r',encoding='utf-8') as file:
            lines = file.readlines()
        rocm_version=lines[0][:-2].replace(".", "")
        version += ".dtk" + rocm_version

    # torch version
zhuwenwen's avatar
zhuwenwen committed
385
    version += ".torch" + torch.__version__[:5]
zhuwenwen's avatar
zhuwenwen committed
386
    
387
    new_version_content = f"""
388
389
390
391
392
393
import warnings

try:
    import vllm.commit_id
    __commit__ = vllm.commit_id.__commit__
except Exception as e:
zhuwenwen's avatar
zhuwenwen committed
394
    warnings.warn(f"Failed to read commit hash:\\n + str(e)",
395
396
397
398
399
400
401
                  RuntimeWarning,
                  stacklevel=2)
    __commit__ = "COMMIT_HASH_PLACEHOLDER"

__version__ = "0.5.4"
__dcu_version__ = f'0.5.4+{version}' 

402
"""
zhuwenwen's avatar
zhuwenwen committed
403
    
zhuwenwen's avatar
zhuwenwen committed
404
    with open(add_version_path, encoding="utf-8",mode="w") as file:
zhuwenwen's avatar
zhuwenwen committed
405
        file.write(new_version_content)
zhuwenwen's avatar
zhuwenwen committed
406
407
408
409
410
411
412
413
414
415
416
    file.close()
    
    
def get_version():
    get_version_add()
    version_file = 'vllm/version.py'
    with open(version_file, encoding='utf-8') as f:
        exec(compile(f.read(), version_file, 'exec'))
    return locals()['__dcu_version__']


417
def get_vllm_version() -> str:
418
    # version = find_version(get_path("vllm", "version.py"))
419

420
    if _is_cuda():
bnellnm's avatar
bnellnm committed
421
        cuda_version = str(get_nvcc_cuda_version())
422
423
424
425
        if cuda_version != MAIN_CUDA_VERSION:
            cuda_version_str = cuda_version.replace(".", "")[:3]
            version += f"+cu{cuda_version_str}"
    elif _is_hip():
426
        # Get the HIP version
zhuwenwen's avatar
zhuwenwen committed
427
428
429
        # hipcc_version = get_hipcc_rocm_version()
        # if hipcc_version != MAIN_CUDA_VERSION:
        #     rocm_version_str = hipcc_version.replace(".", "")[:3]
zhuwenwen's avatar
zhuwenwen committed
430
431
        #     version += f"+rocm{rocm_version_str}"
        version = get_version()
432
433
    elif _is_neuron():
        # Get the Neuron version
bnellnm's avatar
bnellnm committed
434
        neuron_version = str(get_neuronxcc_version())
435
436
437
        if neuron_version != MAIN_CUDA_VERSION:
            neuron_version_str = neuron_version.replace(".", "")[:3]
            version += f"+neuron{neuron_version_str}"
438
439
    elif _is_openvino():
        version += "+openvino"
440
441
    elif _is_tpu():
        version += "+tpu"
442
443
    elif _is_cpu():
        version += "+cpu"
444
445
    elif _is_xpu():
        version += "+xpu"
446
    else:
447
        raise RuntimeError("Unknown runtime environment")
448

449
450
451
    return version


452
def read_readme() -> str:
Stephen Krider's avatar
Stephen Krider committed
453
454
455
456
457
458
    """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 ""
459
460


461
462
def get_requirements() -> List[str]:
    """Get Python package dependencies from requirements.txt."""
463
464
465

    def _read_requirements(filename: str) -> List[str]:
        with open(get_path(filename)) as f:
466
            requirements = f.read().strip().split("\n")
467
468
469
470
471
472
473
474
475
476
        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")
477
        cuda_major, cuda_minor = torch.version.cuda.split(".")
478
479
        modified_requirements = []
        for req in requirements:
480
481
            if ("vllm-flash-attn" in req
                    and not (cuda_major == "12" and cuda_minor == "1")):
482
483
484
485
                # vllm-flash-attn is built only for CUDA 12.1.
                # Skip for other versions.
                continue
            modified_requirements.append(req)
486
        requirements = modified_requirements
487
    elif _is_hip():
488
        requirements = _read_requirements("requirements-rocm.txt")
489
    elif _is_neuron():
490
        requirements = _read_requirements("requirements-neuron.txt")
491
492
    elif _is_openvino():
        requirements = _read_requirements("requirements-openvino.txt")
493
494
    elif _is_tpu():
        requirements = _read_requirements("requirements-tpu.txt")
495
    elif _is_cpu():
496
        requirements = _read_requirements("requirements-cpu.txt")
497
498
    elif _is_xpu():
        requirements = _read_requirements("requirements-xpu.txt")
499
    else:
500
        raise ValueError(
501
502
            "Unsupported platform, please use CUDA, ROCm, Neuron, "
            "OpenVINO, or CPU.")
503
504
505
    return requirements


bnellnm's avatar
bnellnm committed
506
507
ext_modules = []

508
509
510
if _build_core_ext():
    ext_modules.append(CMakeExtension(name="vllm._core_C"))

511
if _is_cuda() or _is_hip():
bnellnm's avatar
bnellnm committed
512
513
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))

514
if _build_custom_ops():
bnellnm's avatar
bnellnm committed
515
516
    ext_modules.append(CMakeExtension(name="vllm._C"))

517
518
519
package_data = {
    "vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]
}
520
if envs.VLLM_USE_PRECOMPILED:
521
    ext_modules = []
Simon Mo's avatar
Simon Mo committed
522
523
    package_data["vllm"].append("*.so")

bnellnm's avatar
bnellnm committed
524
setup(
Woosuk Kwon's avatar
Woosuk Kwon committed
525
    name="vllm",
526
    version=get_vllm_version(),
Woosuk Kwon's avatar
Woosuk Kwon committed
527
    author="vLLM Team",
528
    license="Apache 2.0",
Woosuk Kwon's avatar
Woosuk Kwon committed
529
530
    description=("A high-throughput and memory-efficient inference and "
                 "serving engine for LLMs"),
531
532
    long_description=read_readme(),
    long_description_content_type="text/markdown",
533
    url="https://github.com/vllm-project/vllm",
534
    project_urls={
535
536
        "Homepage": "https://github.com/vllm-project/vllm",
        "Documentation": "https://vllm.readthedocs.io/en/latest/",
537
538
539
540
541
    },
    classifiers=[
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
Woosuk Kwon's avatar
Woosuk Kwon committed
542
        "Programming Language :: Python :: 3.11",
543
        "Programming Language :: Python :: 3.12",
544
545
546
        "License :: OSI Approved :: Apache Software License",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
bnellnm's avatar
bnellnm committed
547
    packages=find_packages(exclude=("benchmarks", "csrc", "docs", "examples",
548
                                    "tests*")),
549
550
    python_requires=">=3.8",
    install_requires=get_requirements(),
Woosuk Kwon's avatar
Woosuk Kwon committed
551
    ext_modules=ext_modules,
552
    extras_require={
553
        "tensorizer": ["tensorizer>=2.9.0"],
554
    },
555
    cmdclass={"build_ext": cmake_build_ext} if len(ext_modules) > 0 else {},
Simon Mo's avatar
Simon Mo committed
556
    package_data=package_data,
Ethan Xu's avatar
Ethan Xu committed
557
558
559
560
561
    entry_points={
        "console_scripts": [
            "vllm=vllm.scripts:main",
        ],
    },
Woosuk Kwon's avatar
Woosuk Kwon committed
562
)