Fix installer selecting ROCm torch on NVIDIA Linux hosts (#6174)
* fix: prevent ROCm torch from installing on NVIDIA Linux hosts
NVIDIA's open kernel module (driver 560+) registers GPU topology nodes in
the KFD sysfs hierarchy with non-zero gpu_id values. The _has_amd_rocm_gpu
(install.sh) and _has_rocm_gpu (install_python_stack.py) sysfs fallbacks
previously treated any non-zero gpu_id as proof of an AMD GPU, so an
NVIDIA-only host with the open kernel driver was misrouted to the ROCm
install path, replacing the correctly-installed CUDA torch with ROCm wheels.
Fixes:
1. install.sh _has_amd_rocm_gpu sysfs fallback: require vendor_id 4098
(AMD 0x1002) in the KFD node properties file before declaring an AMD
GPU present. NVIDIA KFD nodes carry vendor_id 4318 (0x10DE) and are
now skipped.
2. install_python_stack.py _has_rocm_gpu sysfs fallback: same vendor_id
guard. Also preserves the existing fallback for older kernels that
don't ship a properties file (trusts gpu_id alone there).
3. install.sh now exports UNSLOTH_TORCH_BACKEND ("cuda"/"rocm"/"cpu")
immediately after get_torch_index_url() resolves the wheel family.
install_python_stack.py reads this as _TORCH_BACKEND and short-circuits
_ensure_rocm_torch() entirely on cuda/cpu hosts, providing a second
layer of defense that is independent of subprocess GPU detection.
Tests: 9 new cases in TestHasRocmGpuKfdVendorGuard,
TestEnsureRocmTorch, and TestInstallShStructure cover all three changes.
Full test_rocm_support.py suite: 289 passed, 2 skipped, 0 failed.
Closes #6172
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix: show actual torch backend in progress step labels
The 'ROCm torch check' and 'ROCm torch (final)' step labels were
hardcoded regardless of whether the installer was targeting CUDA, ROCm,
or CPU. On NVIDIA hosts they showed 'ROCm' even though no ROCm wheels
were being installed, which was misleading.
Add _torch_step_label(suffix) which reads UNSLOTH_TORCH_BACKEND (set by
install.sh) and formats the label as e.g. 'torch check (cuda)' or
'torch final (rocm)'. Falls back to live GPU detection for standalone
studio update runs that bypass install.sh.
* fix: make KFD sysfs vendor check conservative -- skip if no properties file
The previous implementation fell through to `return True` when the KFD
node's properties file was missing (OSError), intending to support older
kernels. But NVIDIA open driver KFD nodes can also lack a properties file
on some kernel versions, so the fallback still produced a false positive.
Change the `except OSError: pass` to `continue` so any node without a
readable properties file is skipped rather than trusted. KFD properties
files exist on every kernel version that actually exposes /sys/class/kfd,
so this does not regress real AMD GPU detection -- if the directory exists
at all, properties files will be present for genuine GPU nodes.
* fix: bulletproof NVIDIA vs AMD GPU detection
Four changes that together ensure ROCm torch can never be installed on an
NVIDIA host regardless of which detection path fires:
1. _has_rocm_gpu() (Python): NVIDIA guard at the top -- returns False
immediately when _has_usable_nvidia_gpu() is True, blocking rocminfo,
amd-smi, and KFD sysfs from producing a false positive even when ROCm
tools are co-installed alongside the NVIDIA driver.
2. _has_amd_rocm_gpu() (install.sh): same NVIDIA guard -- calls
_has_usable_nvidia_gpu first and returns 1 if it succeeds.
3. _has_usable_nvidia_gpu() (Python): adds /proc/driver/nvidia/gpus/
sysfs fallback. The NVIDIA driver populates this directory on Linux
regardless of nvidia-smi state, so a subprocess PATH gap, timeout, or
driver initialisation race can no longer silence NVIDIA detection.
4. _has_usable_nvidia_gpu() (install.sh): same /proc/driver/nvidia/gpus
fallback, tried after nvidia-smi -L rather than instead of it.
Together: NVIDIA wins at every decision point. If nvidia-smi works, it
confirms NVIDIA. If it fails, /proc/driver/nvidia confirms NVIDIA. If
somehow both fail, _has_rocm_gpu still checks NVIDIA first before any AMD
path runs.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix: two KFD/proc-only corner cases from Codex review
1. KFD awk state not reset per node file (Ryzen+NVIDIA false positive):
The awk glob processes all topology node properties files in one pass.
Without FNR==1 reset, a Ryzen+NVIDIA host where an AMD CPU-agent node
sets amd=1 (vendor_id 4098, gpu_id 0) can combine with a later NVIDIA
node setting gpu=1 (gpu_id > 0), triggering found=1 before vendor_id
4318 is seen. Added FNR==1{ gpu=0; amd=0 } to reset per file.
2. proc-only NVIDIA not reaching CUDA wheel selection:
_has_usable_nvidia_gpu returning true via /proc/driver/nvidia fallback
left _smi empty, so get_torch_index_url entered the AMD/CPU branch and
selected CPU wheels despite NVIDIA being confirmed. Introduced
_nvidia_detected flag (separate from _smi) so the AMD branch is skipped
whenever NVIDIA is confirmed by any path, while _cuda_ver reads from
_smi when available (with the existing cu126 fallback when _smi is absent).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Leo Borcherding
GitHub
parent
ec46a55568
commit
bf2cd745b1
3 changed files with 366 additions and 31 deletions
49
install.sh
49
install.sh
|
|
@ -1717,8 +1717,15 @@ _ensure_rocm_probe_env() {
|
|||
|
||||
# Returns 0 if an AMD GPU is present. Checks rocminfo, amd-smi, then sysfs
|
||||
# KFD topology (env-var-independent fallback for when HIP/ROCR_VISIBLE_DEVICES hides devices).
|
||||
# Always returns 1 (false) when an NVIDIA GPU is present: blocks every
|
||||
# detection path (rocminfo, amd-smi, KFD sysfs) from producing a false
|
||||
# positive on NVIDIA-only or NVIDIA-primary hosts, even when ROCm tools
|
||||
# are co-installed.
|
||||
_has_amd_rocm_gpu() {
|
||||
_ensure_rocm_probe_env
|
||||
if _has_usable_nvidia_gpu; then
|
||||
return 1
|
||||
fi
|
||||
if command -v rocminfo >/dev/null 2>&1 && \
|
||||
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then
|
||||
return 0
|
||||
|
|
@ -1726,27 +1733,42 @@ _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 '/gpu_id/{ if ($2+0 > 0) found=1 } END{ exit !found }' \
|
||||
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 }' \
|
||||
/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.
|
||||
return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
# ── NVIDIA usable-GPU helper ──
|
||||
# Returns 0 (true) only if nvidia-smi is present AND actually lists a GPU.
|
||||
# Prevents AMD-only hosts with a stale nvidia-smi on PATH from being routed
|
||||
# into the CUDA branch.
|
||||
# Returns 0 (true) if an NVIDIA GPU is present and usable.
|
||||
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
|
||||
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
|
||||
# -- handles PATH gaps, subprocess timeouts, and driver init races that
|
||||
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
|
||||
_has_usable_nvidia_gpu() {
|
||||
_nvsmi=""
|
||||
if command -v nvidia-smi >/dev/null 2>&1; then
|
||||
_nvsmi="nvidia-smi"
|
||||
elif [ -x "/usr/bin/nvidia-smi" ]; then
|
||||
_nvsmi="/usr/bin/nvidia-smi"
|
||||
else
|
||||
return 1
|
||||
fi
|
||||
"$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'
|
||||
if [ -n "$_nvsmi" ]; then
|
||||
if "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
|
||||
return 0
|
||||
fi
|
||||
fi
|
||||
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
|
||||
if [ -d /proc/driver/nvidia/gpus ] && \
|
||||
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
|
||||
return 0
|
||||
fi
|
||||
return 1
|
||||
}
|
||||
|
||||
# ── Detect GPU and choose PyTorch index URL ──
|
||||
|
|
@ -1763,14 +1785,16 @@ get_torch_index_url() {
|
|||
# packages) is not sufficient: otherwise an AMD-only host would
|
||||
# silently install CUDA wheels.
|
||||
_smi=""
|
||||
_nvidia_detected=0
|
||||
if _has_usable_nvidia_gpu; then
|
||||
_nvidia_detected=1
|
||||
if command -v nvidia-smi >/dev/null 2>&1; then
|
||||
_smi="nvidia-smi"
|
||||
elif [ -x "/usr/bin/nvidia-smi" ]; then
|
||||
_smi="/usr/bin/nvidia-smi"
|
||||
fi
|
||||
fi
|
||||
if [ -z "$_smi" ]; then
|
||||
if [ "$_nvidia_detected" -eq 0 ]; then
|
||||
# No NVIDIA GPU -- check for AMD ROCm GPU.
|
||||
# PyTorch only publishes ROCm wheels for linux-x86_64; skip the
|
||||
# ROCm branch entirely on aarch64 / arm64 / other architectures
|
||||
|
|
@ -2098,6 +2122,15 @@ _maybe_bootstrap_rocm_wsl || true
|
|||
|
||||
TORCH_INDEX_URL=$(get_torch_index_url)
|
||||
|
||||
# Export the resolved torch backend ("cuda", "rocm", or "cpu") so that
|
||||
# downstream scripts (setup.sh -> install_python_stack.py) know what was
|
||||
# chosen here and can skip ROCm-specific repair steps on CUDA/CPU hosts.
|
||||
case "$TORCH_INDEX_URL" in
|
||||
*/rocm*|*/gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;;
|
||||
*/cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;;
|
||||
*) export UNSLOTH_TORCH_BACKEND="cuda" ;;
|
||||
esac
|
||||
|
||||
# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it.
|
||||
# All other ROCm tags and CUDA stay within <2.11.0.
|
||||
case "$TORCH_INDEX_URL" in
|
||||
|
|
|
|||
|
|
@ -557,7 +557,15 @@ def _persist_bnb_rocm_version(version: str) -> bool:
|
|||
|
||||
|
||||
def _has_rocm_gpu() -> bool:
|
||||
"""Return True only if an actual AMD GPU is visible (not just ROCm tools installed)."""
|
||||
"""Return True only if an actual AMD GPU is visible (not just ROCm tools installed).
|
||||
|
||||
Always returns False when an NVIDIA GPU is present -- NVIDIA takes
|
||||
priority on mixed hosts and prevents every detection path below
|
||||
(rocminfo, amd-smi, KFD sysfs) from producing a false positive even
|
||||
if ROCm tools are installed alongside the NVIDIA driver.
|
||||
"""
|
||||
if _has_usable_nvidia_gpu():
|
||||
return False
|
||||
for cmd, check_fn in (
|
||||
# rocminfo: look for a real gfx GPU id (3-4 chars, nonzero first digit).
|
||||
# gfx000 is the CPU agent; ROCm 6.1+ also emits generic ISA lines like
|
||||
|
|
@ -598,6 +606,13 @@ def _has_rocm_gpu() -> bool:
|
|||
# runtime-only detection. On minimal package-managed installs (no
|
||||
# rocminfo / no amd-smi tools), the kernel exposes AMD GPUs via
|
||||
# /sys/class/kfd so `studio update` can still detect and repair.
|
||||
#
|
||||
# Guard: reject any KFD node whose properties file reports a non-AMD
|
||||
# vendor. With the NVIDIA open kernel module (driver 560+), NVIDIA GPUs
|
||||
# can register KFD topology nodes with a non-zero gpu_id; those nodes
|
||||
# have vendor_id 4318 (0x10DE) rather than the AMD value 4098 (0x1002).
|
||||
# Without this check the fallback returns True on NVIDIA-only systems,
|
||||
# causing _ensure_rocm_torch to install ROCm wheels on NVIDIA hardware.
|
||||
if sys.platform != "win32":
|
||||
try:
|
||||
kfd_nodes = "/sys/class/kfd/kfd/topology/nodes"
|
||||
|
|
@ -609,29 +624,61 @@ def _has_rocm_gpu() -> bool:
|
|||
gpu_id = fh.read().strip()
|
||||
except OSError:
|
||||
continue
|
||||
if gpu_id and gpu_id != "0": # gpu_id 0 = CPU node
|
||||
return True
|
||||
if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node
|
||||
continue
|
||||
# Require AMD vendor_id 4098 (0x1002) in the properties file.
|
||||
# KFD properties files exist on every kernel that exposes
|
||||
# /sys/class/kfd, so absence of the file means we cannot
|
||||
# confirm AMD ownership -- skip the node rather than risk a
|
||||
# false positive (e.g. NVIDIA open driver KFD nodes that
|
||||
# lack a properties file on some kernel versions).
|
||||
props_path = os.path.join(kfd_nodes, entry, "properties")
|
||||
try:
|
||||
with open(props_path) as fh:
|
||||
props = fh.read()
|
||||
except OSError:
|
||||
continue # can't confirm vendor -- skip
|
||||
if not re.search(r"\bvendor_id\s+4098\b", props):
|
||||
continue
|
||||
return True
|
||||
except OSError:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def _has_usable_nvidia_gpu() -> bool:
|
||||
"""Return True only when nvidia-smi exists AND reports at least one GPU."""
|
||||
"""Return True when an NVIDIA GPU is present and usable.
|
||||
|
||||
Primary probe: nvidia-smi -L (subprocess).
|
||||
Fallback: /proc/driver/nvidia/gpus/ sysfs (Linux only) -- handles the
|
||||
case where nvidia-smi is present but the subprocess fails (PATH gap,
|
||||
timeout, driver initialisation race). If either probe confirms an
|
||||
NVIDIA GPU the function returns True so _has_rocm_gpu() is blocked.
|
||||
"""
|
||||
exe = shutil.which("nvidia-smi")
|
||||
if not exe:
|
||||
return False
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[exe, "-L"],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.DEVNULL,
|
||||
text = True,
|
||||
timeout = 10,
|
||||
)
|
||||
except Exception:
|
||||
return False
|
||||
return result.returncode == 0 and "GPU " in result.stdout
|
||||
if exe:
|
||||
try:
|
||||
result = subprocess.run(
|
||||
[exe, "-L"],
|
||||
stdout = subprocess.PIPE,
|
||||
stderr = subprocess.DEVNULL,
|
||||
text = True,
|
||||
timeout = 10,
|
||||
)
|
||||
if result.returncode == 0 and "GPU " in result.stdout:
|
||||
return True
|
||||
except Exception:
|
||||
pass
|
||||
# Fallback: the NVIDIA driver exposes one subdirectory per GPU under
|
||||
# /proc/driver/nvidia/gpus/ on Linux regardless of nvidia-smi state.
|
||||
if sys.platform != "win32":
|
||||
try:
|
||||
gpu_dir = "/proc/driver/nvidia/gpus"
|
||||
if os.path.isdir(gpu_dir) and os.listdir(gpu_dir):
|
||||
return True
|
||||
except OSError:
|
||||
pass
|
||||
return False
|
||||
|
||||
|
||||
def _detect_amd_gfx_codes() -> list[str]:
|
||||
|
|
@ -749,6 +796,13 @@ def _ensure_rocm_torch() -> None:
|
|||
Uses pip_install() to respect uv, constraints, and --python targeting.
|
||||
"""
|
||||
global _rocm_windows_torch_installed
|
||||
# install.sh sets UNSLOTH_TORCH_BACKEND to the resolved wheel family
|
||||
# ("cuda", "rocm", "cpu"). Skip ROCm operations entirely when install.sh
|
||||
# already selected a non-ROCm backend -- this is the authoritative signal
|
||||
# and avoids re-running GPU detection in a subprocess that may see a
|
||||
# different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.).
|
||||
if _TORCH_BACKEND in ("cuda", "cpu"):
|
||||
return
|
||||
# setup.ps1 sets this after installing AMD wheels; skip the probe only when
|
||||
# torch is actually importable as ROCm. If the venv was wiped between runs,
|
||||
# the stale env-var would suppress a needed reinstall.
|
||||
|
|
@ -1088,6 +1142,29 @@ def _infer_no_torch() -> bool:
|
|||
|
||||
NO_TORCH = _infer_no_torch()
|
||||
|
||||
# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() so
|
||||
# that this script knows which torch variant was selected without re-running
|
||||
# GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown
|
||||
# (standalone `unsloth studio update` runs, where we re-detect normally).
|
||||
_TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower()
|
||||
|
||||
|
||||
def _torch_step_label(suffix: str) -> str:
|
||||
"""Return a progress label like 'torch check (cuda)' using the known backend.
|
||||
|
||||
Falls back to GPU detection when UNSLOTH_TORCH_BACKEND is not set (e.g.
|
||||
standalone `unsloth studio update` runs that bypass install.sh).
|
||||
"""
|
||||
backend = _TORCH_BACKEND
|
||||
if not backend:
|
||||
if _has_usable_nvidia_gpu():
|
||||
backend = "cuda"
|
||||
elif _has_rocm_gpu():
|
||||
backend = "rocm"
|
||||
else:
|
||||
backend = "cpu"
|
||||
return f"torch {suffix} ({backend})"
|
||||
|
||||
|
||||
# -- Verbosity control ----------------------------------------------------------
|
||||
# By default the installer shows a minimal in-place one-line progress bar.
|
||||
|
|
@ -1770,7 +1847,7 @@ def install_python_stack() -> int:
|
|||
# venv got CPU-only torch (common when pip resolves torch from PyPI).
|
||||
# Must follow base packages so torch is present for inspection.
|
||||
if not IS_MACOS and not NO_TORCH:
|
||||
_progress("ROCm torch check")
|
||||
_progress(_torch_step_label("check"))
|
||||
_ensure_rocm_torch()
|
||||
|
||||
# Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python
|
||||
|
|
@ -1955,7 +2032,7 @@ def install_python_stack() -> int:
|
|||
# Running the repair last ensures ROCm torch is in place at runtime,
|
||||
# whichever intermediate step clobbered it.
|
||||
if not IS_WINDOWS and not IS_MACOS and not NO_TORCH:
|
||||
_progress("ROCm torch (final)")
|
||||
_progress(_torch_step_label("final"))
|
||||
_ensure_rocm_torch()
|
||||
|
||||
# 14. Final check (silent; third-party conflicts are expected)
|
||||
|
|
|
|||
|
|
@ -719,6 +719,143 @@ class TestEnsureRocmTorch:
|
|||
_ensure_rocm_torch()
|
||||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True)
|
||||
def test_torch_backend_cuda_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip):
|
||||
"""UNSLOTH_TORCH_BACKEND=cuda must short-circuit before any GPU probe."""
|
||||
with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cuda"}):
|
||||
# Reload _TORCH_BACKEND from the patched environment.
|
||||
with patch.object(stack_mod, "_TORCH_BACKEND", "cuda"):
|
||||
_ensure_rocm_torch()
|
||||
mock_pip.assert_not_called()
|
||||
|
||||
@patch.object(stack_mod, "pip_install")
|
||||
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
|
||||
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True)
|
||||
def test_torch_backend_cpu_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_pip):
|
||||
"""UNSLOTH_TORCH_BACKEND=cpu must short-circuit before any GPU probe."""
|
||||
with patch.dict(os.environ, {"UNSLOTH_TORCH_BACKEND": "cpu"}):
|
||||
with patch.object(stack_mod, "_TORCH_BACKEND", "cpu"):
|
||||
_ensure_rocm_torch()
|
||||
mock_pip.assert_not_called()
|
||||
|
||||
|
||||
# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard
|
||||
|
||||
|
||||
class TestHasRocmGpuKfdVendorGuard:
|
||||
"""Verify that the KFD sysfs fallback rejects non-AMD (NVIDIA) KFD nodes.
|
||||
|
||||
These tests are source-level: they verify the regex and logic present in
|
||||
the _has_rocm_gpu implementation rather than running the sysfs traversal
|
||||
(which requires Linux path conventions).
|
||||
"""
|
||||
|
||||
def _src(self) -> str:
|
||||
"""Return the source of _has_rocm_gpu from install_python_stack.py."""
|
||||
import inspect
|
||||
return inspect.getsource(stack_mod._has_rocm_gpu)
|
||||
|
||||
def test_vendor_id_check_present(self):
|
||||
"""_has_rocm_gpu sysfs fallback must check vendor_id 4098 (AMD 0x1002)."""
|
||||
src = self._src()
|
||||
assert "vendor_id" in src, (
|
||||
"_has_rocm_gpu KFD sysfs fallback must read the properties file "
|
||||
"to check vendor_id and exclude NVIDIA KFD nodes"
|
||||
)
|
||||
assert "4098" in src, (
|
||||
"_has_rocm_gpu must require AMD vendor_id 4098 (0x1002) in the "
|
||||
"KFD node properties to avoid false positives on NVIDIA systems"
|
||||
)
|
||||
|
||||
def test_vendor_regex_pattern_anchored(self):
|
||||
"""The vendor_id regex must use a word boundary to avoid partial matches."""
|
||||
import re as _re
|
||||
|
||||
src = self._src()
|
||||
# The pattern should have a word boundary before and after the number
|
||||
# so "vendor_id 41098" doesn't match "vendor_id 4098".
|
||||
assert (
|
||||
_re.search(r"\\b.*vendor_id.*\\b", src) or "\\bvendor_id" in src
|
||||
), "_has_rocm_gpu vendor_id check should use word boundary anchors"
|
||||
|
||||
def test_sysfs_fallback_guarded_by_non_win32(self):
|
||||
"""KFD sysfs fallback must be Linux-only (guarded by sys.platform != 'win32')."""
|
||||
src = self._src()
|
||||
assert "win32" in src, "_has_rocm_gpu sysfs fallback must be guarded by sys.platform check"
|
||||
|
||||
def test_cpu_node_excluded(self):
|
||||
"""gpu_id == '0' must be excluded (CPU topology nodes)."""
|
||||
src = self._src()
|
||||
assert (
|
||||
'!= "0"' in src or "== '0'" in src or "!= '0'" in src or '"0"' in src
|
||||
), "_has_rocm_gpu must skip gpu_id 0 nodes (CPU nodes)"
|
||||
|
||||
def test_install_sh_has_vendor_check(self):
|
||||
"""_has_amd_rocm_gpu in install.sh sysfs fallback must also check vendor_id 4098."""
|
||||
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 "vendor_id" in func_body, "_has_amd_rocm_gpu sysfs fallback must check vendor_id"
|
||||
assert "4098" in func_body, "_has_amd_rocm_gpu must require AMD vendor_id 4098 (0x1002)"
|
||||
|
||||
def test_has_rocm_gpu_returns_false_when_nvidia_present(self):
|
||||
"""_has_rocm_gpu must return False immediately when _has_usable_nvidia_gpu is True.
|
||||
|
||||
This is the primary guard: even if rocminfo, amd-smi, or KFD sysfs
|
||||
produce a false positive, an NVIDIA GPU always wins.
|
||||
"""
|
||||
with patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True):
|
||||
with patch("shutil.which", return_value = "/usr/bin/rocminfo"):
|
||||
# Simulate rocminfo claiming an AMD GPU is present
|
||||
mock_result = MagicMock()
|
||||
mock_result.returncode = 0
|
||||
mock_result.stdout = "Name: gfx1100\n"
|
||||
with patch("subprocess.run", return_value = mock_result):
|
||||
assert not stack_mod._has_rocm_gpu(), (
|
||||
"_has_rocm_gpu must return False when NVIDIA GPU is detected, "
|
||||
"regardless of what rocminfo reports"
|
||||
)
|
||||
|
||||
def test_install_sh_has_rocm_gpu_nvidia_guard(self):
|
||||
"""_has_amd_rocm_gpu in install.sh must call _has_usable_nvidia_gpu and return 1 if true."""
|
||||
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 (
|
||||
"_has_usable_nvidia_gpu" in func_body
|
||||
), "_has_amd_rocm_gpu must call _has_usable_nvidia_gpu to block NVIDIA hosts"
|
||||
assert (
|
||||
"return 1" in func_body
|
||||
), "_has_amd_rocm_gpu must return 1 (false) when NVIDIA GPU is detected"
|
||||
|
||||
def test_has_usable_nvidia_gpu_proc_fallback_present(self):
|
||||
"""`_has_usable_nvidia_gpu` must have a /proc/driver/nvidia fallback."""
|
||||
import inspect
|
||||
|
||||
src = inspect.getsource(stack_mod._has_usable_nvidia_gpu)
|
||||
assert "/proc/driver/nvidia" in src, (
|
||||
"_has_usable_nvidia_gpu must fall back to /proc/driver/nvidia/gpus when "
|
||||
"nvidia-smi subprocess fails, to handle PATH gaps and driver init races"
|
||||
)
|
||||
|
||||
def test_install_sh_has_usable_nvidia_gpu_proc_fallback(self):
|
||||
"""_has_usable_nvidia_gpu in install.sh must also have a /proc/driver/nvidia fallback."""
|
||||
sh_path = PACKAGE_ROOT / "install.sh"
|
||||
source = sh_path.read_text(encoding = "utf-8")
|
||||
func_start = source.find("_has_usable_nvidia_gpu()")
|
||||
func_end = source.find("\n}", func_start)
|
||||
func_body = source[func_start:func_end]
|
||||
assert "/proc/driver/nvidia" in func_body, (
|
||||
"_has_usable_nvidia_gpu in install.sh must fall back to "
|
||||
"/proc/driver/nvidia/gpus when nvidia-smi fails"
|
||||
)
|
||||
|
||||
|
||||
# TEST: install_python_stack.py -- _ROCM_TORCH_INDEX mapping
|
||||
|
||||
|
|
@ -927,16 +1064,21 @@ class TestInstallShStructure:
|
|||
source = sh_path.read_text(encoding = "utf-8")
|
||||
body = _extract_sh_function_body(source, "get_torch_index_url")
|
||||
nvidia_call = body.find("_has_usable_nvidia_gpu")
|
||||
no_nvidia_branch = body.find('if [ -z "$_smi" ]')
|
||||
# Gate changed from [ -z "$_smi" ] to [ "$_nvidia_detected" -eq 0 ] to
|
||||
# handle proc-only NVIDIA hosts where nvidia-smi is absent but _has_usable_nvidia_gpu
|
||||
# returns true via /proc/driver/nvidia/gpus.
|
||||
no_nvidia_branch = body.find('if [ "$_nvidia_detected" -eq 0 ]')
|
||||
if no_nvidia_branch < 0:
|
||||
no_nvidia_branch = body.find('if [ -z "$_smi" ]')
|
||||
rocm_call = body.find("_has_amd_rocm_gpu")
|
||||
assert nvidia_call >= 0, "get_torch_index_url should call _has_usable_nvidia_gpu"
|
||||
assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia-smi"
|
||||
assert no_nvidia_branch >= 0, "get_torch_index_url should gate ROCm on no-nvidia branch"
|
||||
assert (
|
||||
rocm_call > no_nvidia_branch
|
||||
), "ROCm detection should sit inside the 'no nvidia-smi' branch"
|
||||
), "ROCm detection should sit inside the 'no NVIDIA' branch"
|
||||
assert (
|
||||
nvidia_call < no_nvidia_branch
|
||||
), "NVIDIA detection should run before the no-nvidia-smi branch"
|
||||
), "NVIDIA detection should run before the no-NVIDIA branch"
|
||||
|
||||
def test_bitsandbytes_amd_install(self):
|
||||
"""install.sh should install bitsandbytes for AMD when ROCm detected."""
|
||||
|
|
@ -1018,6 +1160,89 @@ class TestInstallShStructure:
|
|||
rocm_pos = func_body.find("amd-smi")
|
||||
assert darwin_pos < rocm_pos, "macOS check should come before ROCm detection"
|
||||
|
||||
def test_unsloth_torch_backend_exported_after_get_torch_index_url(self):
|
||||
"""install.sh must export UNSLOTH_TORCH_BACKEND after TORCH_INDEX_URL is set.
|
||||
|
||||
This lets install_python_stack.py skip ROCm torch operations on CUDA
|
||||
and CPU hosts without re-running GPU detection in a subprocess.
|
||||
"""
|
||||
sh_path = PACKAGE_ROOT / "install.sh"
|
||||
source = sh_path.read_text(encoding = "utf-8")
|
||||
torch_url_pos = source.find("TORCH_INDEX_URL=$(get_torch_index_url)")
|
||||
backend_pos = source.find("UNSLOTH_TORCH_BACKEND")
|
||||
assert backend_pos > 0, "UNSLOTH_TORCH_BACKEND must be set in install.sh"
|
||||
assert (
|
||||
backend_pos > torch_url_pos
|
||||
), "UNSLOTH_TORCH_BACKEND must be set AFTER TORCH_INDEX_URL is resolved"
|
||||
# Verify all three cases are covered
|
||||
assert '"cuda"' in source[backend_pos : backend_pos + 500]
|
||||
assert '"rocm"' in source[backend_pos : backend_pos + 500]
|
||||
assert '"cpu"' in source[backend_pos : backend_pos + 500]
|
||||
# Must be exported so subprocesses (setup.sh, install_python_stack.py) see it
|
||||
assert "export UNSLOTH_TORCH_BACKEND" in source
|
||||
|
||||
def test_kfd_sysfs_amd_vendor_check_in_has_amd_rocm_gpu(self):
|
||||
"""_has_amd_rocm_gpu sysfs fallback must require AMD vendor_id 4098.
|
||||
|
||||
NVIDIA open kernel module (560+) registers KFD nodes with vendor_id
|
||||
4318 (0x10DE). Without the vendor check, _has_amd_rocm_gpu returns 0
|
||||
(true) on NVIDIA-only hosts that have the nvidia-open driver, causing
|
||||
get_torch_index_url to select a ROCm wheel index.
|
||||
"""
|
||||
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 (
|
||||
"vendor_id" in func_body
|
||||
), "_has_amd_rocm_gpu sysfs fallback must check vendor_id to exclude NVIDIA KFD nodes"
|
||||
assert (
|
||||
"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).
|
||||
|
||||
Without the reset, a Ryzen+NVIDIA host where node 0 is an AMD CPU
|
||||
agent (vendor_id 4098, gpu_id 0) and node 1 is an NVIDIA GPU
|
||||
(gpu_id > 0, vendor_id 4318) can produce a false positive: node 0
|
||||
sets amd=1, node 1 sets gpu=1, and the combined state triggers found=1
|
||||
before vendor_id 4318 is seen on node 1.
|
||||
"""
|
||||
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"
|
||||
)
|
||||
|
||||
def test_get_torch_index_url_uses_nvidia_detected_flag(self):
|
||||
"""get_torch_index_url must track NVIDIA detection independently of _smi.
|
||||
|
||||
When _has_usable_nvidia_gpu returns true via /proc/driver/nvidia fallback
|
||||
but nvidia-smi is not on PATH, _smi stays empty. Without a separate
|
||||
_nvidia_detected flag, the function falls into the AMD/CPU branch even
|
||||
though NVIDIA was confirmed, silently installing CPU wheels instead of CUDA.
|
||||
"""
|
||||
sh_path = PACKAGE_ROOT / "install.sh"
|
||||
source = sh_path.read_text(encoding = "utf-8")
|
||||
func_start = source.find("get_torch_index_url()")
|
||||
func_end = source.find("\n}", func_start)
|
||||
func_body = source[func_start:func_end]
|
||||
assert "_nvidia_detected" in func_body, (
|
||||
"get_torch_index_url must use a _nvidia_detected flag (separate from "
|
||||
"_smi) so that proc-only NVIDIA detection still selects CUDA wheels"
|
||||
)
|
||||
# The AMD/ROCm branch must be gated on _nvidia_detected being 0, not on
|
||||
# _smi being empty.
|
||||
assert (
|
||||
'_nvidia_detected" -eq 0' in func_body or "_nvidia_detected" in func_body
|
||||
), "get_torch_index_url AMD branch must be skipped when _nvidia_detected=1"
|
||||
|
||||
|
||||
# TEST: Live regression on current host (NVIDIA B200 expected)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue