setup.py 12.1 KB
Newer Older
1
import io
2
import logging
3
4
import os
import re
5
import subprocess
bnellnm's avatar
bnellnm committed
6
import sys
7
from shutil import which
bnellnm's avatar
bnellnm committed
8
from typing import List
9

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

ROOT_DIR = os.path.dirname(__file__)
17
logger = logging.getLogger(__name__)
18

19
20
21
22
# vLLM only supports Linux platform
assert sys.platform.startswith(
    "linux"), "vLLM only supports Linux platform (including WSL)."

23
24
MAIN_CUDA_VERSION = "12.1"

bnellnm's avatar
bnellnm committed
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

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:
        super().__init__(name, sources=[], **kwa)
        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.
    did_config = {}

    #
    # Determine number of compilation jobs and optionally nvcc compile threads.
    #
    def compute_num_jobs(self):
59
60
61
62
63
64
65
66
67
68
69
        num_jobs = os.environ.get("MAX_JOBS", None)
        if num_jobs is not None:
            num_jobs = int(num_jobs)
            logger.info(f"Using MAX_JOBS={num_jobs} as the number of jobs.")
        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
70

71
72
73
74
75
76
        nvcc_threads = None
        if _is_cuda():
            nvcc_cuda_version = get_nvcc_cuda_version()
            if nvcc_cuda_version >= Version("11.2"):
                nvcc_threads = int(os.getenv("NVCC_THREADS", 8))
                num_jobs = max(1, round(num_jobs / (nvcc_threads / 4)))
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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
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

        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"
        cfg = os.getenv("CMAKE_BUILD_TYPE", default_cfg)

        # 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),
        ]

        verbose = bool(int(os.getenv('VERBOSE', '0')))
        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)

        # Build all the extensions
        for ext in self.extensions:
            self.configure(ext)

            ext_target_name = remove_prefix(ext.name, "vllm.")
            num_jobs, _ = self.compute_num_jobs()

            build_args = [
                '--build', '.', '--target', ext_target_name, '-j',
                str(num_jobs)
            ]

            subprocess.check_call(['cmake', *build_args], cwd=self.build_temp)
175
176


177
def _is_cuda() -> bool:
178
    return torch.version.cuda is not None and not _is_neuron()
179
180


181
182
183
184
def _is_hip() -> bool:
    return torch.version.hip is not None


185
186
187
188
def _is_neuron() -> bool:
    torch_neuronx_installed = True
    try:
        subprocess.run(["neuron-ls"], capture_output=True, check=True)
189
    except (FileNotFoundError, PermissionError, subprocess.CalledProcessError):
190
191
192
193
        torch_neuronx_installed = False
    return torch_neuronx_installed


bnellnm's avatar
bnellnm committed
194
195
def _install_punica() -> bool:
    return bool(int(os.getenv("VLLM_INSTALL_PUNICA_KERNELS", "0")))
196

197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217

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
218

219

220
221
222
def get_neuronxcc_version():
    import sysconfig
    site_dir = sysconfig.get_paths()["purelib"]
223
224
    version_file = os.path.join(site_dir, "neuronxcc", "version",
                                "__init__.py")
225
226
227
228
229
230
231
232
233
234
235
236
237
238

    # 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
239
def get_nvcc_cuda_version() -> Version:
240
241
242
243
    """Get the CUDA version from nvcc.

    Adapted from https://github.com/NVIDIA/apex/blob/8b7a1ff183741dd8f9b87e7bafd04cfde99cea28/setup.py
    """
