A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 22:35 UTC
Revert precompile wheel changes (#22055)
Simon Mo committed
Aug 1, 2025 at 01:26 UTC
da31f6ad3dacea8579adfb36d64d28759dc5c095
4 files changed
+107
-137
docker/Dockerfile
+10
-17
@@ -206,7 +206,16 @@ ARG SCCACHE_REGION_NAME=us-west-2
206
ARG SCCACHE_S3_NO_CREDENTIALS=0
207
208
# Flag to control whether to use pre-built vLLM wheels
209
-ARG VLLM_USE_PRECOMPILED=""
209
+ARG VLLM_USE_PRECOMPILED
210
+# TODO: in setup.py VLLM_USE_PRECOMPILED is sensitive to truthiness, it will take =0 as "true", this should be fixed
211
+ENV VLLM_USE_PRECOMPILED=""
212
+RUN if [ "${VLLM_USE_PRECOMPILED}" = "1" ]; then \
213
+ export VLLM_USE_PRECOMPILED=1 && \
214
+ echo "Using precompiled wheels"; \
215
+ else \
216
+ unset VLLM_USE_PRECOMPILED && \
217
+ echo "Leaving VLLM_USE_PRECOMPILED unset to build wheels from source"; \
218
+ fi
219
220
# if USE_SCCACHE is set, use sccache to speed up compilation
221
RUN --mount=type=cache,target=/root/.cache/uv \
@@ -223,8 +232,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
232
&& export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \
233
&& export SCCACHE_IDLE_TIMEOUT=0 \
234
&& export CMAKE_BUILD_TYPE=Release \
226
- && export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" \
227
- && export VLLM_DOCKER_BUILD_CONTEXT=1 \
235
&& sccache --show-stats \
236
&& python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
237
&& sccache --show-stats; \
@@ -238,22 +245,9 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
245
# Clean any existing CMake artifacts
246
rm -rf .deps && \
247
mkdir -p .deps && \
241
- export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" && \
242
- export VLLM_DOCKER_BUILD_CONTEXT=1 && \
248
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
249
fi
250
246
-# When using precompiled wheels, keep only the newest manylinux1 wheel and delete others
247
-RUN if [ "$VLLM_USE_PRECOMPILED" = "1" ]; then \
248
- echo "Cleaning up extra wheels in dist/..." && \
249
- # Identify the most recent manylinux1_x86_64 wheel
250
- KEEP_WHEEL=$(ls -t dist/*manylinux1_x86_64.whl 2>/dev/null | head -n1) && \
251
- if [ -n "$KEEP_WHEEL" ]; then \
252
- echo "Keeping wheel: $KEEP_WHEEL"; \
253
- find dist/ -type f -name "*.whl" ! -path "${KEEP_WHEEL}" -delete; \
254
- fi; \
255
- fi
256
-
251
# Check the size of the wheel if RUN_WHEEL_CHECK is true
252
COPY .buildkite/check-wheel-size.py check-wheel-size.py
253
# sync the default value with .buildkite/check-wheel-size.py
@@ -369,7 +363,6 @@ RUN --mount=type=cache,target=/root/.cache/uv \
363
fi
364
365
# Install vllm wheel first, so that torch etc will be installed.
372
-# !bang
366
RUN --mount=type=bind,from=build,src=/workspace/dist,target=/vllm-workspace/dist \
367
--mount=type=cache,target=/root/.cache/uv \
368
uv pip install --system dist/*.whl --verbose \
requirements/test.txt
+8
-16
@@ -22,7 +22,9 @@ aiohttp==3.10.11
22
aiohttp-cors==0.8.1
23
# via ray
24
aiosignal==1.3.1
25
- # via aiohttp
25
+ # via
26
+ # aiohttp
27
+ # ray
28
albucore==0.0.16
29
# via terratorch
30
albumentations==1.4.6
@@ -137,7 +139,7 @@ contourpy==1.3.0
139
# via matplotlib
140
cramjam==2.9.0
141
# via fastparquet
140
-cupy-cuda12x==13.5.1
142
+cupy-cuda12x==13.3.0
143
# via ray
144
cycler==0.12.1
145
# via matplotlib
@@ -224,6 +226,7 @@ frozenlist==1.5.0
226
# via
227
# aiohttp
228
# aiosignal
229
+ # ray
230
fsspec==2024.9.0
231
# via
232
# datasets
@@ -600,18 +603,10 @@ opencv-python-headless==4.11.0.86
603
opentelemetry-api==1.35.0
604
# via
605
# mlflow-skinny
603
- # opentelemetry-exporter-prometheus
606
# opentelemetry-sdk
607
# opentelemetry-semantic-conventions
606
-opentelemetry-exporter-prometheus==0.56b0
607
- # via ray
608
-opentelemetry-proto==1.36.0
609
- # via ray
608
opentelemetry-sdk==1.35.0
611
- # via
612
- # mlflow-skinny
613
- # opentelemetry-exporter-prometheus
614
- # ray
609
+ # via mlflow-skinny
610
opentelemetry-semantic-conventions==0.56b0
611
# via opentelemetry-sdk
612
packaging==24.2
@@ -702,9 +697,7 @@ pqdm==0.2.0
697
pretrainedmodels==0.7.4
698
# via segmentation-models-pytorch
699
prometheus-client==0.22.0
705
- # via
706
- # opentelemetry-exporter-prometheus
707
- # ray
700
+ # via ray
701
propcache==0.2.0
702
# via yarl
703
proto-plus==1.26.1
@@ -714,7 +707,6 @@ protobuf==5.28.3
707
# google-api-core
708
# googleapis-common-protos
709
# mlflow-skinny
717
- # opentelemetry-proto
710
# proto-plus
711
# ray
712
# tensorboardx
@@ -862,7 +854,7 @@ rasterio==1.4.3
854
# rioxarray
855
# terratorch
856
# torchgeo
865
-ray==2.48.0
857
+ray==2.43.0
858
# via -r requirements/test.in
859
redis==5.2.0
860
# via tensorizer
setup.py
+87
-95
@@ -7,7 +7,6 @@ import json
7
import logging
8
import os
9
import re
10
-import shutil
10
import subprocess
11
import sys
12
from pathlib import Path
@@ -282,69 +281,10 @@ class cmake_build_ext(build_ext):
281
self.copy_file(file, dst_file)
282
283
285
-class precompiled_wheel_utils:
284
+class repackage_wheel(build_ext):
285
"""Extracts libraries and other files from an existing wheel."""
286
288
- @staticmethod
289
- def extract_precompiled_and_patch_package(wheel_url_or_path: str) -> dict:
290
- import tempfile
291
- import zipfile
292
-
293
- temp_dir = None
294
- try:
295
- if not os.path.isfile(wheel_url_or_path):
296
- wheel_filename = wheel_url_or_path.split("/")[-1]
297
- temp_dir = tempfile.mkdtemp(prefix="vllm-wheels")
298
- wheel_path = os.path.join(temp_dir, wheel_filename)
299
- print(f"Downloading wheel from {wheel_url_or_path} "
300
- f"to {wheel_path}")
301
- from urllib.request import urlretrieve
302
- urlretrieve(wheel_url_or_path, filename=wheel_path)
303
- else:
304
- wheel_path = wheel_url_or_path
305
- print(f"Using existing wheel at {wheel_path}")
306
-
307
- package_data_patch = {}
308
-
309
- with zipfile.ZipFile(wheel_path) as wheel:
310
- files_to_copy = [
311
- "vllm/_C.abi3.so",
312
- "vllm/_moe_C.abi3.so",
313
- "vllm/_flashmla_C.abi3.so",
314
- "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
315
- "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
316
- "vllm/cumem_allocator.abi3.so",
317
- ]
318
-
319
- compiled_regex = re.compile(
320
- r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py")
321
- file_members = list(
322
- filter(lambda x: x.filename in files_to_copy,
323
- wheel.filelist))
324
- file_members += list(
325
- filter(lambda x: compiled_regex.match(x.filename),
326
- wheel.filelist))
327
-
328
- for file in file_members:
329
- print(f"[extract] {file.filename}")
330
- target_path = os.path.join(".", file.filename)
331
- os.makedirs(os.path.dirname(target_path), exist_ok=True)
332
- with wheel.open(file.filename) as src, open(
333
- target_path, "wb") as dst:
334
- shutil.copyfileobj(src, dst)
335
-
336
- pkg = os.path.dirname(file.filename).replace("/", ".")
337
- package_data_patch.setdefault(pkg, []).append(
338
- os.path.basename(file.filename))
339
-
340
- return package_data_patch
341
- finally:
342
- if temp_dir is not None:
343
- print(f"Removing temporary directory {temp_dir}")
344
- shutil.rmtree(temp_dir)
345
-
346
- @staticmethod
347
- def get_base_commit_in_main_branch() -> str:
287
+ def get_base_commit_in_main_branch(self) -> str:
288
# Force to use the nightly wheel. This is mainly used for CI testing.
289
if envs.VLLM_TEST_USE_PRECOMPILED_NIGHTLY_WHEEL:
290
return "nightly"
@@ -357,10 +297,6 @@ class precompiled_wheel_utils:
297
]).decode("utf-8")
298
upstream_main_commit = json.loads(resp_json)["sha"]
299
360
- # In Docker build context, .git may be immutable or missing.
361
- if envs.VLLM_DOCKER_BUILD_CONTEXT:
362
- return upstream_main_commit
363
-
300
# Check if the upstream_main_commit exists in the local repo
301
try:
302
subprocess.check_output(
@@ -393,15 +329,92 @@ class precompiled_wheel_utils:
329
"wheel may not be compatible with your dev branch: %s", err)
330
return "nightly"
331
332
+ def run(self) -> None:
333
+ assert _is_cuda(
334
+ ), "VLLM_USE_PRECOMPILED is only supported for CUDA builds"
335
+
336
+ wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION", None)
337
+ if wheel_location is None:
338
+ base_commit = self.get_base_commit_in_main_branch()
339
+ wheel_location = f"https://wheels.vllm.ai/{base_commit}/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl"
340
+ # Fallback to nightly wheel if latest commit wheel is unavailable,
341
+ # in this rare case, the nightly release CI hasn't finished on main.
342
+ if not is_url_available(wheel_location):
343
+ wheel_location = "https://wheels.vllm.ai/nightly/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl"
344
+
345
+ import zipfile
346
+
347
+ if os.path.isfile(wheel_location):
348
+ wheel_path = wheel_location
349
+ print(f"Using existing wheel={wheel_path}")
350
+ else:
351
+ # Download the wheel from a given URL, assume
352
+ # the filename is the last part of the URL
353
+ wheel_filename = wheel_location.split("/")[-1]
354
+
355
+ import tempfile
356
+
357
+ # create a temporary directory to store the wheel
358
+ temp_dir = tempfile.mkdtemp(prefix="vllm-wheels")
359
+ wheel_path = os.path.join(temp_dir, wheel_filename)
360
+
361
+ print(f"Downloading wheel from {wheel_location} to {wheel_path}")
362
+
363
+ from urllib.request import urlretrieve
364
+
365
+ try:
366
+ urlretrieve(wheel_location, filename=wheel_path)
367
+ except Exception as e:
368
+ from setuptools.errors import SetupError
369
+
370
+ raise SetupError(
371
+ f"Failed to get vLLM wheel from {wheel_location}") from e
372
+
373
+ with zipfile.ZipFile(wheel_path) as wheel:
374
+ files_to_copy = [
375
+ "vllm/_C.abi3.so",
376
+ "vllm/_moe_C.abi3.so",
377
+ "vllm/_flashmla_C.abi3.so",
378
+ "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
379
+ "vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
380
+ "vllm/cumem_allocator.abi3.so",
381
+ # "vllm/_version.py", # not available in nightly wheels yet
382
+ ]
383
+
384
+ file_members = list(
385
+ filter(lambda x: x.filename in files_to_copy, wheel.filelist))
386
+
387
+ # vllm_flash_attn python code:
388
+ # Regex from
389
+ # `glob.translate('vllm/vllm_flash_attn/**/*.py', recursive=True)`
390
+ compiled_regex = re.compile(
391
+ r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py")
392
+ file_members += list(
393
+ filter(lambda x: compiled_regex.match(x.filename),
394
+ wheel.filelist))
395
+
396
+ for file in file_members:
397
+ print(f"Extracting and including {file.filename} "
398
+ "from existing wheel")
399
+ package_name = os.path.dirname(file.filename).replace("/", ".")
400
+ file_name = os.path.basename(file.filename)
401
+
402
+ if package_name not in package_data:
403
+ package_data[package_name] = []
404
+
405
+ wheel.extract(file)
406
+ if file_name.endswith(".py"):
407
+ # python files shouldn't be added to package_data
408
+ continue
409
+
410
+ package_data[package_name].append(file_name)
411
+
412
413
def _no_device() -> bool:
414
return VLLM_TARGET_DEVICE == "empty"
415
416
417
def _is_cuda() -> bool:
402
- # Allow forced CUDA in Docker/precompiled builds, even without torch.cuda
403
- if envs.VLLM_USE_PRECOMPILED and envs.VLLM_DOCKER_BUILD_CONTEXT:
404
- return True
418
has_cuda = torch.version.cuda is not None
419
return (VLLM_TARGET_DEVICE == "cuda" and has_cuda
420
and not (_is_neuron() or _is_tpu()))
@@ -626,37 +639,16 @@ package_data = {
639
]
640
}
641
629
-# If using precompiled, extract and patch package_data (in advance of setup)
630
-if envs.VLLM_USE_PRECOMPILED:
631
- assert _is_cuda(), "VLLM_USE_PRECOMPILED is only supported for CUDA builds"
632
- wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION", None)
633
- if wheel_location is not None:
634
- wheel_url = wheel_location
635
- else:
636
- base_commit = precompiled_wheel_utils.get_base_commit_in_main_branch()
637
- wheel_url = f"https://wheels.vllm.ai/{base_commit}/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl"
638
- from urllib.request import urlopen
639
- try:
640
- with urlopen(wheel_url) as resp:
641
- if resp.status != 200:
642
- wheel_url = "https://wheels.vllm.ai/nightly/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl"
643
- except Exception as e:
644
- print(f"[warn] Falling back to nightly wheel: {e}")
645
- wheel_url = "https://wheels.vllm.ai/nightly/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl"
646
-
647
- patch = precompiled_wheel_utils.extract_precompiled_and_patch_package(
648
- wheel_url)
649
- for pkg, files in patch.items():
650
- package_data.setdefault(pkg, []).extend(files)
651
-
642
if _no_device():
643
ext_modules = []
644
655
-if not ext_modules or envs.VLLM_USE_PRECOMPILED:
656
- # Disable build_ext when using precompiled wheel
645
+if not ext_modules:
646
cmdclass = {}
647
else:
659
- cmdclass = {"build_ext": cmake_build_ext}
648
+ cmdclass = {
649
+ "build_ext":
650
+ repackage_wheel if envs.VLLM_USE_PRECOMPILED else cmake_build_ext
651
+ }
652
653
setup(
654
# static metadata should rather go in pyproject.toml
vllm/envs.py
+2
-9
@@ -68,7 +68,6 @@ if TYPE_CHECKING:
68
MAX_JOBS: Optional[str] = None
69
NVCC_THREADS: Optional[str] = None
70
VLLM_USE_PRECOMPILED: bool = False
71
- VLLM_DOCKER_BUILD_CONTEXT: bool = False
71
VLLM_TEST_USE_PRECOMPILED_NIGHTLY_WHEEL: bool = False
72
VLLM_NO_DEPRECATION_WARNING: bool = False
73
VLLM_KEEP_ALIVE_ON_ENGINE_DEATH: bool = False
@@ -228,14 +227,8 @@ environment_variables: dict[str, Callable[[], Any]] = {
227
228
# If set, vllm will use precompiled binaries (*.so)
229
"VLLM_USE_PRECOMPILED":
231
- lambda: os.environ.get("VLLM_USE_PRECOMPILED", "").strip().lower() in
232
- ("1", "true") or bool(os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION")),
233
-
234
- # Used to mark that setup.py is running in a Docker build context,
235
- # in order to force the use of precompiled binaries.
236
- "VLLM_DOCKER_BUILD_CONTEXT":
237
- lambda: os.environ.get("VLLM_DOCKER_BUILD_CONTEXT", "").strip().lower() in
238
- ("1", "true"),
230
+ lambda: bool(os.environ.get("VLLM_USE_PRECOMPILED")) or bool(
231
+ os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION")),
232
233
# Whether to force using nightly wheel in python build.
234
# This is used for testing the nightly wheel in python build.