Detect ROCm wheel variant from environment for precompiled wheels. (#49365)

Signed-off-by: aarushjain29 <Aarushi.Jain2@amd.com> Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: tjtanaa <tunjian.tan@embeddedllm.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
2026-08-13 21:19:20 -07:00
parent 653cc6faca
commit 59d1af5698
2 changed files with 302 additions and 12 deletions
@@ -507,20 +507,69 @@ class precompiled_wheel_utils:

    @staticmethod
    def fetch_metadata_for_variant(
-        commit: str, variant: str | None
+        commit: str,
+        variant: str | None,
+        *,
+        rocm: bool = False,
    ) -> tuple[list[dict], str]:
        """
        Fetches metadata for a specific variant of the precompiled wheel.
+
+        For non-ROCm, fetches vllm metadata.
+
+        For ROCm, discovers all first-level packages and combines their
+        metadata into a single list.
        """
-        variant_dir = f"{variant}/" if variant is not None else ""
-        repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/"
-        meta_url = repo_url + "metadata.json"
-        print(f"Trying to fetch nightly build metadata from {meta_url}")
+        import json
+        from html.parser import HTMLParser
        from urllib.request import urlopen

-        with urlopen(meta_url) as resp:
-            # urlopen raises HTTPError on unexpected status code
-            wheels = json.loads(resp.read().decode("utf-8"))
+        variant_dir = f"{variant}/" if variant is not None else ""
+
+        if not rocm:
+            # Keep original behavior
+            repo_url = f"https://wheels.vllm.ai/{commit}/{variant_dir}vllm/"
+            meta_url = repo_url + "metadata.json"
+            print(f"Trying to fetch nightly build metadata from {meta_url}")
+            with urlopen(meta_url) as resp:
+                wheels = json.loads(resp.read().decode("utf-8"))
+
+            return wheels, repo_url
+
+        # ROCm: discover all packages under the variant directory.
+        repo_url = f"https://wheels.vllm.ai/rocm/{commit}/{variant_dir}"
+
+        class LinkParser(HTMLParser):
+            def __init__(self):
+                super().__init__()
+                self.links: list[str] = []
+
+            def handle_starttag(self, tag, attrs):
+                if tag == "a":
+                    href = dict(attrs).get("href")
+                    if href:
+                        self.links.append(href)
+
+        with urlopen(repo_url) as resp:
+            parser = LinkParser()
+            parser.feed(resp.read().decode("utf-8"))
+
+        packages = [
+            href.rstrip("/")
+            for href in parser.links
+            if href.endswith("/")
+            and href not in ("../", "./", "/")
+            and "/" not in href.rstrip("/")
+        ]
+
+        wheels: list[dict] = []
+
+        for package in packages:
+            meta_url = f"{repo_url}{package}/metadata.json"
+            print(f"Trying to fetch nightly build metadata from {meta_url}")
+            with urlopen(meta_url) as resp:
+                package_wheels = json.loads(resp.read().decode("utf-8"))
+            wheels.extend(package_wheels)
        return wheels, repo_url

    @staticmethod
@@ -582,6 +631,141 @@ class precompiled_wheel_utils:
        print(f"Detected CUDA {cuda_version}, using variant {variant}")
        return variant

+    @staticmethod
+    def rocm_version_to_variant(rocm_version: str) -> str:
+        """Convert a ROCm version string to a wheel variant, e.g. 7.2.3 -> rocm723."""
+        return "rocm" + rocm_version.replace(".", "")
+
+    @staticmethod
+    def detect_system_rocm_variant() -> str | None:
+        """Auto-detect the ROCm wheel variant from the installed ROCm stack."""
+        rocm_version = get_rocm_version()
+        if not rocm_version:
+            try:
+                import torch
+
+                rocm_version = torch.version.hip
+            except Exception:
+                pass
+        if not rocm_version:
+            return None
+        variant = precompiled_wheel_utils.rocm_version_to_variant(rocm_version)
+        print(f"Detected ROCm {rocm_version}, using variant {variant}")
+        return variant
+
+    @staticmethod
+    def fetch_available_rocm_variants(commit: str) -> list[str]:
+        """List ROCm wheel variants published for a commit on wheels.vllm.ai."""
+        from urllib.request import urlopen
+
+        index_url = f"https://wheels.vllm.ai/rocm/{commit}/"
+        print(f"Fetching available ROCm variants from {index_url}")
+        try:
+            with urlopen(index_url) as resp:
+                html = resp.read().decode("utf-8")
+        except Exception as e:
+            logger.warning(
+                "Failed to fetch ROCm variant index for commit %s: %s", commit, e
+            )
+            return []
+        variants = sorted(set(re.findall(r"rocm\d+", html)))
+        print(f"Available ROCm variants for commit {commit}: {variants}")
+        return variants
+
+    @staticmethod
+    def resolve_rocm_wheel_variant(
+        commit: str, variant_override: str | None
+    ) -> str | None:
+        """Resolve a ROCm wheel variant for a commit from wheels.vllm.ai."""
+        env_variant = precompiled_wheel_utils.detect_system_rocm_variant()
+        available = precompiled_wheel_utils.fetch_available_rocm_variants(commit)
+
+        if variant_override is not None:
+            if env_variant and variant_override != env_variant:
+                logger.warning(
+                    "VLLM_PRECOMPILED_WHEEL_VARIANT=%s does not match the "
+                    "detected environment ROCm variant %s; refusing to use a "
+                    "different ROCm patch wheel",
+                    variant_override,
+                    env_variant,
+                )
+                return None
+            if available and variant_override not in available:
+                logger.warning(
+                    "Requested ROCm variant %s is not available for commit %s "
+                    "(available: %s)",
+                    variant_override,
+                    commit,
+                    available,
+                )
+                return None
+            return variant_override
+
+        if env_variant is None:
+            if available:
+                print(
+                    "Could not detect ROCm version from the environment; "
+                    f"available variants for commit {commit}: {available}"
+                )
+            return None
+
+        if env_variant in available:
+            return env_variant
+
+        logger.warning(
+            "Environment ROCm variant %s is not available for commit %s "
+            "(available: %s). Precompiled wheels may not be compatible.",
+            env_variant,
+            commit,
+            available,
+        )
+        return None
+
+    @staticmethod
+    def warn_if_rocm_torch_version_mismatch(
+        wheels: list[dict], repo_url: str, arch: str
+    ) -> None:
+        """Warn if installed torch differs from the custom ROCm build on
+        wheels.vllm.ai and suggest the correct install command."""
+        try:
+            installed = torch.__version__
+        except Exception:
+            return
+
+        def _wheel_version(pkg: str) -> str | None:
+            for w in wheels:
+                if w.get("package_name") == pkg and arch in w.get("platform_tag", ""):
+                    v = w["version"]
+                    if w.get("variant") and "+" not in v:
+                        v = f"{v}+{w['variant']}"
+                    return v
+            return None
+
+        expected = _wheel_version("torch")
+        if expected is None or installed == expected:
+            return
+
+        pkgs = f"torch=={expected}"
+        triton_ver = _wheel_version("triton")
+        if triton_ver:
+            pkgs += f" triton=={triton_ver}"
+
+        logger.warning(
+            "Installed PyTorch %s does not match the custom build %s "
+            "shipped with vLLM ROCm wheels. The ABI may differ from "
+            "official releases. If you hit extension load errors, "
+            "reinstall from the vLLM index:\n"
+            "  pip install %s --extra-index-url %s\n"
+            "  uv pip install %s --extra-index-url %s "
+            "--index-strategy unsafe-best-match",
+            installed,
+            expected,
+            pkgs,
+            repo_url,
+            pkgs,
+            repo_url,
+        )
+
    @staticmethod
    def find_local_rocm_wheel() -> str | None:
        """Search for a local vllm wheel in common locations."""
@@ -639,6 +823,57 @@ class precompiled_wheel_utils:
            print(f"Found local ROCm wheel: {local_wheel}")
            return local_wheel, None

+        import platform
+
+        arch = platform.machine()
+        commit = os.getenv("VLLM_PRECOMPILED_WHEEL_COMMIT", "").lower()
+        if not commit or len(commit) != 40:
+            print(
+                f"VLLM_PRECOMPILED_WHEEL_COMMIT not valid: {commit}"
+                ", trying to fetch base commit in main branch"
+            )
+            commit = precompiled_wheel_utils.get_base_commit_in_main_branch()
+        variant = precompiled_wheel_utils.resolve_rocm_wheel_variant(
+            commit, os.getenv("VLLM_PRECOMPILED_WHEEL_VARIANT", None)
+        )
+        print(f"Using precompiled ROCm wheel commit {commit} with variant {variant}")
+        wheels, repo_url = None, None
+        if variant is not None:
+            try:
+                wheels, repo_url = precompiled_wheel_utils.fetch_metadata_for_variant(
+                    commit, variant, rocm=True
+                )
+                precompiled_wheel_utils.warn_if_rocm_torch_version_mismatch(
+                    wheels, repo_url, arch
+                )
+            except Exception as e:
+                logger.warning(
+                    "Failed to fetch ROCm wheel metadata for variant %s: %s",
+                    variant,
+                    e,
+                )
+        if wheels is not None and repo_url is not None:
+            from urllib.parse import urljoin
+
+            for wheel in wheels:
+                if wheel.get("package_name") == "vllm" and arch in wheel.get(
+                    "platform_tag", ""
+                ):
+                    print(f"Found precompiled wheel metadata: {wheel}")
+                    if "path" not in wheel:
+                        raise ValueError(f"Wheel metadata missing path: {wheel}")
+                    wheel_url = urljoin(f"{repo_url}vllm/", wheel["path"])
+                    download_filename = wheel.get("filename")
+                    print(f"Using precompiled wheel URL: {wheel_url}")
+                    return wheel_url, download_filename
+            logger.warning(
+                "No precompiled vllm wheel found for architecture %s "
+                "from repo %s. All available wheels: %s",
+                arch,
+                repo_url,
+                wheels,
+            )
+
        # Fall back to AMD's PyPI index
        index_url = os.getenv(
            "VLLM_ROCM_WHEEL_INDEX", "https://pypi.amd.com/vllm-rocm/simple"
@@ -669,8 +904,7 @@ class precompiled_wheel_utils:
            print(f"Using user-specified precompiled wheel location: {wheel_location}")
            return wheel_location, None
        else:
-            # ROCm: use local wheel or AMD's PyPI index
-            # TODO: When we have ROCm nightly wheels, we can update this logic.
+            # ROCm: resolve wheels from wheels.vllm.ai with environment-matched variant
            if precompiled_wheel_utils.is_rocm_system():
                return precompiled_wheel_utils.determine_wheel_url_rocm()

@@ -28,8 +28,64 @@ fi
 # test whether the metadata.json url is valid, retry each 3 minutes up to 5 times
 # this avoids cumbersome error messages & manual retries in case the precompiled wheel
 # for the given commit is still being built in the release pipeline
-meta_json_url="https://wheels.vllm.ai/$merge_base_commit/vllm/metadata.json"
-echo "INFO: will use metadata.json from $meta_json_url"
+_vllm_target_lower="$(printf '%s' "${VLLM_TARGET_DEVICE:-}" | tr '[:upper:]' '[:lower:]')"
+if [[ "${_vllm_target_lower}" == "rocm" ]] || [[ -d /opt/rocm ]] || command -v rocminfo >/dev/null 2>&1; then
+    _rocm_env_variant="$(python3 - <<'PY'
+import ctypes
+import os
+from pathlib import Path
+
+def get_rocm_version() -> str | None:
+    rocm_home = os.environ.get("ROCM_HOME") or os.environ.get("ROCM_PATH") or "/opt/rocm"
+    try:
+        librocm_core = Path(rocm_home) / "lib" / "librocm-core.so"
+        if not librocm_core.is_file():
+            return None
+        librocm = ctypes.CDLL(str(librocm_core))
+        get_rocm_core_version = librocm.getROCmVersion
+        major = ctypes.c_uint32()
+        minor = ctypes.c_uint32()
+        patch = ctypes.c_uint32()
+        if get_rocm_core_version(
+            ctypes.byref(major), ctypes.byref(minor), ctypes.byref(patch)
+        ) == 0:
+            return f"{major.value}.{minor.value}.{patch.value}"
+    except Exception:
+        return None
+    return None
+
+version = get_rocm_version()
+if version:
+    print(f"rocm{version.replace('.', '')}", end="")
+PY
+)"
+    _available_variants="$(curl -sf "https://wheels.vllm.ai/rocm/${merge_base_commit}/" \
+        | grep -oP 'rocm\d+' | sort -u | tr '\n' ' ' || true)"
+    if [[ -n "${VLLM_PRECOMPILED_WHEEL_VARIANT:-}" ]]; then
+        _rocm_variant="${VLLM_PRECOMPILED_WHEEL_VARIANT}"
+        if [[ -n "${_rocm_env_variant}" && "${_rocm_variant}" != "${_rocm_env_variant}" ]]; then
+            echo "ERROR: VLLM_PRECOMPILED_WHEEL_VARIANT=${_rocm_variant} does not match detected environment ROCm variant ${_rocm_env_variant}" >&2
+            exit 1
+        fi
+    else
+        _rocm_variant="${_rocm_env_variant}"
+    fi
+    if [[ -z "${_rocm_variant}" ]]; then
+        echo "ERROR: Could not detect ROCm variant from the environment for commit ${merge_base_commit}" >&2
+        exit 1
+    fi
+    if [[ -z "${_available_variants}" ]] \
+            || [[ " ${_available_variants} " != *" ${_rocm_variant} "* ]]; then
+        echo "ERROR: Environment ROCm variant '${_rocm_variant}' is not published for commit ${merge_base_commit} (available:${_available_variants:-none})" >&2
+        exit 1
+    fi
+    meta_json_url="https://wheels.vllm.ai/rocm/${merge_base_commit}/${_rocm_variant}/vllm/metadata.json"
+    unset -v _rocm_env_variant _available_variants _rocm_variant
+else
+    meta_json_url="https://wheels.vllm.ai/${merge_base_commit}/vllm/metadata.json"
+fi
+unset -v _vllm_target_lower
+echo "INFO: will use metadata.json from ${meta_json_url}"

 for i in {1..5}; do
    echo "Checking metadata.json URL (attempt $i)..."