diff --git a/install.ps1 b/install.ps1 index 8c667df079..fc005e2c9a 100644 --- a/install.ps1 +++ b/install.ps1 @@ -61,7 +61,11 @@ function Install-UnslothStudio { function Get-TauriGpuBranch { param([string]$TorchIndexFamily) if ($SkipTorch) { return "no_torch" } - if ($TorchIndexFamily -like "cu*") { return "cuda" } + # Require digits after "cu" (cu118/cu128/...) so an odd mirror leaf like + # "custom"/"current" is not mis-branded CUDA. $TorchIndexFamily is already + # normalised by Get-TauriTorchIndexFamily today, but keep the guard narrow + # to match the ^cu[0-9] rule in setup.ps1 / install_python_stack.py. + if ($TorchIndexFamily -match '^cu[0-9]') { return "cuda" } if ($TorchIndexFamily -like "rocm*") { return "rocm" } if ($TorchIndexFamily -eq "cpu") { return "cpu" } return "unknown" @@ -1948,6 +1952,17 @@ exit 0 # 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 when the caller names the wheel index + # (headless / CI / cross-install). Matches install.sh::get_torch_index_url and + # install_python_stack.py: UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); + # UNSLOTH_TORCH_INDEX_FAMILY is the convenience leaf (cpu, cu128, rocm6.4, ...) + # appended to the mirror base so UNSLOTH_PYTORCH_MIRROR is still honoured. + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { + return $env:UNSLOTH_TORCH_INDEX_URL.Trim().TrimEnd('/') + } + 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 @@ -1968,6 +1983,30 @@ exit 0 return "$baseUrl/cu126" } + # ── Torch-index marker ─────────────────────────────────────────────────── + # After a successful torch install, record the exact wheel --index-url used at + # a stable per-venv path so a later `unsloth studio update` (setup.ps1 / + # install_python_stack.py) can detect a pin change by an EXACT string compare + # rather than the wheel version-tag heuristic. Path/format MUST match + # install.sh, studio/setup.ps1 and install_python_stack.py: + # \.unsloth-torch-index (single line = the resolved index URL) + # install.ps1 only WRITES the marker; setup.ps1 reads it during stale detection. + function Write-TorchIndexMarker { + param([string]$VenvDir, [string]$IndexUrl) + if ([string]::IsNullOrWhiteSpace($VenvDir)) { return } + if ([string]::IsNullOrWhiteSpace($IndexUrl)) { return } + if (-not (Test-Path -LiteralPath $VenvDir -PathType Container)) { return } + $marker = Join-Path $VenvDir ".unsloth-torch-index" + $tmp = "$marker.$PID.tmp" + try { + # Single LF-terminated line, no BOM (parity with the sh/py writers). + [System.IO.File]::WriteAllText($tmp, ($IndexUrl.Trim() + "`n"), (New-Object System.Text.UTF8Encoding($false))) + Move-Item -LiteralPath $tmp -Destination $marker -Force -ErrorAction Stop + } catch { + try { if (Test-Path -LiteralPath $tmp) { Remove-Item -LiteralPath $tmp -Force -ErrorAction SilentlyContinue } } catch {} + } + } + # ── 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. @@ -2025,6 +2064,11 @@ exit 0 } catch { return $null } } + # An explicit UNSLOTH_TORCH_INDEX_URL / _FAMILY pin is authoritative: the AMD + # ROCm reroute below must not rewrite it (e.g. a deliberate cpu pin on an AMD + # host, or a pinned ROCm family we already resolved in Get-TorchIndexUrl). + $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 ── @@ -2036,7 +2080,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 @@ -2086,6 +2132,42 @@ exit 0 } } + # An explicit 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 torch._C._grouped_mm null-ptr + # bug). Route a pinned ROCm index through the ROCm install path with the same + # 2.11 floor/companions the unpinned reroute derives from the gfx arch. + if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) { + $_pinLeaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLower() + $_pinRocm211 = $false + if ($_pinLeaf -match '^rocm(\d+)\.(\d+)') { + # Only KNOWN-2.11 rocm indexes (rocm7.2) get the 2.11 floor; do not floor + # an unknown newer rocm speculatively (rocm7.3 does not exist). Matches + # install.sh's rocm7.2 KNOWN-2.11 leaf, setup.ps1's Test-RocmKnown211Version + # and _ROCM_KNOWN_TORCH211_VERSIONS in install_python_stack.py. + $_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2) + } + # Only the gfx families the AMD arch map above pins to torch 2.11 need the + # floor here (gfx120X-all, gfx1151, gfx1150 -- the _grouped_mm bug arches). + # Other per-arch indexes (gfx110X-all, gfx90a, gfx908) publish <2.11 wheels + # and the automatic path leaves them bare, so an override to one of those + # must NOT force a 2.11 floor the normal path intentionally avoids (which + # would fail resolution or drag in mismatched companions). + $_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 -like 'gfx*' -or $_pinLeaf -like 'rocm*') { + # Other gfx per-arch indexes and older rocm (<=7.1) ship torch <2.11; + # route via the ROCm path with bare specs so the pinned family resolves + # its own wheels (matches the automatic path's bare floor for these). + $ROCmIndexUrl = $TorchIndexUrl + } + } + if ($ROCmIndexUrl) { $TorchIndexFamily = "rocm" } else { @@ -2198,18 +2280,22 @@ exit 0 $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 --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { # Transient AMD-index failure: fall back to a CPU base so the install # still completes; Studio setup retries ROCm afterwards. + # Use an explicit CPU index: in the unpinned AMD path $TorchIndexUrl is + # already */cpu, but for a pinned ROCm index it IS the ROCm mirror, so + # reusing it here would just retry the failing ROCm index, not fall back. + $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, Studio 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 --index-url $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $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) @@ -2319,8 +2405,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 --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { @@ -2349,6 +2435,17 @@ exit 0 } } + # ── Record the resolved torch wheel index (marker) ── + # Torch is now resolved; write the exact --index-url used so setup.ps1 / + # install_python_stack.py can detect a later pin change by an EXACT string + # compare instead of the version-tag heuristic. Reflects the installed family: + # $ROCmIndexUrl when the ROCm path ran, else the CUDA/CPU/pinned $TorchIndexUrl. + # Skipped for --no-torch (nothing installed). Matches install.sh / setup.ps1. + if (-not $SkipTorch) { + $MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } else { $TorchIndexUrl } + Write-TorchIndexMarker -VenvDir $VenvDir -IndexUrl $MarkerIndexUrl + } + # Overlay Tauri-bundled studio fixes that may be ahead of PyPI. Skipped # for --local: the editable install above already makes _PACKAGE_ROOT in # unsloth_cli/commands/studio.py resolve to the repo (PEP 660 __file__). diff --git a/install.sh b/install.sh index 0370559540..f748aee08c 100755 --- a/install.sh +++ b/install.sh @@ -334,7 +334,11 @@ _tauri_gpu_branch() { return fi case "$_diag_family" in - cu*) echo "cuda" ;; + # Require a digit after cu (cu118/cu128/...) so an odd leaf like custom / + # current is not branded CUDA -- matches the ^cu[0-9] rule in setup.ps1 / + # install_python_stack.py. $_diag_family is already normalised by + # _tauri_torch_index_family, but keep the guard narrow for parity. + cu[0-9]*) echo "cuda" ;; rocm*) if [ "$_diag_radeon" = true ]; then echo "rocm_radeon" @@ -1892,6 +1896,14 @@ if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; t TORCH_CONSTRAINT="torch>=2.6,<2.11.0" fi fi +# Companion (torchvision/torchaudio) constraints. Bare by default: the pytorch.org +# cu*/cpu/rocmX.Y indexes are curated so uv resolves an ABI-consistent trio from a +# bare name. They are pinned alongside TORCH_CONSTRAINT only for the torch-2.11 +# AMD paths (rocm7.2 / per-gfx index / Strix), where AMD publishes each wheel +# independently and can ship a newer torchvision/torchaudio (built against torch +# 2.12) before removing the 2.11-matched one -- see the rocm7.2/gfx case below. +TORCHVISION_CONSTRAINT="torchvision" +TORCHAUDIO_CONSTRAINT="torchaudio" # ── Resolve repo root (for --local installs) ── _REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)" @@ -2015,6 +2027,33 @@ _has_usable_nvidia_gpu() { get_torch_index_url() { _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" _base="${_base%/}" + # Explicit override -- skip ALL GPU probing when the caller pins the wheel + # index. Headless / container / CI builds (and anyone cross-installing for a + # different target) must not let the build host's GPU -- or the lack of one -- + # decide the wheel family. This is the same "tell the build, don't ask the + # hardware" approach the Docker base image and vLLM/SGLang's Dockerfiles take. + # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); UNSLOTH_TORCH_INDEX_FAMILY + # is the convenience form (cpu, cu124, cu126, cu128, cu130, rocm6.4, ...) + # appended to the mirror base so UNSLOTH_PYTORCH_MIRROR is still honoured. + # Trim leading/trailing whitespace so a whitespace-only value is treated as + # unset (parity with the Python .strip() and PowerShell IsNullOrWhiteSpace + # paths); otherwise " " would pass -n and yield an invalid index URL. + _url="${UNSLOTH_TORCH_INDEX_URL:-}" + _url="${_url#"${_url%%[![:space:]]*}"}"; _url="${_url%"${_url##*[![:space:]]}"}" + if [ -n "$_url" ]; then + # Strip ALL trailing slashes (match the Python side's .rstrip("/") and the + # Strix mirror handling below) -- a double/triple-slash URL 404s on strict + # pip proxies (artifactory, sonatype). + while [ "${_url%/}" != "$_url" ]; do _url="${_url%/}"; done + 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. @@ -2145,9 +2184,13 @@ _torch_flavor_tag() { # 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. +# Lowercase the leaf first so gfx120X-all (capital X) / any cased mirror leaf +# still classifies, keeping the cuXXX tag comparison against _torch_flavor_tag +# (which emits lowercase cuXXX) case-consistent. _expected_torch_flavor_tag() { _u="${1%/}" _leaf="${_u##*/}" + _leaf=$(printf '%s' "$_leaf" | tr '[:upper:]' '[:lower:]') case "$_leaf" in cu[0-9]*) echo "$_leaf" ;; cpu) echo "cpu" ;; @@ -2161,15 +2204,72 @@ _expected_torch_flavor_tag() { # resolves (torch + every transitive dep) via --index-url -- the same URLs the # 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. +# Lowercase the leaf first so a cased leaf (e.g. gfx120X-all) is recognised. _torch_index_repairable() { _u="${1%/}" _leaf="${_u##*/}" + _leaf=$(printf '%s' "$_leaf" | tr '[:upper:]' '[:lower:]') case "$_leaf" in cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;; *) echo "no" ;; esac } +# ── Torch-index marker ─────────────────────────────────────────────────────── +# After a successful torch install this records the exact wheel --index-url used +# at a stable per-venv path so `unsloth studio update` / setup.ps1 / +# install_python_stack.py can make the "did the pinned index change?" decision by +# an EXACT string compare instead of inferring it from the wheel version tag +# (which cannot encode the AMD per-arch gfx family). The path and format MUST match +# install_python_stack.py, setup.ps1 and install.ps1: +# /.unsloth-torch-index (single line = the resolved index URL) +_TORCH_INDEX_MARKER_NAME=".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_python_stack.py / setup.ps1 / install.ps1. +_normalize_index_url() { + _n_url="$1" + # Trim leading/trailing whitespace. + _n_url="${_n_url#"${_n_url%%[![:space:]]*}"}"; _n_url="${_n_url%"${_n_url##*[![:space:]]}"}" + [ -n "$_n_url" ] || { printf '%s' ""; return; } + # Strip all trailing slashes. + while [ "${_n_url%/}" != "$_n_url" ]; do _n_url="${_n_url%/}"; done + [ -n "$_n_url" ] || { printf '%s' ""; return; } + case "$_n_url" in + */*) + _n_head="${_n_url%/*}" + _n_leaf="${_n_url##*/}" + _n_leaf=$(printf '%s' "$_n_leaf" | tr '[:upper:]' '[:lower:]') + printf '%s/%s' "$_n_head" "$_n_leaf" + ;; + *) + printf '%s' "$_n_url" | tr '[:upper:]' '[:lower:]' + ;; + esac +} + +# Write the resolved torch --index-url ($2) into the marker under venv dir ($1), +# atomically (temp file + mv). Best-effort: a write failure never aborts the +# install (the repair path then falls back to the version-tag heuristics). A blank +# URL is ignored (nothing meaningful to record). +_write_torch_index_marker() { + _wm_venv="$1" + _wm_url="$2" + [ -n "$_wm_venv" ] || return 0 + [ -d "$_wm_venv" ] || return 0 + _wm_url="${_wm_url#"${_wm_url%%[![:space:]]*}"}"; _wm_url="${_wm_url%"${_wm_url##*[![:space:]]}"}" + [ -n "$_wm_url" ] || return 0 + _wm_marker="$_wm_venv/$_TORCH_INDEX_MARKER_NAME" + _wm_tmp="$_wm_marker.$$.tmp" + if printf '%s\n' "$_wm_url" > "$_wm_tmp" 2>/dev/null; then + mv -f "$_wm_tmp" "$_wm_marker" 2>/dev/null || rm -f "$_wm_tmp" 2>/dev/null || true + else + rm -f "$_wm_tmp" 2>/dev/null || true + fi + return 0 +} + 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/, @@ -2411,7 +2511,26 @@ _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 downstream: skip the WSL ROCm bootstrap (which can run +# sudo + large downloads after probing /dev/dxg) and the Radeon/Strix rerouting +# below (which would re-probe the GPU and overwrite the pinned URL). A headless / +# container / CI build must get exactly the index it asked for. +# Trim leading/trailing whitespace before the -n check (parity with +# get_torch_index_url): a whitespace-only override (e.g. from a CI template) is +# treated as unset there and falls through to normal CPU/ROCm detection, so it +# must NOT flip _torch_index_pinned true here -- otherwise the WSL bootstrap and +# Radeon/Strix reroutes below are skipped while detection still picks the normal +# index, silently disabling them. +_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) @@ -2422,24 +2541,63 @@ 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). +# Lowercase the leaf once here so every gfx*/rocm*/cu* allowlist below matches +# regardless of case. The canonical AMD RDNA4 leaf is gfx120X-all (capital X, +# from the arch maps in install_python_stack.py / install.ps1); without this a +# pin to gfx120X-all would miss the lowercase gfx120x-all allowlist entries. +# The allowlists stay lowercase; inputs are normalised to lowercase. +# CUDA is branded only on a real cu[0-9]* leaf (^cu[0-9]) -- NOT a bare cu*/catch- +# all -- so a full-override mirror leaf like /current or /custom does NOT commit a +# CUDA backend. An unknown leaf leaves the backend var unset so the stack probes +# the GPU instead of returning early in _ensure_rocm_torch on AMD hosts. +# Matches _is_cuda_family_leaf (Python) / Test-CudaFamilyLeaf (PowerShell). _torch_index_leaf="${TORCH_INDEX_URL%/}" _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, e.g. /current, /custom): do NOT commit a backend. + # Unset so a stale inherited value can't leak and the stack probes the GPU. + *) unset UNSLOTH_TORCH_BACKEND ;; esac -# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. -# All other ROCm tags and CUDA stay within <2.11.0. -case "$TORCH_INDEX_URL" in - */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; +# rocm7.2 and the AMD per-gfx indexes with the torch._C._grouped_mm bug on <2.11 +# (repo.amd.com/.../gfx120X-all, gfx1151, gfx1150) ship torch 2.11.0 -- raise the +# constraint to allow it. This also covers a pinned full-URL or family override +# (e.g. UNSLOTH_TORCH_INDEX_URL=.../gfx1151) that returns early above and so never +# hits the Strix reroute that otherwise raises this constraint. Pin the companions +# to the matching 2.11 range too: the per-gfx index publishes torchvision/torchaudio +# independently and a bare name can resolve a 2.12-built wheel (ABI mismatch). +# Matches setup.ps1's *FloorMap and _ROCM_GFX_TORCH211_LEAVES in install_python_stack.py. +# Only these gfx families need the floor: the other per-arch indexes (gfx110X-all, +# gfx90a, gfx908) publish <2.11 wheels, so a pinned override to one of those must NOT +# be forced onto the 2.11 line (which would fail resolution / pull mismatched +# companions). All other ROCm tags and CUDA stay within <2.11.0. +# Match on the FINAL path segment ($_torch_index_leaf, computed above), NOT the +# whole URL: a custom UNSLOTH_PYTORCH_MIRROR whose base path contains a "gfx" or +# "rocm7.2" segment (e.g. https://mirror.local/gfx-cache) with a cu*/cpu family +# must not be treated as an AMD per-arch index and pushed to the 2.11 line. This +# mirrors the leaf-only backend classification just above. +# _torch_index_leaf is already lowercased above, so the canonical gfx120X-all +# (capital X, from the arch maps) pins here via the lowercase gfx120x-all entry. +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 # 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 entirely when the index is pinned: an explicit override (even a ROCm +# one like UNSLOTH_TORCH_INDEX_FAMILY=rocm6.4) 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 && \ @@ -2516,10 +2674,15 @@ 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 the 2.11 range (per-gfx index publishes them + # independently); mirrors the rocm7.2/gfx case above. + 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) _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" @@ -2703,7 +2866,7 @@ if [ "$_MIGRATED" = true ]; then 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 \ + "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ --index-url "$TORCH_INDEX_URL" \ --force-reinstall fi @@ -2829,7 +2992,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then [ "$_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 \ + "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ --index-url "$TORCH_INDEX_URL" else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." @@ -2852,18 +3015,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then 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 \ + "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ --index-url "$TORCH_INDEX_URL" 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 \ + "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ --index-url "$TORCH_INDEX_URL" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." - 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" \ --index-url "$TORCH_INDEX_URL" fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). @@ -2923,7 +3086,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then 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 \ + "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ --index-url "$TORCH_INDEX_URL" \ --force-reinstall fi @@ -2965,7 +3128,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then && [ "$(_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 \ + "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \ --index-url "$TORCH_INDEX_URL" \ --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) @@ -2977,11 +3140,20 @@ 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 --index-url $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\" --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" fi fi fi +# ── Record the resolved torch wheel index (marker) ── +# Torch is now fully resolved; write the exact --index-url used so `unsloth studio +# update` (install_python_stack.py / setup.ps1) can detect a later pin change by an +# exact string compare rather than the version-tag heuristic. Only when torch was +# actually installed from a resolved index (skip --no-torch / no-URL fallback). +if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then + _write_torch_index_marker "$VENV_DIR" "$TORCH_INDEX_URL" +fi + # ── Run studio setup ── tauri_log "STEP" "Running Studio setup" # When --local, use the repo's own setup.sh directly. diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 439d3ffe7b..90aac9533b 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -74,6 +74,23 @@ (6, 0): "rocm6.0", } +# AMD per-arch index leaves that need the torch 2.11 floor (the torch._C._grouped_mm +# null-ptr bug lives in the <2.11 wheels for these arches). Mirrors the gfx keys in +# _WINDOWS_ROCM_TORCH_PKG_SPECS and the *FloorMap sets in install.ps1 / setup.ps1. +# Other per-arch indexes (gfx110X-all, gfx90a, gfx908) publish <2.11 wheels and must +# stay bare, so an override to one of them must NOT be forced onto the 2.11 line. +_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"}) + +# The pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (verified against +# download.pytorch.org): rocm7.2 -> torch 2.11.0 is the ONLY stable 2.11 rocm +# index today (rocm7.1 -> 2.10.0, rocm6.4 -> 2.9.1; rocm7.3 / torch 2.12 do NOT +# exist as stable). Do NOT treat an unknown newer rocm (rocm7.3, rocm8.0, ...) as +# 2.11 speculatively -- that is exactly the mismatch bug tracked in the review. +# MUST match the rocm leaf in the KNOWN-2.11 case in install.sh / setup.ps1 / +# install.ps1 (rocm7.2 there too). Bump alongside those when a new stable rocm +# index publishes torch 2.11+. Stored as (major, minor) tuples for exact compares. +_ROCM_KNOWN_TORCH211_VERSIONS: frozenset[tuple[int, int]] = frozenset({(7, 2)}) + # Per-tag pip specs; rocm7.2 ships torch 2.11.0 (older tags cap at 2.10.x). _ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = { "rocm7.2": ( @@ -103,6 +120,127 @@ os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl" ).rstrip("/") +# ── Torch-index marker ───────────────────────────────────────────────────────── +# After a successful torch install/reinstall, record the exact wheel --index-url +# used at a stable per-venv path. On `studio update`/repair the marker turns the +# "did the pinned index change?" decision into an EXACT string compare instead of +# inferring it from the wheel's +rocm/+cu version tag (which cannot encode the AMD +# per-arch gfx family: two gfx 2.11 indexes both install +rocm7.13.0). This +# dissolves the per-arch-switch (gfx1151 -> gfx120X-all) and custom-URL +# (/simple, /current) cases the version-tag heuristics cannot see. +# +# The path/format MUST match install.sh, install.ps1 and setup.ps1: +# /.unsloth-torch-index (single line = the resolved index URL) +# Written atomically (temp file + os.replace). A missing/empty/corrupt marker is +# treated as absent, so old venvs (and torch installed out-of-band) fall back to +# the existing +rocm/version-tag heuristics -- backward compatibility is required. +_TORCH_INDEX_MARKER_NAME = ".unsloth-torch-index" + + +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 + empty/whitespace-only input. Pure function. + """ + if url is None: + return None + url = url.strip() + if not url: + return None + url = url.rstrip("/") + if not url: + return None + head, sep, leaf = url.rpartition("/") + if sep: + return f"{head}/{leaf.lower()}" + return url.lower() + + +def _torch_index_marker_path() -> Path: + """Path to the per-venv torch-index marker (see _TORCH_INDEX_MARKER_NAME). + + Anchored at sys.prefix (the venv the installer targets via --python + sys.executable), so it matches install.sh's $VENV_DIR and the PowerShell + $VenvDir marker location. + """ + return Path(sys.prefix) / _TORCH_INDEX_MARKER_NAME + + +def _read_torch_index_marker() -> "str | None": + """Return the recorded torch --index-url from the marker, else None. + + None when the marker is missing, empty, or unreadable (corrupt/permission) -- + all treated as "no marker" so the caller falls back to the version-tag + heuristics. The stored URL is returned VERBATIM (not normalized); callers + normalize both sides before comparing. + """ + try: + text = _torch_index_marker_path().read_text(encoding = "utf-8") + except (OSError, ValueError): + return None + line = text.strip() + return line or None + + +def _write_torch_index_marker(index_url: "str | None") -> None: + """Record the resolved torch wheel --index-url at the per-venv marker path. + + Best-effort and atomic (temp file in the same dir + os.replace). Never raises: + a marker write failure must not abort an otherwise-successful install (the + repair path then falls back to the heuristics, same as an old venv). A blank + index_url is ignored (nothing meaningful to record). + """ + if not index_url or not index_url.strip(): + return + marker = _torch_index_marker_path() + payload = index_url.strip() + "\n" + try: + marker.parent.mkdir(parents = True, exist_ok = True) + fd, tmp = tempfile.mkstemp(prefix = ".unsloth-torch-index.", dir = str(marker.parent)) + try: + with os.fdopen(fd, "w", encoding = "utf-8") as fh: + fh.write(payload) + os.replace(tmp, str(marker)) + except OSError: + try: + os.unlink(tmp) + except OSError: + pass + raise + except OSError: + # Non-fatal: fall back to writing directly, then give up silently. + try: + marker.write_text(payload, encoding = "utf-8") + except OSError: + pass + + +def _marker_pin_mismatch(pin_url: str) -> "bool | None": + """Compare an explicit torch-index pin against the recorded marker. + + Returns: + * True -> the pin resolves to a DIFFERENT index than the marker records + (exact compare after _normalize_index_url) -> reinstall needed. + * False -> the pin matches the marker exactly -> no reinstall (no loop). + * None -> no usable marker (missing/empty/corrupt) -> the caller must fall + back to the +rocm/version-tag heuristics (backward compatibility). + + This is the EXACT signal the version-tag heuristics cannot provide: switching + one 2.11 gfx pin to another (gfx1151 -> gfx120X-all) changes the pinned URL but + not the installed +rocm7.13.0 wheel tag. Pure w.r.t. its arg (reads the marker). + """ + marker = _read_torch_index_marker() + if marker is None: + return None + return _normalize_index_url(pin_url) != _normalize_index_url(marker) + + # CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are # pinned to the torch<2.11 family rather than left bare: the install uses an # exclusive --index-url (no PyPI fallback), so a bare name could resolve a @@ -115,6 +253,12 @@ "torchaudio>=2.4,<2.11.0", ) +# CPU torch repair specs (see _ensure_cpu_torch). Same bounds/reasoning as the +# CUDA spec above: the /cpu index now also publishes torch 2.11+, so a bare trio +# from the exclusive --index-url would resolve outside the repo's supported +# <2.11 range (and can pick a torchvision built against a different torch major). +_CPU_TORCH_PKG_SPEC: tuple[str, str, str] = _CUDA_TORCH_PKG_SPEC + # torchao's C++ extensions are built against ONE exact torch release; a newer # torch makes torchao skip its cpp kernels ("Skipping import of cpp extensions # due to incompatible torch version ...") and fall back to slow Python. Because @@ -1002,11 +1146,20 @@ def _detect_cuda_torch_index_url() -> str: Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update` repairs to the same wheel family a fresh `curl | sh` install would pick. - Probes nvidia-smi (PATH, then /usr/bin/nvidia-smi) and parses both the - legacy "CUDA Version:" and the newer "CUDA UMD Version:" spellings. - Defaults to cu126 when nvidia-smi is missing or the version is unreadable - (e.g. NVIDIA detected only via the /proc/driver/nvidia/gpus fallback). + Honours the same explicit overrides first (UNSLOTH_TORCH_INDEX_URL / + UNSLOTH_TORCH_INDEX_FAMILY) so a headless / container / CI install never lets + the host GPU decide the wheel family. Otherwise probes nvidia-smi (PATH, then + /usr/bin/nvidia-smi) and parses both the legacy "CUDA Version:" and the newer + "CUDA UMD Version:" spellings. Defaults to cu126 when nvidia-smi is missing or + the version is unreadable (e.g. NVIDIA detected only via the + /proc/driver/nvidia/gpus fallback). """ + _override_url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + if _override_url: + return _override_url.rstrip("/") + _override_family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + if _override_family: + return f"{_PYTORCH_WHL_BASE}/{_override_family.strip('/')}" exe = shutil.which("nvidia-smi") if not exe and os.path.isfile("/usr/bin/nvidia-smi"): exe = "/usr/bin/nvidia-smi" @@ -1041,6 +1194,209 @@ def _detect_cuda_torch_index_url() -> str: return f"{_PYTORCH_WHL_BASE}/{tag}" +def _explicit_torch_index_url() -> "str | None": + """The wheel index URL pinned via UNSLOTH_TORCH_INDEX_URL / _FAMILY, else None. + + Lets the CUDA/ROCm repair helpers honour the exact pinned family/URL instead + of re-probing the GPU. Mirrors install.sh::get_torch_index_url's override. + """ + url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + if url: + return url.rstrip("/") + family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + if family: + return f"{_PYTORCH_WHL_BASE}/{family.strip('/')}" + return None + + +def _explicit_rocm_torch_index_url() -> "str | None": + """The pinned wheel index URL when it names a ROCm family (rocm*/gfx*), else None.""" + url = _explicit_torch_index_url() + if url is None: + return None + leaf = url.rstrip("/").rsplit("/", 1)[-1].lower() + return url if leaf.startswith(("rocm", "gfx")) else None + + +def _rocm_pin_family_mismatch(pin_url: str, installed_ver: str) -> bool: + """True when an explicit ROCm pin names a different ROCm family than the + already-installed ROCm torch, so the pin needs a reinstall to be applied. + + Mirrors setup.ps1's stale-venv ROCm comparison. The pin leaf classifies into + three cases, matching the install-spec path in _ensure_rocm_torch: + * rocmX.Y leaf -> compare the pinned rocm version to the installed one + exactly when both are readable; else fall back to the torch 2.11 line + (rocm>=7.2 serves 2.11, older rocm does not). + * gfx leaf in _ROCM_GFX_TORCH211_LEAVES (gfx120x-all/gfx1151/gfx1150) -> + the install path pulls AMD's per-arch wheel (tagged with a THREE-part + +rocmA.B.C local version, e.g. 2.11.0+rocm7.13.0). A generic pytorch.org + rocm wheel (two-part +rocmA.B, e.g. +rocm7.2) or any pre-2.11 build is a + mismatch even when both are torch 2.11 -- the user asked for the per-arch + index. An already-installed per-arch wheel (three-part tag) is NOT a + mismatch, so a satisfied gfx pin does not reinstall-loop. + * gfx leaf NOT in the 2.11 allowlist (gfx110X-all/gfx90a/gfx908) -> the + install path uses the default <2.11 specs, so a correct 2.10+rocm wheel + must NOT be flagged. Mismatch only when the installed torch is 2.11+. + A pin that resolves to the same family as what is installed is NOT a mismatch, + so a correct ROCm venv is never needlessly reinstalled. Pure function. + """ + leaf = pin_url.rstrip("/").rsplit("/", 1)[-1].lower() + # Pinned ROCm version (from a rocmX.Y leaf). + _pin_rocm = re.match(r"^rocm(\d+)\.(\d+)", leaf) + _pin_ver = (int(_pin_rocm.group(1)), int(_pin_rocm.group(2))) if _pin_rocm else None + # Installed ROCm version (+rocmX.Y) and whether the installed wheel carries a + # THREE-part local version (+rocmA.B.C) -- the AMD per-arch signature that + # distinguishes a repo.amd.com/gfx* wheel from a two-part pytorch.org one. + _inst_rocm = re.search(r"\+rocm(\d+)\.(\d+)", installed_ver) + _inst_ver = (int(_inst_rocm.group(1)), int(_inst_rocm.group(2))) if _inst_rocm else None + _inst_is_perarch = re.search(r"\+rocm\d+\.\d+\.\d+", installed_ver) is not None + # A ROCm build MUST carry a +rocm local tag. An untagged CPU/CUDA wheel (no + # +rocm, e.g. "2.10.0" / "2.11.0") never satisfies a ROCm pin -- always a + # mismatch -- mirroring setup.ps1's Get-RocmPinStaleTags. (In practice + # _ensure_rocm_torch only calls this when has_hip_torch is True, but keep the + # pure function correct for any input so it stays in lockstep with the PS side.) + _inst_has_rocm = re.search(r"\+rocm", installed_ver) is not None + # Whether the installed torch RELEASE (before "+") is 2.11+. + _inst_rel = re.match(r"^(\d+)\.(\d+)", installed_ver) + _inst_is_211 = ( + (int(_inst_rel.group(1)), int(_inst_rel.group(2))) >= (2, 11) if _inst_rel else False + ) + + if leaf.startswith("gfx"): + # gfx per-arch pin: only the _grouped_mm-bug arches (the 2.11 allowlist) + # pull the AMD per-arch wheel; other gfx leaves stay on the default + # <2.11 specs (see _ROCM_TORCH_PKG_SPECS selection below). + if leaf in _ROCM_GFX_TORCH211_LEAVES: + # Expect the AMD per-arch wheel (three-part +rocmA.B.C, torch 2.11+). + # A satisfied per-arch install is NOT a mismatch (no reinstall loop); + # a generic rocm wheel or any pre-2.11 build IS a mismatch even at 2.11. + return not (_inst_is_211 and _inst_is_perarch) + # Non-2.11 gfx leaf: install path uses default <2.11 specs, so a correct + # <2.11 wheel must stay. An untagged (no +rocm) wheel never satisfies the + # pin -> mismatch; otherwise mismatch only when the installed torch is 2.11+. + return (not _inst_has_rocm) or _inst_is_211 + + # rocmX.Y pin. Only the KNOWN-2.11 rocm indexes are treated as the 2.11 line; + # an unknown newer rocm (rocm7.3, rocm8.0, ...) is NOT floored speculatively + # (it does not exist yet). Aligns with install.sh / setup.ps1 / install.ps1. + _pin_is_211 = _pin_ver in _ROCM_KNOWN_TORCH211_VERSIONS if _pin_ver is not None else False + if _pin_ver is not None and _inst_ver is not None: + # Both ROCm versions readable: exact (major, minor) comparison. A generic + # rocm7.2 pin over the AMD per-arch (+rocm7.13.x) wheel compares (7, 2) vs + # (7, 13) -> mismatch, which correctly reinstalls the generic wheel the + # user pinned instead of leaving the per-arch one in place. + return _pin_ver != _inst_ver + # rocm pin with an unreadable installed version: compare on the torch 2.11 line, + # but an untagged (no +rocm) wheel never satisfies a rocmX.Y pin -> mismatch. + if not _inst_has_rocm: + return True + return _pin_is_211 != _inst_is_211 + + +def _explicit_cpu_torch_index_url() -> "str | None": + """The pinned wheel index URL when it names the CPU family (leaf == cpu), else None. + + An explicit CPU pin (UNSLOTH_TORCH_INDEX_FAMILY=cpu or a URL ending in /cpu) + is authoritative -- see _ensure_cpu_torch. + """ + url = _explicit_torch_index_url() + if url is None: + return None + leaf = url.rstrip("/").rsplit("/", 1)[-1].lower() + return url if leaf == "cpu" else None + + +def _is_cuda_family_leaf(leaf: str) -> bool: + """True only for a real CUDA wheel-family leaf: "cu" followed by digits + (cu118, cu126, cu128, cu130, ...). + + A bare startswith("cu") wrongly matches arbitrary mirror leaves like "custom" + or "current", which would let _ensure_cuda_torch treat a generic mirror pin as + CUDA authority and force a CUDA reinstall over a CPU/ROCm venv on a non-NVIDIA + host -- exactly what _explicit_cuda_torch_index_url's contract forbids. + """ + return re.match(r"^cu[0-9]", leaf) is not None + + +def _explicit_cuda_torch_index_url() -> "str | None": + """The pinned wheel index URL when it names a CUDA family (leaf cuXXX), else None. + + Mirrors _explicit_rocm/cpu_torch_index_url so _ensure_cuda_torch only treats a + *CUDA* pin as authority to override the NVIDIA-presence gate. An arbitrary + mirror URL (or a ROCm/CPU pin) must not force a CUDA reinstall over a working + ROCm/CPU venv on a non-NVIDIA host, so match cuXXX (cu + digits) narrowly + rather than any leaf starting with "cu" (which would catch custom/current). + """ + url = _explicit_torch_index_url() + if url is None: + return None + leaf = url.rstrip("/").rsplit("/", 1)[-1].lower() + return url if _is_cuda_family_leaf(leaf) else None + + +def _explicit_unknown_family_torch_index_url() -> "str | None": + """The pinned index URL when its leaf names NO known torch family, else None. + + A "known" family leaf is rocm* / gfx* / cpu / cuXXX -- the ones the dedicated + _explicit_{rocm,cpu,cuda}_torch_index_url helpers already classify. Anything + else (a private PEP 503 mirror ending in /simple, /current, /custom, ...) is an + UNKNOWN family: the version-tag heuristics cannot infer whether it is stale, so + the marker drives the decision and, when it differs (or is absent), the URL is + reinstalled VERBATIM -- "URL wins verbatim". Matches the unknown-leaf branch in + install.sh / setup.ps1 / install.ps1. + """ + url = _explicit_torch_index_url() + if url is None: + return None + leaf = url.rstrip("/").rsplit("/", 1)[-1].lower() + if leaf.startswith(("rocm", "gfx")) or leaf == "cpu" or _is_cuda_family_leaf(leaf): + return None + return url + + +def _ensure_verbatim_torch_index() -> None: + """Reinstall torch/vision/audio VERBATIM from an explicit custom index pin. + + Handles the "URL wins verbatim" case for an explicit UNSLOTH_TORCH_INDEX_URL + (or _FAMILY) whose leaf names no known family (e.g. a private mirror ending in + /simple or /current). The other _ensure_* helpers all return None for such a + pin (it is neither rocm/gfx nor cpu nor cuXXX), so without this the pin would be + silently ignored and the GPU-probed default index used instead. + + Fires ONLY when the marker exists and records a DIFFERENT index than the pin + (or after this reinstalls, the marker is rewritten to match). With NO marker it + is a no-op: an old venv (or torch installed out-of-band) must not be blindly + force-reinstalled from an unverified custom index -- backward compatibility. + macOS/no-torch: skipped (no torch to repair). The install uses the pinned URL + exclusively (--index-url) with bare specs so it "wins verbatim". + """ + if NO_TORCH or IS_MACOS: + return + pin = _explicit_unknown_family_torch_index_url() + if pin is None: + return + _mismatch = _marker_pin_mismatch(pin) + if _mismatch is not True: + # None -> no marker (fall back / do nothing); False -> already this index. + return + print( + f" explicit torch index pin ({pin}) differs from the recorded index -- " + f"reinstalling torch verbatim from it" + ) + pip_install( + "torch (pinned custom index)", + "--force-reinstall", + "--no-cache-dir", + "torch", + "torchvision", + "torchaudio", + "--index-url", + pin, + constrain = False, + ) + _write_torch_index_marker(pin) + + def _ensure_cuda_torch() -> None: """Repair a venv whose torch is a ROCm build on an NVIDIA host. @@ -1065,15 +1421,23 @@ def _ensure_cuda_torch() -> None: # Never undo a deliberate ROCm install (setup.ps1 sets this marker). if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": return + # An explicit CUDA wheel-index pin (headless / container / CI cross-install) + # commits to CUDA wheels and, like install.sh's get_torch_index_url override, + # skips ALL host-GPU probing -- so it must clear BOTH the CUDA_VISIBLE_DEVICES + # hide gate and the NVIDIA-presence gate below, not just the latter. Otherwise + # `CUDA_VISIBLE_DEVICES=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update` + # (the exact GPU-less CI case this override targets) would still bail here. + _cuda_pinned = _explicit_cuda_torch_index_url() is not None # CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU (for # example a mixed AMD+NVIDIA host that runs ROCm torch on the AMD card); - # never force CUDA wheels over that choice. + # never force CUDA wheels over that choice unless a CUDA index is pinned. _cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if _cvd is not None and _cvd.strip() in ("", "-1"): + if not _cuda_pinned and _cvd is not None and _cvd.strip() in ("", "-1"): return # Only NVIDIA hosts should carry CUDA torch. _has_usable_nvidia_gpu() # covers the /proc/driver/nvidia/gpus fallback when nvidia-smi is absent. - if not _has_usable_nvidia_gpu(): + # The explicit CUDA pin overrides the GPU-presence gate too. + if not _cuda_pinned and not _has_usable_nvidia_gpu(): return # Classify the installed torch: "hip" (ROCm build -- the poisoning @@ -1086,11 +1450,13 @@ def _ensure_cuda_torch() -> None: sys.executable, "-c", ( - "import torch; " + "import torch, re; " "hip = getattr(torch.version, 'hip', '') or ''; " "cuda = getattr(torch.version, 'cuda', '') or ''; " "ver = getattr(torch, '__version__', '').lower(); " - "print('hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'))" + "m = re.search(r'\\+(cu\\d+)', ver); " + "marker = 'hip' if (hip or 'rocm' in ver) else ('cuda' if cuda else 'cpu'); " + "print(marker + '|' + (m.group(1) if m else ''))" ), ], stdout = subprocess.PIPE, @@ -1106,16 +1472,42 @@ def _ensure_cuda_torch() -> None: _marker_lines = [ line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip() ] - if not _marker_lines or _marker_lines[-1] != "hip": - return # healthy CUDA torch, or a deliberate CPU wheel -- leave as-is + if not _marker_lines: + return + _marker, _, _installed_cu = _marker_lines[-1].partition("|") + # Reinstall CUDA torch when the venv carries a ROCm build on an NVIDIA host + # (the poisoning signature), or when an explicit CUDA index is pinned but the + # venv has the wrong family -- a CPU wheel, or a different cuXXX than pinned. + # This covers the headless cross-install (`studio update` with + # UNSLOTH_TORCH_INDEX_FAMILY=cu128): the update path preserves torch rather + # than preinstalling it from install.sh, so without this an explicit CUDA pin + # stays ineffective. A healthy CUDA torch matching the pin, or a CPU wheel + # with no CUDA pin, is deliberate and left alone. + _pin = _explicit_torch_index_url() + _pin_leaf = _pin.rstrip("/").rsplit("/", 1)[-1].lower() if _pin else "" + _pinned_cuda = _is_cuda_family_leaf(_pin_leaf) + if _marker == "hip": + _why = "torch is a ROCm build on an NVIDIA host" + elif _marker == "cpu" and _pinned_cuda: + _why = "torch is a CPU build but an explicit CUDA index is pinned" + elif _marker == "cuda" and _pinned_cuda and _installed_cu != _pin_leaf: + # Mismatch when the installed cuXXX differs from the pin. An UNTAGGED cuda + # build (empty _installed_cu -- e.g. torch re-resolved from default PyPI + # into a CUDA wheel with no +cuXXX local tag) also counts: the family + # cannot be confirmed to match the pin, so reinstall to enforce it. The + # reinstall targets the pinned family, so an already-matching untagged + # 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}" + else: + return # healthy CUDA torch matching the pin, or a deliberate CPU wheel index_url = _detect_cuda_torch_index_url() _torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC print( - f" torch is a ROCm build on an NVIDIA host -- reinstalling " - f"CUDA torch from {index_url}\n" - f" (set UNSLOTH_TORCH_BACKEND=rocm to keep a deliberate ROCm torch " - f"on a mixed AMD+NVIDIA host)" + f" {_why} -- reinstalling CUDA torch from {index_url}\n" + f" (set UNSLOTH_TORCH_BACKEND=rocm or cpu to keep a deliberate " + f"non-CUDA torch)" ) pip_install( "CUDA torch repair", @@ -1128,6 +1520,78 @@ def _ensure_cuda_torch() -> None: index_url, constrain = False, ) + _write_torch_index_marker(index_url) + + +def _ensure_cpu_torch() -> None: + """Reinstall CPU torch when an explicit CPU pin is set but the venv has a GPU build. + + Counterpart to _ensure_cuda_torch / _ensure_rocm_torch for the explicit-CPU + case (UNSLOTH_TORCH_INDEX_FAMILY=cpu or a URL ending in /cpu). Those helpers + both treat a CPU backend as a skip signal, so on a standalone `unsloth studio + update` -- which does not run install.sh's post-install flavor enforcement -- + an existing CUDA/ROCm torch satisfies the version constraint and is never + replaced, ignoring the authoritative CPU pin. Only fires for an EXPLICIT pin: + a CPU backend that came from auto-detection (genuine CPU host via install.sh) + already installed CPU wheels, so there is nothing to repair. + """ + if NO_TORCH: + return + pin = _explicit_cpu_torch_index_url() + if pin is None: + return + + # Classify the installed torch family. A non-zero exit means torch is missing + # or un-importable; the base install step handles that, so leave it alone. + try: + probe = subprocess.run( + [ + sys.executable, + "-c", + ( + "import torch, re; " + "hip = getattr(torch.version, 'hip', '') or ''; " + "cuda = getattr(torch.version, 'cuda', '') or ''; " + "ver = getattr(torch, '__version__', '').lower(); " + "gpu = bool(hip) or 'rocm' in ver or bool(cuda) or bool(re.search(r'\\+cu\\d+', ver)); " + "print('gpu' if gpu else 'cpu')" + ), + ], + stdout = subprocess.PIPE, + stderr = subprocess.DEVNULL, + timeout = 90, + ) + except (OSError, subprocess.TimeoutExpired): + return + if probe.returncode != 0: + return + _lines = [ + line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip() + ] + if not _lines or _lines[-1] != "gpu": + return # already CPU (or unreadable) -- nothing to repair + + print( + f" torch is a GPU build but an explicit CPU index is pinned -- " + f"reinstalling CPU torch from {pin}" + ) + # Pin to the supported torch<2.11 family (same bounds as the CUDA/ROCm repair + # specs). The /cpu index now also serves torch 2.11+, so a bare trio off the + # exclusive --index-url could resolve outside the supported range or drag in + # an ABI-mismatched torchvision/torchaudio. + _torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC + pip_install( + "CPU torch repair", + "--force-reinstall", + "--no-cache-dir", + _torch_pkg, + _vision_pkg, + _audio_pkg, + "--index-url", + pin, + constrain = False, + ) + _write_torch_index_marker(pin) def _ensure_rocm_torch() -> None: @@ -1242,6 +1706,7 @@ def _ensure_rocm_torch() -> None: "later to retry ROCm." ) return + _write_torch_index_marker(index_url) # ROCm torch is installed (or already was); flag it so later phases # do not overwrite it with the generic CPU torch wheel. BNB is a # separate dependency -- a BNB install failure must NOT roll back the @@ -1260,26 +1725,48 @@ def _ensure_rocm_torch() -> None: # ── Linux x86_64 only: PyTorch ROCm wheels are not published for aarch64 ── if platform.machine().lower() not in {"x86_64", "amd64"}: return - # NVIDIA takes precedence on mixed hosts -- but only if a GPU is usable - if _has_usable_nvidia_gpu(): - return - # Use _has_rocm_gpu() (rocminfo / amd-smi GPU data rows) as the - # authoritative "is this an AMD ROCm host?" signal. The old gate required - # /opt/rocm or hipcc to exist, which breaks runtime-only ROCm installs - # (minimal package-managed installs, Radeon software) that ship - # amd-smi/rocminfo without /opt/rocm or hipcc, leaving `unsloth studio - # update` unable to repair a CPU-only venv on those systems. - if not _has_rocm_gpu(): - return # no AMD GPU visible + # An explicit ROCm wheel-index pin (UNSLOTH_TORCH_INDEX_URL/_FAMILY naming a + # rocm*/gfx* leaf) commits to ROCm wheels regardless of which GPU is visible + # here -- the headless / container / CI cross-install case. Mirror + # _ensure_cuda_torch's explicit-pin bypass: skip the NVIDIA-present / no-AMD / + # unreadable-ROCm gates so the pinned index is honoured. Without this, a + # standalone `studio update` with an explicit ROCm pin on an NVIDIA-only or + # GPU-less box returned here and left the CPU/CUDA torch in place. + _rocm_pin = _explicit_rocm_torch_index_url() + if _rocm_pin is None: + # NVIDIA takes precedence on mixed hosts -- but only if a GPU is usable. + if _has_usable_nvidia_gpu(): + return + # Use _has_rocm_gpu() (rocminfo / amd-smi GPU data rows) as the + # authoritative "is this an AMD ROCm host?" signal. The old gate required + # /opt/rocm or hipcc to exist, which breaks runtime-only ROCm installs + # (minimal package-managed installs, Radeon software) that ship + # amd-smi/rocminfo without /opt/rocm or hipcc, leaving `unsloth studio + # update` unable to repair a CPU-only venv on those systems. + if not _has_rocm_gpu(): + return # no AMD GPU visible ver = _detect_rocm_version() if ver is None: - print(" ROCm detected but version unreadable -- skipping torch reinstall") - return - - # Probe whether torch already links against HIP (ROCm already working). + if _rocm_pin is None: + print(" ROCm detected but version unreadable -- skipping torch reinstall") + return + # Explicit pin: the pinned index leaf (not the host ROCm version) drives + # the install below, so a missing/unreadable host ROCm version is fine. + # Use a sentinel so the version-gated Strix reroute (skipped for explicit + # pins anyway) and any ver comparisons stay well-defined. + ver = (0, 0) + + # Probe whether torch already links against HIP (ROCm already working), and + # capture the installed ROCm build tag so a pin mismatch can be detected. # Do NOT skip for CUDA-only builds: they are unusable on AMD-only hosts # (the NVIDIA check above already handled mixed AMD+NVIDIA setups). + # Emit ONE "|" line (mirrors _ensure_cuda_torch) so the + # parse is positional and robust: the HIP marker is the field before "|" (HIP + # version, "rocm" sentinel, or empty for CPU/CUDA torch), the installed wheel + # version (e.g. "2.10.0+rocm6.4") is the field after it. Do NOT filter empty + # LINES and take slot 0 -- for CPU/CUDA torch the marker field IS empty, and + # dropping it would shift the version into slot 0 and wrongly flag has_hip_torch. try: probe = subprocess.run( [ @@ -1289,10 +1776,11 @@ def _ensure_rocm_torch() -> None: "import torch; " "hip=getattr(torch.version,'hip','') or ''; " "ver=getattr(torch,'__version__','').lower(); " - # Print the HIP version when present (back-compat), else a - # "rocm" sentinel when only torch.__version__ flags ROCm - # (AMD SDK / Radeon wheels). Empty string = CPU/CUDA. - "print(hip if hip else ('rocm' if 'rocm' in ver else ''))" + # HIP version when present (back-compat), else a "rocm" + # sentinel when only torch.__version__ flags ROCm (AMD SDK / + # Radeon wheels). Empty marker before "|" = CPU/CUDA torch. + "marker=hip if hip else ('rocm' if 'rocm' in ver else ''); " + "print(marker + '|' + ver)" ), ], stdout = subprocess.PIPE, @@ -1301,11 +1789,46 @@ def _ensure_rocm_torch() -> None: ) except (OSError, subprocess.TimeoutExpired): probe = None - has_hip_torch = ( - probe is not None and probe.returncode == 0 and probe.stdout.decode().strip() != "" + # Take the last non-empty stdout line so stray sitecustomize / import-hook + # output cannot mask the marker; then split positionally on the FIRST "|" -- + # the empty HIP-marker field for CPU/CUDA torch is preserved (has_hip_torch + # is driven by that field, not by "first non-empty line"). + _marker_lines = ( + [ln.strip() for ln in probe.stdout.decode(errors = "replace").splitlines() if ln.strip()] + if (probe is not None and probe.returncode == 0) + else [] ) + _hip_marker, _sep, _installed_torch_ver = ( + _marker_lines[-1].partition("|") if _marker_lines else ("", "", "") + ) + # A "|"-delimited marker line is required: without the separator the probe + # output is unrecognised (old torch, injected noise), so treat HIP as absent + # and fall through to a reinstall rather than trusting an ambiguous string. + has_hip_torch = bool(_sep) and _hip_marker != "" + + # An explicit ROCm pin whose family differs from the already-installed ROCm + # torch must reinstall, mirroring _ensure_cuda_torch (installed cuXXX != pin). + # Without this, `studio update` with UNSLOTH_TORCH_INDEX_FAMILY=rocm7.2 (or a + # gfx* URL) on a venv that already carries an OLDER ROCm build (+rocm6.4 / + # +rocm7.1) short-circuits on has_hip_torch and never applies the override. + # + # Prefer the torch-index MARKER when present: an EXACT compare of the pinned + # index against the index the last install recorded. This is the ONLY signal + # that catches a per-arch switch between two 2.11 gfx indexes (gfx1151 -> + # gfx120X-all) -- both install a +rocm7.13.0 wheel, so the version-tag + # heuristic below sees no difference and would leave the old arch in place. A + # matching marker also guarantees a correctly-pinned venv does NOT reinstall + # (no loop). When there is NO marker (old venv, or torch installed out-of-band) + # fall back to the +rocm/version-tag heuristic -- backward compatibility. + _rocm_pin_mismatch = False + if has_hip_torch and _rocm_pin is not None: + _marker_verdict = _marker_pin_mismatch(_rocm_pin) + if _marker_verdict is None: + _rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver) + else: + _rocm_pin_mismatch = _marker_verdict - rocm_torch_ready = has_hip_torch + rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch # Strix Halo / Strix Point (gfx1151 / gfx1150) segfault under ROCm 7.1 # in torch._grouped_mm. AMD's per-gfx repo ships torch 2.11.0+rocm7.13.0 @@ -1316,7 +1839,9 @@ def _ensure_rocm_torch() -> None: # an incompatible wheel. Use HIP_VISIBLE_DEVICES for the runtime target. _strix_override_url: "str | None" = None _strix_override_pkgs: "tuple[str, str, str] | None" = None - if ver < (7, 2): + # An explicit ROCm wheel-index pin is authoritative: never auto-reroute it to + # the AMD per-gfx index (the caller already chose the family/URL). + if ver < (7, 2) and _explicit_rocm_torch_index_url() is None: gfx_codes = _detect_amd_gfx_codes() _strix_gfx = {"gfx1151", "gfx1150"} _detected_strix = _strix_gfx.intersection(gfx_codes) @@ -1377,25 +1902,48 @@ def _ensure_rocm_torch() -> None: index_url, constrain = False, ) + _write_torch_index_marker(index_url) rocm_torch_ready = True - elif not has_hip_torch: - # Select best matching wheel tag (newest ROCm version <= installed) - tag = next( - ( - t - for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) - if ver >= (maj, mn) - ), - None, - ) - if tag is None: - print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- " f"skipping torch reinstall") + elif not has_hip_torch or _rocm_pin_mismatch: + # Reinstall when torch is not ROCm yet, OR when a ROCm build is present but + # its family differs from an explicit pin (_rocm_pin_mismatch -- the ROCm + # analogue of _ensure_cuda_torch's installed-cuXXX != pin reinstall). + # Honour an explicit ROCm wheel-index pin verbatim instead of re-detecting + # the host ROCm version; otherwise select the best wheel tag (newest ROCm + # version <= installed). gfx*/rocm7.2 indexes serve torch 2.11+, so match + # the constraints to the index leaf when overridden. + _override_idx = _explicit_rocm_torch_index_url() + if _override_idx is not None: + index_url = _override_idx + tag = index_url.rstrip("/").rsplit("/", 1)[-1].lower() else: - index_url = f"{_PYTORCH_WHL_BASE}/{tag}" - print(f" ROCm {ver[0]}.{ver[1]} -- installing torch from {index_url}") - _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get( - tag, _ROCM_TORCH_PKG_SPECS["_default"] + tag = next( + ( + t + for (maj, mn), t in sorted(_ROCM_TORCH_INDEX.items(), reverse = True) + if ver >= (maj, mn) + ), + None, ) + if tag is None: + print(f" No PyTorch wheel for ROCm {ver[0]}.{ver[1]} -- skipping torch reinstall") + else: + if _override_idx is None: + index_url = f"{_PYTORCH_WHL_BASE}/{tag}" + print(f" ROCm torch -- installing from {index_url}") + # Only the gfx arches with the _grouped_mm bug (gfx120X-all, gfx1151, + # gfx1150) need the torch 2.11 spec; other gfx per-arch indexes + # (gfx110X-all, gfx90a, gfx908) publish <2.11 wheels, so a pinned + # override to one of those stays on the default range. Matches the + # gfx floor gating in install.ps1 / setup.ps1. + if tag in _ROCM_GFX_TORCH211_LEAVES: + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["rocm7.2"] + elif tag.startswith("gfx"): + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS["_default"] + else: + _torch_pkg, _vision_pkg, _audio_pkg = _ROCM_TORCH_PKG_SPECS.get( + tag, _ROCM_TORCH_PKG_SPECS["_default"] + ) pip_install( f"ROCm torch ({tag})", "--force-reinstall", @@ -1407,6 +1955,7 @@ def _ensure_rocm_torch() -> None: index_url, constrain = False, ) + _write_torch_index_marker(index_url) rocm_torch_ready = True # Install bitsandbytes only when torch links against ROCm. Prefers the @@ -1489,6 +2038,28 @@ def _infer_no_torch() -> bool: # GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown # (standalone `unsloth studio update` runs, where we re-detect normally). _TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower() +# When install.sh did not run (standalone `unsloth studio update`) but the caller +# pinned the wheel index explicitly, derive the backend from that override so the +# CUDA/ROCm repair helpers honour it instead of re-probing the GPU and possibly +# reinstalling a different family. Classify on the final URL/family segment, +# mirroring install.sh's UNSLOTH_TORCH_BACKEND case. +if not _TORCH_BACKEND: + _idx_override = ( + os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + or os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + ) + _idx_leaf = _idx_override.rstrip("/").rsplit("/", 1)[-1].lower() + if _idx_leaf.startswith(("rocm", "gfx")): + _TORCH_BACKEND = "rocm" + elif _idx_leaf == "cpu": + _TORCH_BACKEND = "cpu" + elif _is_cuda_family_leaf(_idx_leaf): + # Require a digit after "cu" (^cu[0-9]) so a full-override URL ending in + # /current or /custom is NOT branded CUDA. A wrong "cuda" backend makes + # _ensure_rocm_torch() return early on AMD hosts and leaves a CPU/wrong + # torch unrepaired; falling through here keeps _TORCH_BACKEND="" so the + # helpers probe the GPU instead. + _TORCH_BACKEND = "cuda" def _torch_step_label(suffix: str) -> str: @@ -2218,6 +2789,8 @@ def install_python_stack() -> int: _progress(_torch_step_label("check")) _ensure_cuda_torch() _ensure_rocm_torch() + _ensure_cpu_torch() + _ensure_verbatim_torch_index() # Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python # version or unknown ROCm version). @@ -2412,6 +2985,8 @@ def _win_amd_smi_has_gpu(stdout: str) -> bool: _progress(_torch_step_label("final")) _ensure_cuda_torch() _ensure_rocm_torch() + _ensure_cpu_torch() + _ensure_verbatim_torch_index() # 14. Final check (silent; third-party conflicts are expected) subprocess.run( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index bb1e88cc4e..4e28af3d93 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -402,6 +402,205 @@ function Get-PytorchCudaTag { return "cu126" } +# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY), shared by the +# stale-venv check and the install selection below so a pinned wheel index wins +# over GPU probing -- matching install.sh, install.ps1 and install_python_stack.py. +# UNSLOTH_TORCH_INDEX_URL is verbatim (full URL); _FAMILY is the leaf (cpu, cu128, +# rocm6.4, ...) joined to the mirror base so UNSLOTH_PYTORCH_MIRROR is honoured. +function Get-PinnedTorchIndexUrl { + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { + return $env:UNSLOTH_TORCH_INDEX_URL.Trim().TrimEnd('/') + } + 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 +} + +# The last path segment of a wheel index URL (cu128 / cpu / rocm6.4 / gfx1151). +function Get-TorchIndexLeaf { + param([string]$Url) + if ([string]::IsNullOrWhiteSpace($Url)) { return $null } + return ($Url.TrimEnd('/') -split '/')[-1].ToLowerInvariant() +} + +# ── Torch-index marker ─────────────────────────────────────────────────────── +# After a successful torch install this records the exact wheel --index-url used +# at a stable per-venv path so a later `unsloth studio update` can detect a pin +# change by an EXACT string compare rather than the wheel version-tag heuristic +# (which cannot encode the AMD per-arch gfx family). Path/format MUST match +# install.sh, install.ps1 and install_python_stack.py: +# \.unsloth-torch-index (single line = the resolved index URL) +$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. +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() } + $head = $u.Substring(0, $idx) + $leaf = $u.Substring($idx + 1).ToLowerInvariant() + return "$head/$leaf" +} + +# Path to the per-venv torch-index marker. +function Get-TorchIndexMarkerPath { + param([string]$VenvDir) + if ([string]::IsNullOrWhiteSpace($VenvDir)) { return $null } + return Join-Path $VenvDir $TorchIndexMarkerName +} + +# Return the recorded torch --index-url from the marker, else $null (missing / +# empty / unreadable -> "no marker", so the caller falls back to the heuristics). +function Read-TorchIndexMarker { + param([string]$VenvDir) + $marker = Get-TorchIndexMarkerPath -VenvDir $VenvDir + if (-not $marker) { return $null } + if (-not (Test-Path -LiteralPath $marker -PathType Leaf)) { return $null } + try { + $line = (Get-Content -LiteralPath $marker -Raw -ErrorAction Stop).Trim() + } catch { return $null } + if ([string]::IsNullOrWhiteSpace($line)) { return $null } + return $line +} + +# Record the resolved torch wheel --index-url at the marker path, atomically +# (temp file + Move). Best-effort: a write failure never aborts the install. A +# blank URL is ignored (nothing to record). +function Write-TorchIndexMarker { + param([string]$VenvDir, [string]$IndexUrl) + if ([string]::IsNullOrWhiteSpace($VenvDir)) { return } + if ([string]::IsNullOrWhiteSpace($IndexUrl)) { return } + if (-not (Test-Path -LiteralPath $VenvDir -PathType Container)) { return } + $marker = Get-TorchIndexMarkerPath -VenvDir $VenvDir + $tmp = "$marker.$PID.tmp" + try { + # Write a single line, LF-terminated, no BOM (parity with the sh/py writers). + [System.IO.File]::WriteAllText($tmp, ($IndexUrl.Trim() + "`n"), (New-Object System.Text.UTF8Encoding($false))) + Move-Item -LiteralPath $tmp -Destination $marker -Force -ErrorAction Stop + } catch { + try { if (Test-Path -LiteralPath $tmp) { Remove-Item -LiteralPath $tmp -Force -ErrorAction SilentlyContinue } } catch {} + } +} + +# Compare an explicit torch-index pin against the recorded marker. +# $true -> pin resolves to a DIFFERENT index than the marker -> reinstall. +# $false -> pin matches the marker exactly -> no reinstall (no loop). +# $null -> no usable marker -> caller falls back to the version-tag heuristics. +function Test-MarkerPinMismatch { + param([string]$VenvDir, [string]$PinUrl) + $marker = Read-TorchIndexMarker -VenvDir $VenvDir + if ($null -eq $marker) { return $null } + return (Get-NormalizedIndexUrl $PinUrl) -ne (Get-NormalizedIndexUrl $marker) +} + +# The AMD per-arch index leaves that need the torch 2.11 floor (the _grouped_mm +# null-ptr bug lives in the <2.11 wheels for these arches). MUST match the +# $_pinGfx211 allowlist in the install-spec path below (and install.ps1 / +# install_python_stack.py). Other per-arch leaves (gfx110X-all/gfx90a/gfx908) +# publish <2.11 wheels and stay on default specs, so a pin to one of those must +# NOT be judged stale against the 2.11 line. +function Test-RocmGfx211Leaf { + param([string]$Leaf) + return @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $Leaf +} + +# The pytorch.org rocmX.Y versions KNOWN to ship torch 2.11 (verified against +# download.pytorch.org): rocm7.2 is the ONLY stable 2.11 rocm index today. Do NOT +# treat an unknown newer rocm (rocm7.3, rocm8.0, ...) as 2.11 speculatively -- it +# does not exist yet, and that speculative floor is the mismatch bug being fixed. +# MUST match _ROCM_KNOWN_TORCH211_VERSIONS (Python) and the rocm7.2 KNOWN-2.11 leaf +# in install.sh / install.ps1. $Major / $Minor are integers. +function Test-RocmKnown211Version { + param([int]$Major, [int]$Minor) + return ($Major -eq 7 -and $Minor -eq 2) +} + +# True only for a real CUDA wheel-family leaf: "cu" followed by digits (cu118, +# cu126, cu128, cu130, ...). Mirrors install_python_stack.py::_is_cuda_family_leaf. +# A bare -like 'cu*' wrongly matches arbitrary mirror leaves like "custom" / +# "current", which would set an expected tag no installed flavor can equal and +# rebuild the venv on every run instead of trusting the unknown custom index. +function Test-CudaFamilyLeaf { + param([string]$Leaf) + if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false } + return $Leaf -match '^cu[0-9]' +} + +# Stale-venv ROCm comparison for a pinned gfx*/rocm* index. Returns a hashtable +# @{ Expected = ; Installed = } so the caller rebuilds when they +# differ. Mirrors install_python_stack.py::_rocm_pin_family_mismatch: +# * rocmX.Y pin -> compare exact rocm versions when both readable, else the +# torch 2.11 line (rocm>=7.2 serves 2.11, older rocm does not). +# * gfx pin in the 2.11 allowlist -> expect AMD's per-arch wheel (three-part +# +rocmA.B.C local version). A generic (two-part +rocmA.B) or pre-2.11 wheel +# is stale; an already-installed per-arch wheel is NOT (no rebuild loop). +# * gfx pin NOT in the allowlist (gfx110X-all/gfx90a/gfx908) -> default <2.11 +# specs, so a 2.10+rocm wheel is correct; only a 2.11+ build is stale. +# A ROCm pin (gfx* or rocmX.Y) is satisfied ONLY by an installed wheel carrying a +# +rocm local tag: an untagged CPU/CUDA wheel (e.g. 2.10.0 / 2.11.0) never +# satisfies a ROCm pin, so it is reported stale and the pin is (re)applied. +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 } + # 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 local tag. Without it the wheel is a CPU/CUDA + # build that cannot satisfy any ROCm pin, regardless of its release line. + $_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 (the + # three-part tag already implies +rocm). + $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: default <2.11 spec. An untagged (no +rocm) wheel never + # satisfies the pin -> stale. Otherwise stale only when the build is 2.11+. + $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 + } + } + + # rocmX.Y pin. + if ($_pinVer -and $_instVer) { + # Both rocm versions readable: exact comparison. A readable $_instVer already + # implies a +rocm tag, so no separate tag check is needed here. + return @{ Expected = "rocm$_pinVer"; Installed = "rocm$_instVer" } + } + $_pinNeeds211 = $false + if ($_pinRocm.Success) { + # Only the KNOWN-2.11 rocm indexes (rocm7.2) are on the 2.11 line; an unknown + # newer rocm is NOT floored speculatively. 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, but an + # untagged (no +rocm) wheel never satisfies a rocmX.Y pin -> report it 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 { @@ -2517,6 +2716,11 @@ 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 ROCm + # version is repaired later by install_python_stack.py; here we + # only need the flavor so a correct ROCm venv is not marked stale. + $installedTorchTag = "rocm" } elseif ($torchVer -match '\+cpu') { $installedTorchTag = "cpu" } else { @@ -2536,8 +2740,86 @@ 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 + # Torch-index marker: when the last install recorded an index, compare the + # pin against it EXACTLY. This is the only signal that catches a per-arch + # switch between two 2.11 gfx indexes (gfx1151 -> gfx120X-all -- both + # install a +rocm7.13.0 wheel, so the tag heuristic sees no difference) and + # a custom-URL change (/simple -> /current). $null = no usable marker -> + # fall back to the version-tag heuristic (old venv; backward compatible). + $_markerMismatch = Test-MarkerPinMismatch -VenvDir $VenvDir -PinUrl $_pinnedIdx + if ($null -ne $_markerMismatch) { + # Drive the rebuild decision purely off the marker compare. + $_expectedKnown = $true + if ($_markerMismatch) { + $expectedTorchTag = "pinned:$_pinnedIdx" + $installedTorchTag = "marker-mismatch" + } else { + $expectedTorchTag = "pinned:$_pinnedIdx" + $installedTorchTag = "pinned:$_pinnedIdx" + } + } + # cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds. + elseif ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { + # Do NOT collapse a pinned ROCm/gfx leaf to a generic "rocm": that + # would match any installed +rocm wheel and mask a pin change from + # one ROCm family to another (e.g. rocm6.4 -> gfx1151, or rocm6.4 + # -> rocm6.3), leaving the requested index unapplied. Get-RocmPinStaleTags + # compares what the wheel tag exposes and uses the SAME 2.11 allowlist + # (Test-RocmGfx211Leaf) as the install-spec path, so a gfx110X-all / + # gfx90a / gfx908 pin on a valid <2.11 wheel is NOT judged stale, and a + # 2.11-allowlist gfx pin over a generic (two-part +rocm) wheel IS. + $_rocmTags = Get-RocmPinStaleTags -PinLeaf $_pinLeaf -TorchVersion $torchVer + $expectedTorchTag = $_rocmTags.Expected + $installedTorchTag = $_rocmTags.Installed + } elseif ((Test-CudaFamilyLeaf $_pinLeaf) -or $_pinLeaf -eq 'cpu') { + # Require digits after "cu" (cu118/cu128/...) so a mirror leaf like + # /custom or /current is NOT treated as a CUDA flavor -- otherwise the + # expected tag would be an arbitrary word no installed flavor equals, + # rebuilding the venv every run. Such leaves fall through to the + # trust-unknown-index branch below. + $expectedTorchTag = $_pinLeaf + } else { + # Custom index whose final segment is not a torch flavor (e.g. a + # PEP 503 mirror ending in /simple) and no marker to compare against. + # We cannot infer the flavor, so trust the pinned URL and do not + # rebuild on a bogus tag comparison. + $_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 the + # correct build, not stale. (gfx arch counts even when $HasROCm is + # false -- name-inferred Adrenalin hosts still get ROCm torch below.) + # Without this an unpinned ROCm venv compares "rocm" != "cpu" and is + # needlessly rebuilt, and an installer-managed setup exits as stale. + # But only the arches the install path below maps to a repo.amd.com + # wheel index (the $archFamilyMap set) actually get ROCm torch; an + # unmapped/unknown arch (e.g. name-inferred RDNA 2 gfx103X) or an + # unreadable arch -- even with $HasROCm -- falls back to CPU torch + # there (see the $archFamily null branch below). For those, expect + # "cpu" so a correct CPU venv is not marked stale and rebuilt on every + # update (or aborted under installer-managed setup). + $_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)) { + $expectedTorchTag = "rocm" + } else { + $expectedTorchTag = "cpu" + } + } else { + $expectedTorchTag = "cpu" + } + if ($_expectedKnown -and $installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { $shouldRebuild = $true } } @@ -2734,7 +3016,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" @@ -2755,7 +3043,7 @@ $ROCmIndexUrl = $null # SDK -- which flips Studio 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 @@ -2805,8 +3093,53 @@ 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). Otherwise the CUDA branch below installs bare torch / +# torchvision / torchaudio from the ROCm index and resolves a known-bad <2.11 +# wheel or ABI-mismatched companions for gfx115x / gfx120x / rocm>=7.2. +if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) { + $_pinLeaf = Get-TorchIndexLeaf $PinnedTorchIndexUrl + $_pinRocm211 = $false + if ($_pinLeaf -match '^rocm(\d+)\.(\d+)') { + # Only KNOWN-2.11 rocm indexes (rocm7.2) get the 2.11 floor; do not floor an + # unknown newer rocm speculatively. Matches install.sh's rocm7.2 KNOWN-2.11 + # leaf and Test-RocmKnown211Version / _ROCM_KNOWN_TORCH211_VERSIONS. + $_pinRocm211 = Test-RocmKnown211Version -Major ([int]$Matches[1]) -Minor ([int]$Matches[2]) + } + # Only the gfx families the AMD arch map above pins to torch 2.11 need the + # floor here (gfx120X-all, gfx1151, gfx1150 -- the _grouped_mm bug arches). + # Other per-arch indexes (gfx110X-all, gfx90a, gfx908) publish <2.11 wheels + # and the automatic path leaves them bare, so an override to one of those + # must NOT force a 2.11 floor the normal path intentionally avoids. 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 ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { + # Other gfx per-arch indexes and older rocm (<=7.1) ship torch <2.11; + # route via the ROCm path with bare specs (matches the automatic path's + # bare floor for these arches). + $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 UNSLOTH_TORCH_INDEX_URL pin is used verbatim; a family pin already set +# $CuTag, so $PyTorchWhlBase/$CuTag is the requested family index. The CPU/CUDA +# install branches below pull from this instead of re-joining mirror + tag. +# A pinned ROCm install goes through $ROCmIndexUrl; if that fails, the fallback +# must use the CPU wheel index, not retry the ROCm mirror left in the pin. +$TorchInstallIndexUrl = if ($ROCmIndexUrl) { "$PyTorchWhlBase/cpu" } elseif ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" } + $ROCmCpuFallback = $false if ($ROCmIndexUrl) { substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..." @@ -2833,21 +3166,25 @@ 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. + # The $ROCmCpuFallback term matters when a PINNED ROCm index failed: $CuTag is + # still the rocm/gfx leaf (not "cpu"), so without it execution would fall through + # to the CUDA branch and install from the CPU index WITHOUT --force-reinstall, + # leaving the partial ROCm torch in place. # 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. $cpuForce = @() if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") } if ($script:UnslothVerbose) { - Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" + Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl $torchInstallExit = $LASTEXITCODE $output = "" } else { - $output = Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" | Out-String + $output = Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl | Out-String $torchInstallExit = $LASTEXITCODE } if ($torchInstallExit -ne 0) { @@ -2859,11 +3196,11 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") { substep "installing PyTorch with CUDA support ($CuTag)..." substep "(This download is ~2.8 GB -- may take a few minutes)" if ($script:UnslothVerbose) { - Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" + Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl $torchInstallExit = $LASTEXITCODE $output = "" } else { - $output = Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" | Out-String + $output = Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl | Out-String $torchInstallExit = $LASTEXITCODE } if ($torchInstallExit -ne 0) { @@ -2890,6 +3227,16 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") { } } +# ── Record the resolved torch wheel index (marker) ── +# Torch was just installed; write the exact --index-url used so a later `unsloth +# studio update` can detect a pin change by an EXACT string compare (see the +# stale-venv check above) instead of the version-tag heuristic. Reflects the actual +# installed family: $ROCmIndexUrl when the ROCm path ran, else $TorchInstallIndexUrl +# (which is the pinned/CUDA/CPU index, or the CPU index after a ROCm fallback). +# Path/format matches install.sh, install.ps1 and install_python_stack.py. +$MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } else { $TorchInstallIndexUrl } +Write-TorchIndexMarker -VenvDir $VenvDir -IndexUrl $MarkerIndexUrl + # No unsloth.exe rename needed. setup.ps1 runs *via* unsloth.exe, so renaming the # running launcher only ever failed (WinError 32) and printed a scary warning. It's # also unnecessary: install.ps1 sets SKIP_STUDIO_BASE=1 (base never reinstalled) and diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 666ea7ce10..d3c1f02709 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -10,6 +10,8 @@ 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,282 @@ def test_install_ps1_preserves_timeout_override(self): 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 (parity with install.sh's + # get_torch_index_url override), so the Python CUDA repair must let the pin + # clear the CUDA_VISIBLE_DEVICES hide gate, not just the NVIDIA-presence + # gate. Otherwise CVD=-1 UNSLOTH_TORCH_INDEX_FAMILY=cu128 studio update + # (the GPU-less CI case) would bail before repairing. + 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) + # The CVD hide-gate return must be guarded by the CUDA-pin flag. + 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 not install a bare torch trio: the /cpu + # index now also serves torch 2.11+, so a bare install off the exclusive + # --index-url can resolve outside the repo's supported <2.11 range or pull + # an ABI-mismatched companion. It must use the bounded CPU/CUDA spec. + 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-venv check must only expect ROCm torch for arches the install + # path actually maps to a repo.amd.com wheel index. An unmapped arch + # (name-inferred RDNA 2 gfx103X) or an unreadable arch installs CPU torch, + # so expecting "rocm" there marks a correct CPU venv stale and rebuilds it + # every update (or aborts under installer-managed setup). + 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 leaves that carry the torch 2.11 floor (gfx120X-all / + gfx1151 / gfx1150) must be the SAME set in every installer AND in each + installer's stale/mismatch check. When these diverged, a pinned gfx110X-all / + gfx90a / gfx908 wheel (which stays <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 the + # install-spec path must reuse it, so the stale check and install spec can + # never disagree again. + 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") + assert re.search( + r'r"\^cu\[0-9\]"', text + ), "install_python_stack.py _is_cuda_family_leaf must match ^cu[0-9]" + + def test_setup_ps1_requires_cu_digit(self): + text = SETUP_PS1.read_text(encoding = "utf-8") + assert re.search( + r"'\^cu\[0-9\]'", text + ), "setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9], not a bare cu* glob" + # 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") + # The UNSLOTH_TORCH_BACKEND export must brand CUDA only on cu[0-9]* -- a + # bare catch-all *) -> cuda would mis-brand /current, /custom mirror pins + # as CUDA and make the stack skip ROCm repair on AMD hosts (comment #2's + # bug via install.sh instead of standalone studio update). + 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 TestTorchIndexMarkerParity: + """All four installers must agree on the torch-index marker (PR #6692): + the same filename, the same write points, and the same read helpers.""" + + MARKER = ".unsloth-torch-index" + + def test_all_installers_use_same_marker_filename(self): + # The exact marker filename must appear in every installer so bash / py / + # ps write and read the same per-venv path. + for path, label in ( + (INSTALL_SH, "install.sh"), + (INSTALL_PS1, "install.ps1"), + (SETUP_PS1, "setup.ps1"), + (STACK_PY, "install_python_stack.py"), + ): + text = path.read_text(encoding = "utf-8") + assert self.MARKER in text, f"{label} must reference the marker '{self.MARKER}'" + + def test_all_installers_write_the_marker(self): + # install.sh + PowerShell use a _write_torch_index_marker / Write-TorchIndexMarker + # helper; the Python stack calls _write_torch_index_marker at its install sites. + assert "_write_torch_index_marker" in INSTALL_SH.read_text(encoding = "utf-8") + assert "Write-TorchIndexMarker" in INSTALL_PS1.read_text(encoding = "utf-8") + assert "Write-TorchIndexMarker" in SETUP_PS1.read_text(encoding = "utf-8") + assert "_write_torch_index_marker" in STACK_PY.read_text(encoding = "utf-8") + + def test_setup_ps1_and_python_read_the_marker(self): + # The repair/update side (setup.ps1 stale check, python _ensure_rocm_torch) + # must READ the marker to make the exact pin-change decision. + assert "Read-TorchIndexMarker" in SETUP_PS1.read_text(encoding = "utf-8") + assert "Test-MarkerPinMismatch" in SETUP_PS1.read_text(encoding = "utf-8") + stack = STACK_PY.read_text(encoding = "utf-8") + assert "_read_torch_index_marker" in stack + assert "_marker_pin_mismatch" in stack + + def test_all_installers_normalize_index_url(self): + # The exact-compare normalization (trim, strip trailing slash, lowercase + # leaf) must exist in every language so the compare is identical. + assert "_normalize_index_url" in INSTALL_SH.read_text(encoding = "utf-8") + assert "Get-NormalizedIndexUrl" in SETUP_PS1.read_text(encoding = "utf-8") + assert "_normalize_index_url" in STACK_PY.read_text(encoding = "utf-8") + + def test_marker_written_atomically(self): + # bash uses a temp file + mv; PowerShell a temp file + Move-Item; Python + # tempfile + os.replace. Confirm the atomic-write intent in each. + assert re.search(r"mv -f \"\$_wm_tmp\"", INSTALL_SH.read_text(encoding = "utf-8")) + for ps in (INSTALL_PS1, SETUP_PS1): + assert "Move-Item" in ps.read_text(encoding = "utf-8") + assert "os.replace" in STACK_PY.read_text(encoding = "utf-8") + + +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_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}" diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index 6656142625..0c538f0a12 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -379,6 +379,50 @@ _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 -> that +# family (not the cpu fallback detection would pick). +_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 wins over a real detection: a host whose nvidia-smi reports +# 12.6 still gets cu128. This is the exact Docker-build case (the builder sees +# the host driver but must publish 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" + rm -f "$_FUNC_FILE" rm -rf "$_FAKE_SMI_DIR" rm -rf "$_TOOLS_DIR" diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh index 293a709360..7528c8c9db 100644 --- a/tests/sh/test_torch_constraint.sh +++ b/tests/sh/test_torch_constraint.sh @@ -271,6 +271,62 @@ 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 must lowercase _torch_index_leaf before the floor case, +# so the canonical AMD RDNA4 leaf 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 the exact normalization + floor case from install.sh and +# assert both gfx120X-all (capital X, canonical) and gfx120x-all get the floor, +# while non-2.11 leaves (gfx110X-all, rocm6.4, cu128, cpu) 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 # ====================================================================== diff --git a/tests/sh/test_torch_index_marker.sh b/tests/sh/test_torch_index_marker.sh new file mode 100755 index 0000000000..c8f63a5dc7 --- /dev/null +++ b/tests/sh/test_torch_index_marker.sh @@ -0,0 +1,92 @@ +#!/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 torch-index MARKER helpers (_normalize_index_url, +# _write_torch_index_marker). These converge the ROCm/gfx pin-change detection +# across install.sh / install_python_stack.py / setup.ps1 / install.ps1 by +# recording the resolved wheel --index-url so a later update can compare exactly. +# Helpers are extracted from install.sh and sourced (parity with test_torch_flavor.sh). +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +PASS=0 +FAIL=0 + +_TORCH_INDEX_MARKER_NAME=".unsloth-torch-index" + +# Extract the marker helpers from install.sh and source them. +_FUNC_FILE=$(mktemp) +{ + sed -n '/^_normalize_index_url()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_write_torch_index_marker()/,/^}/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 +} + +echo "=== _normalize_index_url ===" +assert_eq "trailing slashes stripped + leaf lowered" \ + "https://repo.amd.com/rocm/whl/gfx120x-all" \ + "$(_normalize_index_url 'https://repo.amd.com/rocm/whl/gfx120X-all///')" +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" \ + "$(_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" \ + "$(_normalize_index_url 'https://repo.amd.com/rocm/whl/gfx120x-all')" \ + "$(_normalize_index_url 'https://repo.amd.com/rocm/whl/gfx120X-all')" +assert_eq "empty -> empty" "" "$(_normalize_index_url ' ')" +# rocm7.2 KNOWN-2.11 leaf normalises to itself. +assert_eq "rocm7.2 unchanged" \ + "https://download.pytorch.org/whl/rocm7.2" \ + "$(_normalize_index_url 'https://download.pytorch.org/whl/rocm7.2/')" + +echo "=== _write_torch_index_marker ===" +_VD=$(mktemp -d) +_write_torch_index_marker "$_VD" "https://download.pytorch.org/whl/rocm7.2" +assert_eq "marker written verbatim (single line)" \ + "https://download.pytorch.org/whl/rocm7.2" \ + "$(cat "$_VD/$_TORCH_INDEX_MARKER_NAME" 2>/dev/null)" + +# Overwrite (per-arch switch gfx1151 -> gfx120X-all): the marker is replaced. +_write_torch_index_marker "$_VD" "https://repo.amd.com/rocm/whl/gfx120X-all" +assert_eq "marker overwritten on re-install" \ + "https://repo.amd.com/rocm/whl/gfx120X-all" \ + "$(cat "$_VD/$_TORCH_INDEX_MARKER_NAME" 2>/dev/null)" + +# Blank URL is ignored (nothing meaningful to record) -- existing marker kept. +_write_torch_index_marker "$_VD" " " +assert_eq "blank url leaves prior marker intact" \ + "https://repo.amd.com/rocm/whl/gfx120X-all" \ + "$(cat "$_VD/$_TORCH_INDEX_MARKER_NAME" 2>/dev/null)" + +# No stray temp files left behind by the atomic write. +_leftover=$(find "$_VD" -maxdepth 1 -name "$_TORCH_INDEX_MARKER_NAME.*.tmp" 2>/dev/null | wc -l | tr -d ' ') +assert_eq "no stray temp file left" "0" "$_leftover" + +# Missing venv dir -> no-op, no error, no file created. +_MISSING="$_VD/does_not_exist_dir" +_write_torch_index_marker "$_MISSING" "https://x/cu128" +assert_eq "missing venv dir -> no marker" \ + "absent" \ + "$([ -e "$_MISSING/$_TORCH_INDEX_MARKER_NAME" ] && echo present || echo absent)" + +rm -rf "$_VD" + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index cea4383268..8f064f1a03 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -64,21 +64,36 @@ 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": return smi_path return None + import tempfile as _tempfile + + # Isolate the torch-index marker so _ensure_cuda_torch's post-reinstall marker + # write lands in a throwaway dir, never the real venv (and never leaks state). + _marker_dir = _tempfile.mkdtemp(prefix = "unsloth-marker-test-") + _marker_path = Path(_marker_dir) / ".unsloth-torch-index" + with ( patch.object(stack_mod, "_TORCH_BACKEND", backend), patch.object(stack_mod, "IS_MACOS", is_macos), @@ -87,6 +102,7 @@ def _which(name, *a, **k): patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = nvidia), patch.object(stack_mod.shutil, "which", side_effect = _which), patch.object(stack_mod.os.path, "isfile", return_value = bool(smi_path)), + patch.object(stack_mod, "_torch_index_marker_path", return_value = _marker_path), patch.object(stack_mod, "pip_install") as mock_pip, patch.object( stack_mod.subprocess, @@ -99,6 +115,10 @@ def _which(name, *a, **k): 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 @@ -128,6 +148,57 @@ def test_rocm_in_version_string_triggers_repair(self): 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 / container / CI cross-install: an explicit cu* index pin + # commits to CUDA wheels even though no NVIDIA GPU is visible here, 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): + # CUDA_VISIBLE_DEVICES=-1 hides the GPU, but an explicit cu* pin skips ALL + # host-GPU probing (like install.sh's get_torch_index_url override), so the + # CVD hide gate must not suppress the repair. This is the exact GPU-less CI + # case the override targets: CVD=-1 UNSLOTH_TORCH_INDEX_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 (torch re-resolved from default PyPI, no +cuXXX + # local tag -> empty installed cu) cannot be confirmed to match the pin, + # so the pin is enforced with a reinstall to the pinned family. + 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) + # No-op cases. @@ -191,6 +262,110 @@ def test_cvd_explicit_device_still_repairs(self): 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 generic mirror URL whose leaf starts with "cu" but is not cuXXX + # (e.g. .../custom, .../current) must NOT be treated as a CUDA pin, so it + # cannot bypass the NVIDIA gate to force CUDA over a CPU/ROCm venv. + 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 exact derivation the module does at import time, using the + # module's own _is_cuda_family_leaf so this stays in lockstep with it. + 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. diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index c8f2053946..10d198fd85 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -557,6 +557,17 @@ def test_hipconfig_timeout(self, tmp_path): class TestEnsureRocmTorch: """Verify ROCm torch reinstall logic.""" + @pytest.fixture(autouse = True) + def _isolate_torch_index_marker(self, tmp_path): + """Point the torch-index marker at a per-test tmp path so _ensure_rocm_torch's + marker writes never touch the real venv and stale markers never leak between + tests. The file is ABSENT by default -> these tests exercise the no-marker + fallback to the +rocm/version-tag heuristic (backward compatibility). The + path is exposed as self._marker_path for tests that seed a marker.""" + self._marker_path = tmp_path / ".unsloth-torch-index" + with patch.object(stack_mod, "_torch_index_marker_path", return_value = self._marker_path): + yield + @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) def test_no_rocm_skips(self, mock_nvidia, mock_pip): @@ -569,19 +580,27 @@ def test_no_rocm_skips(self, mock_nvidia, mock_pip): _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 + # New 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 +610,32 @@ def test_torch_already_has_hip_skips(self, mock_ver, mock_gpu, mock_nvidia, mock """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, "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): + """Regression: a CPU torch probe emits an EMPTY HIP marker before "|"; the + positional parse must keep that empty field so has_hip_torch stays False + (an earlier parse dropped the empty line and shifted the version into the + marker slot, wrongly reporting HIP and skipping the ROCm reinstall).""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"|2.10.0+cpu\n" + # has_hip_torch False -> the AMD host reinstalls ROCm; assert we do NOT skip. + with patch("os.path.isdir", return_value = True): + with patch.object(stack_mod, "pip_install_try", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _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 +719,239 @@ def test_rocm_72_selects_72_tag(self, mock_ver, mock_gpu, mock_nvidia, mock_pip) 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 + + @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 + + @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_installed_perarch_no_reinstall( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """A gfx1151 pin over an already-installed AMD per-arch (+rocm7.13.0) wheel + must NOT reinstall torch -- once the correct per-arch wheel is present the + pin is satisfied, so `studio update` does not reinstall-loop.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"7.13.0|2.11.0+rocm7.13.0\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() + assert not any( + any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list + ) + + 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 + # 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 is a + # mismatch -- the user wants AMD's per-arch (three-part) wheel, not generic. + 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 (no reinstall loop); + # a 2.11 build is (the arch's index does not publish 2.11 wheels). + 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 + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @@ -732,6 +1004,267 @@ def test_torch_backend_cpu_env_skips_entirely(self, mock_nvidia, mock_gpu, mock_ mock_pip.assert_not_called() +# TEST: install_python_stack.py -- torch-index MARKER mechanism (PR #6692) + + +class TestTorchIndexMarkerHelpers: + """Pure marker helpers: normalization, read/write round-trip, exact compare.""" + + def test_normalize_index_url_lowercases_only_leaf(self): + f = stack_mod._normalize_index_url + # Trailing slashes stripped; ONLY the final segment lowercased. + 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" + # Whitespace trimmed. + assert f(" https://download.pytorch.org/whl/cu128 ") == ( + "https://download.pytorch.org/whl/cu128" + ) + # gfx120X-all (capital X) and AMD's lowercase pip leaf compare equal. + assert f("https://repo.amd.com/rocm/whl/gfx120X-all") == ( + f("https://repo.amd.com/rocm/whl/gfx120x-all") + ) + # Empty / whitespace-only -> None. + assert f(" ") is None + assert f(None) is None + + def test_marker_write_read_round_trip(self, tmp_path): + marker = tmp_path / ".unsloth-torch-index" + with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker): + stack_mod._write_torch_index_marker("https://repo.amd.com/rocm/whl/gfx1151") + # Recorded verbatim (single stripped line). + assert marker.read_text().strip() == ("https://repo.amd.com/rocm/whl/gfx1151") + assert stack_mod._read_torch_index_marker() == ("https://repo.amd.com/rocm/whl/gfx1151") + + def test_marker_write_is_atomic_and_ignores_blank(self, tmp_path): + marker = tmp_path / ".unsloth-torch-index" + with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker): + # Blank / whitespace-only is ignored (nothing to record). + stack_mod._write_torch_index_marker("") + stack_mod._write_torch_index_marker(" ") + assert not marker.exists() + # No stray temp files left behind by the atomic write. + stack_mod._write_torch_index_marker("https://x/cu128") + leftovers = [p.name for p in tmp_path.iterdir() if p.name != ".unsloth-torch-index"] + assert leftovers == [] + + def test_missing_or_corrupt_marker_reads_as_absent(self, tmp_path): + marker = tmp_path / ".unsloth-torch-index" + with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker): + # Absent. + assert stack_mod._read_torch_index_marker() is None + # Empty file -> absent. + marker.write_text("") + assert stack_mod._read_torch_index_marker() is None + # Whitespace-only -> absent. + marker.write_text(" \n") + assert stack_mod._read_torch_index_marker() is None + + def test_marker_pin_mismatch_exact_compare(self, tmp_path): + marker = tmp_path / ".unsloth-torch-index" + with patch.object(stack_mod, "_torch_index_marker_path", return_value = marker): + # No marker -> None (caller falls back to the heuristic). + assert stack_mod._marker_pin_mismatch("https://repo.amd.com/rocm/whl/gfx1151") is None + # Marker gfx1151, pin gfx120X-all -> mismatch (True). This is the exact + # per-arch switch the version-tag heuristic cannot see (#2543). + marker.write_text("https://repo.amd.com/rocm/whl/gfx1151\n") + assert ( + stack_mod._marker_pin_mismatch("https://repo.amd.com/rocm/whl/gfx120X-all") is True + ) + # Marker matches the pin (case-insensitive leaf, trailing slash) -> False. + assert stack_mod._marker_pin_mismatch("https://repo.amd.com/rocm/whl/gfx1151/") is False + + def test_known_211_versions_only_rocm72(self): + # KNOWN-2.11 rocm set is exactly {rocm7.2} -- rocm7.1/rocm7.3 are NOT in it. + known = stack_mod._ROCM_KNOWN_TORCH211_VERSIONS + assert (7, 2) in known + assert (7, 1) not in known + assert (7, 3) not in known + assert (8, 0) not in known + + def test_rocm_pin_family_mismatch_known_211_set(self): + # The rocmX.Y unreadable-installed fallback uses the KNOWN-2.11 set, so a + # (hypothetical) rocm7.3 pin is NOT treated as the 2.11 line speculatively. + f = stack_mod._rocm_pin_family_mismatch + base = "https://download.pytorch.org/whl" + # rocm7.3 pin (unknown -> <2.11 line) over an unreadable-version +rocm wheel + # that is <2.11: not a mismatch (both on the non-2.11 line). + assert f(f"{base}/rocm7.3", "2.10.0+rocm") is False + # rocm7.2 pin (KNOWN-2.11) over the same <2.11 unreadable wheel: mismatch. + assert f(f"{base}/rocm7.2", "2.10.0+rocm") is True + + +class TestEnsureRocmTorchMarker: + """_ensure_rocm_torch marker integration (the #2543 per-arch switch fix).""" + + @pytest.fixture(autouse = True) + def _marker(self, tmp_path): + self._marker_path = tmp_path / ".unsloth-torch-index" + with patch.object(stack_mod, "_torch_index_marker_path", return_value = self._marker_path): + yield + + def _seed(self, url): + self._marker_path.write_text(url + "\n") + + @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_marker_gfx_switch_reinstalls_despite_same_wheel_tag( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """Marker records gfx1151; user re-pins to gfx120X-all. Both indexes install + a +rocm7.13.0 per-arch wheel, so the version-tag heuristic sees NO difference + and would leave the old arch in place. The marker's exact compare catches it + and reinstalls from the new gfx index (#2543).""" + self._seed("https://repo.amd.com/rocm/whl/gfx1151") + mock_probe = MagicMock() + mock_probe.returncode = 0 + # has_hip_torch True + an already-installed AMD per-arch (three-part) wheel: + # _rocm_pin_family_mismatch(gfx120X-all, ...+rocm7.13.0) would return False. + mock_probe.stdout = b"7.13.0|2.11.0+rocm7.13.0\n" + env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx120X-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): + 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 "gfx120X-all" in torch_call or "gfx120x-all" in torch_call + # Marker rewritten to the new index. + assert "gfx120X-all" in self._marker_path.read_text() + + @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_marker_matches_pin_no_reinstall( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """Marker matches the current pin exactly -> NO torch reinstall (no loop), + even when the wheel tag would otherwise look ambiguous.""" + self._seed("https://repo.amd.com/rocm/whl/gfx1151") + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"7.13.0|2.11.0+rocm7.13.0\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() + # No torch reinstall (marker matches). bnb-only calls may still happen. + 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_no_marker_falls_back_to_heuristic_no_forced_reinstall( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """NO marker + a healthy ROCm wheel that already satisfies the pin -> the + heuristic (_rocm_pin_family_mismatch) decides, and a correct venv is NOT + force-reinstalled (backward compatibility for old venvs).""" + # No marker seeded (autouse fixture leaves the file absent). + mock_probe = MagicMock() + mock_probe.returncode = 0 + # rocm6.4 pin over an installed +rocm6.4 wheel -> heuristic says no mismatch. + mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n" + env = {"UNSLOTH_TORCH_INDEX_FAMILY": "rocm6.4"} + 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() + 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") + def test_verbatim_custom_url_reinstall_on_marker_mismatch(self, mock_pip): + """Marker records .../simple; user pins a custom .../current index (leaf is + neither rocm/gfx/cu/cpu). _ensure_verbatim_torch_index reinstalls torch + VERBATIM from that URL -- "URL wins verbatim" (#2544).""" + self._seed("https://mirror.local/simple") + env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/current"} + with patch.object(stack_mod, "NO_TORCH", False): + with patch.object(stack_mod, "IS_MACOS", False): + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + stack_mod._ensure_verbatim_torch_index() + assert mock_pip.call_count == 1 + call = str(mock_pip.call_args_list[0]) + assert "https://mirror.local/current" in call + assert "torch" in call + # Marker rewritten to the pinned custom URL. + assert "current" in self._marker_path.read_text() + + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install") + def test_verbatim_custom_url_no_marker_is_noop(self, mock_pip): + """No marker + a custom-URL pin -> _ensure_verbatim_torch_index does NOTHING + (an old venv must not be blindly force-reinstalled from an unverified URL).""" + env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/current"} + with patch.object(stack_mod, "NO_TORCH", False): + with patch.object(stack_mod, "IS_MACOS", False): + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + stack_mod._ensure_verbatim_torch_index() + mock_pip.assert_not_called() + + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install") + def test_verbatim_custom_url_matching_marker_no_reinstall(self, mock_pip): + """Marker matches the custom URL pin -> no reinstall (idempotent, no loop).""" + self._seed("https://mirror.local/current") + env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/current/"} + with patch.object(stack_mod, "NO_TORCH", False): + with patch.object(stack_mod, "IS_MACOS", False): + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + stack_mod._ensure_verbatim_torch_index() + mock_pip.assert_not_called() + + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install") + def test_verbatim_skips_known_family_pins(self, mock_pip): + """A known-family pin (rocm/gfx/cu/cpu) is NOT handled by the verbatim path -- + the dedicated _ensure_{rocm,cuda,cpu} helpers own those. cu128 stays CUDA.""" + self._seed("https://download.pytorch.org/whl/cu126") + for pin in ( + "https://download.pytorch.org/whl/cu128", + "https://download.pytorch.org/whl/rocm7.2", + "https://repo.amd.com/rocm/whl/gfx1151", + "https://download.pytorch.org/whl/cpu", + ): + mock_pip.reset_mock() + env = {"UNSLOTH_TORCH_INDEX_URL": pin} + with patch.object(stack_mod, "NO_TORCH", False): + with patch.object(stack_mod, "IS_MACOS", False): + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + stack_mod._ensure_verbatim_torch_index() + mock_pip.assert_not_called() + + # TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard @@ -2846,6 +3379,26 @@ def test_torch_constraint_updated_for_strix_amd_index(self): 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 on the + # full $TORCH_INDEX_URL (which the earlier, buggy version matched with + # */gfx* and would false-positive on a mirror base path). Only the gfx + # families with the <2.11 _grouped_mm bug (gfx120X-all / gfx1151 / gfx1150) + # are pushed to 2.11 -- a bare gfx* would also floor gfx110X-all/gfx90a/ + # gfx908, which the automatic AMD path intentionally leaves bare. + 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") diff --git a/tests/studio/test_setup_pin_stale.ps1 b/tests/studio/test_setup_pin_stale.ps1 new file mode 100644 index 0000000000..7335044987 --- /dev/null +++ b/tests/studio/test_setup_pin_stale.ps1 @@ -0,0 +1,101 @@ +#!/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 -- no GPU/venv needed. Mirrors the Python +# _rocm_pin_family_mismatch / _is_cuda_family_leaf tests so both stay in lockstep. +# 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")) +# rocm pin vs an untagged (no +rocm) wheel: a CPU/CUDA build never satisfies a +# ROCm pin, regardless of its release line -> always stale (needs reinstall). +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 (no +rocm) 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 leaf is lowercased by Get-TorchIndexLeaf before this helper; +# the caller passes the normalised leaf, so the 2.11-allowlist branch fires and a +# generic/untagged wheel is stale (the per-arch wheel is not). +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") + +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 (both on the non-2.11 line); +# rocm7.2 (KNOWN-2.11) over the same wheel IS stale. This is the #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 diff --git a/tests/studio/test_torch_index_marker.ps1 b/tests/studio/test_torch_index_marker.ps1 new file mode 100644 index 0000000000..9cd12ad248 --- /dev/null +++ b/tests/studio/test_torch_index_marker.ps1 @@ -0,0 +1,101 @@ +#!/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 studio/setup.ps1's torch-index MARKER helpers +# (Get-NormalizedIndexUrl, Read-TorchIndexMarker, Write-TorchIndexMarker, +# Test-MarkerPinMismatch, Test-RocmKnown211Version). These converge the ROCm/gfx +# pin-change detection across install.sh / install_python_stack.py / setup.ps1 / +# install.ps1. Pure helpers, AST-extracted and run in-process -- no GPU/venv. +# Run: pwsh -NoProfile -File tests/studio/test_torch_index_marker.ps1 + +$ErrorActionPreference = "Stop" +$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1") +$setupPath = (Resolve-Path $setupPath).Path + +# Parse setup.ps1 (also 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" } + +# The marker filename constant the helpers close over. +$TorchIndexMarkerName = ".unsloth-torch-index" + +foreach ($name in @( + "Get-NormalizedIndexUrl", "Get-TorchIndexMarkerPath", "Read-TorchIndexMarker", + "Write-TorchIndexMarker", "Test-MarkerPinMismatch", "Test-RocmKnown211Version" +)) { + $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 "Get-NormalizedIndexUrl (trim / strip trailing slash / lowercase leaf)" +Check "trailing slashes + leaf lowered" ` + ((Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all///") -eq "https://repo.amd.com/rocm/whl/gfx120x-all") +Check "whitespace trimmed" ` + ((Get-NormalizedIndexUrl " https://download.pytorch.org/whl/cu128 ") -eq "https://download.pytorch.org/whl/cu128") +Check "host case preserved, leaf lowered" ` + ((Get-NormalizedIndexUrl "https://Mirror.Local/Simple/") -eq "https://Mirror.Local/simple") +Check "gfx120X-all == gfx120x-all after normalize" ` + ((Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all") -eq (Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120x-all")) +Check "empty -> null" ($null -eq (Get-NormalizedIndexUrl " ")) + +Write-Host "Write/Read-TorchIndexMarker (round trip, atomic, blank ignored)" +$venv = Join-Path ([System.IO.Path]::GetTempPath()) ("unsloth-marker-" + [System.Guid]::NewGuid().ToString("N")) +New-Item -ItemType Directory -Path $venv | Out-Null +try { + Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://repo.amd.com/rocm/whl/gfx1151" + Check "marker written verbatim" ` + ((Read-TorchIndexMarker -VenvDir $venv) -eq "https://repo.amd.com/rocm/whl/gfx1151") + # Per-arch switch overwrites the marker. + Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all" + Check "marker overwritten on re-install" ` + ((Read-TorchIndexMarker -VenvDir $venv) -eq "https://repo.amd.com/rocm/whl/gfx120X-all") + # Blank URL ignored -- prior marker kept. + Write-TorchIndexMarker -VenvDir $venv -IndexUrl " " + Check "blank url leaves prior marker intact" ` + ((Read-TorchIndexMarker -VenvDir $venv) -eq "https://repo.amd.com/rocm/whl/gfx120X-all") + # No stray temp file left behind. + $tmpLeft = @(Get-ChildItem -LiteralPath $venv -Filter "$TorchIndexMarkerName.*.tmp" -ErrorAction SilentlyContinue).Count + Check "no stray temp file left" ($tmpLeft -eq 0) + + Write-Host "Test-MarkerPinMismatch (exact compare; null when no marker)" + # Marker gfx120X-all now recorded. A gfx1151 pin differs -> mismatch (#2543). + Check "gfx marker vs gfx1151 pin -> mismatch" ` + ((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://repo.amd.com/rocm/whl/gfx1151") -eq $true) + # Same pin (trailing slash, case) -> not a mismatch (no reinstall loop). + Check "same pin (slash/case) -> no mismatch" ` + ((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://repo.amd.com/rocm/whl/gfx120x-all/") -eq $false) + # Custom URL change /simple -> /current is a mismatch (#2544). + Write-TorchIndexMarker -VenvDir $venv -IndexUrl "https://mirror.local/simple" + Check "custom /simple marker vs /current pin -> mismatch" ` + ((Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://mirror.local/current") -eq $true) + + Write-Host "Read-TorchIndexMarker (missing / empty -> null)" + Remove-Item -LiteralPath (Get-TorchIndexMarkerPath -VenvDir $venv) -Force + Check "missing marker -> null" ($null -eq (Read-TorchIndexMarker -VenvDir $venv)) + Set-Content -LiteralPath (Get-TorchIndexMarkerPath -VenvDir $venv) -Value "" -NoNewline + Check "empty marker -> null" ($null -eq (Read-TorchIndexMarker -VenvDir $venv)) + Check "no marker -> Test-MarkerPinMismatch null" ` + ($null -eq (Test-MarkerPinMismatch -VenvDir $venv -PinUrl "https://x/rocm7.2")) +} finally { + Remove-Item -LiteralPath $venv -Recurse -Force -ErrorAction SilentlyContinue +} + +Write-Host "Test-RocmKnown211Version (KNOWN-2.11 set == rocm7.2 only)" +Check "rocm7.2 -> true" (Test-RocmKnown211Version -Major 7 -Minor 2) +Check "rocm7.1 -> false" (-not (Test-RocmKnown211Version -Major 7 -Minor 1)) +Check "rocm7.3 -> false" (-not (Test-RocmKnown211Version -Major 7 -Minor 3)) +Check "rocm8.0 -> false" (-not (Test-RocmKnown211Version -Major 8 -Minor 0)) + +Write-Host "" +if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 } +Write-Host "All checks passed" -ForegroundColor Green