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
For VLLM_USE_PRECOMPILED, only compiled .so files should be extracted (#21964)
Doug Smith committed
Jul 30, 2025 at 16:04 UTC
b9b753e7a7d95311186bbfc2b30b643a2f9e6ca1
1 file changed
+44
-35
setup.py
+44
-35
@@ -371,40 +371,31 @@ class repackage_wheel(build_ext):
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)
374
+ # Set the dist_dir for Docker build context
375
+ dist_dir = ("/workspace/dist"
376
+ if envs.VLLM_DOCKER_BUILD_CONTEXT else "dist")
377
+ os.makedirs(dist_dir, exist_ok=True)
378
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
379
+ # Extract only necessary compiled .so files from precompiled wheel
380
with zipfile.ZipFile(wheel_path) as wheel:
381
+ # Get version from METADATA (optional, mostly useful for logging)
382
+ metadata_file = next((n for n in wheel.namelist()
383
+ if n.endswith(".dist-info/METADATA")), None)
384
+ if not metadata_file:
385
+ raise RuntimeError(
386
+ "Could not find METADATA in precompiled wheel.")
387
+ metadata = wheel.read(metadata_file).decode()
388
+ version_line = next((line for line in metadata.splitlines()
389
+ if line.startswith("Version: ")), None)
390
+ if not version_line:
391
+ raise RuntimeError(
392
+ "Could not determine version from METADATA.")
393
+ version = version_line.split(": ")[1].strip()
394
+
395
+ print(f"Extracting precompiled kernels from vLLM wheel version: "
396
+ f"{version}")
397
+
398
+ # List of compiled shared objects to extract
399
files_to_copy = [
400
"vllm/_C.abi3.so",
401
"vllm/_moe_C.abi3.so",
@@ -413,6 +404,7 @@ class repackage_wheel(build_ext):
404
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
405
"vllm/cumem_allocator.abi3.so",
406
]
407
+
408
file_members = list(
409
filter(lambda x: x.filename in files_to_copy, wheel.filelist))
410
compiled_regex = re.compile(
@@ -430,9 +422,26 @@ class repackage_wheel(build_ext):
422
if package_name not in package_data:
423
package_data[package_name] = []
424
433
- wheel.extract(file)
434
- if not file_name.endswith(".py"):
435
- package_data[package_name].append(file_name)
425
+ output_base = (dist_dir
426
+ if envs.VLLM_DOCKER_BUILD_CONTEXT else ".")
427
+ target_path = os.path.join(output_base, file.filename)
428
+ os.makedirs(os.path.dirname(target_path), exist_ok=True)
429
+ with wheel.open(file.filename) as src, open(target_path,
430
+ "wb") as dst:
431
+ shutil.copyfileobj(src, dst)
432
+
433
+ package_data[package_name].append(file_name)
434
+
435
+ # Copy wheel into dist dir for Docker to consume (e.g., via --mount)
436
+ if envs.VLLM_DOCKER_BUILD_CONTEXT:
437
+ arch_tag = "cp38-abi3-manylinux1_x86_64"
438
+ corrected_wheel_name = f"vllm-{version}-{arch_tag}.whl"
439
+ final_wheel_path = os.path.join(dist_dir, corrected_wheel_name)
440
+
441
+ print(
442
+ "Docker build context detected, copying precompiled wheel to "
443
+ f"{final_wheel_path}")
444
+ shutil.copy2(wheel_path, final_wheel_path)
445
446
447
def _no_device() -> bool: