Revert "fix(setup): improve precompiled wheel setup for Docker builds"

[kebe7jun]

Reverts #22025

Fix #22008

[gemini-code-assist]

Code Review

This pull request reverts a previous change that aimed to improve the precompiled wheel setup for Docker builds, in order to fix an issue with those builds. The changes primarily involve setup.py and dependency versions in requirements/test.txt. A potential resource leak has been identified in setup.py.

[gemini-code-assist]

This run method creates a temporary directory using tempfile.mkdtemp when downloading a wheel, but it may not always clean it up. This can lead to temporary files being left on the system after the build process, causing a disk space leak.

To fix this, ensure the temporary directory is cleaned up in a finally block.

    def run(self) -> None:
        assert _is_cuda(
        ), "VLLM_USE_PRECOMPILED is only supported for CUDA builds"

        wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION", None)
        if wheel_location is None:
            base_commit = self.get_base_commit_in_main_branch()
            wheel_location = f"https://wheels.vllm.ai/{base_commit}/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl"
            # Fallback to nightly wheel if latest commit wheel is unavailable,
            # in this rare case, the nightly release CI hasn't finished on main.
            if not is_url_available(wheel_location):
                wheel_location = "https://wheels.vllm.ai/nightly/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl"

        import zipfile
        import tempfile

        temp_dir = None
        try:
            if os.path.isfile(wheel_location):
                wheel_path = wheel_location
                print(f"Using existing wheel={wheel_path}")
            else:
                # Download the wheel from a given URL, assume
                # the filename is the last part of the URL
                wheel_filename = wheel_location.split("/")[-1]

                # create a temporary directory to store the wheel
                temp_dir = tempfile.mkdtemp(prefix="vllm-wheels")
                wheel_path = os.path.join(temp_dir, wheel_filename)
                print(f"Downloading wheel from {wheel_location} to {wheel_path}")
                from urllib.request import urlretrieve
                try:
                    urlretrieve(wheel_location, filename=wheel_path)
                except Exception as e:
                    from setuptools.errors import SetupError
                    raise SetupError(
                        f"Failed to get vLLM wheel from {wheel_location}") from e

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

            # 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",
                    "vllm/_flashmla_C.abi3.so",
                    "vllm/vllm_flash_attn/_vllm_fa2_C.abi3.so",
                    "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(
                    r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\\.)[^/]*\\.py")
                file_members += list(
                    filter(lambda x: compiled_regex.match(x.filename),
                           wheel.filelist))

                for file in file_members:
                    print(f"Extracting and including {file.filename} "
                          "from existing wheel")
                    package_name = os.path.dirname(file.filename).replace("/", ".")
                    file_name = os.path.basename(file.filename)

                    if package_name not in package_data:
                        package_data[package_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)
        finally:
            if temp_dir is not None:
                print(f"Removing temporary directory {temp_dir}")
                shutil.rmtree(temp_dir)

[kebe7jun]

@DarkLight1337 @mgoin PTAL.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[DarkLight1337]

Is this a critical issue? Without that PR, CI might not be run in the latest commit of each PR. I prefer having CI actually do its job even with one failing test

[kebe7jun]

Is this a critical issue? Without that PR, CI might not be run in the latest commit of each PR. I prefer having CI actually do its job even with one failing test

This may not be very friendly for local development scenarios.

[simon-mo]

Will revert more than one commit, thx. #22055