install: let UNSLOTH_TORCH_INDEX_FAMILY / _URL override CUDA wheel detection (#6692)

* install: let UNSLOTH_TORCH_INDEX_FAMILY / _URL override CUDA wheel detection

get_torch_index_url (and the studio-update mirror _detect_cuda_torch_index_url)
chose the torch wheel family solely by probing the host GPU, with no override.
In a headless / container / CI build the host driver is visible via the
/proc/driver/nvidia/gpus fallback but nvidia-smi cannot report a CUDA version,
so the function fell back to its cu126 default and installed the wrong wheels
(e.g. a cu128 image got cu126 torch).

Add an explicit override checked before any probing, in both the shell installer
and the Python studio-update path:
  - UNSLOTH_TORCH_INDEX_URL   full index URL, used verbatim (wins)
  - UNSLOTH_TORCH_INDEX_FAMILY family (cpu, cu128, rocm6.4, ...) appended to the
                               mirror base (UNSLOTH_PYTORCH_MIRROR still honoured)

This matches how the published GPU images select CUDA -- vLLM and SGLang take the
CUDA version from an explicit build ARG rather than detecting it, and the Unsloth
Docker base image already pins the cu128 index directly. Desktop installs are
unchanged: with no override set, detection runs exactly as before.

Adds test_get_torch_index_url.sh cases for the override (family, full URL,
precedence, mirror base, trailing-slash strip, empty-ignored).

* install: make the torch-index override authoritative across ROCm paths

Address review feedback on the override added in this PR so a pinned index is
honoured everywhere, not just in get_torch_index_url:

- Skip the WSL ROCm bootstrap (root privilege + large downloads, probes
  /dev/dxg) when UNSLOTH_TORCH_INDEX_URL / _FAMILY is set; it previously ran
  before the override was consulted.
- Skip the Radeon/Strix rerouting (which re-probes the GPU and overwrites the
  resolved URL with repo.radeon.com / repo.amd.com) when the index is pinned, so
  an explicit ROCm override (e.g. UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4) is kept.
- install_python_stack.py: derive _TORCH_BACKEND from the override when
  UNSLOTH_TORCH_BACKEND is unset (standalone studio update), so _ensure_rocm_torch
  / _ensure_cuda_torch repair to the requested family instead of re-detecting.
- Strip ALL leading/trailing slashes in the shell override to match the Python
  side (avoids 404s on strict pip proxies).

Adds test cases for double-slash and leading/trailing-slash overrides.

* install: honor pinned torch index in CUDA/ROCm repair paths

Follow-up to the override work in this PR: the get_torch_index_url / install.sh
reroute already respect a pinned UNSLOTH_TORCH_INDEX_URL / _FAMILY, but the
Python repair helpers in install_python_stack.py still re-probed the GPU and
could overwrite the pinned family. Make the pin authoritative there too:

- _ensure_cuda_torch: an explicit cu* pin commits to CUDA wheels, so repair a
  ROCm-poisoned venv even when no NVIDIA GPU is visible here (headless /
  container / CI cross-install), instead of bailing on the GPU-presence gate.
- _ensure_rocm_torch: skip the AMD per-gfx (Strix) reroute when a ROCm index is
  pinned, and in the generic reinstall path install from the pinned URL verbatim
  rather than re-detecting the host ROCm version. gfx*/rocm7.2 indexes serve
  torch 2.11+, so select the 2.11 package specs for a gfx leaf.
- install.sh: raise the torch constraint to 2.11 for */gfx* indexes too, matching
  rocm7.2, so a pinned full-URL/family override that returns early keeps a valid
  constraint.

Add _explicit_torch_index_url / _explicit_rocm_torch_index_url helpers and tests
covering the no-GPU CUDA pin repair and the explicit gfx index honored verbatim.

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

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

* install: honor torch-index override on the Windows installers too

The pinned-index work landed for install.sh and install_python_stack.py, but the
Windows installers still picked the wheel index from GPU probing. Extend the same
UNSLOTH_TORCH_INDEX_URL / _FAMILY contract so a pinned index wins on every platform:

- install.ps1: Get-TorchIndexUrl returns the pinned URL/family before nvidia-smi
  probing; the AMD ROCm reroute is skipped when the index is pinned, so an explicit
  cpu/cu* pin on an AMD host is not overwritten.
- studio/setup.ps1: add shared Get-PinnedTorchIndexUrl / Get-TorchIndexLeaf helpers;
  the stale-venv check, the install selection and the AMD reroute all honor the pin,
  and the CPU/CUDA install pulls from the resolved index URL.
- tests: parity test that all four installers read both override vars and the two
  Windows installers gate the AMD reroute on the pinned flag.

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

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

* install: complete pinned-index handling for ROCm/Windows edge cases

Follow-ups to the override work flagged in review:

- install.ps1: a pinned gfx*/rocm>=7.2 index previously skipped the AMD reroute
  that sets the torch>=2.11 floor, so the generic install used torch>=2.4,<2.11
  and could resolve the known-bad _grouped_mm wheel. Route a pinned ROCm index
  through the ROCm install path with the 2.11 floor + companions, and guard the
  companion-spec lookup so a skipped reroute block cannot null-deref.
- studio/setup.ps1: the stale-venv check compared the installed flavor (cuXXX/cpu,
  with +rocm misread as cpu) against the raw pinned leaf (gfx1151 / rocm6.4), so a
  correct pinned ROCm venv was always marked stale. Classify +rocm wheels as the
  generic 'rocm' flavor and normalize a pinned rocm*/gfx* leaf to 'rocm' before
  comparing (cu* stays specific so cu126-vs-cu128 still rebuilds).
- install_python_stack.py: _ensure_cuda_torch now also reinstalls from a pinned
  CUDA index when the venv carries a CPU wheel (headless CPU-venv-to-CUDA
  cross-install via 'studio update'), not only when it finds a ROCm build.
- tests: parity assertions already cover all four installers honoring the override.

* install: finish pinned ROCm/CUDA edge cases on Windows + repair path

Follow-ups to the previous round:

- studio/setup.ps1: a pinned gfx*/rocm>=7.2 index now routes through the ROCm
  install path with the 2.11 floor + companions (it previously fell through to the
  CUDA branch with bare torch/torchvision/torchaudio against the ROCm index). The
  CPU/CUDA fallback index is forced to the CPU wheel index when a ROCm index is
  active, so a failed pinned-ROCm install does not retry the ROCm mirror.
- studio/setup.ps1: the stale-venv check no longer treats an unrecognized pinned
  URL leaf (e.g. a PEP 503 mirror ending in /simple) as a torch flavor tag, which
  was marking a correct venv stale; cu*/cpu/rocm/gfx leaves are still compared.
- install.ps1: the post-failure CPU fallback uses an explicit CPU index instead of
  , which for a pinned ROCm index was the ROCm mirror itself (so the
  'fallback' just retried the failing index and aborted the installer).
- install_python_stack.py: _ensure_cuda_torch now also reinstalls when the venv's
  CUDA family differs from a pinned one (installed cu126 vs pinned cu128), not only
  CPU->CUDA; the probe reports the installed cuXXX tag for the comparison.

* install: keep the ROCm to CPU fallback install inside the retry-helper window

The pinned-ROCm CPU fallback computes an explicit CPU index, but the comment
explaining why it cannot reuse $TorchIndexUrl pushed the actual
Invoke-InstallCommandRetry / --force-reinstall call more than 600 chars past the
"ROCm PyTorch install failed" message, so test_pr5940_followups's window check
no longer saw the retry helper. Move the CPU-index computation and its comment
above the failure substep so the retrying force-reinstall stays adjacent to the
message. No behavior change: same explicit CPU index, same retry, same
--force-reinstall.

* install: address #6692 review round 5 (ROCm/CPU pin edge cases)

setup.ps1:
- Stale-venv check: treat an AMD/ROCm host (HasROCm or a resolved gfx arch) with
  no explicit pin as expecting "rocm", not "cpu", so a healthy +rocm venv is not
  flagged stale (which made installer-managed setup exit and direct update rebuild).
- Pinned-ROCm install failure now routes into the force-reinstall CPU branch:
  CuTag stays the rocm/gfx leaf on failure, so the condition also checks
  ROCmCpuFallback; otherwise the CUDA branch installed from the CPU index without
  --force-reinstall and kept the partial ROCm torch.
- Explicit ROCm pin compare no longer collapses gfx*/rocm* to a generic "rocm":
  it compares the +rocmX.Y version (and the torch 2.11 line for gfx pins) so
  changing the pinned family (e.g. rocm6.4 -> gfx1151) rebuilds and applies it.

install_python_stack.py:
- _ensure_rocm_torch: an explicit ROCm wheel-index pin now bypasses the
  NVIDIA-present / no-AMD-GPU / unreadable-ROCm gates (headless/container/CI
  cross-install), mirroring the explicit-CUDA-pin bypass in _ensure_cuda_torch.
- Add _ensure_cpu_torch: an explicit CPU pin (FAMILY=cpu or /cpu URL) now has a
  repair path that reinstalls CPU torch over an existing CUDA/ROCm build on a
  standalone update (which skips install.sh's flavor enforcement).

install.sh:
- Pin torchvision/torchaudio companions alongside torch for the rocm7.2 / per-gfx
  index and the Strix reroute (those AMD indexes publish companions independently
  and a bare name can resolve a torch-2.12-built wheel, an ABI mismatch).

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

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

* torch-index override: classify CUDA pin by leaf; trim blank shell overrides

_ensure_cuda_torch only overrode the NVIDIA-presence gate for *any* pinned index,
so a non-CUDA mirror URL (or a ROCm/CPU pin) on a non-NVIDIA host with ROCm torch
could force a CUDA reinstall over a working ROCm venv. Add
_explicit_cuda_torch_index_url() (leaf cu*), matching the ROCm/CPU helpers, and
gate on it instead.

install.sh::get_torch_index_url treated a whitespace-only UNSLOTH_TORCH_INDEX_URL
/ _FAMILY as authoritative (yielding an invalid index), unlike the Python .strip()
and PowerShell IsNullOrWhiteSpace paths; trim leading/trailing whitespace first.

* install: honor pinned torch index over CVD/GPU gates and fix leaf-based ROCm classification

- install_python_stack.py: an explicit cu* pin now clears the CUDA_VISIBLE_DEVICES
  empty/-1 hide gate as well as the NVIDIA-presence gate, so
  CVD=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update repairs to CUDA wheels
  (parity with install.sh's get_torch_index_url override, which skips all GPU
  probing). Unpinned CVD=-1 still skips.
- install_python_stack.py: _ensure_cpu_torch installs the bounded _CPU_TORCH_PKG_SPEC
  instead of a bare torch/torchvision/torchaudio trio; the /cpu index now also
  serves torch 2.11+, which is outside the supported <2.11 range.
- install.sh: the torch>=2.11 constraint case matches the index leaf (rocm7.2|gfx*)
  instead of the whole URL, so a mirror base path containing a gfx/rocm7.2 segment
  with a cu*/cpu family is not false-matched onto the 2.11 line.
- setup.ps1: the stale-venv check expects rocm torch only for arches the install
  path maps to a repo.amd.com wheel index; an unmapped/unreadable arch installs
  CPU, so a correct CPU venv is no longer marked stale.
- Tests for each of the above.

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

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

* install: tighten pinned torch-index override edge cases

- install.sh: trim whitespace-only UNSLOTH_TORCH_INDEX_URL/_FAMILY before the
  _torch_index_pinned guard, matching get_torch_index_url, so a blank override no
  longer skips the WSL bootstrap and Radeon/Strix reroutes while detection still
  picks the normal index.
- install.sh / install.ps1 / setup.ps1 / install_python_stack.py: force the torch
  2.11 floor only for the gfx families with the <2.11 _grouped_mm bug (gfx120X-all,
  gfx1151, gfx1150). A pinned override to gfx110X-all/gfx90a/gfx908 stays on the
  default range, matching the automatic AMD path.
- install_python_stack.py _ensure_cuda_torch: treat an untagged CUDA build under a
  CUDA pin as a family mismatch (reinstall), and match cuXXX pins narrowly (cu +
  digits) so a custom/current mirror leaf no longer forces CUDA over a CPU/ROCm venv.
- install_python_stack.py _ensure_rocm_torch: reinstall when an explicit ROCm pin
  names a different ROCm family than the already-installed ROCm torch (the ROCm
  analogue of the CUDA cuXXX mismatch repair).

Adds tests for each case.

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

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

* install: fix second-order edge cases in pinned torch-index ROCm/CUDA handling

Parse the ROCm torch probe positionally so an empty HIP marker is kept:
CPU/CUDA torch no longer reads as HIP, so the ROCm reinstall is not skipped.
Emit one "<marker>|<version>" line (like the CUDA probe) for a robust parse.

Limit the gfx torch 2.11 expectation to the install allowlist
(gfx120X-all/gfx1151/gfx1150). A pinned gfx110X-all/gfx90a/gfx908 index stays
on the default <2.11 specs, so a correct 2.10+rocm wheel is no longer judged a
mismatch and force-reinstalled every update.

Distinguish an AMD per-arch wheel (three-part +rocmA.B.C) from a generic
pytorch.org wheel (two-part +rocmA.B): a gfx per-arch pin over a generic 2.11
wheel now reinstalls the per-arch wheel, while an already-installed per-arch
wheel is not re-flagged (no reinstall loop).

Mirror all of the above in setup.ps1 via new Test-RocmGfx211Leaf /
Test-CudaFamilyLeaf / Get-RocmPinStaleTags helpers, reused by both the
install-spec path and the stale-venv check so they cannot diverge again.
Require a digit after "cu" (^cu[0-9]) in setup.ps1, install.ps1 and install.sh
so a mirror leaf like /custom or /current is not branded CUDA and does not
rebuild the venv every run.

Add tests: CPU/CUDA probe -> has_hip_torch False; gfx110X-all pin + 2.10 wheel
not stale; gfx1151 pin + generic 2.11 wheel stale; gfx1151 pin + per-arch wheel
not stale; /custom and /current not CUDA; plus cross-language allowlist and
cu-digit parity guards, and a PowerShell unit test for the new setup.ps1 helpers.

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

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

* Fix ROCm/gfx pin case normalization, ROCm-tag requirement, and CUDA-leaf classification

Normalize torch-index leaves to lowercase before the gfx*/rocm*/cu* allowlist
matches so the canonical gfx120X-all (capital X) gets the torch 2.11 floor in
install.sh (leaf, flavor and repairable helpers). Require an installed +rocm
local tag before a rocmX.Y or non-2.11 gfx pin is judged satisfied in
setup.ps1 Get-RocmPinStaleTags and the Python _rocm_pin_family_mismatch, so an
untagged CPU/CUDA wheel never leaves the pin unapplied. Classify a leaf as CUDA
only via ^cu[0-9]: the Python _TORCH_BACKEND derivation now uses
_is_cuda_family_leaf, and install.sh brands cuda only on cu[0-9]* (unset on an
unknown /current /custom mirror leaf) so the stack probes the GPU instead of
skipping ROCm repair. Add bash, Python and PowerShell tests for capital
gfx120X-all floor, current/custom not-cuda, and untagged-wheel ROCm pins.

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

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

* install: converge torch-index pin detection via a per-venv marker

Introduce a torch-index MARKER that records the exact wheel --index-url used
after each successful torch install, so `unsloth studio update` / repair makes
the "did the pinned index change?" decision by an EXACT string compare rather
than inferring it from the wheel +rocm/+cu version tag. The tag cannot encode
the AMD per-arch gfx family (two 2.11 gfx indexes both install +rocm7.13.0), so
the tag heuristic missed a gfx1151 -> gfx120X-all switch and a custom-URL swap.

Marker path is per-venv (.unsloth-torch-index), one line = the resolved index
URL, written atomically (temp + rename). Path, format and normalization are
shared across all four installers (install.sh, install_python_stack.py,
setup.ps1, install.ps1).

- Reapply gfx pins on a per-arch target change: the marker's exact compare
  reinstalls when the pinned index differs, even when both wheels share a tag.
- Honor custom ROCm URL pins during repair: an explicit index whose leaf is not
  rocm/gfx/cu/cpu (e.g. simple, current) now reinstalls torch VERBATIM from the
  pin when it differs from the marker ("URL wins verbatim").
- Align the KNOWN-2.11 rocm/gfx set to exactly rocm7.2 plus the gfx allowlist
  gfx120x-all/gfx1151/gfx1150 in every language; stop treating an unknown newer
  rocm (rocm7.3, which does not exist) as the 2.11 line speculatively.

Backward compatible: with no marker (old venvs, torch installed out-of-band) the
existing +rocm/version-tag heuristics still decide, and a matching marker never
reinstall-loops. A cu128 CUDA pin stays a CUDA pin; custom and current leaves are
not CUDA. Adds marker tests (py/sh/ps) plus cross-installer parity checks.

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

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

* install: keep the torch-index marker additive to flavor validation

Three narrow fixes in the marker-based stale-venv detection:

- setup.ps1: a matching marker no longer overwrites the detected installed
  flavor. The marker compare is now an additional rebuild trigger, so a stale
  wheel (torch swapped to a +cpu build while the marker still records a cuXXX
  pin) is still caught by the flavor check instead of being masked as up to date.

- setup.ps1: a supported AMD arch carrying CPU torch is no longer marked stale
  and wiped. The downstream AMD Windows ROCm override upgrades CPU torch to ROCm
  in place, so wiping first would delete the venv and abort with "Virtual
  environment not found". Only a genuinely wrong CUDA wheel still rebuilds.

- install.sh: the Radeon --find-links path records its repo.radeon.com base in
  the marker instead of the generic pytorch.org ROCm fallback index, so a later
  pin to that generic family correctly reinstalls rather than comparing equal.
  Mirrors install.ps1/setup.ps1, which already record the real AMD index.

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

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

* install: honor custom pins and repair pinned venvs in place

Four follow-ups to the torch-index marker work:

- install_python_stack.py: _ensure_cuda_torch/_ensure_rocm_torch now bail when an
  explicit custom-index pin names no known torch family, so a verbatim URL override
  (a private/simple mirror) is not clobbered by auto-detected CUDA/ROCm wheels
  before _ensure_verbatim_torch_index applies it.

- install_python_stack.py: the ROCm marker is additive, not a substitute -- a
  matching marker still runs the family/version check so a wheel swapped after the
  marker was written is caught. Mirrors setup.ps1.

- setup.ps1: a stale venv under an explicit pin, whose torch still imports, is
  repaired in place (force-reinstall torch from the pin in the dependency pass)
  instead of wiped. The wipe path only delegates to install.ps1, so on a direct
  update it stranded the user at "Virtual environment not found" instead of
  applying the new pin. A broken venv or unpinned drift still wipes/delegates.

- install.ps1: when a pinned ROCm install fails over to a CPU base, the marker now
  records the CPU index actually used instead of the ROCm pin, so the next managed
  setup does not see CPU torch under a ROCm pin and abort as stale.

* setup.ps1: keep the ROCm CPU-fallback force line the pr5940 test guards

5c93ffd4 folded the pin-change force-reinstall into the ROCm CPU-fallback
condition on one line, so the exact literal that test_pr5940_followups.py checks
(if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") }) no longer appeared
and the test failed. Split the two conditions into separate if lines: the ROCm
fallback line is restored verbatim and the pin-change force is its own line. Both
still set $cpuForce to the array, so @splat passes one arg.

* install: honor exact CUDA/custom index URL pins in the torch-index marker

Address three Codex review findings on the torch-index marker mechanism:

- install.sh: after the ROCm CPU repair reinstalls torch from the generic
  $TORCH_INDEX_URL, record that as the marker source. A Radeon --find-links
  install set _TORCH_MARKER_INDEX_URL to its repo.radeon.com base earlier, so
  leaving it made the marker misreport Radeon wheels and a later Radeon pin would
  compare equal and skip a needed reinstall.

- install_python_stack.py: _ensure_cuda_torch now consults the exact-URL marker
  (_marker_pin_mismatch) when the installed +cuXXX tag matches the pinned leaf,
  so a same-leaf CUDA mirror change (official cu128 to an internal cu128 mirror)
  is reinstalled and re-recorded instead of skipped.

- _normalize_index_url / _normalize_family_leaf (install.sh, setup.ps1,
  install_python_stack.py): lowercase only KNOWN wheel-family leaves (rocm/gfx/
  cpu/cuXXX) so gfx120X-all still matches gfx120x-all, while a custom
  (unknown-family) leaf keeps its case so a verbatim URL pin like /Current does
  not compare equal to /current. Tests updated to assert the refined behavior.

* install: fix 3 torch-index marker edge cases (CPU mirror pin, Radeon leaf, migrated venv)

Addresses three review findings on the torch-index override path:

1. CPU index URL change on an already-CPU venv. _ensure_cpu_torch returned
   early whenever torch was already a CPU build, so a standalone update that
   moved the pin (official /cpu -> a private UNSLOTH_PYTORCH_MIRROR /cpu, same
   +cpu tag) never reinstalled. It now consults the exact-URL marker and
   reinstalls only when _marker_pin_mismatch reports a different index,
   mirroring the CUDA/ROCm same-family handling. A matching marker (or none)
   still leaves CPU torch untouched, so there is no reinstall loop.

2. Radeon find-links directory misclassified as a pip ROCm family. A
   repo.radeon.com/.../rocm-rel-7.2.1 leaf starts with "rocm" but is a
   find-links listing, not a pip --index-url. The old startswith(("rocm",
   "gfx")) test routed it into a --index-url reinstall that fails against
   find-links. New _is_pip_rocm_family_leaf gates on ^rocm\d / gfx (matching
   install.sh's rocm[0-9]* and setup.ps1's ^(rocm[0-9]|gfx)), so a Radeon URL
   routes to the verbatim/marker path instead.

3. Migrated venv rewriting its marker to a pin it did not install. install.sh
   and install.ps1 write the marker unconditionally, so a migration that
   preserves existing torch recorded the newly requested pin and a later
   update then found a matching marker and skipped the reinstall the pin
   needs (e.g. a per-arch gfx1151 -> gfx120X-all switch, identical +rocm tag).
   Both now track _TORCH_INSTALLED_THIS_RUN and write the marker only when
   torch was actually installed or repaired this run.

Also add Get-NormalizedFamilyLeaf to the setup.ps1 helper-extraction list in
test_torch_index_marker.ps1 (it was added to setup.ps1 and the shell test in an
earlier round but missed here) and add two unit tests covering findings 1 and 2.

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

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

* install: keep pinned torch repairs on the pinned index

Two fixes for explicit index pins (UNSLOTH_TORCH_INDEX_FAMILY / _URL):

1. install_python_stack.py's repair paths ran uv without clearing the
   inherited uv index env vars. uv resolves the default index (--index-url
   or --default-index) at the LOWEST priority, so a UV_INDEX or
   UV_EXTRA_INDEX_URL mirror in the environment won for any package it
   served: a cu128-pinned repair could install torch from the mirror and
   then record the cu128 marker it never used. Verified empirically: with
   UV_EXTRA_INDEX_URL=.../cu126 exported, uv pip install torch
   --index-url .../cu128 resolves torch 2.13.0+cu126. Strip the four uv
   index env vars for pinned-index commands only, mirroring the gate
   install.sh, install.ps1 and setup.ps1 already have; non-pinned installs
   keep the user's mirror.

2. install.ps1 routed any pinned leaf matching rocm* through the ROCm
   --default-index path, so a custom find-links leaf like rocm-rel-7.2.1
   was treated as a PEP 503 ROCm index and could silently fall back to CPU
   torch on resolution failure. Require a digit after rocm, matching
   install.sh's rocm[0-9]* and install_python_stack.py's ^rocm\d.

Adds parity + unit tests for both (11 new tests).

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

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

* install: keep pinned repairs off UV_TORCH_BACKEND and narrow setup.ps1's rocm pin match

Round 2 of the pinned-index hardening:

1. _build_uv_cmd converted UV_TORCH_BACKEND into --torch-backend before the
   new env isolation could act, and uv's torch backend redirects torch
   resolution to its own per-backend index even when --index-url is given
   (verified: a cu128-pinned dry run with UV_TORCH_BACKEND=cpu resolves
   torch 2.13.0+cpu). Pinned-index commands now never receive the flag and
   UV_TORCH_BACKEND joins the stripped env vars, so uv cannot re-read it.

2. setup.ps1's pinned reroute had the same bare rocm* glob install.ps1 had:
   a custom find-links leaf like rocm-rel-7.2.1 was routed through the ROCm
   --index-url path instead of the verbatim unknown-pin path. Now requires
   a digit after rocm, matching install.ps1, install.sh and
   _is_pip_rocm_family_leaf.

3. The marker test's case-normalization checks used -eq, which is
   case-insensitive in PowerShell, making them vacuous, and the unknown-leaf
   expectation was written lowercased while the implementation deliberately
   preserves custom-leaf case. Tightened to -ceq with the case-preserving
   expected value.

Adds unit + parity tests for 1 and 2 (5 new tests).

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

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

* install: extend the pinned-index guards to every remaining surface

Round 3 of the pinned-index hardening, closing the same holes on the
surfaces the earlier rounds missed:

1. install.sh's pinned-install env scrub now clears UV_TORCH_BACKEND (uv's
   torch backend redirects torch resolution to its own per-backend index
   even against --default-index), and both PowerShell wrappers clear it in
   their pinned-install scrubs, matching install_python_stack.py.

2. setup.ps1's marker stale check still classified any rocm* leaf as a
   PyTorch ROCm family while the install selection is digit-gated, so a
   custom rocm-current / rocm-rel-7.2.1 pin stale-compared as
   not-rocm vs rocm and force-reinstalled on every studio update. The
   stale check now uses the same ^rocm\d gate.

3. install_python_stack.py's pinned-command scrub also strips
   PIP_EXTRA_INDEX_URL for the pip fallback: pip adds the env extra index
   in addition to --index-url, so an inherited mirror could satisfy torch
   off the pin while the marker recorded the pinned URL. PIP_INDEX_URL
   needs no strip since the explicit --index-url flag overrides it.

Parity + unit tests extended (4 new tests).

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

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

* install: scrub find-links and carry the pinned scrub through pip fallbacks

Round 4 of the pinned-index hardening:

1. UV_FIND_LINKS joins every pinned-install scrub (install.sh, install.ps1,
   setup.ps1, install_python_stack.py): uv's --find-links locations can
   satisfy torch off the pinned index the same way an extra index does.

2. setup.ps1's Fast-Install restored the scrubbed vars in its finally
   BEFORE the pip fallback ran, and never touched the pip env vars at all,
   so a failed uv attempt fell back to python -m pip with an inherited
   PIP_EXTRA_INDEX_URL / PIP_FIND_LINKS able to win over the pinned
   --index-url. The scrub now wraps the whole function (uv attempt + pip
   fallback) and includes the pip vars; restore happens after both.

3. install_python_stack.py's scrub also strips PIP_FIND_LINKS for its own
   pip fallback, completing the PIP_EXTRA_INDEX_URL fix from round 3.

Parity tests extended (2 new tests).

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

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

* install: digit-gate rocm leaves in marker normalization and ROCm side effects

Round 5 of the pinned-index hardening (three custom-rocm-leaf edge cases):

1. _normalize_family_leaf lowercased every leaf starting with rocm, so a
   custom mirror leaf like rocm-Current compared equal to its lowercase form
   and a case-only pin change was skipped. URL paths can be case-sensitive.
   The rocm prefix is now digit-gated (rocm[0-9]*, matching
   _is_pip_rocm_family_leaf) in install.sh, setup.ps1 and
   install_python_stack.py, so only true family leaves (rocm7.2) are
   lowercased; a custom rocm-* leaf keeps its case.

2. setup.ps1 Test-MarkerPinMismatch compared normalized URLs with -ne, which
   is case-insensitive in PowerShell, so a case-only marker change (Simple
   vs simple) was treated as matching and the reinstall skipped. Now -cne.

3. install.sh gated the AMD bitsandbytes install and the "repair ROCm torch"
   --default-index reinstall on a bare whole-URL rocm glob, so a custom
   CPU/CUDA/private index whose leaf merely starts with rocm (rocm-current)
   was force-repaired from the wrong ROCm-only path whenever torch.version.hip
   was empty. Both now gate on _torch_index_is_rocm_family, computed once from
   the digit-gated leaf (rocm[0-9]*/gfx*).

Tests: 4 new parity assertions plus 2 case-sensitivity marker checks.

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

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

* install: apply an explicit custom torch-index pin on the first update

Round 6: an explicitly-set custom (unknown-family) UNSLOTH_TORCH_INDEX_URL
was silently ignored on the first `studio update` of a venv that predates
the marker feature, on both platforms, because the no-marker case was
treated as "do nothing" and the version-tag heuristics cannot judge an
unknown leaf.

1. install_python_stack.py _ensure_verbatim_torch_index now reinstalls
   verbatim when the marker is ABSENT (None), not only when it differs, and
   short-circuits only when the marker already records this exact pin. It
   then writes the marker, so every later update is a no-op. A user who did
   not set the override gets pin=None and is untouched, so an out-of-band
   torch install is never clobbered.

2. setup.ps1: for an unknown-family pin on a marker-less venv the stale-venv
   check now sets PinChangedForceReinstall so the torch block reinstalls in
   place from the pin. It deliberately does NOT set shouldRebuild, which
   would wipe the venv and strand a direct `studio update`.

3. setup.sh (the Linux `studio update` entry point) skipped
   install_python_stack.py entirely when unsloth was already current, so the
   marker-driven reinstall (both the verbatim custom pin and the cu/rocm
   flavor and family-change repair, e.g. gfx1151 to gfx120X-all) never ran.
   It now forces the dependency pass when a torch-index pin env var is set;
   the pass is idempotent and no-ops when the marker already matches. This
   mirrors setup.ps1's stale-venv pre-check.

Tests: 3 new parity assertions.

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

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

* test: expect first-update reinstall for a no-marker custom index pin

Follow-up to d671d8fb2: _ensure_verbatim_torch_index now applies an
explicit unknown-family URL pin verbatim on the first update when the
marker is absent (instead of no-op), so the old
test_verbatim_custom_url_no_marker_is_noop assertion was stale. Rewritten
as test_verbatim_custom_url_no_marker_reinstalls_once: asserts the one
verbatim reinstall from the pinned URL, that the marker is written, and
that a second call with the pin still set is idempotent (no reinstall
loop).

* install: gate the pinned update pass on the marker and record a pin baseline

Round 8, two follow-ups to the round-6 first-update pin fix:

1. setup.sh forced the full dependency pass on EVERY `studio update` while a
   torch-index pin stayed exported, even after the marker already recorded the
   same pin, turning quick updates into the expensive pass every time. It now
   probes install_python_stack.py --torch-pin-needs-apply (which reuses the
   exact marker normalization) and forces the pass only when the pin is not yet
   applied (marker absent or different); an already-applied persistent pin keeps
   the fast path. A probe error fails safe toward running the pass. setup.ps1
   gets the same probe in its fast path for parity.

2. A known-family full-URL pin on a venv predating the marker (e.g. an installed
   cu128 build and UNSLOTH_TORCH_INDEX_URL pointing at a same-family mirror) left
   the marker absent forever: the _ensure_* helpers deliberately do not force a
   multi-GB reinstall of identical-family wheels on an old venv, so nothing
   recorded the pin and every update re-entered the pass. _record_torch_index_pin_baseline
   now records the resolved pin as a baseline after the ensure sequence when the
   family already matches and no marker exists, so the pin is tracked (a later
   genuine change is detected and applied) and the update loop is broken, without
   the redundant reinstall.

Tests: 3 new baseline unit tests, 4 new parity assertions, and the CLI probe.

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

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

* setup.sh: keep the pin probe's exit 1 from killing the update under set -e

The --torch-pin-needs-apply probe deliberately exits 1 for the common
steady-state answer (pin already recorded, keep the fast path), but it ran
as a bare command under set -euo pipefail, so the whole studio update
aborted before the exit code was even captured. Absorb the status with
|| _PIN_NEEDS_APPLY=$? and pre-seed 0 so all three outcomes route as
documented: 0 runs the pass, 1 keeps the fast path, anything else fails
safe into the pass. Parity test asserts the guard.

* install: strip pin credentials, disable uv config discovery, bound verbatim installs

Four verified fix groups from a 12-reviewer audit of the torch-index
override feature, each reproduced before fixing:

1. Credential persistence: all four marker writers stored the raw pin URL,
   so an authenticated pin (https://user:token@mirror/simple) persisted its
   credentials in .unsloth-torch-index (mode 0644 under a default POSIX
   umask) and install_python_stack.py printed pin URLs verbatim in repair
   messages. Userinfo is now stripped before persisting and in every
   log/substep that interpolates a pin, via lockstep helpers
   (_strip_index_url_credentials in install.sh / install_python_stack.py,
   Remove-IndexUrlCredentials in install.ps1 / setup.ps1). The three
   normalizers strip too, so an OLD marker that already carries credentials
   still compares equal to the same pin: no reinstall loop on upgrade.
   Query strings deliberately stay in the marker; two indexes distinguished
   only by query must not compare equal.

2. uv configuration discovery beat the explicit pin: with a discovered
   uv.toml declaring torch-backend = "cpu" or a [[index]] entry, uv 0.10.12
   resolves torch 2.13.0+cpu against an explicit --index-url/.../cu126 pin;
   UV_NO_CONFIG=1 restores +cu126 (reproduced both ways). The pinned-install
   scrub in all four installers now sets UV_NO_CONFIG=1 and drops
   UV_CONFIG_FILE.

3. The verbatim custom-index update path installed a bare, unconstrained
   torch trio while fresh installs from the same unknown-leaf pin apply the
   supported range; _ensure_verbatim_torch_index now installs the bounded
   trio spec, closing the fresh-vs-update asymmetry.

4. Query-bearing pins (.../cu128?token=x) classified by raw leaf split and
   force-reinstalled on every update (the installed cu128 never equals
   cu128?token=x). Query/fragment are now stripped before leaf
   classification in all four implementations; the marker comparison keeps
   the query per (1).

Rejected after verification (no change): the pin-baseline record cannot
produce a wrong later decision (every pin change still mismatches and
reinstalls from the new pin); the venv temp-file symlink scenarios require
an attacker who already owns the environment; pathological inputs like
" / cu128 / " have no realistic caller and fail loudly.

Parity, stack, rocm-support, marker (sh + ps1), pin-stale, index-url and
flavor suites all pass (455 python + full shell/ps1 batteries).

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

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

* install: harden custom-pin repair against clobber, broken torch, and pip config

Four follow-ups to the pinned-index audit fixes:

1. setup.ps1 routed an unknown-leaf custom pin through the CUDA branch with
   a bare torch trio while install.ps1 (fresh) and the Python verbatim path
   bound the supported range; the pinned unknown-leaf route now applies the
   same torch>=2.4,<2.11.0 bound. Known cu* leaves and unpinned runs are
   unchanged.

2. The final torch safety pass could not repair a clobbered unknown-family
   pin: intermediate dependency steps can pull torch from PyPI (the pass
   exists for exactly that reason), but the verbatim helper short-circuited
   on marker==pin and no flavor tag exists to probe. The helper now keeps a
   per-run snapshot of the installed trio (taken after a verbatim reinstall
   or on the first matching-marker pass) and reinstalls from the pin when
   the final pass sees the trio drifted. Probe failure skips the
   comparison; a reinstall refreshes the snapshot, so no loop.

3. _record_torch_index_pin_baseline could freeze a known-family pin as
   applied on a venv whose torch is missing or broken (every family helper
   returns without reinstalling when its probe fails), making
   --torch-pin-needs-apply report done forever. The baseline now probes the
   installed flavor and records only on a match: a cuXXX pin requires the
   matching +cuXXX tag, cpu requires a cpu build, rocm/gfx requires hip;
   probe failure records nothing.

4. The pinned pip fallback stripped PIP_* env vars but user/site pip config
   files still applied (a configured global.extra-index-url can satisfy
   torch off the pin). PIP_CONFIG_FILE is now pointed at the null device
   for pinned commands (pip loads no config files then), in
   _install_env_for_cmd and setup.ps1's Fast-Install pinned scrub.
   install.sh / install.ps1 have no pip fallback (uv-only), verified.

Tests: 7 new rocm_support tests (snapshot reset fixture), 1 stack test,
2 parity tests. Full battery green (464 python, sh and ps1 suites).

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

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

* install: complete the pin-repair coverage across the fast path and platforms

Three cross-platform follow-ups to the round-2 pin-repair fixes:

1. The --torch-pin-needs-apply probe only compared marker==pin, so a torch
   trio clobbered to the wrong family (a cpu wheel replacing cu128 via a
   later pip install) with a still-matching marker reported "already
   applied" and the _ensure_{cuda,rocm,cpu} repair never ran on the Linux
   fast path. The probe is now a testable _torch_pin_needs_apply() that also
   checks the installed flavor against a known-family pin (via a shared
   _torch_flavor_matches_pin() helper, so the baseline and the probe cannot
   drift). An unknown-family pin has no flavor to validate and a failed
   probe cannot prove drift, so both keep the fast path.

2. macOS ARM (real CPU/MPS torch, not NO_TORCH) never applied an unknown-
   family custom pin on update: both the verbatim path and the baseline
   returned on IS_MACOS while fresh install.sh honors the pin, so the marker
   was never written and setup.sh forced the dependency pass on every update
   forever. The guards are now IS_MAC_INTEL (Intel mac is already NO_TORCH),
   and the final pass applies the pin on macOS ARM.

3. The round-2 final verbatim repair sat in the step-13 sequence guarded
   not IS_WINDOWS, so on Windows a dependency step that clobbered torch after
   the pin was applied was masked by the matching marker (setup.ps1 does not
   re-validate the main venv's torch after calling this script -- verified).
   Step 13 now runs the verbatim snapshot-drift repair on Windows and macOS
   ARM too; the Linux-oriented cuda/rocm/cpu family helpers stay Linux-only.

Tests: 13 new rocm_support cases (flavor drift, macOS ARM, Windows repair),
parity updates. Full battery green (475 python, sh and ps1 suites).

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

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

* install: strip query tokens from the marker and tighten the pin-drift probe

Four follow-ups to the round-3 pin-repair fixes:

1. The credential stripper feeding the torch-index marker and the logged repair
   messages dropped only user:pass@ userinfo, so a private feed that carries its
   auth token in the query string (.../simple?token=SECRET) persisted the token
   in the world-readable marker (mode 0644 under a default umask) and printed it
   in substep output. All four strippers (install.sh, install.ps1,
   studio/setup.ps1, install_python_stack.py) now drop the query and fragment
   before building the sanitized URL. A query is not part of a PEP 503 index's
   identity, so this also stops a rotated token from spuriously mismatching the
   marker and forcing a needless reinstall.

2. The --torch-pin-needs-apply fast-path probe accepted an untagged CUDA build
   (no +cuXXX local tag) under a specific cuXXX pin, but _ensure_cuda_torch
   reinstalls exactly that build to enforce the pin. The probe was more lenient
   than the repair, so the repair pass was skipped on the fast path.
   _torch_flavor_matches_pin now reports a mismatch for an untagged build under a
   cuXXX pin, forcing the pass.

3. The probe's ROCm branch accepted any HIP build for a rocm/gfx pin, while
   _ensure_rocm_torch decides a reinstall with the per-arch
   _rocm_pin_family_mismatch predicate (a generic +rocm7.2 wheel under a per-arch
   gfx pin, or a wrong ROCm version, is a mismatch). The probe now reuses that
   predicate, so it is as strict as the repair. This needs the installed torch
   version, so _probe_torch_flavor now returns (marker, cutag, version) and
   _torch_flavor_matches_pin takes the pin URL (extracting the leaf internally).

4. On Windows a known-family cu*/cpu pin is applied to the main venv by setup.ps1
   before install_python_stack.py runs; a later dependency step can clobber it,
   and the GPU-aware _ensure_{cuda,cpu}_torch self-skip on Windows while the
   verbatim helper handles only unknown-family pins, so nothing repaired the
   clobber (setup.ps1 does not re-validate the main venv's torch afterward,
   verified). New _ensure_pinned_known_family_torch reinstalls a drifted cu*/cpu
   pin in the step-13 Windows/macOS-ARM branch; rocm/gfx per-arch specs stay owned
   by setup.ps1, unknown-family by the verbatim helper.

A speculative ROCm 2.11 floor was also raised but is unreachable: the rocm7.2
index publishes no 2.x wheel below 2.11.0, and an unknown newer rocm is not
floored speculatively.

Tests: query/fragment strip cases in the sh + ps1 marker suites and the Python
strip/marker tests; the tri-state helper and the probe/baseline harnesses moved
to the (marker, cutag, version) flavor with matching versions; new probe cases
(untagged CUDA, generic-rocm-under-gfx) and 8 _ensure_pinned_known_family_torch
tests; a four-way query-strip parity assertion. Full battery green (1150 python,
sh 26/26 marker, ps1 marker/flavor/pin-stale).

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

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

* install: reinstall markerless gfx pins and cap custom-index updates at torch 2.11

Two follow-ups from the pin-marker audit:

1. A markerless venv with a gfx per-arch 2.11 pin trusted the wheel version
   tag, which is byte-identical (+rocm7.13.0) across gfx120X-all / gfx1151 /
   gfx1150. A pre-marker install holding one gfx arch's wheel that is now
   pinned to a DIFFERENT gfx index was therefore never switched:
   _rocm_pin_family_mismatch returns no-mismatch for any three-part +rocm
   2.11 wheel, and _ensure_rocm_torch's absent-marker branch fell through to
   that heuristic. _ensure_rocm_torch now forces a one-time reinstall when the
   marker is absent AND the pin leaf is a 2.11 gfx per-arch index; the reinstall
   writes the marker, so the next update compares exactly and does not loop
   (the correctly-pinned no-reinstall guarantee then comes from the exact marker
   compare, not the ambiguous tag). Non-gfx-2.11 pins (rocmX.Y, non-2.11 gfx)
   stay on the tag heuristic -- their tags are distinguishable.

2. The verbatim custom-index update path used _CUDA_TORCH_PKG_SPEC (torch
   <2.12.0) while a FRESH install of the same unknown leaf caps torch at
   <2.11.0 (install.sh's default TORCH_CONSTRAINT, and setup.ps1's custom-pin
   branch), so a private /simple mirror publishing torch 2.11 could upgrade a
   `studio update` to a state the fresh installer never produces. Added
   _CUSTOM_INDEX_TORCH_PKG_SPEC (torch>=2.4,<2.11.0), used only by the verbatim
   path; companions stay pinned for the same exclusive --index-url ABI reason
   as _CUDA_TORCH_PKG_SPEC (a bare name could pull a torch-2.12-built
   torchvision). _CUDA_TORCH_PKG_SPEC is unchanged (known-family cu/cpu repair
   correctly tracks install.sh's widened cu ceiling).

Tests: 2 new markerless-gfx cases (one-time reinstall + marker write + no-loop
second run, and the rocmX.Y absent-marker no-op), the pre-existing markerless
gfx no-reinstall test flipped to assert the one-time reinstall (it had encoded
the old tag-trusting behavior), and the custom-index bound assertions. 488
passed.

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

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

* install: a matching marker must not mask a broken, clobbered, or misclassified torch

Four round-6 follow-ups, all closing cases where a matching torch-index
marker wrongly vouched for a torch that is not actually the pinned one:

1. _is_cuda_family_leaf matched cu+digits by PREFIX (^cu[0-9]), so a custom
   mirror leaf like cu128-private classified as CUDA family; the flavor check
   then compared the installed cu128 tag to the whole leaf cu128-private and
   forced a reinstall on EVERY update (never converging). The cu family is
   now matched EXACTLY (re.fullmatch cu[0-9]+), so a cu-suffixed custom leaf
   routes through the verbatim/unknown path with a stable marker. Mirrored in
   install.sh (_normalize_family_leaf: strip cu, require an all-digit
   remainder) and setup.ps1 / install.ps1 (^cu[0-9]+$).

2. _torch_pin_needs_apply returned False on a failed torch probe (missing or
   unimportable) under a matching marker, so setup.sh kept the fast path and
   a broken torch was never repaired. A failed probe now forces the pass: the
   marker cannot vouch for a torch that does not import, forcing is idempotent,
   and once torch imports again the probe succeeds and the forcing stops
   (self-resolving). Reverses the round-4 conservative choice for this case.

3. _ensure_verbatim_torch_index snapshotted the installed trio on the first
   pass with a matching marker and treated an unimportable torch (snapshot
   None) as "no drift, skip", so a torch clobbered to a broken state before
   the run was masked. A None snapshot now reapplies the pin. A torch
   clobbered to a WORKING-but-wrong build under an unknown-family pin remains
   undetectable from metadata (no flavor tag; reinstalling every update would
   be the loop this avoids) and is documented as a known limitation.

4. The step-13 Windows final repair reran only the verbatim (unknown-family)
   and known-family cu*/cpu paths, so a clobbered explicit rocm/gfx pin (the
   wheel setup.ps1 installed from AMD's per-arch index) was left in place. The
   branch now also runs _ensure_rocm_torch on Windows for an explicit rocm/gfx
   pin; it has a Windows path and no-ops when torch already links HIP, so it
   only reinstalls a genuinely clobbered ROCm venv (loop-safe).

Tests: the round-4 failed-probe-trusts-marker test flipped to force the pass;
new cases for the cu-suffix no-loop, the broken-torch verbatim reinstall, and
the Windows rocm final-repair structure; item-2 exact-cu parity assertions.
490 passed. sh/ps1 marker + flavor + pin-stale suites all green.

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

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

* install: repair Windows ROCm pins from the pinned URL and honor NO_TORCH

Four round-7 review items, two of them regressions in the round-6 work:

1. _torch_pin_needs_apply ignored UNSLOTH_NO_TORCH. With a torch-index env
   var set and no marker, the failed-probe branch forced the dependency pass
   on every `studio update`, and the pass (which also honors NO_TORCH) never
   installs torch or writes a marker, so nothing could ever stop the forcing.
   It now returns False immediately under NO_TORCH: the pin only matters once
   torch is actually installed.

2. The step-13 Windows final repair (round-6) restored a clobbered explicit
   rocm/gfx pin by calling _ensure_rocm_torch, whose Windows path reinstalls
   from the arch AUTO-DETECTED via hipinfo, not from the pin. A user pinning a
   different gfx family or a private mirror was restored from the wrong source
   (and the wrong marker written), and a headless box was skipped entirely
   (the arch probe returns nothing). The repair now goes through
   _ensure_pinned_known_family_torch, which reinstalls from the PINNED url with
   the same per-arch floor setup.ps1 uses (2.11-line gfx leaves) or a bare trio
   (older arches, rocmN mirrors). It is gated on IS_WINDOWS since macOS ARM has
   no ROCm, and the existing flavor check keeps it loop-safe (a matching HIP
   wheel is left alone).

3. _ensure_verbatim_torch_index's broken-torch check (round-6) used
   "_installed_trio_snapshot() is None", but that helper reports a REMOVED torch
   as "torch==absent" (a non-None tuple) and a broken import as the stale
   on-disk version, so a missing or unimportable torch under a matching marker
   was read as "no drift" and skipped. The matching-marker path now confirms
   torch health with an import probe (_probe_torch_flavor): a torch that does
   not import reapplies the pin, while a healthy torch keeps the snapshot-based
   intra-run drift detection.

4. A unit test for _ensure_cpu_torch did not pin NO_TORCH False like its
   siblings, so a suite run with UNSLOTH_NO_TORCH=1 in the environment made the
   guard return early and the reinstall assertions fail spuriously.

Tests: the round-6 broken-torch verbatim test re-encodes the non-None
"torch==absent" snapshot case (the exact state the old "is None" check missed);
new Windows-ROCm pinned-repair cases (reinstall from the pin, per-arch floor vs
bare spec, matching-wheel no-op, off-Windows no-op); a NO_TORCH fast-path probe
case; the parity test now asserts the Windows final branch does not auto-detect
the ROCm index and that the helper reinstalls from the explicit pin. 494 passed.

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

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

* install: floor the rocm7.2 index in the Windows pin repair; isolate marker tests

Three round-8 review items, two of them downstream of the round-7 changes:

1. _ensure_pinned_known_family_torch gave a rocm<d> index leaf a bare
   torch/torchvision/torchaudio trio while flooring only gfx* leaves, so a
   Windows venv clobbered under an explicit rocm7.2 pin could reinstall an
   unbounded or ABI-mismatched trio from that exclusive --index-url. It now
   mirrors the spec the initial ROCm paths pin: the rocm7.2 floor for 2.11-line
   gfx leaves and rocm<d> leaves that serve torch 2.11, the <2.11 default for
   older rocm versions, and a bare trio only for older gfx per-arch leaves
   (which publish no floor), matching _ROCM_TORCH_PKG_SPECS / _ensure_rocm_torch.

2. test_verbatim_custom_url_no_marker_reinstalls_once called
   _ensure_verbatim_torch_index twice; the second call now hits the
   matching-marker health probe, and with pip_install mocked torch never becomes
   importable, so in a no-torch environment _probe_torch_flavor returned None and
   forced another reinstall, failing the idempotence assertion. The test now pins
   a healthy flavor so the idempotence check is about the marker, not ambient
   torch.

3. The TestEnsureRocmTorchMarker fixture patched os.environ per test but not
   _TORCH_BACKEND, which install_python_stack.py computes once at import from
   UNSLOTH_TORCH_BACKEND. A runner starting with a cuda/cpu backend made
   _ensure_rocm_torch early-return and skip the mocked repair these tests
   exercise. The fixture now neutralizes _TORCH_BACKEND so the marker tests are
   independent of the caller's installer-pin environment.

Tests: the Windows floor-spec test now asserts a rocm7.2 mirror pin uses the
rocm7.2 floor (not bare), plus a new rocm7.1 case that must fall back to the
<2.11 default; the marker suite passes under a hostile
UNSLOTH_TORCH_BACKEND=cuda / UNSLOTH_TORCH_INDEX_URL env. 495 passed.

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

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

* install: apply same-flavor pin repoints, keep ROCm fallback nonfatal, bound custom companions

Four round-9 review items, two of them regressions in the round-7 pin helper:

1. _ensure_pinned_known_family_torch returned as satisfied whenever the installed
   flavor matched the pin, so a same-flavor SOURCE change (one /cpu or /cu128
   mirror to another, or a gfx1151 -> gfx120x-all per-arch switch, both carrying
   the same wheel tag) was never applied, while _torch_pin_needs_apply kept forcing
   the pass on the marker mismatch forever. It now also reinstalls when the marker
   records a DIFFERENT index of the same flavor, rewriting the marker so the next
   update matches (no loop), exactly as the Linux _ensure_{cuda,cpu}_torch helpers
   do. An absent marker on an already-matching venv is still left to the baseline
   recorder (no forced reinstall of a correct pre-marker venv).

2. That helper reinstalled a Windows ROCm pin with the FATAL pip_install, so when
   setup.ps1 had taken its CPU fallback (the pinned AMD index unavailable), the
   final repair re-hit the same missing index and aborted the whole install. The
   ROCm reinstall is now nonfatal (pip_install_try): on failure it leaves the CPU
   base in place and writes no ROCm marker, so the install completes -- matching
   _ensure_rocm_torch's Windows path. cu*/cpu pins stay fatal (authoritative source).

3. install.sh left torchvision/torchaudio bare for a pinned custom/unknown-leaf
   index (a private /simple mirror), unlike the Python update path's
   _CUSTOM_INDEX_TORCH_PKG_SPEC, so a mirror also exposing newer companion wheels
   could resolve a torch-2.12-built torchvision against the capped <2.11 torch. It
   now bounds the companions (torchvision>=0.19,<0.26.0 / torchaudio>=2.4,<2.11.0)
   for a custom leaf, gated on an empty _expected_torch_flavor_tag so known families
   keep their curated bare/floored companions.

4. install.sh's _expected_torch_flavor_tag matched cu[0-9]* by prefix, so a custom
   leaf like cu128-private classified as the cu128 family and force-reinstalled a
   correct +cu128 wheel on every run. It now requires exact cu+digits (routing the
   suffixed leaf to the custom path), matching the Python re.fullmatch(cu[0-9]+) and
   PowerShell, and feeding item 3's custom-leaf detection.

Tests: new cases for the same-flavor marker-change reinstall, the nonfatal ROCm
fallback (no marker on failure), the rocm7.2/older-rocm floor selection now split
across the nonfatal path, cu-suffixed custom leaves in test_torch_flavor.sh, and the
custom-leaf companion bounds in test_torch_constraint.sh. 497 python + 143 shell
assertions pass; the marker suite still passes under a hostile
UNSLOTH_TORCH_BACKEND=cuda env.

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

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

* install: bound custom-pin companions on the Windows setup path; isolate pin-probe tests

Two round-10 review items:

1. setup.ps1's custom/unknown-leaf pin branch capped only torch ($cudaTorchSpec)
   and still asked the exclusive index for bare torchvision/torchaudio, so a
   private mirror that also serves newer companion wheels could install a
   torch<2.11 wheel alongside a torchvision>=0.26 / torchaudio>=2.11 built for a
   newer torch ABI, after which the marker records the pin as applied. It now
   bounds the whole trio (torch>=2.4,<2.11.0 / torchvision>=0.19,<0.26.0 /
   torchaudio>=2.4,<2.11.0) for a pinned non-cu-family leaf, matching install.sh,
   install.ps1's fresh pinned install, and install_python_stack.py's
   _CUSTOM_INDEX_TORCH_PKG_SPEC. This completes the companion-bounds fix across all
   three installers; known cu* leaves keep bare specs (the family index bounds them).

2. The _torch_pin_needs_apply probe tests did not pin NO_TORCH False, so a test
   process launched with UNSLOTH_NO_TORCH=1 short-circuited the probe (the round-7
   guard) and returned False for cases that expect the pass to run. The _needs_apply
   helper now patches NO_TORCH (default False) around the call, and the dedicated
   no-torch case passes no_torch=True explicitly.

Tests: the cross-platform parity test now asserts setup.ps1 bounds the full trio
(not just torch) for a custom leaf; the pin-probe suite passes under a hostile
UNSLOTH_NO_TORCH=1 environment. setup.ps1 parses clean; 497 python + shell suites
green.

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

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

* install: bound custom rocm-* pins, redact diag tokens, snapshot custom pins before base update

Three round-11 review items, all reproduced before fixing:

1. install.sh's custom-index companion bounds gated on _expected_torch_flavor_tag
   returning empty, but that helper returned "rocm" for ANY rocm* leaf, so a custom
   mirror whose leaf starts with rocm but is not a pip family (a private rocm-current
   mirror, a Radeon find-links rocm-rel-7.2.1) escaped the bounds and installed bare
   torchvision/torchaudio. It now digit-gates rocm to rocm[0-9]* (matching the Python
   _is_pip_rocm_family_leaf ^rocm\d), so those custom leaves return "" and the <2.11
   companion caps apply; real rocm7.2 / gfx per-arch indexes still classify as rocm.

2. _tauri_torch_index_family classified by the raw last path segment, so a pinned URL
   carrying auth in the query (.../rocm7.2?token=SECRET) had the token echoed verbatim
   into the emitted [TAURI:DIAG] line. It now strips query/fragment before classifying
   (mirroring the marker/log credential stripping), so no token reaches the diagnostic
   output; as a side effect .../cu128?token=x now classifies as cu128 instead of auto.

3. On studio update, the core package step (a newer unsloth can require a torch the
   custom pin does not satisfy, pulling a default PyPI trio) runs BEFORE the step-2b
   verbatim check, which then recorded the already-clobbered trio as the baseline for a
   matching marker and left the pin unapplied. A new _capture_verbatim_baseline() records
   the pre-clobber trio before the core step, so the verbatim pass detects the drift and
   reapplies the pin. Captures only for a matching custom pin with importable torch; a
   mismatched/absent marker or broken torch is left to _ensure_verbatim_torch_index.

Tests: _expected_torch_flavor_tag rocm-current / rocm-rel cases; _tauri_torch_index_family
token/fragment redaction with a no-leak regression guard; _capture_verbatim_baseline
record/skip cases plus an end-to-end clobber-detection scenario; a structural guard that
the capture runs before the core step. 501 python + shell suites pass; install.sh bash -n
clean, shellcheck unchanged from base.

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

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

* install: match rocm family leaves exactly, enforce the rocm7.2 torch line, repair a broken pinned torch

A pinned index is a pip ROCm --index-url family only when its leaf is an exact
rocm<digits> / rocm<digits>.<digits> (rocm7.2) or a gfx* per-arch leaf. The prior
^rocm[0-9] prefix match also caught suffixed private-mirror leaves (rocm7.2-private,
rocm7-current), routing them through the ROCm/companion-family path instead of the
verbatim pin: the companion bounds were skipped and, on a pre-marker venv with a
compatible +rocm wheel, the pin was never applied. Match the family exactly through one
shared helper at every site:
  - install_python_stack.py: _is_pip_rocm_family_leaf (re.fullmatch), plus the two other
    loose gates it feeds (_normalize_family_leaf, _torch_flavor_matches_pin).
  - install.sh: a new _is_pip_rocm_family_leaf routes _expected_torch_flavor_tag,
    _torch_index_repairable, _normalize_family_leaf and the ROCm side-effect gate.
  - setup.ps1: a new Test-PipRocmFamilyLeaf routes Get-NormalizedFamilyLeaf and both
    pinned reroutes; install.ps1 anchors its reroute regex.

_rocm_pin_family_mismatch (and its setup.ps1 mirror Get-RocmPinStaleTags) compared only
the ROCm version, so a +rocm7.2 wheel whose torch release drifted off the 2.11 line
(2.12/2.13 from an out-of-band upgrade or a custom rocm7.2 mirror) satisfied the family
check while violating _ROCM_TORCH_PKG_SPECS['rocm7.2'] (torch>=2.11,<2.12). Flag it stale
so the repair reinstalls to floor; >=2.11 alone is not enough, so the release is compared
exactly against the 2.11 line for a KNOWN-2.11 rocm pin.

_ensure_pinned_known_family_torch returned on a failed import probe, but
_torch_pin_needs_apply forces the dependency pass on that same failed probe: a broken
torch under a known-family pin was left in place and the pass was forced on every update.
Treat an unimportable torch as drift and reinstall the pinned trio (the spec and marker
derive from the pinned leaf, not the absent flavor); once it lands the probe succeeds and
the fast path returns.

Tests: exact-match cases across test_torch_flavor.sh, test_rocm_support.py,
test_cross_platform_parity.py and the two .ps1 helper suites; the rocm7.2 release-line
and broken-probe-reinstall cases; extraction lists updated for the new helpers.

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

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

* install: anchor the PS pinned-ROCm floor gate and bound install.ps1 custom-pin companions

Round 12 made every family CLASSIFIER exact, but the Windows install-flow floor gate reads
$_pinRocm211 directly from the raw pinned leaf with an unanchored -match '^rocm(\d+)\.(\d+)'
BEFORE any exact classification runs. A suffixed custom leaf (rocm7.2-private) matches that
rocm7.2 prefix, so it takes the 2.11-floor branch and is force-routed through the ROCm
install path before the exact-match elseif can send it to the verbatim install. Anchor the
match ($) in both install.ps1 and setup.ps1 so only an exact rocmX.Y leaf is floored; a
suffixed or newer-suffix leaf falls through to the verbatim path. The Python floor
selection is already exact (dict lookups gated on _is_pip_rocm_family_leaf), so only the two
PS scripts needed this.

install.ps1's custom (non-cu-family) pinned-torch install bounded torch>=2.4,<2.11.0 but
left torchvision/torchaudio bare, so a private mirror serving newer companions could pull a
wheel built for a newer torch ABI while the marker records the pin as applied. Bound both
companions (torchvision>=0.19,<0.26.0 / torchaudio>=2.4,<2.11.0) when the leaf is not a
cu<digits> family index (a cu index bounds its own resolution), matching setup.ps1's
Test-CudaFamilyLeaf gate and _CUSTOM_INDEX_TORCH_PKG_SPEC.

Tests: parity guards for the anchored floor gate in both PS scripts and for install.ps1's
bounded custom-pin companions.

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

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

* install: tighten comments in the torch-index-override paths

Collapse the verbose comment and docstring blocks added across the installer
scripts and their tests to fewer, clearer lines without changing behaviour.
Remove a duplicated CUDA-spec comment block. Comments/docstrings only; no code
changes (AST-verified).

* install: repair a broken pinned torch on Linux, strip trailing slash in tauri family, count the final step

_ensure_cuda_torch / _ensure_cpu_torch returned on a failed import probe (torch present but
unimportable). With an explicit CUDA/CPU pin, _torch_pin_needs_apply forces the dependency
pass on that same failed probe, and the base package update does not force-reinstall an
already-installed torch distribution, so the broken torch was left in place and the pass
reran every update without repairing it. Treat a failed probe under a pin as drift and
reinstall from the pinned index (the reinstall rewrites the marker and the next probe
imports, so no loop). This is the Linux counterpart of the known-family repair fix.

_tauri_torch_index_family stripped the query/fragment before classifying but not a trailing
slash, so a token-authenticated pin like .../cu128/?token=x collapsed to .../cu128/ and fell
through the exact-suffix */cu128 and */cpu arms to "auto". Strip a trailing slash too,
mirroring _torch_index_url_leaf.

The Windows / macOS-ARM final torch-repair step (_ensure_pinned_known_family_torch) runs a
progress step that base_total never counted (the final-step increment was gated to Linux),
so _STEP ran one past _TOTAL on those platforms. Add the missing increment.

Tests: broken-probe reinstall for the CUDA (family and URL pins) and CPU paths; trailing
slash / slash+token cases for _tauri_torch_index_family; a full-flow progress-count guard
asserting _STEP == _TOTAL on Windows and Linux.

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

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

* install: tighten comments in the torch-index-override paths

* install: harden the torch-index pin across all four installers

Redact index-URL credentials from captured install logs before they print on
failure. uv/pip failure text embeds the failing --index-url verbatim, so a
user:token@ or ?token= secret could leak into the console. Add a shared
redaction pass (_redact_install_output / Redact-InstallOutput) wired into the
error-output dump in install.sh, install.ps1, setup.ps1 and
install_python_stack.py. Verbose mode still streams live uncaptured output, so
it is intentionally left unredacted (developer opt-in).

Trim trailing slashes on the PATH only for a verbatim UNSLOTH_TORCH_INDEX_URL
override, preserving a ?query/#fragment token. A whole-URL rstrip corrupted a
base64 token ending in "/", and a single-slash strip left .../cu128//
classifying as an empty leaf. Add _trim_index_path_slashes /
Trim-IndexPathSlashes and route the override through it; strip ALL trailing
slashes in the backend-branding leaf classifier so a double slash still yields
the real leaf.

Reject a trailing-dot ROCm leaf (rocm7.) in the bash family validator so it
matches Python re.fullmatch(rocm\d+(?:\.\d+)?) and the PowerShell regex: both the
major and the minor must be non-empty digits, so rocm7. is a custom verbatim pin,
not a pip ROCm family.

Scrub PIP_NO_INDEX and PIP_INDEX_URL for a pinned install in the two installers
that have a plain-pip fallback (install_python_stack.py, setup.ps1):
PIP_NO_INDEX=1 makes the fallback ignore every index including the pinned
--index-url, and PIP_INDEX_URL replaces it. install.sh and install.ps1 install
via uv --default-index (which ignores pip config/env), so they are unaffected.

Add unit tests (bash, Python, PowerShell) and cross-platform parity tests
covering credential redaction, path-only slash trimming, the rocm7. validator,
the double-slash leaf, and the PIP_NO_INDEX/PIP_INDEX_URL scrub.

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

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

* install: redact captured torch-install output and warn on a failed pinned ROCm repair

Close a redaction gap the earlier pass missed: setup.ps1's direct
`Fast-Install ... | Out-String` branches (ROCm from $ROCmIndexUrl, CPU/CUDA from
$TorchInstallIndexUrl, plus the Triton and T5 sub-venv installs) printed the
captured $output verbatim on failure, bypassing Redact-InstallOutput. A private
index carrying userinfo or a ?token= in the pin could leak into Windows Studio
setup logs. Route every `Write-Host $output` through Redact-InstallOutput.

Warn on a failed pinned Windows ROCm reinstall in
_ensure_pinned_known_family_torch: the branch printed "reinstalling from it" then
called pip_install_try, but had no else, so a failure continued silently and left
the user believing the pin was applied while the old CPU/wrong torch survived.
Mirror the auto-ROCm Windows path and warn, telling the user to retry.

* install: redact captured output on the pip fallback and optional-install failure paths

The uv install path already redacted its captured output, but pip_install's pip
fallback runs through run(), which printed result.stdout verbatim on failure, and
_print_optional_install_failure did the same. A pinned --index-url carrying
userinfo or a ?token= could still leak there when uv is unavailable or the pip
fallback also fails. Route both through _redact_install_output. The verbose
pip_install_try path stays raw (developer opt-in), matching the other installers.

* install: split the survive-updates marker subsystem into a follow-up

The torch-index override PR grew a persisted per-venv marker plus repair
machinery (stale-pin detection, verbatim re-apply, update-time reinstall
triggers) that roughly doubled it. That subsystem is orthogonal to the core
feature and is being reworked in a follow-up (versioned/hashed marker,
full-URL pin baseline), so it moves there wholesale instead of shipping
twice.

What this PR still does: UNSLOTH_TORCH_INDEX_URL / UNSLOTH_TORCH_INDEX_FAMILY
pick the torch wheel index at install time in all four installers, with the
exact rocm/gfx/cpu/cu leaf classification, the torch 2.11 floor for the
per-arch AMD indexes, bounded companions for custom leaves, credential
redaction of captured installer output, path-only slash trimming, and the
uv/pip index env scrubs. Flavor-based repair keeps honoring the pin: a wrong
family under an explicit pin still reinstalls from the pinned URL, and
setup.ps1 repairs a pinned stale venv in place instead of wiping it.

What moves to the follow-up: the .unsloth-torch-index marker file and its
writers/readers/normalizers, exact-URL pin-change detection on update
(same-tag gfx switches, custom-mirror repoints), the verbatim trio snapshot
and clobber re-apply, the pin-baseline recorder, and the
--torch-pin-needs-apply fast-path probe in setup.sh / setup.ps1. Their tests
(the marker sh/ps1 suites, the stale-pin suite, and the marker classes in the
rocm/cuda/parity suites) move with them; the removed code is preserved on a
local archive branch to seed that PR.

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

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

* install: re-apply a ROCm pin over an existing HIP wheel via the version tag

The subsystem split left an explicit ROCm/gfx pin unenforced on `studio
update` whenever the venv already imported ANY ROCm torch: the pinned
reinstall lived inside the `elif not has_hip_torch` branch, so a rocm6.4 to
rocm7.2 switch, a gfx1151 pin over a generic +rocm7.2 wheel, or a broken
2.12+rocm7.2 drift never re-applied the pin.

Restore the markerless half of that detection: _rocm_pin_family_mismatch
compares the pinned leaf against the installed wheel tag (exact rocmX.Y
compare, the 2.11 gfx per-arch allowlist, the untagged-wheel rule), the HIP
probe emits "<hip_marker>|<version>" again so the installed tag is available,
and _ensure_rocm_torch reinstalls from the pinned URL when the tag mismatches
even though HIP torch is present. setup.ps1 mirrors it: the stale-venv check
routes a pinned rocm/gfx leaf through Get-RocmPinStaleTags instead of
collapsing it to a generic "rocm" flavor, and the existing pinned in-place
repair (no wipe) applies the change.

What still waits for the follow-up marker PR, by design: pin changes the
wheel tag cannot see -- a per-arch switch between two 2.11 gfx indexes
(identical +rocm7.13.0 tag), a custom-mirror URL repoint under the same
family leaf, and unknown-family verbatim pins. Those need the persisted
index record.

Tests restored with the code: the _rocm_pin_family_mismatch table, the five
update-path cases (older-rocm reinstall, gfx-over-pre-2.11 reinstall,
matching-pin no-reinstall, non-2.11 gfx no-reinstall, gfx-over-generic-2.11
reinstall), the "|" probe-format guards, and the AST-extracted
Get-RocmPinStaleTags suite for setup.ps1.

* install: compare major-only rocm pins, redact URL fragments, bound pinned CPU trio

Three review fixes on the restored pin-repair path.

The family classifier accepts a major-only rocm<d> leaf (rocm7), but the
mismatch comparators only parsed rocmX.Y, so a rocm7 pin fell through to the
2.11-line fallback and INVERTED both verdicts: an installed +rocm6.4 wheel
compared as satisfied (pin never re-applied) while a matching +rocm7.2 wheel
compared as stale (reinstall loop). Major-only pins now compare on the major
alone in _rocm_pin_family_mismatch and Get-RocmPinStaleTags: rocm6.x under a
rocm7 pin is a mismatch, any rocm7.x satisfies it, an untagged wheel never
does, and a bare +rocm tag with an unreadable version is accepted (matching
the existing lenient unreadable fallback).

The output redactors scrubbed userinfo and ?query= values but not #fragments,
so a pin like https://mirror/whl/cu128#token=secret leaked the secret in
captured uv/pip failure text -- inconsistent with the URL handling itself,
which already treats fragments as sensitive. All four redactors gain a
URL-anchored fragment rule (anchored so a bare "# comment" line in tool
output is never touched).

setup.ps1's CPU branch installed a bare torch/torchvision/torchaudio trio;
fine for the unpinned host default, but a PINNED cpu index routes through the
same branch and the /cpu index serves newer torch, so a fresh pinned CPU
install could land an unsupported trio that _ensure_cpu_torch then keeps
(it accepts any CPU build). Under a pin the branch now installs the bounded
trio mirroring _CPU_TORCH_PKG_SPEC (torch>=2.4,<2.12.0 and matching
companions); the unpinned path is unchanged.

Tests: major-only rows in the Python mismatch table and the AST-extracted
setup.ps1 suite; fragment + query-plus-fragment + bare-hash-comment cases in
all four redactor suites; a parity check that the pinned CPU trio bounds
exist, are gated on the pin, and mirror the Python repair spec.

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

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

* install: tighten comments in the torch index override paths

* tests: track the moved pass-through inheritance in the gguf order check

Main moved the llama_extra_args pass-through inheritance out of the
GGUF branch into _resolve_inherited_extra_args, which runs before it,
so the source-order assertion's "if request.llama_extra_args is None"
anchor no longer exists inside the branch and the check failed after
the main merge. The test now asserts the same property in the current
shape: inheritance before the GGUF branch (a carried --no-mmproj still
shapes the hub guard's companion requirement), and marker, hub guard,
unload in order within the branch. Full file passes (32 tests).

* tests: anchor the inheritance order check on the call, not the definition

source.index("_resolve_inherited_extra_args(") matched the function
definition, which always precedes the endpoint, so the ordering
assertion was vacuously true. Anchoring on "= _resolve_inherited_
extra_args(" pins the first call site inside the load endpoint (line
4505), which is the statement whose position relative to the GGUF
branch the test is meant to guard. 32 tests pass.

* tests: align the gguf order test with main

Main fixed the stale ordering assertion in PR 7252; adopting its
version verbatim removes this file from the branch diff entirely and
avoids a conflict on the next main merge. 32 tests pass.

* install: bound the companion constraints to torch's window everywhere

A full platform x vendor validation matrix over this branch surfaced a
real trio mismatch on the cpu/mac paths: torch is capped <2.11 (installs
2.10.0+cpu) but the bare torchaudio companion resolves 2.11.0+cpu,
because torchaudio 2.11 dropped its exact torch pin. Reproduced in a
sandboxed end to end cpu install. torchvision still exact-pins torch and
self-corrected.

The default companion constraints are now bounded to torch's window
(<0.26 / <2.11) and widen together with the cu* torch window (<0.27 /
<2.12), so every leaf resolves a paired trio. Verified with uv dry-runs
on the cpu, cu130, and rocm6.4 leaves (2.10.0/0.25.0/2.10.0,
2.11.0/0.26.0/2.11.0, 2.9.1/0.24.1/2.9.1) and a rerun of the sandboxed
cpu install, which now lands torch 2.10.0+cpu with torchaudio
2.10.0+cpu.

The Strix WSL reroute now also forwards UNSLOTH_TORCH_INDEX_URL and
UNSLOTH_TORCH_INDEX_FAMILY into the rerouted 24.04 distro; dropping
them silently reverted the child install to auto-detection, defeating
the pin this branch introduces.

test_torch_constraint.sh updated: the bounded companions must appear at
the defaults and the custom-leaf block, no bare companion may remain,
and the cu* widen must carry the companions with it.

* install: harden the override path against reroute drift and credential leaks

Review sweep focused on default-path idempotency found no defects on the
unset path; these fixes cover the override path and failure reporting.

install.sh:
- The early WSL Strix Halo distro reroute now honors an explicit index
  pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY): the pin is used in the current
  distro instead of probing the GPU and re-entering another distribution,
  matching the contract of the later Radeon and Strix guards. Whitespace
  only values do not gate, in parity with get_torch_index_url.
- Verbose mode now streams installer output through the credential
  redactor; it previously bypassed the redaction the quiet path applies.
  The exit code survives the pipe via an rc file since the script runs
  under plain sh with no pipefail.
- The kept-release fallback warning now strips credentials from the
  index URL before printing it.

install.ps1:
- Bounded torchvision and torchaudio next to every capped torch install
  (custom pin, ROCm CPU fallback, CUDA flavor repair). torchaudio 2.11
  dropped its exact torch pin from the wheel metadata, so a bare
  companion beside torch<2.11 can resolve a mismatched 2.11.0 build,
  cu family indexes included. Mirrors the install.sh companion bounds.

studio/install_python_stack.py:
- The verbose failure path now redacts index URLs in pip and uv output
  before printing, matching every other output site in the file.

All sh, ps1 and python installer test suites pass (the host-defaults
suite has a known pre-existing failure unrelated to this change).

* install: redact verbose Windows installer output and repair the parity tests

Follow-ups to the override-hardening commit, from review:

- install.ps1 Invoke-InstallCommand and setup.ps1 Invoke-SetupCommand now
  pipe verbose output through Redact-InstallOutput per record, and the
  three verbose Fast-Install torch call sites (ROCm, CPU, CUDA) do the
  same: uv and pip echo the pinned index URL, credentials included, in
  their errors, and verbose mode previously bypassed the redaction the
  quiet paths apply. ForEach-Object and Out-Host leave $LASTEXITCODE
  untouched, verified with a native command exiting 7 behind the pipe.

- test_cross_platform_parity.py: the install.ps1 companion-bounds
  assertion now matches the implemented behavior (bounds on every index,
  no cu-family exemption, since torchaudio 2.11 dropped its exact torch
  pin) instead of requiring the removed $_pinCuLeaf gate.

- test_rocm_support.py: the WSL reroute guard test slices the whole
  function body to its closing brace instead of a fixed 1200-character
  window, which the new pin-gate preamble had outgrown.

428 tests pass across the parity, install stack and rocm support suites;
the sh and ps1 installer suites pass unchanged.

* install: tighten comments in the torch-index and ROCm/CUDA repair paths

* install: digit-gate the gfx family leaf and honor ROCm pins in the Windows repair

Two review follow-ups on the override path:

- The pip ROCm family predicate accepted ANY gfx-prefixed leaf, so a
  custom verbatim pin like /gfx-private classified as a ROCm family and
  enabled the ROCm-only side effects (AMD bitsandbytes, ROCm torch
  repair) on a mirror that may serve CPU/CUDA wheels. gfx now requires a
  following digit (gfx90a, gfx1151, gfx120X-all), consistently in
  install.sh, install_python_stack.py, install.ps1 (family gate and
  expected-flavor classifier) and setup.ps1, matching the strictness the
  rocm side already had (rocm7.2-private stays verbatim). The broader
  backend BRANDING globs are unchanged on purpose: radeon repo leaves
  (rocm-rel-X.Y) must still brand the rocm backend without being
  force-repaired as a family.

- The Windows branch of the ROCm torch repair always installed from the
  public per-arch index, ignoring an explicit ROCm-family pin: after a
  pinned setup.ps1 install failed to a CPU base, the repair retried
  repo.amd.com instead of the pinned index. The branch now resolves
  _explicit_rocm_torch_index_url() first, uses it as the install index
  when set, and mirrors the Linux pin contract by skipping the NVIDIA
  and gfx-detection gates a pin is documented to override.

Source-assertion tests updated to the tightened predicate and the new
repair label. 1165 tests pass across the parity, install stack and
studio install suites; the sh and ps1 suites pass; both PowerShell
installers parse clean.

* Remove scratch archives accidentally committed with the comment pass

The temp/ archive copies of installer and test files were working
scratch, not PR content, and inflated the diff by about nine thousand
lines.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-07-20 00:58:52 -07:00 committed by GitHub
commit 3ab8dce97a
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 3382 additions and 386 deletions

View file

@ -53,7 +53,8 @@ function Install-UnslothStudio {
param([string]$TorchIndexUrl)
if ($SkipTorch) { return "none" }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" }
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
# Drop query/fragment first so a token-authenticated pin classifies by family.
$leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf }
if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf }
return "auto"
@ -62,7 +63,8 @@ function Install-UnslothStudio {
function Get-TauriGpuBranch {
param([string]$TorchIndexFamily)
if ($SkipTorch) { return "no_torch" }
if ($TorchIndexFamily -like "cu*") { return "cuda" }
# Require a digit after "cu" so /current or /custom isn't branded CUDA (parity ^cu[0-9]).
if ($TorchIndexFamily -match '^cu[0-9]') { return "cuda" }
if ($TorchIndexFamily -like "rocm*") { return "rocm" }
if ($TorchIndexFamily -eq "cpu") { return "cpu" }
return "unknown"
@ -467,22 +469,35 @@ function Install-UnslothStudio {
}
}
# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer
# output before printing on failure; uv/pip errors echo the failing --index-url verbatim.
# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted.
function Redact-InstallOutput {
param([string]$Text)
if (-not $Text) { return $Text }
$Text = $Text -replace '(https?://)[^/@\s`]+@', '$1<redacted>@'
$Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=<redacted>'
# A #token=... fragment is as sensitive as a query; URL-anchored.
return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#<redacted>'
}
# Run native commands quietly by default to match install.sh behavior.
# Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1.
function Invoke-InstallCommand {
param(
[Parameter(Mandatory = $true)][ScriptBlock]$Command
)
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when the command pins an index, clear every uv index env var so
# it wins, then restore in finally. Other installs keep the user's mirror.
# Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898):
# for --default-index, clear the uv index env vars (restore in finally) and set
# UV_NO_CONFIG=1 so a uv.toml/pyproject index can't outrank the CLI pin (uv 0.10).
$savedUvIndex = $null
if ($Command.ToString() -match '--default-index') {
$savedUvIndex = @{}
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'UV_CONFIG_FILE', 'UV_NO_CONFIG') {
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
$env:UV_NO_CONFIG = '1'
}
$prevEap = $ErrorActionPreference
$ErrorActionPreference = "Continue"
@ -493,17 +508,23 @@ function Install-UnslothStudio {
# Merge stderr into stdout so progress/warning output stays visible
# without flipping $? on successful native commands (PS 5.1 treats
# stderr records as errors that set $? = $false even on exit code 0).
& $Command 2>&1 | Out-Host
# Redact per record: uv echoes index URLs (credentials and all) in
# its errors, and verbose mode must not bypass the quiet path's
# redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched.
& $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
} else {
$output = & $Command 2>&1 | Out-String
if ($LASTEXITCODE -ne 0) {
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
}
}
return [int]$LASTEXITCODE
} finally {
$ErrorActionPreference = $prevEap
if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } }
if ($savedUvIndex) {
Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue
foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } }
}
}
}
@ -1960,10 +1981,31 @@ exit 0
# On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint.
if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint }
# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL
# TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared.
function Trim-IndexPathSlashes {
param([string]$Url)
$value = $Url.Trim()
$idx = $value.IndexOfAny([char[]]@('?', '#'))
if ($idx -lt 0) {
return $value.TrimEnd('/')
}
return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx)
}
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
# Mirrors Get-PytorchCudaTag in setup.ps1.
function Get-TorchIndexUrl {
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
# Explicit pin -- skip ALL GPU probing (headless / CI / cross-install).
# UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf appended
# to the mirror base. Matches install.sh / install_python_stack.py.
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL)
}
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
}
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = Invoke-NvidiaSmiBounded $NvidiaSmiExe
@ -1984,6 +2026,25 @@ exit 0
return "$baseUrl/cu126"
}
# Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with
# _strip_index_url_credentials (install.sh / py / setup.ps1).
function Remove-IndexUrlCredentials {
param([string]$Url)
$sep = $Url.IndexOf('://')
if ($sep -lt 0) { return $Url }
$scheme = $Url.Substring(0, $sep)
$rest = $Url.Substring($sep + 3)
# Drop query / fragment (may hold auth tokens).
$q = $rest.IndexOfAny([char[]]('?', '#'))
if ($q -ge 0) { $rest = $rest.Substring(0, $q) }
$slash = $rest.IndexOf('/')
$authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest }
$at = $authority.LastIndexOf('@')
$host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority }
if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" }
return "${scheme}://${host_}"
}
# ── 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.
@ -2002,11 +2063,13 @@ exit 0
param([string]$TorchIndexUrl, [string]$ROCmIndexUrl)
if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null }
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
# Drop query/fragment first so .../cu128?token=x classifies as cu128 (else it reinstalls every run).
$leaf = (($TorchIndexUrl -split '[?#]', 2)[0].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' }
# gfx must be followed by a digit (an architecture leaf); gfx-private is custom.
if ($leaf -match '^gfx[0-9]') { return 'rocm' }
return $null
}
@ -2041,6 +2104,10 @@ exit 0
} catch { return $null }
}
# An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it
# (e.g. a deliberate cpu pin on an AMD host).
$TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or `
(-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY))
$TorchIndexUrl = Get-TorchIndexUrl
# ── GPU arch → newest compatible Windows ROCm wheel release ──
@ -2052,7 +2119,9 @@ exit 0
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
$ROCmIndexUrl = $null
$ROCmTorchFloor = $null
if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
$PinnedRocmVisionSpec = $null
$PinnedRocmAudioSpec = $null
if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
@ -2102,6 +2171,32 @@ exit 0
}
}
# A gfx*/rocm pin skips the auto-reroute above, but the generic CPU/CUDA install below
# would use torch>=2.4,<2.11 and pull a known-bad wheel on the gfx115x/gfx120x/rocm>=7.2
# indexes (the _grouped_mm bug). Route a pinned ROCm index through the ROCm path.
if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) {
$_pinLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLower()
$_pinRocm211 = $false
# Anchor ($) so a suffixed custom leaf (rocm7.2-private) falls through to verbatim.
if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') {
# Only KNOWN-2.11 rocm (rocm7.2) gets the floor. Matches Test-RocmKnown211Version.
$_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2)
}
# Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare.
$_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf
if ($_pinGfx211 -or $_pinRocm211) {
$ROCmIndexUrl = $TorchIndexUrl
$ROCmTorchFloor = "torch>=2.11.0,<2.12.0"
$PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0"
$PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0"
substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan"
} elseif ($_pinLeaf -match '^gfx[0-9]' -or $_pinLeaf -match '^rocm[0-9]+(\.[0-9]+)?$') {
# Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with
# bare specs. Only EXACT rocm<digits>/gfx* are families; a suffixed leaf is verbatim.
$ROCmIndexUrl = $TorchIndexUrl
}
}
if ($ROCmIndexUrl) {
$TorchIndexFamily = "rocm"
} else {
@ -2164,8 +2259,8 @@ exit 0
}
if ($_Migrated) {
# Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
# in the new venv location, while preserving existing torch/CUDA
# Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
# existing torch/CUDA unless the flavor repair below re-lands it.
Write-TauriLog "STEP" "Installing unsloth"
substep "upgrading unsloth in migrated environment..."
if ($SkipTorch) {
@ -2210,22 +2305,24 @@ exit 0
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
} elseif ($ROCmIndexUrl) {
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
substep "installing PyTorch from $ROCmIndexUrl..."
substep "installing PyTorch from $(Remove-IndexUrlCredentials $ROCmIndexUrl)..."
$torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
# Pin the companions to match $torchSpec; bare names can resolve an
# ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
if ($torchInstallExit -ne 0) {
# Transient AMD-index failure: fall back to a CPU base so the install
# still completes; Unsloth setup retries ROCm afterwards.
# Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries
# ROCm). Use an explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS
# the ROCm mirror, so reusing it would just retry it.
$CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" }
substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow"
# --force-reinstall: a failed ROCm install can leave an unpinned ROCm
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
# torch>= range, so without it uv would keep the ROCm build and only swap
# the companions -- a mismatched venv the flavor-repair block won't fix.
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2238,8 +2335,14 @@ exit 0
}
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..."
# Bound the companions to the capped torch on EVERY index, cu<digits>
# families included: torchaudio 2.11 dropped its exact torch pin from
# the wheel metadata, so a bare companion next to torch<2.11 can
# resolve a mismatched 2.11.0 build. Mirrors install.sh.
$_pinVisionSpec = "torchvision>=0.19,<0.26.0"
$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2335,8 +2438,8 @@ exit 0
$rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
# Pin companions like the fresh ROCm path (bare names can pull an
# ABI-incompatible torchvision/torchaudio from the per-arch index).
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
if ($torchFixExit -ne 0) {
@ -2347,7 +2450,7 @@ exit 0
} 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 --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $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)

View file

@ -159,18 +159,58 @@ run_maybe_quiet() {
fi
}
# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL
# strip corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared.
_trim_index_path_slashes() {
_tips_v="$1"
case "$_tips_v" in
*[?#]*)
_tips_head="${_tips_v%%[?#]*}"
_tips_tail="${_tips_v#"$_tips_head"}"
;;
*)
_tips_head="$_tips_v"
_tips_tail=""
;;
esac
while [ -n "$_tips_head" ] && [ "${_tips_head%/}" != "$_tips_head" ]; do
_tips_head="${_tips_head%/}"
done
printf '%s%s' "$_tips_head" "$_tips_tail"
}
# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer
# output before printing on failure; uv/pip errors echo the failing --index-url verbatim.
# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted.
_redact_install_output() {
sed -E \
-e 's#(https?://)[^/@[:space:]`]+@#\1<redacted>@#g' \
-e 's#([?&][^=[:space:]&`]+)=[^&#[:space:]`]+#\1=<redacted>#g' \
-e 's|(https?://[^[:space:]`#]+)#[^[:space:]`]+|\1#<redacted>|g' \
"$@"
}
run_install_cmd() {
_label="$1"
shift
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when we pass --default-index, neutralize every uv index env var so
# the pinned index wins. Other installs keep the user's mirror.
# Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898):
# for --default-index, neutralize the uv index/backend/config vars (UV_TORCH_BACKEND
# redirects torch; UV_NO_CONFIG=1 + dropping UV_CONFIG_FILE stops a uv.toml/pyproject
# index outranking the CLI pin, uv 0.10).
case " $* " in
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL -u UV_TORCH_BACKEND -u UV_FIND_LINKS -u UV_CONFIG_FILE UV_NO_CONFIG=1 "$@" ;;
esac
if _is_verbose; then
"$@" && return 0
_rc=$?
# Stream through the redactor: uv echoes index URLs (credentials and
# all) in its errors, and verbose mode previously bypassed the
# redaction the quiet path applies. The rc file preserves the
# command's exit code across the pipe without relying on pipefail
# (this script runs under plain sh).
_rcf=$(mktemp)
{ "$@" 2>&1; printf '%s' "$?" > "$_rcf"; } | _redact_install_output
_rc=$(cat "$_rcf" 2>/dev/null || echo 1)
rm -f "$_rcf"
[ "${_rc:-1}" -eq 0 ] 2>/dev/null && return 0
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
return "$_rc"
fi
@ -178,7 +218,7 @@ run_install_cmd() {
"$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; }
_rc=$?
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
cat "$_log" >&2
_redact_install_output "$_log" >&2
rm -f "$_log"
return $_rc
}
@ -257,7 +297,7 @@ _install_bnb_rocm() {
fi
_bnb_rc=$?
if _is_verbose; then
cat "$_bnb_log" >&2
_redact_install_output "$_bnb_log" >&2
fi
rm -f "$_bnb_log"
step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2
@ -310,6 +350,11 @@ _tauri_torch_index_family() {
return
fi
_diag_url="${1:-}"
# Strip query/fragment AND a trailing slash before classifying (like _torch_index_url_leaf):
# a token isn't echoed into [TAURI:DIAG], and .../cu128/?token=x still classifies as cu128.
_diag_url="${_diag_url%%\?*}"
_diag_url="${_diag_url%%#*}"
_diag_url="${_diag_url%/}"
case "$_diag_url" in
*/cu118) echo "cu118" ;;
*/cu124) echo "cu124" ;;
@ -343,7 +388,8 @@ _tauri_gpu_branch() {
return
fi
case "$_diag_family" in
cu*) echo "cuda" ;;
# Require a digit after cu so /current or /custom isn't branded CUDA (parity ^cu[0-9]).
cu[0-9]*) echo "cuda" ;;
rocm*)
if [ "$_diag_radeon" = true ]; then
echo "rocm_radeon"
@ -1575,6 +1621,12 @@ _has_usable_nvidia_gpu() {
# the STUDIO_HOME mkdir/venv so the origin distro is untouched.
_maybe_reroute_strixhalo_to_2404() {
[ "${OS:-}" = "wsl" ] || return 0
# An explicit index pin skips every GPU-driven reroute (same contract as
# the later Radeon/Strix guard): the pin is honored in THIS distro rather
# than probing the GPU and switching distributions. Whitespace-only
# overrides do not gate (parity with get_torch_index_url).
_rr_pin=$(printf '%s' "${UNSLOTH_TORCH_INDEX_URL:-}${UNSLOTH_TORCH_INDEX_FAMILY:-}" | tr -d '[:space:]')
[ -n "$_rr_pin" ] && return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
[ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0
@ -1636,6 +1688,10 @@ _maybe_reroute_strixhalo_to_2404() {
# Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the
# GPU instead of falling back to the desktop-app prompt path.
[ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1"
# Forward a pinned torch index into the rerouted distro; dropping it would
# silently revert the child install to auto-detection.
[ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_URL=$(_rr_q "$UNSLOTH_TORCH_INDEX_URL")"
[ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_FAMILY=$(_rr_q "$UNSLOTH_TORCH_INDEX_FAMILY")"
[ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1"
_rr_args=""
[ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")"
@ -2001,6 +2057,15 @@ if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; t
TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
fi
fi
# Companion (torchvision/torchaudio) constraints, bounded to torch's window.
# torchaudio 2.11 dropped its exact torch pin, so a bare companion next to a
# <2.11-capped torch resolves torchaudio 2.11 (verified: cpu leaf installed
# torch 2.10.0+cpu with torchaudio 2.11.0+cpu). torchvision still exact-pins
# torch and self-corrects, but is bounded for symmetry. Widened alongside the
# cu* torch window below; the torch-2.11 AMD paths (rocm7.2 / per-gfx / Strix)
# pin their own trio.
TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"
# ── Resolve repo root (for --local installs) ──
_REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)"
@ -2069,6 +2134,24 @@ _has_amd_rocm_gpu() {
get_torch_index_url() {
_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
_base="${_base%/}"
# Explicit override -- skip ALL GPU probing (headless / container / CI / cross-install).
# UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf (cpu, cu128, ...)
# appended to the mirror base. Trim whitespace so a whitespace-only value is unset.
_url="${UNSLOTH_TORCH_INDEX_URL:-}"
_url="${_url#"${_url%%[![:space:]]*}"}"; _url="${_url%"${_url##*[![:space:]]}"}"
if [ -n "$_url" ]; then
# Trim trailing PATH slashes (a multi-slash path 404s on strict pip proxies) while
# preserving a ?query/#fragment token (a whole-URL strip would eat a "/"-ending token).
_url=$(_trim_index_path_slashes "$_url")
echo "$_url"; return
fi
_family="${UNSLOTH_TORCH_INDEX_FAMILY:-}"
_family="${_family#"${_family%%[![:space:]]*}"}"; _family="${_family%"${_family##*[![:space:]]}"}"
if [ -n "$_family" ]; then
while [ "${_family#/}" != "$_family" ]; do _family="${_family#/}"; done
while [ "${_family%/}" != "$_family" ]; do _family="${_family%/}"; done
echo "$_base/$_family"; return
fi
# macOS: always CPU (no CUDA support)
case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac
# Try nvidia-smi -- require the binary to actually list a usable GPU.
@ -2197,6 +2280,45 @@ _torch_flavor_tag() {
esac
}
# Final path segment of a wheel index URL ($1), lowercased, query/fragment stripped first
# so a token-authenticated pin (.../cu128?token=x) classifies as cu128 (else it reinstalls
# every update). Classification only. Shared with the py / ps1 leaf extractors.
_torch_index_url_leaf() {
_tl_u="${1%%\?*}"
_tl_u="${_tl_u%%#*}"
# Strip ALL trailing slashes, not one: .../rocm7.2// must yield rocm7.2, not an empty leaf.
while [ -n "$_tl_u" ] && [ "${_tl_u%/}" != "$_tl_u" ]; do
_tl_u="${_tl_u%/}"
done
printf '%s' "${_tl_u##*/}" | tr '[:upper:]' '[:lower:]'
}
# True (exit 0) when a lowercased leaf is an EXACT pip ROCm family: rocm<digits>[.<digits>]
# or a gfx ARCHITECTURE leaf (gfx followed by a digit: gfx90a, gfx1151, gfx120x-all). A leaf
# that merely starts with rocm/gfx (rocm7.2-private, gfx-private) is a custom verbatim pin.
# Matches the py / ps1 sides.
_is_pip_rocm_family_leaf() {
case "$1" in
gfx[0-9]*) return 0 ;;
rocm[0-9]*)
# Exact rocm<digits>[.<digits>]: both major and minor must be non-empty all-digits
# (rocm7., rocm7.2.1, rocm7.2-private are all custom pins, not a family).
_rocm_rest="${1#rocm}"
case "$_rocm_rest" in
*.*.*) return 1 ;;
*.*)
_rocm_minor="${_rocm_rest#*.}"
case "${_rocm_rest%%.*}" in "" | *[!0-9]*) return 1 ;; esac
case "$_rocm_minor" in "" | *[!0-9]*) return 1 ;; esac
;;
*[!0-9]*) return 1 ;;
esac
return 0
;;
*) return 1 ;;
esac
}
# Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2
# ("torch>=A.B[.C],<D.E.F"). Compares at major.minor granularity, which is exact
# for the windows this script uses (ceilings are always X.Y.0); a non-.0 ceiling
@ -2277,14 +2399,14 @@ _install_torch_default_index() {
esac
if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$_itdi_tv" "$_itdi_ta" \
--default-index "$TORCH_INDEX_URL" "$@"; then
substep "[WARN] $_PREV_TORCH_PIN is not installable from $TORCH_INDEX_URL -- installing the newest supported release instead" "$C_WARN"
substep "[WARN] $_PREV_TORCH_PIN is not installable from $(_strip_index_url_credentials "$TORCH_INDEX_URL") -- installing the newest supported release instead" "$C_WARN"
TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT"
_PREV_TORCH_PIN=""
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
--default-index "$TORCH_INDEX_URL" "$@"
fi
else
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
--default-index "$TORCH_INDEX_URL" "$@"
fi
}
@ -2292,13 +2414,21 @@ _install_torch_default_index() {
# 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##*/}"
_leaf=$(_torch_index_url_leaf "$1")
case "$_leaf" in
cu[0-9]*) echo "$_leaf" ;;
cpu) echo "cpu" ;;
rocm*|gfx*) echo "rocm" ;;
*) echo "" ;;
cu[0-9]*)
# Exact cu + digits only; a cu*-suffixed leaf (cu128-private) -> "" (custom),
# else a correct +cu128 wheel is force-reinstalled every run.
case "${_leaf#cu}" in
*[!0-9]*) echo "" ;;
*) echo "$_leaf" ;;
esac
;;
cpu) echo "cpu" ;;
# Exact rocm/gfx families only; a custom rocm*-suffixed leaf -> "" (custom).
*)
if _is_pip_rocm_family_leaf "$_leaf"; then echo "rocm"; else echo ""; fi
;;
esac
}
@ -2308,14 +2438,42 @@ _expected_torch_flavor_tag() {
# 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##*/}"
_leaf=$(_torch_index_url_leaf "$1")
case "$_leaf" in
cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;;
*) echo "no" ;;
cu[0-9]*) echo "yes" ;;
# Only EXACT rocm/gfx families resolve via --default-index; a suffixed leaf is verbatim.
*)
if _is_pip_rocm_family_leaf "$_leaf"; then echo "yes"; else echo "no"; fi
;;
esac
}
# Remove credentials from a wheel index URL ($1) so an authenticated pin never leaks:
# drops userinfo AND query/fragment; scheme/host/path stay exact. Shared with py / ps1.
_strip_index_url_credentials() {
_sic_url="$1"
case "$_sic_url" in
*://*) ;;
*) printf '%s' "$_sic_url"; return ;;
esac
_sic_scheme="${_sic_url%%://*}"
_sic_rest="${_sic_url#*://}"
# Drop query / fragment (may hold auth tokens).
_sic_rest="${_sic_rest%%\?*}"
_sic_rest="${_sic_rest%%#*}"
_sic_auth="${_sic_rest%%/*}"
# Drop user:pass@ userinfo if present.
case "$_sic_auth" in
*@*) _sic_host="${_sic_auth##*@}" ;;
*) _sic_host="$_sic_auth" ;;
esac
if [ "$_sic_auth" = "$_sic_rest" ]; then
printf '%s://%s' "$_sic_scheme" "$_sic_host"
else
printf '%s://%s/%s' "$_sic_scheme" "$_sic_host" "${_sic_rest#*/}"
fi
}
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/,
@ -2561,7 +2719,19 @@ _maybe_bootstrap_rocm_wsl() {
[ -n "$_rw_tmp" ] && rm -f "$_rw_tmp"
return 0
}
_maybe_bootstrap_rocm_wsl || true
# When the caller pins the wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY), honour it
# everywhere: skip the WSL ROCm bootstrap and the Radeon/Strix reroute below (which would
# re-probe the GPU and overwrite the pin). Trim whitespace first (parity with
# get_torch_index_url): a whitespace-only override is unset there, so must not flip this true.
_torch_index_pinned=false
_ti_url_trim="${UNSLOTH_TORCH_INDEX_URL:-}"
_ti_url_trim="${_ti_url_trim#"${_ti_url_trim%%[![:space:]]*}"}"; _ti_url_trim="${_ti_url_trim%"${_ti_url_trim##*[![:space:]]}"}"
_ti_family_trim="${UNSLOTH_TORCH_INDEX_FAMILY:-}"
_ti_family_trim="${_ti_family_trim#"${_ti_family_trim%%[![:space:]]*}"}"; _ti_family_trim="${_ti_family_trim%"${_ti_family_trim##*[![:space:]]}"}"
if [ -n "$_ti_url_trim" ] || [ -n "$_ti_family_trim" ]; then
_torch_index_pinned=true
fi
[ "$_torch_index_pinned" = true ] || _maybe_bootstrap_rocm_wsl || true
TORCH_INDEX_URL=$(get_torch_index_url)
@ -2572,29 +2742,74 @@ TORCH_INDEX_URL=$(get_torch_index_url)
# whose base path happens to contain "rocm" or "gfx" must not mislabel a
# cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix
# overrides in gfxNNNN/, so the trailing slash is stripped first).
_torch_index_leaf="${TORCH_INDEX_URL%/}"
# Lowercase the leaf so every gfx*/rocm*/cu* arm matches regardless of case (canonical AMD
# RDNA4 leaf is gfx120X-all). CUDA is branded only on a real cu[0-9]* leaf, so a mirror
# leaf (/current) does NOT commit a CUDA backend; an unknown leaf leaves the var unset so
# the stack probes the GPU. Query/fragment dropped first, then ALL trailing slashes (in
# lockstep with the shared _torch_index_url_leaf extractor).
_torch_index_leaf="${TORCH_INDEX_URL%%\?*}"
_torch_index_leaf="${_torch_index_leaf%%#*}"
# Strip ALL trailing slashes, not one: .../cu128// must yield cu128, not an empty leaf.
while [ -n "$_torch_index_leaf" ] && [ "${_torch_index_leaf%/}" != "$_torch_index_leaf" ]; do
_torch_index_leaf="${_torch_index_leaf%/}"
done
_torch_index_leaf="${_torch_index_leaf##*/}"
_torch_index_leaf=$(printf '%s' "$_torch_index_leaf" | tr '[:upper:]' '[:lower:]')
case "$_torch_index_leaf" in
rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;;
cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;;
*) export UNSLOTH_TORCH_BACKEND="cuda" ;;
cu[0-9]*) export UNSLOTH_TORCH_BACKEND="cuda" ;;
# Unknown leaf (odd mirror, /current): unset so a stale inherited value can't leak and
# the stack probes the GPU.
*) unset UNSLOTH_TORCH_BACKEND ;;
esac
# rocm7.2 and the CUDA cu12x/cu13x indexes now ship torch 2.11.x, so widen the
# ceiling to <2.12.0 (matches the base image and _CUDA_TORCH_PKG_SPEC in
# studio/install_python_stack.py). Keep the >=2.4 floor so an older CUDA index
# (e.g. cu118) still resolves. Match on _torch_index_leaf, not the full URL, so
# a mirror whose base path contains cu*/rocm7.2 but resolves to a cpu/older-rocm
# leaf keeps the default <2.11.0.
# Whether TORCH_INDEX_URL names an actual pip ROCm family (rocm<digit>* / gfx*), gating the
# ROCm-only side effects below (AMD bitsandbytes, ROCm-torch repair). Digit-gated so a leaf
# merely STARTING with "rocm" isn't force-repaired from the wrong path.
if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then
_torch_index_is_rocm_family=true
else
_torch_index_is_rocm_family=false
fi
# rocm7.2 and the per-gfx indexes with the _grouped_mm <2.11 bug (gfx120X-all, gfx1151,
# gfx1150) ship torch 2.11.0 -- raise the floor (also covers a pinned override that skipped
# the Strix reroute). Pin the companions too: the per-gfx index publishes them independently
# and a bare name can resolve a 2.12 ABI-mismatched wheel. Match on the FINAL leaf so a
# custom mirror with a gfx/rocm7.2 path segment but a cu*/cpu family isn't forced.
case "$_torch_index_leaf" in
rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
rocm7.2|gfx120x-all|gfx1151|gfx1150)
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
;;
# CUDA cu12x/cu13x indexes ship torch 2.11.x: widen the ceiling to <2.12.0 (matches
# _CUDA_TORCH_PKG_SPEC) and widen the companions with it so the trio stays paired.
cu[0-9]*)
TORCH_CONSTRAINT="torch>=2.4,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0"
;;
esac
# A pinned custom/unknown-leaf index (/simple, /current, /cu128-private) has no curated
# companion set, so bound torchvision/torchaudio to the same <2.11 range the Python path pins
# (else a mirror with newer companions resolves a 2.12 ABI-mismatched wheel). Known families
# keep their curated companions above (_expected_torch_flavor_tag returns "" only for custom).
if [ "$_torch_index_pinned" = true ] && \
[ -z "$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")" ]; then
TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"
fi
# Auto-detect GPU for AMD ROCm based
# get_torch_index_url must have chosen */rocm*
# (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon".
# Skipped when the index is pinned: an explicit override must not be rerouted to the
# Radeon/Strix repos by GPU probing.
_amd_gpu_radeon=false
if [ "$_torch_index_pinned" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \
@ -2671,10 +2886,14 @@ case "$TORCH_INDEX_URL" in
done
TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/"
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
# Pin companions to 2.11 (per-gfx index publishes them independently).
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
_amd_gpu_radeon=false
fi
;;
esac
fi # _torch_index_pinned guard (Radeon + Strix reroute)
# Re-run over an existing install: keep the previous venv's torch RELEASE; the fresh
# index above supplies the right flavor for this machine. Evaluated HERE, after every
# index/constraint decision including the Strix reroute, so the window checked is the
@ -2821,7 +3040,7 @@ case "$TORCH_INDEX_URL" in
if [ "$_amd_gpu_radeon" = true ]; then
substep "wheels: repo.radeon.com (Radeon)"
else
substep "wheels: $TORCH_INDEX_URL"
substep "wheels: $(_strip_index_url_credentials "$TORCH_INDEX_URL")"
fi
;;
esac
@ -2867,8 +3086,8 @@ for _p in ('torch', 'torchvision', 'torchaudio'):
}
if [ "$_MIGRATED" = true ]; then
# Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
# in the new venv location, while preserving existing torch/CUDA
# Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
# existing torch/CUDA unless the ROCm repair below fires.
substep "upgrading unsloth in migrated environment..."
if [ "$SKIP_TORCH" = true ]; then
# No-torch: install unsloth + unsloth-zoo with --no-deps (current
@ -2909,18 +3128,14 @@ if [ "$_MIGRATED" = true ]; then
# AMD ROCm: install bitsandbytes even in migrated environments so
# existing ROCm installs gain the AMD bitsandbytes build without a
# fresh reinstall.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/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)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
_install_torch_default_index --force-reinstall
fi
;;
esac
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
_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)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
_install_torch_default_index --force-reinstall
fi
fi
elif [ -n "$TORCH_INDEX_URL" ]; then
# Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac)
@ -3074,7 +3289,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \
[ "$_radeon_versions_match" != true ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN"
_install_torch_default_index
else
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
@ -3095,7 +3310,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
fi
else
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN"
_install_torch_default_index
fi
else
@ -3103,19 +3318,15 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_install_torch_default_index
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
substep "installing PyTorch ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))..."
_install_torch_default_index
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
# host stays in GGUF-only mode rather than pulling in bitsandbytes,
# which is only useful once torch is present for training.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*|*/gfx*)
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
;;
esac
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
fi
# Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed
tauri_log "STEP" "Installing Unsloth"
@ -3161,16 +3372,12 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_UNSLOTH_TORCH_OVERRIDES=""
# AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/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)..."
_install_torch_default_index --force-reinstall
fi
;;
esac
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
_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)..."
_install_torch_default_index --force-reinstall
fi
fi
else
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
@ -3217,7 +3424,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; 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 --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" \"$TORCHVISION_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --default-index $(_strip_index_url_credentials "$TORCH_INDEX_URL") --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
fi
fi
fi

