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

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

Browse source 
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.
This commit is contained in:
danielhanchen 2026-07-06 10:44:41 +00:00
commit 5ed073755c
5 changed files with 75 additions and 19 deletions
Download patch file Download diff file Expand all files Collapse all files

26
install.sh
View file

@ -2225,9 +2225,20 @@ _torch_index_repairable() {
# <venv_dir>/.unsloth-torch-index (single line = the resolved index URL)
_TORCH_INDEX_MARKER_NAME=".unsloth-torch-index"
# Lowercase ONLY a known wheel-family leaf (rocm* / gfx* / cpu / cuXXX); a custom
# mirror leaf keeps its case so a verbatim URL pin is not falsely matched equal.
# Mirrors _normalize_family_leaf in install_python_stack.py / setup.ps1.
_normalize_family_leaf() {
_l_low=$(printf '%s' "$1" | tr '[:upper:]' '[:lower:]')
case "$_l_low" in
rocm*|gfx*|cpu|cu[0-9]*) printf '%s' "$_l_low" ;;
*) printf '%s' "$1" ;;
esac
}
# Normalise a wheel index URL for exact marker/pin comparison: trim whitespace,
# strip ALL trailing slashes, lowercase ONLY the final path segment (the leaf).
# Mirrors _normalize_index_url in install_python_stack.py / setup.ps1 / install.ps1.
# strip ALL trailing slashes, lowercase ONLY a known-family final path segment.
# Mirrors _normalize_index_url in install_python_stack.py / setup.ps1.
_normalize_index_url() {
_n_url="$1"
# Trim leading/trailing whitespace.
@ -2240,11 +2251,11 @@ _normalize_index_url() {
*/*)
_n_head="${_n_url%/*}"
_n_leaf="${_n_url##*/}"
_n_leaf=$(printf '%s' "$_n_leaf" | tr '[:upper:]' '[:lower:]')
_n_leaf=$(_normalize_family_leaf "$_n_leaf")
printf '%s/%s' "$_n_head" "$_n_leaf"
;;
*)
printf '%s' "$_n_url" | tr '[:upper:]' '[:lower:]'
_normalize_family_leaf "$_n_url"
;;
esac
}
@ -3098,6 +3109,13 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
"$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
--index-url "$TORCH_INDEX_URL" \
--force-reinstall
# The repair reinstalled torch from $TORCH_INDEX_URL (the generic
# ROCm index), not the Radeon --find-links repo, so record THAT as
# the marker source. A Radeon --find-links install set
# _TORCH_MARKER_INDEX_URL to its repo.radeon.com base earlier;
# leaving it would make the marker misreport Radeon wheels and let a
# later Radeon pin compare-equal and skip a needed reinstall.
_TORCH_MARKER_INDEX_URL="$TORCH_INDEX_URL"
fi
;;
esac

39
studio/install_python_stack.py
View file

