diff --git a/install.sh b/install.sh index befcaaa186..44d1c39d02 100755 --- a/install.sh +++ b/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 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 0166bfd505..0460580922 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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) diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 973c0a1e81..6bc94e3d0f 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -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)