[Build] Find compatible published wheels for precompiled installs (#58871)

Signed-off-by: greg pereira <grpereir@redhat.com>
Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

.buildkite/test_areas/misc.yaml

@@ -556,6 +556,8 @@ steps:
   - tests/multimodal
   - tests/renderers
   - tests/standalone_tests/lazy_imports.py
+  - tests/standalone_tests/test_precompiled_wheel.py
+  - setup.py
   - tests/reasoning
   - tools/pre_commit/check_test_tethering.py
   - tools/pre_commit/test_tethering_allowlist.txt
@@ -572,6 +574,7 @@ steps:
   commands:
   - pip install cmake # test_cmake_utils.py shells out to cmake; the CPU image has none
   - python3 standalone_tests/lazy_imports.py
+  - python3 standalone_tests/test_precompiled_wheel.py
   - pytest -v -s device_allocator
   - pytest -v -s test_uvicorn_logging.py
   - pytest -v -s test_attention_backend_registry.py

docs/getting_started/installation/gpu.cuda.inc.md

@@ -117,18 +117,21 @@ If you need to recompile the `vllm-rs` Rust frontend binary, you can rebuild and
 
     This will install the required Rust toolchain if needed, build the binary, and place it in `vllm/vllm-rs`.
 
-In case you see an error about wheel not found when running the above command, it might be because the commit you based on in the `main` branch was just merged and its precompiled wheel is not available yet. You can wait around an hour and retry, or set `VLLM_PRECOMPILED_WHEEL_COMMIT=nightly` to automatically select the most recent already-built commit on `main`.
+Wheels are published an hour or two after a commit lands on `main`. If the wheel for
+your merge-base with upstream `main` is not available yet, the installer uses the
+nearest older commit that has one, searching up to 20 commits back. It stops with an
+error if compiled sources or build configuration (`csrc/`, `cmake/`, `CMakeLists.txt`,
+`pyproject.toml`, `vllm/_custom_ops.py`) changed between that commit and your
+merge-base, and warns if `rust/` or `setup.py` did. This is a source-level check, not
+a guarantee of binary compatibility.
 
-```bash
-export VLLM_PRECOMPILED_WHEEL_COMMIT=nightly
-export VLLM_USE_PRECOMPILED=1
-uv pip install --editable .
-```
+Setting `VLLM_PRECOMPILED_WHEEL_LOCATION` or a full SHA in
+`VLLM_PRECOMPILED_WHEEL_COMMIT` skips this search and check.
 
 There are more environment variables to control the behavior of Python-only build:
 
 - `VLLM_PRECOMPILED_WHEEL_LOCATION`: specify the exact wheel URL or local file path of a pre-compiled wheel to use. All other logic to find the wheel will be skipped.
-- `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the commit hash to download the pre-compiled wheel. It can be `nightly` to use the last **already built** commit on the main branch.
+- `VLLM_PRECOMPILED_WHEEL_COMMIT`: override the commit to download the pre-compiled wheel from. Must be a full 40-character SHA.
 - `VLLM_PRECOMPILED_WHEEL_VARIANT`: specify the variant subdirectory to use on the nightly index, e.g., `cu129`, `cu130`, `cpu`. If not specified, the variant is auto-detected based on your system's CUDA version (from PyTorch or nvidia-smi). You can also set `VLLM_MAIN_CUDA_VERSION` to override auto-detection.
 
 You can find more information about vLLM's wheels in [Install the latest code](#install-the-latest-code).

setup.py

@@ -896,6 +896,93 @@ class precompiled_wheel_utils:
         print(f"Using ROCm precompiled wheel: {wheel_url}")
         return wheel_url, download_filename
 
+    # Publication lags main by roughly 5-10 commits; anything much older is
+    # unlikely to still match the checkout's compiled sources.
+    WHEEL_SEARCH_DEPTH = 20
+    # Changes here make an older wheel's extensions unusable.
+    WHEEL_BLOCKING_PATHS = (
+        "csrc",
+        "cmake",
+        "CMakeLists.txt",
+        "pyproject.toml",
+        "vllm/_custom_ops.py",
+    )
+    # Changes here are usually harmless for an older wheel.
+    WHEEL_WARNING_PATHS = ("rust", "setup.py")
+
+    @staticmethod
+    def find_published_wheel_commit(
+        base_commit: str, variant: str | None, arch: str
+    ) -> str:
+        """Find the nearest ancestor of base_commit that has a published wheel.
+
+        Raises:
+            ValueError: If no wheel is published within WHEEL_SEARCH_DEPTH
+                commits, or compiled sources changed since the nearest one.
+
+        """
+        from urllib.error import HTTPError
+
+        utils = precompiled_wheel_utils
+        candidates = subprocess.check_output(
+            [
+                "git",
+                "rev-list",
+                "--first-parent",
+                f"--max-count={utils.WHEEL_SEARCH_DEPTH}",
+                base_commit,
+            ],
+            text=True,
+        ).splitlines()
+        for distance, commit in enumerate(candidates):
+            try:
+                wheels, _ = utils.fetch_metadata_for_variant(commit, variant)
+            except HTTPError as err:
+                if err.code != 404:
+                    raise
+                continue
+            if not any(
+                wheel.get("package_name") == "vllm"
+                and arch in wheel.get("platform_tag", "")
+                for wheel in wheels
+            ):
+                continue
+            if distance == 0:
+                return commit
+
+            diff = ["git", "diff", "--name-only", commit, base_commit, "--"]
+            blocking, warning = (
+                subprocess.check_output([*diff, *paths], text=True).strip()
+                for paths in (utils.WHEEL_BLOCKING_PATHS, utils.WHEEL_WARNING_PATHS)
+            )
+            if blocking:
+                raise ValueError(
+                    f"No precompiled wheel is published yet for {base_commit}, and "
+                    f"the nearest one ({commit}, {distance} commits older) was "
+                    f"built from different compiled sources:\n{blocking}\n"
+                    "Wait for wheel publication, rebase onto a commit with a "
+                    "published wheel, build from source (unset "
+                    "VLLM_USE_PRECOMPILED), or set VLLM_PRECOMPILED_WHEEL_COMMIT "
+                    "to a full commit SHA to skip this check."
+                )
+            if warning:
+                logger.warning(
+                    "Precompiled wheel %s predates changes to:\n%s\n"
+                    "The Rust frontend and packaged files may be stale.",
+                    commit,
+                    warning,
+                )
+            print(
+                f"No precompiled wheel is published yet for {base_commit}; using "
+                f"{commit} ({distance} commits older)"
+            )
+            return commit
+        raise ValueError(
+            f"No published precompiled wheel for variant {variant} and "
+            f"architecture {arch} within {utils.WHEEL_SEARCH_DEPTH} commits of "
+            f"{base_commit}. Wait for wheel publication or build vLLM from source."
+        )
+
     @staticmethod
     def determine_wheel_url() -> tuple[str, str | None]:
         """Try to determine the precompiled wheel URL or path to use.
@@ -906,8 +993,8 @@ class precompiled_wheel_utils:
            or CUDA variant selected from VLLM_MAIN_CUDA_VERSION, torch, or nvidia-smi
 
         If downloading from the nightly repo, the commit can be specified via
-        VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, the head commit in the main branch
-        is used.
+        VLLM_PRECOMPILED_WHEEL_COMMIT; otherwise, local checkouts search backwards
+        from their merge-base with upstream main for a compatible published wheel.
         """
         wheel_location = os.getenv("VLLM_PRECOMPILED_WHEEL_LOCATION", None)
         if wheel_location is not None:
@@ -933,6 +1020,14 @@ class precompiled_wheel_utils:
                     ", trying to fetch base commit in main branch"
                 )
                 commit = precompiled_wheel_utils.get_base_commit_in_main_branch()
+                if (
+                    envs.VLLM_USE_PRECOMPILED
+                    and not envs.VLLM_DOCKER_BUILD_CONTEXT
+                    and commit != "nightly"
+                ):
+                    commit = precompiled_wheel_utils.find_published_wheel_commit(
+                        commit, variant, arch
+                    )
             print(f"Using precompiled wheel commit {commit} with variant {variant}")
             download_filename = None
             try:
@@ -1149,17 +1244,9 @@ class precompiled_wheel_utils:
                     ["git", "fetch", "https://github.com/vllm-project/vllm", "main"]
                 )
 
-            # Then get the commit hash of the current branch that is the same as
-            # the upstream main commit.
-            current_branch = (
-                subprocess.check_output(["git", "branch", "--show-current"])
-                .decode("utf-8")
-                .strip()
-            )
-
             base_commit = (
                 subprocess.check_output(
-                    ["git", "merge-base", f"{upstream_main_commit}", current_branch]
+                    ["git", "merge-base", upstream_main_commit, "HEAD"]
                 )
                 .decode("utf-8")
                 .strip()

tests/standalone_tests/test_precompiled_wheel.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Exercise setup-time wheel selection without importing torch or running setup."""
+
+import ast
+import logging
+import os
+import subprocess
+import tempfile
+import unittest
+from email.message import Message
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import patch
+from urllib.error import HTTPError, URLError
+
+# setup.py executes setup() on import; load only the wheel utility class.
+setup_path = Path(__file__).resolve().parents[2] / "setup.py"
+node = next(
+    node
+    for node in ast.parse(setup_path.read_text()).body
+    if isinstance(node, ast.ClassDef) and node.name == "precompiled_wheel_utils"
+)
+namespace: dict[str, Any] = {
+    "subprocess": subprocess,
+    "os": os,
+    "logger": logging.getLogger("setup"),
+    "envs": SimpleNamespace(VLLM_DOCKER_BUILD_CONTEXT=False, VLLM_USE_PRECOMPILED=True),
+}
+exec(
+    compile(ast.Module(body=[node], type_ignores=[]), str(setup_path), "exec"),
+    namespace,
+)
+wheel_utils = namespace["precompiled_wheel_utils"]
+WHEEL = {
+    "package_name": "vllm",
+    "platform_tag": "manylinux_x86_64",
+    "path": "wheel.whl",
+    "filename": "wheel.whl",
+}
+NOT_FOUND = HTTPError("url", 404, "missing", Message(), None)
+
+
+class TestPublishedWheelSearch(unittest.TestCase):
+    """Walk back from the merge-base in a real Git repo; only the index is mocked."""
+
+    def setUp(self):
+        tmp = tempfile.TemporaryDirectory()
+        self.addCleanup(tmp.cleanup)
+        self.repo = Path(tmp.name)
+        self.addCleanup(os.chdir, os.getcwd())
+        os.chdir(self.repo)
+        self.git("init", "-q")
+        self.published = {self.commit("vllm/a.py")}
+
+    def git(self, *args: str) -> str:
+        identity = ["-c", "user.name=t", "-c", "user.email=t@t"]
+        return subprocess.check_output(["git", *identity, *args], text=True).strip()
+
+    def commit(self, path: str) -> str:
+        file = self.repo / path
+        file.parent.mkdir(parents=True, exist_ok=True)
+        file.write_text(file.read_text() + "x" if file.exists() else "x")
+        self.git("add", "-A")
+        self.git("commit", "-q", "-m", path)
+        return self.git("rev-parse", "HEAD")
+
+    def search(self, base: str, arch: str = "x86_64") -> str:
+        def fetch(commit, variant):
+            if commit not in self.published:
+                raise NOT_FOUND
+            return [WHEEL], "url"
+
+        with patch.object(wheel_utils, "fetch_metadata_for_variant", fetch):
+            return wheel_utils.find_published_wheel_commit(base, "cu130", arch)
+
+    def test_published_base_is_used_despite_local_native_edits(self):
+        base = self.commit("csrc/ops.cpp")
+        self.published.add(base)
+        self.commit("csrc/ops.cpp")
+        (self.repo / "pyproject.toml").write_text("dirty")
+        self.assertEqual(self.search(base), base)
+
+    def test_publication_lag_uses_nearest_published_ancestor(self):
+        published = self.commit("vllm/b.py")
+        self.published.add(published)
+        self.commit("vllm/c.py")
+        self.assertEqual(self.search(self.commit("vllm/d.py")), published)
+
+    def test_compiled_changes_since_wheel_are_rejected(self):
+        (published,) = self.published
+        for path in wheel_utils.WHEEL_BLOCKING_PATHS:
+            with self.subTest(path=path):
+                self.git("reset", "-q", "--hard", published)
+                file = path if "." in path else f"{path}/file"
+                with self.assertRaisesRegex(ValueError, file):
+                    self.search(self.commit(file))
+
+    def test_rust_changes_since_wheel_only_warn(self):
+        (published,) = self.published
+        base = self.commit("rust/src/main.rs")
+        with self.assertLogs("setup", "WARNING") as logs:
+            self.assertEqual(self.search(base), published)
+        self.assertIn("rust/src/main.rs", logs.output[0])
+
+    def test_wrong_architecture_is_not_a_match(self):
+        with self.assertRaisesRegex(ValueError, "No published precompiled wheel"):
+            self.search("HEAD", arch="aarch64")
+
+    def test_search_stops_at_depth_limit(self):
+        with patch.object(wheel_utils, "WHEEL_SEARCH_DEPTH", 2):
+            self.commit("vllm/b.py")
+            with self.assertRaisesRegex(ValueError, "within 2 commits"):
+                self.search(self.commit("vllm/c.py"))
+
+    def test_network_and_server_errors_do_not_trigger_walkback(self):
+        for error in (
+            URLError("offline"),
+            HTTPError("url", 503, "busy", Message(), None),
+        ):
+            with (
+                self.subTest(error=error),
+                patch.object(
+                    wheel_utils, "fetch_metadata_for_variant", side_effect=error
+                ) as fetch,
+                self.assertRaises(type(error)),
+            ):
+                wheel_utils.find_published_wheel_commit("HEAD", "cu130", "x86_64")
+            fetch.assert_called_once()
+
+
+class TestDetermineWheelUrl(unittest.TestCase):
+    def determine(self, base: str, **env: str):
+        with (
+            patch.dict(
+                os.environ,
+                {"VLLM_PRECOMPILED_WHEEL_VARIANT": "cu130", **env},
+                clear=True,
+            ),
+            patch("platform.machine", return_value="x86_64"),
+            patch.object(wheel_utils, "is_rocm_system", return_value=False),
+            patch.object(
+                wheel_utils, "get_base_commit_in_main_branch", return_value=base
+            ),
+            patch.object(
+                wheel_utils,
+                "fetch_metadata_for_variant",
+                return_value=([WHEEL], "https://wheels.test/"),
+            ) as fetch,
+            patch.object(
+                wheel_utils, "find_published_wheel_commit", return_value="older"
+            ) as search,
+        ):
+            url = wheel_utils.determine_wheel_url()
+        return url, fetch, search
+
+    def test_automatic_install_downloads_the_searched_commit(self):
+        url, fetch, search = self.determine("base")
+        self.assertEqual(url, ("https://wheels.test/wheel.whl", "wheel.whl"))
+        search.assert_called_once_with("base", "cu130", "x86_64")
+        fetch.assert_called_once_with("older", "cu130")
+
+    def test_search_is_skipped_when_it_cannot_apply(self):
+        cases = {
+            "nightly fallback": ("nightly", {}, {}),
+            "explicit commit": (
+                "base",
+                {"VLLM_PRECOMPILED_WHEEL_COMMIT": "a" * 40},
+                {},
+            ),
+            "rust only": ("base", {}, {"VLLM_USE_PRECOMPILED": False}),
+            "docker": ("base", {}, {"VLLM_DOCKER_BUILD_CONTEXT": True}),
+        }
+        for name, (base, env, flags) in cases.items():
+            with self.subTest(name), patch.dict(vars(namespace["envs"]), flags):
+                _, fetch, search = self.determine(base, **env)
+                search.assert_not_called()
+                commit = env.get("VLLM_PRECOMPILED_WHEEL_COMMIT", base)
+                fetch.assert_called_once_with(commit, "cu130")
+
+    def test_explicit_wheel_location_bypasses_selection(self):
+        for location in ("/tmp/local.whl", "https://example.com/custom.whl"):
+            with self.subTest(location=location):
+                url, fetch, search = self.determine(
+                    "base", VLLM_PRECOMPILED_WHEEL_LOCATION=location
+                )
+                self.assertEqual(url, (location, None))
+                fetch.assert_not_called()
+                search.assert_not_called()
+
+
+if __name__ == "__main__":
+    unittest.main()