setup.py 16 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
from shutil import which
9
from typing import Dict, List
10

Woosuk Kwon's avatar
Woosuk Kwon committed
11
import torch
12
13
14
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
15
from torch.utils.cpp_extension import CUDA_HOME
16

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

21
22
23
24
25
26
27
28
29

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


30
ROOT_DIR = os.path.dirname(__file__)
31
logger = logging.getLogger(__name__)
32
33
34
35
36
37

# 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
38

39
40
41
# vLLM only supports Linux platform
assert sys.platform.startswith(
    "linux"), "vLLM only supports Linux platform (including WSL)."
42

43
44
MAIN_CUDA_VERSION = "12.1"

bnellnm's avatar
bnellnm committed
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66

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:
67
        super().__init__(name, sources=[], py_limited_api=True, **kwa)
bnellnm's avatar
bnellnm committed
68
69
70
71
72
        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.
73
    did_config: Dict[str, bool] = {}
bnellnm's avatar
bnellnm committed
74
75
76
77
78

    #
    # Determine number of compilation jobs and optionally nvcc compile threads.
    #
    def compute_num_jobs(self):
79
80
        # `num_jobs` is either the value of the MAX_JOBS environment variable
        # (if defined) or the number of CPUs available.
81
        num_jobs = envs.MAX_JOBS
82
83
        if num_jobs is not None:
            num_jobs = int(num_jobs)
84
            logger.info("Using MAX_JOBS=%d as the number of jobs.", num_jobs)
85
86
87
88
89
90
91
        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
92

93
        nvcc_threads = None
94
95
96
97
98
        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.
99
            nvcc_threads = envs.NVCC_THREADS
100
101
            if nvcc_threads is not None:
                nvcc_threads = int(nvcc_threads)
102
103
104
                logger.info(
                    "Using NVCC_THREADS=%d as the number of nvcc threads.",
                    nvcc_threads)
105
106
107
            else:
                nvcc_threads = 1
            num_jobs = max(1, num_jobs // nvcc_threads)
bnellnm's avatar
bnellnm committed
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124

        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"
125
        cfg = envs.CMAKE_BUILD_TYPE or default_cfg
bnellnm's avatar
bnellnm committed
126
127
128
129
130
131
132
133
134
135

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

139
        verbose = envs.VERBOSE
bnellnm's avatar
bnellnm committed
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
        if verbose:
            cmake_args += ['-DCMAKE_VERBOSE_MAKEFILE=ON']

        if is_sccache_available():
            cmake_args += [
                '-DCMAKE_CXX_COMPILER_LAUNCHER=sccache',
                '-DCMAKE_CUDA_COMPILER_LAUNCHER=sccache',
            ]
        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)

194
        targets = []
bnellnm's avatar
bnellnm committed
195
196
197
        # Build all the extensions
        for ext in self.extensions:
            self.configure(ext)
198
            targets.append(remove_prefix(ext.name, "vllm."))
bnellnm's avatar
bnellnm committed
199

200
        num_jobs, _ = self.compute_num_jobs()
bnellnm's avatar
bnellnm committed
201

202
203
204
205
206
207
        build_args = [
            "--build",
            ".",
            f"-j={num_jobs}",
            *[f"--target={name}" for name in targets],
        ]
bnellnm's avatar
bnellnm committed
208

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

211

212
def _is_cuda() -> bool:
213
214
215
    return VLLM_TARGET_DEVICE == "cuda" \
            and torch.version.cuda is not None \
            and not _is_neuron()
216
217
218


def _is_hip() -> bool:
219
220
    return (VLLM_TARGET_DEVICE == "cuda"
            or VLLM_TARGET_DEVICE == "rocm") and torch.version.hip is not None
221
222


223
224
225
226
def _is_neuron() -> bool:
    torch_neuronx_installed = True
    try:
        subprocess.run(["neuron-ls"], capture_output=True, check=True)
227
    except (FileNotFoundError, PermissionError, subprocess.CalledProcessError):
228
        torch_neuronx_installed = False
229
    return torch_neuronx_installed or VLLM_TARGET_DEVICE == "neuron"
