Harden ROCm detection, fix VRAM heuristic, and expand RDNA2 coverage
- Windows ROCm detection: validate actual GPU presence via hipinfo/amd-smi
output markers instead of just checking tool existence on PATH
- _ensure_rocm_torch: validate nvidia-smi actually reports a GPU before
giving NVIDIA precedence (fixes AMD-only hosts with stale NVIDIA tools)
- amd.py _parse_numeric: handle dict-shaped metric objects from newer
amd-smi versions ({"value": 10, "unit": "W"}) and strip MiB/GiB units
- amd.py VRAM heuristic: raise threshold from 100k to 10M to correctly
handle MI300X (192 GB = 196608 MB) and other high-VRAM GPUs
- amd.py visible GPU: use AMD-reported GPU IDs instead of enumerate index
so non-dense sets like CUDA_VISIBLE_DEVICES=1,3 report correctly
- install.sh: add ROCm <6.0 minimum version guard (no PyTorch wheels
exist for older versions); fix rocm7.1* glob to not match rocm7.10+
- is_rdna: add gfx1033-1036 for RDNA2 mobile GPUs (RX 6600M etc.)
- worker.py: increase ROCm source build timeout from 600s to 1800s;
fix success log message for ROCm source builds
- Tests: update mocks for _has_usable_nvidia_gpu, add RDNA2 target asserts
This commit is contained in:
Daniel Han
parent
14823265d2
commit
134638dd0e
7 changed files with 108 additions and 41 deletions
|
|
@ -1019,12 +1019,16 @@ get_torch_index_url() {
|
|||
*) _rocm_tag="" ;; # reject malformed (empty, garbled, or major=0)
|
||||
esac
|
||||
if [ -n "$_rocm_tag" ]; then
|
||||
# Minimum supported: ROCm 6.0 (no PyTorch wheels exist for older)
|
||||
case "$_rocm_tag" in
|
||||
rocm[1-5].*) echo "$_base/cpu"; return ;;
|
||||
esac
|
||||
# ROCm 7.2 only has torch 2.11.0 which exceeds current bounds (<2.11.0).
|
||||
# Fall back to rocm7.1 index which has torch 2.10.0.
|
||||
# TODO: uncomment the next line when torch upper bound is bumped to >=2.11.0
|
||||
# echo "$_base/$_rocm_tag"; return
|
||||
case "$_rocm_tag" in
|
||||
rocm6.*|rocm7.0*|rocm7.1*)
|
||||
rocm6.*|rocm7.0|rocm7.0.*|rocm7.1|rocm7.1.*)
|
||||
echo "$_base/$_rocm_tag" ;;
|
||||
*)
|
||||
# ROCm 7.2+ (including future 10.x+): cap to rocm7.1
|
||||
|
|
|
|||
|
|
@ -284,8 +284,8 @@ def _install_package_wheel_first(
|
|||
f"{pypi_name}=={pypi_version}",
|
||||
]
|
||||
|
||||
# Source compilation on ROCm can take 5-10 minutes; use a generous timeout
|
||||
timeout = 600 if is_hip else 300
|
||||
# Source compilation on ROCm can take 10-30 minutes; use a generous timeout
|
||||
timeout = 1800 if is_hip else 300
|
||||
|
||||
try:
|
||||
result = _sp.run(
|
||||
|
|
@ -331,7 +331,10 @@ def _install_package_wheel_first(
|
|||
)
|
||||
return
|
||||
|
||||
logger.info("Installed %s from PyPI", display_name)
|
||||
if is_hip:
|
||||
logger.info("Compiled and installed %s from source for ROCm", display_name)
|
||||
else:
|
||||
logger.info("Installed %s from PyPI", display_name)
|
||||
|
||||
|
||||
def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None:
|
||||
|
|
|
|||
|
|
@ -17,7 +17,7 @@ from loggers import get_logger
|
|||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[dict]:
|
||||
def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
|
||||
"""Run amd-smi with the given arguments and return parsed JSON, or None."""
|
||||
try:
|
||||
result = subprocess.run(
|
||||
|
|
@ -43,11 +43,17 @@ def _parse_numeric(value: Any) -> Optional[float]:
|
|||
"""Extract a numeric value from amd-smi output (may be str, int, float, or dict)."""
|
||||
if value is None:
|
||||
return None
|
||||
# Newer amd-smi versions emit {"value": 10, "unit": "W"}
|
||||
if isinstance(value, dict):
|
||||
return _parse_numeric(value.get("value"))
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
import math
|
||||
f = float(value)
|
||||
return f if math.isfinite(f) else None
|
||||
if isinstance(value, str):
|
||||
# Strip units like "W", "C", "%", "MB" etc.
|
||||
cleaned = value.strip().rstrip("WCMBGb% ").strip()
|
||||
# Strip units like "W", "C", "%", "MB", "MiB", "GB", "GiB" etc.
|
||||
import re
|
||||
cleaned = re.sub(r'\s*[A-Za-z/%]+$', '', value.strip())
|
||||
if not cleaned or cleaned.lower() in ("n/a", "none", "unknown"):
|
||||
return None
|
||||
try:
|
||||
|
|
@ -112,16 +118,18 @@ def _extract_gpu_metrics(gpu_data: dict) -> dict[str, Any]:
|
|||
vram_used_bytes = None
|
||||
vram_total_bytes = None
|
||||
|
||||
# Convert VRAM from bytes to MB if values are large (>10000 = likely bytes)
|
||||
# Convert VRAM from bytes to MB if values are very large.
|
||||
# amd-smi typically reports in MB, but some versions report bytes.
|
||||
# Threshold: 10 million -- no GPU has <10 MB, and even 10 TB = 10M MB.
|
||||
vram_used_mb = None
|
||||
vram_total_mb = None
|
||||
if vram_used_bytes is not None:
|
||||
if vram_used_bytes > 100000: # Likely bytes
|
||||
if vram_used_bytes > 10_000_000: # Likely bytes (>10M)
|
||||
vram_used_mb = vram_used_bytes / (1024 * 1024)
|
||||
else: # Likely already MB
|
||||
vram_used_mb = vram_used_bytes
|
||||
if vram_total_bytes is not None:
|
||||
if vram_total_bytes > 100000: # Likely bytes
|
||||
if vram_total_bytes > 10_000_000: # Likely bytes (>10M)
|
||||
vram_total_mb = vram_total_bytes / (1024 * 1024)
|
||||
else: # Likely already MB
|
||||
vram_total_mb = vram_total_bytes
|
||||
|
|
@ -211,12 +219,18 @@ def get_visible_gpu_utilization(
|
|||
"index_kind": "physical",
|
||||
}
|
||||
|
||||
gpu_list = data if isinstance(data, list) else data.get("gpus", [data])
|
||||
gpu_list = data if isinstance(data, list) else data.get("gpus", data.get("gpu", [data]))
|
||||
visible_set = set(parent_visible_ids)
|
||||
ordinal_map = {gpu_id: ordinal for ordinal, gpu_id in enumerate(parent_visible_ids)}
|
||||
|
||||
devices = []
|
||||
for idx, gpu_data in enumerate(gpu_list):
|
||||
for fallback_idx, gpu_data in enumerate(gpu_list):
|
||||
# Use AMD-reported GPU ID when available, fall back to enumeration index
|
||||
raw_id = gpu_data.get("gpu", gpu_data.get("gpu_id", gpu_data.get("id", fallback_idx))) if isinstance(gpu_data, dict) else fallback_idx
|
||||
try:
|
||||
idx = int(raw_id)
|
||||
except (TypeError, ValueError):
|
||||
idx = fallback_idx
|
||||
if idx not in visible_set:
|
||||
continue
|
||||
metrics = _extract_gpu_metrics(gpu_data)
|
||||
|
|
|
|||
|
|
@ -1450,15 +1450,22 @@ def detect_host() -> HostInfo:
|
|||
has_rocm = True
|
||||
break
|
||||
elif is_windows:
|
||||
# Windows: check for HIP runtime DLL or hipinfo tool
|
||||
if shutil.which("hipinfo") or shutil.which("amd-smi"):
|
||||
has_rocm = True
|
||||
elif any(
|
||||
Path(d).joinpath("amdhip64.dll").exists()
|
||||
for d in os.environ.get("PATH", "").split(os.pathsep)
|
||||
if d
|
||||
# Windows: validate actual AMD GPU presence (not just tool/DLL existence)
|
||||
for _cmd, _marker in (
|
||||
(["hipinfo"], "gcnarchname"),
|
||||
(["amd-smi", "list"], "gpu"),
|
||||
):
|
||||
has_rocm = True
|
||||
_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 in _result.stdout.lower():
|
||||
has_rocm = True
|
||||
break
|
||||
|
||||
return HostInfo(
|
||||
system = system,
|
||||
|
|
|
|||
|
|
@ -109,6 +109,24 @@ def _has_rocm_gpu() -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _has_usable_nvidia_gpu() -> bool:
|
||||
"""Return True only when nvidia-smi exists AND reports at least one GPU."""
|
||||
exe = shutil.which("nvidia-smi")
|
||||
if not exe:
|
||||
return False
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[exe, "-L"],
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.DEVNULL,
|
||||
text=True,
|
||||
timeout=10,
|
||||
)
|
||||
except Exception:
|
||||
return False
|
||||
return result.returncode == 0 and "GPU " in result.stdout
|
||||
|
||||
|
||||
def _ensure_rocm_torch() -> None:
|
||||
"""Reinstall torch with ROCm wheels when the venv received CPU-only torch.
|
||||
|
||||
|
|
@ -118,8 +136,8 @@ def _ensure_rocm_torch() -> None:
|
|||
(NVIDIA takes precedence).
|
||||
Uses pip_install() to respect uv, constraints, and --python targeting.
|
||||
"""
|
||||
# NVIDIA takes precedence on mixed hosts
|
||||
if shutil.which("nvidia-smi"):
|
||||
# NVIDIA takes precedence on mixed hosts -- but only if an actual GPU is usable
|
||||
if _has_usable_nvidia_gpu():
|
||||
return
|
||||
rocm_root = os.environ.get("ROCM_PATH") or "/opt/rocm"
|
||||
if not os.path.isdir(rocm_root) and not shutil.which("hipcc"):
|
||||
|
|
@ -707,7 +725,7 @@ def install_python_stack() -> int:
|
|||
|
||||
# Windows + AMD GPU: PyTorch does not publish ROCm wheels for Windows.
|
||||
# Detect and warn so users know manual steps are needed for GPU training.
|
||||
if IS_WINDOWS and not NO_TORCH:
|
||||
if IS_WINDOWS and not NO_TORCH and not _has_usable_nvidia_gpu():
|
||||
if shutil.which("hipinfo") or shutil.which("amd-smi"):
|
||||
_safe_print(
|
||||
_dim(" Note:"),
|
||||
|
|
|
|||
|
|
@ -49,6 +49,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
|
||||
_has_usable_nvidia_gpu = stack_mod._has_usable_nvidia_gpu
|
||||
_ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX
|
||||
|
||||
|
||||
|
|
@ -396,20 +397,20 @@ class TestHostInfoRocm:
|
|||
assert host.has_rocm is False
|
||||
assert host.has_usable_nvidia is True
|
||||
|
||||
def test_detect_host_with_rocm_path_env(self):
|
||||
"""detect_host() checks ROCM_PATH env var for ROCm detection."""
|
||||
# Verify the detect_host function source references ROCM_PATH
|
||||
def test_detect_host_has_rocm_detection_logic(self):
|
||||
"""detect_host() should have ROCm GPU detection logic."""
|
||||
import inspect
|
||||
|
||||
source = inspect.getsource(prebuilt_mod.detect_host)
|
||||
assert "ROCM_PATH" in source or "rocm" in source.lower()
|
||||
# Must probe for actual GPU, not just tool presence
|
||||
assert "rocminfo" in source or "amd-smi" in source
|
||||
|
||||
def test_detect_host_windows_rocm_detection(self):
|
||||
"""detect_host() source should have Windows-specific HIP detection."""
|
||||
"""detect_host() source should have Windows-specific ROCm GPU detection."""
|
||||
import inspect
|
||||
|
||||
source = inspect.getsource(prebuilt_mod.detect_host)
|
||||
assert "hipinfo" in source or "amdhip64" in source
|
||||
assert "hipinfo" in source or "amd-smi" in source
|
||||
|
||||
|
||||
# =============================================================================
|
||||
|
|
@ -520,7 +521,8 @@ class TestEnsureRocmTorch:
|
|||
"""Verify ROCm torch reinstall logic."""
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
def test_no_rocm_skips(self, mock_pip):
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
def test_no_rocm_skips(self, mock_nvidia, mock_pip):
|
||||
"""No ROCm toolchain should skip entirely."""
|
||||
with patch("os.path.isdir", return_value = False):
|
||||
with patch("shutil.which", return_value = None):
|
||||
|
|
@ -528,9 +530,10 @@ class TestEnsureRocmTorch:
|
|||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@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_gpu, mock_pip):
|
||||
def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
|
||||
"""If torch already has CUDA, should skip ROCm reinstall."""
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
|
|
@ -541,9 +544,10 @@ class TestEnsureRocmTorch:
|
|||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@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_gpu, mock_pip):
|
||||
def test_torch_already_has_hip_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
|
||||
"""If torch already has HIP, should skip ROCm reinstall."""
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
|
|
@ -554,9 +558,10 @@ class TestEnsureRocmTorch:
|
|||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@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_gpu, mock_pip):
|
||||
def test_cpu_torch_gets_rocm_reinstall(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
|
||||
"""CPU-only torch on ROCm host should trigger reinstall."""
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
|
|
@ -572,9 +577,10 @@ class TestEnsureRocmTorch:
|
|||
assert "bitsandbytes" in str(bnb_call)
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@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_gpu, mock_pip):
|
||||
def test_rocm_63_selects_correct_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
|
||||
"""ROCm 6.3 should select rocm6.3 tag."""
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
|
|
@ -586,9 +592,10 @@ class TestEnsureRocmTorch:
|
|||
assert "rocm6.3" in str(torch_call)
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@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_gpu, mock_pip):
|
||||
def test_old_rocm_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
|
||||
"""ROCm version too old (below 6.0) should skip."""
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
|
|
@ -599,10 +606,11 @@ class TestEnsureRocmTorch:
|
|||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@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_gpu, mock_pip, capsys
|
||||
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, capsys
|
||||
):
|
||||
"""ROCm detected but version unreadable should print warning and skip."""
|
||||
with patch("os.path.isdir", return_value = True):
|
||||
|
|
@ -612,9 +620,10 @@ class TestEnsureRocmTorch:
|
|||
assert "unreadable" in captured.out
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@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_gpu, mock_pip):
|
||||
def test_rocm_72_selects_71_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
|
||||
"""ROCm 7.2 should select rocm7.1 tag (capped, not in mapping)."""
|
||||
mock_probe = MagicMock()
|
||||
mock_probe.returncode = 0
|
||||
|
|
@ -626,9 +635,10 @@ class TestEnsureRocmTorch:
|
|||
assert "rocm7.1" in str(torch_call)
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@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_gpu, mock_pip):
|
||||
def test_probe_timeout_triggers_reinstall(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
|
||||
"""Probe subprocess timeout should not crash; should proceed to reinstall."""
|
||||
with patch("os.path.isdir", return_value = True):
|
||||
with patch(
|
||||
|
|
@ -640,8 +650,9 @@ class TestEnsureRocmTorch:
|
|||
assert "rocm7.1" in str(mock_pip.call_args_list[0])
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
|
||||
@patch.object(stack_mod, "_has_rocm_gpu", return_value = False)
|
||||
def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_pip):
|
||||
def test_no_gpu_with_rocm_tools_skips(self, mock_gpu, mock_nvidia, mock_pip):
|
||||
"""ROCm tools present but no actual AMD GPU should skip entirely."""
|
||||
with patch("os.path.isdir", return_value = True):
|
||||
_ensure_rocm_torch()
|
||||
|
|
@ -865,7 +876,9 @@ class TestInstallShStructure:
|
|||
source = sh_path.read_text()
|
||||
assert 'echo "$_base/rocm7.1"' in source # fallback for unknown versions
|
||||
# Allowlisted versions should pass through directly
|
||||
assert "rocm6.*|rocm7.0*|rocm7.1*)" in source
|
||||
assert "rocm6.*" in source
|
||||
assert "rocm7.0" in source
|
||||
assert "rocm7.1" in source
|
||||
|
||||
def test_rocm_tag_validation_guard_exists(self):
|
||||
"""install.sh should validate _rocm_tag with a case guard."""
|
||||
|
|
@ -1269,6 +1282,10 @@ class TestIsRdnaExpansion:
|
|||
assert "gfx1030" in func_body
|
||||
assert "gfx1031" in func_body
|
||||
assert "gfx1032" in func_body
|
||||
assert "gfx1033" in func_body
|
||||
assert "gfx1034" in func_body
|
||||
assert "gfx1035" in func_body
|
||||
assert "gfx1036" in func_body
|
||||
|
||||
def test_is_rdna_source_has_rdna3(self):
|
||||
"""is_rdna() should include RDNA3 architectures."""
|
||||
|
|
|
|||
|
|
@ -94,6 +94,10 @@ def is_rdna():
|
|||
"gfx1030",
|
||||
"gfx1031",
|
||||
"gfx1032",
|
||||
"gfx1033",
|
||||
"gfx1034",
|
||||
"gfx1035",
|
||||
"gfx1036",
|
||||
# RDNA3 (Navi 31-33)
|
||||
"gfx1100",
|
||||
"gfx1101",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue