ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Installer: repair stale/CPU-only PyTorch and warn on silent CPU fallback (NVIDIA + AMD, Win/Linux/Mac/WSL) (#5942)

Browse source 
* Windows installer: repair a stale CPU PyTorch instead of looping forever

A Windows machine with an NVIDIA CUDA 13 driver (e.g. RTX 6000 Pro on enterprise
drivers) could get permanently stuck at:

  Stale venv detected (torch cpu != required cu130).
  [ERROR] The existing Studio environment needs repair.
          Re-run install.ps1 so it can replace the environment safely with rollback.

Re-running install.ps1 did not help. install.ps1 installs torch with
"torch>=2.4,<2.11.0" --index-url .../cu130 but no --force-reinstall, so when a
torch==X+cpu is already present uv treats it as satisfying the range (PEP 440
ignores the +cpu/+cuXXX local label) and makes no change -- the CPU wheel is
never replaced. setup.ps1 then rejects the venv as cpu != cu130 and exits, but it
cannot create a venv or install torch, so the loop never resolves. The migrated-
venv branch also preserves existing torch and never reinstalls it.

After the install step, detect the installed torch flavor (cuXXX/cpu/rocm) and,
when it does not match the tag implied by the selected index, force-reinstall the
torch/torchvision/torchaudio triplet from the correct index via three
--reinstall-package flags. No-op on a healthy matching venv; skipped for
--no-torch, ROCm (already --force-reinstalls), and CPU-only machines.

Adds two pure helpers (ConvertTo-TorchFlavorTag, Get-ExpectedTorchFlavorTag), a
PowerShell unit test (tests/studio/test_torch_flavor.ps1), and a CI parse gate for
install.ps1 (previously unparsed).

* install.sh: repair a stale CPU PyTorch on Linux too (parity with install.ps1)

install.sh has the same latent bug as the Windows installer: the CUDA torch
install uses "torch>=2.4,<2.11.0" --index-url .../cuXXX with no
--force-reinstall, so an already-present torch==X+cpu satisfies the version
range (PEP 440 ignores the +cpu/+cuXXX local label) and uv leaves it in place.
The migrated-venv branch also preserves existing torch. Unlike Windows there is
no stale-venv check in setup.sh, so on Linux the symptom is silent CPU training
rather than a hard loop -- same root cause.

Mirror the install.ps1 fix: after the install block, detect the installed torch
flavor (_torch_flavor_tag) and, when it does not match the index tag
(_expected_torch_flavor_tag), force-reinstall the torch/torchvision/torchaudio
triplet from the selected index via --reinstall-package. No-op on a healthy
matching venv; skipped for --no-torch, ROCm (its own repair force-reinstalls),
and CPU-only / macOS hosts. Adds tests/sh/test_torch_flavor.sh (run in
studio-backend-ci and run_all.sh).

* Installer: catch CPU-fallback on AMD/WSL too (repair ROCm, warn when unfixable)

Extend the torch-flavor safety net beyond NVIDIA:

- install.sh now auto-repairs a stale CPU torch on standard pytorch.org ROCm
  indexes too (the rocm-index install path lacked --force-reinstall, unlike the
  Windows ROCm install). Reuses the rocm-adjusted $TORCH_CONSTRAINT + rocm index,
  so it pulls the correct ROCm wheels.
- Both installers gain a universal post-install warning: when a GPU build was
  expected (cuXXX / rocm, including the repo.amd.com gfx* arch indexes) but torch
  is still CPU-only, warn loudly instead of silently training on CPU. This catches
  the cases auto-repair cannot safely fix (AMD gfx arch indexes that need
  --find-links, a migrated AMD venv on Windows where the ROCm install was skipped).
- Mac / Intel / CPU-only hosts resolve to the cpu index -> expected == installed
  -> no-op, no false warning. WSL uses install.sh, so the NVIDIA repair + warning
  apply there.

Adds Get-InstalledTorchTag (ps1) and _torch_index_repairable (sh) helpers and
extends both unit tests. gfx*/AMD indexes now map to the 'rocm' expected flavor.

* Installer: tighten torch-flavor comments (no logic change)

Condense the rationale comments added for the stale/CPU PyTorch repair in
install.ps1, install.sh and the two helper unit tests; same intent, fewer
lines. Comment-only: AST parse of install.ps1/setup.ps1 clean, helper unit
tests (15 ps1, 24 sh under bash and dash) and the integration sims
(24 ps1, 28 sh) still pass, banner markers the sims slice on are unchanged.

* Installer: bound torch probe, auto-repair gfx, fix ROCm gate parity

install.ps1: in Get-InstalledTorchTag, call WaitForExit(30000) and drain stdout
and stderr asynchronously instead of reading stdout synchronously first, so a
hung or noisy "import torch" (a wedged CUDA/driver, the exact failure this PR
targets) can no longer block the probe past the timeout.

install.sh and install.ps1: treat the repo.amd.com gfx* indexes as plain
--index-url reinstallable. They are PEP 503 simple indexes uv resolves in full
(torch plus every transitive dep) via --index-url, the same URLs the fresh
ROCm install paths already use, so a stale CPU torch on AMD Strix now auto-repairs
to the correct ROCm build instead of only warning.

install.sh: include */gfx* alongside */rocm* in the bitsandbytes install and
ROCm torch repair gates, so a custom UNSLOTH_AMD_ROCM_MIRROR whose path lacks
/rocm/ still installs the AMD bitsandbytes build and repairs ROCm torch.

tests/sh/test_torch_flavor.sh: gfx indexes now assert repairable, plus a
gfx1151 case and an unknown-mirror not-repairable case.

* install.ps1: guard Get-InstalledTorchTag against an empty python path

Make the early return explicit for an empty $PythonExe instead of relying on
Test-Path -LiteralPath '' returning false, so the probe stays safe under
Set-StrictMode or a future refactor that drops the [string] annotation.
This commit is contained in:
Daniel Han 2026-06-18 08:57:17 -07:00 • committed by GitHub
commit a65672d947
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 313 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

3
.github/workflows/studio-backend-ci.yml vendored
View file

@ -224,7 +224,8 @@ jobs:
tests/sh/test_mac_intel_compat.sh \
tests/sh/test_nvcc_meets_llama_minimum.sh \
tests/sh/test_tauri_install_exit_order.sh \
tests/sh/test_torch_constraint.sh; do
tests/sh/test_torch_constraint.sh \
tests/sh/test_torch_flavor.sh; do
echo "::group::$s"
bash "$s"
echo "::endgroup::"

19
.github/workflows/studio-windows-inference-smoke.yml vendored
View file

@ -65,17 +65,20 @@ jobs:
with:
persist-credentials: false
# Fast GPU-free gate: parse setup.ps1 and run the Resolve-CudaToolkit unit
# test (deferred Windows CUDA Toolkit check) before the heavy GGUF smoke.
- name: setup.ps1 unit test (Resolve-CudaToolkit)
# Fast GPU-free gate: parse install.ps1 + setup.ps1 and run the PowerShell
# unit tests (CUDA-toolkit + torch-flavor helpers) before the heavy GGUF smoke.
- name: PowerShell installer unit tests
shell: pwsh
run: |
$errs = $null
[void][System.Management.Automation.Language.Parser]::ParseFile(
(Resolve-Path studio/setup.ps1).Path, [ref]$null, [ref]$errs)
if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 }
Write-Host "setup.ps1 parsed with no errors"
foreach ($f in @('install.ps1', 'studio/setup.ps1')) {
$errs = $null
[void][System.Management.Automation.Language.Parser]::ParseFile(
(Resolve-Path $f).Path, [ref]$null, [ref]$errs)
if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 }
Write-Host "$f parsed with no errors"
}
pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1
pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:

