setup.py 18.3 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
74
75
76
if not sys.platform.startswith("linux"):
    logger.warning(
        "vLLM only supports Linux platform (including WSL). "
        "Building on %s, "
        "so vLLM may not be able to run correctly", sys.platform)
    VLLM_TARGET_DEVICE = "empty"
77

78
79
MAIN_CUDA_VERSION = "12.1"

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

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

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

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

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

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

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

        if is_sccache_available():
            cmake_args += [
                '-DCMAKE_CXX_COMPILER_LAUNCHER=sccache',
                '-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache',
182
                '-DCMAKE_C_COMPILER_LAUNCHER=sccache',
bnellnm's avatar
bnellnm committed
183
184
185
186
187
188
189
190
191
192
193
            ]
        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)]

194
195
196
197
        # Pass the python path to cmake so it can reuse the build dependencies
        # on subsequent calls to python.
        cmake_args += ['-DVLLM_PYTHON_PATH={}'.format(":".join(sys.path))]

bnellnm's avatar
bnellnm committed
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
223
224
225
226
227
228
229
        #
        # 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)

230
        targets = []
bnellnm's avatar
bnellnm committed
231
232
233
        # Build all the extensions
        for ext in self.extensions:
            self.configure(ext)
234
            targets.append(remove_prefix(ext.name, "vllm."))
bnellnm's avatar
bnellnm committed
235

236
        num_jobs, _ = self.compute_num_jobs()
bnellnm's avatar
bnellnm committed
237

238
239
240
241
242
243
        build_args = [
            "--build",
            ".",
            f"-j={num_jobs}",
            *[f"--target={name}" for name in targets],
        ]
bnellnm's avatar
bnellnm committed
244

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

247

248
249
250
251
def _no_device() -> bool:
    return VLLM_TARGET_DEVICE == "empty"


252
def _is_cuda() -> bool:
253
254
255
    has_cuda = torch.version.cuda is not None
    return (VLLM_TARGET_DEVICE == "cuda" and has_cuda
            and not (_is_neuron() or _is_tpu()))
256
257
258


def _is_hip() -> bool:
259
260
    return (VLLM_TARGET_DEVICE == "cuda"
            or VLLM_TARGET_DEVICE == "rocm") and torch.version.hip is not None
261
262


263
264
265
266
def _is_neuron() -> bool:
    torch_neuronx_installed = True
    try:
        subprocess.run(["neuron-ls"], capture_output=True, check=True)
267
    except (FileNotFoundError, PermissionError, subprocess.CalledProcessError):
268
        torch_neuronx_installed = False
269
    return torch_neuronx_installed or VLLM_TARGET_DEVICE == "neuron"
270
271


272
273
274
275
def _is_tpu() -> bool:
    return VLLM_TARGET_DEVICE == "tpu"


276
277
def _is_cpu() -> bool:
    return VLLM_TARGET_DEVICE == "cpu"
278
279


280
281
282
283
def _is_openvino() -> bool:
    return VLLM_TARGET_DEVICE == "openvino"


284
285
286
287
def _is_xpu() -> bool:
    return VLLM_TARGET_DEVICE == "xpu"


288
289
290
291
def _build_custom_ops() -> bool:
    return _is_cuda() or _is_hip() or _is_cpu()


292
def _build_core_ext() -> bool:
Kunshang Ji's avatar
Kunshang Ji committed
293
    return not (_is_neuron() or _is_tpu() or _is_openvino() or _is_xpu())
294

295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315

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
316

317

318
319
320
def get_neuronxcc_version():
    import sysconfig
    site_dir = sysconfig.get_paths()["purelib"]
321
322
    version_file = os.path.join(site_dir, "neuronxcc", "version",
                                "__init__.py")
323
324
325
326
327
328
329
330
331
332
333
334
335
336

    # 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
337
def get_nvcc_cuda_version() -> Version:
338
339
340
341
    """Get the CUDA version from nvcc.

    Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
    """
342
    assert CUDA_HOME is not None, "CUDA_HOME is not set"
bnellnm's avatar
bnellnm committed
343
    nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"],
344
345
346
347
348
349
350
                                          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


351
352
353
354
def get_path(*filepath) -> str:
    return os.path.join(ROOT_DIR, *filepath)


