ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions
unsloth/tests/sh/test_get_torch_index_url.sh
Daniel Han 2db9fad4b5
 Installer: GPU detection follow-ups after #6174 (poisoned venv repair, llama.cpp routing, probe bounds) (#6183) 
* Installer: harden GPU detection follow-ups after #6174

Ports the NVIDIA-priority and /proc/driver/nvidia/gpus hardening from #6174
to the remaining pathways and adds recovery for already-poisoned venvs:

- install_python_stack.py: add _ensure_cuda_torch so 'unsloth studio update'
  force-reinstalls CUDA torch when the venv carries a ROCm build on an NVIDIA
  Linux host (the pre-#6174 poisoning signature). Honors UNSLOTH_TORCH_BACKEND,
  UNSLOTH_ROCM_TORCH_INSTALLED, and CUDA_VISIBLE_DEVICES=-1/'' opt-outs; never
  touches healthy CUDA, deliberate CPU wheels, macOS, or Windows.
- install_llama_prebuilt.py: detect_host gains the /proc NVIDIA fallback and
  skips ROCm probes when NVIDIA is usable; forwarded --rocm-gfx/--has-rocm
  overrides still win.
- setup.sh: GPU summary classifies NVIDIA first through a timeout-bounded
  probe with the /proc fallback; AMD probes are bounded and gain a KFD
  vendor_id 4098 fallback; the llama.cpp source build only selects
  GGML_CUDA/GGML_HIP when the matching GPU is actually detected.
- install.sh: bound both nvidia-smi calls with a 10s timeout (no behavior
  change when healthy or when the timeout binary is absent); classify the
  exported UNSLOTH_TORCH_BACKEND on the final index path segment so custom
  mirrors containing 'rocm'/'gfx' in their base path are not mislabeled.
- install.ps1 + setup.ps1: NVIDIA probes now require a real 'GPU N:' row from
  nvidia-smi -L under a 10s bound instead of bare exit code 0; later CUDA
  version and compute_cap queries are bounded too.

Tests: 3 new test files (50+ tests), suite at 788 passed.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix Resolve-CudaToolkit driver probe for extracted-function unit test

tests/studio/test_resolve_cuda_toolkit.ps1 extracts Resolve-CudaToolkit alone
into a child pwsh and stubs nvidia-smi with a .ps1 script. The bounded runner
is not in scope there (and ProcessStartInfo cannot dispatch .ps1 stubs), so
the DriverMaxCuda parse silently returned nothing and the major-mismatch
scenarios failed. Fall back to direct invocation when Invoke-NvidiaSmiBounded
is unavailable; production setup.ps1 always has it defined and keeps the
10s bound.

* Treat CUDA_VISIBLE_DEVICES empty or -1 as hidden in NVIDIA-first guards

The NVIDIA-first guards added in this branch only special-cased
CUDA_VISIBLE_DEVICES=-1 at two setup.sh gates and ignored the empty-string
form entirely, while the Python detector (install_llama_prebuilt.py)
already treats both as hidden. On a mixed AMD+NVIDIA host steered to the
AMD card via CUDA_VISIBLE_DEVICES, the guards suppressed the AMD probes,
so setup.sh fell to a CPU llama.cpp build and install.sh picked CUDA
wheels instead of ROCm.

Move the policy into the helpers so every consumer agrees:

- install.sh: new _cvd_hides_nvidia checked first in _has_usable_nvidia_gpu
- studio/setup.sh: same via _setup_cvd_hides_nvidia; the two ad-hoc
  CUDA_VISIBLE_DEVICES=-1 gate conditions are now redundant and removed
- studio/install_python_stack.py: _has_usable_nvidia_gpu returns False
  when CUDA_VISIBLE_DEVICES is set to  or -1 (whitespace tolerated)

Tests: 5 new sh scenarios (hidden via , -1, padded -1, visible device,
and mixed host with hidden NVIDIA restoring the ROCm route) plus a pytest
class covering all three implementations behaviourally.

Addresses the review comment on the NVIDIA-first setup.sh block.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Retrigger CI after PyPI 503 outage during the previous run

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-11 05:06:02 -07:00

386 lines
14 KiB
Bash
Executable file
Raw Permalink Blame History

#!/bin/bash
# Unit tests for get_torch_index_url() from install.sh
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
# Extract get_torch_index_url and its helper functions from install.sh.
# Also replace the hardcoded /usr/bin/nvidia-smi fallback with a
# controllable path so we can test the "no GPU" scenario on GPU machines.
_FUNC_FILE=$(mktemp)
_FAKE_SMI_DIR=$(mktemp -d)
{
sed -n '/^_run_bounded()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_cvd_hides_nvidia()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_has_amd_rocm_gpu()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^get_torch_index_url()/,/^}/p' "$INSTALL_SH"
} | sed "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \
> "$_FUNC_FILE"
# Save system PATH so we always have basic tools (uname, grep, head, etc.)
_SYS_PATH="/usr/local/bin:/usr/bin:/bin"
assert_eq() {
_label="$1"; _expected="$2"; _actual="$3"
if [ "$_actual" = "$_expected" ]; then
echo " PASS: $_label"
PASS=$((PASS + 1))
else
echo " FAIL: $_label (expected '$_expected', got '$_actual')"
FAIL=$((FAIL + 1))
fi
}
# Helper: create a mock nvidia-smi that prints a given CUDA version string.
# Handles both default output (version header) and -L (GPU listing) so that
# _has_usable_nvidia_gpu sees a valid GPU.
make_mock_smi() {
_dir=$(mktemp -d)
cat > "$_dir/nvidia-smi" <<MOCK
#!/bin/sh
case "\$1" in
-L)
echo "GPU 0: NVIDIA GeForce RTX 3090 (UUID: GPU-fake-uuid)"
;;
*)
cat <<'SMI_OUT'
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15 Driver Version: 550.54.15 CUDA Version: $1 |
+-----------------------------------------------------------------------------------------+
SMI_OUT
;;
esac
MOCK
chmod +x "$_dir/nvidia-smi"
echo "$_dir"
}
# Helper: create a mock nvidia-smi that prints the new "CUDA UMD Version" header
# layout used by newer NVIDIA drivers (e.g. 610.x on Windows). See issue #5812.
make_mock_smi_umd() {
_dir=$(mktemp -d)
cat > "$_dir/nvidia-smi" <<MOCK
#!/bin/sh
case "\$1" in
-L)
echo "GPU 0: NVIDIA GeForce RTX 5090 Laptop GPU (UUID: GPU-fake-uuid)"
;;
*)
cat <<'SMI_OUT'
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 610.47 KMD Version: 610.47 CUDA UMD Version: $1 |
+-----------------------------------------------------------------------------------------+
SMI_OUT
;;
esac
MOCK
chmod +x "$_dir/nvidia-smi"
echo "$_dir"
}
# Helper: create a mock amd-smi that prints a given ROCm version string
# Supports both "amd-smi version" and "amd-smi list" subcommands so that
# the GPU presence check (amd-smi list) also succeeds in tests.
make_mock_amd_smi() {
_dir=$(mktemp -d)
cat > "$_dir/amd-smi" <<MOCK
#!/bin/sh
case "\$1" in
list)
printf 'GPU: 0\\n BDF: 0000:03:00.0\\n NAME: gfx1100\\n'
;;
*)
cat <<AMD_OUT
AMDSMI Tool: 25.0.1+2b74356 | AMDSMI Library version: 25.0.1.0 | ROCm version: $1
AMD_OUT
;;
esac
MOCK
chmod +x "$_dir/amd-smi"
echo "$_dir"
}
# Build a minimal tools directory with symlinks to essential commands
# (uname, grep, head, etc.) but WITHOUT nvidia-smi or amd-smi.
_TOOLS_DIR=$(mktemp -d)
for _cmd in uname grep sed head sh bash cat awk printf tr; do
_real=$(command -v "$_cmd" 2>/dev/null || true)
[ -n "$_real" ] && ln -sf "$_real" "$_TOOLS_DIR/$_cmd"
done
# Helper: run get_torch_index_url with a custom PATH
# $1 = directory with mock nvidia-smi (prepended to PATH), or "none" for no-GPU test
run_func() {
_mock_dir="$1"
# Default: strip CUDA_VISIBLE_DEVICES so the host environment cannot leak
# in; a second argument sets it explicitly (hidden-GPU scenarios).
if [ "$#" -ge 2 ]; then
_cvd_setup="export CUDA_VISIBLE_DEVICES='$2'"
else
_cvd_setup="unset CUDA_VISIBLE_DEVICES"
fi
if [ "$_mock_dir" = "none" ]; then
# Minimal PATH with only basic tools, no nvidia-smi anywhere
PATH="$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
else
# Put mock nvidia-smi dir first, then basic tools
PATH="$_mock_dir:$_TOOLS_DIR" bash -c "$_cvd_setup; . '$_FUNC_FILE'; get_torch_index_url" 2>/dev/null
fi
}
echo "=== test_get_torch_index_url ==="
# 1) No nvidia-smi available -> cpu
_result=$(run_func "none")
assert_eq "no nvidia-smi -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
# 2) CUDA 12.6 -> cu126
_dir=$(make_mock_smi "12.6")
_result=$(run_func "$_dir")
assert_eq "CUDA 12.6 -> cu126" "https://download.pytorch.org/whl/cu126" "$_result"
rm -rf "$_dir"
# 3) CUDA 12.8 -> cu128
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir")
assert_eq "CUDA 12.8 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
rm -rf "$_dir"
# 4) CUDA 13.0 -> cu130
_dir=$(make_mock_smi "13.0")
_result=$(run_func "$_dir")
assert_eq "CUDA 13.0 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
# 5) CUDA 12.4 -> cu124
_dir=$(make_mock_smi "12.4")
_result=$(run_func "$_dir")
assert_eq "CUDA 12.4 -> cu124" "https://download.pytorch.org/whl/cu124" "$_result"
rm -rf "$_dir"
# 6) CUDA 11.8 -> cu118
_dir=$(make_mock_smi "11.8")
_result=$(run_func "$_dir")
assert_eq "CUDA 11.8 -> cu118" "https://download.pytorch.org/whl/cu118" "$_result"
rm -rf "$_dir"
# 7) CUDA 10.2 (too old) -> cpu
_dir=$(make_mock_smi "10.2")
_result=$(run_func "$_dir")
assert_eq "CUDA 10.2 -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# 8) Unparseable nvidia-smi version but valid GPU listing -> cu126 default
_dir=$(mktemp -d)
cat > "$_dir/nvidia-smi" <<'MOCK'
#!/bin/sh
case "$1" in
-L) echo "GPU 0: NVIDIA GeForce RTX 3090 (UUID: GPU-fake-uuid)" ;;
*) echo "something completely unexpected" ;;
esac
MOCK
chmod +x "$_dir/nvidia-smi"
_result=$(run_func "$_dir")
assert_eq "unparseable -> cu126" "https://download.pytorch.org/whl/cu126" "$_result"
rm -rf "$_dir"
# 9) ROCm 6.3 (no nvidia-smi) -> rocm6.3
_dir=$(make_mock_amd_smi "6.3")
_result=$(run_func "$_dir")
assert_eq "ROCm 6.3 -> rocm6.3" "https://download.pytorch.org/whl/rocm6.3" "$_result"
rm -rf "$_dir"
# 10) ROCm 7.1 (no nvidia-smi) -> rocm7.1
_dir=$(make_mock_amd_smi "7.1")
_result=$(run_func "$_dir")
assert_eq "ROCm 7.1 -> rocm7.1" "https://download.pytorch.org/whl/rocm7.1" "$_result"
rm -rf "$_dir"
# 11) ROCm 7.2 (no nvidia-smi) -> rocm7.2
_dir=$(make_mock_amd_smi "7.2")
_result=$(run_func "$_dir")
assert_eq "ROCm 7.2 -> rocm7.2" "https://download.pytorch.org/whl/rocm7.2" "$_result"
rm -rf "$_dir"
# 12) Both nvidia-smi and amd-smi present -> CUDA takes precedence
_cuda_dir=$(make_mock_smi "12.6")
_amd_dir=$(make_mock_amd_smi "6.3")
_combined_dir=$(mktemp -d)
ln -sf "$_cuda_dir/nvidia-smi" "$_combined_dir/nvidia-smi"
ln -sf "$_amd_dir/amd-smi" "$_combined_dir/amd-smi"
_result=$(run_func "$_combined_dir")
assert_eq "CUDA+ROCm -> CUDA precedence" "https://download.pytorch.org/whl/cu126" "$_result"
rm -rf "$_cuda_dir" "$_amd_dir" "$_combined_dir"
# 13) No nvidia-smi, no amd-smi -> cpu (duplicate of test 1, confirms ROCm didn't break it)
_result=$(run_func "none")
assert_eq "no GPU -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
# 14) ROCm 6.1 (no nvidia-smi) -> rocm6.1
_dir=$(make_mock_amd_smi "6.1")
_result=$(run_func "$_dir")
assert_eq "ROCm 6.1 -> rocm6.1" "https://download.pytorch.org/whl/rocm6.1" "$_result"
rm -rf "$_dir"
# 15) ROCm 6.4 (no nvidia-smi) -> rocm6.4
_dir=$(make_mock_amd_smi "6.4")
_result=$(run_func "$_dir")
assert_eq "ROCm 6.4 -> rocm6.4" "https://download.pytorch.org/whl/rocm6.4" "$_result"
rm -rf "$_dir"
# 16) ROCm 7.0 (no nvidia-smi) -> rocm7.0
_dir=$(make_mock_amd_smi "7.0")
_result=$(run_func "$_dir")
assert_eq "ROCm 7.0 -> rocm7.0" "https://download.pytorch.org/whl/rocm7.0" "$_result"
rm -rf "$_dir"
# 17) ROCm 8.0 (future, no nvidia-smi) -> rocm7.2 (capped to latest known)
_dir=$(make_mock_amd_smi "8.0")
_result=$(run_func "$_dir")
assert_eq "ROCm 8.0 -> rocm7.2 (capped)" "https://download.pytorch.org/whl/rocm7.2" "$_result"
rm -rf "$_dir"
# 18) Malformed amd-smi output (empty version field) -> cpu
_dir=$(mktemp -d)
cat > "$_dir/amd-smi" <<'MOCK'
#!/bin/sh
echo "AMDSMI Tool: 25.0.1 | AMDSMI Library version: 25.0.1.0 | ROCm version: "
MOCK
chmod +x "$_dir/amd-smi"
_result=$(run_func "$_dir")
assert_eq "empty amd-smi version -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# 19) amd-smi with "N/A" version -> cpu
_dir=$(mktemp -d)
cat > "$_dir/amd-smi" <<'MOCK'
#!/bin/sh
echo "AMDSMI Tool: 25.0.1 | AMDSMI Library version: 25.0.1.0 | ROCm version: N/A"
MOCK
chmod +x "$_dir/amd-smi"
_result=$(run_func "$_dir")
assert_eq "N/A amd-smi version -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# 20) ROCm version with trailing text (e.g. "6.3.1-beta") -> rocm6.3
_dir=$(make_mock_amd_smi "6.3.1-beta")
_result=$(run_func "$_dir")
assert_eq "ROCm 6.3.1-beta -> rocm6.3" "https://download.pytorch.org/whl/rocm6.3" "$_result"
rm -rf "$_dir"
# 22) CUDA 12.6 still works after ROCm changes (regression check)
_dir=$(make_mock_smi "12.6")
_result=$(run_func "$_dir")
assert_eq "CUDA 12.6 regression -> cu126" "https://download.pytorch.org/whl/cu126" "$_result"
rm -rf "$_dir"
# 23) CUDA 13.0 still works after ROCm changes (regression check)
_dir=$(make_mock_smi "13.0")
_result=$(run_func "$_dir")
assert_eq "CUDA 13.0 regression -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
# 24) CUDA 12.8 still works after ROCm changes (regression check)
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir")
assert_eq "CUDA 12.8 regression -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
rm -rf "$_dir"
# 25) UNSLOTH_PYTORCH_MIRROR overrides base URL (CUDA case)
_dir=$(make_mock_smi "12.6")
_result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl" run_func "$_dir")
assert_eq "mirror env + CUDA 12.6 -> mirror/cu126" "https://mirror.example.com/whl/cu126" "$_result"
rm -rf "$_dir"
# 26) UNSLOTH_PYTORCH_MIRROR overrides base URL (CPU case)
_result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl" run_func "none")
assert_eq "mirror env + no GPU -> mirror/cpu" "https://mirror.example.com/whl/cpu" "$_result"
# 27) Empty UNSLOTH_PYTORCH_MIRROR falls back to official URL
_result=$(UNSLOTH_PYTORCH_MIRROR="" run_func "none")
assert_eq "empty mirror env -> official/cpu" "https://download.pytorch.org/whl/cpu" "$_result"
# 28) Trailing slash in UNSLOTH_PYTORCH_MIRROR is stripped
_result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl/" run_func "none")
assert_eq "trailing slash stripped -> mirror/cpu" "https://mirror.example.com/whl/cpu" "$_result"
# 29) "CUDA UMD Version: 13.3" header (newer NVIDIA driver layout, issue #5812)
# -> cu130, not the cu126 fallback.
_dir=$(make_mock_smi_umd "13.3")
_result=$(run_func "$_dir")
assert_eq "CUDA UMD Version 13.3 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
# 30) "CUDA UMD Version: 12.8" header (newer layout on a 12.x driver) -> cu128
_dir=$(make_mock_smi_umd "12.8")
_result=$(run_func "$_dir")
assert_eq "CUDA UMD Version 12.8 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
rm -rf "$_dir"
# 31) "CUDA UMD Version: 11.8" header (newer layout on an older driver) -> cu118
_dir=$(make_mock_smi_umd "11.8")
_result=$(run_func "$_dir")
assert_eq "CUDA UMD Version 11.8 -> cu118" "https://download.pytorch.org/whl/cu118" "$_result"
rm -rf "$_dir"
# 32) Driver-reported "CUDA Version: 13.3" (legacy header) -> cu130.
_dir=$(make_mock_smi "13.3")
_result=$(run_func "$_dir")
assert_eq "CUDA Version 13.3 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
# 33) "CUDA Version: 13.7" -> cu130 (until a cu137 wheel index exists).
_dir=$(make_mock_smi "13.7")
_result=$(run_func "$_dir")
assert_eq "CUDA Version 13.7 -> cu130" "https://download.pytorch.org/whl/cu130" "$_result"
rm -rf "$_dir"
# 34) CUDA_VISIBLE_DEVICES="" hides the NVIDIA GPU -> cpu (no AMD present)
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir" "")
assert_eq "CVD='' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# 35) CUDA_VISIBLE_DEVICES=-1 hides the NVIDIA GPU -> cpu (no AMD present)
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir" "-1")
assert_eq "CVD=-1 hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# 36) Mixed AMD+NVIDIA host with NVIDIA hidden -> ROCm route is restored
_cuda_dir=$(make_mock_smi "12.6")
_amd_dir=$(make_mock_amd_smi "6.4")
_combined_dir=$(mktemp -d)
ln -sf "$_cuda_dir/nvidia-smi" "$_combined_dir/nvidia-smi"
ln -sf "$_amd_dir/amd-smi" "$_combined_dir/amd-smi"
_result=$(run_func "$_combined_dir" "-1")
assert_eq "CUDA+ROCm with CVD=-1 -> rocm6.4" "https://download.pytorch.org/whl/rocm6.4" "$_result"
rm -rf "$_cuda_dir" "$_amd_dir" "$_combined_dir"
# 37) CUDA_VISIBLE_DEVICES=0 (a visible device) must NOT hide the GPU
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir" "0")
assert_eq "CVD=0 keeps NVIDIA -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
rm -rf "$_dir"
# 38) Whitespace-padded "-1" still hides the GPU
_dir=$(make_mock_smi "12.8")
_result=$(run_func "$_dir" " -1 ")
assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
rm -f "$_FUNC_FILE"
rm -rf "$_FAKE_SMI_DIR"
rm -rf "$_TOOLS_DIR"
echo ""
echo "Results: $PASS passed, $FAIL failed"
[ "$FAIL" -eq 0 ] || exit 1
Reference in a new issue View git blame Copy permalink