diff --git a/setup.py b/setup.py index 24bff16ff6..87d3dacd4f 100644 --- a/setup.py +++ b/setup.py @@ -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() diff --git a/tests/standalone_tests/python_only_compile.sh b/tests/standalone_tests/python_only_compile.sh index 14f563ae6f..8e3d08d456 100644 --- a/tests/standalone_tests/python_only_compile.sh +++ b/tests/standalone_tests/python_only_compile.sh @@ -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)..."