@ -137,15 +137,31 @@ _PYTORCH_WHL_BASE = (
_TORCH_INDEX_MARKER_NAME = ".unsloth-torch-index"
def _normalize_family_leaf(leaf: str) -> str:
"""Lowercase ONLY a known wheel-family leaf (rocm* / gfx* / cpu / cuXXX).
The canonical gfx120X-all (capital X) must match AMD's lowercase gfx120x-all, so
known-family leaves are lowercased. A custom mirror leaf (/Current, /simple, ...)
keeps its case: an unknown-family URL pin is applied verbatim, so /Current and
/current must NOT compare equal. Same known-family set as
_explicit_unknown_family_torch_index_url. Mirrors the gate in install.sh /
setup.ps1. Pure function.
"""
low = leaf.lower()
if low.startswith(("rocm", "gfx")) or low == "cpu" or re.match(r"^cu[0-9]", low):
return low
return leaf
def _normalize_index_url(url: "str | None") -> "str | None":
"""Canonicalise a wheel index URL for exact marker/pin comparison.
Trims surrounding whitespace, strips ALL trailing slashes, and lowercases only
the FINAL path segment (the wheel-family leaf: cu128 / cpu / rocm7.2 / gfx1151 /
gfx120X-all). The host part is left untouched (it may be case-sensitive on some
mirrors); the leaf is lowercased so the canonical gfx120X-all (capital X) and
AMD's lowercase pip leaf gfx120x-all compare equal. MUST match the same
normalization in install.sh / setup.ps1 / install.ps1. Returns None for an
Trims surrounding whitespace, strips ALL trailing slashes, and lowercases the
FINAL path segment ONLY when it is a known wheel-family leaf (cu128 / cpu /
rocm7.2 / gfx1151 / gfx120X-all) -- see _normalize_family_leaf. The host part is
left untouched (case-sensitive on some mirrors), and a custom (unknown-family)
leaf keeps its case so a verbatim URL pin is not falsely matched equal. MUST
match the same normalization in install.sh / setup.ps1. Returns None for an
empty/whitespace-only input. Pure function.
"""
if url is None:
@ -158,8 +174,8 @@ def _normalize_index_url(url: "str | None") -> "str | None":
return None
head, sep, leaf = url.rpartition("/")
if sep:
return f"{head}/{leaf.lower()}"
return url.lower()
return f"{head}/{_normalize_family_leaf(leaf)}"
return _normalize_family_leaf(url)
def _torch_index_marker_path() -> Path:
@ -1504,6 +1520,13 @@ def _ensure_cuda_torch() -> None:
# build simply re-lands on the same family (idempotent).
_installed_desc = _installed_cu if _installed_cu else "an untagged CUDA build"
_why = f"torch is {_installed_desc} but the pinned CUDA index is {_pin_leaf}"
elif _marker == "cuda" and _pinned_cuda and _marker_pin_mismatch(_pin) is True:
# Same cuXXX leaf but the marker records a DIFFERENT full index URL (e.g. the
# official cu128 index vs an internal mirror's cu128). The +cuXXX tag cannot
# see the host change, so consult the exact-URL marker and reinstall from the
# pinned URL (via _detect_cuda_torch_index_url, which honours the override) so
# an explicit mirror pin is applied and re-recorded, not skipped.
_why = f"the pinned CUDA index URL differs from the recorded marker (leaf {_pin_leaf})"
else:
return # healthy CUDA torch matching the pin, or a deliberate CPU wheel

16
studio/setup.ps1
View file

@ -436,16 +436,26 @@ $TorchIndexMarkerName = ".unsloth-torch-index"
# Normalise a wheel index URL for exact marker/pin comparison: trim whitespace,
# strip ALL trailing slashes, lowercase ONLY the final path segment (the leaf).
# Mirrors _normalize_index_url in install.sh / install_python_stack.py / install.ps1.
# Lowercase ONLY a known wheel-family leaf (rocm* / gfx* / cpu / cuXXX); a custom
# mirror leaf keeps its case so a verbatim URL pin is not falsely matched equal.
# Mirrors _normalize_family_leaf in install.sh / install_python_stack.py.
function Get-NormalizedFamilyLeaf {
param([string]$Leaf)
$low = $Leaf.ToLowerInvariant()
if ($low -match '^(rocm|gfx)' -or $low -eq 'cpu' -or $low -match '^cu[0-9]') { return $low }
return $Leaf
}
# Mirrors _normalize_index_url in install.sh / install_python_stack.py.
function Get-NormalizedIndexUrl {
param([string]$Url)
if ([string]::IsNullOrWhiteSpace($Url)) { return $null }
$u = $Url.Trim().TrimEnd('/')
if ([string]::IsNullOrWhiteSpace($u)) { return $null }
$idx = $u.LastIndexOf('/')
if ($idx -lt 0) { return $u.ToLowerInvariant() }
if ($idx -lt 0) { return (Get-NormalizedFamilyLeaf $u) }
$head = $u.Substring(0, $idx)
$leaf = $u.Substring($idx + 1).ToLowerInvariant()
$leaf = Get-NormalizedFamilyLeaf ($u.Substring($idx + 1))
return "$head/$leaf"
}

6
tests/sh/test_torch_index_marker.sh
View file

@ -18,6 +18,8 @@ _TORCH_INDEX_MARKER_NAME=".unsloth-torch-index"
# Extract the marker helpers from install.sh and source them.
_FUNC_FILE=$(mktemp)
{
sed -n '/^_normalize_family_leaf()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_normalize_index_url()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_write_torch_index_marker()/,/^}/p' "$INSTALL_SH"
@ -42,8 +44,8 @@ assert_eq "trailing slashes stripped + leaf lowered" \
assert_eq "whitespace trimmed" \
"https://download.pytorch.org/whl/cu128" \
"$(_normalize_index_url ' https://download.pytorch.org/whl/cu128 ')"
assert_eq "host case preserved, only leaf lowered" \
"https://Mirror.Local/simple" \
assert_eq "host + custom (unknown-family) leaf case preserved" \
"https://Mirror.Local/Simple" \
"$(_normalize_index_url 'https://Mirror.Local/Simple/')"
# gfx120X-all (capital X) and AMD's lowercase pip leaf normalise equal.
assert_eq "gfx120X-all == gfx120x-all after normalize" \

7
tests/studio/install/test_rocm_support.py
View file

@ -1016,8 +1016,11 @@ class TestTorchIndexMarkerHelpers:
assert f("https://repo.amd.com/rocm/whl/gfx120X-all///") == (
"https://repo.amd.com/rocm/whl/gfx120x-all"
)
# Host case preserved (only the leaf is lowered).
assert f("https://Mirror.Local/Simple/") == "https://Mirror.Local/simple"
# Host case preserved; a custom (unknown-family) leaf keeps its case so a
# verbatim URL pin is not falsely matched equal (only known families lower).
assert f("https://Mirror.Local/Simple/") == "https://Mirror.Local/Simple"
# A custom mirror leaf differing only in case must NOT compare equal.
assert f("https://mirror.local/Current") != f("https://mirror.local/current")
# Whitespace trimmed.
assert f(" https://download.pytorch.org/whl/cu128 ") == (
"https://download.pytorch.org/whl/cu128"