102
install.ps1
View file

@ -1903,6 +1903,64 @@ exit 0
substep "could not determine CUDA version from nvidia-smi, defaulting to cu126" "Yellow"
return "$baseUrl/cu126"
}
# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
# torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu,
# matching setup.ps1's stale-venv parse.
function ConvertTo-TorchFlavorTag {
param([string]$TorchVersion)
if (-not $TorchVersion) { return $null }
if ($TorchVersion -match '\+(cu\d+)') { return $Matches[1] }
if ($TorchVersion -match '\+rocm') { return 'rocm' }
if ($TorchVersion -match '\+cpu') { return 'cpu' }
return 'cpu'
}
# Expected tag from the index leaf: cuXXX / cpu / rocm ($ROCmIndexUrl or a
# gfx* leaf -> rocm). $null on an unknown leaf (odd mirror) so repair no-ops.
function Get-ExpectedTorchFlavorTag {
param([string]$TorchIndexUrl, [string]$ROCmIndexUrl)
if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null }
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if ($leaf -match '^cu\d+$') { return $leaf }
if ($leaf -eq 'cpu') { return 'cpu' }
if ($leaf -match '^rocm') { return 'rocm' }
if ($leaf -match '^gfx') { return 'rocm' }
return $null
}
# Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses
# ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference.
function Get-InstalledTorchTag {
param([string]$PythonExe)
if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null }
try {
$psi = New-Object System.Diagnostics.ProcessStartInfo
$psi.FileName = $PythonExe
$psi.Arguments = '-c "import torch; print(torch.__version__)"'
$psi.RedirectStandardOutput = $true
$psi.RedirectStandardError = $true
$psi.UseShellExecute = $false
$psi.CreateNoWindow = $true
$proc = [System.Diagnostics.Process]::Start($psi)
# Drain BOTH streams async, then WaitForExit. A synchronous ReadToEnd()
# before the wait would block forever if a wedged "import torch" never
# closes stdout; leaving the redirected stderr undrained would deadlock a
# child that floods it past the pipe buffer. Async reads let a noisy-but-
# exiting probe finish, while a truly hung one still hits the 30s timeout
# and is killed -- bounded either way.
$outTask = $proc.StandardOutput.ReadToEndAsync()
$errTask = $proc.StandardError.ReadToEndAsync()
$finished = $proc.WaitForExit(30000)
if (-not $finished) { try { $proc.Kill() } catch {}; return $null }
$torchVer = $outTask.GetAwaiter().GetResult().Trim()
[void]$errTask.GetAwaiter().GetResult()
if ($proc.ExitCode -ne 0 -or -not $torchVer) { return $null }
return ConvertTo-TorchFlavorTag $torchVer
} catch { return $null }
}
$TorchIndexUrl = Get-TorchIndexUrl
# ── GPU arch → newest compatible Windows ROCm wheel release ──
@ -2149,6 +2207,50 @@ exit 0
}
}
# ── Enforce the installed torch flavor matches the detected GPU build ──
# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
# keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
# "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is
# expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx*
# is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install
# above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops.
if (-not $SkipTorch) {
$expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl
if ($expectedTorchTag -and $expectedTorchTag -ne 'cpu') {
$installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython
if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
if ($expectedTorchTag -eq 'rocm' -and $ROCmIndexUrl) {
# AMD: a migrated venv can keep a stale CPU torch the fresh ROCm path
# would have force-reinstalled. Repair from the same repo.amd.com index.
$rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec torchvision torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
}
$installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython
} elseif ($expectedTorchTag -ne 'rocm') {
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)
}
$installedTorchTag = Get-InstalledTorchTag -PythonExe $VenvPython
}
}
# Safety net (incl. AMD): GPU build expected but still CPU -> warn loudly.
if ($installedTorchTag -eq 'cpu') {
Write-Host ""
Write-Host " [WARN] PyTorch is CPU-only but a $expectedTorchTag GPU build was expected for this machine." -ForegroundColor Yellow
Write-Host " [WARN] Training and GPU inference will run on CPU until this is fixed." -ForegroundColor Yellow
Write-Host " [WARN] Re-run this installer, or reinstall the GPU build manually for your GPU." -ForegroundColor Yellow
}
}
}
# Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped
# for --local: the editable install above already makes _PACKAGE_ROOT in
# unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__).

