* 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.
a65672d947