ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Require actual AMD GPU presence before selecting ROCm paths

Browse source 
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
Download patch file Download diff file Expand all files Collapse all files

15
install.sh
View file

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

25
studio/install_llama_prebuilt.py
View file

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

32
studio/install_python_stack.py
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:

11
tests/sh/test_get_torch_index_url.sh
View file

@ -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" <<MOCK
#!/bin/sh
cat <<AMD_OUT
case "\$1" in
list)
printf 'GPU: 0\\n BDF: 0000:03:00.0\\n NAME: gfx1100\\n'
;;
*)
cat <<AMD_OUT
AMDSMI Tool: 25.0.1+2b74356 | AMDSMI Library version: 25.0.1.0 | ROCm version: $1
AMD_OUT
;;
esac
MOCK
chmod +x "$_dir/amd-smi"
echo "$_dir"

45
tests/studio/install/test_rocm_support.py
View file

@ -47,6 +47,7 @@ _STACK_SPEC.loader.exec_module(stack_mod)
_detect_rocm_version = stack_mod._detect_rocm_version
_ensure_rocm_torch = stack_mod._ensure_rocm_torch
_has_rocm_gpu = stack_mod._has_rocm_gpu
_ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX
@ -519,8 +520,9 @@ class TestEnsureRocmTorch:
mock_pip.assert_not_called()
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_torch_already_has_cuda_skips(self, mock_ver, mock_pip):
def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_pip):
"""If torch already has CUDA, should skip ROCm reinstall."""
mock_probe = MagicMock()
mock_probe.returncode = 0
@ -531,8 +533,9 @@ class TestEnsureRocmTorch:
mock_pip.assert_not_called()
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_torch_already_has_hip_skips(self, mock_ver, mock_pip):
def test_torch_already_has_hip_skips(self, mock_ver, mock_gpu, mock_pip):
"""If torch already has HIP, should skip ROCm reinstall."""
mock_probe = MagicMock()
mock_probe.returncode = 0
@ -543,8 +546,9 @@ class TestEnsureRocmTorch:
mock_pip.assert_not_called()
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_cpu_torch_gets_rocm_reinstall(self, mock_ver, mock_pip):
def test_cpu_torch_gets_rocm_reinstall(self, mock_ver, mock_gpu, mock_pip):
"""CPU-only torch on ROCm host should trigger reinstall."""
mock_probe = MagicMock()
mock_probe.returncode = 0
@ -560,8 +564,9 @@ class TestEnsureRocmTorch:
assert "bitsandbytes" in str(bnb_call)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 3))
def test_rocm_63_selects_correct_tag(self, mock_ver, mock_pip):
def test_rocm_63_selects_correct_tag(self, mock_ver, mock_gpu, mock_pip):
"""ROCm 6.3 should select rocm6.3 tag."""
mock_probe = MagicMock()
mock_probe.returncode = 0
@ -573,8 +578,9 @@ class TestEnsureRocmTorch:
assert "rocm6.3" in str(torch_call)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (5, 0))
def test_old_rocm_skips(self, mock_ver, mock_pip):
def test_old_rocm_skips(self, mock_ver, mock_gpu, mock_pip):
"""ROCm version too old (below 6.0) should skip."""
mock_probe = MagicMock()
mock_probe.returncode = 0
@ -585,8 +591,9 @@ class TestEnsureRocmTorch:
mock_pip.assert_not_called()
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = None)
def test_version_unreadable_prints_warning(self, mock_ver, mock_pip, capsys):
def test_version_unreadable_prints_warning(self, mock_ver, mock_gpu, mock_pip, capsys):
"""ROCm detected but version unreadable should print warning and skip."""
with patch("os.path.isdir", return_value = True):
_ensure_rocm_torch()
@ -595,8 +602,9 @@ class TestEnsureRocmTorch:
assert "unreadable" in captured.out
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
def test_rocm_72_selects_71_tag(self, mock_ver, mock_pip):
def test_rocm_72_selects_71_tag(self, mock_ver, mock_gpu, mock_pip):
"""ROCm 7.2 should select rocm7.1 tag (capped, not in mapping)."""
mock_probe = MagicMock()
mock_probe.returncode = 0
@ -608,8 +616,9 @@ class TestEnsureRocmTorch:
assert "rocm7.1" in str(torch_call)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_probe_timeout_triggers_reinstall(self, mock_ver, mock_pip):
def test_probe_timeout_triggers_reinstall(self, mock_ver, mock_gpu, mock_pip):
"""Probe subprocess timeout should not crash; should proceed to reinstall."""
with patch("os.path.isdir", return_value = True):
with patch(
@ -620,6 +629,14 @@ class TestEnsureRocmTorch:
assert mock_pip.call_count == 2
assert "rocm7.1" in str(mock_pip.call_args_list[0])
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_rocm_gpu", return_value = False)
def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_pip):
"""ROCm tools present but no actual AMD GPU should skip entirely."""
with patch("os.path.isdir", return_value = True):
_ensure_rocm_torch()
mock_pip.assert_not_called()
# =============================================================================
# TEST: install_python_stack.py -- _ROCM_TORCH_INDEX mapping
@ -657,7 +674,11 @@ class TestRocmTorchIndex:
"""ROCm 7.2 (not in map) should select rocm7.1 via >= 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"