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
This commit is contained in:
Daniel Han 2026-03-31 09:25:56 +00:00
commit 9e33c25eac
5 changed files with 105 additions and 23 deletions

View file

@ -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: