From 62912e106923340b4742f36f74ce52ea365334fa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 08:29:15 +0000 Subject: [PATCH 1/6] Add AMD ROCm detection to install.sh with validation and shell tests Extend get_torch_index_url() to detect AMD ROCm when nvidia-smi is not found. The detection chain tries, in order: amd-smi, /opt/rocm version file, hipconfig, dpkg-query (rocm-core), and rpm (rocm-core). Key design decisions: - NVIDIA always takes precedence (ROCm detection only runs when nvidia-smi is absent) - ROCm 7.2+ is capped to rocm7.1 index because torch 2.11.0 (the only version on the rocm7.2 index) exceeds the current upper bound <2.11.0 - macOS returns CPU immediately before any ROCm check - A validation guard rejects malformed _rocm_tag values (e.g. "rocm." from garbled amd-smi output with empty version fields) - Debian epoch prefixes (e.g. "2:6.2.0") are stripped from dpkg-query output before parsing Also adds ROCm status messaging (CPU-only hint mentions AMD, ROCm detected message), and installs bitsandbytes for AMD when a ROCm PyTorch index is selected. Shell tests expanded from 8 to 23 cases covering all CUDA tiers, ROCm versions 6.0-8.0, CUDA+ROCm precedence, malformed amd-smi output (empty version, "N/A", trailing text), and CUDA regression checks. --- install.sh | 54 +++++++++++- tests/sh/test_get_torch_index_url.sh | 119 ++++++++++++++++++++++++++- 2 files changed, 169 insertions(+), 4 deletions(-) diff --git a/install.sh b/install.sh index 9ea80bc161..b8f49abf51 100755 --- a/install.sh +++ b/install.sh @@ -982,7 +982,44 @@ get_torch_index_url() { elif [ -x "/usr/bin/nvidia-smi" ]; then _smi="/usr/bin/nvidia-smi" fi - if [ -z "$_smi" ]; then echo "$_base/cpu"; return; fi + if [ -z "$_smi" ]; then + # No NVIDIA GPU -- check for AMD ROCm + _rocm_tag="" + _rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \ + amd-smi version 2>/dev/null | awk -F'ROCm version: ' \ + 'NF>1{gsub(/[^0-9.]/, "", $2); split($2,a,"."); print "rocm"a[1]"."a[2]; ok=1; exit} END{exit !ok}'; } || \ + { [ -r /opt/rocm/.info/version ] && \ + awk -F. '{print "rocm"$1"."$2; exit}' /opt/rocm/.info/version; } || \ + { command -v hipconfig >/dev/null 2>&1 && \ + hipconfig --version 2>/dev/null | awk 'NR==1{split($1,a,"."); if(a[1]+0>0) print "rocm"a[1]"."a[2]}'; } || \ + { command -v dpkg-query >/dev/null 2>&1 && \ + ver="$(dpkg-query -W -f='${Version}\n' rocm-core 2>/dev/null)" && \ + [ -n "$ver" ] && \ + printf '%s\n' "$ver" | sed 's/^[0-9]*://' | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; } || \ + { command -v rpm >/dev/null 2>&1 && \ + ver="$(rpm -q --qf '%{VERSION}\n' rocm-core 2>/dev/null)" && \ + [ -n "$ver" ] && \ + printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null + # Validate _rocm_tag: must match "rocmX.Y" with leading digits + case "$_rocm_tag" in + rocm[0-9]*.[0-9]*) : ;; # valid + *) _rocm_tag="" ;; # reject malformed (empty version, garbled output) + esac + if [ -n "$_rocm_tag" ]; then + # 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 + rocm7.2*|rocm7.3*|rocm7.4*|rocm7.5*|rocm8*|rocm9*) + echo "$_base/rocm7.1" ;; + *) + echo "$_base/$_rocm_tag" ;; + esac + return + fi + echo "$_base/cpu"; return + fi # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P) _cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \ | sed -n 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \ @@ -1007,13 +1044,19 @@ case "$TORCH_INDEX_URL" in */cpu) if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then echo "" - echo " NOTE: No NVIDIA GPU detected (nvidia-smi not found)." + echo " NOTE: No GPU detected (nvidia-smi and ROCm not found)." echo " Installing CPU-only PyTorch. If you only need GGUF chat/inference," echo " re-run with --no-torch for a faster, lighter install:" echo " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch" + echo " AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd" echo "" fi ;; + */rocm*) + echo "" + echo " AMD ROCm detected -- installing ROCm-enabled PyTorch ($TORCH_INDEX_URL)" + echo "" + ;; esac # ── Install unsloth directly into the venv (no activation needed) ── @@ -1051,6 +1094,13 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "installing PyTorch ($TORCH_INDEX_URL)..." run_install_cmd "install PyTorch" uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \ --index-url "$TORCH_INDEX_URL" + # AMD ROCm: install bitsandbytes with AMD support + case "$TORCH_INDEX_URL" in + */rocm*) + substep "installing bitsandbytes for AMD ROCm..." + run_install_cmd "install bitsandbytes (AMD)" uv pip install --python "$_VENV_PY" "bitsandbytes>=0.49.1" + ;; + esac fi # Fresh: Step 2 - install unsloth, preserving pre-installed torch substep "installing unsloth (this may take a few minutes)..." diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index 6387922712..81da79aa32 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -45,10 +45,23 @@ MOCK echo "$_dir" } +# Helper: create a mock amd-smi that prints a given ROCm version string +make_mock_amd_smi() { + _dir=$(mktemp -d) + cat > "$_dir/amd-smi" </dev/null || true) [ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd" done @@ -119,6 +132,108 @@ _result=$(run_func "$_dir") assert_eq "unparseable -> cu126" "https://download.pytorch.org/whl/cu126" "$_result" rm -rf "$_dir" +# 9) ROCm 6.3 (no nvidia-smi) -> rocm6.3 +_dir=$(make_mock_amd_smi "6.3") +_result=$(run_func "$_dir") +assert_eq "ROCm 6.3 -> rocm6.3" "https://download.pytorch.org/whl/rocm6.3" "$_result" +rm -rf "$_dir" + +# 10) ROCm 7.1 (no nvidia-smi) -> rocm7.1 +_dir=$(make_mock_amd_smi "7.1") +_result=$(run_func "$_dir") +assert_eq "ROCm 7.1 -> rocm7.1" "https://download.pytorch.org/whl/rocm7.1" "$_result" +rm -rf "$_dir" + +# 11) ROCm 7.2 (no nvidia-smi) -> rocm7.1 (capped due to torch <2.11.0) +_dir=$(make_mock_amd_smi "7.2") +_result=$(run_func "$_dir") +assert_eq "ROCm 7.2 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result" +rm -rf "$_dir" + +# 12) Both nvidia-smi and amd-smi present -> CUDA takes precedence +_cuda_dir=$(make_mock_smi "12.6") +_amd_dir=$(make_mock_amd_smi "6.3") +_combined_dir=$(mktemp -d) +ln -sf "$_cuda_dir/nvidia-smi" "$_combined_dir/nvidia-smi" +ln -sf "$_amd_dir/amd-smi" "$_combined_dir/amd-smi" +_result=$(run_func "$_combined_dir") +assert_eq "CUDA+ROCm -> CUDA precedence" "https://download.pytorch.org/whl/cu126" "$_result" +rm -rf "$_cuda_dir" "$_amd_dir" "$_combined_dir" + +# 13) No nvidia-smi, no amd-smi -> cpu (duplicate of test 1, confirms ROCm didn't break it) +_result=$(run_func "none") +assert_eq "no GPU -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" + +# 14) ROCm 6.1 (no nvidia-smi) -> rocm6.1 +_dir=$(make_mock_amd_smi "6.1") +_result=$(run_func "$_dir") +assert_eq "ROCm 6.1 -> rocm6.1" "https://download.pytorch.org/whl/rocm6.1" "$_result" +rm -rf "$_dir" + +# 15) ROCm 6.4 (no nvidia-smi) -> rocm6.4 +_dir=$(make_mock_amd_smi "6.4") +_result=$(run_func "$_dir") +assert_eq "ROCm 6.4 -> rocm6.4" "https://download.pytorch.org/whl/rocm6.4" "$_result" +rm -rf "$_dir" + +# 16) ROCm 7.0 (no nvidia-smi) -> rocm7.0 +_dir=$(make_mock_amd_smi "7.0") +_result=$(run_func "$_dir") +assert_eq "ROCm 7.0 -> rocm7.0" "https://download.pytorch.org/whl/rocm7.0" "$_result" +rm -rf "$_dir" + +# 17) ROCm 8.0 (future, no nvidia-smi) -> rocm7.1 (capped) +_dir=$(make_mock_amd_smi "8.0") +_result=$(run_func "$_dir") +assert_eq "ROCm 8.0 -> rocm7.1 (capped)" "https://download.pytorch.org/whl/rocm7.1" "$_result" +rm -rf "$_dir" + +# 18) Malformed amd-smi output (empty version field) -> cpu +_dir=$(mktemp -d) +cat > "$_dir/amd-smi" <<'MOCK' +#!/bin/sh +echo "AMDSMI Tool: 25.0.1 | AMDSMI Library version: 25.0.1.0 | ROCm version: " +MOCK +chmod +x "$_dir/amd-smi" +_result=$(run_func "$_dir") +assert_eq "empty amd-smi version -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + +# 19) amd-smi with "N/A" version -> cpu +_dir=$(mktemp -d) +cat > "$_dir/amd-smi" <<'MOCK' +#!/bin/sh +echo "AMDSMI Tool: 25.0.1 | AMDSMI Library version: 25.0.1.0 | ROCm version: N/A" +MOCK +chmod +x "$_dir/amd-smi" +_result=$(run_func "$_dir") +assert_eq "N/A amd-smi version -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" +rm -rf "$_dir" + +# 20) ROCm version with trailing text (e.g. "6.3.1-beta") -> rocm6.3 +_dir=$(make_mock_amd_smi "6.3.1-beta") +_result=$(run_func "$_dir") +assert_eq "ROCm 6.3.1-beta -> rocm6.3" "https://download.pytorch.org/whl/rocm6.3" "$_result" +rm -rf "$_dir" + +# 22) CUDA 12.6 still works after ROCm changes (regression check) +_dir=$(make_mock_smi "12.6") +_result=$(run_func "$_dir") +assert_eq "CUDA 12.6 regression -> cu126" "https://download.pytorch.org/whl/cu126" "$_result" +rm -rf "$_dir" + +# 23) CUDA 13.0 still works after ROCm changes (regression check) +_dir=$(make_mock_smi "13.0") +_result=$(run_func "$_dir") +assert_eq "CUDA 13.0 regression -> cu130" "https://download.pytorch.org/whl/cu130" "$_result" +rm -rf "$_dir" + +# 24) CUDA 12.8 still works after ROCm changes (regression check) +_dir=$(make_mock_smi "12.8") +_result=$(run_func "$_dir") +assert_eq "CUDA 12.8 regression -> cu128" "https://download.pytorch.org/whl/cu128" "$_result" +rm -rf "$_dir" + rm -f "$_FUNC_FILE" rm -rf "$_FAKE_SMI_DIR" rm -rf "$_TOOLS_DIR" From 485dcaae3718f6810f6cecf1b391b1acf3d25337 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 08:30:09 +0000 Subject: [PATCH 2/6] Add ROCm torch reinstall support to install_python_stack.py When a Linux host has ROCm installed but the venv received CPU-only torch (common when pip resolves torch from PyPI without an explicit index URL), the new _ensure_rocm_torch() function detects this and reinstalls torch with the correct ROCm wheels. New functions: - _detect_rocm_version(): probes /opt/rocm version files, hipconfig, and ROCM_PATH to determine (major, minor) of the installed ROCm stack - _ensure_rocm_torch(): checks whether torch is GPU-enabled; if not and ROCm is present, reinstalls torch from the best-matching ROCm wheel index and installs bitsandbytes for AMD The ROCm version-to-wheel mapping covers ROCm 6.0 through 7.1. ROCm 7.2 is excluded from the active mapping because the only torch build on that index (2.11.0) exceeds the current upper bound (<2.11.0); ROCm 7.2 hosts get the rocm7.1 wheels via the >= fallback. The torch GPU probe subprocess has a 30-second timeout to prevent hangs if torch import stalls (e.g. broken CUDA/HIP driver). The step count is incremented on Linux (non-macOS, non-no-torch) to account for the new "ROCm torch check" progress step. --- studio/install_python_stack.py | 124 +++++++++++++++++++++++++++++++++ 1 file changed, 124 insertions(+) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index f2981ea665..1f17a2e79c 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -25,6 +25,120 @@ IS_WINDOWS = sys.platform == "win32" IS_MACOS = sys.platform == "darwin" IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64" +# ── ROCm / AMD GPU support ───────────────────────────────────────────────────── +# Mapping from detected ROCm (major, minor) to the best PyTorch wheel tag on +# download.pytorch.org. Entries are checked newest-first (>=). +# ROCm 7.2 only has torch 2.11.0 on download.pytorch.org, which exceeds the +# current torch upper bound (<2.11.0). Fall back to rocm7.1 (torch 2.10.0). +# TODO: uncomment rocm7.2 when torch upper bound is bumped to >=2.11.0 +_ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { + # (7, 2): "rocm7.2", # torch 2.11.0 -- requires torch>=2.11 + (7, 1): "rocm7.1", + (7, 0): "rocm7.0", + (6, 4): "rocm6.4", + (6, 3): "rocm6.3", + (6, 2): "rocm6.2", + (6, 1): "rocm6.1", + (6, 0): "rocm6.0", +} +_PYTORCH_WHL_BASE = "https://download.pytorch.org/whl" + + +def _detect_rocm_version() -> tuple[int, int] | None: + """Return (major, minor) of the installed ROCm stack, or None.""" + # Check /opt/rocm/.info/version or ROCM_PATH equivalent + rocm_root = os.environ.get("ROCM_PATH", "/opt/rocm") + for path in ( + os.path.join(rocm_root, ".info", "version"), + os.path.join(rocm_root, "lib", "rocm_version"), + ): + try: + parts = open(path).read().strip().split("-")[0].split(".") + return int(parts[0]), int(parts[1]) + except Exception: + pass + + # Try hipconfig --version (outputs bare version like "6.3.21234.2") + hipconfig = shutil.which("hipconfig") + if hipconfig: + try: + result = subprocess.run( + [hipconfig, "--version"], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 5, + ) + if result.returncode == 0: + raw = result.stdout.decode().strip().split("\n")[0] + parts = raw.split(".") + if len(parts) >= 2 and parts[0].isdigit(): + return int(parts[0]), int(parts[1]) + except Exception: + pass + + return None + + +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. + Uses pip_install() to respect uv, constraints, and --python targeting. + """ + rocm_root = os.environ.get("ROCM_PATH", "/opt/rocm") + if not os.path.isdir(rocm_root) and not shutil.which("hipcc"): + return # no ROCm toolchain + + ver = _detect_rocm_version() + if ver is None: + print(" ROCm detected but version unreadable -- skipping torch reinstall") + return + + # Skip if torch is already GPU-enabled (HIP or CUDA) + probe = subprocess.run( + [ + sys.executable, + "-c", + "import torch; print(torch.version.hip or torch.version.cuda or '')", + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 30, + ) + if probe.returncode == 0 and probe.stdout.decode().strip(): + return # torch already GPU-enabled + + # Select best matching wheel tag (newest ROCm version <= installed) + tag = next( + (t for (maj, mn), t in _ROCM_TORCH_INDEX.items() if ver >= (maj, mn)), + None, + ) + if tag is None: + print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- skipping") + return + + index_url = f"{_PYTORCH_WHL_BASE}/{tag}" + print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}") + pip_install( + f"ROCm torch ({tag})", + "--force-reinstall", + "--no-cache-dir", + "torch", + "torchvision", + "torchaudio", + "--index-url", + index_url, + constrain = False, + ) + # Also install bitsandbytes for AMD + pip_install( + "bitsandbytes (AMD)", + "--no-cache-dir", + "bitsandbytes>=0.49.1", + constrain = False, + ) + def _infer_no_torch() -> bool: """Determine whether to run in no-torch (GGUF-only) mode. @@ -414,6 +528,9 @@ def install_python_stack() -> int: base_total = 10 if IS_WINDOWS else 11 if IS_MACOS: base_total -= 1 # triton step is skipped on macOS + # ROCm torch check step (Linux only, non-macOS, non-no-torch) + if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: + base_total += 1 _TOTAL = (base_total - 1) if skip_base else base_total # 1. Try to use uv for faster installs (must happen before pip upgrade @@ -537,6 +654,13 @@ def install_python_stack() -> int: req = REQ_ROOT / "base.txt", ) + # 2b. AMD ROCm: reinstall torch with HIP wheels if the host has ROCm but the + # venv received CPU-only torch (common when pip resolves torch from PyPI). + # Must come immediately after base packages so torch is present for inspection. + if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: + _progress("ROCm torch check") + _ensure_rocm_torch() + # 3. Extra dependencies _progress("unsloth extras") pip_install( From cc904e9b640139d01ffd0724e205b82d57bb30cf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 08:30:53 +0000 Subject: [PATCH 3/6] Add ROCm support to llama.cpp prebuilt installer Extend the prebuilt installer to detect AMD ROCm GPUs and select the appropriate llama.cpp binary: - Add has_rocm field (default False) to HostInfo dataclass - Extend detect_host() to probe for ROCm via hipcc, amd-smi, rocm-smi, /opt/rocm, and ROCM_PATH env var (skipped on macOS) - Linux x86_64 ROCm hosts (without NVIDIA) try the upstream ROCm 7.2 prebuilt first, with a log warning that it may fall back to source build on other ROCm versions - Windows x86_64 ROCm hosts try the HIP prebuilt; if not found, a log message is printed before falling through to CPU - Add "linux-rocm" and "windows-hip" install kinds to runtime_patterns_for_choice() with the correct shared library globs (libggml-hip.so* for Linux, *.dll for Windows) The ROCm path is only entered when has_usable_nvidia is False, so NVIDIA always takes precedence on dual-GPU systems. The source build fallback (via setup.sh with -DGGML_HIP=ON) compiles against the exact GPU target via rocminfo, which is more reliable for consumer GPUs (e.g. gfx1151) that may not be in the prebuilt. --- studio/install_llama_prebuilt.py | 58 ++++++++++++++++++++++++++++++-- 1 file changed, 56 insertions(+), 2 deletions(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 516dc4b6a4..cbc8f6d15f 100755 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -88,6 +88,7 @@ class HostInfo: visible_cuda_devices: str | None has_physical_nvidia: bool has_usable_nvidia: bool + has_rocm: bool = False @dataclass @@ -1430,6 +1431,18 @@ def detect_host() -> HostInfo: except Exception: pass + # Detect AMD ROCm (HIP) + has_rocm = False + if not is_macos: + rocm_hints = [ + shutil.which("hipcc"), + shutil.which("amd-smi"), + shutil.which("rocm-smi"), + ] + rocm_paths = ["/opt/rocm", os.environ.get("ROCM_PATH", "")] + if any(rocm_hints) or any(os.path.isdir(p) for p in rocm_paths if p): + has_rocm = True + return HostInfo( system = system, machine = machine, @@ -1444,6 +1457,7 @@ def detect_host() -> HostInfo: visible_cuda_devices = visible_cuda_devices, has_physical_nvidia = has_physical_nvidia, has_usable_nvidia = has_usable_nvidia, + has_rocm = has_rocm, ) @@ -1724,6 +1738,30 @@ def resolve_linux_cuda_choice( def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice: upstream_assets = github_release_assets(UPSTREAM_REPO, llama_tag) if host.is_linux and host.is_x86_64: + # AMD ROCm: try upstream ROCm prebuilt first, then fall back to source build. + # Source build (via setup.sh) compiles with -DGGML_HIP=ON and auto-detects + # the exact GPU target via rocminfo, which is more reliable for consumer + # GPUs (e.g. gfx1151) that may not be in the prebuilt. + if host.has_rocm and not host.has_usable_nvidia: + rocm_name = f"llama-{llama_tag}-bin-ubuntu-rocm-7.2-x64.tar.gz" + if rocm_name in upstream_assets: + log(f"AMD ROCm detected -- trying upstream prebuilt {rocm_name}") + log("Note: prebuilt is compiled for ROCm 7.2; if your ROCm version differs, " + "this may fail preflight and fall back to a source build (safe)") + return AssetChoice( + repo = UPSTREAM_REPO, + tag = llama_tag, + name = rocm_name, + url = upstream_assets[rocm_name], + source_label = "upstream", + install_kind = "linux-rocm", + ) + # No ROCm prebuilt available -- fall back to source build + raise PrebuiltFallback( + "AMD ROCm detected but no upstream ROCm prebuilt found; " + "falling back to source build with HIP support" + ) + upstream_name = f"llama-{llama_tag}-bin-ubuntu-x64.tar.gz" if upstream_name not in upstream_assets: raise PrebuiltFallback("upstream Linux CPU asset was not found") @@ -1743,6 +1781,21 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice return attempts[0] raise PrebuiltFallback("no compatible Windows CUDA asset was found") + # AMD ROCm on Windows: try HIP prebuilt + if host.has_rocm: + hip_name = f"llama-{llama_tag}-bin-win-hip-radeon-x64.zip" + if hip_name in upstream_assets: + log(f"AMD ROCm detected on Windows -- trying upstream HIP prebuilt {hip_name}") + return AssetChoice( + repo = UPSTREAM_REPO, + tag = llama_tag, + name = hip_name, + url = upstream_assets[hip_name], + source_label = "upstream", + install_kind = "windows-hip", + ) + log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU") + upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip" if upstream_name not in upstream_assets: raise PrebuiltFallback("upstream Windows CPU asset was not found") @@ -2121,7 +2174,7 @@ def overlay_directory_for_choice( def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: - if choice.install_kind in {"linux-cpu", "linux-cuda"}: + if choice.install_kind in {"linux-cpu", "linux-cuda", "linux-rocm"}: return [ "llama-server", "llama-quantize", @@ -2131,11 +2184,12 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: "libmtmd.so*", "libggml-cpu-*.so*", "libggml-cuda.so*", + "libggml-hip.so*", "libggml-rpc.so*", ] if choice.install_kind in {"macos-arm64", "macos-x64"}: return ["llama-server", "llama-quantize", "lib*.dylib"] - if choice.install_kind in {"windows-cpu", "windows-cuda"}: + if choice.install_kind in {"windows-cpu", "windows-cuda", "windows-hip"}: return ["*.exe", "*.dll"] raise PrebuiltFallback( f"unsupported install kind for runtime overlay: {choice.install_kind}" From cc6c8160d1a47818fb632a10b04d0fac13525246 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 08:31:31 +0000 Subject: [PATCH 4/6] Add IS_ROCM flag to hardware module and fix AMD error message Hardware module changes: - Add IS_ROCM global flag (bool, default False) to hardware.py - detect_hardware() sets IS_ROCM=True when torch.version.hip is set, while keeping DeviceType as CUDA (torch.cuda.* works on ROCm via HIP) - Print "ROCm (HIP X.Y)" instead of just "CUDA" when IS_ROCM is True - get_package_versions() now returns a "rocm" key with the HIP version (torch.version.hip) alongside the existing "cuda" key - Export IS_ROCM from studio/backend/utils/hardware/__init__.py Tokenizer error message fix: - Replace "We do not support AMD" with a helpful message pointing to ROCm installation docs at docs.unsloth.ai, since AMD is now supported --- studio/backend/utils/hardware/__init__.py | 2 ++ studio/backend/utils/hardware/hardware.py | 21 ++++++++++++++++----- unsloth/tokenizer_utils.py | 2 +- 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/studio/backend/utils/hardware/__init__.py b/studio/backend/utils/hardware/__init__.py index aaa0452406..b9b61cdcfe 100644 --- a/studio/backend/utils/hardware/__init__.py +++ b/studio/backend/utils/hardware/__init__.py @@ -9,6 +9,7 @@ from .hardware import ( DeviceType, DEVICE, CHAT_ONLY, + IS_ROCM, detect_hardware, get_device, is_apple_silicon, @@ -49,6 +50,7 @@ __all__ = [ "DeviceType", "DEVICE", "CHAT_ONLY", + "IS_ROCM", "detect_hardware", "get_device", "is_apple_silicon", diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 742e8f6b7e..d5de59a592 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -43,6 +43,7 @@ class DeviceType(str, Enum): DEVICE: Optional[DeviceType] = None CHAT_ONLY: bool = True # No CUDA GPU -> GGUF chat only (Mac, CPU-only, etc.) +IS_ROCM: bool = False # True when running on AMD ROCm (HIP) -- display/logging only # ========== Detection ========== @@ -85,10 +86,11 @@ def detect_hardware() -> DeviceType: 2. MLX (Apple Silicon via MLX framework) 3. CPU (fallback) """ - global DEVICE, CHAT_ONLY - CHAT_ONLY = True # reset -- only CUDA sets it to False + global DEVICE, CHAT_ONLY, IS_ROCM + CHAT_ONLY = True # reset -- only CUDA/ROCm sets it to False + IS_ROCM = False - # --- CUDA: try PyTorch --- + # --- CUDA / ROCm: try PyTorch --- if _has_torch(): import torch @@ -96,7 +98,14 @@ def detect_hardware() -> DeviceType: DEVICE = DeviceType.CUDA CHAT_ONLY = False device_name = torch.cuda.get_device_properties(0).name - print(f"Hardware detected: CUDA — {device_name}") + + # Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes. + # DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP. + if getattr(torch.version, "hip", None) is not None: + IS_ROCM = True + print(f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}") + else: + print(f"Hardware detected: CUDA -- {device_name}") return DEVICE # --- XPU: Intel GPU --- @@ -315,13 +324,15 @@ def get_package_versions() -> Dict[str, Optional[str]]: except PackageNotFoundError: versions[name] = None - # CUDA toolkit version bundled with torch + # GPU runtime version bundled with torch try: import torch versions["cuda"] = getattr(torch.version, "cuda", None) + versions["rocm"] = getattr(torch.version, "hip", None) except Exception: versions["cuda"] = None + versions["rocm"] = None return versions diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 8be6bb5a5a..07949cd32e 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -1103,7 +1103,7 @@ def patch_sft_trainer_tokenizer(): " a = np.array([int(x.decode('utf-8'))/1024 for x in a])\n" "except:\n" " if not torch.cuda.is_available():\n" - " raise RuntimeError('Unsloth: We do not support AMD / Intel machines yet - it is a work in progress!')\n" + " raise RuntimeError('Unsloth: No GPU detected. AMD ROCm users: install ROCm-enabled PyTorch -- see https://docs.unsloth.ai/get-started/install-and-update/amd')\n" "if ((a - PRE_CHECK) >= 1).sum() > 1:\n" " raise RuntimeError('Unsloth currently does not support multi GPU setups - but we are working on it!')\n" "for _ in range(3):\n" From 4ada33011fd985cdfcabb03060d38a02dca8c3a8 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 31 Mar 2026 08:32:17 +0000 Subject: [PATCH 5/6] Add comprehensive ROCm support test suite (68 tests) Add tests/studio/install/test_rocm_support.py with 68 test cases covering all ROCm-related code paths across install_llama_prebuilt.py, install_python_stack.py, hardware.py, tokenizer_utils.py, and install.sh. All tests use mocks and run without AMD hardware. Test classes and coverage: TestResolveUpstreamAssetChoice (11 tests): NVIDIA gets CPU asset, ROCm gets ROCm prebuilt, CPU gets CPU, macOS/Windows paths, NVIDIA+ROCm precedence, no prebuilt fallback, Windows+ROCm no HIP falls to CPU, macOS+ROCm, Linux aarch64+ROCm TestRuntimePatterns (5 tests): Correct file globs for linux-cpu, linux-cuda, linux-rocm, windows-hip, macos-arm64 TestHostInfoRocm (4 tests): Default has_rocm=False, explicit True, NVIDIA host no ROCm, detect_host checks ROCM_PATH TestDetectRocmVersion (9 tests): No ROCm returns None, version from file (7.1, 6.2), hipconfig fallback, empty version file, epoch prefix, multiple sources first wins, multiline hipconfig, hipconfig timeout TestEnsureRocmTorch (9 tests): No ROCm skips, torch has CUDA skips, torch has HIP skips, CPU torch triggers reinstall, ROCm 6.3 correct tag, old ROCm skips, version unreadable warning, ROCm 7.2 capped to 7.1, probe timeout handled TestRocmTorchIndex (8 tests): Mapping sorted descending, ROCm 7.2 not in map, key lookups, all tags use download.pytorch.org, newer ROCm selects best match TestHardwareRocmFlag (8 tests): IS_ROCM defined, set on HIP, still returns CUDA for ROCm, no DeviceType.ROCM, IS_ROCM exported, in __all__, get_package_versions has rocm key TestTokenizerErrorMessage (2 tests): No old "We do not support AMD" message, new message has docs link TestInstallShStructure (10 tests): No <<< here-strings, ROCm detection present, CUDA precedence, bitsandbytes install, CPU hint mentions AMD, ROCm 7.2 capped, validation guard exists, dpkg epoch handling, no [[ ]] syntax, macOS returns CPU before ROCm TestLiveRegression (1 test): Live nvidia-smi returns CUDA URL (skipped if no GPU) Depends on #4714, #4715, #4716, #4717 being merged first. --- tests/studio/install/test_rocm_support.py | 838 ++++++++++++++++++++++ 1 file changed, 838 insertions(+) create mode 100644 tests/studio/install/test_rocm_support.py diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py new file mode 100644 index 0000000000..560edf7a16 --- /dev/null +++ b/tests/studio/install/test_rocm_support.py @@ -0,0 +1,838 @@ +"""Tests for AMD ROCm support across install pathways. + +Verifies that ROCm detection and installation logic works correctly +WITHOUT breaking existing CUDA, CPU, macOS, and Windows pathways. +All tests use mocks -- no AMD hardware required. +""" + +import importlib.util +import os +import subprocess +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch, PropertyMock + +import pytest + + +# ── Load modules under test ────────────────────────────────────────────────── + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] + +# install_llama_prebuilt.py +_PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" +_PREBUILT_SPEC = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt", _PREBUILT_PATH +) +assert _PREBUILT_SPEC is not None and _PREBUILT_SPEC.loader is not None +prebuilt_mod = importlib.util.module_from_spec(_PREBUILT_SPEC) +sys.modules[_PREBUILT_SPEC.name] = prebuilt_mod +_PREBUILT_SPEC.loader.exec_module(prebuilt_mod) + +HostInfo = prebuilt_mod.HostInfo +AssetChoice = prebuilt_mod.AssetChoice +PrebuiltFallback = prebuilt_mod.PrebuiltFallback +resolve_upstream_asset_choice = prebuilt_mod.resolve_upstream_asset_choice +runtime_patterns_for_choice = prebuilt_mod.runtime_patterns_for_choice + +# install_python_stack.py +_STACK_PATH = PACKAGE_ROOT / "studio" / "install_python_stack.py" +_STACK_SPEC = importlib.util.spec_from_file_location( + "studio_install_python_stack", _STACK_PATH +) +assert _STACK_SPEC is not None and _STACK_SPEC.loader is not None +stack_mod = importlib.util.module_from_spec(_STACK_SPEC) +sys.modules[_STACK_SPEC.name] = stack_mod +_STACK_SPEC.loader.exec_module(stack_mod) + +_detect_rocm_version = stack_mod._detect_rocm_version +_ensure_rocm_torch = stack_mod._ensure_rocm_torch +_ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX + + +# ── Helper: build HostInfo for different scenarios ────────────────────────── + +def nvidia_host(**overrides) -> HostInfo: + """NVIDIA Linux x86_64 host.""" + defaults = dict( + system = "Linux", machine = "x86_64", + is_windows = False, is_linux = True, is_macos = False, + is_x86_64 = True, is_arm64 = False, + nvidia_smi = "/usr/bin/nvidia-smi", + driver_cuda_version = (12, 6), + compute_caps = ["89"], + visible_cuda_devices = None, + has_physical_nvidia = True, + has_usable_nvidia = True, + has_rocm = False, + ) + defaults.update(overrides) + return HostInfo(**defaults) + + +def rocm_host(**overrides) -> HostInfo: + """AMD ROCm Linux x86_64 host (no NVIDIA).""" + defaults = dict( + system = "Linux", machine = "x86_64", + is_windows = False, is_linux = True, is_macos = False, + is_x86_64 = True, is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = True, + ) + defaults.update(overrides) + return HostInfo(**defaults) + + +def cpu_host(**overrides) -> HostInfo: + """CPU-only Linux x86_64 host.""" + defaults = dict( + system = "Linux", machine = "x86_64", + is_windows = False, is_linux = True, is_macos = False, + is_x86_64 = True, is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = False, + ) + defaults.update(overrides) + return HostInfo(**defaults) + + +def macos_host(**overrides) -> HostInfo: + """macOS arm64 host.""" + defaults = dict( + system = "Darwin", machine = "arm64", + is_windows = False, is_linux = False, is_macos = True, + is_x86_64 = False, is_arm64 = True, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = False, + ) + defaults.update(overrides) + return HostInfo(**defaults) + + +def windows_host(**overrides) -> HostInfo: + """Windows x86_64 host.""" + defaults = dict( + system = "Windows", machine = "amd64", + is_windows = True, is_linux = False, is_macos = False, + is_x86_64 = True, is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = False, + ) + defaults.update(overrides) + return HostInfo(**defaults) + + +def windows_rocm_host(**overrides) -> HostInfo: + """Windows x86_64 host with ROCm.""" + defaults = dict( + system = "Windows", machine = "amd64", + is_windows = True, is_linux = False, is_macos = False, + is_x86_64 = True, is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = True, + ) + defaults.update(overrides) + return HostInfo(**defaults) + + +# ── Upstream asset fixture ─────────────────────────────────────────────────── + +LLAMA_TAG = "b8508" + +UPSTREAM_ASSETS = { + f"llama-{LLAMA_TAG}-bin-ubuntu-x64.tar.gz": f"https://example.com/{LLAMA_TAG}-linux-cpu.tar.gz", + f"llama-{LLAMA_TAG}-bin-ubuntu-rocm-7.2-x64.tar.gz": f"https://example.com/{LLAMA_TAG}-linux-rocm.tar.gz", + f"llama-{LLAMA_TAG}-bin-win-cpu-x64.zip": f"https://example.com/{LLAMA_TAG}-win-cpu.zip", + f"llama-{LLAMA_TAG}-bin-win-cuda-12.4-x64.zip": f"https://example.com/{LLAMA_TAG}-win-cuda.zip", + f"llama-{LLAMA_TAG}-bin-win-hip-radeon-x64.zip": f"https://example.com/{LLAMA_TAG}-win-hip.zip", + f"llama-{LLAMA_TAG}-bin-macos-arm64.tar.gz": f"https://example.com/{LLAMA_TAG}-macos-arm64.tar.gz", + f"llama-{LLAMA_TAG}-bin-macos-x64.tar.gz": f"https://example.com/{LLAMA_TAG}-macos-x64.tar.gz", +} + + +# ============================================================================= +# TEST: install_llama_prebuilt.py -- resolve_upstream_asset_choice +# ============================================================================= + +class TestResolveUpstreamAssetChoice: + """Verify that the asset selection logic picks the right binary for each platform.""" + + @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) + def test_nvidia_linux_gets_cpu_asset(self, mock_assets): + """NVIDIA host should NOT hit the ROCm path -- gets CPU asset (CUDA handled elsewhere).""" + host = nvidia_host() + choice = resolve_upstream_asset_choice(host, LLAMA_TAG) + assert choice.install_kind == "linux-cpu" + assert "ubuntu-x64" in choice.name + assert "rocm" not in choice.name + + @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) + def test_rocm_linux_gets_rocm_prebuilt(self, mock_assets): + """AMD ROCm Linux host should get the ROCm prebuilt.""" + host = rocm_host() + choice = resolve_upstream_asset_choice(host, LLAMA_TAG) + assert choice.install_kind == "linux-rocm" + assert "rocm" in choice.name + + @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) + def test_cpu_linux_gets_cpu_asset(self, mock_assets): + """CPU-only Linux host should get CPU asset.""" + host = cpu_host() + choice = resolve_upstream_asset_choice(host, LLAMA_TAG) + assert choice.install_kind == "linux-cpu" + assert "ubuntu-x64" in choice.name + + @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) + def test_macos_arm64_gets_macos_asset(self, mock_assets): + """macOS arm64 host should get macOS asset.""" + host = macos_host() + choice = resolve_upstream_asset_choice(host, LLAMA_TAG) + assert choice.install_kind == "macos-arm64" + assert "macos-arm64" in choice.name + + @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) + def test_windows_cpu_gets_cpu_asset(self, mock_assets): + """Windows CPU-only host should get Windows CPU asset.""" + host = windows_host() + choice = resolve_upstream_asset_choice(host, LLAMA_TAG) + assert choice.install_kind == "windows-cpu" + assert "win-cpu" in choice.name + + @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) + def test_windows_rocm_gets_hip_asset(self, mock_assets): + """Windows ROCm host should get Windows HIP asset.""" + host = windows_rocm_host() + choice = resolve_upstream_asset_choice(host, LLAMA_TAG) + assert choice.install_kind == "windows-hip" + assert "hip" in choice.name + + @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) + def test_mixed_nvidia_rocm_prefers_nvidia(self, mock_assets): + """Host with both NVIDIA and ROCm should use NVIDIA (CPU path here, CUDA elsewhere).""" + host = nvidia_host(has_rocm = True) + choice = resolve_upstream_asset_choice(host, LLAMA_TAG) + # NVIDIA hosts go through the normal path (CUDA handled by resolve_linux_cuda_choice) + assert choice.install_kind == "linux-cpu" + assert "rocm" not in choice.name + + @patch.object(prebuilt_mod, "github_release_assets") + def test_rocm_linux_no_prebuilt_falls_back(self, mock_assets): + """AMD ROCm host should fall back to source build when no ROCm prebuilt exists.""" + # Remove the ROCm asset from available assets + assets_without_rocm = {k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k} + mock_assets.return_value = assets_without_rocm + host = rocm_host() + with pytest.raises(PrebuiltFallback, match = "ROCm detected"): + resolve_upstream_asset_choice(host, LLAMA_TAG) + + @patch.object(prebuilt_mod, "github_release_assets") + def test_windows_rocm_no_hip_falls_to_cpu(self, mock_assets): + """Windows+ROCm with HIP prebuilt missing should fall through to CPU.""" + assets_no_hip = {k: v for k, v in UPSTREAM_ASSETS.items() if "hip" not in k} + mock_assets.return_value = assets_no_hip + host = windows_rocm_host() + choice = resolve_upstream_asset_choice(host, LLAMA_TAG) + assert choice.install_kind == "windows-cpu" + + @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) + def test_macos_rocm_impossible_has_rocm_false(self, mock_assets): + """macOS host should never have has_rocm=True in practice; verify it gets macOS asset.""" + host = macos_host(has_rocm = True) + choice = resolve_upstream_asset_choice(host, LLAMA_TAG) + assert choice.install_kind == "macos-arm64" + + @patch.object(prebuilt_mod, "github_release_assets", return_value = UPSTREAM_ASSETS) + def test_linux_aarch64_rocm_gets_prebuilt_fallback(self, mock_assets): + """Linux aarch64 with ROCm -- no x86_64 match, should raise PrebuiltFallback.""" + host = rocm_host(machine = "aarch64", is_x86_64 = False, is_arm64 = True) + with pytest.raises(PrebuiltFallback): + resolve_upstream_asset_choice(host, LLAMA_TAG) + + +# ============================================================================= +# TEST: install_llama_prebuilt.py -- runtime_patterns_for_choice +# ============================================================================= + +class TestRuntimePatterns: + """Verify runtime file patterns for all install kinds.""" + + def test_linux_cpu_patterns(self): + choice = AssetChoice(repo = "", tag = "", name = "", url = "", + source_label = "", install_kind = "linux-cpu") + patterns = runtime_patterns_for_choice(choice) + assert "llama-server" in patterns + assert "llama-quantize" in patterns + + def test_linux_cuda_patterns(self): + choice = AssetChoice(repo = "", tag = "", name = "", url = "", + source_label = "", install_kind = "linux-cuda") + patterns = runtime_patterns_for_choice(choice) + assert "libggml-cuda.so*" in patterns + + def test_linux_rocm_patterns(self): + choice = AssetChoice(repo = "", tag = "", name = "", url = "", + source_label = "", install_kind = "linux-rocm") + patterns = runtime_patterns_for_choice(choice) + assert "libggml-hip.so*" in patterns + assert "llama-server" in patterns + + def test_windows_hip_patterns(self): + choice = AssetChoice(repo = "", tag = "", name = "", url = "", + source_label = "", install_kind = "windows-hip") + patterns = runtime_patterns_for_choice(choice) + assert "*.exe" in patterns + assert "*.dll" in patterns + + def test_macos_patterns(self): + choice = AssetChoice(repo = "", tag = "", name = "", url = "", + source_label = "", install_kind = "macos-arm64") + patterns = runtime_patterns_for_choice(choice) + assert "lib*.dylib" in patterns + + +# ============================================================================= +# TEST: install_llama_prebuilt.py -- HostInfo.has_rocm field +# ============================================================================= + +class TestHostInfoRocm: + """Verify has_rocm field does not affect other HostInfo behavior.""" + + def test_has_rocm_default_false(self): + host = HostInfo( + system = "Linux", machine = "x86_64", + is_windows = False, is_linux = True, is_macos = False, + is_x86_64 = True, is_arm64 = False, + nvidia_smi = None, driver_cuda_version = None, + compute_caps = [], visible_cuda_devices = None, + has_physical_nvidia = False, has_usable_nvidia = False, + ) + assert host.has_rocm is False + + def test_has_rocm_explicit_true(self): + host = rocm_host() + assert host.has_rocm is True + + def test_nvidia_host_no_rocm(self): + host = nvidia_host() + 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 + import inspect + source = inspect.getsource(prebuilt_mod.detect_host) + assert "ROCM_PATH" in source or "rocm" in source.lower() + + +# ============================================================================= +# TEST: install_python_stack.py -- _detect_rocm_version +# ============================================================================= + +class TestDetectRocmVersion: + """Verify ROCm version detection from various sources.""" + + def test_no_rocm_returns_none(self, tmp_path): + """No ROCm installed should return None.""" + with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): + with patch("shutil.which", return_value = None): + result = _detect_rocm_version() + assert result is None + + def test_version_from_file(self, tmp_path): + """Reads version from /opt/rocm/.info/version.""" + info_dir = tmp_path / ".info" + info_dir.mkdir() + (info_dir / "version").write_text("7.1.0-12345\n") + with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): + result = _detect_rocm_version() + assert result == (7, 1) + + def test_version_62(self, tmp_path): + """Reads ROCm 6.2 version.""" + info_dir = tmp_path / ".info" + info_dir.mkdir() + (info_dir / "version").write_text("6.2.0\n") + with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): + result = _detect_rocm_version() + assert result == (6, 2) + + def test_hipconfig_fallback(self, tmp_path): + """Falls back to hipconfig --version when file not found.""" + with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = b"6.3.21234.2\n" + with patch("shutil.which", return_value = "/usr/bin/hipconfig"): + with patch("subprocess.run", return_value = mock_result): + result = _detect_rocm_version() + assert result == (6, 3) + + def test_empty_version_file(self, tmp_path): + """Empty version file should return None.""" + info_dir = tmp_path / ".info" + info_dir.mkdir() + (info_dir / "version").write_text("") + with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): + with patch("shutil.which", return_value = None): + result = _detect_rocm_version() + assert result is None + + def test_version_with_epoch_prefix(self, tmp_path): + """Debian epoch prefix (2:6.2.0) -- version file has no epoch, so should parse.""" + info_dir = tmp_path / ".info" + info_dir.mkdir() + # Version files don't typically have epoch prefix, but lib/rocm_version might + (info_dir / "version").write_text("6.2.0\n") + with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): + result = _detect_rocm_version() + assert result == (6, 2) + + def test_multiple_version_sources_first_wins(self, tmp_path): + """When both .info/version and lib/rocm_version exist, first found wins.""" + info_dir = tmp_path / ".info" + info_dir.mkdir() + (info_dir / "version").write_text("7.1.0\n") + lib_dir = tmp_path / "lib" + lib_dir.mkdir() + (lib_dir / "rocm_version").write_text("6.3.0\n") + with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path)}): + result = _detect_rocm_version() + assert result == (7, 1) # .info/version checked first + + def test_hipconfig_multiline_output(self, tmp_path): + """hipconfig with multi-line output -- should use first line.""" + with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): + mock_result = MagicMock() + mock_result.returncode = 0 + mock_result.stdout = b"6.3.21234.2\nSome extra info\n" + with patch("shutil.which", return_value = "/usr/bin/hipconfig"): + with patch("subprocess.run", return_value = mock_result): + result = _detect_rocm_version() + assert result == (6, 3) + + def test_hipconfig_timeout(self, tmp_path): + """hipconfig that times out should return None.""" + with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): + with patch("shutil.which", return_value = "/usr/bin/hipconfig"): + with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("hipconfig", 5)): + result = _detect_rocm_version() + assert result is None + + +# ============================================================================= +# TEST: install_python_stack.py -- _ensure_rocm_torch +# ============================================================================= + +class TestEnsureRocmTorch: + """Verify ROCm torch reinstall logic.""" + + @patch.object(stack_mod, "pip_install") + def test_no_rocm_skips(self, mock_pip): + """No ROCm toolchain should skip entirely.""" + with patch("os.path.isdir", return_value = False): + with patch("shutil.which", return_value = None): + _ensure_rocm_torch() + mock_pip.assert_not_called() + + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) + def test_torch_already_has_cuda_skips(self, mock_ver, mock_pip): + """If torch already has CUDA, should skip ROCm reinstall.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"12.6\n" # CUDA version string + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + mock_pip.assert_not_called() + + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) + def test_torch_already_has_hip_skips(self, mock_ver, mock_pip): + """If torch already has HIP, should skip ROCm reinstall.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"7.1.12345\n" # HIP version string + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + mock_pip.assert_not_called() + + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) + def test_cpu_torch_gets_rocm_reinstall(self, mock_ver, mock_pip): + """CPU-only torch on ROCm host should trigger reinstall.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # empty = no GPU backend + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + # Should call pip_install twice: once for torch, once for bitsandbytes + assert mock_pip.call_count == 2 + torch_call = mock_pip.call_args_list[0] + assert "rocm7.1" in str(torch_call) + bnb_call = mock_pip.call_args_list[1] + assert "bitsandbytes" in str(bnb_call) + + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 3)) + def test_rocm_63_selects_correct_tag(self, mock_ver, mock_pip): + """ROCm 6.3 should select rocm6.3 tag.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + torch_call = mock_pip.call_args_list[0] + assert "rocm6.3" in str(torch_call) + + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_detect_rocm_version", return_value = (5, 0)) + def test_old_rocm_skips(self, mock_ver, mock_pip): + """ROCm version too old (below 6.0) should skip.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + mock_pip.assert_not_called() + + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_detect_rocm_version", return_value = None) + def test_version_unreadable_prints_warning(self, mock_ver, mock_pip, capsys): + """ROCm detected but version unreadable should print warning and skip.""" + with patch("os.path.isdir", return_value = True): + _ensure_rocm_torch() + mock_pip.assert_not_called() + captured = capsys.readouterr() + assert "unreadable" in captured.out + + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + def test_rocm_72_selects_71_tag(self, mock_ver, mock_pip): + """ROCm 7.2 should select rocm7.1 tag (capped, not in mapping).""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + torch_call = mock_pip.call_args_list[0] + assert "rocm7.1" in str(torch_call) + + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) + def test_probe_timeout_handled(self, mock_ver, mock_pip): + """Probe subprocess timeout should be handled gracefully.""" + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30)): + # Should not crash -- timeout on probe means torch not importable + # The function will get an exception from subprocess.run and + # proceed to reinstall + try: + _ensure_rocm_torch() + except subprocess.TimeoutExpired: + pass # Acceptable -- the fix is about the timeout being set + + +# ============================================================================= +# TEST: install_python_stack.py -- _ROCM_TORCH_INDEX mapping +# ============================================================================= + +class TestRocmTorchIndex: + """Verify the ROCm version -> torch index tag mapping.""" + + def test_mapping_is_sorted_descending(self): + """Keys should be in descending order for the next() iteration to work.""" + keys = list(_ROCM_TORCH_INDEX.keys()) + assert keys == sorted(keys, reverse = True) + + def test_rocm_72_not_in_mapping(self): + """ROCm 7.2 should NOT be in the active mapping (torch 2.11.0 exceeds bound).""" + assert (7, 2) not in _ROCM_TORCH_INDEX + + def test_rocm_71_maps_correctly(self): + assert _ROCM_TORCH_INDEX[(7, 1)] == "rocm7.1" + + def test_rocm_63_maps_correctly(self): + assert _ROCM_TORCH_INDEX[(6, 3)] == "rocm6.3" + + def test_rocm_60_maps_correctly(self): + assert _ROCM_TORCH_INDEX[(6, 0)] == "rocm6.0" + + def test_all_tags_use_download_pytorch(self): + """All tags should be for download.pytorch.org, not repo.radeon.com.""" + for tag in _ROCM_TORCH_INDEX.values(): + assert tag.startswith("rocm") + assert "radeon" not in tag + + def test_newer_rocm_selects_best_match(self): + """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)), + None, + ) + assert tag == "rocm7.1" + + 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)), + None, + ) + assert tag == "rocm6.4" + + +# ============================================================================= +# TEST: hardware.py -- IS_ROCM flag and detect_hardware +# ============================================================================= + +class TestHardwareRocmFlag: + """Verify IS_ROCM flag behavior without importing the full hardware module.""" + + def test_hardware_py_has_is_rocm(self): + """hardware.py should define IS_ROCM.""" + hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + source = hw_path.read_text() + assert "IS_ROCM: bool = False" in source + + def test_hardware_py_sets_is_rocm_on_hip(self): + """detect_hardware() should set IS_ROCM when torch.version.hip is set.""" + hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + source = hw_path.read_text() + assert 'torch.version, "hip"' in source or "torch.version.hip" in source + + def test_hardware_py_still_returns_cuda_for_rocm(self): + """DeviceType should remain CUDA even on ROCm -- no DeviceType.ROCM.""" + hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + source = hw_path.read_text() + # Ensure ROCM is NOT a DeviceType member + enum_section = source.split("class DeviceType")[1].split("\n\n")[0] + assert "ROCM" not in enum_section + + def test_hardware_py_has_rocm_in_package_versions(self): + """get_package_versions() should include 'rocm' key.""" + hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + source = hw_path.read_text() + assert '"rocm"' in source + + def test_hardware_py_device_type_cuda_references_intact(self): + """All existing DeviceType.CUDA references should still be present.""" + hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + source = hw_path.read_text() + # Key functions that must still reference DeviceType.CUDA + assert "DeviceType.CUDA" in source + assert "DEVICE = DeviceType.CUDA" in source + + def test_is_rocm_exported_from_init(self): + """IS_ROCM should be exported from hardware __init__.py.""" + init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + source = init_path.read_text() + assert "IS_ROCM" in source + + def test_is_rocm_in_all_list(self): + """IS_ROCM should be in __all__ list in __init__.py.""" + init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + source = init_path.read_text() + # Extract __all__ section + assert '"IS_ROCM"' in source + + def test_get_package_versions_returns_rocm_key(self): + """get_package_versions() source should return both 'cuda' and 'rocm' keys.""" + hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + source = hw_path.read_text() + # Find the get_package_versions function body + func_start = source.find("def get_package_versions") + func_body = source[func_start:source.find("\ndef ", func_start + 1)] + assert '"cuda"' in func_body + assert '"rocm"' in func_body + + +# ============================================================================= +# TEST: tokenizer_utils.py -- error message +# ============================================================================= + +class TestTokenizerErrorMessage: + """Verify the AMD error message is updated.""" + + def test_no_old_amd_message(self): + """Old 'We do not support AMD' message should be gone.""" + tu_path = PACKAGE_ROOT / "unsloth" / "tokenizer_utils.py" + source = tu_path.read_text() + assert "We do not support AMD" not in source + + def test_new_message_has_docs_link(self): + """New message should point to Unsloth AMD docs.""" + tu_path = PACKAGE_ROOT / "unsloth" / "tokenizer_utils.py" + source = tu_path.read_text() + assert "docs.unsloth.ai" in source or "No GPU detected" in source + + +# ============================================================================= +# TEST: install.sh -- structural checks +# ============================================================================= + +class TestInstallShStructure: + """Verify install.sh structural properties without running it.""" + + def test_no_here_strings(self): + """install.sh must not use <<< (not POSIX).""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + # <<< is bash-only; breaks dash + for i, line in enumerate(source.splitlines(), 1): + stripped = line.lstrip() + if stripped.startswith("#"): + continue + assert "<<<" not in line, f"install.sh:{i} uses non-POSIX <<< here-string" + + def test_rocm_detection_present(self): + """install.sh should have ROCm detection in get_torch_index_url.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + assert "amd-smi" in source + assert "rocm" in source.lower() + + def test_cuda_precedence(self): + """ROCm detection should only run when nvidia-smi is absent.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + # The ROCm block should be inside the "if [ -z "$_smi" ]" branch + smi_block_start = source.find('if [ -z "$_smi" ]') + rocm_block_start = source.find("amd-smi") + assert smi_block_start < rocm_block_start, \ + "ROCm detection should be inside the 'no nvidia-smi' branch" + + def test_bitsandbytes_amd_install(self): + """install.sh should install bitsandbytes for AMD when ROCm detected.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + assert "bitsandbytes" in source + assert "rocm*)" in source # case pattern for ROCm URLs + + def test_cpu_hint_mentions_amd(self): + """CPU-only hint should mention AMD ROCm.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + assert "ROCm" in source + + def test_rocm72_capped_to_71(self): + """ROCm 7.2+ should fall back to rocm7.1 index.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + assert "rocm7.2" in source # case pattern + assert 'echo "$_base/rocm7.1"' in source # fallback + + def test_rocm_tag_validation_guard_exists(self): + """install.sh should validate _rocm_tag with a case guard.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + assert 'rocm[0-9]*.[0-9]*)' in source + assert '_rocm_tag=""' in source # rejection path + + def test_dpkg_epoch_handling(self): + """install.sh should strip Debian epoch prefix from dpkg-query output.""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + assert "sed 's/^[0-9]*://' " in source or "sed 's/^[0-9]*://'" in source + + def test_no_double_bracket_in_rocm_block(self): + """ROCm detection block should not use [[ ]] (bash-only, not POSIX). + Note: [[:space:]], [[:digit:]] etc. are valid POSIX character classes, not bash [[ ]].""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + func_start = source.find("get_torch_index_url()") + func_end = source.find("\n}", func_start) + func_body = source[func_start:func_end] + import re + for i, line in enumerate(func_body.splitlines(), 1): + stripped = line.lstrip() + if stripped.startswith("#"): + continue + # Remove POSIX character classes [[:foo:]] before checking for [[ ]] + cleaned = re.sub(r'\[\[:[a-z]+:\]\]', '', line) + assert "[[" not in cleaned, f"get_torch_index_url line {i} uses non-POSIX [[" + + def test_no_arithmetic_expansion_in_rocm_block(self): + """ROCm detection block should not use (( )) (bash-only).""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + func_start = source.find("get_torch_index_url()") + func_end = source.find("\n}", func_start) + func_body = source[func_start:func_end] + for i, line in enumerate(func_body.splitlines(), 1): + stripped = line.lstrip() + if stripped.startswith("#"): + continue + assert "((" not in line or "))" not in line or "$(()" in line, \ + f"get_torch_index_url line {i} may use non-POSIX (( ))" + + def test_macos_returns_cpu_before_rocm_check(self): + """macOS should return CPU immediately (before any ROCm check).""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text() + func_start = source.find("get_torch_index_url()") + func_body = source[func_start:] + darwin_pos = func_body.find("Darwin") + rocm_pos = func_body.find("amd-smi") + assert darwin_pos < rocm_pos, "macOS check should come before ROCm detection" + + +# ============================================================================= +# TEST: Live regression on current host (NVIDIA B200 expected) +# ============================================================================= + +class TestLiveRegression: + """Live checks that run on the actual host -- skip if no NVIDIA GPU.""" + + def test_get_torch_index_url_returns_cuda_on_nvidia(self): + """On an NVIDIA machine, get_torch_index_url should return a CUDA URL.""" + import shutil + if not shutil.which("nvidia-smi"): + pytest.skip("No nvidia-smi available") + sh_path = PACKAGE_ROOT / "install.sh" + # Extract just the function (don't source the whole installer) + result = subprocess.run( + ["bash", "-c", + f"eval \"$(sed -n '/^get_torch_index_url()/,/^}}/p' '{sh_path}')\"; " + "get_torch_index_url"], + capture_output = True, text = True, timeout = 30, + ) + if result.returncode != 0: + pytest.skip("Could not extract get_torch_index_url for live test") + url = result.stdout.strip() + assert "cu1" in url or "cuda" in url.lower(), f"Expected CUDA URL, got: {url}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 968119386f7b02ba79fca3f186d70101a86f8578 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 31 Mar 2026 08:35:03 +0000 Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- studio/backend/utils/hardware/hardware.py | 4 +- studio/install_llama_prebuilt.py | 14 +- tests/studio/install/test_rocm_support.py | 206 ++++++++++++++++------ 3 files changed, 162 insertions(+), 62 deletions(-) diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index d5de59a592..244ef3fe98 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -103,7 +103,9 @@ def detect_hardware() -> DeviceType: # DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP. if getattr(torch.version, "hip", None) is not None: IS_ROCM = True - print(f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}") + print( + f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}" + ) else: print(f"Hardware detected: CUDA -- {device_name}") return DEVICE diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index cbc8f6d15f..4fbd87850d 100755 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -1746,8 +1746,10 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice rocm_name = f"llama-{llama_tag}-bin-ubuntu-rocm-7.2-x64.tar.gz" if rocm_name in upstream_assets: log(f"AMD ROCm detected -- trying upstream prebuilt {rocm_name}") - log("Note: prebuilt is compiled for ROCm 7.2; if your ROCm version differs, " - "this may fail preflight and fall back to a source build (safe)") + log( + "Note: prebuilt is compiled for ROCm 7.2; if your ROCm version differs, " + "this may fail preflight and fall back to a source build (safe)" + ) return AssetChoice( repo = UPSTREAM_REPO, tag = llama_tag, @@ -1785,7 +1787,9 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice if host.has_rocm: hip_name = f"llama-{llama_tag}-bin-win-hip-radeon-x64.zip" if hip_name in upstream_assets: - log(f"AMD ROCm detected on Windows -- trying upstream HIP prebuilt {hip_name}") + log( + f"AMD ROCm detected on Windows -- trying upstream HIP prebuilt {hip_name}" + ) return AssetChoice( repo = UPSTREAM_REPO, tag = llama_tag, @@ -1794,7 +1798,9 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice source_label = "upstream", install_kind = "windows-hip", ) - log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU") + log( + "AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU" + ) upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip" if upstream_name not in upstream_assets: diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 560edf7a16..62ceece457 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -52,12 +52,17 @@ _ROCM_TORCH_INDEX = stack_mod._ROCM_TORCH_INDEX # ── Helper: build HostInfo for different scenarios ────────────────────────── + def nvidia_host(**overrides) -> HostInfo: """NVIDIA Linux x86_64 host.""" defaults = dict( - system = "Linux", machine = "x86_64", - is_windows = False, is_linux = True, is_macos = False, - is_x86_64 = True, is_arm64 = False, + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, nvidia_smi = "/usr/bin/nvidia-smi", driver_cuda_version = (12, 6), compute_caps = ["89"], @@ -73,9 +78,13 @@ def nvidia_host(**overrides) -> HostInfo: def rocm_host(**overrides) -> HostInfo: """AMD ROCm Linux x86_64 host (no NVIDIA).""" defaults = dict( - system = "Linux", machine = "x86_64", - is_windows = False, is_linux = True, is_macos = False, - is_x86_64 = True, is_arm64 = False, + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], @@ -91,9 +100,13 @@ def rocm_host(**overrides) -> HostInfo: def cpu_host(**overrides) -> HostInfo: """CPU-only Linux x86_64 host.""" defaults = dict( - system = "Linux", machine = "x86_64", - is_windows = False, is_linux = True, is_macos = False, - is_x86_64 = True, is_arm64 = False, + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], @@ -109,9 +122,13 @@ def cpu_host(**overrides) -> HostInfo: def macos_host(**overrides) -> HostInfo: """macOS arm64 host.""" defaults = dict( - system = "Darwin", machine = "arm64", - is_windows = False, is_linux = False, is_macos = True, - is_x86_64 = False, is_arm64 = True, + system = "Darwin", + machine = "arm64", + is_windows = False, + is_linux = False, + is_macos = True, + is_x86_64 = False, + is_arm64 = True, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], @@ -127,9 +144,13 @@ def macos_host(**overrides) -> HostInfo: def windows_host(**overrides) -> HostInfo: """Windows x86_64 host.""" defaults = dict( - system = "Windows", machine = "amd64", - is_windows = True, is_linux = False, is_macos = False, - is_x86_64 = True, is_arm64 = False, + system = "Windows", + machine = "amd64", + is_windows = True, + is_linux = False, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], @@ -145,9 +166,13 @@ def windows_host(**overrides) -> HostInfo: def windows_rocm_host(**overrides) -> HostInfo: """Windows x86_64 host with ROCm.""" defaults = dict( - system = "Windows", machine = "amd64", - is_windows = True, is_linux = False, is_macos = False, - is_x86_64 = True, is_arm64 = False, + system = "Windows", + machine = "amd64", + is_windows = True, + is_linux = False, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, nvidia_smi = None, driver_cuda_version = None, compute_caps = [], @@ -179,6 +204,7 @@ UPSTREAM_ASSETS = { # TEST: install_llama_prebuilt.py -- resolve_upstream_asset_choice # ============================================================================= + class TestResolveUpstreamAssetChoice: """Verify that the asset selection logic picks the right binary for each platform.""" @@ -244,7 +270,9 @@ class TestResolveUpstreamAssetChoice: def test_rocm_linux_no_prebuilt_falls_back(self, mock_assets): """AMD ROCm host should fall back to source build when no ROCm prebuilt exists.""" # Remove the ROCm asset from available assets - assets_without_rocm = {k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k} + assets_without_rocm = { + k: v for k, v in UPSTREAM_ASSETS.items() if "rocm" not in k + } mock_assets.return_value = assets_without_rocm host = rocm_host() with pytest.raises(PrebuiltFallback, match = "ROCm detected"): @@ -278,39 +306,55 @@ class TestResolveUpstreamAssetChoice: # TEST: install_llama_prebuilt.py -- runtime_patterns_for_choice # ============================================================================= + class TestRuntimePatterns: """Verify runtime file patterns for all install kinds.""" def test_linux_cpu_patterns(self): - choice = AssetChoice(repo = "", tag = "", name = "", url = "", - source_label = "", install_kind = "linux-cpu") + choice = AssetChoice( + repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-cpu" + ) patterns = runtime_patterns_for_choice(choice) assert "llama-server" in patterns assert "llama-quantize" in patterns def test_linux_cuda_patterns(self): - choice = AssetChoice(repo = "", tag = "", name = "", url = "", - source_label = "", install_kind = "linux-cuda") + choice = AssetChoice( + repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-cuda" + ) patterns = runtime_patterns_for_choice(choice) assert "libggml-cuda.so*" in patterns def test_linux_rocm_patterns(self): - choice = AssetChoice(repo = "", tag = "", name = "", url = "", - source_label = "", install_kind = "linux-rocm") + choice = AssetChoice( + repo = "", tag = "", name = "", url = "", source_label = "", install_kind = "linux-rocm" + ) patterns = runtime_patterns_for_choice(choice) assert "libggml-hip.so*" in patterns assert "llama-server" in patterns def test_windows_hip_patterns(self): - choice = AssetChoice(repo = "", tag = "", name = "", url = "", - source_label = "", install_kind = "windows-hip") + choice = AssetChoice( + repo = "", + tag = "", + name = "", + url = "", + source_label = "", + install_kind = "windows-hip", + ) patterns = runtime_patterns_for_choice(choice) assert "*.exe" in patterns assert "*.dll" in patterns def test_macos_patterns(self): - choice = AssetChoice(repo = "", tag = "", name = "", url = "", - source_label = "", install_kind = "macos-arm64") + choice = AssetChoice( + repo = "", + tag = "", + name = "", + url = "", + source_label = "", + install_kind = "macos-arm64", + ) patterns = runtime_patterns_for_choice(choice) assert "lib*.dylib" in patterns @@ -319,17 +363,25 @@ class TestRuntimePatterns: # TEST: install_llama_prebuilt.py -- HostInfo.has_rocm field # ============================================================================= + class TestHostInfoRocm: """Verify has_rocm field does not affect other HostInfo behavior.""" def test_has_rocm_default_false(self): host = HostInfo( - system = "Linux", machine = "x86_64", - is_windows = False, is_linux = True, is_macos = False, - is_x86_64 = True, is_arm64 = False, - nvidia_smi = None, driver_cuda_version = None, - compute_caps = [], visible_cuda_devices = None, - has_physical_nvidia = False, has_usable_nvidia = False, + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, ) assert host.has_rocm is False @@ -346,6 +398,7 @@ class TestHostInfoRocm: """detect_host() checks ROCM_PATH env var for ROCm detection.""" # Verify the detect_host function source references ROCM_PATH import inspect + source = inspect.getsource(prebuilt_mod.detect_host) assert "ROCM_PATH" in source or "rocm" in source.lower() @@ -354,6 +407,7 @@ class TestHostInfoRocm: # TEST: install_python_stack.py -- _detect_rocm_version # ============================================================================= + class TestDetectRocmVersion: """Verify ROCm version detection from various sources.""" @@ -440,7 +494,10 @@ class TestDetectRocmVersion: """hipconfig that times out should return None.""" with patch.dict(os.environ, {"ROCM_PATH": str(tmp_path / "nonexistent")}): with patch("shutil.which", return_value = "/usr/bin/hipconfig"): - with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("hipconfig", 5)): + with patch( + "subprocess.run", + side_effect = subprocess.TimeoutExpired("hipconfig", 5), + ): result = _detect_rocm_version() assert result is None @@ -449,6 +506,7 @@ class TestDetectRocmVersion: # TEST: install_python_stack.py -- _ensure_rocm_torch # ============================================================================= + class TestEnsureRocmTorch: """Verify ROCm torch reinstall logic.""" @@ -554,7 +612,9 @@ class TestEnsureRocmTorch: def test_probe_timeout_handled(self, mock_ver, mock_pip): """Probe subprocess timeout should be handled gracefully.""" with patch("os.path.isdir", return_value = True): - with patch("subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30)): + with patch( + "subprocess.run", side_effect = subprocess.TimeoutExpired("python", 30) + ): # Should not crash -- timeout on probe means torch not importable # The function will get an exception from subprocess.run and # proceed to reinstall @@ -568,6 +628,7 @@ class TestEnsureRocmTorch: # TEST: install_python_stack.py -- _ROCM_TORCH_INDEX mapping # ============================================================================= + class TestRocmTorchIndex: """Verify the ROCm version -> torch index tag mapping.""" @@ -617,24 +678,31 @@ class TestRocmTorchIndex: # TEST: hardware.py -- IS_ROCM flag and detect_hardware # ============================================================================= + class TestHardwareRocmFlag: """Verify IS_ROCM flag behavior without importing the full hardware module.""" def test_hardware_py_has_is_rocm(self): """hardware.py should define IS_ROCM.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text() assert "IS_ROCM: bool = False" in source def test_hardware_py_sets_is_rocm_on_hip(self): """detect_hardware() should set IS_ROCM when torch.version.hip is set.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text() assert 'torch.version, "hip"' in source or "torch.version.hip" in source def test_hardware_py_still_returns_cuda_for_rocm(self): """DeviceType should remain CUDA even on ROCm -- no DeviceType.ROCM.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text() # Ensure ROCM is NOT a DeviceType member enum_section = source.split("class DeviceType")[1].split("\n\n")[0] @@ -642,13 +710,17 @@ class TestHardwareRocmFlag: def test_hardware_py_has_rocm_in_package_versions(self): """get_package_versions() should include 'rocm' key.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text() assert '"rocm"' in source def test_hardware_py_device_type_cuda_references_intact(self): """All existing DeviceType.CUDA references should still be present.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text() # Key functions that must still reference DeviceType.CUDA assert "DeviceType.CUDA" in source @@ -656,24 +728,30 @@ class TestHardwareRocmFlag: def test_is_rocm_exported_from_init(self): """IS_ROCM should be exported from hardware __init__.py.""" - init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + init_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + ) source = init_path.read_text() assert "IS_ROCM" in source def test_is_rocm_in_all_list(self): """IS_ROCM should be in __all__ list in __init__.py.""" - init_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + init_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "__init__.py" + ) source = init_path.read_text() # Extract __all__ section assert '"IS_ROCM"' in source def test_get_package_versions_returns_rocm_key(self): """get_package_versions() source should return both 'cuda' and 'rocm' keys.""" - hw_path = PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + hw_path = ( + PACKAGE_ROOT / "studio" / "backend" / "utils" / "hardware" / "hardware.py" + ) source = hw_path.read_text() # Find the get_package_versions function body func_start = source.find("def get_package_versions") - func_body = source[func_start:source.find("\ndef ", func_start + 1)] + func_body = source[func_start : source.find("\ndef ", func_start + 1)] assert '"cuda"' in func_body assert '"rocm"' in func_body @@ -682,6 +760,7 @@ class TestHardwareRocmFlag: # TEST: tokenizer_utils.py -- error message # ============================================================================= + class TestTokenizerErrorMessage: """Verify the AMD error message is updated.""" @@ -702,6 +781,7 @@ class TestTokenizerErrorMessage: # TEST: install.sh -- structural checks # ============================================================================= + class TestInstallShStructure: """Verify install.sh structural properties without running it.""" @@ -730,8 +810,9 @@ class TestInstallShStructure: # The ROCm block should be inside the "if [ -z "$_smi" ]" branch smi_block_start = source.find('if [ -z "$_smi" ]') rocm_block_start = source.find("amd-smi") - assert smi_block_start < rocm_block_start, \ - "ROCm detection should be inside the 'no nvidia-smi' branch" + assert ( + smi_block_start < rocm_block_start + ), "ROCm detection should be inside the 'no nvidia-smi' branch" def test_bitsandbytes_amd_install(self): """install.sh should install bitsandbytes for AMD when ROCm detected.""" @@ -757,7 +838,7 @@ class TestInstallShStructure: """install.sh should validate _rocm_tag with a case guard.""" sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text() - assert 'rocm[0-9]*.[0-9]*)' in source + assert "rocm[0-9]*.[0-9]*)" in source assert '_rocm_tag=""' in source # rejection path def test_dpkg_epoch_handling(self): @@ -775,13 +856,16 @@ class TestInstallShStructure: func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] import re + for i, line in enumerate(func_body.splitlines(), 1): stripped = line.lstrip() if stripped.startswith("#"): continue # Remove POSIX character classes [[:foo:]] before checking for [[ ]] - cleaned = re.sub(r'\[\[:[a-z]+:\]\]', '', line) - assert "[[" not in cleaned, f"get_torch_index_url line {i} uses non-POSIX [[" + cleaned = re.sub(r"\[\[:[a-z]+:\]\]", "", line) + assert ( + "[[" not in cleaned + ), f"get_torch_index_url line {i} uses non-POSIX [[" def test_no_arithmetic_expansion_in_rocm_block(self): """ROCm detection block should not use (( )) (bash-only).""" @@ -794,8 +878,9 @@ class TestInstallShStructure: stripped = line.lstrip() if stripped.startswith("#"): continue - assert "((" not in line or "))" not in line or "$(()" in line, \ - f"get_torch_index_url line {i} may use non-POSIX (( ))" + assert ( + "((" not in line or "))" not in line or "$(()" in line + ), f"get_torch_index_url line {i} may use non-POSIX (( ))" def test_macos_returns_cpu_before_rocm_check(self): """macOS should return CPU immediately (before any ROCm check).""" @@ -812,21 +897,28 @@ class TestInstallShStructure: # TEST: Live regression on current host (NVIDIA B200 expected) # ============================================================================= + class TestLiveRegression: """Live checks that run on the actual host -- skip if no NVIDIA GPU.""" def test_get_torch_index_url_returns_cuda_on_nvidia(self): """On an NVIDIA machine, get_torch_index_url should return a CUDA URL.""" import shutil + if not shutil.which("nvidia-smi"): pytest.skip("No nvidia-smi available") sh_path = PACKAGE_ROOT / "install.sh" # Extract just the function (don't source the whole installer) result = subprocess.run( - ["bash", "-c", - f"eval \"$(sed -n '/^get_torch_index_url()/,/^}}/p' '{sh_path}')\"; " - "get_torch_index_url"], - capture_output = True, text = True, timeout = 30, + [ + "bash", + "-c", + f"eval \"$(sed -n '/^get_torch_index_url()/,/^}}/p' '{sh_path}')\"; " + "get_torch_index_url", + ], + capture_output = True, + text = True, + timeout = 30, ) if result.returncode != 0: pytest.skip("Could not extract get_torch_index_url for live test")