From a65672d94752f20ca00762b65a78cbcae47e0600 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 18 Jun 2026 08:57:17 -0700 Subject: [PATCH] Installer: repair stale/CPU-only PyTorch and warn on silent CPU fallback (NVIDIA + AMD, Win/Linux/Mac/WSL) (#5942) * 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. --- .github/workflows/studio-backend-ci.yml | 3 +- .../studio-windows-inference-smoke.yml | 19 ++-- install.ps1 | 102 ++++++++++++++++++ install.sh | 80 +++++++++++++- tests/run_all.sh | 1 + tests/sh/test_torch_flavor.sh | 68 ++++++++++++ tests/studio/test_torch_flavor.ps1 | 52 +++++++++ 7 files changed, 313 insertions(+), 12 deletions(-) create mode 100755 tests/sh/test_torch_flavor.sh create mode 100644 tests/studio/test_torch_flavor.ps1 diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 88c7344683..b394b308e4 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -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::" diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index a772a6d102..a6f1401067 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -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: diff --git a/install.ps1 b/install.ps1 index cbb736b371..47851ed92e 100644 --- a/install.ps1 +++ b/install.ps1 @@ -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__). diff --git a/install.sh b/install.sh index f80d9b63b2..a0d8637d94 100755 --- a/install.sh +++ b/install.sh @@ -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. diff --git a/tests/run_all.sh b/tests/run_all.sh index d84c930392..4311aa4276 100755 --- a/tests/run_all.sh +++ b/tests/run_all.sh @@ -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 ===" diff --git a/tests/sh/test_torch_flavor.sh b/tests/sh/test_torch_flavor.sh new file mode 100755 index 0000000000..1f25ad807c --- /dev/null +++ b/tests/sh/test_torch_flavor.sh @@ -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 ] diff --git a/tests/studio/test_torch_flavor.ps1 b/tests/studio/test_torch_flavor.ps1 new file mode 100644 index 0000000000..86303af16c --- /dev/null +++ b/tests/studio/test_torch_flavor.ps1 @@ -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