80
install.sh
View file

@ -1990,6 +1990,45 @@ get_torch_index_url() {
else echo "$_base/cpu"; fi
}
# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
# torch.__version__ ($1) -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu.
_torch_flavor_tag() {
case "$1" in
*+cu[0-9]*) printf '%s\n' "$1" | sed -n 's/.*+\(cu[0-9][0-9]*\).*/\1/p' ;;
*+rocm*) echo "rocm" ;;
*+cpu*) echo "cpu" ;;
"") echo "" ;;
*) echo "cpu" ;;
esac
}
# Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* ->
# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops.
_expected_torch_flavor_tag() {
_u="${1%/}"
_leaf="${_u##*/}"
case "$_leaf" in
cu[0-9]*) echo "$_leaf" ;;
cpu) echo "cpu" ;;
rocm*|gfx*) echo "rocm" ;;
*) echo "" ;;
esac
}
# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX /
# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv
# resolves (torch + every transitive dep) via --index-url -- the same URLs the
# fresh-install paths above already use -- so a stale wheel is auto-repairable.
# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
_torch_index_repairable() {
_u="${1%/}"
_leaf="${_u##*/}"
case "$_leaf" in
cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;;
*) echo "no" ;;
esac
}
get_radeon_wheel_url() {
# Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing
# contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/,
@ -2513,7 +2552,7 @@ if [ "$_MIGRATED" = true ]; then
# fresh reinstall.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
*/rocm*|*/gfx*)
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
# Repair ROCm torch if overwritten during migrated install
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
@ -2689,7 +2728,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# which is only useful once torch is present for training.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
*/rocm*|*/gfx*)
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
;;
esac
@ -2735,7 +2774,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
*/rocm*|*/gfx*)
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
@ -2764,6 +2803,41 @@ else
fi
fi
# ── Enforce the installed torch flavor matches the detected GPU build ──
# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
# keeps a stale torch==X+cpu against a GPU index and the venv silently trains on
# CPU. Reinstall the right wheel triplet when a GPU build is expected; if it
# can't be reinstalled, warn loudly. --no-torch / CPU-only / macOS: no-op.
if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
_expected_torch_tag=$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")
# Only act when a GPU build is expected (cuXXX / rocm); cpu and unknown skip.
if [ -n "$_expected_torch_tag" ] && [ "$_expected_torch_tag" != "cpu" ]; then
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
[ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
# Repair when flavor is wrong AND the index is plain --index-url reinstallable
# (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only.
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
[ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
fi
# Safety net (incl. AMD/WSL): GPU build expected but still CPU -> warn loudly.
if [ "$_installed_torch_tag" = "cpu" ]; then
substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
fi
fi
fi
# ── Run studio setup ──
tauri_log "STEP" "Running Studio setup"
# When --local, use the repo's own setup.sh directly.

1
tests/run_all.sh
View file

@ -9,6 +9,7 @@ sh "$TESTS_DIR/sh/test_get_torch_index_url.sh"
sh "$TESTS_DIR/sh/test_mac_intel_compat.sh"
sh "$TESTS_DIR/sh/test_torch_constraint.sh"
sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh"
sh "$TESTS_DIR/sh/test_torch_flavor.sh"
echo ""
echo "=== Python tests ==="

68
tests/sh/test_torch_flavor.sh Executable file
View file

@ -0,0 +1,68 @@
#!/bin/bash
# Unit tests for install.sh's torch-flavor helpers (_torch_flavor_tag,
# _expected_torch_flavor_tag, _torch_index_repairable) that drive the
# stale-CPU-PyTorch repair. Helpers are extracted from install.sh and sourced.
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
# Extract the three helper functions from install.sh and source them.
_FUNC_FILE=$(mktemp)
{
sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_expected_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_torch_index_repairable()/,/^}/p' "$INSTALL_SH"
} > "$_FUNC_FILE"
# shellcheck disable=SC1090
. "$_FUNC_FILE"
rm -f "$_FUNC_FILE"
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
}
echo "=== _torch_flavor_tag ==="
assert_eq "cu130 wheel" "cu130" "$(_torch_flavor_tag '2.10.0+cu130')"
assert_eq "cu128 wheel" "cu128" "$(_torch_flavor_tag '2.8.0+cu128')"
assert_eq "cu124 wheel" "cu124" "$(_torch_flavor_tag '2.5.1+cu124')"
assert_eq "cu118 wheel" "cu118" "$(_torch_flavor_tag '2.4.0+cu118')"
assert_eq "cpu wheel" "cpu" "$(_torch_flavor_tag '2.10.0+cpu')"
assert_eq "untagged -> cpu" "cpu" "$(_torch_flavor_tag '2.10.0')"
assert_eq "rocm wheel" "rocm" "$(_torch_flavor_tag '2.11.0+rocm7.1')"
assert_eq "nightly cu130" "cu130" "$(_torch_flavor_tag '2.10.0.dev20250601+cu130')"
assert_eq "cu130 with suffix" "cu130" "$(_torch_flavor_tag '2.10.0+cu130.post1')"
assert_eq "empty -> empty" "" "$(_torch_flavor_tag '')"
assert_eq "garbage -> cpu" "cpu" "$(_torch_flavor_tag 'not-a-version')"
echo "=== _expected_torch_flavor_tag ==="
assert_eq "cu130 index" "cu130" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cu130')"
assert_eq "cu130 trailing /" "cu130" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cu130/')"
assert_eq "cu128 index" "cu128" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cu128')"
assert_eq "cpu index" "cpu" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/cpu')"
assert_eq "rocm index" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7.2')"
assert_eq "amd gfx index" "rocm" "$(_expected_torch_flavor_tag 'https://repo.amd.com/rocm/whl/gfx120X-all/')"
assert_eq "mirror cu130 leaf" "cu130" "$(_expected_torch_flavor_tag 'https://my.mirror/pytorch/whl/cu130')"
assert_eq "unrecognized leaf" "" "$(_expected_torch_flavor_tag 'https://my.mirror/whl/simple')"
assert_eq "empty url" "" "$(_expected_torch_flavor_tag '')"
echo "=== _torch_index_repairable ==="
assert_eq "cu130 repairable" "yes" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cu130')"
assert_eq "rocm7.2 repairable" "yes" "$(_torch_index_repairable 'https://download.pytorch.org/whl/rocm7.2')"
assert_eq "gfx repairable" "yes" "$(_torch_index_repairable 'https://repo.amd.com/rocm/whl/gfx120X-all/')"
assert_eq "gfx1151 repairable" "yes" "$(_torch_index_repairable 'https://repo.amd.com/rocm/whl/gfx1151/')"
assert_eq "cpu NOT repairable" "no" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cpu')"
assert_eq "unknown NOT repair" "no" "$(_torch_index_repairable 'https://my.mirror/whl/simple')"
echo ""
echo "Results: $PASS passed, $FAIL failed"
[ "$FAIL" -eq 0 ]

