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"