diff --git a/install.sh b/install.sh index d5498e7f7a..ea77ecd4d9 100755 --- a/install.sh +++ b/install.sh @@ -2225,9 +2225,20 @@ _torch_index_repairable() { # /.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 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index d5b0aa15af..b68aebb328 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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 diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 2c865d7056..b05dcf2ff9 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -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" } diff --git a/tests/sh/test_torch_index_marker.sh b/tests/sh/test_torch_index_marker.sh index c8f63a5dc7..ad7b1512da 100755 --- a/tests/sh/test_torch_index_marker.sh +++ b/tests/sh/test_torch_index_marker.sh @@ -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" \ diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index 10d198fd85..45b819f98f 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -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"