From 9e33c25eac8368eb19e1866b0ecf1d5e787207db Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 09:25:56 +0000 Subject: [PATCH] Require actual AMD GPU presence before selecting ROCm paths All 8 reviewers across 2 cycles independently flagged that ROCm detection used toolkit/filesystem hints (hipcc, /opt/rocm, rocm-core) as a proxy for GPU presence, which would misroute CPU-only or NVIDIA hosts that happen to have ROCm tools installed. Now all 3 detection points (install.sh, install_python_stack.py, install_llama_prebuilt.py) probe for an actual AMD GPU before entering the ROCm path: - install.sh: check rocminfo for gfx* GPU names, or amd-smi list for device rows, before version detection - install_python_stack.py: new _has_rocm_gpu() function probes rocminfo and amd-smi list before _ensure_rocm_torch() proceeds - install_llama_prebuilt.py: detect_host() probes rocminfo/amd-smi list instead of just checking tool existence or directory paths Also: - Shell test mock amd-smi now handles "list" subcommand - Python tests updated to mock _has_rocm_gpu where needed - Added test_no_gpu_with_rocm_tools_skips to verify the new guard - Test index lookups now use sorted() to match production code --- install.sh | 15 +++++++- studio/install_llama_prebuilt.py | 25 ++++++++----- studio/install_python_stack.py | 32 +++++++++++++++- tests/sh/test_get_torch_index_url.sh | 11 +++++- tests/studio/install/test_rocm_support.py | 45 ++++++++++++++++++----- 5 files changed, 105 insertions(+), 23 deletions(-) diff --git a/install.sh b/install.sh index 85268b14a4..4d121bf331 100755 --- a/install.sh +++ b/install.sh @@ -983,7 +983,20 @@ get_torch_index_url() { _smi="/usr/bin/nvidia-smi" fi if [ -z "$_smi" ]; then - # No NVIDIA GPU -- check for AMD ROCm + # No NVIDIA GPU -- check for AMD ROCm GPU + # First confirm an actual AMD GPU is present (not just ROCm tools installed) + _has_rocm_gpu=false + if command -v rocminfo >/dev/null 2>&1 && \ + rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[0-9]/{found=1} END{exit !found}'; then + _has_rocm_gpu=true + elif command -v amd-smi >/dev/null 2>&1 && \ + amd-smi list 2>/dev/null | awk 'NR>1 && NF{found=1} END{exit !found}'; then + _has_rocm_gpu=true + fi + if [ "$_has_rocm_gpu" != true ]; then + echo "$_base/cpu"; return + fi + # AMD GPU confirmed -- detect ROCm version _rocm_tag="" _rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \ amd-smi version 2>/dev/null | awk -F'ROCm version: ' \ diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 43f3362b81..3b46373eb0 100755 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -1431,17 +1431,24 @@ def detect_host() -> HostInfo: except Exception: pass - # Detect AMD ROCm (HIP) + # Detect AMD ROCm (HIP) -- require actual GPU, not just tools installed has_rocm = False if not is_macos: - rocm_hints = [ - shutil.which("hipcc"), - shutil.which("amd-smi"), - shutil.which("rocm-smi"), - ] - rocm_paths = [p for p in ("/opt/rocm", os.environ.get("ROCM_PATH")) if p] - if any(rocm_hints) or any(os.path.isdir(p) for p in rocm_paths): - has_rocm = True + for _cmd, _marker in ( + (["rocminfo"], "gfx"), + (["amd-smi", "list"], None), + ): + _exe = shutil.which(_cmd[0]) + if not _exe: + continue + try: + _result = run_capture([_exe, *_cmd[1:]], timeout = 10) + except Exception: + continue + if _result.returncode == 0 and _result.stdout.strip(): + if _marker is None or _marker in _result.stdout.lower(): + has_rocm = True + break return HostInfo( system = system, diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 16b9877da9..48d95c9b03 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -80,16 +80,44 @@ def _detect_rocm_version() -> tuple[int, int] | None: return None +def _has_rocm_gpu() -> bool: + """Return True only if an actual AMD GPU is visible (not just ROCm tools installed).""" + for cmd, marker in ( + (["rocminfo"], "gfx"), + (["amd-smi", "list"], None), + ): + exe = shutil.which(cmd[0]) + if not exe: + continue + try: + result = subprocess.run( + [exe, *cmd[1:]], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + text = True, + timeout = 10, + ) + except Exception: + continue + if result.returncode == 0 and result.stdout.strip(): + if marker is None or marker in result.stdout.lower(): + return True + return False + + def _ensure_rocm_torch() -> None: """Reinstall torch with ROCm wheels when the venv received CPU-only torch. - Runs only on Linux hosts where ROCm is installed. No-op when torch already - links against HIP (ROCm) or CUDA (NVIDIA). Skips on Windows/macOS. + Runs only on Linux hosts where ROCm is installed and an AMD GPU is + present. No-op when torch already links against HIP (ROCm) or CUDA + (NVIDIA). Skips on Windows/macOS. Uses pip_install() to respect uv, constraints, and --python targeting. """ rocm_root = os.environ.get("ROCM_PATH") or "/opt/rocm" if not os.path.isdir(rocm_root) and not shutil.which("hipcc"): return # no ROCm toolchain + if not _has_rocm_gpu(): + return # ROCm tools present but no AMD GPU ver = _detect_rocm_version() if ver is None: diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index 81da79aa32..adbdbe4e14 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -46,13 +46,22 @@ MOCK } # Helper: create a mock amd-smi that prints a given ROCm version string +# Supports both "amd-smi version" and "amd-smi list" subcommands so that +# the GPU presence check (amd-smi list) also succeeds in tests. make_mock_amd_smi() { _dir=$(mktemp -d) cat > "$_dir/amd-smi" <= comparison.""" ver = (7, 2) tag = next( - (t for (maj, mn), t in _ROCM_TORCH_INDEX.items() if ver >= (maj, mn)), + ( + t + for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) + if ver >= (maj, mn) + ), None, ) assert tag == "rocm7.1" @@ -665,7 +686,11 @@ class TestRocmTorchIndex: def test_rocm_64_selects_64(self): ver = (6, 4) tag = next( - (t for (maj, mn), t in _ROCM_TORCH_INDEX.items() if ver >= (maj, mn)), + ( + t + for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) + if ver >= (maj, mn) + ), None, ) assert tag == "rocm6.4"