bnellnm's avatar
bnellnm committed
244
    nvcc_output = subprocess.check_output([CUDA_HOME + "/bin/nvcc", "-V"],
245
246
247
248
249
250
251
                                          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


252
253
254
255
def get_path(*filepath) -> str:
    return os.path.join(ROOT_DIR, *filepath)


256
def find_version(filepath: str) -> str:
257
258
259
260
261
    """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
262
263
        version_match = re.search(r"^__version__ = ['\"]([^'\"]*)['\"]",
                                  fp.read(), re.M)
264
265
266
267
268
        if version_match:
            return version_match.group(1)
        raise RuntimeError("Unable to find version string.")


269
270
def get_vllm_version() -> str:
    version = find_version(get_path("vllm", "__init__.py"))
271

272
    if _is_cuda():
bnellnm's avatar
bnellnm committed
273
        cuda_version = str(get_nvcc_cuda_version())
274
275
276
277
        if cuda_version != MAIN_CUDA_VERSION:
            cuda_version_str = cuda_version.replace(".", "")[:3]
            version += f"+cu{cuda_version_str}"
    elif _is_hip():
278
279
280
281
282
        # Get the HIP version
        hipcc_version = get_hipcc_rocm_version()
        if hipcc_version != MAIN_CUDA_VERSION:
            rocm_version_str = hipcc_version.replace(".", "")[:3]
            version += f"+rocm{rocm_version_str}"
283
284
    elif _is_neuron():
        # Get the Neuron version
bnellnm's avatar
bnellnm committed
285
        neuron_version = str(get_neuronxcc_version())
286
287
288
        if neuron_version != MAIN_CUDA_VERSION:
            neuron_version_str = neuron_version.replace(".", "")[:3]
            version += f"+neuron{neuron_version_str}"
289
    else:
290
        raise RuntimeError("Unknown runtime environment")
291

292
293
294
    return version


295
def read_readme() -> str:
Stephen Krider's avatar
Stephen Krider committed
296
297
298
299
300
301
    """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 ""
302
303


304
305
def get_requirements() -> List[str]:
    """Get Python package dependencies from requirements.txt."""
306
    if _is_cuda():
307
308
        with open(get_path("requirements.txt")) as f:
            requirements = f.read().strip().split("\n")
bnellnm's avatar
bnellnm committed
309
        if get_nvcc_cuda_version() <= Version("11.8"):
310
311
312
313
314
            # replace cupy-cuda12x with cupy-cuda11x for cuda 11.x
            for i in range(len(requirements)):
                if requirements[i].startswith("cupy-cuda12x"):
                    requirements[i] = "cupy-cuda11x"
                    break
315
316
317
318
319
320
321
322
323
324
    elif _is_hip():
        with open(get_path("requirements-rocm.txt")) as f:
            requirements = f.read().strip().split("\n")
    elif _is_neuron():
        with open(get_path("requirements-neuron.txt")) as f:
            requirements = f.read().strip().split("\n")
    else:
        raise ValueError(
            "Unsupported platform, please use CUDA, ROCM or Neuron.")

325
326
327
    return requirements


bnellnm's avatar
bnellnm committed
328
329
330
331
332
333
334
335
336
337
338
ext_modules = []

if _is_cuda():
    ext_modules.append(CMakeExtension(name="vllm._moe_C"))

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

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

339
340
341
package_data = {
    "vllm": ["py.typed", "model_executor/layers/fused_moe/configs/*.json"]
}
Simon Mo's avatar
Simon Mo committed
342
343
344
if os.environ.get("VLLM_USE_PRECOMPILED"):
    package_data["vllm"].append("*.so")

bnellnm's avatar
bnellnm committed
345
setup(
Woosuk Kwon's avatar
Woosuk Kwon committed
346
    name="vllm",
347
    version=get_vllm_version(),
Woosuk Kwon's avatar
Woosuk Kwon committed
348
    author="vLLM Team",
349
    license="Apache 2.0",
Woosuk Kwon's avatar
Woosuk Kwon committed
350
351
    description=("A high-throughput and memory-efficient inference and "
                 "serving engine for LLMs"),
352
353
    long_description=read_readme(),
    long_description_content_type="text/markdown",
354
    url="https://github.com/vllm-project/vllm",
355
    project_urls={
356
357
        "Homepage": "https://github.com/vllm-project/vllm",
        "Documentation": "https://vllm.readthedocs.io/en/latest/",
358
359
360
361
362
    },
    classifiers=[
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
Woosuk Kwon's avatar
Woosuk Kwon committed
363
        "Programming Language :: Python :: 3.11",
364
365
366
        "License :: OSI Approved :: Apache Software License",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
    ],
bnellnm's avatar
bnellnm committed
367
368
    packages=find_packages(exclude=("benchmarks", "csrc", "docs", "examples",
                                    "tests")),
369
370
    python_requires=">=3.8",
    install_requires=get_requirements(),
Woosuk Kwon's avatar
Woosuk Kwon committed
371
    ext_modules=ext_modules,
bnellnm's avatar
bnellnm committed
372
    cmdclass={"build_ext": cmake_build_ext} if not _is_neuron() else {},
Simon Mo's avatar
Simon Mo committed
373
    package_data=package_data,
Woosuk Kwon's avatar
Woosuk Kwon committed
374
)