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

gaoqiong's avatar
gaoqiong committed
21
add_git_version = False
22
23
24
if int(os.environ.get('ADD_GIT_VERSION', '0')) == 1:
    add_git_version = True

25
26
27
28
29
30
31
32
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


33
ROOT_DIR = os.path.dirname(__file__)
34
logger = logging.getLogger(__name__)
35
36
37
38
39
40

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

42
43
44
# vLLM only supports Linux platform
assert sys.platform.startswith(
    "linux"), "vLLM only supports Linux platform (including WSL)."
45

46
47
MAIN_CUDA_VERSION = "12.1"

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

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

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

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

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

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

142
        verbose = envs.VERBOSE
bnellnm's avatar
bnellnm committed
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
194
195
196
        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)

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

203
        num_jobs, _ = self.compute_num_jobs()
bnellnm's avatar
bnellnm committed
204

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

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

214

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


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


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


235
236
def _is_cpu() -> bool:
    return VLLM_TARGET_DEVICE == "cpu"
237
238


bnellnm's avatar
bnellnm committed
239
def _install_punica() -> bool:
240
    return envs.VLLM_INSTALL_PUNICA_KERNELS
241

242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262

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
263

264

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

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

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


298
299
300
301
def get_path(*filepath) -> str:
    return os.path.join(ROOT_DIR, *filepath)


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


315
316
317
318
319
320
321
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'


zhuwenwen's avatar
zhuwenwen committed
322
323
324
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")
325
326
327
328
329
    if add_git_version:
        if sha != 'Unknown':
            if sha is None:
                sha = get_sha(vllm_root)
    else:
zhuwenwen's avatar
zhuwenwen committed
330
        version = 'das.opt1'
zhuwenwen's avatar
zhuwenwen committed
331
332
333
334
335
336
337

    # 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
338
        rocm_version=lines[0].replace(".", "")
zhuwenwen's avatar
zhuwenwen committed
339
340
341
        version += ".dtk" + rocm_version

    with open(add_version_path, encoding="utf-8",mode="w") as file:
zhuwenwen's avatar
zhuwenwen committed
342
343
        file.write("__version__='0.5.0'\n")
        file.write("__dcu_version__='0.5.0+{}'\n".format(version))
zhuwenwen's avatar
zhuwenwen committed
344
345
346
347
348
349
350
351
352
353
354
    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__']


355
356
def get_vllm_version() -> str:
    version = find_version(get_path("vllm", "__init__.py"))
357

358
    if _is_cuda():
bnellnm's avatar
bnellnm committed
359
        cuda_version = str(get_nvcc_cuda_version())
360
361
362
363
        if cuda_version != MAIN_CUDA_VERSION:
            cuda_version_str = cuda_version.replace(".", "")[:3]
            version += f"+cu{cuda_version_str}"
    elif _is_hip():
364
        # Get the HIP version
zhuwenwen's avatar
zhuwenwen committed
365
366
367
        # hipcc_version = get_hipcc_rocm_version()
        # if hipcc_version != MAIN_CUDA_VERSION:
        #     rocm_version_str = hipcc_version.replace(".", "")[:3]
zhuwenwen's avatar
zhuwenwen committed
368
369
        #     version += f"+rocm{rocm_version_str}"
        version = get_version()
370
371
    elif _is_neuron():
        # Get the Neuron version
bnellnm's avatar
bnellnm committed
372
        neuron_version = str(get_neuronxcc_version())
373
374
375
        if neuron_version != MAIN_CUDA_VERSION:
            neuron_version_str = neuron_version.replace(".", "")[:3]
            version += f"+neuron{neuron_version_str}"
376
377
    elif _is_cpu():
        version += "+cpu"
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
    elif _is_cpu():
424
        requirements = _read_requirements("requirements-cpu.txt")
425
    else:
426
        raise ValueError(
427
            "Unsupported platform, please use CUDA, ROCm, Neuron, or CPU.")
428
429
430
    return requirements


bnellnm's avatar
bnellnm committed
431
432
ext_modules = []

433
if _is_cuda() or _is_hip():
bnellnm's avatar
bnellnm committed
434
435
436
437
438
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))

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

439
440
441
    if _install_punica():
        ext_modules.append(CMakeExtension(name="vllm._punica_C"))

442
443
444
package_data = {
    "vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]
}
445
if envs.VLLM_USE_PRECOMPILED:
446
    ext_modules = []
Simon Mo's avatar
Simon Mo committed
447
448
    package_data["vllm"].append("*.so")

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