355
def find_version(filepath: str) -> str:
356
357
358
359
360
    """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
361
362
        version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
                                  fp.read(), re.M)
363
364
365
366
367
        if version_match:
            return version_match.group(1)
        raise RuntimeError("Unable to find version string.")


zhuwenwen's avatar
zhuwenwen committed
368
369
370
371
372
373
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'

374

zhuwenwen's avatar
zhuwenwen committed
375
376
377
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")
378
379
380
381
    if add_git_version:
        if sha != 'Unknown':
            if sha is None:
                sha = get_sha(vllm_root)
zhuwenwen's avatar
zhuwenwen committed
382
            version = 'das.opt1' + sha[:7]
383
    else:
zhuwenwen's avatar
zhuwenwen committed
384
        version = 'das.opt1'
zhuwenwen's avatar
zhuwenwen committed
385
386
387
388
389
390
391

    # 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()
zhuwenwen's avatar
zhuwenwen committed
392
        rocm_version=lines[0].replace(".", "")
zhuwenwen's avatar
zhuwenwen committed
393
        version += ".dtk" + rocm_version
zhuwenwen's avatar
zhuwenwen committed
394
    
395
    new_version_content = f"""
396
397
398
399
400
401
import warnings

try:
    import vllm.commit_id
    __commit__ = vllm.commit_id.__commit__
except Exception as e:
zhuwenwen's avatar
zhuwenwen committed
402
    warnings.warn(f"Failed to read commit hash:\\n + str(e)",
403
404
405
406
                  RuntimeWarning,
                  stacklevel=2)
    __commit__ = "COMMIT_HASH_PLACEHOLDER"

407
408
__version__ = "0.5.5"
__dcu_version__ = f'0.5.5+{version}' 
409

410
"""
zhuwenwen's avatar
zhuwenwen committed
411
    
zhuwenwen's avatar
zhuwenwen committed
412
    with open(add_version_path, encoding="utf-8",mode="w") as file:
zhuwenwen's avatar
zhuwenwen committed
413
        file.write(new_version_content)
zhuwenwen's avatar
zhuwenwen committed
414
415
416
417
418
419
420
421
422
423
424
    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__']


425
def get_vllm_version() -> str:
426
    # version = find_version(get_path("vllm", "version.py"))
427

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

459
460
461
    return version


462
def read_readme() -> str:
Stephen Krider's avatar
Stephen Krider committed
463
464
465
466
467
468
    """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 ""
469
470


471
472
def get_requirements() -> List[str]:
    """Get Python package dependencies from requirements.txt."""
473
474
475

    def _read_requirements(filename: str) -> List[str]:
        with open(get_path(filename)) as f:
476
            requirements = f.read().strip().split("\n")
477
478
479
480
481
482
483
484
        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

485
486
487
    if _no_device():
        requirements = _read_requirements("requirements-cuda.txt")
    elif _is_cuda():
488
        requirements = _read_requirements("requirements-cuda.txt")
489
        cuda_major, cuda_minor = torch.version.cuda.split(".")
490
491
        modified_requirements = []
        for req in requirements:
492
493
            if ("vllm-flash-attn" in req
                    and not (cuda_major == "12" and cuda_minor == "1")):
494
495
496
497
                # vllm-flash-attn is built only for CUDA 12.1.
                # Skip for other versions.
                continue
            modified_requirements.append(req)
498
        requirements = modified_requirements
499
    elif _is_hip():
500
        requirements = _read_requirements("requirements-rocm.txt")
501
    elif _is_neuron():
502
        requirements = _read_requirements("requirements-neuron.txt")
503
504
    elif _is_openvino():
        requirements = _read_requirements("requirements-openvino.txt")
505
506
    elif _is_tpu():
        requirements = _read_requirements("requirements-tpu.txt")
507
    elif _is_cpu():
508
        requirements = _read_requirements("requirements-cpu.txt")
509
510
    elif _is_xpu():
        requirements = _read_requirements("requirements-xpu.txt")
511
    else:
512
        raise ValueError(
513
514
            "Unsupported platform, please use CUDA, ROCm, Neuron, "
            "OpenVINO, or CPU.")
515
516
517
    return requirements


bnellnm's avatar
bnellnm committed
518
519
ext_modules = []

520
521
522
if _build_core_ext():
    ext_modules.append(CMakeExtension(name="vllm._core_C"))

523
if _is_cuda() or _is_hip():
bnellnm's avatar
bnellnm committed
524
525
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))

526
if _build_custom_ops():
bnellnm's avatar
bnellnm committed
527
528
    ext_modules.append(CMakeExtension(name="vllm._C"))

529
530
531
package_data = {
    "vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]
}
532
if envs.VLLM_USE_PRECOMPILED:
533
    ext_modules = []
Simon Mo's avatar
Simon Mo committed
534
535
    package_data["vllm"].append("*.so")

536
537
538
if _no_device():
    ext_modules = []

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