File diff suppressed because it is too large Load diff

View file

@ -431,6 +431,167 @@ function Get-PytorchCudaTag {
return "cu126"
}
# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL
# TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared.
function Trim-IndexPathSlashes {
param([string]$Url)
$value = $Url.Trim()
$idx = $value.IndexOfAny([char[]]@('?', '#'))
if ($idx -lt 0) {
return $value.TrimEnd('/')
}
return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx)
}
# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY), shared by the stale-venv check
# and install selection so a pinned index wins over GPU probing (parity with the other
# installers). URL is verbatim; _FAMILY is the leaf joined to the mirror base.
function Get-PinnedTorchIndexUrl {
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL)
}
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
$base = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
return "$base/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
}
return $null
}
# Last path segment of a wheel index URL, query/fragment dropped first so a token-authenticated
# pin (.../cu128?token=x) classifies as cu128 (else it reinstalls every update). Classification
# only. Shared with the py / install.sh leaf extractors.
function Get-TorchIndexLeaf {
param([string]$Url)
if ([string]::IsNullOrWhiteSpace($Url)) { return $null }
$path = ($Url -split '[?#]', 2)[0]
if ([string]::IsNullOrWhiteSpace($path)) { return $null }
return ($path.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
}
# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer
# output before printing on failure; uv/pip errors echo the failing --index-url verbatim.
# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted.
function Redact-InstallOutput {
param([string]$Text)
if (-not $Text) { return $Text }
$Text = $Text -replace '(https?://)[^/@\s`]+@', '$1<redacted>@'
$Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=<redacted>'
# A #token=... fragment is as sensitive as a query; URL-anchored.
return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#<redacted>'
}
# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug). MUST match
# the install-spec path below and the other installers; other leaves ship <2.11 and stay default.
function Test-RocmGfx211Leaf {
param([string]$Leaf)
return @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $Leaf
}
# rocmX.Y versions KNOWN to ship torch 2.11: rocm7.2 only today. Do NOT floor an unknown newer
# rocm speculatively. MUST match _ROCM_KNOWN_TORCH211_VERSIONS and the rocm7.2 leaf elsewhere.
function Test-RocmKnown211Version {
param([int]$Major, [int]$Minor)
return ($Major -eq 7 -and $Minor -eq 2)
}
# True only for a real CUDA family leaf: "cu" + digits (cu118, cu128, ...). A bare -like 'cu*'
# would match "custom"/"current" and rebuild the venv every run. Mirrors _is_cuda_family_leaf.
function Test-CudaFamilyLeaf {
param([string]$Leaf)
if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false }
# EXACT cu+digits: cu128-private routes through the unknown-leaf path instead.
return $Leaf -match '^cu[0-9]+$'
}
# True only for a real pip ROCm family leaf: EXACT rocm<digits>[.<digits>] or a gfx leaf. A leaf
# that merely STARTS with rocm (rocm-rel-7.2.1, rocm7.2-private) is a custom pin the verbatim
# path owns, so anchor the match. Mirrors _is_pip_rocm_family_leaf / install.sh.
function Test-PipRocmFamilyLeaf {
param([string]$Leaf)
if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false }
# gfx must be followed by a digit (an architecture leaf); gfx-private is custom.
return ($Leaf -match '^gfx[0-9]') -or ($Leaf -match '^rocm[0-9]+(\.[0-9]+)?$')
}
# Stale-venv ROCm comparison for a pinned gfx*/rocm* index. Returns @{ Expected; Installed } so
# the caller rebuilds when they differ. Mirrors _rocm_pin_family_mismatch (same rocmX.Y / gfx
# cases). An untagged (no +rocm) wheel never satisfies a ROCm pin -> stale.
function Get-RocmPinStaleTags {
param([string]$PinLeaf, [string]$TorchVersion)
$_pinRocm = [regex]::Match($PinLeaf, '^rocm(\d+)\.(\d+)')
$_pinVer = if ($_pinRocm.Success) { "$($_pinRocm.Groups[1].Value).$($_pinRocm.Groups[2].Value)" } else { $null }
# The family classifier accepts a major-only rocm<d> leaf too (rocm7).
$_pinMajorOnly = [regex]::Match($PinLeaf, '^rocm(\d+)$')
# Installed rocm version and whether the wheel is a per-arch (three-part) build.
$_instRocm = [regex]::Match($TorchVersion, '\+rocm(\d+)\.(\d+)')
$_instVer = if ($_instRocm.Success) { "$($_instRocm.Groups[1].Value).$($_instRocm.Groups[2].Value)" } else { $null }
$_instPerArch = [regex]::IsMatch($TorchVersion, '\+rocm\d+\.\d+\.\d+')
# A ROCm build MUST carry a +rocm tag; an untagged wheel can't satisfy any ROCm pin.
$_instHasRocm = [regex]::IsMatch($TorchVersion, '\+rocm')
$_instRel = [regex]::Match($TorchVersion, '^(\d+)\.(\d+)')
$_instIs211 = $false
if ($_instRel.Success) {
$_instIs211 = ([int]$_instRel.Groups[1].Value -gt 2) -or ([int]$_instRel.Groups[1].Value -eq 2 -and [int]$_instRel.Groups[2].Value -ge 11)
}
if ($PinLeaf -like 'gfx*') {
if (Test-RocmGfx211Leaf $PinLeaf) {
# Expect the AMD per-arch (three-part) 2.11 wheel: satisfied only when BOTH
# a 2.11 release AND a three-part rocm tag are installed.
$installed = if ($_instIs211 -and $_instPerArch) { "rocm-perarch(torch>=2.11)" } else { "rocm-generic-or-old" }
return @{ Expected = "rocm-perarch(torch>=2.11)"; Installed = $installed }
}
# Non-2.11 gfx leaf (<2.11 spec): stale on an untagged wheel or a 2.11+ build.
$installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
return @{
Expected = "rocm(torch<2.11)"
Installed = $installed
}
}
# Major-only rocm pin (rocm7): compare majors only -- a +rocm6.4 wheel under a rocm7
# pin is stale, any +rocm7.x wheel satisfies it (no pinned minor to compare, and the
# 2.11-line fallback below would invert both verdicts). Mirrors _rocm_pin_family_mismatch.
if ($_pinMajorOnly.Success) {
$_pinMaj = [int]$_pinMajorOnly.Groups[1].Value
if ($_instVer) {
$_instMaj = [int]$_instRocm.Groups[1].Value
$expected = if ($_instMaj -eq $_pinMaj) { "rocm$_instVer" } else { "rocm$_pinMaj.x" }
return @{ Expected = $expected; Installed = "rocm$_instVer" }
}
# Untagged wheel never satisfies a ROCm pin; a +rocm tag with an unreadable
# version is accepted (matches the lenient unreadable fallback below).
$installed = if ($_instHasRocm) { "rocm" } else { "not-rocm" }
return @{ Expected = "rocm"; Installed = $installed }
}
# rocmX.Y pin.
if ($_pinVer -and $_instVer) {
# Both readable: exact compare. When they match AND the pin is KNOWN-2.11, the
# installed release must also be 2.11 (a +rocm7.2 wheel drifted to 2.12 shares the
# tag but violates the spec), so fold the release into the tag. Mirrors _rocm_pin_family_mismatch.
$_pinKnown211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value)
$_instOn211 = $_instRel.Success -and [int]$_instRel.Groups[1].Value -eq 2 -and [int]$_instRel.Groups[2].Value -eq 11
if ($_pinKnown211 -and -not $_instOn211) {
return @{ Expected = "rocm$_pinVer(torch2.11)"; Installed = "rocm$_instVer(torch-off-2.11)" }
}
return @{ Expected = "rocm$_pinVer"; Installed = "rocm$_instVer" }
}
$_pinNeeds211 = $false
if ($_pinRocm.Success) {
# Only KNOWN-2.11 rocm (rocm7.2) is on the 2.11 line (no speculative floor).
# Matches _ROCM_KNOWN_TORCH211_VERSIONS.
$_pinNeeds211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value)
}
# Fallback (installed rocm version unreadable): compare on the 2.11 line; an untagged
# wheel never satisfies a rocmX.Y pin -> stale.
$installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
return @{
Expected = if ($_pinNeeds211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
Installed = $installed
}
}
# VS generator -> MSBuild BuildCustomizations dir; toolset tracks the VS major
# (18->v180, 17->v170), defaulting to v170 when unparseable.
function Get-VcBuildCustomizationsDir {
@ -813,11 +974,14 @@ function Invoke-SetupCommand {
# Merge stderr into stdout so progress/warning output stays visible
# without flipping $? on successful native commands (PS 5.1 treats
# stderr records as errors that set $? = $false even on exit code 0).
& $Command 2>&1 | Out-Host
# Redact per record: uv/pip echo index URLs (credentials and all) in
# their errors, and verbose mode must not bypass the quiet path's
# redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched.
& $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
} else {
$output = & $Command 2>&1 | Out-String
if ($LASTEXITCODE -ne 0) {
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
}
}
return [int]$LASTEXITCODE
@ -2535,6 +2699,8 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
$VenvPyExe = Join-Path $VenvDir "Scripts\python.exe"
$installedTorchTag = $null
$shouldRebuild = $false
# Set when a stale venv under a pin is repaired in place (force-reinstall) not wiped.
$script:PinChangedForceReinstall = $false
if (Test-Path -LiteralPath $VenvPyExe) {
try {
@ -2551,10 +2717,14 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
if ($finished -and $proc.ExitCode -eq 0 -and $torchVer) {
if ($torchVer -match '\+(cu\d+)') {
$installedTorchTag = $Matches[1]
} elseif ($torchVer -match '\+rocm') {
# Any +rocm / gfx wheel -> generic "rocm" flavor (the exact version is
# repaired later by install_python_stack.py; here we only need the flavor).
$installedTorchTag = "rocm"
} elseif ($torchVer -match '\+cpu') {
$installedTorchTag = "cpu"
} else {
# Untagged wheel (plain "2.x.y" from PyPI) -- treat as cpu
# Untagged wheel (plain "2.x.y" from PyPI) -> cpu.
$installedTorchTag = "cpu"
}
} else {
@ -2570,12 +2740,71 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
}
if (-not $shouldRebuild) {
$expectedTorchTag = if ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" }
if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
$_pinnedIdx = Get-PinnedTorchIndexUrl
$_expectedKnown = $true
if ($_pinnedIdx) {
$_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx
# Digit-gated like the install selection: a custom rocm-* leaf (rocm-current /
# rocm-rel-7.2.1) is NOT a ROCm family and must not be stale-compared.
if (Test-PipRocmFamilyLeaf $_pinLeaf) {
# Don't collapse a pinned ROCm/gfx leaf to a generic "rocm" (would mask a family
# change, rocm6.4 -> gfx1151). Get-RocmPinStaleTags uses the SAME 2.11 allowlist
# as the install path, so a gfx110X-all/gfx90a/gfx908 pin on a <2.11 wheel is NOT stale.
$_rocmTags = Get-RocmPinStaleTags -PinLeaf $_pinLeaf -TorchVersion $torchVer
$expectedTorchTag = $_rocmTags.Expected
$installedTorchTag = $_rocmTags.Installed
} elseif ((Test-CudaFamilyLeaf $_pinLeaf) -or $_pinLeaf -eq 'cpu') {
# cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds;
# /custom and /current fall through to the unknown-index branch below.
$expectedTorchTag = $_pinLeaf
} else {
# Custom index whose leaf is not a torch flavor (a /simple mirror): the
# flavor can't be inferred, so never treat the venv as stale over it.
$_expectedKnown = $false
$expectedTorchTag = $installedTorchTag
}
} elseif ($HasNvidiaSmi) {
$expectedTorchTag = Get-PytorchCudaTag
} elseif ($HasROCm -or $script:ROCmGfxArch) {
# AMD/ROCm host with no explicit pin: an existing +rocm wheel is correct (gfx arch
# counts even when $HasROCm is false). But only the arches the install path maps to a
# repo.amd.com index get ROCm torch; an unmapped arch installs CPU, so expect "cpu"
# for those or a correct CPU venv rebuilds every update.
$_rocmWheelArches = @(
"gfx1201", "gfx1200", # RDNA 4
"gfx1151", "gfx1150", # RDNA 3.5 (Strix Halo/Point)
"gfx1103", "gfx1102", "gfx1101", "gfx1100", # RDNA 3
"gfx90a", "gfx908" # MI200 / MI100
)
if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) {
# A correct +rocm wheel is not stale. A CPU wheel on a supported AMD arch is
# NOT wiped either (the AMD Windows ROCm override below upgrades it in place);
# expect "cpu" for that case. A wrong CUDA wheel still rebuilds.
if ($installedTorchTag -eq "cpu") {
$expectedTorchTag = "cpu"
} else {
$expectedTorchTag = "rocm"
}
} else {
$expectedTorchTag = "cpu"
}
} else {
$expectedTorchTag = "cpu"
}
if ($_expectedKnown -and $installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
$shouldRebuild = $true
}
}
# A stale venv under a pin whose torch still imports is repaired IN PLACE (the dependency
# pass force-reinstalls from the pin). The rebuild path wipes the venv and would strand a
# direct `studio update`; only a broken venv or an unpinned drift wipes.
if ($shouldRebuild -and $_pinnedIdx -and $installedTorchTag) {
substep "Torch-index pin changed ($installedTorchTag) -- reinstalling torch from the pin in place." "Cyan"
$script:PinChangedForceReinstall = $true
$shouldRebuild = $false
}
if ($shouldRebuild) {
$reason = if ($installedTorchTag) { "torch $installedTorchTag != required $expectedTorchTag" } else { "torch could not be imported" }
if ($InstallerManagedSetup) {
@ -2653,23 +2882,41 @@ if (Get-Command uv -ErrorAction SilentlyContinue) {
# Helper: install a package, preferring uv with pip fallback
function Fast-Install {
param([Parameter(ValueFromRemainingArguments=$true)]$Args_)
if ($UseUv) {
$VenvPy = (Get-Command python).Source
# An explicit --index-url must win. Inherited uv index env vars otherwise
# override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop
# them only for index-pinned installs; mirrors still apply elsewhere.
$saved = @{}
if (@($Args_) -contains '--index-url') {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
$saved[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
# An explicit --index-url must win: inherited uv index vars otherwise pull CPU torch over
# the CUDA/ROCm build (#6898), so drop them for pinned installs (scrub covers the whole
# function since the pip fallback honours PIP_* too). UV_TORCH_BACKEND / UV_FIND_LINKS also
# reroute; UV_NO_CONFIG=1 (+ dropping UV_CONFIG_FILE) stops a uv.toml index outranking the
# pin (uv 0.10); PIP_NO_INDEX / PIP_INDEX_URL would defeat the pinned --index-url in pip.
$saved = @{}
$pinned = @($Args_) -contains '--index-url'
if ($pinned) {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL',
'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'PIP_EXTRA_INDEX_URL', 'PIP_FIND_LINKS',
'PIP_NO_INDEX', 'PIP_INDEX_URL',
'UV_CONFIG_FILE', 'UV_NO_CONFIG', 'PIP_CONFIG_FILE') {
$saved[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 }
finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } }
if ($LASTEXITCODE -eq 0) { return }
$env:UV_NO_CONFIG = '1'
# A `pip config` global.extra-index-url still adds indexes to the pip FALLBACK;
# PIP_CONFIG_FILE = 'nul' (Windows devnull) loads NO config (uv ignores pip config).
$env:PIP_CONFIG_FILE = 'nul'
}
try {
if ($UseUv) {
$VenvPy = (Get-Command python).Source
$result = & uv pip install --python $VenvPy @Args_ 2>&1
if ($LASTEXITCODE -eq 0) { return }
}
& python -m pip install @Args_ 2>&1
}
finally {
if ($pinned) {
Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue
Remove-Item "Env:PIP_CONFIG_FILE" -ErrorAction SilentlyContinue
}
foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } }
}
& python -m pip install @Args_ 2>&1
}
# ── Check if Python deps need updating ──
@ -2752,6 +2999,10 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1)
# pip install unsloth 2>&1 | Out-Null
# }
# A torch-index pin change repairs in place: force the dependency pass so the torch install
# below force-reinstalls from the new pin (else the fast path keeps the old wheel).
if ($script:PinChangedForceReinstall) { $SkipPythonDeps = $false }
if (-not $SkipPythonDeps) {
if ($script:UnslothVerbose) {
@ -2779,7 +3030,13 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir
[Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User')
substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)"
if ($HasNvidiaSmi) {
# Explicit pin (URL or family) wins over GPU probing and suppresses the AMD reroute below;
# matches install.sh / install.ps1 / install_python_stack.py.
$PinnedTorchIndexUrl = Get-PinnedTorchIndexUrl
$TorchIndexPinned = [bool]$PinnedTorchIndexUrl
if ($PinnedTorchIndexUrl) {
$CuTag = Get-TorchIndexLeaf $PinnedTorchIndexUrl
} elseif ($HasNvidiaSmi) {
$CuTag = Get-PytorchCudaTag
} else {
$CuTag = "cpu"
@ -2800,7 +3057,7 @@ $ROCmIndexUrl = $null
# SDK -- which flips Unsloth out of chat-only (CHAT_ONLY) and enables Train/Export.
# Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed
# ROCm install still falls back to CPU below, so this is safe.
if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
@ -2850,8 +3107,45 @@ if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
}
}
# A pinned gfx*/rocm index skips the auto-reroute above; route it through the ROCm install path
# with the same floor/companions the unpinned AMD path uses (mirrors install.ps1), else the CUDA
# branch installs bare torch and resolves a known-bad wheel for gfx115x/gfx120x/rocm>=7.2.
if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) {
$_pinLeaf = Get-TorchIndexLeaf $PinnedTorchIndexUrl
$_pinRocm211 = $false
# Anchor the match ($) so a suffixed custom leaf (rocm7.2-private) falls through to the
# verbatim install instead of being floored by its rocm7.2 prefix.
if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') {
# Only KNOWN-2.11 rocm (rocm7.2) gets the floor (no speculative floor). Matches
# Test-RocmKnown211Version / _ROCM_KNOWN_TORCH211_VERSIONS.
$_pinRocm211 = Test-RocmKnown211Version -Major ([int]$Matches[1]) -Minor ([int]$Matches[2])
}
# Only the 2.11 gfx arches need the floor; others publish <2.11 and stay bare. Reuse
# Test-RocmGfx211Leaf so this allowlist and the stale-venv check never diverge.
$_pinGfx211 = Test-RocmGfx211Leaf $_pinLeaf
if ($_pinGfx211 -or $_pinRocm211) {
$ROCmIndexUrl = $PinnedTorchIndexUrl
$ROCmTorchSpec = "torch>=2.11.0,<2.12.0"
$ROCmVisionSpec = "torchvision>=0.26.0,<0.27.0"
$ROCmAudioSpec = "torchaudio>=2.11.0,<2.12.0"
substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchSpec" "Cyan"
} elseif (Test-PipRocmFamilyLeaf $_pinLeaf) {
# Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with
# bare specs. Only EXACT rocm<digits> and gfx* are --index-url families; a suffixed
# leaf stays on the verbatim path. Mirrors install.ps1 / _is_pip_rocm_family_leaf.
$ROCmIndexUrl = $PinnedTorchIndexUrl
$ROCmTorchSpec = "torch"
$ROCmVisionSpec = "torchvision"
$ROCmAudioSpec = "torchaudio"
}
}
$PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
# A full URL pin is used verbatim; a family pin already set $CuTag. A pinned ROCm install
# goes through $ROCmIndexUrl; on failure the fallback uses the CPU index, not the ROCm pin.
$TorchInstallIndexUrl = if ($ROCmIndexUrl) { "$PyTorchWhlBase/cpu" } elseif ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" }
$ROCmCpuFallback = $false
if ($ROCmIndexUrl) {
substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..."
@ -2859,7 +3153,7 @@ if ($ROCmIndexUrl) {
substep " enforcing $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec (known _grouped_mm bug in older wheels)" "Cyan"
}
if ($script:UnslothVerbose) {
Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl
Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
@ -2868,7 +3162,7 @@ if ($ROCmIndexUrl) {
}
if ($torchInstallExit -ne 0) {
Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow
Write-Host $output -ForegroundColor Yellow
Write-Host (Redact-InstallOutput $output) -ForegroundColor Yellow
$ROCmIndexUrl = $null
$ROCmCpuFallback = $true
} else {
@ -2878,42 +3172,70 @@ if ($ROCmIndexUrl) {
}
}
if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) {
substep "installing PyTorch (CPU-only)..."
# After an AMD ROCm fallback, force-reinstall so a partially-installed ROCm torch
# (which still satisfies the CPU torch>= range) is replaced by the CPU build. Skip
# the forced reinstall on a genuine CPU-only host so the common path stays fast.
# Build the array directly: an if-expression collapses @("x") to a scalar string,
# which @splat would then enumerate char-by-char into broken single-letter args.
# After an AMD ROCm fallback, force-reinstall so a partial ROCm torch (which satisfies the
# CPU torch>= range) is replaced by the CPU build; skip on a genuine CPU host to stay fast.
# $ROCmCpuFallback matters when a PINNED ROCm index failed ($CuTag is still the rocm leaf).
# Build the array directly: an if-expression collapses @("x") to a scalar @splat would
# enumerate char-by-char.
$cpuForce = @()
if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") }
# --force-reinstall on a pin change: a stale +cu / +rocm wheel still satisfies the CPU
# torch>= range, so uv would keep it and only swap companions.
if ($script:PinChangedForceReinstall) { $cpuForce = @("--force-reinstall") }
# A PINNED cpu index installs the bounded trio (parity with _CPU_TORCH_PKG_SPEC): the /cpu
# index serves newer torch, and _ensure_cpu_torch keeps any CPU build, so a bare trio could
# land an unsupported version. Unpinned CPU hosts keep the bare trio (pre-pin behavior).
$cpuTorchSpec = "torch"; $cpuVisionSpec = "torchvision"; $cpuAudioSpec = "torchaudio"
if ($TorchIndexPinned) {
$cpuTorchSpec = "torch>=2.4,<2.12.0"
$cpuVisionSpec = "torchvision>=0.19,<0.27.0"
$cpuAudioSpec = "torchaudio>=2.4,<2.12.0"
}
if ($script:UnslothVerbose) {
Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu"
Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
$output = Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" | Out-String
$output = Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | Out-String
$torchInstallExit = $LASTEXITCODE
}
if ($torchInstallExit -ne 0) {
Write-Host "[FAILED] PyTorch install failed (exit code $torchInstallExit)" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
exit 1
}
} elseif (-not $ROCmIndexUrl) {
substep "installing PyTorch with CUDA support ($CuTag)..."
substep "(This download is ~2.8 GB -- may take a few minutes)"
# --force-reinstall on a pin change: an installed cuXXX wheel satisfies the bare torch
# requirement (PEP 440 ignores the +cuXXX tag), so without it a changed CUDA pin (cu126
# -> cu128) never applies.
$cudaForce = @()
if ($script:PinChangedForceReinstall) { $cudaForce = @("--force-reinstall") }
# An unknown-leaf custom pin (/simple, /current) routes here with $CuTag as that leaf. Bound
# the trio like the fresh custom-pin paths so a mirror can't pull an ABI-newer companion
# against the capped torch. Known cu* leaves keep bare specs.
$cudaTorchSpec = "torch"
$cudaVisionSpec = "torchvision"
$cudaAudioSpec = "torchaudio"
if ($TorchIndexPinned -and -not (Test-CudaFamilyLeaf $CuTag)) {
$cudaTorchSpec = "torch>=2.4,<2.11.0"
$cudaVisionSpec = "torchvision>=0.19,<0.26.0"
$cudaAudioSpec = "torchaudio>=2.4,<2.11.0"
}
if ($script:UnslothVerbose) {
Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag"
Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
$output = Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" | Out-String
$output = Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | Out-String
$torchInstallExit = $LASTEXITCODE
}
if ($torchInstallExit -ne 0) {
Write-Host "[FAILED] PyTorch CUDA install failed (exit code $torchInstallExit)" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
exit 1
}
@ -2929,7 +3251,7 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
}
if ($tritonInstallExit -ne 0) {
substep "Triton install failed -- torch.compile may not work" "Yellow"
Write-Host $output -ForegroundColor Yellow
Write-Host (Redact-InstallOutput $output) -ForegroundColor Yellow
} else {
substep "Triton for Windows installed (enables torch.compile)"
}
@ -3026,7 +3348,7 @@ foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.8.0", "hf_xet==1.4
}
if ($t5PkgExit -ne 0) {
Write-Host "[FAIL] Could not install $pkg into .venv_t5_530/" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
$ErrorActionPreference = $prevEAP_t5
exit 1
}
@ -3061,7 +3383,7 @@ foreach ($pkg in @("transformers==5.5.0", "huggingface_hub==1.8.0", "hf_xet==1.4
}
if ($t5PkgExit -ne 0) {
Write-Host "[FAIL] Could not install $pkg into .venv_t5_550/" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
$ErrorActionPreference = $prevEAP_t5
exit 1
}
@ -3096,7 +3418,7 @@ foreach ($pkg in @("transformers==5.10.2", "huggingface_hub==1.8.0", "hf_xet==1.
}
if ($t5PkgExit -ne 0) {
Write-Host "[FAIL] Could not install $pkg into .venv_t5_510/" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
$ErrorActionPreference = $prevEAP_t5
exit 1
}

View file

@ -10,6 +10,8 @@ import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
INSTALL_SH = REPO_ROOT / "install.sh"
INSTALL_PS1 = REPO_ROOT / "install.ps1"
SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1"
STACK_PY = REPO_ROOT / "studio" / "install_python_stack.py"
class TestNoTorchBackendAutoInInstallSh:
@ -180,3 +182,607 @@ class TestUvBytecodeCompileTimeout:
assert (
'$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text
), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT"
class TestTorchIndexOverrideParity:
"""Every installer must honor UNSLOTH_TORCH_INDEX_URL / _FAMILY so a pinned wheel
index wins over GPU probing on all platforms (no asymmetric, per-OS coverage)."""
@pytest.mark.parametrize(
"path",
[INSTALL_SH, INSTALL_PS1, SETUP_PS1, STACK_PY],
ids = ["install.sh", "install.ps1", "setup.ps1", "install_python_stack.py"],
)
def test_installer_reads_override_env(self, path):
text = path.read_text(encoding = "utf-8")
for var in ("UNSLOTH_TORCH_INDEX_URL", "UNSLOTH_TORCH_INDEX_FAMILY"):
assert var in text, f"{path.name} does not honor {var}"
@pytest.mark.parametrize(
"path",
[INSTALL_PS1, SETUP_PS1],
ids = ["install.ps1", "setup.ps1"],
)
def test_amd_reroute_guarded_when_pinned(self, path):
# The AMD ROCm reroute must be skipped when the index is explicitly pinned,
# so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten.
text = path.read_text(encoding = "utf-8")
assert (
"TorchIndexPinned" in text
), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
def test_cuda_pin_overrides_cvd_hide_gate(self):
# A pinned cu* index skips ALL host-GPU probing, so the CUDA repair must clear the
# CUDA_VISIBLE_DEVICES hide gate too (else the GPU-less CI case bails).
text = STACK_PY.read_text(encoding = "utf-8")
m = re.search(r"def _ensure_cuda_torch\(\).*?(?=\ndef )", text, re.DOTALL)
assert m, "could not locate _ensure_cuda_torch"
body = m.group(0)
assert "_cuda_pinned" in body, (
"_ensure_cuda_torch should compute a CUDA-pin flag so the pin can "
"override the CVD hide gate"
)
assert re.search(
r"if not _cuda_pinned and _cvd is not None", body
), "the CVD hide gate must be bypassed when a CUDA index is pinned"
def test_cpu_repair_pins_supported_torch_range(self):
# The explicit-CPU repair must use the bounded CPU/CUDA spec, not a bare trio (the
# /cpu index serves torch 2.11+, so a bare install could resolve out of range).
text = STACK_PY.read_text(encoding = "utf-8")
m = re.search(r"def _ensure_cpu_torch\(\).*?(?=\ndef )", text, re.DOTALL)
assert m, "could not locate _ensure_cpu_torch"
body = m.group(0)
assert "_CPU_TORCH_PKG_SPEC" in body, (
"_ensure_cpu_torch should install the bounded _CPU_TORCH_PKG_SPEC, "
"not a bare torch/torchvision/torchaudio trio"
)
def test_setup_ps1_stale_check_gates_rocm_on_supported_arch(self):
# The stale check must expect ROCm torch only for arches the install path maps to a
# repo.amd.com index; expecting "rocm" for an unmapped arch marks a good CPU venv stale.
text = SETUP_PS1.read_text(encoding = "utf-8")
assert "_rocmWheelArches" in text, (
"setup.ps1 stale check should restrict the ROCm expected-tag to the "
"supported gfx wheel arches"
)
class TestGfx211AllowlistParity:
"""The gfx per-arch 2.11-floor leaves (gfx120X-all / gfx1151 / gfx1150) must be the
SAME set in every installer and its stale/mismatch check. When they diverged, a
pinned gfx110X-all / gfx90a / gfx908 wheel (<2.11) was force-reinstalled every update."""
EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150"}
def test_install_sh_allowlist(self):
text = INSTALL_SH.read_text(encoding = "utf-8").lower()
# install.sh: the TORCH_CONSTRAINT case (rocm7.2|gfx120x-all|gfx1151|gfx1150).
m = re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150", text)
assert m, "install.sh gfx-2.11 allowlist case not found / changed"
def test_install_ps1_allowlist(self):
text = INSTALL_PS1.read_text(encoding = "utf-8").lower()
m = re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text)
assert m, "install.ps1 $_pinGfx211 allowlist not found / changed"
def test_setup_ps1_defines_single_allowlist_helper(self):
# setup.ps1 must define the allowlist once (Test-RocmGfx211Leaf) and reuse it, so
# the stale check and install spec can't disagree.
text = SETUP_PS1.read_text(encoding = "utf-8")
assert (
"function Test-RocmGfx211Leaf" in text
), "setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper"
assert re.search(
r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower()
), "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist"
assert "$_pinGfx211 = Test-RocmGfx211Leaf" in text, (
"setup.ps1 install-spec path should reuse Test-RocmGfx211Leaf, not "
"re-hardcode the allowlist (they must not diverge)"
)
def test_stack_py_allowlist(self):
text = STACK_PY.read_text(encoding = "utf-8").lower()
assert (
'"gfx120x-all", "gfx1151", "gfx1150"' in text
), "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed"
class TestCudaLeafDigitParity:
"""A wheel-family leaf is CUDA only when it is "cu" + digits (cu118/cu128/...).
A bare cu* glob wrongly catches mirror leaves like /custom or /current; when
that happened the venv was marked stale and rebuilt on every run. Every
installer must require a digit after "cu" in its family/CUDA classification."""
def test_stack_py_requires_cu_digit(self):
text = STACK_PY.read_text(encoding = "utf-8")
# EXACT cu+digits: a custom leaf like cu128-private must route to the
# verbatim/unknown path, not be compared against the installed +cu128 tag.
assert re.search(
r'r"cu\[0-9\]\+"', text
), "install_python_stack.py _is_cuda_family_leaf must fullmatch cu[0-9]+"
def test_setup_ps1_requires_cu_digit(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
# EXACT cu+digits: cu128-private must not classify as CUDA (it would become
# the expected tag and rebuild the venv on every update).
assert re.search(
r"'\^cu\[0-9\]\+\$'", text
), "setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9]+$, not a cu* prefix"
# The stale-venv branch must go through the digit-guarded helper.
assert (
"Test-CudaFamilyLeaf $_pinLeaf" in text
), "setup.ps1 stale check should classify CUDA via Test-CudaFamilyLeaf"
def test_install_ps1_requires_cu_digit_in_gpu_branch(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
assert re.search(
r"'\^cu\[0-9\]'", text
), "install.ps1 Get-TauriGpuBranch must require a digit after cu"
def test_install_sh_requires_cu_digit_in_gpu_branch(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The _tauri_gpu_branch cuda case must be cu[0-9]*, not a bare cu*.
assert re.search(
r"cu\[0-9\]\*\)\s*echo \"cuda\"", text
), "install.sh _tauri_gpu_branch cuda case must be cu[0-9]*, not cu*"
def test_install_sh_backend_export_requires_cu_digit(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# Brand CUDA only on cu[0-9]*; a bare catch-all *) -> cuda would mis-brand
# /current, /custom pins and skip ROCm repair on AMD hosts.
assert re.search(
r'cu\[0-9\]\*\)\s*export UNSLOTH_TORCH_BACKEND="cuda"', text
), "install.sh backend export must brand cuda only on cu[0-9]*"
# An unknown leaf must NOT commit a cuda backend (it unsets instead).
assert re.search(
r"\*\)\s*unset UNSLOTH_TORCH_BACKEND", text
), "install.sh backend export must unset (not force cuda) on an unknown leaf"
def test_install_sh_lowercases_backend_leaf(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The leaf feeding both the backend case and the 2.11 floor case must be
# lowercased so the canonical gfx120X-all (capital X) matches.
assert re.search(
r"_torch_index_leaf=\$\(printf '%s' \"\$_torch_index_leaf\" \| tr '\[:upper:\]' '\[:lower:\]'\)",
text,
), "install.sh must lowercase _torch_index_leaf before the gfx/rocm/cu case matches"
class TestKnown211SetParity:
"""The KNOWN-2.11 rocm/gfx set must be identical across all four installers:
exactly {rocm7.2} plus the gfx allowlist {gfx120x-all, gfx1151, gfx1150}.
rocm7.3 / torch 2.12 do not exist, so no side may floor them speculatively."""
def test_install_sh_known_211_leaf_is_rocm72_and_gfx_allowlist(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The 2.11 floor case matches exactly rocm7.2 + the three gfx leaves.
assert re.search(
r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text
), "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150"
# No speculative rocm7.3 anywhere.
assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3"
def test_python_known_211_versions_is_only_rocm72(self):
text = STACK_PY.read_text(encoding = "utf-8")
assert "_ROCM_KNOWN_TORCH211_VERSIONS" in text
# The frozenset literal is exactly {(7, 2)}.
m = re.search(r"_ROCM_KNOWN_TORCH211_VERSIONS[^=]*=\s*frozenset\(\{([^}]*)\}\)", text)
assert m is not None, "install_python_stack.py must define _ROCM_KNOWN_TORCH211_VERSIONS"
assert "(7, 2)" in m.group(1)
assert "7, 3" not in m.group(1) and "7, 1" not in m.group(1)
def test_setup_ps1_known_211_helper_is_only_rocm72(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
assert "Test-RocmKnown211Version" in text
# The predicate is Major -eq 7 -and Minor -eq 2 (only rocm7.2).
assert re.search(
r"Test-RocmKnown211Version[\s\S]{0,400}\$Major -eq 7 -and \$Minor -eq 2", text
), "setup.ps1 Test-RocmKnown211Version must accept only rocm7.2"
def test_install_ps1_pin_floor_is_only_rocm72(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
# The pinned-ROCm install-spec floor must be Major -eq 7 -and Minor -eq 2,
# not the speculative >= 2 that would floor a non-existent rocm7.3.
assert re.search(
r"\$_pinRocm211 = \(\[int\]\$Matches\[1\] -eq 7 -and \[int\]\$Matches\[2\] -eq 2\)",
text,
), "install.ps1 pinned-ROCm floor must be rocm7.2 only (no speculative >= 2)"
def test_ps1_pin_floor_gate_is_anchored(self):
"""The floor-selection gate that reads $_pinRocm211 from the raw leaf must anchor
the rocm match ($), or a suffixed custom leaf (rocm7.2-private) matches the rocm7.2
prefix, takes the 2.11-floor branch, and is force-routed through the ROCm path
before the exact-match elseif can send it to the verbatim install (Codex P2)."""
for path, label in ((INSTALL_PS1, "install.ps1"), (SETUP_PS1, "setup.ps1")):
text = path.read_text(encoding = "utf-8")
assert "-match '^rocm(\\d+)\\.(\\d+)$'" in text, (
f"{label} floor gate must anchor the rocm match (^rocm(\\d+)\\.(\\d+)$) so a "
"suffixed custom leaf is not floored/routed as rocm7.2"
)
assert (
"-match '^rocm(\\d+)\\.(\\d+)'\n" not in text
), f"{label} floor gate must not use the unanchored ^rocm(\\d+)\\.(\\d+) prefix"
def test_install_ps1_bounds_unknown_leaf_pinned_torch(self):
"""install.ps1's pinned-torch install must bound BOTH companions on EVERY
index, cu<digits> families included: torchaudio 2.11 dropped its exact torch
pin from the wheel metadata, so a bare companion beside torch<2.11 can
resolve a mismatched 2.11.0 build (Codex P2, then unconditional per the
torchaudio 2.11 unpinning)."""
text = INSTALL_PS1.read_text(encoding = "utf-8")
assert (
'$_pinVisionSpec = "torchvision>=0.19,<0.26.0"' in text
), "install.ps1 custom-pin install must bound torchvision (>=0.19,<0.26.0)"
assert (
'$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"' in text
), "install.ps1 custom-pin install must bound torchaudio (>=2.4,<2.11.0)"
# No cu-family exemption: the bounds apply unconditionally.
assert (
"$_pinCuLeaf" not in text
), "install.ps1 must bound companions on every index (no cu-family exemption)"
# The bounded companions must actually be passed to the install command.
assert re.search(
r'"torch>=2\.4,<2\.11\.0" \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl',
text,
), "install.ps1 custom-pin install must pass the bounded companion specs to uv"
def test_gfx_allowlist_matches_across_installers(self):
# The gfx 2.11 allowlist {gfx120x-all, gfx1151, gfx1150} must appear in each.
gfx = ("gfx120x-all", "gfx1151", "gfx1150")
for path, label in (
(INSTALL_SH, "install.sh"),
(INSTALL_PS1, "install.ps1"),
(SETUP_PS1, "setup.ps1"),
(STACK_PY, "install_python_stack.py"),
):
low = path.read_text(encoding = "utf-8").lower()
for g in gfx:
assert g in low, f"{label} missing gfx 2.11 allowlist member {g}"
class TestPinnedRocmLeafDigitParity:
"""A pinned index is a pip ROCm --default-index family only when its leaf is an
EXACT rocm+digits (rocm7 / rocm7.2) or gfx*. A ^rocm[0-9] PREFIX (or a bare rocm*
glob) wrongly catches a custom mirror / find-links leaf (rocm-current /
rocm-rel-7.2.1) AND a suffixed private-mirror leaf (rocm7.2-private / rocm7-current),
routing it through the ROCm install path (which silently falls back to CPU on
failure) or skipping the custom-index companion bounds, instead of the verbatim
--default-index install. All installers must match the family EXACTLY: Python and
install.sh via a shared _is_pip_rocm_family_leaf, setup.ps1 via Test-PipRocmFamilyLeaf,
install.ps1 via an anchored ^rocm[0-9]+(\\.[0-9]+)?$ reroute."""
def test_install_ps1_pinned_reroute_requires_rocm_digit(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
# The pinned gfx*/rocm reroute must match rocm EXACTLY (anchored), so a suffixed
# rocm7.2-private / rocm-current falls through to the verbatim --default-index path.
assert "-match '^rocm[0-9]+(\\.[0-9]+)?$'" in text, (
"install.ps1 pinned-index reroute must anchor the rocm match "
"(^rocm[0-9]+(\\.[0-9]+)?$), not a bare -like 'rocm*' or an unanchored ^rocm\\d"
)
# Neither the broad glob nor the unanchored prefix may drive that reroute.
assert (
"-like 'rocm*'" not in text
), "install.ps1 must not route a pinned index on a bare -like 'rocm*' glob"
assert (
"-match '^rocm\\d'" not in text
), "install.ps1 must not route a pinned index on an unanchored -match '^rocm\\d'"
def test_setup_ps1_pinned_reroute_requires_rocm_digit(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
# setup.ps1 routes every family decision through Test-PipRocmFamilyLeaf, which
# anchors the rocm match so a suffixed custom leaf stays on the verbatim path.
assert (
"function Test-PipRocmFamilyLeaf" in text
), "setup.ps1 must define Test-PipRocmFamilyLeaf (the exact rocm/gfx family gate)"
assert "'^rocm[0-9]+(\\.[0-9]+)?$'" in text, (
"setup.ps1 Test-PipRocmFamilyLeaf must anchor the rocm match "
"(^rocm[0-9]+(\\.[0-9]+)?$) so rocm7.2-private / rocm-current stay verbatim"
)
pinned_block = text[text.find("$_pinGfx211 = Test-RocmGfx211Leaf") :][:2000]
assert (
"-like 'rocm*'" not in pinned_block
), "setup.ps1 pinned reroute must not route on a bare -like 'rocm*' glob"
def test_install_sh_repairable_requires_rocm_digit(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# _torch_index_repairable routes rocm/gfx through the exact-match helper.
assert (
"_is_pip_rocm_family_leaf" in text
), "install.sh must define/use _is_pip_rocm_family_leaf for the exact rocm gate"
# gfx needs a following digit: gfx-private / gfxfoo are custom verbatim pins.
assert re.search(
r'case "\$1" in\n\s*gfx\[0-9\]\*\) return 0', text
), "install.sh _is_pip_rocm_family_leaf must treat only gfx<digit>* as a family"
assert not re.search(
r'case "\$1" in\n\s*gfx\*\) return 0', text
), "install.sh _is_pip_rocm_family_leaf must not family-match a bare gfx* glob"
def test_stack_py_pip_rocm_family_requires_digit(self):
text = STACK_PY.read_text(encoding = "utf-8")
assert re.search(
r'fullmatch\(r"rocm\\d\+\(\?:\\\.\\d\+\)\?", leaf\)', text
), "install_python_stack.py _is_pip_rocm_family_leaf must fullmatch rocm\\d+(?:\\.\\d+)?"
# The unanchored prefix must be gone from the family/flavor gates.
assert (
're.match(r"^rocm\\d"' not in text
), "install_python_stack.py must not gate a family on an unanchored re.match(^rocm\\d)"
def test_install_sh_rocm_side_effects_digit_gated(self):
"""The AMD bitsandbytes + 'repair ROCm torch' side effects must fire only on
an EXACT ROCm family (rocm7.2/gfx*), not a bare */rocm* whole-URL glob nor a
^rocm[0-9] prefix that catches a custom CPU/CUDA index like /rocm-current or a
suffixed /rocm7.2-private and force-repairs it from the wrong --default-index."""
text = INSTALL_SH.read_text(encoding = "utf-8")
assert (
'if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then\n _torch_index_is_rocm_family=true'
in text
), "install.sh must set _torch_index_is_rocm_family from the exact-match helper"
assert (
'[ "$_torch_index_is_rocm_family" = true ]' in text
), "install.sh ROCm bnb/repair hooks must gate on _torch_index_is_rocm_family"
assert (
"*/rocm*|*/gfx*)\n _install_bnb_rocm" not in text
), "install.sh must not gate _install_bnb_rocm on a bare */rocm* whole-URL glob"
class TestPinnedIndexClearsUvEnvParity:
"""Every installer must neutralise the uv index env vars for a pinned torch
install (#6898). uv treats the default index (--index-url / --default-index) as
lowest priority, so an inherited UV_INDEX / UV_EXTRA_INDEX_URL mirror would win
under uv's first-index strategy and pull torch from the wrong index -- after
which the pinned wheel index is silently never used."""
UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL")
def test_install_sh_clears_uv_index_vars(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
assert (
"env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in text
), "install.sh run_install_cmd must clear the uv index vars for --default-index installs"
def test_install_ps1_clears_uv_index_vars(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
for var in self.UV_VARS:
assert var in text, f"install.ps1 must clear {var} for pinned installs"
def test_setup_ps1_clears_uv_index_vars(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
for var in self.UV_VARS:
assert var in text, f"setup.ps1 must clear {var} for pinned installs"
def test_stack_py_clears_uv_index_vars(self):
text = STACK_PY.read_text(encoding = "utf-8")
assert "_install_env_for_cmd" in text, (
"install_python_stack.py must scrub inherited uv index vars for pinned "
"installs via _install_env_for_cmd (parity with install.sh #6898)"
)
for var in self.UV_VARS:
assert var in text, f"install_python_stack.py must clear {var} for pinned installs"
def test_all_installers_clear_uv_torch_backend(self):
"""uv's torch backend redirects torch resolution to its own per-backend
index even against an explicit pin, so every installer's pinned-install
scrub must clear UV_TORCH_BACKEND too."""
sh = INSTALL_SH.read_text(encoding = "utf-8")
assert "-u UV_TORCH_BACKEND" in sh, "install.sh pinned scrub must clear UV_TORCH_BACKEND"
for path in (INSTALL_PS1, SETUP_PS1):
text = path.read_text(encoding = "utf-8")
assert (
"'UV_TORCH_BACKEND'" in text
), f"{path.name} pinned scrub must clear UV_TORCH_BACKEND"
stack = STACK_PY.read_text(encoding = "utf-8")
assert (
'"UV_TORCH_BACKEND",' in stack
), "install_python_stack.py strip tuple must include UV_TORCH_BACKEND"
def test_stack_py_strips_pip_extra_index_for_pip_fallback(self):
"""The pip fallback honours PIP_EXTRA_INDEX_URL (pip adds it IN ADDITION
to --index-url), so the pinned-command scrub must strip it."""
stack = STACK_PY.read_text(encoding = "utf-8")
assert (
'"PIP_EXTRA_INDEX_URL",' in stack
), "install_python_stack.py strip tuple must include PIP_EXTRA_INDEX_URL"
def test_all_installers_scrub_find_links(self):
"""uv's --find-links (env UV_FIND_LINKS) adds candidate locations that can
satisfy torch off a pinned index; every pinned-install scrub must clear it."""
sh = INSTALL_SH.read_text(encoding = "utf-8")
assert "-u UV_FIND_LINKS" in sh
for path in (INSTALL_PS1, SETUP_PS1):
assert "'UV_FIND_LINKS'" in path.read_text(encoding = "utf-8"), path.name
stack = STACK_PY.read_text(encoding = "utf-8")
assert '"UV_FIND_LINKS",' in stack and '"PIP_FIND_LINKS",' in stack
def test_setup_ps1_scrub_covers_pip_fallback(self):
"""setup.ps1's Fast-Install must keep the scrub active through the pip
fallback (pip honours PIP_EXTRA_INDEX_URL / PIP_FIND_LINKS in addition to
--index-url); restoring the vars before the fallback reopens the hole."""
text = SETUP_PS1.read_text(encoding = "utf-8")
fi = text[text.find("function Fast-Install") :][:2500]
assert "'PIP_EXTRA_INDEX_URL'" in fi and "'PIP_FIND_LINKS'" in fi
# the pip fallback must sit INSIDE the try whose finally restores the vars
assert fi.find("python -m pip install") < fi.find(
"finally"
), "pip fallback must run before the scrub is restored"
def test_all_installers_disable_uv_config_for_pinned_installs(self):
"""A DISCOVERED uv.toml / pyproject [tool.uv] outranks the CLI pin
(verified with uv 0.10: [pip] torch-backend = "cpu" and a non-default
[[index]] both resolve torch+cpu against an explicit --index-url /
--default-index cu126 pin; UV_NO_CONFIG=1 restores the pin). Every
installer's pinned scrub must set UV_NO_CONFIG=1 and drop UV_CONFIG_FILE."""
sh = INSTALL_SH.read_text(encoding = "utf-8")
assert "-u UV_CONFIG_FILE UV_NO_CONFIG=1" in sh, (
"install.sh run_install_cmd must set UV_NO_CONFIG=1 and drop "
"UV_CONFIG_FILE for --default-index installs"
)
for path in (INSTALL_PS1, SETUP_PS1):
text = path.read_text(encoding = "utf-8")
assert "'UV_CONFIG_FILE'" in text, f"{path.name} must drop UV_CONFIG_FILE"
assert (
"$env:UV_NO_CONFIG = '1'" in text
), f"{path.name} must set UV_NO_CONFIG=1 for pinned installs"
stack = STACK_PY.read_text(encoding = "utf-8")
assert (
'"UV_CONFIG_FILE",' in stack
), "install_python_stack.py strip tuple must include UV_CONFIG_FILE"
assert (
'env["UV_NO_CONFIG"] = "1"' in stack
), "_install_env_for_cmd must set UV_NO_CONFIG=1 for pinned installs"
def test_pip_fallbacks_disable_pip_config_files(self):
"""The pip FALLBACK (uv missing/failed) honours user/site pip config files
even with the PIP_* env vars stripped: `pip config set
global.extra-index-url` still adds indexes to a pinned install. pip loads
NO configuration files when PIP_CONFIG_FILE is the platform devnull, so
the two installers that HAVE a pip fallback (install_python_stack.py and
setup.ps1's Fast-Install) must set it in their pinned scrub. install.sh
and install.ps1 are uv-only (no python -m pip fallback) and need no
equivalent."""
stack = STACK_PY.read_text(encoding = "utf-8")
assert 'env["PIP_CONFIG_FILE"] = os.devnull' in stack, (
"_install_env_for_cmd must point PIP_CONFIG_FILE at os.devnull for "
"pinned installs (pip fallback isolation)"
)
setup = SETUP_PS1.read_text(encoding = "utf-8")
assert "$env:PIP_CONFIG_FILE = 'nul'" in setup, (
"setup.ps1 Fast-Install pinned scrub must point PIP_CONFIG_FILE at nul "
"(Windows devnull) so the pip fallback ignores user/site pip config"
)
assert (
"'PIP_CONFIG_FILE'" in setup
), "setup.ps1 must save/restore PIP_CONFIG_FILE around the pinned scrub"
def test_setup_ps1_bounds_unknown_leaf_pinned_torch(self):
"""A first-time/changed unknown-leaf custom pin routes through setup.ps1's
CUDA branch; install.ps1's fresh pinned install, install.sh, and the Python
verbatim path bound the WHOLE trio, so the Windows update path must too -- a
private mirror serving newer torch OR newer companions must not lift the venv
above the supported range under the pin."""
text = SETUP_PS1.read_text(encoding = "utf-8")
# The custom-leaf branch bounds torch AND both companions (parity with the
# other installers' custom-pin trio bounds), gated on a non-cu-family leaf.
for spec in (
'$cudaTorchSpec = "torch>=2.4,<2.11.0"',
'$cudaVisionSpec = "torchvision>=0.19,<0.26.0"',
'$cudaAudioSpec = "torchaudio>=2.4,<2.11.0"',
):
assert spec in text, f"setup.ps1 must bound the custom-leaf trio: {spec}"
assert (
"if ($TorchIndexPinned -and -not (Test-CudaFamilyLeaf $CuTag)) {" in text
), "the custom-leaf trio bounds must be gated on a pinned non-cu-family leaf"
assert (
"Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec" in text
), "setup.ps1's CUDA branch must install via the bounded spec variables"
def test_setup_ps1_bounds_pinned_cpu_torch(self):
"""setup.ps1's CPU branch must bound the trio under an explicit pin (parity with
_CPU_TORCH_PKG_SPEC): the /cpu index serves newer torch, and _ensure_cpu_torch
keeps any CPU build, so a bare pinned trio could land an unsupported version.
An unpinned CPU host keeps the bare trio (pre-pin behavior unchanged)."""
text = SETUP_PS1.read_text(encoding = "utf-8")
for spec in (
'$cpuTorchSpec = "torch>=2.4,<2.12.0"',
'$cpuVisionSpec = "torchvision>=0.19,<0.27.0"',
'$cpuAudioSpec = "torchaudio>=2.4,<2.12.0"',
):
assert spec in text, f"setup.ps1 must bound the pinned CPU trio: {spec}"
assert (
"if ($TorchIndexPinned) {" in text
), "the CPU trio bounds must be gated on an explicit pin"
assert (
"Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce" in text
), "setup.ps1's CPU branch must install via the spec variables"
# The ceilings mirror the Python repair spec exactly.
stack = STACK_PY.read_text(encoding = "utf-8")
spec_block = re.search(r"_CUDA_TORCH_PKG_SPEC[^(]*\(\s*(.*?)\)", stack, re.DOTALL)
assert spec_block and '"torch>=2.4,<2.12.0"' in spec_block.group(1), (
"_CPU_TORCH_PKG_SPEC (via _CUDA_TORCH_PKG_SPEC) must keep the torch<2.12 "
"ceiling the setup.ps1 pinned CPU branch mirrors"
)
def test_setup_ps1_stale_check_requires_rocm_digit(self):
"""The stale-venv check must use the same EXACT rocm/gfx gate as the install
selection (Test-PipRocmFamilyLeaf), or a custom rocm-* / suffixed rocm7.2-private
leaf is stale-compared as a family and force-reinstalls on every studio update."""
text = SETUP_PS1.read_text(encoding = "utf-8")
anchor = text.find("$_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx")
assert anchor >= 0, "setup.ps1 stale check must classify the pinned leaf"
stale = text[anchor:][:2500]
assert (
"Test-PipRocmFamilyLeaf" in stale
), "setup.ps1 stale check must gate rocm leaves via the exact Test-PipRocmFamilyLeaf"
assert (
stale.count("-like 'rocm*'") == 0
), "setup.ps1 stale check must not use a bare -like 'rocm*' glob"
assert (
"-match '^rocm\\d'" not in stale
), "setup.ps1 stale check must not use an unanchored -match '^rocm\\d'"
class TestIndexPathSlashTrimParity:
"""Every installer must trim trailing PATH slashes only on the verbatim
UNSLOTH_TORCH_INDEX_URL override, preserving a ?query/#fragment token: a whole-URL
strip corrupts a base64 token ending in "/", a single strip leaves a double-slash leaf
empty. The helper must be DEFINED and WIRED into the override return in all four."""
def test_helper_defined_in_all_installers(self):
assert "def _trim_index_path_slashes(" in STACK_PY.read_text(encoding = "utf-8")
assert "_trim_index_path_slashes()" in INSTALL_SH.read_text(encoding = "utf-8")
assert "function Trim-IndexPathSlashes" in INSTALL_PS1.read_text(encoding = "utf-8")
assert "function Trim-IndexPathSlashes" in SETUP_PS1.read_text(encoding = "utf-8")
def test_helper_wired_into_override_in_all_installers(self):
assert "_trim_index_path_slashes(url)" in STACK_PY.read_text(encoding = "utf-8")
assert '_url=$(_trim_index_path_slashes "$_url")' in INSTALL_SH.read_text(encoding = "utf-8")
assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in INSTALL_PS1.read_text(
encoding = "utf-8"
)
assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in SETUP_PS1.read_text(
encoding = "utf-8"
)
class TestInstallOutputRedactionParity:
"""uv/pip failure text embeds the failing --index-url verbatim, so a captured install
log dumped on error can leak a user:token@ or ?token= secret. Every installer must
DEFINE a redaction helper and WIRE it into the captured-output print path."""
def test_helper_defined_in_all_installers(self):
assert "def _redact_install_output(" in STACK_PY.read_text(encoding = "utf-8")
assert "_redact_install_output()" in INSTALL_SH.read_text(encoding = "utf-8")
assert "function Redact-InstallOutput" in INSTALL_PS1.read_text(encoding = "utf-8")
assert "function Redact-InstallOutput" in SETUP_PS1.read_text(encoding = "utf-8")
def test_helper_wired_into_failure_print(self):
# install.sh dumps the captured log through the redactor on failure.
assert '_redact_install_output "$_log"' in INSTALL_SH.read_text(encoding = "utf-8")
# Both ps1 installers redact the captured $output before Write-Host on non-zero exit.
assert (
"Write-Host (Redact-InstallOutput $output) -ForegroundColor Red"
in INSTALL_PS1.read_text(encoding = "utf-8")
)
assert (
"Write-Host (Redact-InstallOutput $output) -ForegroundColor Red"
in SETUP_PS1.read_text(encoding = "utf-8")
)
# Python redacts the captured stdout before printing.
assert "_redact_install_output(" in STACK_PY.read_text(encoding = "utf-8")
class TestPipNoIndexScrubParity:
"""The plain-pip fallback honours PIP_*: PIP_NO_INDEX=1 makes it ignore ALL indexes
(defeating the pinned --index-url) and PIP_INDEX_URL replaces the pin. The two installers
that HAVE a plain-pip fallback (Python + setup.ps1) must scrub both for a pinned install.
install.sh / install.ps1 are uv-only (--default-index), which ignores pip config/env."""
def test_python_scrubs_pip_no_index_and_pip_index_url(self):
text = STACK_PY.read_text(encoding = "utf-8")
assert '"PIP_NO_INDEX"' in text
assert '"PIP_INDEX_URL"' in text
def test_setup_ps1_scrubs_pip_no_index_and_pip_index_url(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
assert "'PIP_NO_INDEX'" in text
assert "'PIP_INDEX_URL'" in text

View file

@ -54,6 +54,24 @@ class TestBuildUvCmdTorchBackend:
a.startswith("--torch-backend") for a in cmd
), f"Empty UV_TORCH_BACKEND should not add flag, got: {cmd}"
def test_uv_torch_backend_skipped_for_pinned_index(self):
"""A pinned-index command must NOT get --torch-backend: uv's torch backend
redirects torch resolution to its own per-backend index even when
--index-url is given (verified: cu128 pin + backend cpu installs
torch+cpu), defeating the pin."""
for pin_flag in ("--index-url", "--default-index"):
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
cmd = self._call(("torch", pin_flag, "https://download.pytorch.org/whl/cu128"))
assert not any(
a.startswith("--torch-backend") for a in cmd
), f"{pin_flag} command must not carry --torch-backend, got: {cmd}"
def test_uv_torch_backend_kept_for_unpinned(self):
"""Non-pinned commands still honour UV_TORCH_BACKEND."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
cmd = self._call(("somepackage",))
assert "--torch-backend=cpu" in cmd
class TestUvSafePath:
"""_uv_safe_path hands uv a space-free `-c`/`-r` path (issue #6503)."""
@ -148,3 +166,119 @@ class TestUvSafePathHardening:
assert " " not in value
assert Path(value).read_text() == "transformers>=4.57.6\n"
class TestPinnedIndexClearsUvEnv:
"""A pinned torch install (--index-url / --default-index) must neutralise an
inherited UV_INDEX / UV_EXTRA_INDEX_URL so the pinned wheel index wins.
uv treats the default index (--index-url / --default-index) as LOWEST priority,
so an inherited UV_INDEX / UV_EXTRA_INDEX_URL (a corporate/CPU mirror) would be
searched first and, under uv's default first-index strategy, resolve torch from
the wrong mirror -- after which the marker records a wheel index that was never
used. install.sh (#6898), install.ps1 and setup.ps1 already clear these for
pinned installs; install_python_stack must match (parity across all installers).
"""
UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL")
def test_pinned_index_url_strips_uv_index_vars(self):
cmd = [
"uv",
"pip",
"install",
"--force-reinstall",
"torch",
"torchvision",
"torchaudio",
"--index-url",
"https://download.pytorch.org/whl/cu128",
]
with mock.patch.dict(
os.environ,
{
"UV_INDEX": "https://mirror.corp/simple",
"UV_EXTRA_INDEX_URL": "https://mirror.corp/extra",
"UV_INDEX_URL": "https://mirror.corp/root",
"UV_DEFAULT_INDEX": "https://mirror.corp/default",
},
):
env = ips._install_env_for_cmd(cmd)
assert env is not None, "a --index-url install must run with a scrubbed env"
for var in self.UV_VARS:
assert var not in env, f"{var} must be cleared for a pinned-index install"
def test_pinned_default_index_strips_uv_index_vars(self):
# --default-index must be gated too (matches install.sh / install.ps1).
cmd = ["uv", "pip", "install", "torch", "--default-index", "https://x/cu126"]
with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}):
env = ips._install_env_for_cmd(cmd)
assert env is not None
assert "UV_INDEX" not in env
def test_non_pinned_install_keeps_user_mirror(self):
# A plain install (no --index-url) must NOT scrub the env, so a user's mirror
# still applies to base packages.
cmd = ["uv", "pip", "install", "unsloth", "unsloth-zoo"]
with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}):
env = ips._install_env_for_cmd(cmd)
assert env is None, "non-pinned installs must inherit the caller env unchanged"
def test_scrubbed_env_preserves_other_vars(self):
cmd = ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
with mock.patch.dict(
os.environ,
{"UV_INDEX": "https://mirror.corp/simple", "PATH_SENTINEL_XYZ": "keepme"},
):
env = ips._install_env_for_cmd(cmd)
assert env is not None
assert env.get("PATH_SENTINEL_XYZ") == "keepme", "only uv index vars are removed"
def test_pinned_cmd_strips_pip_extra_index_url(self):
"""PIP_EXTRA_INDEX_URL is stripped for pinned commands so the pip
fallback cannot satisfy torch from an inherited extra index."""
with mock.patch.dict(os.environ, {"PIP_EXTRA_INDEX_URL": "https://mirror/simple"}):
env = ips._install_env_for_cmd(
["pip", "install", "torch", "--index-url", "https://x/cu128"]
)
assert env is not None and "PIP_EXTRA_INDEX_URL" not in env
def test_pinned_cmd_strips_uv_torch_backend(self):
"""UV_TORCH_BACKEND is stripped for pinned commands so uv cannot read it
from the environment and reroute torch off the pinned index."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
env = ips._install_env_for_cmd(
["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
)
assert env is not None and "UV_TORCH_BACKEND" not in env
def test_pinned_cmd_disables_uv_config_discovery(self):
"""A DISCOVERED uv.toml / pyproject [tool.uv] outranks the CLI pin too
(verified with uv 0.10: [pip] torch-backend = "cpu" and a non-default
[[index]] both resolve torch+cpu against an explicit --index-url /
--default-index cu126 pin). Pinned commands must run with UV_NO_CONFIG=1
and without an inherited UV_CONFIG_FILE."""
with mock.patch.dict(os.environ, {"UV_CONFIG_FILE": "/etc/uv/uv.toml"}):
env = ips._install_env_for_cmd(
["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
)
assert env is not None
assert env.get("UV_NO_CONFIG") == "1"
assert "UV_CONFIG_FILE" not in env
def test_pinned_cmd_disables_pip_config_files(self):
"""The pip FALLBACK honours user/site pip config files (pip config set
global.extra-index-url) even with the PIP_* env vars stripped; pip loads
NO configuration files when PIP_CONFIG_FILE is os.devnull. Harmless for
uv, decisive for the fallback."""
env = ips._install_env_for_cmd(
["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
)
assert env is not None
assert env.get("PIP_CONFIG_FILE") == os.devnull
def test_non_pinned_cmd_keeps_uv_config_discovery(self):
"""Non-pinned installs inherit the caller env unchanged, so a user's uv
configuration still applies to base packages."""
env = ips._install_env_for_cmd(["uv", "pip", "install", "unsloth"])
assert env is None

View file

@ -15,6 +15,7 @@ sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh"
sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh"
sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh"
sh "$TESTS_DIR/sh/test_torch_flavor.sh"
sh "$TESTS_DIR/sh/test_redact_install_output.sh"
sh "$TESTS_DIR/sh/test_install_uv_override_space.sh"
echo ""

View file

@ -23,6 +23,8 @@ _FAKE_SMI_DIR=$(mktemp -d)
echo ""
sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_trim_index_path_slashes()/,/^}/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"
@ -379,6 +381,61 @@ _result=$(run_func "$_dir" " -1 ")
assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# --- explicit overrides (headless / container / CI; no GPU probing) ----------
# 39) UNSLOTH_TORCH_INDEX_FAMILY pins the family with no GPU present (not the cpu fallback).
_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
assert_eq "family override (no GPU) -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
# 40) Family override beats real detection: an nvidia-smi 12.6 host still gets cu128
# (the Docker-build case -- builder sees the host driver but publishes a cu128 image).
_dir=$(make_mock_smi "12.6")
_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "$_dir")
assert_eq "family override beats detected 12.6 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
rm -rf "$_dir"
# 41) UNSLOTH_TORCH_INDEX_URL is used verbatim and wins over detection.
_dir=$(make_mock_smi "12.6")
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu999" run_func "$_dir")
assert_eq "url override beats detection -> verbatim" "https://mirror.example.com/whl/cu999" "$_result"
rm -rf "$_dir"
# 42) Family override is appended to UNSLOTH_PYTORCH_MIRROR (mirror still honoured).
_result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
assert_eq "mirror + family override -> mirror/cu128" "https://mirror.example.com/whl/cu128" "$_result"
# 43) Trailing slash in UNSLOTH_TORCH_INDEX_URL is stripped.
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128/" run_func "none")
assert_eq "url override trailing slash stripped" "https://mirror.example.com/whl/cu128" "$_result"
# 44) URL override takes precedence over family override.
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu130" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
assert_eq "url override beats family override -> url" "https://mirror.example.com/whl/cu130" "$_result"
# 45) An empty override is ignored (falls through to normal detection).
_result=$(UNSLOTH_TORCH_INDEX_FAMILY="" UNSLOTH_TORCH_INDEX_URL="" run_func "none")
assert_eq "empty overrides ignored -> detected cpu" "https://download.pytorch.org/whl/cpu" "$_result"
# 46) ALL trailing slashes are stripped from a URL override (not just one).
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128///" run_func "none")
assert_eq "url override double slash stripped" "https://mirror.example.com/whl/cu128" "$_result"
# 47) Leading and trailing slashes stripped from a family override.
_result=$(UNSLOTH_TORCH_INDEX_FAMILY="//cu128//" run_func "none")
assert_eq "family override slashes stripped" "https://download.pytorch.org/whl/cu128" "$_result"
# 48) A ?query token that ends in "/" is PRESERVED: only PATH slashes are trimmed, so a
# base64 token ending in "/" is not corrupted (path-only trim, not whole-URL rstrip).
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128?token=ab12cd/" run_func "none")
assert_eq "url override preserves query token slash" "https://mirror.example.com/whl/cu128?token=ab12cd/" "$_result"
# 49) Double PATH slash before a query is collapsed while the query survives intact.
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128//?token=ab12cd/" run_func "none")
assert_eq "url override path slash trimmed, query kept" "https://mirror.example.com/whl/cu128?token=ab12cd/" "$_result"
# 50) A #fragment ending in "/" is likewise preserved.
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128#anchor/" run_func "none")
assert_eq "url override preserves fragment slash" "https://mirror.example.com/whl/cu128#anchor/" "$_result"
rm -f "$_FUNC_FILE"
rm -rf "$_FAKE_SMI_DIR"
rm -rf "$_TOOLS_DIR"

View file

@ -0,0 +1,89 @@
#!/bin/bash
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Unit tests for install.sh's _redact_install_output helper. uv/pip failure text embeds the
# failing --index-url verbatim, so a captured install log dumped on error can leak a
# user:token@ or ?token= secret. The helper redacts both before printing. Mirrors
# _redact_install_output (install_python_stack.py) / Redact-InstallOutput (install.ps1 /
# setup.ps1).
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
_FUNC_FILE=$(mktemp)
sed -n '/^_redact_install_output()/,/^}/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
}
# Redact from a file (the actual call site passes a captured-log tempfile).
redact_str() {
_rs_tmp=$(mktemp)
printf '%s\n' "$1" > "$_rs_tmp"
_rs_out=$(_redact_install_output "$_rs_tmp")
rm -f "$_rs_tmp"
printf '%s' "$_rs_out"
}
echo "=== _redact_install_output ==="
assert_eq "userinfo user:token@ redacted" \
"ERROR: failed https://<redacted>@download.pytorch.org/whl/cu128" \
"$(redact_str 'ERROR: failed https://alice:s3cr3t@download.pytorch.org/whl/cu128')"
assert_eq "bare-token@ userinfo redacted" \
"fetch https://<redacted>@host/whl/cu128 failed" \
"$(redact_str 'fetch https://ghp_deadbeef@host/whl/cu128 failed')"
assert_eq "single ?token= query redacted" \
"url https://host/whl/cu128?token=<redacted> unreachable" \
"$(redact_str 'url https://host/whl/cu128?token=abcd1234 unreachable')"
assert_eq "multiple query values redacted" \
"https://host/whl/cu128?token=<redacted>&channel=<redacted>" \
"$(redact_str 'https://host/whl/cu128?token=abcd1234&channel=beta')"
assert_eq "http (not https) userinfo redacted" \
"http://<redacted>@host/simple" \
"$(redact_str 'http://u:p@host/simple')"
assert_eq "fragment token redacted" \
"ERROR: could not fetch https://mirror.local/whl/cu128#<redacted> (403)" \
"$(redact_str 'ERROR: could not fetch https://mirror.local/whl/cu128#token=SECRET123 (403)')"
assert_eq "query and fragment both redacted" \
"https://host/whl/cu128?token=<redacted>#<redacted> done" \
"$(redact_str 'https://host/whl/cu128?token=abc#sig=xyz done')"
# Non-secret text is untouched (no false positives on ordinary log lines).
assert_eq "plain line untouched" \
"Resolved 42 packages in 1.2s" \
"$(redact_str 'Resolved 42 packages in 1.2s')"
assert_eq "plain url without creds untouched" \
"downloading https://download.pytorch.org/whl/cu128/torch-2.8.0.whl" \
"$(redact_str 'downloading https://download.pytorch.org/whl/cu128/torch-2.8.0.whl')"
assert_eq "bare hash comment untouched" \
"# retrying with --no-cache-dir" \
"$(redact_str '# retrying with --no-cache-dir')"
# Regression guard: no secret substring survives.
_leak=$(redact_str 'https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET')
case "$_leak" in
*s3cr3t*|*SUPERSECRET*|*ALSOSECRET*) assert_eq "no secret leak" "clean" "leaked:$_leak" ;;
*) assert_eq "no secret leak" "clean" "clean" ;;
esac
echo ""
echo "Results: $PASS passed, $FAIL failed"
[ "$FAIL" -eq 0 ]

View file

@ -108,6 +108,25 @@ assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var"
_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true)
assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded"
# Companions must be bounded to torch's window everywhere: the <2.11 bound appears
# twice (default assignments + the pinned custom-leaf block), never bare. torchaudio
# 2.11 dropped its exact torch pin, so a bare companion next to a <2.11-capped torch
# resolves a mismatched 2.11 build.
_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"' "$INSTALL_SH" || true)
assert_eq "torchvision bounded (<0.26) at default + custom-leaf" "2" "$_count"
_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"' "$INSTALL_SH" || true)
assert_eq "torchaudio bounded (<2.11) at default + custom-leaf" "2" "$_count"
_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision"$' "$INSTALL_SH" || true)
assert_eq "no bare torchvision companion remains" "0" "$_count"
_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio"$' "$INSTALL_SH" || true)
assert_eq "no bare torchaudio companion remains" "0" "$_count"
# The cu* widen must carry the companions with it (torch <2.12 with torchaudio <2.11
# would cap a mismatched pair the other way).
assert_eq "cu widen pairs torchaudio (<2.12)" "1" "$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0"' "$INSTALL_SH" || true)"
_gated=$(grep -c '_expected_torch_flavor_tag "$TORCH_INDEX_URL"' "$INSTALL_SH" || true)
_has_gate=$([ "$_gated" -ge 1 ] && echo "yes" || echo "no")
assert_eq "custom-companion bound gated on empty flavor tag" "yes" "$_has_gate"
# A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x land torch
# 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC).
_cuda_widen=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' "$INSTALL_SH" || true)
@ -285,6 +304,61 @@ bash -c "
_uv_got2=$(cat "$_UV_LOG2" 2>/dev/null || echo "")
assert_contains "mock uv arm64+py312 receives torch>=2.4" "$_uv_got2" "torch>=2.4,<2.11.0"
# ======================================================================
# ROCm 2.11 floor: leaf is lowercased before the gfx*/rocm* allowlist match
# ======================================================================
echo ""
echo "=== ROCm 2.11 floor case (leaf normalization) ==="
# Structural: install.sh lowercases _torch_index_leaf before the floor case, so the
# canonical gfx120X-all (capital X) matches gfx120x-all.
_has_lc=$(grep -c '_torch_index_leaf=$(printf .* | tr .\[:upper:\]. .\[:lower:\].)' "$INSTALL_SH" || true)
_has_lc_ok=$([ "$_has_lc" -ge 1 ] && echo "yes" || echo "no")
assert_eq "install.sh lowercases _torch_index_leaf" "yes" "$_has_lc_ok"
# Runtime: replicate install.sh's normalization + floor case and assert both gfx120X-all
# and gfx120x-all get the floor, while non-2.11 leaves keep the default.
run_floor_case() {
_url="$1"
bash -c '
TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
TORCHVISION_CONSTRAINT="torchvision"
TORCHAUDIO_CONSTRAINT="torchaudio"
_torch_index_leaf="${1%/}"
_torch_index_leaf="${_torch_index_leaf##*/}"
_torch_index_leaf=$(printf "%s" "$_torch_index_leaf" | tr "[:upper:]" "[:lower:]")
case "$_torch_index_leaf" in
rocm7.2|gfx120x-all|gfx1151|gfx1150)
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
;;
esac
echo "$TORCH_CONSTRAINT"
' _ "$_url"
}
assert_eq "gfx120X-all (capital) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all')"
assert_eq "gfx120X-all trailing slash -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all/')"
assert_eq "gfx120x-all (lowercase) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120x-all')"
assert_eq "gfx1151 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1151')"
assert_eq "gfx1150 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1150')"
assert_eq "rocm7.2 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://download.pytorch.org/whl/rocm7.2')"
assert_eq "gfx110X-all -> default (no floor)" "torch>=2.4,<2.11.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx110X-all')"
assert_eq "rocm6.4 -> default (no floor)" "torch>=2.4,<2.11.0" \
"$(run_floor_case 'https://download.pytorch.org/whl/rocm6.4')"
assert_eq "cu128 -> default (no floor)" "torch>=2.4,<2.11.0" \
"$(run_floor_case 'https://download.pytorch.org/whl/cu128')"
assert_eq "cpu -> default (no floor)" "torch>=2.4,<2.11.0" \
"$(run_floor_case 'https://download.pytorch.org/whl/cpu')"
# ======================================================================
# Summary
# ======================================================================

View file

@ -11,14 +11,21 @@ INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
# Extract the three helper functions from install.sh and source them.
# Extract the helper functions from install.sh and source them
# (_torch_index_url_leaf is the shared leaf extractor the classifiers call).
_FUNC_FILE=$(mktemp)
{
sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_torch_index_url_leaf()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_is_pip_rocm_family_leaf()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_expected_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_torch_index_repairable()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_tauri_torch_index_family()/,/^}/p' "$INSTALL_SH"
} > "$_FUNC_FILE"
# shellcheck disable=SC1090
. "$_FUNC_FILE"
@ -56,6 +63,26 @@ assert_eq "amd gfx index" "rocm" "$(_expected_torch_flavor_tag 'https://re
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 '')"
# Query/fragment dropped before classification: .../cu128?token=x classifies as cu128,
# not an opaque leaf that reinstalls every run.
assert_eq "query-bearing cu128" "cu128" "$(_expected_torch_flavor_tag 'https://m/whl/cu128?token=x')"
assert_eq "fragment-bearing cpu" "cpu" "$(_expected_torch_flavor_tag 'https://m/whl/cpu#frag')"
# A cu-suffixed CUSTOM leaf (cu128-private, cu128x) is NOT the cu128 family (exact
# cu+digits only). Mirrors Python re.fullmatch(cu[0-9]+) / PowerShell.
assert_eq "cu-suffix custom leaf" "" "$(_expected_torch_flavor_tag 'https://m/whl/cu128-private')"
assert_eq "cu-alnum custom leaf" "" "$(_expected_torch_flavor_tag 'https://m/whl/cu128x')"
assert_eq "bare cu digits stays" "cu126" "$(_expected_torch_flavor_tag 'https://m/whl/cu126')"
# A custom leaf merely STARTING with rocm (rocm-current, rocm-rel-7.2.1) is NOT a pip
# rocm family -> "" (custom); real families (rocm7.2) and gfx indexes stay "rocm".
assert_eq "custom rocm-current" "" "$(_expected_torch_flavor_tag 'https://mirror/whl/rocm-current')"
assert_eq "radeon rocm-rel leaf" "" "$(_expected_torch_flavor_tag 'https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1')"
assert_eq "real rocm7.2 stays" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7.2')"
# A rocm<digit>-SUFFIX private mirror (rocm7.2-private, rocm7-current) is a custom pin ->
# "" (custom); match the family exactly, not the prefix.
assert_eq "suffixed rocm7.2-private" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7.2-private')"
assert_eq "suffixed rocm7-current" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7-current')"
assert_eq "two-dot rocm7.2.1" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7.2.1')"
assert_eq "bare rocm7 stays" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7')"
echo "=== _torch_index_repairable ==="
assert_eq "cu130 repairable" "yes" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cu130')"
@ -64,6 +91,66 @@ assert_eq "gfx repairable" "yes" "$(_torch_index_repairable 'https://repo.
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')"
# A suffixed rocm leaf is a verbatim pin, not a --default-index repairable family.
assert_eq "rocm-private NOT repair" "no" "$(_torch_index_repairable 'https://co.internal/whl/rocm7.2-private')"
echo "=== _is_pip_rocm_family_leaf ==="
assert_family() {
_label="$1"; _expected="$2"; _leaf="$3"
if _is_pip_rocm_family_leaf "$_leaf"; then _actual="yes"; else _actual="no"; fi
assert_eq "$_label" "$_expected" "$_actual"
}
assert_family "rocm7.2 family" "yes" "rocm7.2"
assert_family "rocm6.4 family" "yes" "rocm6.4"
assert_family "bare rocm7 family" "yes" "rocm7"
assert_family "gfx120x-all family" "yes" "gfx120x-all"
assert_family "gfx1151 family" "yes" "gfx1151"
assert_family "rocm7.2-private custom" "no" "rocm7.2-private"
assert_family "rocm7-current custom" "no" "rocm7-current"
assert_family "rocm-current custom" "no" "rocm-current"
assert_family "rocm-rel-7.2.1 custom" "no" "rocm-rel-7.2.1"
assert_family "rocm7.2.1 custom" "no" "rocm7.2.1"
# A trailing dot (rocm7.) or leading/double dot is NOT a family: both major and minor must
# be non-empty all-digits, matching Python re.fullmatch(rocm\d+(?:\.\d+)?). Bash previously
# accepted rocm7. via a bare %/-style trim while Python rejected it (validator asymmetry).
assert_family "rocm7. trailing-dot custom" "no" "rocm7."
assert_family "rocm.7 leading-dot custom" "no" "rocm.7"
assert_family "rocm7..2 double-dot custom" "no" "rocm7..2"
assert_family "cpu not rocm" "no" "cpu"
assert_family "cu128 not rocm" "no" "cu128"
assert_family "simple not rocm" "no" "simple"
echo "=== _torch_index_url_leaf (ALL trailing slashes stripped -> non-empty leaf) ==="
# A double (or triple) trailing slash must yield the real leaf, not an empty string that
# fails every classifier arm. Python .rstrip("/") drops them all; bash must match (a bare
# %/ left .../cu128// classifying as "").
assert_eq "double slash cu128 leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128//')"
assert_eq "triple slash rocm7.2 leaf" "rocm7.2" "$(_torch_index_url_leaf 'https://m/whl/rocm7.2///')"
assert_eq "double slash + token leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128//?token=x')"
assert_eq "single slash cu128 leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128/')"
# The classifier that consumes the leaf must therefore still tag a double-slash index.
assert_eq "double-slash cu128 tag" "cu128" "$(_expected_torch_flavor_tag 'https://m/whl/cu128//')"
assert_eq "double-slash rocm7.2 tag" "rocm" "$(_expected_torch_flavor_tag 'https://m/whl/rocm7.2//')"
echo "=== _tauri_torch_index_family (credential redaction) ==="
# A token/fragment must be stripped BEFORE classification so it never reaches the
# [TAURI:DIAG] line (the family is the last path segment, which else carries the query).
SKIP_TORCH=false
assert_eq "token stripped from rocm" "rocm7.2" "$(_tauri_torch_index_family 'https://mirror/whl/rocm7.2?token=SECRET')"
assert_eq "token-bearing cu classifies" "cu128" "$(_tauri_torch_index_family 'https://m/whl/cu128?token=x')"
assert_eq "fragment stripped cpu" "cpu" "$(_tauri_torch_index_family 'https://m/whl/cpu#frag')"
assert_eq "plain rocm7.2 unchanged" "rocm7.2" "$(_tauri_torch_index_family 'https://download.pytorch.org/whl/rocm7.2')"
# A trailing slash must be stripped too, or the */cu128 and */cpu arms miss .../cu128/
# and it falls through to "auto".
assert_eq "trailing slash cu128" "cu128" "$(_tauri_torch_index_family 'https://download.pytorch.org/whl/cu128/')"
assert_eq "slash + token cu128" "cu128" "$(_tauri_torch_index_family 'https://m/whl/cu128/?token=x')"
assert_eq "trailing slash cpu" "cpu" "$(_tauri_torch_index_family 'https://m/whl/cpu/')"
# Regression guard: no secret token substring may survive in any classification.
_leak=$(_tauri_torch_index_family 'https://mirror/whl/rocm7.2?token=SECRET')
case "$_leak" in
*SECRET*|*token*) assert_eq "no token leak in family" "clean" "leaked:$_leak" ;;
*) assert_eq "no token leak in family" "clean" "clean" ;;
esac
echo ""
echo "Results: $PASS passed, $FAIL failed"

View file

@ -64,15 +64,23 @@ def _run_cuda_repair(
rocm_marker = False,
smi_path = "/usr/bin/nvidia-smi",
cvd = None,
index_family = None,
index_url = None,
):
"""Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock.
cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it."""
cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it.
index_family sets UNSLOTH_TORCH_INDEX_FAMILY (the explicit wheel-index pin).
index_url sets UNSLOTH_TORCH_INDEX_URL (the full-URL pin form)."""
env = {}
if rocm_marker:
env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1"
if cvd is not None:
env["CUDA_VISIBLE_DEVICES"] = cvd
if index_family is not None:
env["UNSLOTH_TORCH_INDEX_FAMILY"] = index_family
if index_url is not None:
env["UNSLOTH_TORCH_INDEX_URL"] = index_url
def _which(name, *a, **k):
if name == "nvidia-smi":
@ -99,6 +107,10 @@ def _run_cuda_repair(
stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None)
if cvd is None:
stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None)
if index_family is None:
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
if index_url is None:
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
_ensure_cuda_torch()
return mock_pip
@ -123,11 +135,74 @@ class TestCudaRepairFires:
assert mock_pip.call_args.kwargs["constrain"] is False
def test_rocm_in_version_string_triggers_repair(self):
# AMD SDK / Radeon wheels may encode rocm in __version__ without
# torch.version.hip; the probe prints "hip" for both.
# AMD SDK / Radeon wheels may encode rocm in __version__ without torch.version.hip;
# the probe prints "hip" for both.
mock_pip = _run_cuda_repair(torch_state = "hip")
assert mock_pip.call_count == 1
def test_no_gpu_but_explicit_cuda_pin_repairs(self):
# Headless / CI cross-install: an explicit cu* pin commits to CUDA wheels with no
# NVIDIA GPU visible, so a ROCm-poisoned venv is still repaired to the pinned family.
mock_pip = _run_cuda_repair(
nvidia = False,
backend = "cuda",
index_family = "cu128",
torch_state = "hip",
)
assert mock_pip.call_count == 1
assert "cu128" in _index_url(mock_pip)
def test_cvd_hidden_but_explicit_cuda_pin_repairs(self):
# CVD=-1/"" hides the GPU, but an explicit cu* pin skips ALL host-GPU probing, so the
# CVD hide gate must not suppress the repair (GPU-less CI: CVD=-1, FAMILY=cu128).
for _cvd in ("-1", ""):
mock_pip = _run_cuda_repair(
nvidia = False,
backend = "cuda",
cvd = _cvd,
index_family = "cu128",
torch_state = "hip",
)
assert mock_pip.call_count == 1
assert "cu128" in _index_url(mock_pip)
def test_tagged_cuda_mismatch_repairs(self):
# A healthy CUDA torch whose +cuXXX differs from the pin is repaired.
mock_pip = _run_cuda_repair(
index_family = "cu128",
torch_state = "cuda|cu126",
cuda_version = "12.8",
)
assert mock_pip.call_count == 1
assert "cu128" in _index_url(mock_pip)
def test_untagged_cuda_build_under_pin_repairs(self):
# An untagged CUDA build (no +cuXXX tag -> empty installed cu) can't be confirmed
# to match the pin, so the pin is enforced with a reinstall.
mock_pip = _run_cuda_repair(
index_family = "cu128",
torch_state = "cuda", # marker cuda, empty installed cu
cuda_version = "12.8",
)
assert mock_pip.call_count == 1
assert "cu128" in _index_url(mock_pip)
def test_broken_probe_with_cuda_pin_repairs(self):
# torch present but unimportable under a CUDA pin: the base update won't repair a
# broken already-installed torch, so reinstall from the pin instead of stranding it.
mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1, index_family = "cu128")
assert mock_pip.call_count == 1
assert "cu128" in _index_url(mock_pip)
def test_broken_probe_with_cuda_url_pin_repairs(self):
mock_pip = _run_cuda_repair(
torch_state = "cpu",
torch_rc = 1,
index_url = "https://mirror.local/cu128",
)
assert mock_pip.call_count == 1
assert "https://mirror.local/cu128" in _index_url(mock_pip)
# No-op cases.
@ -157,8 +232,9 @@ class TestCudaRepairSkips:
mock_pip = _run_cuda_repair(nvidia = False, torch_state = "hip")
mock_pip.assert_not_called()
def test_torch_missing_skips(self):
# Non-zero probe exit = torch missing / un-importable.
def test_torch_missing_no_pin_skips(self):
# Non-zero probe exit = torch missing/un-importable. With NO CUDA pin the base
# install owns it, so leave it alone (a pinned build reinstalls).
mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1)
mock_pip.assert_not_called()
@ -191,6 +267,109 @@ class TestCudaRepairSkips:
mock_pip = _run_cuda_repair(cvd = "0", torch_state = "hip")
assert mock_pip.call_count == 1
def test_matching_tagged_cuda_pin_no_repair(self):
# Healthy CUDA torch whose +cuXXX already matches the pin: no reinstall.
mock_pip = _run_cuda_repair(
index_family = "cu128",
torch_state = "cuda|cu128",
cuda_version = "12.8",
)
mock_pip.assert_not_called()
def test_custom_mirror_leaf_not_treated_as_cuda_pin(self):
# A mirror leaf starting with "cu" but not cuXXX (.../custom, .../current) must
# NOT be treated as a CUDA pin, so it can't bypass the NVIDIA gate.
for _leaf in ("custom", "current"):
mock_pip = _run_cuda_repair(
nvidia = False,
backend = "cuda",
index_url = f"https://mymirror.example/{_leaf}",
torch_state = "hip",
)
mock_pip.assert_not_called()
def test_explicit_cuda_family_leaf_helper(self):
# _explicit_cuda_torch_index_url matches cuXXX narrowly, not any cu* leaf.
import contextlib
def _with(url):
with patch.dict(stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
return stack_mod._explicit_cuda_torch_index_url()
assert _with("https://download.pytorch.org/whl/cu128") is not None
assert _with("https://download.pytorch.org/whl/cu126") is not None
assert _with("https://mymirror.example/custom") is None
assert _with("https://mymirror.example/current") is None
assert _with("https://download.pytorch.org/whl/cpu") is None
with contextlib.suppress(Exception):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
class TestTorchBackendDerivationFromPin:
"""The module-level _TORCH_BACKEND derivation (standalone `studio update`
with no install.sh-set UNSLOTH_TORCH_BACKEND) must classify the pinned index
leaf via _is_cuda_family_leaf (^cu[0-9]), NOT a bare startswith("cu"). A
full-override URL ending in /current or /custom must fall through to backend
"" (probe the GPU) so _ensure_rocm_torch() still repairs a wrong/CPU torch on
AMD hosts, instead of being wrongly branded "cuda" and returning early."""
@staticmethod
def _derive(env):
# Re-run the module's import-time derivation, using its own _is_cuda_family_leaf
# so this stays in lockstep.
idx_override = (
env.get("UNSLOTH_TORCH_INDEX_URL", "").strip()
or env.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip()
)
backend = env.get("UNSLOTH_TORCH_BACKEND", "").lower()
if not backend:
leaf = idx_override.rstrip("/").rsplit("/", 1)[-1].lower()
if leaf.startswith(("rocm", "gfx")):
backend = "rocm"
elif leaf == "cpu":
backend = "cpu"
elif stack_mod._is_cuda_family_leaf(leaf):
backend = "cuda"
return backend
def test_cu128_pin_is_cuda(self):
assert (
self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://download.pytorch.org/whl/cu128"})
== "cuda"
)
def test_cu128_family_is_cuda(self):
assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "cu128"}) == "cuda"
def test_current_leaf_not_cuda(self):
# ^cu[0-9] rejects /current -> backend stays "" (probe GPU), so an AMD host still
# repairs a CPU/wrong torch instead of short-circuiting.
assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/current"}) == ""
def test_custom_leaf_not_cuda(self):
assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/custom"}) == ""
def test_rocm_and_gfx_pins_are_rocm(self):
assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"}) == "rocm"
assert (
self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx120X-all"})
== "rocm"
)
def test_cpu_pin_is_cpu(self):
assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "cpu"}) == "cpu"
def test_source_uses_helper_not_bare_startswith(self):
# Guard against a regression back to elif _idx_leaf.startswith("cu").
src = _STACK_PATH.read_text(encoding = "utf-8")
assert (
"elif _is_cuda_family_leaf(_idx_leaf):" in src
), "_TORCH_BACKEND derivation must classify CUDA via _is_cuda_family_leaf"
assert (
'elif _idx_leaf.startswith("cu"):' not in src
), "_TORCH_BACKEND derivation must not use a bare startswith('cu')"
# CUDA index ladder.

View file

@ -284,7 +284,9 @@ class TestBackendExportLeafClassification:
def test_export_block_uses_leaf(self, install_src):
anchor = install_src.find("_torch_index_leaf=")
assert anchor >= 0, "backend export must classify on the final path segment"
window = install_src[anchor : anchor + 500]
# Window spans the leaf-normalization prelude (query/frag drop + all-slash trim loop)
# through the export case arms.
window = install_src[anchor : anchor + 900]
assert 'export UNSLOTH_TORCH_BACKEND="rocm"' in window
assert 'export UNSLOTH_TORCH_BACKEND="cpu"' in window
assert 'export UNSLOTH_TORCH_BACKEND="cuda"' in window
@ -459,3 +461,130 @@ class TestHiddenCvdNotUsable:
cvd,
)
assert out == expected
class TestRedactInstallOutput:
"""_redact_install_output scrubs index-URL credentials from a captured install log
before it is printed on failure (uv/pip embeds the failing --index-url verbatim)."""
def test_userinfo_redacted(self):
out = stack_mod._redact_install_output(
"ERROR: failed https://alice:s3cr3t@download.pytorch.org/whl/cu128"
)
assert out == "ERROR: failed https://<redacted>@download.pytorch.org/whl/cu128"
def test_bytes_input_decoded_and_redacted(self):
out = stack_mod._redact_install_output(b"fetch https://ghp_deadbeef@host/whl/cu128 failed")
assert out == "fetch https://<redacted>@host/whl/cu128 failed"
def test_query_values_redacted(self):
out = stack_mod._redact_install_output(
"url https://host/whl/cu128?token=abcd1234&channel=beta unreachable"
)
assert out == "url https://host/whl/cu128?token=<redacted>&channel=<redacted> unreachable"
def test_fragment_redacted(self):
out = stack_mod._redact_install_output(
"ERROR: could not fetch https://mirror.local/whl/cu128#token=SECRET123 (403)"
)
assert out == "ERROR: could not fetch https://mirror.local/whl/cu128#<redacted> (403)"
def test_query_and_fragment_both_redacted(self):
out = stack_mod._redact_install_output("https://host/whl/cu128?token=abc#sig=xyz done")
assert out == "https://host/whl/cu128?token=<redacted>#<redacted> done"
def test_bare_hash_comment_untouched(self):
# The fragment redaction is URL-anchored: a shell comment in tool output survives.
assert (
stack_mod._redact_install_output("# retrying with --no-cache-dir")
== "# retrying with --no-cache-dir"
)
def test_plain_line_untouched(self):
assert (
stack_mod._redact_install_output("Resolved 42 packages in 1.2s")
== "Resolved 42 packages in 1.2s"
)
def test_no_secret_substring_survives(self):
out = stack_mod._redact_install_output(
"https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET"
)
assert "s3cr3t" not in out and "SUPERSECRET" not in out and "ALSOSECRET" not in out
class TestTrimIndexPathSlashes:
"""_trim_index_path_slashes strips trailing PATH slashes only; a ?query/#fragment token
ending in "/" must survive (a whole-URL rstrip would corrupt a base64 token)."""
def test_double_path_slash_collapsed(self):
assert stack_mod._trim_index_path_slashes("https://h/whl/cu128//") == "https://h/whl/cu128"
def test_query_token_slash_preserved(self):
assert (
stack_mod._trim_index_path_slashes("https://h/whl/cu128?token=ab12cd/")
== "https://h/whl/cu128?token=ab12cd/"
)
def test_path_slash_trimmed_query_kept(self):
assert (
stack_mod._trim_index_path_slashes("https://h/whl/cu128//?token=ab12cd/")
== "https://h/whl/cu128?token=ab12cd/"
)
def test_fragment_slash_preserved(self):
assert (
stack_mod._trim_index_path_slashes("https://h/whl/cu128#anchor/")
== "https://h/whl/cu128#anchor/"
)
class TestRocmFamilyLeafParity:
"""_is_pip_rocm_family_leaf must match re.fullmatch(rocm\\d+(?:\\.\\d+)?): a trailing dot
(rocm7.) is a CUSTOM pin, not a family (the historical bash/py validator asymmetry)."""
@pytest.mark.parametrize(
"leaf, expected",
[
("rocm7", True),
("rocm7.2", True),
("gfx1151", True),
("rocm7.", False),
("rocm.7", False),
("rocm7..2", False),
("rocm7.2.1", False),
("rocm7.2-private", False),
("cpu", False),
("cu128", False),
],
)
def test_family_classification(self, leaf, expected):
assert stack_mod._is_pip_rocm_family_leaf(leaf) is expected
class TestTorchIndexLeafAllSlashes:
"""_torch_index_leaf drops query/fragment then strips ALL trailing slashes, so a
double-slash index still yields the real leaf (not an empty string)."""
@pytest.mark.parametrize(
"url, expected",
[
("https://m/whl/cu128//", "cu128"),
("https://m/whl/rocm7.2///", "rocm7.2"),
("https://m/whl/cu128//?token=x", "cu128"),
("https://m/whl/cu128/", "cu128"),
],
)
def test_leaf_never_empty_on_double_slash(self, url, expected):
assert stack_mod._torch_index_leaf(url) == expected
class TestUvIndexEnvVarsScrub:
"""The pinned-install env scrub must drop PIP_NO_INDEX (which makes the pip fallback
ignore ALL indexes, defeating the pin) and PIP_INDEX_URL (replaces the pinned index)."""
def test_pip_no_index_scrubbed(self):
assert "PIP_NO_INDEX" in stack_mod._UV_INDEX_ENV_VARS
def test_pip_index_url_scrubbed(self):
assert "PIP_INDEX_URL" in stack_mod._UV_INDEX_ENV_VARS

View file

@ -784,7 +784,7 @@ def test_install_python_stack_windows_rocm_repair_pins_and_is_nonfatal():
assert re.search(
r'"' + gfx + r'":\s*_ROCM_TORCH_PKG_SPECS\["rocm7\.2"\]', text
), f"{gfx} must pin to the rocm7.2 trio like install.ps1/setup.ps1"
i = text.find('f"ROCm torch (Windows, {gfx_arch})"')
i = text.find("f\"ROCm torch (Windows, {gfx_arch or 'pinned'})\"")
assert i != -1, "Windows ROCm repair pip call not found"
# The nearest preceding call must be the nonfatal pip_install_try, not pip_install.
j = text.rfind("pip_install_try(", 0, i)

View file

@ -569,19 +569,27 @@ class TestEnsureRocmTorch:
_ensure_rocm_torch()
mock_pip.assert_not_called()
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
"""If torch already has CUDA, should skip ROCm reinstall."""
def test_cuda_torch_on_amd_host_reinstalls(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""A CUDA-only torch build is unusable on an AMD-only host, so it must be
reinstalled to ROCm (has_hip_torch is driven by the empty HIP marker, not
by treating the CUDA version string as a HIP marker)."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"12.6\n" # CUDA version
# Single-line probe: empty HIP marker before "|" for a CUDA build.
mock_probe.stdout = b"|2.10.0+cu126\n"
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
mock_pip.assert_not_called()
assert mock_pip.call_count == 1
assert "rocm7.1" in str(mock_pip.call_args_list[0])
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@ -591,12 +599,31 @@ class TestEnsureRocmTorch:
"""If torch already has HIP, should skip ROCm reinstall."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"7.1.12345\n" # HIP version
mock_probe.stdout = b"7.1.12345|2.10.0+rocm7.1\n" # HIP marker + version
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
mock_pip.assert_not_called()
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1))
def test_cpu_torch_probe_line_not_read_as_hip(self, mock_ver, mock_gpu, mock_nvidia, mock_pip):
"""A CPU build's probe line ("|2.10.0+cpu") must not read as HIP: the version
after the "|" separator is data, not a HIP marker, so has_hip_torch stays False
and the reinstall fires."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"|2.10.0+cpu\n"
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
with patch.object(stack_mod, "pip_install_try", return_value = True):
_ensure_rocm_torch()
assert mock_pip.call_count == 1
assert "rocm7.1" in str(mock_pip.call_args_list[0])
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@ -680,6 +707,295 @@ class TestEnsureRocmTorch:
torch_call = mock_pip.call_args_list[0]
assert "rocm7.2" in str(torch_call)
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4))
def test_explicit_gfx_index_honored_and_skips_strix_reroute(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""An explicit gfx wheel-index pin is authoritative: install from it verbatim
with torch 2.11, and never re-probe gfx codes to second-guess it (host ROCm 6.4
would otherwise pick the rocm6.4 wheel / trigger the Strix re-route)."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"\n" # cpu torch -> reinstall
env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"}
with patch.dict(stack_mod.os.environ, env, clear = False):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
# Would raise if the Strix block ran (it is skipped on an explicit pin).
with patch.object(
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
):
_ensure_rocm_torch()
assert mock_pip.call_count == 1
torch_call = str(mock_pip.call_args_list[0])
assert "gfx1151" in torch_call
assert "torch>=2.11.0,<2.12.0" in torch_call
def test_rocm_pin_family_mismatch_helper(self):
"""_rocm_pin_family_mismatch: exact rocm compare, else the 2.11 line."""
f = stack_mod._rocm_pin_family_mismatch
base = "https://download.pytorch.org/whl"
amd = "https://repo.amd.com/rocm/whl"
# Exact rocm version comparison.
assert f(f"{base}/rocm7.2", "2.11.0+rocm7.2") is False
assert f(f"{base}/rocm7.2", "2.10.0+rocm6.4") is True
assert f(f"{base}/rocm6.4", "2.10.0+rocm6.4") is False
# rocm7.2 is KNOWN-2.11. A +rocm7.2 wheel whose RELEASE drifted off 2.11 shares the
# tag but violates the spec -> mismatch (a plain version compare would accept it).
assert f(f"{base}/rocm7.2", "2.12.0+rocm7.2") is True
assert f(f"{base}/rocm7.2", "2.13.0+rocm7.2") is True
assert f(f"{base}/rocm7.2", "2.11.5+rocm7.2") is False # patch on 2.11 is in-spec
# An UNKNOWN newer rocm (not on the 2.11 allowlist) is not floored to 2.11, so a
# matching rocm version at any release line is NOT a mismatch on this branch.
assert f(f"{base}/rocm8.0", "2.12.0+rocm8.0") is False
# gfx pin (2.11 line) vs installed release line.
assert f(f"{amd}/gfx1151", "2.10.0+rocm6.4") is True
assert f(f"{amd}/gfx1151", "2.11.0+rocm7.13.0") is False
# rocm7.2 pin vs an untagged (no +rocm) wheel: a CPU/CUDA build never
# satisfies a ROCm pin, regardless of its release line -> always a mismatch.
assert f(f"{base}/rocm7.2", "2.10.0") is True
assert f(f"{base}/rocm7.2", "2.11.0") is True
assert f(f"{base}/rocm6.4", "2.10.0") is True
# A 2.11-allowlist gfx pin over a GENERIC (two-part +rocm7.2) 2.11 wheel mismatches:
# the user wants AMD's per-arch (three-part) wheel, not the generic one.
assert f(f"{amd}/gfx1151", "2.11.0+rocm7.2") is True
assert f(f"{amd}/gfx120X-all", "2.11.0+rocm7.2") is True
# ...but an already-installed per-arch (three-part) wheel is NOT re-flagged
# (no reinstall loop once the correct gfx wheel is present).
assert f(f"{amd}/gfx120X-all", "2.11.0+rocm7.13.0") is False
assert f(f"{amd}/gfx1150", "2.11.0+rocm7.13.0") is False
# A NON-2.11 gfx pin (gfx110X-all/gfx90a/gfx908) tracks the default <2.11 spec: a
# correct 2.10+rocm wheel is NOT a mismatch, a 2.11 build is.
assert f(f"{amd}/gfx110X-all", "2.10.0+rocm6.4") is False
assert f(f"{amd}/gfx90a", "2.10.0+rocm6.3") is False
assert f(f"{amd}/gfx908", "2.10.0+rocm7.0") is False
assert f(f"{amd}/gfx110X-all", "2.11.0+rocm7.2") is True
# A non-2.11 gfx pin over an untagged (no +rocm) wheel is a mismatch even
# when torch is already <2.11: a CPU/CUDA build never satisfies the ROCm pin.
assert f(f"{amd}/gfx110X-all", "2.10.0") is True
assert f(f"{amd}/gfx90a", "2.10.0") is True
# A major-only rocm pin (rocm7) compares on the major alone: rocm6.x mismatches,
# any rocm7.x satisfies it, an untagged wheel never does, a bare +rocm is lenient.
assert f(f"{base}/rocm7", "2.10.0+rocm6.4") is True
assert f(f"{base}/rocm7", "2.11.0+rocm7.2") is False
assert f(f"{base}/rocm7", "2.11.0+rocm7.13.0") is False
assert f(f"{base}/rocm7", "2.10.0") is True
assert f(f"{base}/rocm7", "2.10.0+rocm") is False
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
def test_rocm_pin_mismatch_over_installed_rocm_reinstalls(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""A rocm7.2 pin over an already-installed OLDER +rocm6.4 build must reinstall,
even though has_hip_torch is True (the ROCm analogue of the CUDA cuXXX mismatch)."""
mock_probe = MagicMock()
mock_probe.returncode = 0
# HIP marker present (has_hip_torch=True) + installed +rocm6.4 wheel.
mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n"
env = {"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"}
with patch.dict(stack_mod.os.environ, env, clear = False):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "rocm7.2" in torch_call
assert "torch>=2.11.0,<2.12.0" in torch_call
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4))
def test_gfx_pin_over_installed_pre211_rocm_reinstalls(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""A gfx* pin (2.11 line) over an installed pre-2.11 +rocm6.4 build reinstalls."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n"
env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"}
with patch.dict(stack_mod.os.environ, env, clear = False):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
with patch.object(
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "gfx1151" in torch_call
assert "torch>=2.11.0,<2.12.0" in torch_call
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
def test_rocm_pin_matches_installed_no_torch_reinstall(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""A rocm7.2 pin over an already-matching +rocm7.2 build must NOT reinstall torch
(no false reinstall of a correct ROCm venv)."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"7.2.12345|2.11.0+rocm7.2\n"
env = {"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"}
with patch.dict(stack_mod.os.environ, env, clear = False):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None)
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
# No torch reinstall: any pip_install call must not target a torch index.
for _call in mock_pip.call_args_list:
_args = [str(a) for a in _call.args]
if "--index-url" in _args:
_url = _args[_args.index("--index-url") + 1]
assert "rocm7.2" not in _url or "torch" not in " ".join(
_args
), "torch must not be reinstalled when the pin already matches"
# A torch reinstall would pass torch>=... as a positional; assert none did.
assert not any(
any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
)
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4))
def test_non211_gfx_pin_over_210_rocm_no_reinstall(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""A gfx110X-all pin (NOT in the 2.11 allowlist) over a correct 2.10+rocm
wheel must NOT be flagged stale -- the install path uses the default <2.11
specs for that arch, so re-flagging would reinstall-loop on every update."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n"
env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx110X-all"}
with patch.dict(stack_mod.os.environ, env, clear = False):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
_ensure_rocm_torch()
# has_hip_torch True + no mismatch -> torch must NOT be reinstalled.
assert not any(
any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list
)
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False)
@patch.object(stack_mod, "_has_rocm_gpu", return_value = True)
@patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2))
def test_gfx_pin_over_generic_rocm211_reinstalls(
self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try
):
"""A gfx1151 pin over a GENERIC (two-part +rocm7.2) 2.11 wheel must reinstall
the AMD per-arch wheel -- even though both are torch 2.11, the generic wheel
is not the per-arch build the user pinned (Strix stays off the generic wheel)."""
mock_probe = MagicMock()
mock_probe.returncode = 0
mock_probe.stdout = b"7.2.12345|2.11.0+rocm7.2\n"
env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"}
with patch.dict(stack_mod.os.environ, env, clear = False):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
with patch("os.path.isdir", return_value = True):
with patch("subprocess.run", return_value = mock_probe):
with patch.object(
stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError
):
_ensure_rocm_torch()
torch_call = str(mock_pip.call_args_list[0])
assert "gfx1151" in torch_call
assert "torch>=2.11.0,<2.12.0" in torch_call
def test_radeon_url_not_classified_as_pip_rocm_family(self):
"""A repo.radeon.com find-links dir (leaf rocm-rel-7.2.1) starts with "rocm" but is
NOT a pip --index-url ROCm family: it must route to the verbatim path, not a
--index-url reinstall that fails against a find-links listing."""
leaf_f = stack_mod._is_pip_rocm_family_leaf
# Real pip ROCm families (download.pytorch.org/whl/rocmX.Y, repo.amd.com gfx).
assert leaf_f("rocm7.2") is True
assert leaf_f("rocm6.4") is True
assert leaf_f("gfx120x-all") is True
assert leaf_f("gfx1151") is True
# A bare rocm<digits> (no minor) is still an exact family.
assert leaf_f("rocm7") is True
# A Radeon find-links dir leaf, a custom mirror, cpu and cuda are NOT pip rocm.
assert leaf_f("rocm-rel-7.2.1") is False
assert leaf_f("simple") is False
assert leaf_f("current") is False
assert leaf_f("cpu") is False
assert leaf_f("cu128") is False
# A rocm<digit>-SUFFIX private mirror shares the family prefix but is a custom pin
# the verbatim path owns: a ^rocm\d PREFIX match would wrongly treat it as a
# --index-url family. Match EXACTLY.
assert leaf_f("rocm7.2-private") is False
assert leaf_f("rocm7-current") is False
assert leaf_f("rocm7.2.1") is False # two-part local suffix -> custom, not rocm7.2
radeon = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1"
pip_rocm = "https://download.pytorch.org/whl/rocm7.2"
amd_gfx = "https://repo.amd.com/rocm/whl/gfx120X-all"
def _classify(url, fn):
with patch.dict(stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
return fn()
rocm_fn = stack_mod._explicit_rocm_torch_index_url
unk_fn = stack_mod._explicit_unknown_family_torch_index_url
# Real pip rocm/gfx pins ARE a ROCm family (reinstallable via --index-url) and
# are NOT "unknown".
assert _classify(pip_rocm, rocm_fn) == pip_rocm
assert _classify(amd_gfx, rocm_fn) == amd_gfx
assert _classify(pip_rocm, unk_fn) is None
assert _classify(amd_gfx, unk_fn) is None
# The Radeon find-links URL is NOT a pip ROCm family (so _ensure_rocm_torch skips
# it) and IS unknown, so the family repair helpers leave it alone.
assert _classify(radeon, rocm_fn) is None
assert _classify(radeon, unk_fn) == radeon
# A rocm<digit>-suffix private mirror routes the same way: NOT a pip rocm family,
# IS an unknown-family (verbatim) pin.
suffixed = "https://co.internal/whl/rocm7.2-private"
assert _classify(suffixed, rocm_fn) is None
assert _classify(suffixed, unk_fn) == suffixed
@patch.object(stack_mod, "pip_install")
def test_ensure_cpu_torch_broken_probe_reinstalls(self, mock_pip):
"""_ensure_cpu_torch: torch present but unimportable (probe exit != 0) under an
explicit CPU pin must reinstall from the pin, not return -- the base update does
not repair a broken installed torch, so returning would strand it (Codex P2)."""
mock_probe = MagicMock()
mock_probe.returncode = 1 # torch present but cannot import
mock_probe.stdout = b""
env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/cpu"}
with patch.dict(stack_mod.os.environ, env, clear = False):
stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None)
with patch("subprocess.run", return_value = mock_probe):
with patch.object(stack_mod, "NO_TORCH", False):
stack_mod._ensure_cpu_torch()
assert mock_pip.call_count == 1
assert "https://mirror.local/cpu" in str(mock_pip.call_args)
@patch.object(stack_mod, "IS_WINDOWS", False)
@patch.object(stack_mod, "pip_install_try", return_value = True)
@patch.object(stack_mod, "pip_install")
@ -732,7 +1048,7 @@ class TestEnsureRocmTorch:
mock_pip.assert_not_called()
# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard
# TEST: install_python_stack.py -- torch-index MARKER mechanism (PR #6692)
class TestHasRocmGpuKfdVendorGuard:
@ -1716,9 +2032,8 @@ class TestDetectWindowsGfxArch:
assert result == "gfx1200"
def test_returns_arch_on_crash_with_gcnarchname_in_output(self):
# Regression #6043: hipinfo may crash (0xC0000005 on RDNA 4) after
# printing gcnArchName. Accept the arch whenever gcnArchName is in
# stdout, regardless of exit code (previously a CPU fallback).
# Regression #6043: hipinfo may crash (0xC0000005 on RDNA 4) after printing
# gcnArchName. Accept the arch whenever gcnArchName is in stdout, any exit code.
mock_result = MagicMock()
mock_result.returncode = -1073741819 # 0xC0000005 STATUS_ACCESS_VIOLATION
mock_result.stdout = b"gcnArchName : gfx1200\nsome other line\n"
@ -2360,6 +2675,49 @@ class TestWindowsRocmTorchaoGuard:
assert not any("torchao" in arg for arg in installed_specs)
class TestProgressStepCountMatchesTotal:
"""The progress bar must reach exactly _TOTAL: every _progress() step is counted in
base_total. Regression for a repair step added without incrementing base_total,
which pushed _STEP past _TOTAL (Codex P2)."""
def _run_stack(self, tmp_path, *, is_windows, is_macos, is_mac_arm):
unstructured_plugin = tmp_path / "unstructured"
github_plugin = tmp_path / "github"
unstructured_plugin.mkdir()
github_plugin.mkdir()
sub = MagicMock()
sub.returncode = 0
sub.stdout = ""
with (
patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}),
patch.object(stack_mod, "IS_WINDOWS", is_windows),
patch.object(stack_mod, "IS_MACOS", is_macos),
patch.object(stack_mod, "IS_MAC_ARM", is_mac_arm),
patch.object(stack_mod, "NO_TORCH", False),
patch.object(stack_mod, "_rocm_windows_torch_installed", False),
patch.object(stack_mod, "_bootstrap_uv", return_value = False),
patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = False),
patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True),
patch.object(stack_mod, "_repair_bad_anyio"),
patch.object(stack_mod, "_ensure_cuda_torch"),
patch.object(stack_mod, "_ensure_rocm_torch"),
patch.object(stack_mod, "_ensure_cpu_torch"),
patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin),
patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin),
patch.object(stack_mod.subprocess, "run", return_value = sub),
):
assert stack_mod.install_python_stack() == 0
return stack_mod._STEP, stack_mod._TOTAL
def test_windows_progress_reaches_total(self, tmp_path):
step, total = self._run_stack(tmp_path, is_windows = True, is_macos = False, is_mac_arm = False)
assert step == total, f"Windows progress {step} != total {total} (final step uncounted)"
def test_linux_progress_reaches_total(self, tmp_path):
step, total = self._run_stack(tmp_path, is_windows = False, is_macos = False, is_mac_arm = False)
assert step == total, f"Linux progress {step} != total {total}"
# TEST: worker.py -- Windows ROCm patches (source-level checks)
@ -2846,6 +3204,24 @@ class TestStrixRocm71Override:
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
assert "TORCH_CONSTRAINT" in source and "2.11" in source
def test_torch_constraint_211_matches_leaf_not_whole_url(self):
"""The 2.11 constraint case must match the index LEAF, not the whole URL.
A custom UNSLOTH_PYTORCH_MIRROR whose base path contains a gfx/rocm7.2
segment (e.g. https://mirror.local/gfx-cache) with a cu*/cpu family must
not be pushed to the torch 2.11 line -- same leaf-only reasoning the
UNSLOTH_TORCH_BACKEND classification uses.
"""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
# The 2.11 constraint block must switch on $_torch_index_leaf, not the full
# $TORCH_INDEX_URL (a */gfx* match false-positives on a mirror base path). Only the
# _grouped_mm-bug gfx families (gfx120X-all / gfx1151 / gfx1150) are pushed to 2.11;
# a bare gfx* would also floor gfx110X-all/gfx90a/gfx908, left bare on purpose.
assert 'case "$_torch_index_leaf" in\n rocm7.2|gfx120x-all|gfx1151|gfx1150)' in source, (
"the torch>=2.11 constraint must match the specific gfx leaves that need "
"it (rocm7.2|gfx120x-all|gfx1151|gfx1150), not a bare gfx* or the whole URL"
)
def test_amd_rocm_mirror_env_var_respected(self):
"""install.sh must honour UNSLOTH_AMD_ROCM_MIRROR for air-gapped installs."""
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
@ -2934,9 +3310,9 @@ class TestServerStartupRocmFixes:
assert '"BNB_ROCM_VERSION" not in os.environ' in source
# ── hipInfo.exe PATH prepend (bitsandbytes arch-probe fix) ────────────────
# bnb's get_rocm_gpu_arch() runs hipinfo.exe via PATH at import; the AMD
# wheel ships it in venv Scripts (on PATH only for activated venvs), so
# without the prepend bnb logs "[WinError 2]" when launched directly.
# bnb's get_rocm_gpu_arch() runs hipinfo.exe via PATH at import; the AMD wheel ships it
# in venv Scripts (on PATH only for activated venvs), so without the prepend bnb logs
# "[WinError 2]" when launched directly.
def test_main_py_prepends_hipinfo_dir_to_path(self):
"""main.py must make hipInfo.exe resolvable before bnb imports."""
@ -3178,11 +3554,10 @@ class TestRocmGfxForwarding:
assert '$HelperReleaseRepo = "unslothai/llama.cpp"' in source
assert "$HelperReleaseRepo = if (" not in source
# The text pins above guard the literal. The tests below *execute* the real
# routing line from setup.sh / setup.ps1 and assert the resolved release repo,
# so a refactor that reintroduces a conditional (or a ggml-org branch) is still
# caught. Inputs are varied -- CPU-only, inferred/forwarded gfx, usable NVIDIA --
# to prove no host slips back onto ggml-org. No GPU, no tooling, no network.
# The text pins above guard the literal. The tests below execute the real routing line
# from setup.sh / setup.ps1 and assert the resolved release repo, so a refactor that
# reintroduces a conditional (or a ggml-org branch) is still caught. Inputs vary
# (CPU-only, inferred/forwarded gfx, usable NVIDIA) to prove no host hits ggml-org.
@staticmethod
def _resolve_setup_sh_repo(
@ -3277,8 +3652,8 @@ class TestRocmGfxForwarding:
# TEST: _pick_rocm_gfx_target -- visible-device selection from rocminfo output.
# Honours CUDA/HIP_VISIBLE_DEVICES so a mixed-arch host installs the prebuilt
# for the selected GPU, not GPU 0.
# Honours CUDA/HIP_VISIBLE_DEVICES so a mixed-arch host installs the prebuilt for the
# selected GPU, not GPU 0.
_pick_rocm_gfx_target = prebuilt_mod._pick_rocm_gfx_target
@ -3454,7 +3829,11 @@ class TestWslRerouteNvidiaGuard:
source = _INSTALL_SH_PATH.read_text(encoding = "utf-8")
start = source.find("_maybe_reroute_strixhalo_to_2404()")
assert start != -1
body = source[start : start + 1200]
# Slice the WHOLE function body (to its closing brace at column 0), not a
# fixed-length window: preamble growth must not push the signals out of view.
end = source.find("\n}", start)
assert end != -1
body = source[start:end]
nv = body.find("_has_usable_nvidia_gpu")
wmi = body.find("_wsl_amd_gpu_name")
assert nv != -1, "reroute must consult _has_usable_nvidia_gpu before deciding to reroute"

View file

@ -0,0 +1,114 @@
#!/usr/bin/env pwsh
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Unit test for studio/setup.ps1's pinned-torch-index stale-venv helpers
# (Test-RocmGfx211Leaf, Test-CudaFamilyLeaf, Get-RocmPinStaleTags). Pure helpers,
# AST-extracted and run in-process. Mirrors the Python _rocm_pin_family_mismatch /
# _is_cuda_family_leaf tests.
# Run: pwsh -NoProfile -File tests/studio/test_setup_pin_stale.ps1
$ErrorActionPreference = "Stop"
$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1")
$setupPath = (Resolve-Path $setupPath).Path
# --- Parse setup.ps1 (also serves as a syntax gate) and extract the helpers ---
$tokens = $null; $errors = $null
$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors)
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" }
foreach ($name in @("Test-RocmGfx211Leaf", "Test-RocmKnown211Version", "Test-CudaFamilyLeaf", "Get-RocmPinStaleTags")) {
$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 setup.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++ }
}
# A pinned gfx/rocm index is stale when Expected != Installed.
function IsStale($leaf, $ver) {
$t = Get-RocmPinStaleTags -PinLeaf $leaf -TorchVersion $ver
return $t.Expected -ne $t.Installed
}
Write-Host "Test-RocmGfx211Leaf (the 2.11 gfx allowlist)"
Check "gfx1151 -> true" (Test-RocmGfx211Leaf "gfx1151")
Check "gfx1150 -> true" (Test-RocmGfx211Leaf "gfx1150")
Check "gfx120x-all -> true" (Test-RocmGfx211Leaf "gfx120x-all")
Check "gfx110x-all -> false" (-not (Test-RocmGfx211Leaf "gfx110x-all"))
Check "gfx90a -> false" (-not (Test-RocmGfx211Leaf "gfx90a"))
Check "gfx908 -> false" (-not (Test-RocmGfx211Leaf "gfx908"))
Write-Host "Test-CudaFamilyLeaf (^cu[0-9])"
Check "cu118 -> true" (Test-CudaFamilyLeaf "cu118")
Check "cu128 -> true" (Test-CudaFamilyLeaf "cu128")
Check "cu130 -> true" (Test-CudaFamilyLeaf "cu130")
Check "custom -> false" (-not (Test-CudaFamilyLeaf "custom"))
Check "current -> false" (-not (Test-CudaFamilyLeaf "current"))
Check "cpu -> false" (-not (Test-CudaFamilyLeaf "cpu"))
Check "empty -> false" (-not (Test-CudaFamilyLeaf ""))
Write-Host "Get-RocmPinStaleTags (mirror of _rocm_pin_family_mismatch)"
# Exact rocm version comparison.
Check "rocm7.2 pin + 2.11.0+rocm7.2 -> not stale" (-not (IsStale "rocm7.2" "2.11.0+rocm7.2"))
Check "rocm7.2 pin + 2.10.0+rocm6.4 -> stale" (IsStale "rocm7.2" "2.10.0+rocm6.4")
Check "rocm6.4 pin + 2.10.0+rocm6.4 -> not stale" (-not (IsStale "rocm6.4" "2.10.0+rocm6.4"))
# rocm7.2 is a KNOWN-2.11 index. A +rocm7.2 wheel whose RELEASE drifted off 2.11 shares
# the tag but violates the spec -> stale (mirror of _rocm_pin_family_mismatch).
Check "rocm7.2 pin + 2.12.0+rocm7.2 -> stale" (IsStale "rocm7.2" "2.12.0+rocm7.2")
Check "rocm7.2 pin + 2.13.0+rocm7.2 -> stale" (IsStale "rocm7.2" "2.13.0+rocm7.2")
Check "rocm7.2 pin + 2.11.5+rocm7.2 -> not stale" (-not (IsStale "rocm7.2" "2.11.5+rocm7.2"))
# An UNKNOWN newer rocm (off the 2.11 allowlist) isn't floored, so a matching version at
# any release line is NOT stale on this exact-compare branch.
Check "rocm8.0 pin + 2.12.0+rocm8.0 -> not stale" (-not (IsStale "rocm8.0" "2.12.0+rocm8.0"))
# An untagged (no +rocm) wheel never satisfies a ROCm pin -> always stale.
Check "rocm7.2 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm7.2" "2.10.0")
Check "rocm7.2 pin + 2.11.0 (untagged) -> stale" (IsStale "rocm7.2" "2.11.0")
Check "rocm6.4 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm6.4" "2.10.0")
# 2.11-allowlist gfx pin: per-arch (three-part) wheel is satisfied, generic is stale.
Check "gfx1151 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx1151" "2.11.0+rocm7.13.0"))
Check "gfx1150 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx1150" "2.11.0+rocm7.13.0"))
Check "gfx120x-all pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx120x-all" "2.11.0+rocm7.13.0"))
Check "gfx1151 pin + 2.11.0+rocm7.2 (generic) -> stale" (IsStale "gfx1151" "2.11.0+rocm7.2")
Check "gfx1151 pin + 2.10.0+rocm6.4 -> stale" (IsStale "gfx1151" "2.10.0+rocm6.4")
# Non-2.11 gfx pin (gfx110X-all/gfx90a/gfx908): a valid <2.11 wheel is NOT stale.
Check "gfx110x-all pin + 2.10.0+rocm6.4 -> not stale" (-not (IsStale "gfx110x-all" "2.10.0+rocm6.4"))
Check "gfx90a pin + 2.10.0+rocm6.3 -> not stale" (-not (IsStale "gfx90a" "2.10.0+rocm6.3"))
Check "gfx908 pin + 2.10.0+rocm7.0 -> not stale" (-not (IsStale "gfx908" "2.10.0+rocm7.0"))
Check "gfx110x-all pin + 2.11.0+rocm7.2 -> stale" (IsStale "gfx110x-all" "2.11.0+rocm7.2")
# Non-2.11 gfx pin over an untagged wheel: never satisfies the pin -> stale, so the
# explicit ROCm index is applied even when torch is already <2.11.
Check "gfx110x-all pin + 2.10.0 (untagged) -> stale" (IsStale "gfx110x-all" "2.10.0")
Check "gfx90a pin + 2.10.0 (untagged) -> stale" (IsStale "gfx90a" "2.10.0")
# Capital gfx120X-all is lowercased by Get-TorchIndexLeaf before this helper, so the
# 2.11-allowlist branch fires and a generic/untagged wheel is stale.
Check "gfx120x-all pin + 2.11.0+rocm7.2 (generic) -> stale" (IsStale "gfx120x-all" "2.11.0+rocm7.2")
Check "gfx120x-all pin + 2.10.0 (untagged) -> stale" (IsStale "gfx120x-all" "2.10.0")
# Major-only rocm pin (rocm7): majors compared alone; mirrors _rocm_pin_family_mismatch.
Check "rocm7 pin + 2.10.0+rocm6.4 -> stale" (IsStale "rocm7" "2.10.0+rocm6.4")
Check "rocm7 pin + 2.11.0+rocm7.2 -> not stale" (-not (IsStale "rocm7" "2.11.0+rocm7.2"))
Check "rocm7 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "rocm7" "2.11.0+rocm7.13.0"))
Check "rocm7 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm7" "2.10.0")
Check "rocm7 pin + 2.10.0+rocm (unreadable) -> not stale" (-not (IsStale "rocm7" "2.10.0+rocm"))
Write-Host "Test-RocmKnown211Version + KNOWN-2.11 fallback (rocm7.2 only; no speculative rocm7.3)"
Check "rocm7.2 -> known 2.11" (Test-RocmKnown211Version -Major 7 -Minor 2)
Check "rocm7.1 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 1))
Check "rocm7.3 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 3))
Check "rocm8.0 -> not known" (-not (Test-RocmKnown211Version -Major 8 -Minor 0))
# Unreadable-installed fallback: a rocm7.3 pin (unknown -> <2.11 line) over a <2.11 +rocm
# wheel with an unreadable version is NOT stale; rocm7.2 (KNOWN-2.11) over the same wheel
# IS stale (#2534 alignment).
Check "rocm7.3 pin + 2.10.0+rocm (unreadable ver) -> not stale" (-not (IsStale "rocm7.3" "2.10.0+rocm"))
Check "rocm7.2 pin + 2.10.0+rocm (unreadable ver) -> stale" (IsStale "rocm7.2" "2.10.0+rocm")
Write-Host ""
if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
Write-Host "All checks passed" -ForegroundColor Green

View file

@ -15,7 +15,7 @@ $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")) {
foreach ($name in @("ConvertTo-TorchFlavorTag", "Get-ExpectedTorchFlavorTag", "Trim-IndexPathSlashes", "Redact-InstallOutput")) {
$fn = $ast.FindAll({ param($n)
$n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name
}, $true)
@ -49,6 +49,19 @@ Check "mirror cu130 leaf -> cu130" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl
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 "Trim-IndexPathSlashes (install.ps1 parity: path-only, token-preserving)"
Check "double path slash collapsed" ((Trim-IndexPathSlashes "https://h/whl/cu128//") -eq "https://h/whl/cu128")
Check "single trailing slash trimmed" ((Trim-IndexPathSlashes "https://h/whl/cu128/") -eq "https://h/whl/cu128")
Check "query token slash preserved" ((Trim-IndexPathSlashes "https://h/whl/cu128?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/")
Check "path slash trimmed, query kept" ((Trim-IndexPathSlashes "https://h/whl/cu128//?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/")
Write-Host "Redact-InstallOutput (install.ps1 parity: credential redaction)"
Check "userinfo redacted" ((Redact-InstallOutput "ERROR https://alice:s3cr3t@download.pytorch.org/whl/cu128") -eq "ERROR https://<redacted>@download.pytorch.org/whl/cu128")
Check "query value redacted" ((Redact-InstallOutput "https://host/whl/cu128?token=abcd1234&channel=beta") -eq "https://host/whl/cu128?token=<redacted>&channel=<redacted>")
Check "fragment token redacted" ((Redact-InstallOutput "ERROR https://mirror.local/whl/cu128#token=SECRET123 (403)") -eq "ERROR https://mirror.local/whl/cu128#<redacted> (403)")
Check "bare hash comment untouched" ((Redact-InstallOutput "# retrying with --no-cache-dir") -eq "# retrying with --no-cache-dir")
Check "plain line untouched" ((Redact-InstallOutput "Resolved 42 packages in 1.2s") -eq "Resolved 42 packages in 1.2s")
Write-Host ""
if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
Write-Host "All checks passed" -ForegroundColor Green

View file

@ -0,0 +1,78 @@
#!/usr/bin/env pwsh
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Unit tests for setup.ps1's torch-index pin-hardening helpers: Trim-IndexPathSlashes
# (path-only slash trim, token-preserving), Redact-InstallOutput (credential redaction of
# captured install logs), Get-TorchIndexLeaf (ALL trailing slashes stripped) and
# Test-PipRocmFamilyLeaf (rocm7. is a custom pin, not a family). Pure helpers, AST-extracted
# and run in-process. Run: pwsh -NoProfile -File tests/studio/test_torch_index_pin_hardening.ps1
$ErrorActionPreference = "Stop"
$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1")
$setupPath = (Resolve-Path $setupPath).Path
$setupText = Get-Content -Raw $setupPath
# --- Parse setup.ps1 (also a syntax gate) and extract the pure helpers ---
$tokens = $null; $errors = $null
$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors)
if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" }
foreach ($name in @("Trim-IndexPathSlashes", "Redact-InstallOutput", "Get-TorchIndexLeaf", "Test-PipRocmFamilyLeaf")) {
$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 setup.ps1, found $($fn.Count)" }
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 "Trim-IndexPathSlashes (path-only, token-preserving)"
Check "double path slash collapsed" ((Trim-IndexPathSlashes "https://h/whl/cu128//") -eq "https://h/whl/cu128")
Check "single trailing slash trimmed" ((Trim-IndexPathSlashes "https://h/whl/cu128/") -eq "https://h/whl/cu128")
Check "no slash unchanged" ((Trim-IndexPathSlashes "https://h/whl/cu128") -eq "https://h/whl/cu128")
Check "query token slash preserved" ((Trim-IndexPathSlashes "https://h/whl/cu128?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/")
Check "path slash trimmed, query kept" ((Trim-IndexPathSlashes "https://h/whl/cu128//?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/")
Check "fragment slash preserved" ((Trim-IndexPathSlashes "https://h/whl/cu128#anchor/") -eq "https://h/whl/cu128#anchor/")
Write-Host "Redact-InstallOutput (credential redaction)"
Check "userinfo redacted" ((Redact-InstallOutput "ERROR https://alice:s3cr3t@download.pytorch.org/whl/cu128") -eq "ERROR https://<redacted>@download.pytorch.org/whl/cu128")
Check "bare-token@ redacted" ((Redact-InstallOutput "fetch https://ghp_deadbeef@host/whl/cu128 failed") -eq "fetch https://<redacted>@host/whl/cu128 failed")
Check "single query value redacted" ((Redact-InstallOutput "url https://host/whl/cu128?token=abcd1234 unreachable") -eq "url https://host/whl/cu128?token=<redacted> unreachable")
Check "multiple query values redacted" ((Redact-InstallOutput "https://host/whl/cu128?token=abcd1234&channel=beta") -eq "https://host/whl/cu128?token=<redacted>&channel=<redacted>")
Check "fragment token redacted" ((Redact-InstallOutput "ERROR https://mirror.local/whl/cu128#token=SECRET123 (403)") -eq "ERROR https://mirror.local/whl/cu128#<redacted> (403)")
Check "query and fragment both redacted" ((Redact-InstallOutput "https://host/whl/cu128?token=abc#sig=xyz done") -eq "https://host/whl/cu128?token=<redacted>#<redacted> done")
Check "bare hash comment untouched" ((Redact-InstallOutput "# retrying with --no-cache-dir") -eq "# retrying with --no-cache-dir")
Check "plain line untouched" ((Redact-InstallOutput "Resolved 42 packages in 1.2s") -eq "Resolved 42 packages in 1.2s")
$leak = Redact-InstallOutput "https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET"
Check "no secret substring survives" (($leak -notmatch "s3cr3t") -and ($leak -notmatch "SUPERSECRET") -and ($leak -notmatch "ALSOSECRET"))
Write-Host "Get-TorchIndexLeaf (ALL trailing slashes stripped)"
Check "double slash cu128 -> cu128" ((Get-TorchIndexLeaf "https://m/whl/cu128//") -eq "cu128")
Check "triple slash rocm7.2 -> rocm7.2" ((Get-TorchIndexLeaf "https://m/whl/rocm7.2///") -eq "rocm7.2")
Check "double slash + token -> cu128" ((Get-TorchIndexLeaf "https://m/whl/cu128//?token=x") -eq "cu128")
Check "single slash cu128 -> cu128" ((Get-TorchIndexLeaf "https://m/whl/cu128/") -eq "cu128")
Write-Host "Test-PipRocmFamilyLeaf (rocm7. is a custom pin, not a family)"
Check "rocm7 family" (Test-PipRocmFamilyLeaf "rocm7")
Check "rocm7.2 family" (Test-PipRocmFamilyLeaf "rocm7.2")
Check "gfx1151 family" (Test-PipRocmFamilyLeaf "gfx1151")
Check "rocm7. trailing-dot NOT family" (-not (Test-PipRocmFamilyLeaf "rocm7."))
Check "rocm.7 leading-dot NOT family" (-not (Test-PipRocmFamilyLeaf "rocm.7"))
Check "rocm7.2.1 two-dot NOT family" (-not (Test-PipRocmFamilyLeaf "rocm7.2.1"))
Check "rocm7.2-private NOT family" (-not (Test-PipRocmFamilyLeaf "rocm7.2-private"))
Check "cu128 NOT family" (-not (Test-PipRocmFamilyLeaf "cu128"))
Write-Host "Fast-Install pinned-install env scrub (source assertion)"
# The pip fallback honours PIP_*; PIP_NO_INDEX=1 would make it ignore the pinned --index-url
# and PIP_INDEX_URL would replace it, so both must be scrubbed for a pinned install.
Check "PIP_NO_INDEX scrubbed" ($setupText -match "'PIP_NO_INDEX'")
Check "PIP_INDEX_URL scrubbed" ($setupText -match "'PIP_INDEX_URL'")
Write-Host ""
if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 }
Write-Host "All checks passed" -ForegroundColor Green