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

22
23
24
25
26
27
28
29
30

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


31
ROOT_DIR = os.path.dirname(__file__)
32
logger = logging.getLogger(__name__)
33

34
35
36

def embed_commit_hash():
    try:
youkaichao's avatar
youkaichao committed
37
38
39
40
41
42
43
        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()

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

62
63
64
65
66
# 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
67

68
69
70
# vLLM only supports Linux platform
assert sys.platform.startswith(
    "linux"), "vLLM only supports Linux platform (including WSL)."
71

72
73
MAIN_CUDA_VERSION = "12.1"

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

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

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

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

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

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

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

        if is_sccache_available():
            cmake_args += [
                '-DCMAKE_CXX_COMPILER_LAUNCHER=sccache',
                '-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache',
176
                '-DCMAKE_C_COMPILER_LAUNCHER=sccache',
bnellnm's avatar
bnellnm committed
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
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)]

        if _install_punica():
            cmake_args += ['-DVLLM_INSTALL_PUNICA_KERNELS=ON']

        #
        # 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()


bnellnm's avatar
bnellnm committed
281
def _install_punica() -> bool:
282
    return envs.VLLM_INSTALL_PUNICA_KERNELS
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
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
def get_abi():
    try:
        command = "echo '#include <string>' | gcc -x c++ -E -dM - | fgrep _GLIBCXX_USE_CXX11_ABI" 
        result = subprocess.run(command, shell=True, capture_output=True, text=True) 
        output = result.stdout.strip() 
        abi = "abi" + output.split(" ")[-1]
        return abi
    except Exception:
        return 'abiUnknown'


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'

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")
    if sha != 'Unknown':
        if sha is None:
            sha = get_sha(vllm_root)
zhuwenwen's avatar
zhuwenwen committed
380
        version = 'das1.1.git' + sha[:7]
zhuwenwen's avatar
zhuwenwen committed
381
382
383
384
385
386
387
388
389
390
391
392
393
394

    # abi version
    version += "." + get_abi()

    # 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
395
    version += ".torch" + torch.__version__[:5]
zhuwenwen's avatar
zhuwenwen committed
396
    
397
398
    new_version_content = f"""
import warnings  
zhuwenwen's avatar
zhuwenwen committed
399
    
400
401
402
403
404
405
406
407
408
409
410
411
try:  
    import vllm.commit_id  
    __commit__ = vllm.commit_id.__commit__  
except Exception as e:  
    warnings.warn(f"Failed to read commit hash:\\n + str(e)",  
                RuntimeWarning,  
                stacklevel=2)  
    __commit__ = "COMMIT_HASH_PLACEHOLDER"  

__version__ = "0.5.3.post1"  
__dcu_version__ = f'0.5.3.post1+{version}'
"""
zhuwenwen's avatar
zhuwenwen committed
412
    
zhuwenwen's avatar
zhuwenwen committed
413
    with open(add_version_path, encoding="utf-8",mode="w") as file:
zhuwenwen's avatar
zhuwenwen committed
414
        file.write(new_version_content)
zhuwenwen's avatar
zhuwenwen committed
415
416
417
418
419
420
421
422
423
424
425
    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__']


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

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

458
459
460
    return version


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


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

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


bnellnm's avatar
bnellnm committed
515
516
ext_modules = []

517
if _is_cuda() or _is_hip():
bnellnm's avatar
bnellnm committed
518
519
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))

520
if _build_custom_ops():
bnellnm's avatar
bnellnm committed
521
522
    ext_modules.append(CMakeExtension(name="vllm._C"))

523
524
525
    if _install_punica():
        ext_modules.append(CMakeExtension(name="vllm._punica_C"))

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

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