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, 23:45 UTC
docker: docker-aware precompiled wheel support (#21127)
Signed-off-by: dougbtv <dosmith@redhat.com>
Doug Smith committed
Jul 29, 2025 at 17:45 UTC
a1873db23dd597930a7e4731a53314ace92baf49
3 files changed
+68
-27
docker/Dockerfile
+16
-10
@@ -209,16 +209,7 @@ ARG SCCACHE_REGION_NAME=us-west-2
209
ARG SCCACHE_S3_NO_CREDENTIALS=0
210
211
# Flag to control whether to use pre-built vLLM wheels
212
-ARG VLLM_USE_PRECOMPILED
213
-# TODO: in setup.py VLLM_USE_PRECOMPILED is sensitive to truthiness, it will take =0 as "true", this should be fixed
214
-ENV VLLM_USE_PRECOMPILED=""
215
-RUN if [ "${VLLM_USE_PRECOMPILED}" = "1" ]; then \
216
- export VLLM_USE_PRECOMPILED=1 && \
217
- echo "Using precompiled wheels"; \
218
- else \
219
- unset VLLM_USE_PRECOMPILED && \
220
- echo "Leaving VLLM_USE_PRECOMPILED unset to build wheels from source"; \
221
- fi
212
+ARG VLLM_USE_PRECOMPILED=""
213
214
# if USE_SCCACHE is set, use sccache to speed up compilation
215
RUN --mount=type=cache,target=/root/.cache/uv \
@@ -235,6 +226,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
226
&& export SCCACHE_S3_NO_CREDENTIALS=${SCCACHE_S3_NO_CREDENTIALS} \
227
&& export SCCACHE_IDLE_TIMEOUT=0 \
228
&& export CMAKE_BUILD_TYPE=Release \
229
+ && export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" \
230
+ && export VLLM_DOCKER_BUILD_CONTEXT=1 \
231
&& sccache --show-stats \
232
&& python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38 \
233
&& sccache --show-stats; \
@@ -248,9 +241,22 @@ RUN --mount=type=cache,target=/root/.cache/ccache \
241
# Clean any existing CMake artifacts
242
rm -rf .deps && \
243
mkdir -p .deps && \
244
+ export VLLM_USE_PRECOMPILED="${VLLM_USE_PRECOMPILED}" && \
245
+ export VLLM_DOCKER_BUILD_CONTEXT=1 && \
246
python3 setup.py bdist_wheel --dist-dir=dist --py-limited-api=cp38; \
247
fi
248
249
+# When using precompiled wheels, keep only the newest manylinux1 wheel and delete others
250
+RUN if [ "$VLLM_USE_PRECOMPILED" = "1" ]; then \
251
+ echo "Cleaning up extra wheels in dist/..." && \
252
+ # Identify the most recent manylinux1_x86_64 wheel
253
+ KEEP_WHEEL=$(ls -t dist/*manylinux1_x86_64.whl 2>/dev/null | head -n1) && \
254
+ if [ -n "$KEEP_WHEEL" ]; then \
255
+ echo "Keeping wheel: $KEEP_WHEEL"; \
256
+ find dist/ -type f -name "*.whl" ! -path "${KEEP_WHEEL}" -delete; \
257
+ fi; \
258
+ fi
259
+
260
# Check the size of the wheel if RUN_WHEEL_CHECK is true
261
COPY .buildkite/check-wheel-size.py check-wheel-size.py
262
# sync the default value with .buildkite/check-wheel-size.py
setup.py
+43
-15
@@ -7,6 +7,7 @@ import json
7
import logging
8
import os
9
import re
10
+import shutil
11
import subprocess
12
import sys
13
from pathlib import Path
@@ -297,6 +298,10 @@ class repackage_wheel(build_ext):
298
]).decode("utf-8")
299
upstream_main_commit = json.loads(resp_json)["sha"]
300
301
+ # In Docker build context, .git may be immutable or missing.
302
+ if envs.VLLM_DOCKER_BUILD_CONTEXT:
303
+ return upstream_main_commit
304
+
305
# Check if the upstream_main_commit exists in the local repo
306
try:
307
subprocess.check_output(
@@ -357,19 +362,48 @@ class repackage_wheel(build_ext):
362
# create a temporary directory to store the wheel
363
temp_dir = tempfile.mkdtemp(prefix="vllm-wheels")
364
wheel_path = os.path.join(temp_dir, wheel_filename)
360
-
365
print(f"Downloading wheel from {wheel_location} to {wheel_path}")
362
-
366
from urllib.request import urlretrieve
364
-
367
try:
368
urlretrieve(wheel_location, filename=wheel_path)
369
except Exception as e:
370
from setuptools.errors import SetupError
369
-
371
raise SetupError(
372
f"Failed to get vLLM wheel from {wheel_location}") from e
373
374
+ # During a docker build: determine correct filename, copy wheel.
375
+ if envs.VLLM_DOCKER_BUILD_CONTEXT:
376
+ dist_dir = "/workspace/dist"
377
+ os.makedirs(dist_dir, exist_ok=True)
378
+ # Determine correct wheel filename from METADATA
379
+ with zipfile.ZipFile(wheel_path, "r") as z:
380
+ metadata_file = next(
381
+ (n for n in z.namelist()
382
+ if n.endswith(".dist-info/METADATA")),
383
+ None,
384
+ )
385
+ if not metadata_file:
386
+ raise RuntimeError(
387
+ "Could not find METADATA in precompiled wheel.")
388
+ metadata = z.read(metadata_file).decode()
389
+ version_line = next((line for line in metadata.splitlines()
390
+ if line.startswith("Version: ")), None)
391
+ if not version_line:
392
+ raise RuntimeError(
393
+ "Could not determine version from METADATA.")
394
+ version = version_line.split(": ")[1].strip()
395
+
396
+ # Build correct filename using internal version
397
+ arch_tag = "cp38-abi3-manylinux1_x86_64"
398
+ corrected_wheel_name = f"vllm-{version}-{arch_tag}.whl"
399
+ final_wheel_path = os.path.join(dist_dir, corrected_wheel_name)
400
+
401
+ print(f"Docker build context detected, copying precompiled wheel "
402
+ f"({version}) to {final_wheel_path}")
403
+ shutil.copy2(wheel_path, final_wheel_path)
404
+ return
405
+
406
+ # Unzip the wheel when not in Docker context
407
with zipfile.ZipFile(wheel_path) as wheel:
408
files_to_copy = [
409
"vllm/_C.abi3.so",
@@ -378,15 +412,9 @@ class repackage_wheel(build_ext):
412
"vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
413
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
414
"vllm/cumem_allocator.abi3.so",
381
- # "vllm/_version.py", # not available in nightly wheels yet
415
]
383
-
416
file_members = list(
417
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)`
418
compiled_regex = re.compile(
419
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py")
420
file_members += list(
@@ -403,11 +431,8 @@ class repackage_wheel(build_ext):
431
package_data[package_name] = []
432
433
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)
434
+ if not file_name.endswith(".py"):
435
+ package_data[package_name].append(file_name)
436
437
438
def _no_device() -> bool:
@@ -415,6 +440,9 @@ def _no_device() -> bool:
440
441
442
def _is_cuda() -> bool:
443
+ # Allow forced CUDA in Docker/precompiled builds, even without torch.cuda
444
+ if envs.VLLM_USE_PRECOMPILED and envs.VLLM_DOCKER_BUILD_CONTEXT:
445
+ return True
446
has_cuda = torch.version.cuda is not None
447
return (VLLM_TARGET_DEVICE == "cuda" and has_cuda
448
and not (_is_neuron() or _is_tpu()))
vllm/envs.py
+9
-2
@@ -68,6 +68,7 @@ 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
72
VLLM_TEST_USE_PRECOMPILED_NIGHTLY_WHEEL: bool = False
73
VLLM_NO_DEPRECATION_WARNING: bool = False
74
VLLM_KEEP_ALIVE_ON_ENGINE_DEATH: bool = False
@@ -222,8 +223,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
223
224
# If set, vllm will use precompiled binaries (*.so)
225
"VLLM_USE_PRECOMPILED":
225
- lambda: bool(os.environ.get("VLLM_USE_PRECOMPILED")) or bool(
226
- os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION")),
226
+ lambda: os.environ.get("VLLM_USE_PRECOMPILED", "").strip().lower() in
227
+ ("1", "true") or bool(os.environ.get("VLLM_PRECOMPILED_WHEEL_LOCATION")),
228
+
229
+ # Used to mark that setup.py is running in a Docker build context,
230
+ # in order to force the use of precompiled binaries.
231
+ "VLLM_DOCKER_BUILD_CONTEXT":
232
+ lambda: os.environ.get("VLLM_DOCKER_BUILD_CONTEXT", "").strip().lower() in
233
+ ("1", "true"),
234
235
# Whether to force using nightly wheel in python build.
236
# This is used for testing the nightly wheel in python build.