230
231


232
233
def _is_cpu() -> bool:
    return VLLM_TARGET_DEVICE == "cpu"
234
235


bnellnm's avatar
bnellnm committed
236
def _install_punica() -> bool:
237
    return envs.VLLM_INSTALL_PUNICA_KERNELS
238

239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259

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
260

261

262
263
264
def get_neuronxcc_version():
    import sysconfig
    site_dir = sysconfig.get_paths()["purelib"]
265
266
    version_file = os.path.join(site_dir, "neuronxcc", "version",
                                "__init__.py")
267
268
269
270
271
272
273
274
275
276
277
278
279
280

    # 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
281
def get_nvcc_cuda_version() -> Version:
282
283
284
285
    """Get the CUDA version from nvcc.

    Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
    """
286
    assert CUDA_HOME is not None, "CUDA_HOME is not set"
bnellnm's avatar
bnellnm committed
287
    nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"],
288
289
290
291
292
293
294
                                          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


295
296
297
298
def get_path(*filepath) -> str:
    return os.path.join(ROOT_DIR, *filepath)


299
def find_version(filepath: str) -> str:
300
301
302
303
304
    """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
305
306
        version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
                                  fp.read(), re.M)
307
308
309
310
311
        if version_match:
            return version_match.group(1)
        raise RuntimeError("Unable to find version string.")


zhuwenwen's avatar
zhuwenwen committed
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
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
335
        version = 'das1.1.git' + sha[:7]
zhuwenwen's avatar
zhuwenwen committed
336
337
338
339
340
341
342
343
344
345
346
347
348
349

    # 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
350
    version += ".torch" + torch.__version__[:5]
zhuwenwen's avatar
zhuwenwen committed
351
352

    with open(add_version_path, encoding="utf-8",mode="w") as file:
zhuwenwen's avatar
zhuwenwen committed
353
354
        file.write("__version__='0.5.0'\n")
        file.write("__dcu_version__='0.5.0+{}'\n".format(version))
zhuwenwen's avatar
zhuwenwen committed
355
356
357
358
359
360
361
362
363
364
365
    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__']


366
367
def get_vllm_version() -> str:
    version = find_version(get_path("vllm", "__init__.py"))
368

369
    if _is_cuda():
bnellnm's avatar
bnellnm committed
370
        cuda_version = str(get_nvcc_cuda_version())
371
372
373
374
        if cuda_version != MAIN_CUDA_VERSION:
            cuda_version_str = cuda_version.replace(".", "")[:3]
            version += f"+cu{cuda_version_str}"
    elif _is_hip():
375
        # Get the HIP version
zhuwenwen's avatar
zhuwenwen committed
376
377
378
        # hipcc_version = get_hipcc_rocm_version()
        # if hipcc_version != MAIN_CUDA_VERSION:
        #     rocm_version_str = hipcc_version.replace(".", "")[:3]
zhuwenwen's avatar
zhuwenwen committed
379
380
        #     version += f"+rocm{rocm_version_str}"
        version = get_version()
381
382
    elif _is_neuron():
        # Get the Neuron version
bnellnm's avatar
bnellnm committed
383
        neuron_version = str(get_neuronxcc_version())
384
385
386
        if neuron_version != MAIN_CUDA_VERSION:
            neuron_version_str = neuron_version.replace(".", "")[:3]
            version += f"+neuron{neuron_version_str}"
387
388
    elif _is_cpu():
        version += "+cpu"
389
    else:
390
        raise RuntimeError("Unknown runtime environment")
391

392
393
394
    return version


395
def read_readme() -> str:
Stephen Krider's avatar
Stephen Krider committed
396
397
398
399
400
401
    """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 ""
402
403


404
405
def get_requirements() -> List[str]:
    """Get Python package dependencies from requirements.txt."""
406
407
408

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


bnellnm's avatar
bnellnm committed
442
443
ext_modules = []

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

if not _is_neuron():
    ext_modules.append(CMakeExtension(name="vllm._C"))

450
451
452
    if _install_punica():
        ext_modules.append(CMakeExtension(name="vllm._punica_C"))

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

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