diff --git a/install.sh b/install.sh index 963107524b..d06fff07c9 100755 --- a/install.sh +++ b/install.sh @@ -2115,13 +2115,16 @@ _has_amd_rocm_gpu() { amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then return 0 elif [ -e /dev/kfd ] && \ - awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \ - gpu && amd { found=1 } END{ exit !found }' \ + awk '/vendor_id/ && $2 == 4098 { found = 1 } END { exit !found }' \ /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then - # vendor_id 4098 = 0x1002 (AMD). NVIDIA open kernel module (driver - # 560+) can register KFD topology nodes with non-zero gpu_id but - # vendor_id 4318 (0x10DE). Require AMD vendor to avoid misrouting - # NVIDIA-only hosts to the ROCm install path. + # vendor_id 4098 = 0x1002 (AMD) marks a GPU node: the KFD CPU node + # reports vendor_id 0, so any 4098 node is an AMD GPU. NVIDIA's open + # kernel module (driver 560+) registers KFD nodes as vendor_id 4318 + # (0x10DE), so this never false-positives on NVIDIA-only hosts. + # The prior check also required a gpu_id line, but gpu_id is a SIBLING + # sysfs file, not a line in properties -- it never matched, so the + # fallback silently missed every ROCm-less AMD host (issue: fresh + # Arch/CachyOS boxes reporting "no GPU detected"). return 0 fi return 1 @@ -2230,6 +2233,30 @@ EOF return 1 } +# Reads the AMD gfx arch for wheel-index decisions: a user-set +# UNSLOTH_ROCM_GFX_ARCH is authoritative (lowercased), else rocminfo, then +# amd-smi. rocminfo/amd-smi honor ROCR/HIP_VISIBLE_DEVICES, so a container mask +# (e.g. ROCR_VISIBLE_DEVICES=-1) would hide a GPU that the env-independent KFD +# detection still sees -- the tool probes run with the masks cleared. Prints the +# gfx token(s) or nothing when unreadable, and always returns 0 (a failing probe +# as the last command would trip set -e in callers' assignments). Shared by +# get_torch_index_url's gfx gate and the runtime-less reroute gate so the two +# can never disagree on what "readable" means. +_probe_amd_gfx_arch() { + _ensure_rocm_probe_env + _pg=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]') + if [ -z "$_pg" ] && command -v rocminfo >/dev/null 2>&1; then + _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + if [ -z "$_pg" ] && command -v amd-smi >/dev/null 2>&1; then + _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + if [ -z "$_pg" ]; then + _pg=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + fi + printf '%s\n' "$_pg" +} + # ── Detect GPU and choose PyTorch index URL ── # Mirrors Get-TorchIndexUrl in install.ps1. # On CPU-only machines this returns the cpu index, avoiding the solver @@ -2283,6 +2310,29 @@ get_torch_index_url() { if ! _has_amd_rocm_gpu; then echo "$_base/cpu"; return fi + # A generic rocm index is only safe when the gfx arch is readable: the + # Strix reroute (gfx1150/1151 -> arch-specific index) learns gfx from + # rocminfo/amd-smi, so if those are missing OR do not enumerate the GPU, an + # unknown-arch box might be Strix and would get the broken _grouped_mm + # wheels. Probe via the shared helper (override first, then rocminfo/amd-smi + # with visibility masks cleared); if the arch is unreadable, never guess a + # rocm index. A KFD-only host whose arch is still inferable from hardware + # IDs (PCI/cpuinfo/lspci) returns the cpu index and lets the runtime-less + # reroute below upgrade it to AMD per-arch wheels -- the reroute gate uses + # this same probe, so the handoff can't misfire. Only when inference fails + # too is CPU final, with the actionable warning. + _amd_gfx_probe=$(_probe_amd_gfx_arch) + if [ -z "$_amd_gfx_probe" ]; then + if _amd_inferred_gfx=$(_infer_linux_amd_gfx_arch 2>/dev/null) && \ + [ -n "$_amd_inferred_gfx" ] && \ + _amd_arch_index_family_for_gfx "$_amd_inferred_gfx" >/dev/null 2>&1; then + echo "[WARN] AMD GPU detected but rocminfo/amd-smi can't read its gfx arch -- inferring $_amd_inferred_gfx from hardware IDs." >&2 + echo "$_base/cpu"; return + fi + echo "[WARN] AMD GPU detected but its gfx arch can't be read (rocminfo/amd-smi missing or not enumerating the GPU) -- installing CPU-only PyTorch." >&2 + echo "[WARN] For GPU PyTorch, install or repair rocminfo/amd-smi (e.g. sudo pacman -S rocm-hip-sdk) and re-run this installer." >&2 + echo "$_base/cpu"; return + fi # AMD GPU confirmed -- detect ROCm version _rocm_tag="" _rocm_tag=$({ command -v amd-smi >/dev/null 2>&1 && \ @@ -2299,7 +2349,11 @@ get_torch_index_url() { { 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 + printf '%s\n' "$ver" | awk -F'[.-]' '{print "rocm"$1"."$2; exit}'; }) 2>/dev/null || _rocm_tag="" + # ^ || guard: when EVERY version source is missing (e.g. rocminfo present + # but rocm-core not installed, so dpkg-query/rpm exit 1), the whole || + # chain fails and set -e would kill the installer BEFORE the actionable + # no-version WARN below -- exactly the fresh-install case it exists for. # Validate _rocm_tag: must match "rocmX.Y" with major >= 1 case "$_rocm_tag" in rocm[1-9]*.[0-9]*) : ;; # valid (major >= 1) @@ -2335,12 +2389,27 @@ get_torch_index_url() { esac return fi - # AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be - # read from any source (amd-smi, /opt/rocm/.info/version, hipconfig, - # dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch. - echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2 - echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2 - echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + # AMD GPU confirmed (rocminfo/amd-smi or the KFD topology fallback) but + # no ROCm/HIP install was found to read the version from (amd-smi, + # /opt/rocm/.info/version, hipconfig, dpkg, rpm). This is the common + # fresh-install case: the GPU is real, but with no ROCm userspace the + # correct PyTorch build can't be selected. Warn with an actionable fix + # rather than silently installing CPU PyTorch. + # A user-set UNSLOTH_ROCM_GFX_ARCH seeded the probe above, so rocminfo/ + # amd-smi may still be unable to see the GPU; when the named arch maps to + # a wheel family, the runtime-less reroute (gated on the override) will + # install the AMD per-arch wheels -- a CPU-only warning here would be + # false for that path. Defer like the inferable-arch branch does. + if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] && \ + _amd_arch_index_family_for_gfx "$_amd_gfx_probe" >/dev/null 2>&1; then + echo "[WARN] AMD GPU detected with no readable ROCm version, but UNSLOTH_ROCM_GFX_ARCH=$_amd_gfx_probe is set -- routing to AMD per-arch wheels." >&2 + echo "$_base/cpu"; return + fi + echo "[WARN] AMD GPU detected, but no ROCm/HIP install was found to select the matching GPU PyTorch build -- falling back to CPU-only PyTorch." >&2 + echo "[WARN] Install the ROCm/HIP SDK, then re-run this installer:" >&2 + echo "[WARN] Arch / CachyOS : sudo pacman -S rocm-hip-sdk" >&2 + echo "[WARN] other distros : https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2 + echo "[WARN] Minimum required for version detection: amd-smi, hipconfig, /opt/rocm/.info/version, or the rocm-core package." >&2 echo "$_base/cpu"; return fi # Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P). @@ -2841,14 +2910,20 @@ TORCH_INDEX_URL=$(get_torch_index_url) # Linux: ROCm runtime missing but a supported AMD gfx arch is inferable (Strix Halo # in /proc/cpuinfo, lspci marketing name, UNSLOTH_ROCM_GFX_ARCH). Route to AMD's # per-arch wheels like install.ps1 does on Windows (unslothai#7301). -# Gated on _has_amd_rocm_gpu being FALSE: a */cpu index on a host whose GPU IS -# visible to the ROCm probes is a deliberate fallback (unsupported/unreadable -# ROCm version, after its own warning), not a missing runtime -- rerouting it -# would contradict that decision. An explicit UNSLOTH_ROCM_GFX_ARCH override -# stays authoritative either way. +# Gated on the runtime probes NOT naming a gfx: either no AMD GPU is detected at +# all (_has_amd_rocm_gpu false), or the GPU is visible only through the +# env-independent KFD topology while rocminfo/amd-smi can't read its arch +# (KFD-only host, unslothai#7314 -- before the KFD detection fix these hosts +# reached this reroute via the false branch, so the empty-probe condition +# preserves that routing). A */cpu index chosen WITH a readable gfx +# (unsupported/unreadable ROCm version, after its own warning) is a deliberate +# fallback -- rerouting it would contradict that decision, and stays excluded +# because the shared probe returns its gfx. An explicit UNSLOTH_ROCM_GFX_ARCH +# override stays authoritative either way. if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \ ! _has_usable_nvidia_gpu && \ - { [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu; } && \ + { [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ] || ! _has_amd_rocm_gpu || \ + [ -z "$(_probe_amd_gfx_arch)" ]; } && \ case "$(uname -s)" in Linux) true ;; *) false ;; esac && \ case "$_ARCH" in x86_64|amd64) true ;; *) false ;; esac; then # ROCm torch wheels are x86_64-only; get_torch_index_url returns CPU on other @@ -2880,7 +2955,13 @@ if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \ ;; esac echo "" >&2 - echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2 + # KFD-only hosts reach this reroute with /dev/kfd present + # (that's what detected them), so don't claim it's missing. + if _has_amd_rocm_gpu; then + echo " [WARN] AMD GPU visible via the kernel driver (KFD) but rocminfo/amd-smi can't read its gfx arch; using $_linux_inferred_gfx." >&2 + else + echo " [WARN] ROCm runtime not visible (/dev/kfd, rocminfo, amd-smi) but $_linux_inferred_gfx inferred." >&2 + fi echo " [WARN] Routing to AMD arch-specific wheels ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))." >&2 echo " [WARN] These wheels bundle their own ROCm runtime; install the kernel stack for native compute:" >&2 echo " [WARN] https://docs.unsloth.ai/get-started/install-and-update/amd" >&2 @@ -3004,8 +3085,10 @@ case "$_torch_index_leaf" in # || true on each probe: no gfx match makes grep exit 1, which under # set -euo pipefail would abort the installer before the next fallback # runs (now that the case matches every rocm* index, not just rocm7.1). - _gfx_all="" - if command -v rocminfo >/dev/null 2>&1; then + # A user-supplied UNSLOTH_ROCM_GFX_ARCH overrides probing (mirrors setup.sh + # and the display block), so a Strix override still reaches the arch index. + _gfx_all=$(printf '%s' "${UNSLOTH_ROCM_GFX_ARCH:-}" | tr '[:upper:]' '[:lower:]') + if [ -z "$_gfx_all" ] && command -v rocminfo >/dev/null 2>&1; then _gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) fi if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then @@ -3016,6 +3099,23 @@ case "$_torch_index_leaf" in _gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) fi fi + # get_torch_index_url reads the arch with ROCR/HIP masks cleared, so a + # mask hiding every agent (e.g. ROCR_VISIBLE_DEVICES=-1) still lands + # here on a generic rocm index; re-probe unmasked or a masked-out Strix + # box keeps the broken generic wheels. Partial masks never get here + # (they enumerate at least one agent above) and keep their selection. + # ${VAR+x} (not :-): a SET-but-empty mask also hides every agent and + # must trigger the re-probe too. + if [ -z "$_gfx_all" ] && [ -n "${ROCR_VISIBLE_DEVICES+x}${HIP_VISIBLE_DEVICES+x}" ]; then + if command -v rocminfo >/dev/null 2>&1; then + _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; rocminfo 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then + _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi list 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + [ -z "$_gfx_all" ] && \ + _gfx_all=$( (unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES; amd-smi static --asic 2>/dev/null) | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true) + fi + fi _runtime_gfx="" if [ -n "$_gfx_all" ]; then _vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}" @@ -3169,6 +3269,17 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then elif [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then # Apple Silicon: PyTorch gets Metal (MPS) acceleration over unified memory, so not CPU-only. step "gpu" "Apple Silicon (Metal, unified memory)" +elif _has_amd_rocm_gpu; then + if [ "$_torch_index_pinned" = true ]; then + # An explicit UNSLOTH_TORCH_INDEX_URL/_FAMILY pin skipped all probing; + # do not claim ROCm is unusable when a CPU/other index was requested. + step "gpu" "AMD GPU (torch index pinned: $_torch_index_leaf)" "$C_WARN" + else + # AMD GPU visible to the kernel but the torch index stayed CPU: no usable + # ROCm userspace to pick a wheel. "none" would repeat the false diagnosis + # this installer used to give. + step "gpu" "AMD GPU (no usable ROCm -- CPU fallback)" "$C_WARN" + fi else step "gpu" "none (CPU-only)" "$C_WARN" fi @@ -3177,8 +3288,17 @@ fi case "$TORCH_INDEX_URL" in */cpu) if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then - substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" - if [ "$OS" = "wsl" ]; then + if [ "$_torch_index_pinned" = true ]; then + # An explicit CPU pin is a request, not a detection failure: + # skip the SDK guidance (ROCm may be perfectly healthy here). + substep "CPU-only PyTorch (index pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY)." + elif _has_amd_rocm_gpu; then + substep "AMD GPU detected, but no usable ROCm/HIP install -- installing CPU-only PyTorch." "$C_WARN" + substep "Install the ROCm/HIP SDK and re-run this installer for GPU PyTorch." "$C_WARN" + else + substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN" + fi + if [ "$OS" = "wsl" ] && [ "$_torch_index_pinned" = false ]; then # WSL + no GPU detected (detection above found nothing). Common # cause: an AMD GPU whose ROCm-on-WSL runtime isn't exposed yet -- # /dev/dxg present (graphics) but no ROCm runtime. diff --git a/studio/setup.sh b/studio/setup.sh index 2a2b41d0f6..0183ef3776 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1101,8 +1101,7 @@ if [ "$_setup_nvidia_usable" != true ]; then _setup_mkt=$(_setup_run_smi amd-smi static --asic 2>/dev/null | awk -F'[:|]' \ '/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true) elif [ -e /dev/kfd ] && \ - awk 'FNR==1{ gpu=0; amd=0 } /gpu_id/{ gpu=($2+0>0) } /vendor_id/{ amd=($2==4098) } \ - gpu && amd { found=1 } END{ exit !found }' \ + awk '/vendor_id/ && $2 == 4098 { found = 1 } END { exit !found }' \ /sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then # KFD sysfs fallback, AMD vendor_id 4098 only (mirrors install.sh # _has_amd_rocm_gpu): covers AMD hosts where rocminfo/amd-smi are @@ -1358,9 +1357,14 @@ else # name-inferred arch). Implies --has-rocm on the installer side. if [ -n "${_setup_gfx:-}" ]; then _PREBUILT_CMD+=(--rocm-gfx "$_setup_gfx") - elif [ "$_setup_amd_detected" = true ]; then - # AMD was detected but gfx resolution failed; tell the installer ROCm is - # present so it can still attempt a prebuilt. Mirrors setup.ps1 behaviour. + elif [ "$_setup_amd_detected" = true ] && \ + { command -v hipcc >/dev/null 2>&1 || [ -x /opt/rocm/bin/hipcc ] || \ + ls /opt/rocm-*/bin/hipcc >/dev/null 2>&1; }; then + # AMD detected but gfx unknown (KFD-only host): forward --has-rocm only when + # hipcc can actually build llama.cpp (incl. a versioned /opt/rocm-*/bin, the + # same paths the source build uses). With no gfx the prebuilt resolver finds + # no ROCm bundle and the source build would fail, so without hipcc fall + # through to the CPU prebuilt instead of breaking the install. _PREBUILT_CMD+=(--has-rocm) fi # UNSLOTH_LLAMA_CPP_BACKEND=cpu (case-insensitive, trimmed) forces the CPU-only diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index cd7b68f4b6..fa76011041 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -1459,6 +1459,59 @@ class TestInstallShStructure: sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") assert "amd-smi" in source + + def test_cpu_index_note_respects_explicit_pin(self): + """An explicit UNSLOTH_TORCH_INDEX_URL/_FAMILY CPU pin is a request, not + a detection failure: the */cpu wheel note must report the pin instead of + claiming ROCm/HIP is unusable, the WSL setup guidance must be skipped, + and the gpu summary must not label a pinned AMD host "no usable ROCm".""" + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text(encoding = "utf-8") + note = source.find('substep "AMD GPU detected, but no usable ROCm/HIP install') + assert note != -1 + assert ( + '[ "$_torch_index_pinned" = true ]' in source[note - 400 : note] + ), "the */cpu note must check the explicit pin before diagnosing ROCm" + assert ( + '[ "$OS" = "wsl" ] && [ "$_torch_index_pinned" = false ]' in source + ), "ROCm-on-WSL guidance is detection advice; skip it for pinned installs" + summary = source.find('step "gpu" "AMD GPU (no usable ROCm -- CPU fallback)"') + assert summary != -1 + assert ( + '[ "$_torch_index_pinned" = true ]' in source[summary - 700 : summary] + ), "the gpu summary must not claim no usable ROCm for a pinned index" + + def test_rocm_version_chain_survives_no_source_under_set_e(self): + """When every ROCm version source is missing (e.g. rocminfo present but + rocm-core not installed, so dpkg-query/rpm exit 1), the _rocm_tag || + chain fails as a whole; without the || guard set -e kills the installer + BEFORE the actionable no-version WARN it feeds. Executed, not text.""" + shell = shutil.which("bash") + if not shell: + pytest.skip("bash needed to execute the version chain") + sh_path = PACKAGE_ROOT / "install.sh" + source = sh_path.read_text(encoding = "utf-8") + chain = re.search( + r'^ _rocm_tag=\$\(\{ command -v amd-smi.*?\|\| _rocm_tag=""\n', + source, + re.S | re.M, + ) + assert chain, "could not extract the guarded _rocm_tag chain" + with tempfile.TemporaryDirectory() as d: + # Tools exist on PATH but yield nothing usable, like a box with the + # probe tools installed and no rocm-core package. + for name in ("amd-smi", "hipconfig", "dpkg-query", "rpm"): + p = os.path.join(d, name) + with open(p, "w", encoding = "utf-8") as f: + f.write("#!/bin/sh\nexit 1\n") + os.chmod(p, 0o755) + script = ( + "set -euo pipefail\n" + chain.group(0) + '\nprintf "SURVIVED:%s\\n" "$_rocm_tag"\n' + ) + env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", "")) + r = subprocess.run([shell, "-c", script], env = env, capture_output = True, text = True) + assert r.returncode == 0, f"version chain aborted under set -e: {r.stderr}" + assert r.stdout.startswith("SURVIVED:"), r.stdout assert "rocm" in source.lower() def test_cuda_precedence(self): @@ -1590,17 +1643,446 @@ class TestInstallShStructure: "4098" in func_body ), "_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098 (0x1002)" - def test_kfd_awk_resets_state_per_file(self): - """KFD sysfs awk must reset gpu/amd state per file (FNR==1) to avoid Ryzen+NVIDIA false positives.""" + def test_kfd_awk_vendor_check_is_per_line(self): + """KFD sysfs awk must decide on a single vendor_id line, with no cross-node state. + + The old awk paired two per-node flags (gpu_id + vendor_id) and needed an FNR==1 + reset so flags from different KFD nodes could not combine into a Ryzen+NVIDIA + false positive. gpu_id is a sibling sysfs file and never appears inside + properties, so that pairing also never matched at all (every ROCm-less AMD host + was reported as no-GPU). The replacement keys on one atomic line: only an AMD + GPU node reports `vendor_id 4098` (KFD CPU nodes report 0, NVIDIA's open kernel + module registers 4318), so there is no cross-file state left to reset. + """ sh_path = PACKAGE_ROOT / "install.sh" source = sh_path.read_text(encoding = "utf-8") func_start = source.find("_has_amd_rocm_gpu()") func_end = source.find("\n}", func_start) func_body = source[func_start:func_end] - assert "FNR==1" in func_body, ( - "_has_amd_rocm_gpu KFD awk must reset state per file with FNR==1 " - "to avoid false positives on Ryzen+NVIDIA hosts with multiple KFD nodes" + assert "$2 == 4098" in func_body, ( + "_has_amd_rocm_gpu KFD awk must match `vendor_id 4098` as a single-line " + "condition so no per-node state can leak across KFD nodes" ) + assert "/gpu_id/" not in func_body, ( + "_has_amd_rocm_gpu KFD awk must not key on a gpu_id line: gpu_id is a " + "sibling sysfs file, not a line in properties, so it never matches there" + ) + + def test_setup_sh_kfd_awk_matches_install_sh(self): + """setup.sh's KFD fallback must use the same per-line vendor_id check as install.sh. + + setup.sh re-probes AMD detection independently of install.sh; if its copy keeps + the dead gpu_id-inside-properties pairing, a host that install.sh routes to ROCm + still gets a CPU llama.cpp from the setup step (_setup_amd_detected stays false). + """ + source = (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8") + assert ( + "$2 == 4098" in source + ), "setup.sh KFD awk must match `vendor_id 4098` as a single-line condition" + assert ( + "/gpu_id/" not in source + ), "setup.sh KFD awk must not key on a gpu_id line inside properties" + + def test_kfd_only_torch_falls_back_to_cpu(self): + """An AMD host whose gfx arch can't be read (rocminfo/amd-smi missing, or + present but not enumerating the GPU) must route torch to CPU, not a generic + rocm index: a Strix box (gfx1150/1151) would otherwise get the broken + _grouped_mm wheels because the reroute has no gfx to correct it.""" + source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + body = _extract_sh_function_body(source, "get_torch_index_url") + probe = body.find("_amd_gfx_probe=$(_probe_amd_gfx_arch)") + assert probe >= 0, "get_torch_index_url must probe the gfx arch before picking a rocm index" + # The shared probe reads gfx (not just tests binary presence), from rocminfo + # AND amd-smi, so an installed-but-not-enumerating probe still falls to CPU. + helper = _extract_sh_function_body(source, "_probe_amd_gfx_arch") + assert helper, "install.sh must define the shared _probe_amd_gfx_arch helper" + assert ( + "rocminfo 2>/dev/null) | grep -oE 'gfx" in helper + ), "probe must read gfx from rocminfo" + assert ( + "amd-smi list 2>/dev/null) | grep -oE 'gfx" in helper + ), "probe must read gfx from amd-smi" + # The probe clears ROCR/HIP_VISIBLE_DEVICES so a container mask + # (ROCR_VISIBLE_DEVICES=-1) can't blind the env-independent KFD detection. + assert ( + "unset ROCR_VISIBLE_DEVICES HIP_VISIBLE_DEVICES" in helper + ), "the gfx probe must clear the visibility masks so a mask can't force CPU" + cpu_guard = body.find('if [ -z "$_amd_gfx_probe" ]') + assert cpu_guard >= 0, "unreadable gfx must fall back to CPU" + assert cpu_guard < body.find( + "_rocm_tag=" + ), "the gfx gate must run before the ROCm version/index selection" + + def test_kfd_only_llama_requires_hipcc(self): + """setup.sh must forward --has-rocm for a gfx-unknown (KFD-only) host only when + hipcc is present. With no gfx the prebuilt resolver finds no ROCm bundle and the + source build would fail, so without a HIP toolchain the host keeps the CPU + prebuilt rather than breaking the llama.cpp install.""" + source = (PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8") + idx = source.find("_PREBUILT_CMD+=(--has-rocm)") + assert idx >= 0, "setup.sh must still be able to forward --has-rocm" + window = source[max(0, idx - 900) : idx] + assert ( + "hipcc" in window + ), "the gfx-unknown --has-rocm branch must gate on hipcc (a usable HIP toolchain)" + assert ( + "command -v hipcc" in window or "/opt/rocm/bin/hipcc" in window + ), "hipcc presence must be checked via command -v or the rocm bin path" + assert ( + "/opt/rocm-*/bin/hipcc" in window + ), "the hipcc gate must also accept a versioned /opt/rocm-*/bin/hipcc toolchain" + + def test_gfx_unknown_guard_honors_override(self): + """A user-set UNSLOTH_ROCM_GFX_ARCH must seed the gfx probe before the CPU + fallback: an air-gapped/rocminfo-less Strix host that names its arch should + still reach a rocm index instead of being forced to CPU.""" + source = (PACKAGE_ROOT / "install.sh").read_text(encoding = "utf-8") + helper = _extract_sh_function_body(source, "_probe_amd_gfx_arch") + assert helper, "install.sh must define the shared _probe_amd_gfx_arch helper" + seed = helper.find("$(printf") + assert seed >= 0, "the gfx probe must seed from UNSLOTH_ROCM_GFX_ARCH" + assert "UNSLOTH_ROCM_GFX_ARCH" in helper[seed : seed + 80] + assert seed < helper.find( + "rocminfo 2>/dev/null) | grep -oE 'gfx" + ), "the override must be read before probing rocminfo" + body = _extract_sh_function_body(source, "get_torch_index_url") + call = body.find("_amd_gfx_probe=$(_probe_amd_gfx_arch)") + assert call >= 0, "get_torch_index_url must call the shared probe" + assert call < body.find( + 'if [ -z "$_amd_gfx_probe" ]; then' + ), "the probe must run before the CPU fallback guard" + + def test_gfx_override_seeds_reroute_without_tools(self): + """The Strix reroute must honour UNSLOTH_ROCM_GFX_ARCH even when rocminfo and + amd-smi are absent, so a manual override reaches the arch index; with no + override and no tools it must stay empty (no false Strix routing).""" + shell = shutil.which("bash") + if not shell: + pytest.skip("bash needed to execute the probe block") + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + block = re.search( + r'^ _gfx_all=\$\(printf[^\n]*\n.*?(?=^ _strix_gfx="")', + source, + re.S | re.M, + ) + assert block, "could not extract the gfx-detection block" + with tempfile.TemporaryDirectory() as d: + # Shim rocminfo/amd-smi to enumerate nothing, so only the override can + # supply a gfx (keeps coreutils on PATH for tr/grep/printf). + for name in ("rocminfo", "amd-smi"): + p = os.path.join(d, name) + with open(p, "w", encoding = "utf-8") as f: + f.write("#!/bin/sh\nexit 0\n") + os.chmod(p, 0o755) + script = ( + 'set -euo pipefail\nHIP_VISIBLE_DEVICES=""\nROCR_VISIBLE_DEVICES=""\n' + + block.group(0) + + '\nprintf "OK:%s\\n" "$_gfx_all"\n' + ) + + def run(**extra): + env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""), **extra) + return subprocess.run( + [shell, "-c", script], env = env, capture_output = True, text = True + ) + + r = run(UNSLOTH_ROCM_GFX_ARCH = "GFX1151") + assert r.returncode == 0, f"override probe aborted: {r.stderr}" + assert "OK:gfx1151" in r.stdout, f"override not honoured/lowercased: {r.stdout!r}" + r2 = run() + assert r2.returncode == 0, f"empty probe aborted: {r2.stderr}" + assert ( + "OK:\n" in r2.stdout or r2.stdout.strip() == "OK:" + ), f"no override + no tools must leave gfx empty: {r2.stdout!r}" + + def test_gfx_probe_ignores_visibility_mask(self): + """A container visibility mask (ROCR_VISIBLE_DEVICES=-1) must not blind the + gfx probe: rocminfo honours the mask and would enumerate nothing, but KFD + detection is env-independent, so the probe clears the mask and still reads + the arch (else a masked host is wrongly forced to CPU).""" + shell = shutil.which("bash") + if not shell: + pytest.skip("bash needed to execute the probe block") + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + probe_fn = _extract_sh_function_body(source, "_probe_amd_gfx_arch") + assert probe_fn, "could not extract _probe_amd_gfx_arch" + with tempfile.TemporaryDirectory() as d: + # rocminfo that mimics ROCR_VISIBLE_DEVICES=-1 hiding all agents. + with open(os.path.join(d, "rocminfo"), "w", encoding = "utf-8") as f: + f.write( + "#!/bin/sh\n" + 'if [ "${ROCR_VISIBLE_DEVICES:-}" = "-1" ]; then echo "no agents"; exit 0; fi\n' + 'echo " Name: gfx1151"\n' + ) + os.chmod(os.path.join(d, "rocminfo"), 0o755) + script = ( + "set -euo pipefail\n" + "_ensure_rocm_probe_env() { :; }\n" + + probe_fn + + '\n_amd_gfx_probe=$(_probe_amd_gfx_arch)\nprintf "OK:%s\\n" "$_amd_gfx_probe"\n' + ) + + def run(**extra): + env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""), **extra) + return subprocess.run( + [shell, "-c", script], env = env, capture_output = True, text = True + ) + + r = run(ROCR_VISIBLE_DEVICES = "-1") + assert r.returncode == 0, f"masked probe aborted: {r.stderr}" + assert ( + "OK:gfx1151" in r.stdout + ), f"a visibility mask must not blind the gfx probe: {r.stdout!r}" + + def test_kfd_only_inferable_gfx_defers_to_reroute(self): + """A KFD-only host (GPU detected, gfx unreadable) whose arch IS inferable + from hardware IDs must not print the 'installing CPU-only PyTorch' warning: + get_torch_index_url returns the cpu index quietly and the runtime-less + reroute upgrades it to AMD per-arch wheels. Only when inference also fails + (or maps to no supported family) is CPU final, with the actionable hint.""" + shell = shutil.which("bash") + if not shell: + pytest.skip("bash needed to execute get_torch_index_url") + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + fn = _extract_sh_function_body(source, "get_torch_index_url") + probe_fn = _extract_sh_function_body(source, "_probe_amd_gfx_arch") + family_fn = _extract_sh_function_body(source, "_amd_arch_index_family_for_gfx") + assert fn and probe_fn and family_fn + with tempfile.TemporaryDirectory() as d: + # uname -> Linux/x86_64 so the AMD branch runs on any dev host; the + # rocminfo/amd-smi shims enumerate nothing (KFD-only host). + with open(os.path.join(d, "uname"), "w", encoding = "utf-8", newline = "\n") as f: + f.write('#!/bin/sh\ncase "${1:-}" in -m) echo x86_64 ;; *) echo Linux ;; esac\n') + for name in ("rocminfo", "amd-smi"): + with open(os.path.join(d, name), "w", encoding = "utf-8", newline = "\n") as f: + f.write("#!/bin/sh\nexit 0\n") + for name in ("uname", "rocminfo", "amd-smi"): + os.chmod(os.path.join(d, name), 0o755) + + def run(infer_stub): + script = ( + "set -euo pipefail\n" + "_ensure_rocm_probe_env() { :; }\n" + "_trim_index_path_slashes() { printf '%s\\n' \"$1\"; }\n" + "_has_usable_nvidia_gpu() { return 1; }\n" + "_has_amd_rocm_gpu() { return 0; }\n" + + infer_stub + + "\n" + + probe_fn + + "\n" + + family_fn + + "\n" + + fn + + "\n" + "get_torch_index_url\n" + ) + # Run from a file, not -c: Windows bash mangles multi-KB -c strings. + sp = os.path.join(d, "gtiu.sh") + with open(sp, "w", encoding = "utf-8", newline = "\n") as f: + f.write(script) + env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", "")) + for var in ( + "UNSLOTH_ROCM_GFX_ARCH", + "UNSLOTH_TORCH_INDEX_URL", + "UNSLOTH_TORCH_INDEX_FAMILY", + "UNSLOTH_PYTORCH_MIRROR", + "ROCR_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + ): + env.pop(var, None) + return subprocess.run( + [shell, sp.replace("\\", "/")], env = env, capture_output = True, text = True + ) + + r = run("_infer_linux_amd_gfx_arch() { echo gfx1100; }") + assert r.returncode == 0, f"inferable case aborted: {r.stderr}" + assert r.stdout.strip().endswith( + "/cpu" + ), f"must hand */cpu to the reroute: {r.stdout!r}" + assert ( + "inferring gfx1100" in r.stderr + ), f"must announce the inference handoff: {r.stderr!r}" + assert ( + "installing CPU-only PyTorch" not in r.stderr + ), f"must not promise a CPU-only install the reroute will override: {r.stderr!r}" + r2 = run("_infer_linux_amd_gfx_arch() { return 1; }") + assert r2.returncode == 0, f"uninferable case aborted: {r2.stderr}" + assert r2.stdout.strip().endswith("/cpu") + assert ( + "installing CPU-only PyTorch" in r2.stderr + ), f"uninferable gfx must keep the actionable CPU warning: {r2.stderr!r}" + r3 = run("_infer_linux_amd_gfx_arch() { echo gfx906; }") + assert r3.returncode == 0, f"unsupported-family case aborted: {r3.stderr}" + assert r3.stdout.strip().endswith("/cpu") + assert ( + "installing CPU-only PyTorch" in r3.stderr + ), f"an inferred arch with no wheel family must keep the CPU warning: {r3.stderr!r}" + + def test_no_version_cpu_warning_respects_gfx_override(self): + """With UNSLOTH_ROCM_GFX_ARCH set on a KFD-only host that has no ROCm + version sources, the gfx probe is seeded by the override, so the + no-version endpoint used to print 'falling back to CPU-only PyTorch' + even though the reroute then installs the per-arch wheels (Codex P3). + A supported override must defer; an unsupported override, or a + readable-gfx host without an override, keeps the CPU warning.""" + shell = shutil.which("bash") + if not shell: + pytest.skip("bash needed to execute get_torch_index_url") + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + fn = _extract_sh_function_body(source, "get_torch_index_url") + probe_fn = _extract_sh_function_body(source, "_probe_amd_gfx_arch") + family_fn = _extract_sh_function_body(source, "_amd_arch_index_family_for_gfx") + assert fn and probe_fn and family_fn + with tempfile.TemporaryDirectory() as d: + with open(os.path.join(d, "uname"), "w", encoding = "utf-8", newline = "\n") as f: + f.write('#!/bin/sh\ncase "${1:-}" in -m) echo x86_64 ;; *) echo Linux ;; esac\n') + # Silence every ROCm version source, not just amd-smi: a dev box with + # a real hipconfig/dpkg would otherwise resolve a version and skip + # the no-version endpoint this test exercises. + with open(os.path.join(d, "amd-smi"), "w", encoding = "utf-8", newline = "\n") as f: + f.write("#!/bin/sh\nexit 0\n") + for name in ("hipconfig", "dpkg-query", "rpm"): + with open(os.path.join(d, name), "w", encoding = "utf-8", newline = "\n") as f: + f.write("#!/bin/sh\nexit 1\n") + for name in ("uname", "amd-smi", "hipconfig", "dpkg-query", "rpm"): + os.chmod(os.path.join(d, name), 0o755) + script = ( + "set -euo pipefail\n" + "_ensure_rocm_probe_env() { :; }\n" + "_trim_index_path_slashes() { printf '%s\\n' \"$1\"; }\n" + "_has_usable_nvidia_gpu() { return 1; }\n" + "_has_amd_rocm_gpu() { return 0; }\n" + "_infer_linux_amd_gfx_arch() { return 1; }\n" + + probe_fn + + "\n" + + family_fn + + "\n" + + fn + + "\n" + "get_torch_index_url\n" + ) + sp = os.path.join(d, "gtiu.sh") + with open(sp, "w", encoding = "utf-8", newline = "\n") as f: + f.write(script) + + def run(rocminfo_body, **extra): + with open(os.path.join(d, "rocminfo"), "w", encoding = "utf-8", newline = "\n") as f: + f.write("#!/bin/sh\n" + rocminfo_body) + os.chmod(os.path.join(d, "rocminfo"), 0o755) + env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""), **extra) + for var in ( + "UNSLOTH_TORCH_INDEX_URL", + "UNSLOTH_TORCH_INDEX_FAMILY", + "UNSLOTH_PYTORCH_MIRROR", + "ROCR_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + ): + env.pop(var, None) + if "UNSLOTH_ROCM_GFX_ARCH" not in extra: + env.pop("UNSLOTH_ROCM_GFX_ARCH", None) + return subprocess.run( + [shell, sp.replace("\\", "/")], env = env, capture_output = True, text = True + ) + + # Supported override on a tool-blind host: defer to the reroute. + r = run("exit 0\n", UNSLOTH_ROCM_GFX_ARCH = "gfx1151") + assert r.returncode == 0, f"override case aborted: {r.stderr}" + assert r.stdout.strip().endswith("/cpu") + assert ( + "falling back to CPU-only PyTorch" not in r.stderr + ), f"a supported override must not get the false CPU warning: {r.stderr!r}" + assert ( + "UNSLOTH_ROCM_GFX_ARCH=gfx1151 is set" in r.stderr + ), f"the override deferral must be announced: {r.stderr!r}" + # Unsupported override: the reroute can't map it -> CPU warning stays. + r2 = run("exit 0\n", UNSLOTH_ROCM_GFX_ARCH = "gfx906") + assert r2.returncode == 0, f"unsupported-override case aborted: {r2.stderr}" + assert ( + "falling back to CPU-only PyTorch" in r2.stderr + ), f"an unmappable override must keep the CPU warning: {r2.stderr!r}" + # Readable gfx, no override, no version: deliberate CPU fallback. + r3 = run('echo " Name: gfx1151"\n') + assert r3.returncode == 0, f"readable-gfx case aborted: {r3.stderr}" + assert ( + "falling back to CPU-only PyTorch" in r3.stderr + ), f"a readable-gfx host without a version keeps the CPU warning: {r3.stderr!r}" + + def test_reroute_gate_covers_kfd_only(self): + """The runtime-less reroute must fire for a KFD-only host: _has_amd_rocm_gpu + is now true via the KFD topology, so the gate also accepts a detected GPU + whose gfx probe is empty (unslothai#7314 P2). A */cpu index chosen with a + READABLE gfx (deliberate ROCm-version fallback) must stay un-rerouted.""" + shell = shutil.which("bash") + if not shell: + pytest.skip("bash needed to execute the reroute block") + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + block = re.search( + r'^if \[ "\$_torch_index_pinned" = false \] && \[ "\$SKIP_TORCH" = false \] && \\\n' + r".*?^fi\n", + source, + re.S | re.M, + ) + assert block, "could not extract the runtime-less reroute block" + family_fn = _extract_sh_function_body(source, "_amd_arch_index_family_for_gfx") + assert family_fn + with tempfile.TemporaryDirectory() as d: + with open(os.path.join(d, "uname"), "w", encoding = "utf-8", newline = "\n") as f: + f.write('#!/bin/sh\ncase "${1:-}" in -m) echo x86_64 ;; *) echo Linux ;; esac\n') + os.chmod(os.path.join(d, "uname"), 0o755) + + def run(gpu_stub, probe_stub): + script = ( + "set -euo pipefail\n" + "_has_usable_nvidia_gpu() { return 1; }\n" + f"_has_amd_rocm_gpu() {{ {gpu_stub}; }}\n" + f"_probe_amd_gfx_arch() {{ {probe_stub}; }}\n" + "_infer_linux_amd_gfx_arch() { echo gfx1100; }\n" + "_strip_index_url_credentials() { printf '%s\\n' \"$1\"; }\n" + family_fn + "\n" + "_torch_index_pinned=false\nSKIP_TORCH=false\n_ARCH=x86_64\n" + "TORCH_INDEX_URL=https://download.pytorch.org/whl/cpu\n" + + block.group(0) + + 'printf "URL:%s GFX:%s\\n" "$TORCH_INDEX_URL" "${UNSLOTH_ROCM_GFX_ARCH:-}"\n' + ) + # Run from a file, not -c: Windows bash mangles multi-KB -c strings. + sp = os.path.join(d, "reroute.sh") + with open(sp, "w", encoding = "utf-8", newline = "\n") as f: + f.write(script) + env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", "")) + for var in ("UNSLOTH_ROCM_GFX_ARCH", "UNSLOTH_AMD_ROCM_MIRROR"): + env.pop(var, None) + return subprocess.run( + [shell, sp.replace("\\", "/")], env = env, capture_output = True, text = True + ) + + # KFD-only: GPU detected, probe empty -> reroute to per-arch wheels. + r = run("return 0", "printf '\\n'") + assert r.returncode == 0, f"kfd-only reroute aborted: {r.stderr}" + assert ( + "URL:https://repo.amd.com/rocm/whl/gfx110X-all/ GFX:gfx1100" in r.stdout + ), f"KFD-only host must reach the AMD arch index: {r.stdout!r}" + # The diagnostic must not claim /dev/kfd is missing: KFD visibility is + # exactly what detected this host (Codex P3). + assert ( + "ROCm runtime not visible" not in r.stderr + ), f"KFD-only reroute must not claim /dev/kfd is missing: {r.stderr!r}" + assert ( + "visible via the kernel driver (KFD)" in r.stderr + ), f"KFD-only reroute must name the tooling gap: {r.stderr!r}" + # Readable gfx: the */cpu index is a deliberate fallback -> untouched. + r2 = run("return 0", "echo gfx1151") + assert r2.returncode == 0, f"readable-gfx case aborted: {r2.stderr}" + assert ( + "URL:https://download.pytorch.org/whl/cpu GFX:" in r2.stdout + ), f"a deliberate CPU fallback must not be rerouted: {r2.stdout!r}" + # No AMD GPU detected at all: the pre-KFD-fix path still reroutes. + r3 = run("return 1", "printf '\\n'") + assert r3.returncode == 0, f"undetected-GPU case aborted: {r3.stderr}" + assert ( + "URL:https://repo.amd.com/rocm/whl/gfx110X-all/ GFX:gfx1100" in r3.stdout + ), f"the original undetected-GPU reroute must keep working: {r3.stdout!r}" + assert ( + "ROCm runtime not visible" in r3.stderr + ), f"a truly runtime-invisible host keeps the original diagnostic: {r3.stderr!r}" def test_get_torch_index_url_uses_nvidia_detected_flag(self): """get_torch_index_url must track NVIDIA via _nvidia_detected (proc-only NVIDIA still picks CUDA).""" @@ -3641,7 +4123,9 @@ class TestStrixRocm71Override: pytest.skip("bash needed to execute the probe block") source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") block = re.search( - r'^ _gfx_all=""\n.*?(?=^ _strix_gfx="")', source, re.S | re.M + r'^ _gfx_all=\$\(printf[^\n]*\n.*?(?=^ _strix_gfx="")', + source, + re.S | re.M, ) assert block, "could not extract the gfx-detection block" with tempfile.TemporaryDirectory() as d: @@ -3661,6 +4145,67 @@ class TestStrixRocm71Override: assert r.returncode == 0, f"probe aborted under set -e: {r.stderr}" assert "OK:gfx1151" in r.stdout, f"amd-smi fallback not reached: {r.stdout!r}" + def test_strix_reroute_reprobes_when_mask_hides_all(self): + """A visibility mask hiding every agent (ROCR_VISIBLE_DEVICES=-1) must not + skip the Strix reroute: get_torch_index_url reads the arch unmasked, so + the reroute must re-probe unmasked too or a masked Strix box gets the + broken generic wheels. A partial mask must keep its per-GPU selection. + Executed with mask-honouring shims, not a text match.""" + shell = shutil.which("bash") + if not shell: + pytest.skip("bash needed to execute the probe block") + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + block = re.search( + r'^ _gfx_all=\$\(printf[^\n]*\n.*?(?=^ _strix_gfx="")', + source, + re.S | re.M, + ) + assert block, "could not extract the gfx-detection block" + with tempfile.TemporaryDirectory() as d: + # rocminfo honours ROCR_VISIBLE_DEVICES like the real tool: -1 and + # set-but-empty hide both agents, 1 renumbers to the dGPU only, + # unset shows both. + rocminfo = ( + "#!/bin/sh\n" + 'case "${ROCR_VISIBLE_DEVICES-__unset__}" in\n' + ' __unset__) printf "Name: gfx1151\\nName: gfx1201\\n" ;;\n' + ' ""|-1) echo "no visible agents" ;;\n' + ' 1) printf "Name: gfx1201\\n" ;;\n' + ' *) printf "Name: gfx1151\\nName: gfx1201\\n" ;;\n' + "esac\n" + ) + for name, body in (("rocminfo", rocminfo), ("amd-smi", "#!/bin/sh\nexit 0\n")): + p = os.path.join(d, name) + with open(p, "w", encoding = "utf-8") as f: + f.write(body) + os.chmod(p, 0o755) + script = ( + "set -euo pipefail\n" + block.group(0) + '\nprintf "OK:%s\\n" "$_runtime_gfx"\n' + ) + + def run(**extra): + env = dict(os.environ, PATH = d + os.pathsep + os.environ.get("PATH", ""), **extra) + env.pop("UNSLOTH_ROCM_GFX_ARCH", None) + env.pop("HIP_VISIBLE_DEVICES", None) + return subprocess.run( + [shell, "-c", script], env = env, capture_output = True, text = True + ) + + # Mask hides everything: re-probe must recover the first GPU (Strix). + r = run(ROCR_VISIBLE_DEVICES = "-1") + assert r.returncode == 0, f"masked probe aborted: {r.stderr}" + assert "OK:gfx1151" in r.stdout, f"reroute blinded by full mask: {r.stdout!r}" + # A SET-but-empty mask also hides every agent and must re-probe too + # (the ${VAR+x} guard, not ${VAR:-}). + r0 = run(ROCR_VISIBLE_DEVICES = "") + assert r0.returncode == 0, f"empty-mask probe aborted: {r0.stderr}" + assert "OK:gfx1151" in r0.stdout, f"reroute blinded by empty mask: {r0.stdout!r}" + # Partial mask: enumeration already reflects it; the dGPU selection + # must survive (no unmasked re-probe overriding the user's pick). + r2 = run(ROCR_VISIBLE_DEVICES = "1") + assert r2.returncode == 0, f"partial-mask probe aborted: {r2.stderr}" + assert "OK:gfx1201" in r2.stdout, f"partial mask selection lost: {r2.stdout!r}" + def test_strix_routing_helpers_cover_rocm714(self): # Reroute for any generic pytorch.org index below the 7.13 arch floor (7.0, # 7.2, a future 7.3+), never at/above it -- mirrors install.sh _rocm_leaf_below.