mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2025-12-09 03:15:00 +08:00
For VLLM_USE_PRECOMPILED, only compiled .so files should be extracted (#21964)
This commit is contained in:
parent
56bd537dde
commit
b9b753e7a7
79
setup.py
79
setup.py
@ -371,40 +371,31 @@ class repackage_wheel(build_ext):
|
||||
raise SetupError(
|
||||
f"Failed to get vLLM wheel from {wheel_location}") from e
|
||||
|
||||
# During a docker build: determine correct filename, copy wheel.
|
||||
if envs.VLLM_DOCKER_BUILD_CONTEXT:
|
||||
dist_dir = "/workspace/dist"
|
||||
os.makedirs(dist_dir, exist_ok=True)
|
||||
# Determine correct wheel filename from METADATA
|
||||
with zipfile.ZipFile(wheel_path, "r") as z:
|
||||
metadata_file = next(
|
||||
(n for n in z.namelist()
|
||||
if n.endswith(".dist-info/METADATA")),
|
||||
None,
|
||||
)
|
||||
if not metadata_file:
|
||||
raise RuntimeError(
|
||||
"Could not find METADATA in precompiled wheel.")
|
||||
metadata = z.read(metadata_file).decode()
|
||||
version_line = next((line for line in metadata.splitlines()
|
||||
if line.startswith("Version: ")), None)
|
||||
if not version_line:
|
||||
raise RuntimeError(
|
||||
"Could not determine version from METADATA.")
|
||||
version = version_line.split(": ")[1].strip()
|
||||
# Set the dist_dir for Docker build context
|
||||
dist_dir = ("/workspace/dist"
|
||||
if envs.VLLM_DOCKER_BUILD_CONTEXT else "dist")
|
||||
os.makedirs(dist_dir, exist_ok=True)
|
||||
|
||||
# Build correct filename using internal version
|
||||
arch_tag = "cp38-abi3-manylinux1_x86_64"
|
||||
corrected_wheel_name = f"vllm-{version}-{arch_tag}.whl"
|
||||
final_wheel_path = os.path.join(dist_dir, corrected_wheel_name)
|
||||
|
||||
print(f"Docker build context detected, copying precompiled wheel "
|
||||
f"({version}) to {final_wheel_path}")
|
||||
shutil.copy2(wheel_path, final_wheel_path)
|
||||
return
|
||||
|
||||
# Unzip the wheel when not in Docker context
|
||||
# Extract only necessary compiled .so files from precompiled wheel
|
||||
with zipfile.ZipFile(wheel_path) as wheel:
|
||||
# Get version from METADATA (optional, mostly useful for logging)
|
||||
metadata_file = next((n for n in wheel.namelist()
|
||||
if n.endswith(".dist-info/METADATA")), None)
|
||||
if not metadata_file:
|
||||
raise RuntimeError(
|
||||
"Could not find METADATA in precompiled wheel.")
|
||||
metadata = wheel.read(metadata_file).decode()
|
||||
version_line = next((line for line in metadata.splitlines()
|
||||
if line.startswith("Version: ")), None)
|
||||
if not version_line:
|
||||
raise RuntimeError(
|
||||
"Could not determine version from METADATA.")
|
||||
version = version_line.split(": ")[1].strip()
|
||||
|
||||
print(f"Extracting precompiled kernels from vLLM wheel version: "
|
||||
f"{version}")
|
||||
|
||||
# List of compiled shared objects to extract
|
||||
files_to_copy = [
|
||||
"vllm/_C.abi3.so",
|
||||
"vllm/_moe_C.abi3.so",
|
||||
@ -413,6 +404,7 @@ class repackage_wheel(build_ext):
|
||||
"vllm/vllm_flash_attn/_vllm_fa3_C.abi3.so",
|
||||
"vllm/cumem_allocator.abi3.so",
|
||||
]
|
||||
|
||||
file_members = list(
|
||||
filter(lambda x: x.filename in files_to_copy, wheel.filelist))
|
||||
compiled_regex = re.compile(
|
||||
@ -430,9 +422,26 @@ class repackage_wheel(build_ext):
|
||||
if package_name not in package_data:
|
||||
package_data[package_name] = []
|
||||
|
||||
wheel.extract(file)
|
||||
if not file_name.endswith(".py"):
|
||||
package_data[package_name].append(file_name)
|
||||
output_base = (dist_dir
|
||||
if envs.VLLM_DOCKER_BUILD_CONTEXT else ".")
|
||||
target_path = os.path.join(output_base, file.filename)
|
||||
os.makedirs(os.path.dirname(target_path), exist_ok=True)
|
||||
with wheel.open(file.filename) as src, open(target_path,
|
||||
"wb") as dst:
|
||||
shutil.copyfileobj(src, dst)
|
||||
|
||||
package_data[package_name].append(file_name)
|
||||
|
||||
# Copy wheel into dist dir for Docker to consume (e.g., via --mount)
|
||||
if envs.VLLM_DOCKER_BUILD_CONTEXT:
|
||||
arch_tag = "cp38-abi3-manylinux1_x86_64"
|
||||
corrected_wheel_name = f"vllm-{version}-{arch_tag}.whl"
|
||||
final_wheel_path = os.path.join(dist_dir, corrected_wheel_name)
|
||||
|
||||
print(
|
||||
"Docker build context detected, copying precompiled wheel to "
|
||||
f"{final_wheel_path}")
|
||||
shutil.copy2(wheel_path, final_wheel_path)
|
||||
|
||||
|
||||
def _no_device() -> bool:
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user