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

Installer: preserve the previous torch release across every flavor and vendor on re-runs (#7250)

Browse source 
* install: preserve the previous torch release across every flavor and vendor

A re-run of curl | sh over an existing install was supposed to keep the
user's validated torch release, but the pin required the old build's
local flavor tag to match the freshly chosen index leaf. That gate was
wrong in practice: a PyPI-sourced torch reports a BARE version (on Linux
the PyPI wheel IS a CUDA build), which classified as cpu and never
matched a cu leaf, so a healthy 2.10 on a cu130 host was silently moved
to 2.11 (reproduced end to end); the same happened for any flavor drift
such as cu128 to cu130 after a driver upgrade, and AMD ROCm leaves were
excluded from preservation entirely.

The rule is now release-based and flavor-agnostic: the probed previous
release is pinned whenever it sits inside the final constraint window,
and the pin installs from the freshly chosen index, so the flavor always
follows the machine (NVIDIA cu*, AMD rocm/gfx, Intel/CPU, mac) while the
release follows the user. The pin is evaluated AFTER every index and
constraint decision including the Strix reroute, so raised floors
(rocm7.2 / Strix gfx need torch 2.11 for the _grouped_mm fix) correctly
reject an older release and win. UNSLOTH_TORCH_UPGRADE=1 still opts out,
out-of-window releases are never kept, and probe noise never becomes a
pin.

The kept-release install with its range fallback (for indexes that do
not carry the exact release) is factored into
_install_torch_default_index and used by every --default-index torch
path: the default NVIDIA/CPU/mac path and all three ROCm-index
fallbacks, which previously bypassed the fallback. The Radeon-repo
direct-wheel path keeps its curated per-arch wheel set (those wheels are
already exact-pinned per rocm release).

Platform coverage: install.sh serves Linux, WSL (including the WoA
fallback), and macOS for all vendors; native Windows install.ps1 still
caps at <2.11.0 everywhere, so the silent 2.10-to-2.11 move cannot occur
there (2.11 alignment is a separate follow-up).

Verified: 35-check unit suite rewritten to the new spec (any-flavor
keep, floor rejection, noise, window edges, opt-out, wiring including
pin-after-reroute and helper coverage); end-to-end matrix against
sandboxed UNSLOTH_STUDIO_HOME installs on a cu130 host covering PyPI
bare, cu128 drift, cu130 same-flavor, out-of-window 2.3, the upgrade
opt-out, the hidden-GPU cpu leaf, and a fresh-install control.

* install: honor the kept torch release on the Radeon direct-wheel path

The Radeon repo path installs an explicit wheel trio selected by
_pick_radeon_wheel, bypassing --default-index, so the kept-release pin
only took effect when the listing failed and the install fell back to
the ROCm index. On a re-run over an in-window Radeon install the trio
search started at the newest common minor and silently moved the user
forward (2.9 to 2.10 whenever the repo offered both).

The trio search now starts at the kept release's minor when
_PREV_TORCH_PIN is set and the listing still offers a torch wheel for
that minor. Radeon wheels are patch-curated per rocm release, so the
minor is the unit of preservation there; the raised rocm7.2 / Strix
floors still win because the pin is window-checked against the final
constraint before this point, and gaps keep the existing downward
search / ROCm-index fallback.

Verified with a simulated listing carrying both a 2.9 and a 2.10 trio:
no pin selects the 2.10 trio, a kept 2.9 release selects the matched
2.9 / 0.24 / 2.9 trio, and an unavailable minor degrades to the newest
trio. Added a structural wiring check to test_previous_torch_pin.sh
(now 36 checks).

* install: tighten comments in the torch preservation paths

* install: exact kept release on the Radeon path, pin fallback in ROCm repairs

The minor-level clamp on the Radeon direct-wheel path still allowed
patch drift (a kept 2.10.0 could become 2.10.1 when the listing carried
both) and the downward gap search could settle below the kept minor,
both breaking the exact preservation guarantee the other vendor paths
honor. The kept release now gets an exact-first trio attempt before the
newest-trio search: pick the kept patch (else the newest patch of the
kept minor, for listings that pruned the exact patch) together with the
paired torchvision/torchaudio wheels for that minor. Any gap warns and
falls back to the unchanged newest-trio search, mirroring
_install_torch_default_index, so a rerun installs either the kept
release or the same set a fresh install would choose, never something
in between.

The two ROCm torch repair sites (torch overwritten by dependency
resolution, on the migrated and fresh paths) installed TORCH_CONSTRAINT
directly, so a pinned release missing from the generic ROCm index would
abort the rerun instead of falling back. Both now route through
_install_torch_default_index, which passes extra uv args through
(--force-reinstall) and clears the pin once the fallback fires so later
paths stay consistent.

Verified against synthetic listings: both patches listed keeps exactly
2.10.0; a kept minor missing vision/audio warns and yields the newest
complete trio rather than a silent undercut; a pruned patch stays on
the kept minor; no pin keeps the existing newest-trio behavior. Unit
suite now 39 checks, all passing.

* install: never pin nightly/dev/source torch builds on a rerun

A survey of published torch version strings (PyPI bare, +cpu, +cu116
through +cu132, +rocmX.Y and +rocmX.Y.Z, +xpu, nightly .devYYYYMMDD,
source a0+git, rc tags) showed one gap: nightly, dev, rc, and source
builds passed the loose release-shape check, producing a pin such as
torch==2.11.0.dev20250704 that no stable index carries. The range
fallback rescued the install, but it printed "keeping it" and then
burned a doomed resolve first. The base must now be a plain numeric
X.Y[.Z] release, so those builds skip the pin and go straight to the
newest supported release.

Added unit checks for +xpu and three-component +rocm7.2.1 tags (both
already preserved correctly) and for nightly, a0 source, and rc builds
(never pinned). Suite now 44 checks, all passing.

* install: pair kept-release companions, protect the flavor repair, note substitutions

Three fixes from a 12-way review pass over the preservation work:

The kept-release install left torchvision and torchaudio unconstrained
next to the exact torch pin. torchvision exact-pins its torch in wheel
metadata so it always paired correctly, but torchaudio no longer does:
a kept torch 2.9.0 on cu130 resolved torchaudio 2.11.0 (verified with
uv dry-runs). The helper now pairs both companions to the kept minor
(torchvision 0.minor+15, torchaudio 2.minor); if the index lacks the
paired set the existing range fallback fires. Verified resolving
correctly on cu130, cu126, and rocm6.4.

The wrong-flavor repair at the end of the install was the one remaining
default-index torch install outside the helper. It runs under set -e,
so a retained pin absent from the repair index (reachable when the
Radeon direct-wheel path installed the kept release and dependency
resolution later overwrote it) aborted the installer at the last step
instead of falling back. It now routes through the helper with its
reinstall flags passed through.

The Radeon kept-release path installed a same-series build silently
when the listing had pruned the exact patch; it now prints what it is
substituting.

Unit suite extended with wiring checks for all three (46 checks, all
passing).
This commit is contained in:
Daniel Han 2026-07-19 07:55:06 -07:00 • committed by GitHub
commit b3c0259cff
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 181 additions and 105 deletions
Download patch file Download diff file Expand all files Collapse all files

180
install.sh
View file

@ -2225,37 +2225,67 @@ _torch_release_in_window() {
echo "no"
}
# Whether a re-run should keep the previous venv's torch: echo "torch==X.Y.Z" when the
# probed previous version ($1) has a flavor tag matching the freshly chosen cu*/cpu index
# leaf ($2) AND sits inside the active constraint window ($3), else "". Re-running
# `curl | sh` rebuilds the venv for clean state, but a healthy torch the user already
# validated must not be silently moved to a newer release (2.10 -> 2.11); a flavor
# change (cpu <-> cuda, cu126 -> cu130) still installs the correct new build, rocm
# leaves keep their floors (rocm7.2 must land 2.11 for the Strix _grouped_mm fix), and
# a release outside the window (2.3.x manual install, 2.12.x manual upgrade) is never
# kept: the installer's own bounds win. Opt out with UNSLOTH_TORCH_UPGRADE=1 to get
# the newest release.
# Keep the previous venv's torch on a re-run: echo "torch==X.Y.Z" when the probed
# version ($1) is inside the active constraint window ($2), else "". The RELEASE is kept
# regardless of flavor tag; the pin installs from the freshly chosen index, so flavor
# follows the machine (cpu <-> cuda, cu126 -> cu130, PyPI bare -> +cu130) while the
# release follows the user. Gating on flavor was wrong: a PyPI torch reports a BARE
# version (on Linux the PyPI wheel IS CUDA), misclassified "cpu", so a healthy 2.10 on a
# cu130 host was moved to 2.11. Per-leaf floors still win (rocm7.2 / gfx >=2.11 for the
# Strix _grouped_mm fix, out-of-window manual installs) and are never pinned; the caller's
# _PREV_FALLBACK_CONSTRAINT installs the newest supported release when the index lacks the
# exact one. Opt out with UNSLOTH_TORCH_UPGRADE=1.
_previous_torch_pin() {
_ptp_ver="$1"
_ptp_leaf="$2"
_ptp_con="$3"
_ptp_con="$2"
[ -n "$_ptp_ver" ] || { echo ""; return; }
[ "${UNSLOTH_TORCH_UPGRADE:-0}" = "1" ] && { echo ""; return; }
case "$_ptp_leaf" in
cu[0-9]*|cpu) ;;
*) echo ""; return ;;
esac
_ptp_base="${_ptp_ver%%+*}"
# The base must look like a release (probe noise / garbage must never become a pin).
# Base must be a plain numeric release (X.Y[.Z]); probe noise and
# nightly/dev/source builds (2.11.0.dev20250704, 2.9.0a0) must never
# become a pin -- no stable index carries them, so pinning would only
# print "keeping it" and then burn a doomed resolve before falling back.
case "$_ptp_base" in
*[!0-9.]* | *..* | .* | *.) echo ""; return ;;
[0-9]*.[0-9]*) ;;
*) echo ""; return ;;
esac
[ "$(_torch_release_in_window "$_ptp_base" "$_ptp_con")" = "yes" ] || { echo ""; return; }
if [ "$(_torch_flavor_tag "$_ptp_ver")" = "$_ptp_leaf" ]; then
echo "torch==$_ptp_base"
echo "torch==$_ptp_base"
}
# Install torch from TORCH_INDEX_URL honoring a kept-release pin: with _PREV_TORCH_PIN
# set, TORCH_CONSTRAINT is the exact previous release; fall back to the supported range
# if the index lacks it (pruned mirror) rather than failing. Used by every --default-index
# path (NVIDIA cu*, AMD rocm/gfx fallbacks, cpu/mac, ROCm repairs) so preservation is
# uniform. Extra args (e.g. --force-reinstall) are passed through to uv.
_install_torch_default_index() {
if [ -n "$_PREV_TORCH_PIN" ]; then
# Pair the companions with the kept torch minor: torchaudio no longer
# exact-pins torch in its metadata, so leaving it unconstrained resolves
# a newer mismatched build (a kept torch 2.9.0 pulled torchaudio 2.11.0).
_itdi_base="${_PREV_TORCH_PIN#torch==}"
_itdi_minor="${_itdi_base#*.}"
_itdi_minor="${_itdi_minor%%.*}"
_itdi_tv="torchvision"
_itdi_ta="torchaudio"
case "$_itdi_base" in
2.*)
_itdi_tv="torchvision==0.$((_itdi_minor + 15)).*"
_itdi_ta="torchaudio==2.${_itdi_minor}.*"
;;
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"
TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT"
_PREV_TORCH_PIN=""
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL" "$@"
fi
else
echo ""
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL" "$@"
fi
}
@ -2561,21 +2591,6 @@ case "$_torch_index_leaf" in
cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
esac
# Re-run over an existing install: keep the previous venv's torch release instead of
# resolving the newest in range. The range stays in _PREV_FALLBACK_CONSTRAINT so the
# install can fall back when the exact release is not on the chosen index (custom
# mirrors may prune old wheels). Skipped for --no-torch (no previous probe runs).
_PREV_TORCH_PIN=""
_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT"
if [ "$SKIP_TORCH" = false ]; then
_prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$_torch_index_leaf" "$TORCH_CONSTRAINT")
if [ -n "$_prev_pin" ]; then
_PREV_TORCH_PIN="$_prev_pin"
TORCH_CONSTRAINT="$_prev_pin"
substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)"
fi
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".
@ -2660,6 +2675,23 @@ case "$TORCH_INDEX_URL" in
fi
;;
esac
# 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
# final one and a raised floor (rocm7.2 / Strix gfx) rejects an older release.
# _PREV_FALLBACK_CONSTRAINT keeps the range so the install can fall back when the exact
# release is not on the chosen index (mirrors may prune old wheels). Skipped for --no-torch.
_PREV_TORCH_PIN=""
_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT"
if [ "$SKIP_TORCH" = false ]; then
_prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$TORCH_CONSTRAINT")
if [ -n "$_prev_pin" ]; then
_PREV_TORCH_PIN="$_prev_pin"
TORCH_CONSTRAINT="$_prev_pin"
substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)"
fi
fi
_TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL")
if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then
_TAURI_TORCH_INDEX_FAMILY="radeon"
@ -2885,10 +2917,7 @@ if [ "$_MIGRATED" = 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)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
_install_torch_default_index --force-reinstall
fi
;;
esac
@ -2953,7 +2982,42 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_ta_ver=$(_extract_version "$_ta_whl" "torchaudio")
_radeon_versions_match=false
if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then
# Kept release (_PREV_TORCH_PIN) wins here too: pick its exact
# patch (else the newest patch of its minor) plus the paired
# vision/audio wheels. Any gap falls back to the newest-trio
# search below, mirroring _install_torch_default_index, so a
# rerun never drifts to another release nor below the kept one.
if [ -n "$_PREV_TORCH_PIN" ]; then
_prev_kept_base="${_PREV_TORCH_PIN#torch==}"
_prev_kept_minor="${_prev_kept_base#*.}"
_prev_kept_minor="${_prev_kept_minor%%.*}"
case "$_prev_kept_minor" in
''|*[!0-9]*) ;;
*)
_kept_torch=$(_pick_radeon_wheel "torch" "${_prev_kept_base}" 2>/dev/null) || _kept_torch=""
[ -z "$_kept_torch" ] && { _kept_torch=$(_pick_radeon_wheel "torch" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_torch=""; }
_kept_tv=$(_pick_radeon_wheel "torchvision" "0.$((_prev_kept_minor + 15))." 2>/dev/null) || _kept_tv=""
_kept_ta=$(_pick_radeon_wheel "torchaudio" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_ta=""
if [ -n "$_kept_torch" ] && [ -n "$_kept_tv" ] && [ -n "$_kept_ta" ]; then
_torch_whl=$_kept_torch
_tv_whl=$_kept_tv
_ta_whl=$_kept_ta
_tri_whl=""
_radeon_versions_match=true
# Say so when the listing pruned the exact patch
# and a same-series build is installed instead.
case "$(printf '%s' "${_kept_torch##*/}" | sed 's/%2[Bb]/+/g')" in
"torch-${_prev_kept_base}"[+-]*) ;;
*) substep "kept release ${_prev_kept_base} is not in the Radeon listing -- installing the closest 2.${_prev_kept_minor} series build instead" ;;
esac
else
substep "[WARN] Radeon repo lacks a complete wheel set for kept $_PREV_TORCH_PIN -- installing the newest compatible set instead" "$C_WARN"
fi
;;
esac
fi
if [ "$_radeon_versions_match" != true ] && \
[ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then
_torch_minor=${_torch_ver#*.}
_ta_minor=${_ta_ver#*.}
_tv_minor=${_tv_ver#*.}
@ -3011,9 +3075,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"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL"
_install_torch_default_index
else
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
# Pass explicit wheel URLs so the matched trio is
@ -3034,32 +3096,15 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
else
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL"
_install_torch_default_index
fi
else
substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL"
_install_torch_default_index
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
if [ -n "$_PREV_TORCH_PIN" ]; then
# Kept previous release: fall back to the supported range if the exact
# release is not resolvable from the chosen index (pruned mirror).
if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--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"
TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL"
fi
else
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL"
fi
_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
@ -3122,10 +3167,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; 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)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
_install_torch_default_index --force-reinstall
fi
;;
esac
@ -3164,9 +3206,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL" \
_install_torch_default_index \
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""

106
tests/sh/test_previous_torch_pin.sh
View file

@ -2,9 +2,12 @@
# 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 _previous_torch_pin, which keeps the previous
# venv's torch release on a re-run (curl | sh over an existing install) instead
# of silently moving the user to a newer release. Helpers are extracted from
# install.sh and sourced.
# venv's torch RELEASE on a re-run instead of moving the user to a newer one.
# The release is kept regardless of the old build's flavor tag (PyPI bare,
# +cuXXX, +rocm, +cpu): the pin installs from the freshly chosen index, so the
# flavor follows the machine while the release follows the user. Per-leaf
# windows still win (rocm7.2 / Strix floors, out-of-window manual installs).
# Helpers are extracted from install.sh and sourced.
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
@ -12,12 +15,9 @@ INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
# Extract _previous_torch_pin and its dependencies _torch_flavor_tag and
# _torch_release_in_window.
# Extract _previous_torch_pin and its dependency _torch_release_in_window.
_FUNC_FILE=$(mktemp)
{
sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_torch_release_in_window()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_previous_torch_pin()/,/^}/p' "$INSTALL_SH"
@ -37,37 +37,43 @@ assert_eq() {
unset UNSLOTH_TORCH_UPGRADE
echo "=== _previous_torch_pin: matching flavor keeps the release ==="
assert_eq "cu126 wheel on cu126 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
assert_eq "cu130 wheel on cu130 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')"
assert_eq "cpu wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cpu' 'cpu' 'torch>=2.4,<2.12.0')"
assert_eq "untagged wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0' 'cpu' 'torch>=2.4,<2.12.0')"
assert_eq "local suffix stripped" "torch==2.9.1" "$(_previous_torch_pin '2.9.1+cu128' 'cu128' 'torch>=2.4,<2.12.0')"
echo "=== _previous_torch_pin: in-window releases are kept, any flavor ==="
assert_eq "cu126 wheel" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu126' 'torch>=2.4,<2.12.0')"
assert_eq "cu130 wheel" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu130' 'torch>=2.4,<2.12.0')"
assert_eq "cpu wheel" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cpu' 'torch>=2.4,<2.12.0')"
assert_eq "PyPI bare version (CUDA build on Linux)" "torch==2.10.0" "$(_previous_torch_pin '2.10.0' 'torch>=2.4,<2.12.0')"
assert_eq "rocm wheel" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+rocm6.4' 'torch>=2.4,<2.11.0')"
assert_eq "rocm three-component tag" "torch==2.9.1" "$(_previous_torch_pin '2.9.1+rocm7.2.1' 'torch>=2.4,<2.12.0')"
assert_eq "Intel xpu wheel" "torch==2.9.0" "$(_previous_torch_pin '2.9.0+xpu' 'torch>=2.4,<2.12.0')"
assert_eq "local suffix stripped" "torch==2.9.1" "$(_previous_torch_pin '2.9.1+cu128' 'torch>=2.4,<2.12.0')"
echo "=== _previous_torch_pin: flavor change installs the new build ==="
assert_eq "cu126 wheel on cu130 leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu130' 'torch>=2.4,<2.12.0')"
assert_eq "cpu wheel on cu126 leaf" "" "$(_previous_torch_pin '2.10.0+cpu' 'cu126' 'torch>=2.4,<2.12.0')"
assert_eq "cu126 wheel on cpu leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cpu' 'torch>=2.4,<2.12.0')"
echo "=== _previous_torch_pin: rocm and unknown leaves never pin ==="
assert_eq "rocm7.2 leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'rocm7.2' 'torch>=2.4,<2.12.0')"
assert_eq "gfx leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'gfx120X-all' 'torch>=2.4,<2.12.0')"
assert_eq "unknown mirror leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'simple' 'torch>=2.4,<2.12.0')"
echo "=== _previous_torch_pin: raised floors reject older releases ==="
# rocm7.2 / Strix gfx leaves raise TORCH_CONSTRAINT to >=2.11.0 BEFORE the pin
# is evaluated, so an old 2.10 is out of window there and the floor wins.
assert_eq "old 2.10 vs rocm7.2 floor" "" "$(_previous_torch_pin '2.10.0+rocm7.1' 'torch>=2.11.0,<2.12.0')"
assert_eq "2.11 passes the rocm7.2 floor" "torch==2.11.0" "$(_previous_torch_pin '2.11.0+rocm7.2' 'torch>=2.11.0,<2.12.0')"
echo "=== _previous_torch_pin: probe noise never becomes a pin ==="
assert_eq "empty version" "" "$(_previous_torch_pin '' 'cu126' 'torch>=2.4,<2.12.0')"
assert_eq "garbage version" "" "$(_previous_torch_pin 'not-a-version' 'cpu' 'torch>=2.4,<2.12.0')"
assert_eq "traceback fragment" "" "$(_previous_torch_pin "ModuleNotFoundError: No module named 'torch'" 'cpu' 'torch>=2.4,<2.12.0')"
assert_eq "empty version" "" "$(_previous_torch_pin '' 'torch>=2.4,<2.12.0')"
assert_eq "garbage version" "" "$(_previous_torch_pin 'not-a-version' 'torch>=2.4,<2.12.0')"
assert_eq "traceback fragment" "" "$(_previous_torch_pin "ModuleNotFoundError: No module named 'torch'" 'torch>=2.4,<2.12.0')"
echo "=== _previous_torch_pin: nightly / dev / source builds never pin ==="
# No stable index carries these, so pinning would print "keeping it" and then
# burn a doomed resolve before the range fallback rescues the install.
assert_eq "nightly dev build" "" "$(_previous_torch_pin '2.11.0.dev20250704+cu128' 'torch>=2.4,<2.12.0')"
assert_eq "source build a0 tag" "" "$(_previous_torch_pin '2.9.0a0+gitabc1234' 'torch>=2.4,<2.12.0')"
assert_eq "release candidate" "" "$(_previous_torch_pin '2.11.0rc1+cu130' 'torch>=2.4,<2.12.0')"
echo "=== _previous_torch_pin: out-of-window releases never pin ==="
assert_eq "2.3.x below the cu floor" "" "$(_previous_torch_pin '2.3.1+cu118' 'cu118' 'torch>=2.4,<2.12.0')"
assert_eq "2.12.x above the cu ceiling" "" "$(_previous_torch_pin '2.12.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')"
assert_eq "floor boundary 2.4.0 kept" "torch==2.4.0" "$(_previous_torch_pin '2.4.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
assert_eq "ceiling-adjacent 2.11.x kept" "torch==2.11.1" "$(_previous_torch_pin '2.11.1+cu130' 'cu130' 'torch>=2.4,<2.12.0')"
assert_eq "cpu window excludes 2.11.x" "" "$(_previous_torch_pin '2.11.0+cpu' 'cpu' 'torch>=2.4,<2.11.0')"
assert_eq "mac floor excludes 2.5.x" "" "$(_previous_torch_pin '2.5.1' 'cpu' 'torch>=2.6,<2.11.0')"
assert_eq "malformed window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch')"
assert_eq "empty window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' '')"
assert_eq "2.3.x below the cu floor" "" "$(_previous_torch_pin '2.3.1+cu118' 'torch>=2.4,<2.12.0')"
assert_eq "2.12.x above the cu ceiling" "" "$(_previous_torch_pin '2.12.0+cu130' 'torch>=2.4,<2.12.0')"
assert_eq "floor boundary 2.4.0 kept" "torch==2.4.0" "$(_previous_torch_pin '2.4.0+cu126' 'torch>=2.4,<2.12.0')"
assert_eq "ceiling-adjacent 2.11.x kept" "torch==2.11.1" "$(_previous_torch_pin '2.11.1+cu130' 'torch>=2.4,<2.12.0')"
assert_eq "cpu window excludes 2.11.x" "" "$(_previous_torch_pin '2.11.0+cpu' 'torch>=2.4,<2.11.0')"
assert_eq "mac floor excludes 2.5.x" "" "$(_previous_torch_pin '2.5.1' 'torch>=2.6,<2.11.0')"
assert_eq "malformed window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'torch')"
assert_eq "empty window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' '')"
echo "=== _torch_release_in_window ==="
assert_eq "in window" "yes" "$(_torch_release_in_window '2.10.0' 'torch>=2.4,<2.12.0')"
@ -80,8 +86,8 @@ assert_eq "no ceiling -> no" "no" "$(_torch_release_in_window '2.10.0' 'tor
assert_eq "garbage minor -> no" "no" "$(_torch_release_in_window '2.x' 'torch>=2.4,<2.12.0')"
echo "=== _previous_torch_pin: UNSLOTH_TORCH_UPGRADE=1 opts out ==="
assert_eq "upgrade env set" "" "$(UNSLOTH_TORCH_UPGRADE=1 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
assert_eq "upgrade env 0" "torch==2.10.0" "$(UNSLOTH_TORCH_UPGRADE=0 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')"
assert_eq "upgrade env set" "" "$(UNSLOTH_TORCH_UPGRADE=1 _previous_torch_pin '2.10.0+cu126' 'torch>=2.4,<2.12.0')"
assert_eq "upgrade env 0" "torch==2.10.0" "$(UNSLOTH_TORCH_UPGRADE=0 _previous_torch_pin '2.10.0+cu126' 'torch>=2.4,<2.12.0')"
echo "=== install.sh wiring ==="
# The probe must run against the OLD venv, before it is moved aside for rollback.
@ -89,9 +95,39 @@ _probe_line=$(grep -n '_PREV_TORCH_VER=\$(' "$INSTALL_SH" | head -1 | cut -d: -f
_move_line=$(grep -n '_start_studio_venv_replacement "\$VENV_DIR"' "$INSTALL_SH" | head -1 | cut -d: -f1)
assert_eq "probe exists" "yes" "$([ -n "$_probe_line" ] && echo yes)"
assert_eq "probe before venv replacement" "yes" "$([ -n "$_probe_line" ] && [ -n "$_move_line" ] && [ "$_probe_line" -lt "$_move_line" ] && echo yes)"
# The pin must be evaluated AFTER the last index/constraint decision (the Strix
# reroute raises the floor), so a raised floor rejects an older kept release.
_pin_line=$(grep -n '_prev_pin=\$(_previous_torch_pin' "$INSTALL_SH" | head -1 | cut -d: -f1)
_strix_line=$(grep -n 'Strix Halo / Strix Point: force rocm7.2 wheels' "$INSTALL_SH" | head -1 | cut -d: -f1)
assert_eq "pin evaluated after the Strix reroute" "yes" "$([ -n "$_pin_line" ] && [ -n "$_strix_line" ] && [ "$_pin_line" -gt "$_strix_line" ] && echo yes)"
# A kept release that vanished from the index must fall back to the supported range.
assert_eq "resolve-failure fallback wired" "yes" "$(grep -q 'TORCH_CONSTRAINT="\$_PREV_FALLBACK_CONSTRAINT"' "$INSTALL_SH" && echo yes)"
assert_eq "pin gated on SKIP_TORCH" "yes" "$(grep -q 'if \[ "\$SKIP_TORCH" = false \]; then' "$INSTALL_SH" && echo yes)"
# Every --default-index torch install path must go through the kept-release
# helper (definition + default path + three ROCm-index fallbacks + two ROCm
# repairs + the flavor repair), so a pinned release missing from the index
# never aborts a rerun.
_helper_uses=$(grep -c '_install_torch_default_index' "$INSTALL_SH")
assert_eq "kept-release helper used by all default-index paths" "yes" "$([ "$_helper_uses" -ge 8 ] && echo yes)"
_repair_uses=$(grep -c '_install_torch_default_index --force-reinstall' "$INSTALL_SH")
assert_eq "ROCm repairs routed through the kept-release helper" "yes" "$([ "$_repair_uses" -ge 2 ] && echo yes)"
# The wrong-flavor repair must use the helper too (it runs under set -e, so a
# direct uv call with an unresolvable pin would abort the whole installer).
assert_eq "flavor repair routed through the kept-release helper" "yes" "$(grep -q '_install_torch_default_index \\' "$INSTALL_SH" && grep -q -- '--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio' "$INSTALL_SH" && echo yes)"
# The kept-release install must pair the companions with the kept minor:
# torchaudio no longer exact-pins torch, so unconstrained it resolves a newer
# mismatched build (verified: torch==2.9.0 pulled torchaudio 2.11.0 on cu130).
assert_eq "kept-release install pairs torchvision/torchaudio to the kept minor" "yes" "$(grep -q 'torchaudio==2.\${_itdi_minor}.\*' "$INSTALL_SH" && grep -q 'torchvision==0.\$((_itdi_minor + 15)).\*' "$INSTALL_SH" && echo yes)"
# The Radeon direct-wheel path must also honor the pin: an exact-first kept-trio
# attempt (exact patch, else the kept minor's newest patch, with paired
# vision/audio) runs BEFORE the newest-trio search, and the newest-trio search
# only runs when that attempt did not produce a match, so a kept release can
# neither drift to another patch/minor nor be undercut by the gap search.
_radeon_kept_line=$(grep -n '_kept_torch=\$(_pick_radeon_wheel "torch" *"\${_prev_kept_base}"' "$INSTALL_SH" | head -1 | cut -d: -f1)
_radeon_loop_line=$(grep -n 'Loop downwards to find the first complete matching trio' "$INSTALL_SH" | head -1 | cut -d: -f1)
assert_eq "Radeon kept-trio attempt before the newest-trio search" "yes" "$([ -n "$_radeon_kept_line" ] && [ -n "$_radeon_loop_line" ] && [ "$_radeon_kept_line" -lt "$_radeon_loop_line" ] && echo yes)"
assert_eq "Radeon newest-trio search gated on no kept match" "yes" "$(grep -q 'if \[ "\$_radeon_versions_match" != true \] &&' "$INSTALL_SH" && echo yes)"
assert_eq "Radeon kept-trio gap falls back with a warning" "yes" "$(grep -q 'lacks a complete wheel set for kept' "$INSTALL_SH" && echo yes)"
echo ""
if [ "$FAIL" -gt 0 ]; then