52
tests/studio/test_torch_flavor.ps1 Normal file
View file

@ -0,0 +1,52 @@
#!/usr/bin/env pwsh
# Unit test for install.ps1's torch-flavor helpers (ConvertTo-TorchFlavorTag,
# Get-ExpectedTorchFlavorTag) that drive the stale-CPU-PyTorch repair. Pure
# helpers, AST-extracted and run in-process -- no GPU/venv needed.
# Run: pwsh -NoProfile -File tests/studio/test_torch_flavor.ps1
$ErrorActionPreference = "Stop"
$installPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "install.ps1")
$installPath = (Resolve-Path $installPath).Path
# --- Parse install.ps1 (also serves as a syntax gate) and extract the helpers ---
$tokens = $null; $errors = $null
$ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors)
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" }
foreach ($name in @("ConvertTo-TorchFlavorTag", "Get-ExpectedTorchFlavorTag")) {
$fn = $ast.FindAll({ param($n)
$n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
}, $true)
if ($fn.Count -ne 1) { throw "expected exactly one $name in install.ps1, found $($fn.Count)" }
# Pure helpers (no exit / external calls) -- safe to define in this scope.
Invoke-Expression $fn[0].Extent.Text
}
$failures = 0
function Check($name, $cond) {
if ($cond) { Write-Host " PASS $name" }
else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ }
}
Write-Host "ConvertTo-TorchFlavorTag"
Check "2.10.0+cu130 -> cu130" ((ConvertTo-TorchFlavorTag "2.10.0+cu130") -eq "cu130")
Check "2.8.0+cu128 -> cu128" ((ConvertTo-TorchFlavorTag "2.8.0+cu128") -eq "cu128")
Check "2.10.0+cpu -> cpu" ((ConvertTo-TorchFlavorTag "2.10.0+cpu") -eq "cpu")
Check "2.10.0 (untagged) -> cpu" ((ConvertTo-TorchFlavorTag "2.10.0") -eq "cpu")
Check "2.11.0+rocm7.1 -> rocm" ((ConvertTo-TorchFlavorTag "2.11.0+rocm7.1") -eq "rocm")
Check "empty -> null" ($null -eq (ConvertTo-TorchFlavorTag ""))
Write-Host "Get-ExpectedTorchFlavorTag"
Check "cu130 index -> cu130" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cu130") -eq "cu130")
Check "trailing slash -> cu130" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cu130/") -eq "cu130")
Check "cpu index -> cpu" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cpu") -eq "cpu")
Check "ROCm url -> rocm" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/cpu" -ROCmIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all/") -eq "rocm")
Check "gfx index leaf -> rocm" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all/") -eq "rocm")
Check "rocm7.2 leaf -> rocm" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://download.pytorch.org/whl/rocm7.2") -eq "rocm")
Check "mirror cu130 leaf -> cu130" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://my.mirror/whl/cu130") -eq "cu130")
Check "unrecognized leaf -> null" ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://my.mirror/whl/simple"))
Check "empty url -> null" ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl ""))
Write-Host ""
if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
Write-Host "All checks passed" -ForegroundColor Green