diff --git a/install.ps1 b/install.ps1 index df49414620..6e059ee0dd 100644 --- a/install.ps1 +++ b/install.ps1 @@ -53,7 +53,8 @@ function Install-UnslothStudio { param([string]$TorchIndexUrl) if ($SkipTorch) { return "none" } if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" } - $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + # Drop query/fragment first so a token-authenticated pin classifies by family. + $leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf } if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf } return "auto" @@ -62,7 +63,8 @@ function Install-UnslothStudio { function Get-TauriGpuBranch { param([string]$TorchIndexFamily) if ($SkipTorch) { return "no_torch" } - if ($TorchIndexFamily -like "cu*") { return "cuda" } + # Require a digit after "cu" so /current or /custom isn't branded CUDA (parity ^cu[0-9]). + if ($TorchIndexFamily -match '^cu[0-9]') { return "cuda" } if ($TorchIndexFamily -like "rocm*") { return "rocm" } if ($TorchIndexFamily -eq "cpu") { return "cpu" } return "unknown" @@ -467,22 +469,35 @@ function Install-UnslothStudio { } } + # Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer + # output before printing on failure; uv/pip errors echo the failing --index-url verbatim. + # Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. + function Redact-InstallOutput { + param([string]$Text) + if (-not $Text) { return $Text } + $Text = $Text -replace '(https?://)[^/@\s`]+@', '$1@' + $Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=' + # A #token=... fragment is as sensitive as a query; URL-anchored. + return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#' + } + # Run native commands quietly by default to match install.sh behavior. # Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1. function Invoke-InstallCommand { param( [Parameter(Mandatory = $true)][ScriptBlock]$Command ) - # Installer-pinned index installs (torch) must beat an inherited uv mirror - # (#6898): when the command pins an index, clear every uv index env var so - # it wins, then restore in finally. Other installs keep the user's mirror. + # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898): + # for --default-index, clear the uv index env vars (restore in finally) and set + # UV_NO_CONFIG=1 so a uv.toml/pyproject index can't outrank the CLI pin (uv 0.10). $savedUvIndex = $null if ($Command.ToString() -match '--default-index') { $savedUvIndex = @{} - foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'UV_CONFIG_FILE', 'UV_NO_CONFIG') { $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n) Remove-Item "Env:$n" -ErrorAction SilentlyContinue } + $env:UV_NO_CONFIG = '1' } $prevEap = $ErrorActionPreference $ErrorActionPreference = "Continue" @@ -493,17 +508,23 @@ function Install-UnslothStudio { # Merge stderr into stdout so progress/warning output stays visible # without flipping $? on successful native commands (PS 5.1 treats # stderr records as errors that set $? = $false even on exit code 0). - & $Command 2>&1 | Out-Host + # Redact per record: uv echoes index URLs (credentials and all) in + # its errors, and verbose mode must not bypass the quiet path's + # redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched. + & $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host } else { $output = & $Command 2>&1 | Out-String if ($LASTEXITCODE -ne 0) { - Write-Host $output -ForegroundColor Red + Write-Host (Redact-InstallOutput $output) -ForegroundColor Red } } return [int]$LASTEXITCODE } finally { $ErrorActionPreference = $prevEap - if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } } + if ($savedUvIndex) { + Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue + foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } + } } } @@ -1960,10 +1981,31 @@ exit 0 # On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint. if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint } + # Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL + # TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. + function Trim-IndexPathSlashes { + param([string]$Url) + $value = $Url.Trim() + $idx = $value.IndexOfAny([char[]]@('?', '#')) + if ($idx -lt 0) { + return $value.TrimEnd('/') + } + return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx) + } + # ── Choose the correct PyTorch index URL based on driver CUDA version ── # Mirrors Get-PytorchCudaTag in setup.ps1. function Get-TorchIndexUrl { $baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } + # Explicit pin -- skip ALL GPU probing (headless / CI / cross-install). + # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf appended + # to the mirror base. Matches install.sh / install_python_stack.py. + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { + return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL) + } + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) { + return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))" + } if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" } try { $output = Invoke-NvidiaSmiBounded $NvidiaSmiExe @@ -1984,6 +2026,25 @@ exit 0 return "$baseUrl/cu126" } + # Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with + # _strip_index_url_credentials (install.sh / py / setup.ps1). + function Remove-IndexUrlCredentials { + param([string]$Url) + $sep = $Url.IndexOf('://') + if ($sep -lt 0) { return $Url } + $scheme = $Url.Substring(0, $sep) + $rest = $Url.Substring($sep + 3) + # Drop query / fragment (may hold auth tokens). + $q = $rest.IndexOfAny([char[]]('?', '#')) + if ($q -ge 0) { $rest = $rest.Substring(0, $q) } + $slash = $rest.IndexOf('/') + $authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest } + $at = $authority.LastIndexOf('@') + $host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority } + if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" } + return "${scheme}://${host_}" + } + # ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ── # torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu, # matching setup.ps1's stale-venv parse. @@ -2002,11 +2063,13 @@ exit 0 param([string]$TorchIndexUrl, [string]$ROCmIndexUrl) if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' } if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null } - $leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant() + # Drop query/fragment first so .../cu128?token=x classifies as cu128 (else it reinstalls every run). + $leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() if ($leaf -match '^cu\d+$') { return $leaf } if ($leaf -eq 'cpu') { return 'cpu' } if ($leaf -match '^rocm') { return 'rocm' } - if ($leaf -match '^gfx') { return 'rocm' } + # gfx must be followed by a digit (an architecture leaf); gfx-private is custom. + if ($leaf -match '^gfx[0-9]') { return 'rocm' } return $null } @@ -2041,6 +2104,10 @@ exit 0 } catch { return $null } } + # An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it + # (e.g. a deliberate cpu pin on an AMD host). + $TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or ` + (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) $TorchIndexUrl = Get-TorchIndexUrl # ── GPU arch → newest compatible Windows ROCm wheel release ── @@ -2052,7 +2119,9 @@ exit 0 # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. $ROCmIndexUrl = $null $ROCmTorchFloor = $null - if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { + $PinnedRocmVisionSpec = $null + $PinnedRocmAudioSpec = $null + if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 @@ -2102,6 +2171,32 @@ exit 0 } } + # A gfx*/rocm pin skips the auto-reroute above, but the generic CPU/CUDA install below + # would use torch>=2.4,<2.11 and pull a known-bad wheel on the gfx115x/gfx120x/rocm>=7.2 + # indexes (the _grouped_mm bug). Route a pinned ROCm index through the ROCm path. + if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) { + $_pinLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLower() + $_pinRocm211 = $false + # Anchor ($) so a suffixed custom leaf (rocm7.2-private) falls through to verbatim. + if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') { + # Only KNOWN-2.11 rocm (rocm7.2) gets the floor. Matches Test-RocmKnown211Version. + $_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2) + } + # Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare. + $_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf + if ($_pinGfx211 -or $_pinRocm211) { + $ROCmIndexUrl = $TorchIndexUrl + $ROCmTorchFloor = "torch>=2.11.0,<2.12.0" + $PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0" + $PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0" + substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan" + } elseif ($_pinLeaf -match '^gfx[0-9]' -or $_pinLeaf -match '^rocm[0-9]+(\.[0-9]+)?$') { + # Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with + # bare specs. Only EXACT rocm/gfx* are families; a suffixed leaf is verbatim. + $ROCmIndexUrl = $TorchIndexUrl + } + } + if ($ROCmIndexUrl) { $TorchIndexFamily = "rocm" } else { @@ -2164,8 +2259,8 @@ exit 0 } if ($_Migrated) { - # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state - # in the new venv location, while preserving existing torch/CUDA + # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving + # existing torch/CUDA unless the flavor repair below re-lands it. Write-TauriLog "STEP" "Installing unsloth" substep "upgrading unsloth in migrated environment..." if ($SkipTorch) { @@ -2210,22 +2305,24 @@ exit 0 substep "skipping PyTorch (--no-torch flag set)." "Yellow" } elseif ($ROCmIndexUrl) { Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)" - substep "installing PyTorch from $ROCmIndexUrl..." + substep "installing PyTorch from $(Remove-IndexUrlCredentials $ROCmIndexUrl)..." $torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin the companions to match $torchSpec; bare names can resolve an # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { - # Transient AMD-index failure: fall back to a CPU base so the install - # still completes; Unsloth setup retries ROCm afterwards. + # Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries + # ROCm). Use an explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS + # the ROCm mirror, so reusing it would just retry it. + $CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" } substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow" # --force-reinstall: a failed ROCm install can leave an unpinned ROCm # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU # torch>= range, so without it uv would keep the ROCm build and only swap # the companions -- a mismatched venv the flavor-repair block won't fix. - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2238,8 +2335,14 @@ exit 0 } } else { Write-TauriLog "STEP" "Installing PyTorch" - substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } + substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..." + # Bound the companions to the capped torch on EVERY index, cu + # families included: torchaudio 2.11 dropped its exact torch pin from + # the wheel metadata, so a bare companion next to torch<2.11 can + # resolve a mismatched 2.11.0 build. Mirrors install.sh. + $_pinVisionSpec = "torchvision>=0.19,<0.26.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2335,8 +2438,8 @@ exit 0 $rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" } # Pin companions like the fresh ROCm path (bare names can pull an # ABI-incompatible torchvision/torchaudio from the per-arch index). - $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } - $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } + $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } + $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { @@ -2347,7 +2450,7 @@ exit 0 } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) diff --git a/install.sh b/install.sh index 7918a2bd23..c02552628f 100755 --- a/install.sh +++ b/install.sh @@ -159,18 +159,58 @@ run_maybe_quiet() { fi } +# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL +# strip corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. +_trim_index_path_slashes() { + _tips_v="$1" + case "$_tips_v" in + *[?#]*) + _tips_head="${_tips_v%%[?#]*}" + _tips_tail="${_tips_v#"$_tips_head"}" + ;; + *) + _tips_head="$_tips_v" + _tips_tail="" + ;; + esac + while [ -n "$_tips_head" ] && [ "${_tips_head%/}" != "$_tips_head" ]; do + _tips_head="${_tips_head%/}" + done + printf '%s%s' "$_tips_head" "$_tips_tail" +} + +# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer +# output before printing on failure; uv/pip errors echo the failing --index-url verbatim. +# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. +_redact_install_output() { + sed -E \ + -e 's#(https?://)[^/@[:space:]`]+@#\1@#g' \ + -e 's#([?&][^=[:space:]&`]+)=[^&#[:space:]`]+#\1=#g' \ + -e 's|(https?://[^[:space:]`#]+)#[^[:space:]`]+|\1#|g' \ + "$@" +} + run_install_cmd() { _label="$1" shift - # Installer-pinned index installs (torch) must beat an inherited uv mirror - # (#6898): when we pass --default-index, neutralize every uv index env var so - # the pinned index wins. Other installs keep the user's mirror. + # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898): + # for --default-index, neutralize the uv index/backend/config vars (UV_TORCH_BACKEND + # redirects torch; UV_NO_CONFIG=1 + dropping UV_CONFIG_FILE stops a uv.toml/pyproject + # index outranking the CLI pin, uv 0.10). case " $* " in - *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;; + *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL -u UV_TORCH_BACKEND -u UV_FIND_LINKS -u UV_CONFIG_FILE UV_NO_CONFIG=1 "$@" ;; esac if _is_verbose; then - "$@" && return 0 - _rc=$? + # Stream through the redactor: uv echoes index URLs (credentials and + # all) in its errors, and verbose mode previously bypassed the + # redaction the quiet path applies. The rc file preserves the + # command's exit code across the pipe without relying on pipefail + # (this script runs under plain sh). + _rcf=$(mktemp) + { "$@" 2>&1; printf '%s' "$?" > "$_rcf"; } | _redact_install_output + _rc=$(cat "$_rcf" 2>/dev/null || echo 1) + rm -f "$_rcf" + [ "${_rc:-1}" -eq 0 ] 2>/dev/null && return 0 step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 return "$_rc" fi @@ -178,7 +218,7 @@ run_install_cmd() { "$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; } _rc=$? step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 - cat "$_log" >&2 + _redact_install_output "$_log" >&2 rm -f "$_log" return $_rc } @@ -257,7 +297,7 @@ _install_bnb_rocm() { fi _bnb_rc=$? if _is_verbose; then - cat "$_bnb_log" >&2 + _redact_install_output "$_bnb_log" >&2 fi rm -f "$_bnb_log" step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2 @@ -310,6 +350,11 @@ _tauri_torch_index_family() { return fi _diag_url="${1:-}" + # Strip query/fragment AND a trailing slash before classifying (like _torch_index_url_leaf): + # a token isn't echoed into [TAURI:DIAG], and .../cu128/?token=x still classifies as cu128. + _diag_url="${_diag_url%%\?*}" + _diag_url="${_diag_url%%#*}" + _diag_url="${_diag_url%/}" case "$_diag_url" in */cu118) echo "cu118" ;; */cu124) echo "cu124" ;; @@ -343,7 +388,8 @@ _tauri_gpu_branch() { return fi case "$_diag_family" in - cu*) echo "cuda" ;; + # Require a digit after cu so /current or /custom isn't branded CUDA (parity ^cu[0-9]). + cu[0-9]*) echo "cuda" ;; rocm*) if [ "$_diag_radeon" = true ]; then echo "rocm_radeon" @@ -1575,6 +1621,12 @@ _has_usable_nvidia_gpu() { # the STUDIO_HOME mkdir/venv so the origin distro is untouched. _maybe_reroute_strixhalo_to_2404() { [ "${OS:-}" = "wsl" ] || return 0 + # An explicit index pin skips every GPU-driven reroute (same contract as + # the later Radeon/Strix guard): the pin is honored in THIS distro rather + # than probing the GPU and switching distributions. Whitespace-only + # overrides do not gate (parity with get_torch_index_url). + _rr_pin=$(printf '%s' "${UNSLOTH_TORCH_INDEX_URL:-}${UNSLOTH_TORCH_INDEX_FAMILY:-}" | tr -d '[:space:]') + [ -n "$_rr_pin" ] && return 0 [ "${SKIP_TORCH:-false}" = "false" ] || return 0 [ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0 [ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0 @@ -1636,6 +1688,10 @@ _maybe_reroute_strixhalo_to_2404() { # Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the # GPU instead of falling back to the desktop-app prompt path. [ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1" + # Forward a pinned torch index into the rerouted distro; dropping it would + # silently revert the child install to auto-detection. + [ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_URL=$(_rr_q "$UNSLOTH_TORCH_INDEX_URL")" + [ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_FAMILY=$(_rr_q "$UNSLOTH_TORCH_INDEX_FAMILY")" [ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1" _rr_args="" [ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")" @@ -2001,6 +2057,15 @@ if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; t TORCH_CONSTRAINT="torch>=2.6,<2.11.0" fi fi +# Companion (torchvision/torchaudio) constraints, bounded to torch's window. +# torchaudio 2.11 dropped its exact torch pin, so a bare companion next to a +# <2.11-capped torch resolves torchaudio 2.11 (verified: cpu leaf installed +# torch 2.10.0+cpu with torchaudio 2.11.0+cpu). torchvision still exact-pins +# torch and self-corrects, but is bounded for symmetry. Widened alongside the +# cu* torch window below; the torch-2.11 AMD paths (rocm7.2 / per-gfx / Strix) +# pin their own trio. +TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" +TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" # ── Resolve repo root (for --local installs) ── _REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)" @@ -2069,6 +2134,24 @@ _has_amd_rocm_gpu() { get_torch_index_url() { _base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}" _base="${_base%/}" + # Explicit override -- skip ALL GPU probing (headless / container / CI / cross-install). + # UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf (cpu, cu128, ...) + # appended to the mirror base. Trim whitespace so a whitespace-only value is unset. + _url="${UNSLOTH_TORCH_INDEX_URL:-}" + _url="${_url#"${_url%%[![:space:]]*}"}"; _url="${_url%"${_url##*[![:space:]]}"}" + if [ -n "$_url" ]; then + # Trim trailing PATH slashes (a multi-slash path 404s on strict pip proxies) while + # preserving a ?query/#fragment token (a whole-URL strip would eat a "/"-ending token). + _url=$(_trim_index_path_slashes "$_url") + echo "$_url"; return + fi + _family="${UNSLOTH_TORCH_INDEX_FAMILY:-}" + _family="${_family#"${_family%%[![:space:]]*}"}"; _family="${_family%"${_family##*[![:space:]]}"}" + if [ -n "$_family" ]; then + while [ "${_family#/}" != "$_family" ]; do _family="${_family#/}"; done + while [ "${_family%/}" != "$_family" ]; do _family="${_family%/}"; done + echo "$_base/$_family"; return + fi # macOS: always CPU (no CUDA support) case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac # Try nvidia-smi -- require the binary to actually list a usable GPU. @@ -2197,6 +2280,45 @@ _torch_flavor_tag() { esac } +# Final path segment of a wheel index URL ($1), lowercased, query/fragment stripped first +# so a token-authenticated pin (.../cu128?token=x) classifies as cu128 (else it reinstalls +# every update). Classification only. Shared with the py / ps1 leaf extractors. +_torch_index_url_leaf() { + _tl_u="${1%%\?*}" + _tl_u="${_tl_u%%#*}" + # Strip ALL trailing slashes, not one: .../rocm7.2// must yield rocm7.2, not an empty leaf. + while [ -n "$_tl_u" ] && [ "${_tl_u%/}" != "$_tl_u" ]; do + _tl_u="${_tl_u%/}" + done + printf '%s' "${_tl_u##*/}" | tr '[:upper:]' '[:lower:]' +} + +# True (exit 0) when a lowercased leaf is an EXACT pip ROCm family: rocm[.] +# or a gfx ARCHITECTURE leaf (gfx followed by a digit: gfx90a, gfx1151, gfx120x-all). A leaf +# that merely starts with rocm/gfx (rocm7.2-private, gfx-private) is a custom verbatim pin. +# Matches the py / ps1 sides. +_is_pip_rocm_family_leaf() { + case "$1" in + gfx[0-9]*) return 0 ;; + rocm[0-9]*) + # Exact rocm[.]: both major and minor must be non-empty all-digits + # (rocm7., rocm7.2.1, rocm7.2-private are all custom pins, not a family). + _rocm_rest="${1#rocm}" + case "$_rocm_rest" in + *.*.*) return 1 ;; + *.*) + _rocm_minor="${_rocm_rest#*.}" + case "${_rocm_rest%%.*}" in "" | *[!0-9]*) return 1 ;; esac + case "$_rocm_minor" in "" | *[!0-9]*) return 1 ;; esac + ;; + *[!0-9]*) return 1 ;; + esac + return 0 + ;; + *) return 1 ;; + esac +} + # Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2 # ("torch>=A.B[.C], # rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. _expected_torch_flavor_tag() { - _u="${1%/}" - _leaf="${_u##*/}" + _leaf=$(_torch_index_url_leaf "$1") case "$_leaf" in - cu[0-9]*) echo "$_leaf" ;; - cpu) echo "cpu" ;; - rocm*|gfx*) echo "rocm" ;; - *) echo "" ;; + cu[0-9]*) + # Exact cu + digits only; a cu*-suffixed leaf (cu128-private) -> "" (custom), + # else a correct +cu128 wheel is force-reinstalled every run. + case "${_leaf#cu}" in + *[!0-9]*) echo "" ;; + *) echo "$_leaf" ;; + esac + ;; + cpu) echo "cpu" ;; + # Exact rocm/gfx families only; a custom rocm*-suffixed leaf -> "" (custom). + *) + if _is_pip_rocm_family_leaf "$_leaf"; then echo "rocm"; else echo ""; fi + ;; esac } @@ -2308,14 +2438,42 @@ _expected_torch_flavor_tag() { # fresh-install paths above already use -- so a stale wheel is auto-repairable. # Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. _torch_index_repairable() { - _u="${1%/}" - _leaf="${_u##*/}" + _leaf=$(_torch_index_url_leaf "$1") case "$_leaf" in - cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;; - *) echo "no" ;; + cu[0-9]*) echo "yes" ;; + # Only EXACT rocm/gfx families resolve via --default-index; a suffixed leaf is verbatim. + *) + if _is_pip_rocm_family_leaf "$_leaf"; then echo "yes"; else echo "no"; fi + ;; esac } +# Remove credentials from a wheel index URL ($1) so an authenticated pin never leaks: +# drops userinfo AND query/fragment; scheme/host/path stay exact. Shared with py / ps1. +_strip_index_url_credentials() { + _sic_url="$1" + case "$_sic_url" in + *://*) ;; + *) printf '%s' "$_sic_url"; return ;; + esac + _sic_scheme="${_sic_url%%://*}" + _sic_rest="${_sic_url#*://}" + # Drop query / fragment (may hold auth tokens). + _sic_rest="${_sic_rest%%\?*}" + _sic_rest="${_sic_rest%%#*}" + _sic_auth="${_sic_rest%%/*}" + # Drop user:pass@ userinfo if present. + case "$_sic_auth" in + *@*) _sic_host="${_sic_auth##*@}" ;; + *) _sic_host="$_sic_auth" ;; + esac + if [ "$_sic_auth" = "$_sic_rest" ]; then + printf '%s://%s' "$_sic_scheme" "$_sic_host" + else + printf '%s://%s/%s' "$_sic_scheme" "$_sic_host" "${_sic_rest#*/}" + fi +} + get_radeon_wheel_url() { # Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing # contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/, @@ -2561,7 +2719,19 @@ _maybe_bootstrap_rocm_wsl() { [ -n "$_rw_tmp" ] && rm -f "$_rw_tmp" return 0 } -_maybe_bootstrap_rocm_wsl || true +# When the caller pins the wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY), honour it +# everywhere: skip the WSL ROCm bootstrap and the Radeon/Strix reroute below (which would +# re-probe the GPU and overwrite the pin). Trim whitespace first (parity with +# get_torch_index_url): a whitespace-only override is unset there, so must not flip this true. +_torch_index_pinned=false +_ti_url_trim="${UNSLOTH_TORCH_INDEX_URL:-}" +_ti_url_trim="${_ti_url_trim#"${_ti_url_trim%%[![:space:]]*}"}"; _ti_url_trim="${_ti_url_trim%"${_ti_url_trim##*[![:space:]]}"}" +_ti_family_trim="${UNSLOTH_TORCH_INDEX_FAMILY:-}" +_ti_family_trim="${_ti_family_trim#"${_ti_family_trim%%[![:space:]]*}"}"; _ti_family_trim="${_ti_family_trim%"${_ti_family_trim##*[![:space:]]}"}" +if [ -n "$_ti_url_trim" ] || [ -n "$_ti_family_trim" ]; then + _torch_index_pinned=true +fi +[ "$_torch_index_pinned" = true ] || _maybe_bootstrap_rocm_wsl || true TORCH_INDEX_URL=$(get_torch_index_url) @@ -2572,29 +2742,74 @@ TORCH_INDEX_URL=$(get_torch_index_url) # whose base path happens to contain "rocm" or "gfx" must not mislabel a # cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix # overrides in gfxNNNN/, so the trailing slash is stripped first). -_torch_index_leaf="${TORCH_INDEX_URL%/}" +# Lowercase the leaf so every gfx*/rocm*/cu* arm matches regardless of case (canonical AMD +# RDNA4 leaf is gfx120X-all). CUDA is branded only on a real cu[0-9]* leaf, so a mirror +# leaf (/current) does NOT commit a CUDA backend; an unknown leaf leaves the var unset so +# the stack probes the GPU. Query/fragment dropped first, then ALL trailing slashes (in +# lockstep with the shared _torch_index_url_leaf extractor). +_torch_index_leaf="${TORCH_INDEX_URL%%\?*}" +_torch_index_leaf="${_torch_index_leaf%%#*}" +# Strip ALL trailing slashes, not one: .../cu128// must yield cu128, not an empty leaf. +while [ -n "$_torch_index_leaf" ] && [ "${_torch_index_leaf%/}" != "$_torch_index_leaf" ]; do + _torch_index_leaf="${_torch_index_leaf%/}" +done _torch_index_leaf="${_torch_index_leaf##*/}" +_torch_index_leaf=$(printf '%s' "$_torch_index_leaf" | tr '[:upper:]' '[:lower:]') case "$_torch_index_leaf" in rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;; cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;; - *) export UNSLOTH_TORCH_BACKEND="cuda" ;; + cu[0-9]*) export UNSLOTH_TORCH_BACKEND="cuda" ;; + # Unknown leaf (odd mirror, /current): unset so a stale inherited value can't leak and + # the stack probes the GPU. + *) unset UNSLOTH_TORCH_BACKEND ;; esac -# rocm7.2 and the CUDA cu12x/cu13x indexes now ship torch 2.11.x, so widen the -# ceiling to <2.12.0 (matches the base image and _CUDA_TORCH_PKG_SPEC in -# studio/install_python_stack.py). Keep the >=2.4 floor so an older CUDA index -# (e.g. cu118) still resolves. Match on _torch_index_leaf, not the full URL, so -# a mirror whose base path contains cu*/rocm7.2 but resolves to a cpu/older-rocm -# leaf keeps the default <2.11.0. +# Whether TORCH_INDEX_URL names an actual pip ROCm family (rocm* / gfx*), gating the +# ROCm-only side effects below (AMD bitsandbytes, ROCm-torch repair). Digit-gated so a leaf +# merely STARTING with "rocm" isn't force-repaired from the wrong path. +if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then + _torch_index_is_rocm_family=true +else + _torch_index_is_rocm_family=false +fi + +# rocm7.2 and the per-gfx indexes with the _grouped_mm <2.11 bug (gfx120X-all, gfx1151, +# gfx1150) ship torch 2.11.0 -- raise the floor (also covers a pinned override that skipped +# the Strix reroute). Pin the companions too: the per-gfx index publishes them independently +# and a bare name can resolve a 2.12 ABI-mismatched wheel. Match on the FINAL leaf so a +# custom mirror with a gfx/rocm7.2 path segment but a cu*/cpu family isn't forced. case "$_torch_index_leaf" in - rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; - cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;; + rocm7.2|gfx120x-all|gfx1151|gfx1150) + TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" + ;; + # CUDA cu12x/cu13x indexes ship torch 2.11.x: widen the ceiling to <2.12.0 (matches + # _CUDA_TORCH_PKG_SPEC) and widen the companions with it so the trio stays paired. + cu[0-9]*) + TORCH_CONSTRAINT="torch>=2.4,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0" + ;; esac +# A pinned custom/unknown-leaf index (/simple, /current, /cu128-private) has no curated +# companion set, so bound torchvision/torchaudio to the same <2.11 range the Python path pins +# (else a mirror with newer companions resolves a 2.12 ABI-mismatched wheel). Known families +# keep their curated companions above (_expected_torch_flavor_tag returns "" only for custom). +if [ "$_torch_index_pinned" = true ] && \ + [ -z "$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")" ]; then + TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0" +fi + # Auto-detect GPU for AMD ROCm based # get_torch_index_url must have chosen */rocm* # (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". +# Skipped when the index is pinned: an explicit override must not be rerouted to the +# Radeon/Strix repos by GPU probing. _amd_gpu_radeon=false +if [ "$_torch_index_pinned" = false ]; then case "$TORCH_INDEX_URL" in */rocm*) if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \ @@ -2671,10 +2886,14 @@ case "$TORCH_INDEX_URL" in done TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/" TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + # Pin companions to 2.11 (per-gfx index publishes them independently). + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" _amd_gpu_radeon=false fi ;; esac +fi # _torch_index_pinned guard (Radeon + Strix reroute) # Re-run over an existing install: keep the previous venv's torch RELEASE; the fresh # index above supplies the right flavor for this machine. Evaluated HERE, after every # index/constraint decision including the Strix reroute, so the window checked is the @@ -2821,7 +3040,7 @@ case "$TORCH_INDEX_URL" in if [ "$_amd_gpu_radeon" = true ]; then substep "wheels: repo.radeon.com (Radeon)" else - substep "wheels: $TORCH_INDEX_URL" + substep "wheels: $(_strip_index_url_credentials "$TORCH_INDEX_URL")" fi ;; esac @@ -2867,8 +3086,8 @@ for _p in ('torch', 'torchvision', 'torchaudio'): } if [ "$_MIGRATED" = true ]; then - # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state - # in the new venv location, while preserving existing torch/CUDA + # Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving + # existing torch/CUDA unless the ROCm repair below fires. substep "upgrading unsloth in migrated environment..." if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps (current @@ -2909,18 +3128,14 @@ if [ "$_MIGRATED" = true ]; then # AMD ROCm: install bitsandbytes even in migrated environments so # existing ROCm installs gain the AMD bitsandbytes build without a # fresh reinstall. - if [ "$SKIP_TORCH" = false ]; then - case "$TORCH_INDEX_URL" in - */rocm*|*/gfx*) - _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" - # Repair ROCm torch if overwritten during migrated install - _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) - if [ -z "$_has_hip" ]; then - substep "repairing ROCm torch (overwritten by dependency resolution)..." - _install_torch_default_index --force-reinstall - fi - ;; - esac + if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" + # Repair ROCm torch if overwritten during migrated install + _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) + if [ -z "$_has_hip" ]; then + substep "repairing ROCm torch (overwritten by dependency resolution)..." + _install_torch_default_index --force-reinstall + fi fi elif [ -n "$TORCH_INDEX_URL" ]; then # Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac) @@ -3074,7 +3289,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \ [ "$_radeon_versions_match" != true ]; then - substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" + substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN" _install_torch_default_index else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." @@ -3095,7 +3310,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi fi else - substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" + substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN" _install_torch_default_index fi else @@ -3103,19 +3318,15 @@ elif [ -n "$TORCH_INDEX_URL" ]; then _install_torch_default_index fi else - substep "installing PyTorch ($TORCH_INDEX_URL)..." + substep "installing PyTorch ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))..." _install_torch_default_index fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm # host stays in GGUF-only mode rather than pulling in bitsandbytes, # which is only useful once torch is present for training. - if [ "$SKIP_TORCH" = false ]; then - case "$TORCH_INDEX_URL" in - */rocm*|*/gfx*) - _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" - ;; - esac + if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then + _install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY" fi # Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed tauri_log "STEP" "Installing Unsloth" @@ -3161,16 +3372,12 @@ elif [ -n "$TORCH_INDEX_URL" ]; then _UNSLOTH_TORCH_OVERRIDES="" # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. - if [ "$SKIP_TORCH" = false ]; then - case "$TORCH_INDEX_URL" in - */rocm*|*/gfx*) - _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) - if [ -z "$_has_hip" ]; then - substep "repairing ROCm torch (overwritten by dependency resolution)..." - _install_torch_default_index --force-reinstall - fi - ;; - esac + if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then + _has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true) + if [ -z "$_has_hip" ]; then + substep "repairing ROCm torch (overwritten by dependency resolution)..." + _install_torch_default_index --force-reinstall + fi fi else # Fallback: GPU detection failed to produce a URL -- let uv resolve torch @@ -3217,7 +3424,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" - substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" \"$TORCHVISION_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --default-index $(_strip_index_url_credentials "$TORCH_INDEX_URL") --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" fi fi fi diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 95c9356d4a..9921b83543 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -44,11 +44,10 @@ IS_MAC_INTEL = IS_MACOS and platform.machine() == "x86_64" IS_MAC_ARM = IS_MACOS and platform.machine() == "arm64" IS_LINUX = sys.platform.startswith("linux") -# DiskPart-prompt suppression: amd-smi auto-elevates on Windows, popping a -# UAC/DiskPart prompt mid-install. This installer only spawns probes and pip/uv -# (none need elevation), so set __COMPAT_LAYER=RunAsInvoker process-wide -- every -# amd-smi subprocess then runs un-elevated, no per-call guard needed. setup.ps1 -# keeps per-call guards since it ALSO spawns winget installers that need elevation. +# amd-smi auto-elevates on Windows (UAC/DiskPart prompt mid-install). This installer +# only spawns probes and pip/uv (no elevation), so set __COMPAT_LAYER=RunAsInvoker +# process-wide; amd-smi then runs un-elevated. setup.ps1 keeps per-call guards (it +# also spawns winget installers that need elevation). if IS_WINDOWS: os.environ.setdefault("__COMPAT_LAYER", "RunAsInvoker") # torchcodec ships wheels only for manylinux_2_28_x86_64, macosx_12_0_arm64, @@ -74,6 +73,14 @@ _ROCM_TORCH_INDEX: dict[tuple[int, int], str] = { (6, 0): "rocm6.0", } +# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug). +# Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare. +_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"}) + +# pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (rocm7.2 only today); don't +# floor an unknown newer rocm speculatively. Match install.sh / setup.ps1 / install.ps1. +_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": ( @@ -81,18 +88,16 @@ _ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = { "torchvision>=0.26.0,<0.27.0", "torchaudio>=2.11.0,<2.12.0", ), - # Default for rocm7.1 and earlier: torch 2.x below 2.11 + # rocm7.1 and earlier: torch 2.x below 2.11 "_default": ( "torch>=2.4,<2.11.0", "torchvision>=0.19,<0.26.0", "torchaudio>=2.4,<2.11.0", ), } -# Windows AMD per-arch companion pins for the repo.amd.com index, mirroring the -# install.ps1 / setup.ps1 floor maps (gfx120X and Strix Halo/Point use the rocm7.2 -# torch 2.11 trio). Pinning the companions keeps AMD's per-arch index -- which -# publishes each independently -- from resolving an ABI-mismatched one. Unlisted -# arches have no published floor, so stay bare. Bump with the PS maps at 2.12.x. +# Windows AMD per-arch companion pins for the repo.amd.com index (mirrors the install.ps1 / +# setup.ps1 floor maps): pinning stops the per-arch index (each published independently) from +# resolving an ABI-mismatched companion. Unlisted arches have no floor, so stay bare. _WINDOWS_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = { "gfx1201": _ROCM_TORCH_PKG_SPECS["rocm7.2"], "gfx1200": _ROCM_TORCH_PKG_SPECS["rocm7.2"], @@ -103,19 +108,79 @@ _PYTORCH_WHL_BASE = ( os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl" ).rstrip("/") -# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed: its -# torchao 0.17 cpp kernels load cleanly (0.16 crashes on cu130), and the flash-attn -# / causal-conv1d / mamba torch2.10 wheels load and pass their upstream suites on -# 2.11 (see wheel_utils._PREBUILT_WHEEL_TORCH_MM). torchvision/torchaudio are pinned -# (not bare) because the install uses an exclusive --index-url (no PyPI fallback), so -# a bare name could resolve one built against a different torch major (e.g. 0.27 for -# torch 2.12) and fail at runtime with an ABI mismatch. + +def _strip_index_url_credentials(url: str) -> str: + """Strip userinfo (user:password@) AND query/fragment from a wheel index URL. + + An authenticated pin must not leak credentials in printed output; query/fragment + may hold tokens and aren't part of the PEP 503 index identity. Host/path stay + exact. MUST match install.sh / setup.ps1 / install.ps1. + """ + scheme, sep, rest = url.partition("://") + if not sep: + return url + rest = rest.split("?", 1)[0].split("#", 1)[0] # drop query / fragment + authority, slash, tail = rest.partition("/") + host = authority.rpartition("@")[2] # drop user:pass@ userinfo + return f"{scheme}://{host}{slash}{tail}" + + +_URL_USERINFO_RE = re.compile(r"(https?://)[^/@\s`]+@") +_URL_QUERY_VALUE_RE = re.compile(r"([?&][^=\s&`]+)=[^&#\s`]+") +# URL-anchored so a bare "#..." (a shell comment in tool output) is never touched. +_URL_FRAGMENT_RE = re.compile(r"(https?://[^\s`#]+)#[^\s`]+") + + +def _redact_install_output(output: "bytes | str") -> str: + """Redact index-URL credentials (userinfo + query values + fragments) from captured + installer output before printing. uv/pip failure text embeds the failing --index-url + verbatim, which can carry a user:token@, ?token= or #token= secret. MUST match + install.sh / setup.ps1 / install.ps1's output sanitizers.""" + text = output.decode(errors = "replace") if isinstance(output, bytes) else output + text = _URL_USERINFO_RE.sub(r"\1@", text) + text = _URL_QUERY_VALUE_RE.sub(r"\1=", text) + return _URL_FRAGMENT_RE.sub(r"\1#", text) + + +def _trim_index_path_slashes(url: str) -> str: + """Trim trailing slashes from the URL PATH only, preserving ?query / #fragment. A + whole-URL rstrip("/") corrupts a token that ends in "/" (e.g. base64 ...abc/) and a + single-slash strip leaves .../cu128// classifying as an empty leaf. MUST match + install.sh / setup.ps1 / install.ps1.""" + value = url.strip() + match = re.fullmatch(r"([^?#]*)([?#].*)?", value) + if match is None: + return value.rstrip("/") + return match.group(1).rstrip("/") + (match.group(2) or "") + + +def _torch_index_leaf(url: str) -> str: + """Final URL path segment, lowercased, query/fragment removed first. + + So a token-authenticated pin (.../cu128?token=x) classifies as cu128 (a raw leaf + keeps the query, never equals the +cu128 tag, and force-reinstalls every update). + CLASSIFICATION only; the install keeps the full URL. MUST match install.sh / + setup.ps1 / install.ps1. + """ + path = url.split("?", 1)[0].split("#", 1)[0] + return path.rstrip("/").rsplit("/", 1)[-1].lower() + + +# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed (torchao +# 0.17 cpp loads cleanly, and the flash-attn/causal-conv1d/mamba wheels pass on 2.11). +# torchvision/torchaudio are pinned (not bare) so the exclusive --index-url can't +# resolve one built against a different torch major -> ABI mismatch. _CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = ( "torch>=2.4,<2.12.0", "torchvision>=0.19,<0.27.0", "torchaudio>=2.4,<2.12.0", ) +# CPU torch repair specs (see _ensure_cpu_torch). Same bounds/reasoning as CUDA: the +# /cpu index also serves newer torch, so a bare trio could resolve out of range or ABI- +# mismatched. +_CPU_TORCH_PKG_SPEC: tuple[str, str, str] = _CUDA_TORCH_PKG_SPEC + # torchao's cpp extensions are pinned to ONE torch release AND CUDA major. A torch # mismatch just skips the cpp kernels (slow Python fallback); a CUDA mismatch fails # to import ("libcudart.so.12: cannot open shared object file"). The torch pin is a @@ -408,9 +473,8 @@ def _detect_rocm_version() -> tuple[int, int] | None: try: with open(path) as fh: parts = fh.read().strip().split("-")[0].split(".") - # Explicit length guard so we don't rely on the broad except - # below to swallow IndexError when the version file has a - # single component (e.g. "6\n" on a partial install). + # Explicit length guard: don't rely on the broad except below to + # swallow IndexError on a single-component version (e.g. "6\n"). if len(parts) >= 2: return int(parts[0]), int(parts[1]) except Exception: @@ -455,11 +519,10 @@ def _detect_rocm_version() -> tuple[int, int] | None: except Exception: pass - # Distro package-manager fallbacks. Package-managed ROCm installs can - # expose GPUs via rocminfo/amd-smi but lack /opt/rocm/.info/version and - # hipconfig, so probe dpkg (Debian/Ubuntu) and rpm (RHEL/Fedora/SUSE) - # for the rocm-core version. Matches install.sh::get_torch_index_url so - # `unsloth studio update` behaves like a fresh `curl | sh` install. + # Distro package-manager fallbacks: package-managed ROCm can expose GPUs via + # rocminfo/amd-smi but lack /opt/rocm/.info/version and hipconfig, so probe + # dpkg (Debian/Ubuntu) and rpm (RHEL/Fedora/SUSE) for the rocm-core version. + # Matches install.sh::get_torch_index_url so `studio update` == fresh install. for cmd in ( ["dpkg-query", "-W", "-f=${Version}\n", "rocm-core"], ["rpm", "-q", "--qf", "%{VERSION}\n", "rocm-core"], @@ -561,11 +624,10 @@ def _detect_windows_gfx_arch() -> str | None: stderr = subprocess.DEVNULL, timeout = 10, ) - # Accept partial output even when hipinfo crashes (e.g. exit code - # 0xC0000005 / STATUS_ACCESS_VIOLATION on some RDNA 4 hosts): if - # gcnArchName is present in stdout the device was enumerated before - # the crash, so the arch is trustworthy. Ignoring it causes a - # silent CPU PyTorch fallback (issue #6043). + # Accept partial output even when hipinfo crashes (e.g. 0xC0000005 / + # STATUS_ACCESS_VIOLATION on some RDNA 4 hosts): a gcnArchName in stdout + # means the device was enumerated pre-crash, so the arch is trustworthy. + # Ignoring it causes a silent CPU PyTorch fallback (issue #6043). text = result.stdout.decode(errors = "replace") # findall gets every gcnArchName line so multi-GPU hosts are # enumerable and HIP_VISIBLE_DEVICES selects correctly. @@ -706,9 +768,8 @@ def _detect_bnb_rocm_dll_ver() -> str | None: m = re.search(r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(dll)) if m: all_vers.append(m.group(1)) - # Pick the highest numeric suffix so e.g. "713" wins over "72" when both - # variants are present. Glob order is not guaranteed, so always sort - # rather than stopping at the first match. + # Highest numeric suffix wins (e.g. "713" over "72"); glob order is not + # guaranteed, so sort rather than take the first match. return max(all_vers, key = lambda v: int(v)) if all_vers else None @@ -825,17 +886,14 @@ def _has_rocm_gpu() -> bool: if result.returncode == 0 and result.stdout.strip(): if check_fn(result.stdout): return True - # sysfs KFD topology fallback (Linux only) -- matches install.sh's - # runtime-only detection. On minimal package-managed installs (no - # rocminfo / no amd-smi tools), the kernel exposes AMD GPUs via - # /sys/class/kfd so `studio update` can still detect and repair. + # sysfs KFD topology fallback (Linux only) -- matches install.sh's runtime-only + # detection. On minimal package-managed installs (no rocminfo / amd-smi), the + # kernel exposes AMD GPUs via /sys/class/kfd so `studio update` can still repair. # - # Guard: reject any KFD node whose properties file reports a non-AMD - # vendor. With the NVIDIA open kernel module (driver 560+), NVIDIA GPUs - # can register KFD topology nodes with a non-zero gpu_id; those nodes - # have vendor_id 4318 (0x10DE) rather than the AMD value 4098 (0x1002). - # Without this check the fallback returns True on NVIDIA-only systems, - # causing _ensure_rocm_torch to install ROCm wheels on NVIDIA hardware. + # Guard: reject any KFD node whose properties file reports a non-AMD vendor. The + # NVIDIA open kernel module (driver 560+) registers KFD nodes with a non-zero + # gpu_id and vendor_id 4318 (0x10DE), not the AMD 4098 (0x1002); without this + # check the fallback returns True on NVIDIA-only hosts, installing ROCm wheels. if sys.platform != "win32": try: kfd_nodes = "/sys/class/kfd/kfd/topology/nodes" @@ -849,12 +907,10 @@ def _has_rocm_gpu() -> bool: continue if not gpu_id or gpu_id == "0": # gpu_id 0 = CPU node continue - # Require AMD vendor_id 4098 (0x1002) in the properties file. - # KFD properties files exist on every kernel that exposes - # /sys/class/kfd, so absence of the file means we cannot - # confirm AMD ownership -- skip the node rather than risk a - # false positive (e.g. NVIDIA open driver KFD nodes that - # lack a properties file on some kernel versions). + # Require AMD vendor_id 4098 (0x1002). KFD properties files exist + # on every kernel exposing /sys/class/kfd, so a missing file means + # AMD ownership is unconfirmed -- skip the node rather than risk a + # false positive (e.g. NVIDIA open-driver KFD nodes lacking it). props_path = os.path.join(kfd_nodes, entry, "properties") try: with open(props_path) as fh: @@ -981,13 +1037,10 @@ def _install_bnb_windows_rocm() -> bool: ) if not _ok: return False - # After install: detect the actual ROCm DLL suffix shipped in the wheel and - # set BNB_ROCM_VERSION so bitsandbytes loads the correct DLL regardless of - # what torch.version.hip reports. The wheel may ship an older suffix (e.g. - # "72") while torch reports a newer HIP version (e.g. 7.13); the env var - # override ensures bitsandbytes does not fail looking for a non-existent DLL. - # The worker subprocess inherits this env var automatically. - # Fall back to "72" if detection fails (e.g. install was a no-op / dry-run). + # Detect the actual ROCm DLL suffix in the wheel and set BNB_ROCM_VERSION so bnb + # loads the right DLL regardless of torch.version.hip (the wheel may ship "72" + # while torch reports 7.13). The worker subprocess inherits it; fall back to "72" + # if detection fails (e.g. a no-op / dry-run install). _env_ver = os.environ.get("BNB_ROCM_VERSION") _env_is_persisted_default = ( os.environ.get(_BNB_ROCM_VERSION_SOURCE_ENV) == _BNB_ROCM_VERSION_SOURCE_SITECUSTOMIZE @@ -1002,13 +1055,11 @@ def _install_bnb_windows_rocm() -> bool: _persist_detected_version = True if _persist_detected_version: _persist_bnb_rocm_version(_ver) - # Make hipInfo.exe (shipped into the venv Scripts dir by the AMD torch - # wheel) resolvable via PATH for this process and every child python the - # installer spawns (import checks, precompile): bitsandbytes runs - # `hipinfo.exe` at import time to detect the GPU arch and logs a scary - # (harmless) ERROR + WARNING on every import when it is missing. The venv - # Scripts dir is on PATH only when the venv is activated, which neither - # Unsloth nor the installer's child processes ever do. + # Make hipInfo.exe (shipped into venv Scripts by the AMD torch wheel) resolvable + # via PATH for this process and every child python (import checks, precompile): + # bitsandbytes runs hipinfo.exe at import to detect the GPU arch and logs a scary + # (harmless) ERROR + WARNING when it is missing. Scripts is on PATH only for an + # activated venv, which neither Unsloth nor the installer's children ever do. _scripts_dir = os.path.dirname(sys.executable) if os.path.isfile(os.path.join(_scripts_dir, "hipInfo.exe")) and not shutil.which( "hipinfo.exe" @@ -1020,13 +1071,18 @@ def _install_bnb_windows_rocm() -> bool: def _detect_cuda_torch_index_url() -> str: """Return the pytorch.org CUDA wheel index URL for the host's NVIDIA driver. - 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). + Mirrors install.sh::get_torch_index_url's CUDA ladder so `studio update` repairs + to the same wheel family a fresh install would pick. Honours the explicit + overrides first (UNSLOTH_TORCH_INDEX_URL / _FAMILY) so a headless / CI install + never lets the host GPU decide. Otherwise probes nvidia-smi (parsing both "CUDA + Version:" and "CUDA UMD Version:"), defaulting to cu126 when unreadable. """ + _override_url = os.environ.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + if _override_url: + return _trim_index_path_slashes(_override_url) + _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" @@ -1061,6 +1117,157 @@ 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 _trim_index_path_slashes(url) + family = os.environ.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + if family: + return f"{_PYTORCH_WHL_BASE}/{family.strip('/')}" + return None + + +def _is_pip_rocm_family_leaf(leaf: str) -> bool: + """True when a lowercased leaf names a pip --index-url ROCm family: an EXACT + rocm[.] leaf or a gfx leaf. A suffixed leaf (rocm-rel-7.2.1, + rocm7.2-private) starts with "rocm" but is a custom pin the verbatim path owns, so + match EXACTLY. Mirrors install.sh / setup.ps1. + """ + # gfx must be followed by a digit (gfx90a, gfx1151, gfx120X-all): a gfx-prefixed + # custom leaf (gfx-private) is a verbatim pin, like rocm7.2-private. + return bool(re.fullmatch(r"rocm\d+(?:\.\d+)?", leaf)) or bool(re.match(r"gfx\d", leaf)) + + +def _explicit_rocm_torch_index_url() -> "str | None": + """The pinned wheel index URL when it names a pip ROCm family (rocm/gfx*), else None.""" + url = _explicit_torch_index_url() + if url is None: + return None + return url if _is_pip_rocm_family_leaf(_torch_index_leaf(url)) 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 installed + ROCm torch, so the pin needs a reinstall. Mirrors setup.ps1's stale-venv comparison; + same three pin-leaf cases as _ensure_rocm_torch. A same-family pin is NOT a mismatch. + """ + leaf = _torch_index_leaf(pin_url) + # Pinned ROCm version. The family classifier accepts a major-only rocm leaf too, + # so parse the minor as optional; a major-only pin compares on the major alone. + _pin_rocm = re.match(r"^rocm(\d+)(?:\.(\d+))?", leaf) + _pin_major = int(_pin_rocm.group(1)) if _pin_rocm else None + _pin_ver = ( + (int(_pin_rocm.group(1)), int(_pin_rocm.group(2))) + if _pin_rocm and _pin_rocm.group(2) is not None + else None + ) + # Installed +rocmX.Y version; a THREE-part +rocmA.B.C tag is the AMD per-arch + # (repo.amd.com/gfx*) signature vs a two-part pytorch.org wheel. + _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 tag; an untagged wheel never satisfies a ROCm pin. + _inst_has_rocm = re.search(r"\+rocm", installed_ver) is not None + # 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"): + # 2.11-allowlist arches expect the AMD per-arch wheel (three-part +rocmA.B.C, + # torch 2.11+); a generic or pre-2.11 build is a mismatch. + if leaf in _ROCM_GFX_TORCH211_LEAVES: + return not (_inst_is_211 and _inst_is_perarch) + # Non-2.11 gfx leaf (<2.11 specs): mismatch on an untagged wheel or torch 2.11+. + return (not _inst_has_rocm) or _inst_is_211 + + # Major-only rocm pin (rocm7): compare majors only -- a +rocm6.4 wheel under a rocm7 + # pin is a mismatch, any +rocm7.x wheel satisfies it (there is no pinned minor to + # compare, and the 2.11-line fallback below would invert both verdicts). + if _pin_major is not None and _pin_ver is None: + if _inst_ver is not None: + return _inst_ver[0] != _pin_major + # Untagged wheel never satisfies a ROCm pin; a +rocm tag with an unreadable + # version is accepted (matches the lenient unreadable fallback below). + return not _inst_has_rocm + + # rocmX.Y pin. Only KNOWN-2.11 rocm is the 2.11 line (no speculative floor). + _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 readable: exact (major, minor) compare (rocm7.2 pin over +rocm7.13.x -> + # mismatch, reinstall the pinned wheel). + if _pin_ver != _inst_ver: + return True + # Same family: a KNOWN-2.11 pin whose release drifted off 2.11 (2.12+rocm7.2) + # violates the spec -> reinstall to floor (exact compare, not >=2.11). + if _pin_is_211 and _inst_rel is not None: + if (int(_inst_rel.group(1)), int(_inst_rel.group(2))) != (2, 11): + return True + return False + # rocm pin, unreadable installed version: compare on the 2.11 line, but an untagged + # 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 + return url if _torch_index_leaf(url) == "cpu" else None + + +def _is_cuda_family_leaf(leaf: str) -> bool: + """True only for a real CUDA wheel-family leaf: "cu" + digits (cu118, cu128, ...). + + A bare startswith("cu") would match "custom"/"current". The match is EXACT so + "cu128-private" is NOT a family leaf and routes to the verbatim path instead. + """ + return re.fullmatch(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 + or a ROCm/CPU pin must not force a CUDA reinstall on a non-NVIDIA host). + """ + url = _explicit_torch_index_url() + if url is None: + return None + return url if _is_cuda_family_leaf(_torch_index_leaf(url)) 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. + + Known = rocm* / gfx* / cpu / cuXXX. Anything else (a private mirror /simple, + /current) is UNKNOWN: version-tag heuristics can't judge it, so the family + repair helpers must leave it alone (the install applied it verbatim). + Matches install.sh / setup.ps1 / install.ps1. + """ + url = _explicit_torch_index_url() + if url is None: + return None + leaf = _torch_index_leaf(url) + if _is_pip_rocm_family_leaf(leaf) or leaf == "cpu" or _is_cuda_family_leaf(leaf): + return None + return url + + def _ensure_cuda_torch() -> None: """Repair a venv whose torch is a ROCm build on an NVIDIA host. @@ -1073,44 +1280,47 @@ def _ensure_cuda_torch() -> None: Only repairs when torch actually links against HIP/ROCm. Healthy CUDA torch and deliberate CPU-only torch are left untouched. """ - # Respect an explicit backend choice from install.sh: only "" (standalone - # `studio update`) or "cuda" should ever force CUDA wheels. "rocm"/"cpu" - # (or any unrecognised value) are deliberate and must not be overridden. + # Respect install.sh's backend: only "" (standalone update) or "cuda" force CUDA + # wheels; "rocm"/"cpu"/unrecognised are deliberate. if _TORCH_BACKEND not in ("", "cuda"): return - # No CUDA torch on macOS; Windows venv/torch lifecycle is owned by - # install.ps1 (and the KFD poisoning bug is Linux-only), so skip both. + # An explicit unknown-family pin was applied VERBATIM at install time; leave it alone. + if _explicit_unknown_family_torch_index_url() is not None: + return + # No CUDA torch on macOS; Windows torch is owned by install.ps1 (KFD bug is Linux-only). if IS_MACOS or IS_WINDOWS or NO_TORCH: return # Never undo a deliberate ROCm install (setup.ps1 sets this marker). if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": return - # 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. + # An explicit CUDA pin (headless / CI cross-install) commits to CUDA wheels and skips ALL + # GPU probing, so it clears both the CUDA_VISIBLE_DEVICES hide gate and the NVIDIA gate below. + _cuda_pinned = _explicit_cuda_torch_index_url() is not None + # CUDA_VISIBLE_DEVICES="" / "-1" deliberately hides the NVIDIA GPU; never force CUDA + # wheels over that 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(): + # Only NVIDIA hosts carry CUDA torch (the CUDA pin overrides this gate too). + if not _cuda_pinned and not _has_usable_nvidia_gpu(): return - # Classify the installed torch: "hip" (ROCm build -- the poisoning - # signature), "cuda" (healthy), or "cpu" (deliberate CPU wheel). A - # non-zero exit means torch is missing or un-importable; the base install - # step handles that, so leave it alone. + # Classify the installed torch: "hip" (ROCm poisoning signature), "cuda" (healthy), + # or "cpu". A non-zero exit means torch is missing/un-importable: without a pin the + # base install owns it, but a pinned CUDA index reinstalls it below. try: probe = subprocess.run( [ 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, @@ -1120,22 +1330,60 @@ def _ensure_cuda_torch() -> None: except (OSError, subprocess.TimeoutExpired): return if probe.returncode != 0: + # torch present but can't import. Without a pin the base install owns it; but an + # explicit CUDA pin forces this pass (failed probe) and the base update won't + # reinstall an already-installed torch, so reinstall from the pin (self-resolving). + if not _cuda_pinned: + return + index_url = _detect_cuda_torch_index_url() + _torch_pkg, _vision_pkg, _audio_pkg = _CUDA_TORCH_PKG_SPEC + print( + f" torch cannot import but an explicit CUDA index is pinned -- reinstalling " + f"CUDA torch from {_strip_index_url_credentials(index_url)}" + ) + pip_install( + "CUDA torch repair", + "--force-reinstall", + "--no-cache-dir", + _torch_pkg, + _vision_pkg, + _audio_pkg, + "--index-url", + index_url, + constrain = False, + ) return - # Take the last non-empty stdout line: stray output from sitecustomize or - # an import hook must not mask the marker (fail-closed either way). + # Last non-empty line: stray sitecustomize/import-hook output must not mask the marker. _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 on a ROCm build on an NVIDIA host (poisoning signature), or when a + # CUDA index is pinned but the venv has the wrong family (CPU or a different cuXXX). A + # healthy match, or a CPU wheel with no CUDA pin, is left alone. + _pin = _explicit_torch_index_url() + _pin_leaf = _torch_index_leaf(_pin) 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: + # Installed cuXXX differs from the pin. An untagged build (empty) counts too: + # the family can't be confirmed, so reinstall to enforce it (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 {_strip_index_url_credentials(index_url)}\n" + f" (set UNSLOTH_TORCH_BACKEND=rocm or cpu to keep a deliberate " + f"non-CUDA torch)" ) pip_install( "CUDA torch repair", @@ -1150,6 +1398,90 @@ def _ensure_cuda_torch() -> None: ) +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/rocm_torch for the explicit-CPU case (those treat a CPU + backend as a skip, so a standalone `studio update` would ignore the authoritative CPU + pin). Only fires for an EXPLICIT pin. + """ + 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 explicit CPU pin reinstalls it below. + 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: + # torch present but can't import. The explicit CPU pin forces this pass (failed + # probe) and the base update won't reinstall an already-installed torch, so + # reinstall from the pin (self-resolving, no loop). + _torch_pkg, _vision_pkg, _audio_pkg = _CPU_TORCH_PKG_SPEC + print( + f" torch cannot import but an explicit CPU index is pinned -- reinstalling " + f"CPU torch from {_strip_index_url_credentials(pin)}" + ) + pip_install( + "CPU torch repair", + "--force-reinstall", + "--no-cache-dir", + _torch_pkg, + _vision_pkg, + _audio_pkg, + "--index-url", + pin, + constrain = False, + ) + return + _lines = [ + line.strip() for line in probe.stdout.decode(errors = "replace").splitlines() if line.strip() + ] + if not _lines: + return # unreadable -- the base install step handles a missing torch + if _lines[-1] != "gpu": + return # already a CPU build + + print( + " torch is a GPU build but an explicit CPU index is pinned -- reinstalling " + f"CPU torch from {_strip_index_url_credentials(pin)}" + ) + # Pin the supported torch<2.11 family (the /cpu index now serves 2.11+, so a bare + # trio could resolve out of range or ABI-mismatched). + _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, + ) + + def _ensure_rocm_torch() -> None: """Reinstall torch with ROCm wheels when the venv received CPU-only torch. @@ -1160,16 +1492,15 @@ def _ensure_rocm_torch() -> None: Uses pip_install() to respect uv, constraints, and --python targeting. """ global _rocm_windows_torch_installed - # install.sh sets UNSLOTH_TORCH_BACKEND to the resolved wheel family - # ("cuda", "rocm", "cpu"). Skip ROCm operations entirely when install.sh - # already selected a non-ROCm backend -- this is the authoritative signal - # and avoids re-running GPU detection in a subprocess that may see a - # different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.). + # install.sh's resolved backend is authoritative: skip ROCm when it already chose a + # non-ROCm family (avoids re-detecting in a subprocess that may see a different env). if _TORCH_BACKEND in ("cuda", "cpu"): return - # setup.ps1 sets this after installing AMD wheels; skip the probe only when - # torch is actually importable as ROCm. If the venv was wiped between runs, - # the stale env-var would suppress a needed reinstall. + # An explicit unknown-family pin was applied VERBATIM at install time; leave it alone. + if _explicit_unknown_family_torch_index_url() is not None: + return + # setup.ps1 sets this after installing AMD wheels; skip only when torch is actually + # importable as ROCm (a wiped venv leaves a stale env-var that must not suppress it). if os.environ.get("UNSLOTH_ROCM_TORCH_INSTALLED") == "1": _torch_ok = False try: @@ -1193,9 +1524,8 @@ def _ensure_rocm_torch() -> None: pass if _torch_ok: _rocm_windows_torch_installed = True - # setup.ps1 already installed ROCm torch, but we still need the AMD - # Windows BNB wheel here -- the PyPI bitsandbytes wheel ships only - # CUDA DLLs and fails to load on ROCm. + # ROCm torch is already installed, but the AMD Windows BNB wheel is still + # needed (the PyPI bitsandbytes ships only CUDA DLLs, fails on ROCm). _install_bnb_windows_rocm() return # torch was wiped between runs; fall through to the full install path @@ -1203,10 +1533,15 @@ def _ensure_rocm_torch() -> None: return if IS_WINDOWS: - if _has_usable_nvidia_gpu(): + # An explicit ROCm-family pin commits to ROCm wheels regardless of the visible + # GPU and overrides the public per-arch index (mirrors the Linux pin handling + # below): after a pinned setup.ps1 install fails to CPU, this repair must retry + # the PINNED index, not repo.amd.com. + _win_rocm_pin = _explicit_rocm_torch_index_url() + if _win_rocm_pin is None and _has_usable_nvidia_gpu(): return gfx_arch = _detect_windows_gfx_arch() - if not gfx_arch: + if not gfx_arch and _win_rocm_pin is None: return # no AMD GPU visible via hipinfo # Probe whether torch already links against HIP. _torch_already_rocm = False @@ -1231,23 +1566,24 @@ def _ensure_rocm_torch() -> None: except (OSError, subprocess.TimeoutExpired): pass if not _torch_already_rocm: - index_url = _windows_rocm_index_url(gfx_arch) + index_url = _win_rocm_pin or _windows_rocm_index_url(gfx_arch) if index_url is None: print(f" No AMD Windows torch index for GPU arch {gfx_arch} -- skipping") return - print(f" {gfx_arch} (Windows) -- installing torch from {index_url}") - # Pin companions for the arches install.ps1/setup.ps1 pin (gfx120X / - # Strix) so the per-arch index resolves an ABI-consistent trio; other - # arches stay bare (no published floor), matching the PowerShell side. + print( + f" {gfx_arch or 'pinned ROCm index'} (Windows) -- installing torch from " + f"{_strip_index_url_credentials(index_url)}" + ) + # Pin companions for the arches install.ps1/setup.ps1 pin (gfx120X / Strix) + # so the per-arch index resolves an ABI-consistent trio; other arches stay bare. _torch_pkg, _vision_pkg, _audio_pkg = _WINDOWS_ROCM_TORCH_PKG_SPECS.get( gfx_arch, ("torch", "torchvision", "torchaudio") ) - # Nonfatal: a transient AMD-index failure must not abort the whole - # install once the PowerShell side has fallen back to CPU torch. - # --force-reinstall resolves before uninstalling, so a failed index - # leaves the existing build intact; keep it and let the user retry. + # Nonfatal: a transient AMD-index failure must not abort the install. + # --force-reinstall resolves before uninstalling, so a failed index keeps the + # existing build intact; let the user retry. if not pip_install_try( - f"ROCm torch (Windows, {gfx_arch})", + f"ROCm torch (Windows, {gfx_arch or 'pinned'})", "--force-reinstall", "--index-url", index_url, @@ -1257,7 +1593,7 @@ def _ensure_rocm_torch() -> None: constrain = False, ): print( - f" Warning: AMD Windows ROCm torch install failed for {gfx_arch}; " + f" Warning: AMD Windows ROCm torch install failed for {gfx_arch or 'the pinned index'}; " "keeping the existing torch build. Re-run 'unsloth studio update' " "later to retry ROCm." ) @@ -1280,26 +1616,30 @@ 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 pin commits to ROCm wheels regardless of the visible GPU (headless / CI). + # Mirror _ensure_cuda_torch: skip the NVIDIA/no-AMD/unreadable gates. + _rocm_pin = _explicit_rocm_torch_index_url() + if _rocm_pin is None: + # NVIDIA takes precedence on mixed hosts (only if a GPU is usable). + if _has_usable_nvidia_gpu(): + return + # _has_rocm_gpu() (rocminfo / amd-smi rows) is the authoritative AMD-host signal; + # the old /opt/rocm-or-hipcc gate broke runtime-only ROCm installs. + 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 + if _rocm_pin is None: + print(" ROCm detected but version unreadable -- skipping torch reinstall") + return + # Explicit pin: the pinned leaf drives the install, so an unreadable host version + # is fine (sentinel keeps ver comparisons defined). + ver = (0, 0) - # Probe whether torch already links against HIP (ROCm already working). - # Do NOT skip for CUDA-only builds: they are unusable on AMD-only hosts - # (the NVIDIA check above already handled mixed AMD+NVIDIA setups). + # Probe whether torch links against HIP, capturing the installed ROCm tag for pin-mismatch + # detection. Emit ONE "|" line: marker (HIP version, "rocm" sentinel, + # or empty for CPU/CUDA) before "|", wheel version after. try: probe = subprocess.run( [ @@ -1309,10 +1649,10 @@ 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 if present, else a "rocm" sentinel when only the + # version string flags ROCm; empty marker = CPU/CUDA torch. + "marker=hip if hip else ('rocm' if 'rocm' in ver else ''); " + "print(marker + '|' + ver)" ), ], stdout = subprocess.PIPE, @@ -1321,29 +1661,42 @@ 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() != "" + # Last non-empty line, split on the FIRST "|" so the empty HIP field is preserved. + _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 line is required; without it treat HIP as absent -> reinstall. + has_hip_torch = bool(_sep) and _hip_marker != "" + + # An explicit ROCm pin whose family differs from the installed torch must reinstall, else a + # rocm7.2/gfx* pin over an older +rocm6.4/7.1 build never applies. Version-tag heuristic + # only: a same-tag per-arch switch (gfx1151 -> gfx120X-all, both +rocm7.13.0) isn't detectable. + _rocm_pin_mismatch = ( + _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver) + if (has_hip_torch and _rocm_pin is not None) + else False ) - 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 - # with the real fix, so route those hosts there instead of the generic - # pytorch.org rocm7.1 wheel. Mirrors install.sh's Strix override. - # On mixed hosts (Strix iGPU + non-Strix dGPU), route to the AMD per-gfx - # index only when HIP's runtime GPU is the Strix one -- else the dGPU gets - # an incompatible wheel. Use HIP_VISIBLE_DEVICES for the runtime target. + # Strix Halo / Point (gfx1151 / gfx1150) segfault under ROCm 7.1 in torch._grouped_mm; + # AMD's per-gfx repo ships 2.11.0+rocm7.13.0 with the fix, so route those hosts there + # (mirrors install.sh). On mixed hosts, reroute only when HIP's runtime GPU is the Strix one. _strix_override_url: "str | None" = None _strix_override_pkgs: "tuple[str, str, str] | None" = None - if ver < (7, 2): + # An explicit ROCm pin is authoritative: never auto-reroute it. + 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) if _detected_strix: - # Pick the runtime-visible GPU: use the HIP_VISIBLE_DEVICES index - # into gfx_codes, else default to the first GPU. Skip the override - # unless the resolved GPU is Strix. + # Runtime-visible GPU (HIP_VISIBLE_DEVICES index into gfx_codes, else first); + # skip the override unless it's Strix. _runtime_gfx = gfx_codes[_pick_visible_index(len(gfx_codes))] if gfx_codes else None if _runtime_gfx in _strix_gfx: _selected_gfx = _runtime_gfx @@ -1353,12 +1706,8 @@ def _ensure_rocm_torch() -> None: _strix_override_url = f"{_amd_mirror}/{_selected_gfx}/" _strix_override_pkgs = ( "torch>=2.11.0,<2.12.0", - # Pin torchvision/torchaudio to the 2.11.x-compatible range. - # The install uses --index-url (exclusive, no PyPI fallback), - # so bare unversioned names risk resolving an AMD-index build - # targeting a different torch major (e.g. 0.27 built against - # torch 2.12), which fails at runtime with an ABI/version - # mismatch. Matches _ROCM_TORCH_CONSTRAINT["rocm7.2"]. + # Pin companions to the 2.11.x range: the exclusive --index-url could + # otherwise resolve a build for a different torch major (ABI mismatch). "torchvision>=0.26.0,<0.27.0", "torchaudio>=2.11.0,<2.12.0", ) @@ -1378,14 +1727,15 @@ def _ensure_rocm_torch() -> None: f" skipping AMD per-gfx index override.\n" ) - # Strix override on ROCm 7.1 must fire even when has_hip_torch is True -- - # an existing torch with `torch.version.hip == "7.1"` is exactly the broken - # combo the override repairs, so skipping it leaves users on the known - # _grouped_mm segfault. + # The Strix override must fire even when has_hip_torch is True: an existing + # torch.version.hip == "7.1" is exactly the broken combo it repairs. if _strix_override_url is not None and _strix_override_pkgs is not None: index_url = _strix_override_url _torch_pkg, _vision_pkg, _audio_pkg = _strix_override_pkgs - print(f" Strix ROCm 7.1 override -- installing torch from {index_url}") + print( + f" Strix ROCm 7.1 override -- installing torch from " + f"{_strip_index_url_credentials(index_url)}" + ) pip_install( "ROCm torch (Strix arch-specific)", "--force-reinstall", @@ -1398,24 +1748,38 @@ def _ensure_rocm_torch() -> None: constrain = False, ) 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 a ROCm build's family differs from a pin. + # Honour a ROCm pin verbatim; else pick the newest wheel tag <= host. + _override_idx = _explicit_rocm_torch_index_url() + if _override_idx is not None: + index_url = _override_idx + tag = _torch_index_leaf(index_url) 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 {_strip_index_url_credentials(index_url)}") + # Only the _grouped_mm-bug gfx arches need the 2.11 spec; other gfx indexes ship + # <2.11 and stay on the default range (matches 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", @@ -1504,11 +1868,26 @@ def _infer_no_torch() -> bool: NO_TORCH = _infer_no_torch() -# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() so -# that this script knows which torch variant was selected without re-running -# GPU detection. Values: "cuda", "rocm", or "cpu". Empty means unknown -# (standalone `unsloth studio update` runs, where we re-detect normally). +# UNSLOTH_TORCH_BACKEND is set by install.sh after get_torch_index_url() ("cuda", "rocm", +# "cpu"; empty = standalone `studio update`, where we re-detect). _TORCH_BACKEND: str = os.environ.get("UNSLOTH_TORCH_BACKEND", "").lower() +# Standalone update with an explicit pin: derive the backend from the override (classify on +# the final URL/family segment, mirroring install.sh) instead of re-probing the GPU. +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 = _torch_index_leaf(_idx_override) + 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" so /current or /custom is NOT branded CUDA (a wrong backend + # makes _ensure_rocm_torch return early on AMD hosts). An unknown leaf keeps "" so the + # helpers probe the GPU. + _TORCH_BACKEND = "cuda" def _torch_step_label(suffix: str) -> str: @@ -1724,12 +2103,15 @@ def run( cmd, stdout = subprocess.PIPE if quiet else None, stderr = subprocess.STDOUT if quiet else None, + env = _install_env_for_cmd(cmd), **_windows_hidden_subprocess_kwargs(), ) if result.returncode != 0: _step("error", f"{label} failed (exit code {result.returncode})", _red) if result.stdout: - print(result.stdout.decode(errors = "replace")) + # Redact before printing: the failing pip command may carry a pinned --index-url + # with userinfo/?token= creds, so raw pip error text would leak them. + print(_redact_install_output(result.stdout)) sys.exit(result.returncode) return result @@ -1737,15 +2119,13 @@ def run( # Packages to skip on Windows (require special build steps) WINDOWS_SKIP_PACKAGES = {"triton_kernels"} -# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode). -# These either *are* torch extensions or have unconditional -# ``Requires-Dist: torch``, so installing them would pull torch back in. -# ``librosa`` is here too despite not requiring torch: upstream ``llvmlite`` -# dropped its macOS x86_64 wheel between 0.42.0 and 0.46.0+ (see -# https://pypi.org/project/llvmlite/0.47.0/#files -- only -# macosx_arm64 / manylinux / win_amd64 remain), so on Intel Mac the -# librosa -> numba -> llvmlite chain triggers a from-source build that fails -# in CI and on hosts without LLVM 14/15 headers. Tracked in unslothai/unsloth#5046. +# Packages to skip when torch is unavailable (Intel Mac GGUF-only mode). These +# either *are* torch extensions or have unconditional ``Requires-Dist: torch``, so +# installing them pulls torch back in. ``librosa`` is here despite not requiring +# torch: upstream ``llvmlite`` dropped its macOS x86_64 wheel (0.46.0+ ships only +# macosx_arm64 / manylinux / win_amd64), so on Intel Mac the librosa -> numba -> +# llvmlite chain triggers a from-source build that fails without LLVM 14/15 headers. +# Tracked in unslothai/unsloth#5046. NO_TORCH_SKIP_PACKAGES = { "torch-stoi", "timm", @@ -1767,7 +2147,8 @@ def _build_flash_attn_wheel_url(env: dict[str, str]) -> str | None: def _print_optional_install_failure(label: str, result: subprocess.CompletedProcess[str]) -> None: _step("warning", f"{label} failed (exit code {result.returncode})", _cyan) if result.stdout: - print(result.stdout.strip()) + # Redact any pinned --index-url credentials before printing captured output. + print(_redact_install_output(result.stdout).strip()) def _flash_attn_install_disabled() -> bool: @@ -1913,15 +2294,60 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]: # Colab and similar). cmd.extend(["--python", sys.executable]) cmd.extend(_translate_pip_args_for_uv(args)) - # Torch is pre-installed by install.sh/setup.ps1. Do not add - # --torch-backend by default -- it can cause solver dead-ends on CPU-only - # machines. Callers that need it can set UV_TORCH_BACKEND. + # Torch is pre-installed, so don't add --torch-backend by default (solver dead-ends on + # CPU-only machines); callers can set UV_TORCH_BACKEND. Never add it to a pinned-index + # command: uv's torch backend redirects torch to its own per-backend index, defeating the pin. _tb = os.environ.get("UV_TORCH_BACKEND", "") - if _tb: + if _tb and not _is_pinned_index_cmd(cmd): cmd.append(f"--torch-backend={_tb}") return cmd +# uv resolves --index-url / --default-index at LOWEST priority, so an inherited UV_INDEX / +# UV_EXTRA_INDEX_URL mirror wins and a pinned torch repair silently ignores the pin. +# Neutralise these for pinned installs (as install.sh #6898 / install.ps1 / setup.ps1 do). +# UV_TORCH_BACKEND redirects torch; PIP_* matter for the pip FALLBACK; UV_CONFIG_FILE is +# stripped + UV_NO_CONFIG=1 (a discovered uv.toml outranks the CLI pin, uv 0.10). +_UV_INDEX_ENV_VARS = ( + "UV_CONFIG_FILE", + "UV_DEFAULT_INDEX", + "UV_INDEX_URL", + "UV_INDEX", + "UV_EXTRA_INDEX_URL", + "UV_TORCH_BACKEND", + "UV_FIND_LINKS", + "PIP_EXTRA_INDEX_URL", + "PIP_FIND_LINKS", + # PIP_NO_INDEX=1 makes the pip fallback ignore ALL indexes (defeating --index-url); + # PIP_INDEX_URL is dropped too so a stale mirror env can't outrank the pin. + "PIP_NO_INDEX", + "PIP_INDEX_URL", +) + + +def _is_pinned_index_cmd(cmd: "list[str] | tuple[str, ...]") -> bool: + """True when the command pins an index via --index-url / --default-index.""" + return any(arg in ("--index-url", "--default-index") for arg in cmd) + + +def _install_env_for_cmd(cmd: "list[str]") -> "dict[str, str] | None": + """Return an env with the uv index vars stripped for a pinned-index install. + + None (inherit env) when the command does NOT pin an index, so ordinary installs honour + the user's mirror. For pinned commands, the uv index/backend vars are removed, + UV_NO_CONFIG=1 set (a discovered uv.toml outranks the CLI pin), and PIP_CONFIG_FILE + pointed at os.devnull for the pip fallback. Mirrors install.sh's gate (#6898). + """ + if not _is_pinned_index_cmd(cmd): + return None + env = os.environ.copy() + for name in _UV_INDEX_ENV_VARS: + env.pop(name, None) + env["UV_NO_CONFIG"] = "1" + env["PIP_CONFIG_FILE"] = os.devnull + return env + + def pip_install_try( label: str, *args: str, @@ -1948,11 +2374,13 @@ def pip_install_try( cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, + env = _install_env_for_cmd(cmd), ) if result.returncode == 0: return True if VERBOSE and result.stdout: - print(result.stdout.decode(errors = "replace")) + # pip/uv echo index URLs (credentials included) in failure output. + print(_redact_install_output(result.stdout)) return False @@ -2000,13 +2428,14 @@ def pip_install( uv_cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, + env = _install_env_for_cmd(uv_cmd), **_windows_hidden_subprocess_kwargs(), ) if result.returncode == 0: return print(_red(f" uv failed, falling back to pip...")) if result.stdout: - print(result.stdout.decode(errors = "replace")) + print(_redact_install_output(result.stdout)) pip_cmd = _build_pip_cmd(args) + constraint_args_pip + req_args_pip run(f"{label} (pip)" if USE_UV else label, pip_cmd) @@ -2054,10 +2483,9 @@ def install_python_stack() -> int: global USE_UV, _STEP, _TOTAL _STEP = 0 - # install.sh (which already installed unsloth) sets SKIP_STUDIO_BASE=1 to - # avoid reinstalling base packages. "unsloth studio update" does NOT set it, - # so base packages (unsloth + unsloth-zoo) are reinstalled to pick up new - # versions. + # install.sh sets SKIP_STUDIO_BASE=1 to avoid reinstalling base packages; + # `studio update` does NOT, so unsloth + unsloth-zoo are reinstalled to pick + # up new versions. skip_base = os.environ.get("SKIP_STUDIO_BASE", "0") == "1" # --package installs a different package name (for testing). package_name = os.environ.get("STUDIO_PACKAGE_NAME", "unsloth") @@ -2067,9 +2495,9 @@ def install_python_stack() -> int: if IS_MACOS: base_total -= 1 # triton step is skipped on macOS if not IS_MACOS and not NO_TORCH: - base_total += 1 # ROCm torch check (line 1526) -- all non-macOS platforms + base_total += 1 # ROCm torch check (step 2b), non-macOS if not IS_WINDOWS: - base_total += 2 # flash-attn (line 1620) + ROCm torch final (line 1705) -- Linux only + base_total += 2 # flash-attn + torch final repair (step 13), Linux _TOTAL = (base_total - 1) if skip_base else base_total # 1. Try uv for faster installs (before pip upgrade -- uv venvs don't @@ -2134,9 +2562,8 @@ def install_python_stack() -> int: if skip_base: pass elif NO_TORCH: - # No-torch update path: install unsloth + unsloth-zoo with --no-deps - # (PyPI metadata still declares torch as a hard dep), then runtime deps - # with --no-deps (avoids transitive torch). + # No-torch update path: install unsloth + unsloth-zoo, then runtime deps, + # both with --no-deps (PyPI metadata declares torch a hard dep; avoid it). _progress("base packages (no torch)") pip_install( f"Updating {package_name} + unsloth-zoo (no-torch mode)", @@ -2149,10 +2576,9 @@ def install_python_stack() -> int: package_name, "unsloth-zoo", ) - # Resolve pydantic WITH deps so pip pins pydantic-core to the exact - # version pydantic's metadata declares. Under --no-deps pip picks the - # latest of each and trips pydantic's _ensure_pydantic_core_version - # check. Transitive deps are torch-free. + # Resolve pydantic WITH deps so pip pins pydantic-core to the exact version + # its metadata declares (under --no-deps pip picks the latest of each and + # trips pydantic's _ensure_pydantic_core_version check). Deps are torch-free. pip_install( "Installing pydantic (with deps for compatible core)", "--no-cache-dir", @@ -2244,6 +2670,7 @@ def install_python_stack() -> int: _progress(_torch_step_label("check")) _ensure_cuda_torch() _ensure_rocm_torch() + _ensure_cpu_torch() # Windows + AMD GPU: warn if ROCm torch was not installed (wrong Python # version or unknown ROCm version). @@ -2309,11 +2736,10 @@ def install_python_stack() -> int: req = REQ_ROOT / "extras-no-deps.txt", ) - # 4. Overrides (torchao) -- force-reinstall. The torchao version is chosen to - # match the torch installed in the venv so its C++ extensions load (see - # _select_torchao_spec). Skip when torch is unavailable (e.g. Intel Mac - # GGUF-only mode): torchao requires torch. Also skipped on Windows ROCm - # (no working build; see below). + # 4. Overrides (torchao) -- force-reinstall to a version matching the venv's + # torch so its C++ extensions load (see _select_torchao_spec). Skipped when + # torch is unavailable (Intel Mac GGUF-only) and on Windows ROCm (no working + # build; see below). if NO_TORCH: _progress("dependency overrides (skipped, no torch)") elif _rocm_windows_torch_installed or _installed_torch_is_windows_rocm(): @@ -2430,14 +2856,12 @@ def install_python_stack() -> int: [sys.executable, str(SINGLE_ENV / "patch_metadata.py")], ) - # 13. AMD ROCm: final torch repair. Several steps above can pull in CUDA - # torch from PyPI (base packages, extras, overrides, studio deps, etc.). - # Running the repair last ensures ROCm torch is in place at runtime, - # whichever intermediate step clobbered it. + # 13. Final torch repair. Steps above can pull CUDA torch from PyPI, so repair last. if not IS_WINDOWS and not IS_MACOS and not NO_TORCH: _progress(_torch_step_label("final")) _ensure_cuda_torch() _ensure_rocm_torch() + _ensure_cpu_torch() # 14. Final check (silent; third-party conflicts are expected) subprocess.run( diff --git a/studio/setup.ps1 b/studio/setup.ps1 index f7d33a1142..f523b9ff14 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -431,6 +431,167 @@ function Get-PytorchCudaTag { return "cu126" } +# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL +# TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared. +function Trim-IndexPathSlashes { + param([string]$Url) + $value = $Url.Trim() + $idx = $value.IndexOfAny([char[]]@('?', '#')) + if ($idx -lt 0) { + return $value.TrimEnd('/') + } + return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx) +} + +# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY), shared by the stale-venv check +# and install selection so a pinned index wins over GPU probing (parity with the other +# installers). URL is verbatim; _FAMILY is the leaf joined to the mirror base. +function Get-PinnedTorchIndexUrl { + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) { + return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL) + } + if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) { + $base = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } + return "$base/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))" + } + return $null +} + +# Last path segment of a wheel index URL, query/fragment dropped first so a token-authenticated +# pin (.../cu128?token=x) classifies as cu128 (else it reinstalls every update). Classification +# only. Shared with the py / install.sh leaf extractors. +function Get-TorchIndexLeaf { + param([string]$Url) + if ([string]::IsNullOrWhiteSpace($Url)) { return $null } + $path = ($Url -split '[?#]', 2)[0] + if ([string]::IsNullOrWhiteSpace($path)) { return $null } + return ($path.TrimEnd('/') -split '/')[-1].ToLowerInvariant() +} + +# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer +# output before printing on failure; uv/pip errors echo the failing --index-url verbatim. +# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted. +function Redact-InstallOutput { + param([string]$Text) + if (-not $Text) { return $Text } + $Text = $Text -replace '(https?://)[^/@\s`]+@', '$1@' + $Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=' + # A #token=... fragment is as sensitive as a query; URL-anchored. + return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#' +} + +# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug). MUST match +# the install-spec path below and the other installers; other leaves ship <2.11 and stay default. +function Test-RocmGfx211Leaf { + param([string]$Leaf) + return @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $Leaf +} + +# rocmX.Y versions KNOWN to ship torch 2.11: rocm7.2 only today. Do NOT floor an unknown newer +# rocm speculatively. MUST match _ROCM_KNOWN_TORCH211_VERSIONS and the rocm7.2 leaf elsewhere. +function Test-RocmKnown211Version { + param([int]$Major, [int]$Minor) + return ($Major -eq 7 -and $Minor -eq 2) +} + +# True only for a real CUDA family leaf: "cu" + digits (cu118, cu128, ...). A bare -like 'cu*' +# would match "custom"/"current" and rebuild the venv every run. Mirrors _is_cuda_family_leaf. +function Test-CudaFamilyLeaf { + param([string]$Leaf) + if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false } + # EXACT cu+digits: cu128-private routes through the unknown-leaf path instead. + return $Leaf -match '^cu[0-9]+$' +} + +# True only for a real pip ROCm family leaf: EXACT rocm[.] or a gfx leaf. A leaf +# that merely STARTS with rocm (rocm-rel-7.2.1, rocm7.2-private) is a custom pin the verbatim +# path owns, so anchor the match. Mirrors _is_pip_rocm_family_leaf / install.sh. +function Test-PipRocmFamilyLeaf { + param([string]$Leaf) + if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false } + # gfx must be followed by a digit (an architecture leaf); gfx-private is custom. + return ($Leaf -match '^gfx[0-9]') -or ($Leaf -match '^rocm[0-9]+(\.[0-9]+)?$') +} + +# Stale-venv ROCm comparison for a pinned gfx*/rocm* index. Returns @{ Expected; Installed } so +# the caller rebuilds when they differ. Mirrors _rocm_pin_family_mismatch (same rocmX.Y / gfx +# cases). An untagged (no +rocm) wheel never satisfies a ROCm pin -> stale. +function Get-RocmPinStaleTags { + param([string]$PinLeaf, [string]$TorchVersion) + $_pinRocm = [regex]::Match($PinLeaf, '^rocm(\d+)\.(\d+)') + $_pinVer = if ($_pinRocm.Success) { "$($_pinRocm.Groups[1].Value).$($_pinRocm.Groups[2].Value)" } else { $null } + # The family classifier accepts a major-only rocm leaf too (rocm7). + $_pinMajorOnly = [regex]::Match($PinLeaf, '^rocm(\d+)$') + # Installed rocm version and whether the wheel is a per-arch (three-part) build. + $_instRocm = [regex]::Match($TorchVersion, '\+rocm(\d+)\.(\d+)') + $_instVer = if ($_instRocm.Success) { "$($_instRocm.Groups[1].Value).$($_instRocm.Groups[2].Value)" } else { $null } + $_instPerArch = [regex]::IsMatch($TorchVersion, '\+rocm\d+\.\d+\.\d+') + # A ROCm build MUST carry a +rocm tag; an untagged wheel can't satisfy any ROCm pin. + $_instHasRocm = [regex]::IsMatch($TorchVersion, '\+rocm') + $_instRel = [regex]::Match($TorchVersion, '^(\d+)\.(\d+)') + $_instIs211 = $false + if ($_instRel.Success) { + $_instIs211 = ([int]$_instRel.Groups[1].Value -gt 2) -or ([int]$_instRel.Groups[1].Value -eq 2 -and [int]$_instRel.Groups[2].Value -ge 11) + } + + if ($PinLeaf -like 'gfx*') { + if (Test-RocmGfx211Leaf $PinLeaf) { + # Expect the AMD per-arch (three-part) 2.11 wheel: satisfied only when BOTH + # a 2.11 release AND a three-part rocm tag are installed. + $installed = if ($_instIs211 -and $_instPerArch) { "rocm-perarch(torch>=2.11)" } else { "rocm-generic-or-old" } + return @{ Expected = "rocm-perarch(torch>=2.11)"; Installed = $installed } + } + # Non-2.11 gfx leaf (<2.11 spec): stale on an untagged wheel or a 2.11+ build. + $installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } + return @{ + Expected = "rocm(torch<2.11)" + Installed = $installed + } + } + + # Major-only rocm pin (rocm7): compare majors only -- a +rocm6.4 wheel under a rocm7 + # pin is stale, any +rocm7.x wheel satisfies it (no pinned minor to compare, and the + # 2.11-line fallback below would invert both verdicts). Mirrors _rocm_pin_family_mismatch. + if ($_pinMajorOnly.Success) { + $_pinMaj = [int]$_pinMajorOnly.Groups[1].Value + if ($_instVer) { + $_instMaj = [int]$_instRocm.Groups[1].Value + $expected = if ($_instMaj -eq $_pinMaj) { "rocm$_instVer" } else { "rocm$_pinMaj.x" } + return @{ Expected = $expected; Installed = "rocm$_instVer" } + } + # Untagged wheel never satisfies a ROCm pin; a +rocm tag with an unreadable + # version is accepted (matches the lenient unreadable fallback below). + $installed = if ($_instHasRocm) { "rocm" } else { "not-rocm" } + return @{ Expected = "rocm"; Installed = $installed } + } + + # rocmX.Y pin. + if ($_pinVer -and $_instVer) { + # Both readable: exact compare. When they match AND the pin is KNOWN-2.11, the + # installed release must also be 2.11 (a +rocm7.2 wheel drifted to 2.12 shares the + # tag but violates the spec), so fold the release into the tag. Mirrors _rocm_pin_family_mismatch. + $_pinKnown211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value) + $_instOn211 = $_instRel.Success -and [int]$_instRel.Groups[1].Value -eq 2 -and [int]$_instRel.Groups[2].Value -eq 11 + if ($_pinKnown211 -and -not $_instOn211) { + return @{ Expected = "rocm$_pinVer(torch2.11)"; Installed = "rocm$_instVer(torch-off-2.11)" } + } + return @{ Expected = "rocm$_pinVer"; Installed = "rocm$_instVer" } + } + $_pinNeeds211 = $false + if ($_pinRocm.Success) { + # Only KNOWN-2.11 rocm (rocm7.2) is on the 2.11 line (no speculative floor). + # Matches _ROCM_KNOWN_TORCH211_VERSIONS. + $_pinNeeds211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value) + } + # Fallback (installed rocm version unreadable): compare on the 2.11 line; an untagged + # wheel never satisfies a rocmX.Y pin -> stale. + $installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } + return @{ + Expected = if ($_pinNeeds211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" } + Installed = $installed + } +} + # VS generator -> MSBuild BuildCustomizations dir; toolset tracks the VS major # (18->v180, 17->v170), defaulting to v170 when unparseable. function Get-VcBuildCustomizationsDir { @@ -813,11 +974,14 @@ function Invoke-SetupCommand { # Merge stderr into stdout so progress/warning output stays visible # without flipping $? on successful native commands (PS 5.1 treats # stderr records as errors that set $? = $false even on exit code 0). - & $Command 2>&1 | Out-Host + # Redact per record: uv/pip echo index URLs (credentials and all) in + # their errors, and verbose mode must not bypass the quiet path's + # redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched. + & $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host } else { $output = & $Command 2>&1 | Out-String if ($LASTEXITCODE -ne 0) { - Write-Host $output -ForegroundColor Red + Write-Host (Redact-InstallOutput $output) -ForegroundColor Red } } return [int]$LASTEXITCODE @@ -2535,6 +2699,8 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode $VenvPyExe = Join-Path $VenvDir "Scripts\python.exe" $installedTorchTag = $null $shouldRebuild = $false + # Set when a stale venv under a pin is repaired in place (force-reinstall) not wiped. + $script:PinChangedForceReinstall = $false if (Test-Path -LiteralPath $VenvPyExe) { try { @@ -2551,10 +2717,14 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode if ($finished -and $proc.ExitCode -eq 0 -and $torchVer) { if ($torchVer -match '\+(cu\d+)') { $installedTorchTag = $Matches[1] + } elseif ($torchVer -match '\+rocm') { + # Any +rocm / gfx wheel -> generic "rocm" flavor (the exact version is + # repaired later by install_python_stack.py; here we only need the flavor). + $installedTorchTag = "rocm" } elseif ($torchVer -match '\+cpu') { $installedTorchTag = "cpu" } else { - # Untagged wheel (plain "2.x.y" from PyPI) -- treat as cpu + # Untagged wheel (plain "2.x.y" from PyPI) -> cpu. $installedTorchTag = "cpu" } } else { @@ -2570,12 +2740,71 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode } if (-not $shouldRebuild) { - $expectedTorchTag = if ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" } - if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { + $_pinnedIdx = Get-PinnedTorchIndexUrl + $_expectedKnown = $true + if ($_pinnedIdx) { + $_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx + # Digit-gated like the install selection: a custom rocm-* leaf (rocm-current / + # rocm-rel-7.2.1) is NOT a ROCm family and must not be stale-compared. + if (Test-PipRocmFamilyLeaf $_pinLeaf) { + # Don't collapse a pinned ROCm/gfx leaf to a generic "rocm" (would mask a family + # change, rocm6.4 -> gfx1151). Get-RocmPinStaleTags uses the SAME 2.11 allowlist + # as the install path, so a gfx110X-all/gfx90a/gfx908 pin on a <2.11 wheel is NOT stale. + $_rocmTags = Get-RocmPinStaleTags -PinLeaf $_pinLeaf -TorchVersion $torchVer + $expectedTorchTag = $_rocmTags.Expected + $installedTorchTag = $_rocmTags.Installed + } elseif ((Test-CudaFamilyLeaf $_pinLeaf) -or $_pinLeaf -eq 'cpu') { + # cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds; + # /custom and /current fall through to the unknown-index branch below. + $expectedTorchTag = $_pinLeaf + } else { + # Custom index whose leaf is not a torch flavor (a /simple mirror): the + # flavor can't be inferred, so never treat the venv as stale over it. + $_expectedKnown = $false + $expectedTorchTag = $installedTorchTag + } + } elseif ($HasNvidiaSmi) { + $expectedTorchTag = Get-PytorchCudaTag + } elseif ($HasROCm -or $script:ROCmGfxArch) { + # AMD/ROCm host with no explicit pin: an existing +rocm wheel is correct (gfx arch + # counts even when $HasROCm is false). But only the arches the install path maps to a + # repo.amd.com index get ROCm torch; an unmapped arch installs CPU, so expect "cpu" + # for those or a correct CPU venv rebuilds every update. + $_rocmWheelArches = @( + "gfx1201", "gfx1200", # RDNA 4 + "gfx1151", "gfx1150", # RDNA 3.5 (Strix Halo/Point) + "gfx1103", "gfx1102", "gfx1101", "gfx1100", # RDNA 3 + "gfx90a", "gfx908" # MI200 / MI100 + ) + if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) { + # A correct +rocm wheel is not stale. A CPU wheel on a supported AMD arch is + # NOT wiped either (the AMD Windows ROCm override below upgrades it in place); + # expect "cpu" for that case. A wrong CUDA wheel still rebuilds. + if ($installedTorchTag -eq "cpu") { + $expectedTorchTag = "cpu" + } else { + $expectedTorchTag = "rocm" + } + } else { + $expectedTorchTag = "cpu" + } + } else { + $expectedTorchTag = "cpu" + } + if ($_expectedKnown -and $installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) { $shouldRebuild = $true } } + # A stale venv under a pin whose torch still imports is repaired IN PLACE (the dependency + # pass force-reinstalls from the pin). The rebuild path wipes the venv and would strand a + # direct `studio update`; only a broken venv or an unpinned drift wipes. + if ($shouldRebuild -and $_pinnedIdx -and $installedTorchTag) { + substep "Torch-index pin changed ($installedTorchTag) -- reinstalling torch from the pin in place." "Cyan" + $script:PinChangedForceReinstall = $true + $shouldRebuild = $false + } + if ($shouldRebuild) { $reason = if ($installedTorchTag) { "torch $installedTorchTag != required $expectedTorchTag" } else { "torch could not be imported" } if ($InstallerManagedSetup) { @@ -2653,23 +2882,41 @@ if (Get-Command uv -ErrorAction SilentlyContinue) { # Helper: install a package, preferring uv with pip fallback function Fast-Install { param([Parameter(ValueFromRemainingArguments=$true)]$Args_) - if ($UseUv) { - $VenvPy = (Get-Command python).Source - # An explicit --index-url must win. Inherited uv index env vars otherwise - # override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop - # them only for index-pinned installs; mirrors still apply elsewhere. - $saved = @{} - if (@($Args_) -contains '--index-url') { - foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { - $saved[$n] = [Environment]::GetEnvironmentVariable($n) - Remove-Item "Env:$n" -ErrorAction SilentlyContinue - } + # An explicit --index-url must win: inherited uv index vars otherwise pull CPU torch over + # the CUDA/ROCm build (#6898), so drop them for pinned installs (scrub covers the whole + # function since the pip fallback honours PIP_* too). UV_TORCH_BACKEND / UV_FIND_LINKS also + # reroute; UV_NO_CONFIG=1 (+ dropping UV_CONFIG_FILE) stops a uv.toml index outranking the + # pin (uv 0.10); PIP_NO_INDEX / PIP_INDEX_URL would defeat the pinned --index-url in pip. + $saved = @{} + $pinned = @($Args_) -contains '--index-url' + if ($pinned) { + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', + 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'PIP_EXTRA_INDEX_URL', 'PIP_FIND_LINKS', + 'PIP_NO_INDEX', 'PIP_INDEX_URL', + 'UV_CONFIG_FILE', 'UV_NO_CONFIG', 'PIP_CONFIG_FILE') { + $saved[$n] = [Environment]::GetEnvironmentVariable($n) + Remove-Item "Env:$n" -ErrorAction SilentlyContinue } - try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 } - finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } } - if ($LASTEXITCODE -eq 0) { return } + $env:UV_NO_CONFIG = '1' + # A `pip config` global.extra-index-url still adds indexes to the pip FALLBACK; + # PIP_CONFIG_FILE = 'nul' (Windows devnull) loads NO config (uv ignores pip config). + $env:PIP_CONFIG_FILE = 'nul' + } + try { + if ($UseUv) { + $VenvPy = (Get-Command python).Source + $result = & uv pip install --python $VenvPy @Args_ 2>&1 + if ($LASTEXITCODE -eq 0) { return } + } + & python -m pip install @Args_ 2>&1 + } + finally { + if ($pinned) { + Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue + Remove-Item "Env:PIP_CONFIG_FILE" -ErrorAction SilentlyContinue + } + foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } } - & python -m pip install @Args_ 2>&1 } # ── Check if Python deps need updating ── @@ -2752,6 +2999,10 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1) # pip install unsloth 2>&1 | Out-Null # } +# A torch-index pin change repairs in place: force the dependency pass so the torch install +# below force-reinstalls from the new pin (else the fast path keeps the old wheel). +if ($script:PinChangedForceReinstall) { $SkipPythonDeps = $false } + if (-not $SkipPythonDeps) { if ($script:UnslothVerbose) { @@ -2779,7 +3030,13 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir [Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User') substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)" -if ($HasNvidiaSmi) { +# Explicit pin (URL or family) wins over GPU probing and suppresses the AMD reroute below; +# matches install.sh / install.ps1 / install_python_stack.py. +$PinnedTorchIndexUrl = Get-PinnedTorchIndexUrl +$TorchIndexPinned = [bool]$PinnedTorchIndexUrl +if ($PinnedTorchIndexUrl) { + $CuTag = Get-TorchIndexLeaf $PinnedTorchIndexUrl +} elseif ($HasNvidiaSmi) { $CuTag = Get-PytorchCudaTag } else { $CuTag = "cpu" @@ -2800,7 +3057,7 @@ $ROCmIndexUrl = $null # SDK -- which flips Unsloth out of chat-only (CHAT_ONLY) and enables Train/Export. # Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed # ROCm install still falls back to CPU below, so this is safe. -if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") { +if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") { $amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" } $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 @@ -2850,8 +3107,45 @@ if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") { } } +# A pinned gfx*/rocm index skips the auto-reroute above; route it through the ROCm install path +# with the same floor/companions the unpinned AMD path uses (mirrors install.ps1), else the CUDA +# branch installs bare torch and resolves a known-bad wheel for gfx115x/gfx120x/rocm>=7.2. +if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) { + $_pinLeaf = Get-TorchIndexLeaf $PinnedTorchIndexUrl + $_pinRocm211 = $false + # Anchor the match ($) so a suffixed custom leaf (rocm7.2-private) falls through to the + # verbatim install instead of being floored by its rocm7.2 prefix. + if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') { + # Only KNOWN-2.11 rocm (rocm7.2) gets the floor (no speculative floor). Matches + # Test-RocmKnown211Version / _ROCM_KNOWN_TORCH211_VERSIONS. + $_pinRocm211 = Test-RocmKnown211Version -Major ([int]$Matches[1]) -Minor ([int]$Matches[2]) + } + # Only the 2.11 gfx arches need the floor; others publish <2.11 and stay bare. Reuse + # Test-RocmGfx211Leaf so this allowlist and the stale-venv check never diverge. + $_pinGfx211 = Test-RocmGfx211Leaf $_pinLeaf + if ($_pinGfx211 -or $_pinRocm211) { + $ROCmIndexUrl = $PinnedTorchIndexUrl + $ROCmTorchSpec = "torch>=2.11.0,<2.12.0" + $ROCmVisionSpec = "torchvision>=0.26.0,<0.27.0" + $ROCmAudioSpec = "torchaudio>=2.11.0,<2.12.0" + substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchSpec" "Cyan" + } elseif (Test-PipRocmFamilyLeaf $_pinLeaf) { + # Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with + # bare specs. Only EXACT rocm and gfx* are --index-url families; a suffixed + # leaf stays on the verbatim path. Mirrors install.ps1 / _is_pip_rocm_family_leaf. + $ROCmIndexUrl = $PinnedTorchIndexUrl + $ROCmTorchSpec = "torch" + $ROCmVisionSpec = "torchvision" + $ROCmAudioSpec = "torchaudio" + } +} + $PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" } +# A full URL pin is used verbatim; a family pin already set $CuTag. A pinned ROCm install +# goes through $ROCmIndexUrl; on failure the fallback uses the CPU index, not the ROCm pin. +$TorchInstallIndexUrl = if ($ROCmIndexUrl) { "$PyTorchWhlBase/cpu" } elseif ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" } + $ROCmCpuFallback = $false if ($ROCmIndexUrl) { substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..." @@ -2859,7 +3153,7 @@ if ($ROCmIndexUrl) { substep " enforcing $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec (known _grouped_mm bug in older wheels)" "Cyan" } if ($script:UnslothVerbose) { - Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl + Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host $torchInstallExit = $LASTEXITCODE $output = "" } else { @@ -2868,7 +3162,7 @@ if ($ROCmIndexUrl) { } if ($torchInstallExit -ne 0) { Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow - Write-Host $output -ForegroundColor Yellow + Write-Host (Redact-InstallOutput $output) -ForegroundColor Yellow $ROCmIndexUrl = $null $ROCmCpuFallback = $true } else { @@ -2878,42 +3172,70 @@ if ($ROCmIndexUrl) { } } -if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") { +if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { substep "installing PyTorch (CPU-only)..." - # After an AMD ROCm fallback, force-reinstall so a partially-installed ROCm torch - # (which still satisfies the CPU torch>= range) is replaced by the CPU build. Skip - # the forced reinstall on a genuine CPU-only host so the common path stays fast. - # Build the array directly: an if-expression collapses @("x") to a scalar string, - # which @splat would then enumerate char-by-char into broken single-letter args. + # After an AMD ROCm fallback, force-reinstall so a partial ROCm torch (which satisfies the + # CPU torch>= range) is replaced by the CPU build; skip on a genuine CPU host to stay fast. + # $ROCmCpuFallback matters when a PINNED ROCm index failed ($CuTag is still the rocm leaf). + # Build the array directly: an if-expression collapses @("x") to a scalar @splat would + # enumerate char-by-char. $cpuForce = @() if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") } + # --force-reinstall on a pin change: a stale +cu / +rocm wheel still satisfies the CPU + # torch>= range, so uv would keep it and only swap companions. + if ($script:PinChangedForceReinstall) { $cpuForce = @("--force-reinstall") } + # A PINNED cpu index installs the bounded trio (parity with _CPU_TORCH_PKG_SPEC): the /cpu + # index serves newer torch, and _ensure_cpu_torch keeps any CPU build, so a bare trio could + # land an unsupported version. Unpinned CPU hosts keep the bare trio (pre-pin behavior). + $cpuTorchSpec = "torch"; $cpuVisionSpec = "torchvision"; $cpuAudioSpec = "torchaudio" + if ($TorchIndexPinned) { + $cpuTorchSpec = "torch>=2.4,<2.12.0" + $cpuVisionSpec = "torchvision>=0.19,<0.27.0" + $cpuAudioSpec = "torchaudio>=2.4,<2.12.0" + } if ($script:UnslothVerbose) { - Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" + Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host $torchInstallExit = $LASTEXITCODE $output = "" } else { - $output = Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" | Out-String + $output = Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | Out-String $torchInstallExit = $LASTEXITCODE } if ($torchInstallExit -ne 0) { Write-Host "[FAILED] PyTorch install failed (exit code $torchInstallExit)" -ForegroundColor Red - Write-Host $output -ForegroundColor Red + Write-Host (Redact-InstallOutput $output) -ForegroundColor Red exit 1 } } elseif (-not $ROCmIndexUrl) { substep "installing PyTorch with CUDA support ($CuTag)..." substep "(This download is ~2.8 GB -- may take a few minutes)" + # --force-reinstall on a pin change: an installed cuXXX wheel satisfies the bare torch + # requirement (PEP 440 ignores the +cuXXX tag), so without it a changed CUDA pin (cu126 + # -> cu128) never applies. + $cudaForce = @() + if ($script:PinChangedForceReinstall) { $cudaForce = @("--force-reinstall") } + # An unknown-leaf custom pin (/simple, /current) routes here with $CuTag as that leaf. Bound + # the trio like the fresh custom-pin paths so a mirror can't pull an ABI-newer companion + # against the capped torch. Known cu* leaves keep bare specs. + $cudaTorchSpec = "torch" + $cudaVisionSpec = "torchvision" + $cudaAudioSpec = "torchaudio" + if ($TorchIndexPinned -and -not (Test-CudaFamilyLeaf $CuTag)) { + $cudaTorchSpec = "torch>=2.4,<2.11.0" + $cudaVisionSpec = "torchvision>=0.19,<0.26.0" + $cudaAudioSpec = "torchaudio>=2.4,<2.11.0" + } if ($script:UnslothVerbose) { - Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" + Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host $torchInstallExit = $LASTEXITCODE $output = "" } else { - $output = Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" | Out-String + $output = Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | Out-String $torchInstallExit = $LASTEXITCODE } if ($torchInstallExit -ne 0) { Write-Host "[FAILED] PyTorch CUDA install failed (exit code $torchInstallExit)" -ForegroundColor Red - Write-Host $output -ForegroundColor Red + Write-Host (Redact-InstallOutput $output) -ForegroundColor Red exit 1 } @@ -2929,7 +3251,7 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") { } if ($tritonInstallExit -ne 0) { substep "Triton install failed -- torch.compile may not work" "Yellow" - Write-Host $output -ForegroundColor Yellow + Write-Host (Redact-InstallOutput $output) -ForegroundColor Yellow } else { substep "Triton for Windows installed (enables torch.compile)" } @@ -3026,7 +3348,7 @@ foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.8.0", "hf_xet==1.4 } if ($t5PkgExit -ne 0) { Write-Host "[FAIL] Could not install $pkg into .venv_t5_530/" -ForegroundColor Red - Write-Host $output -ForegroundColor Red + Write-Host (Redact-InstallOutput $output) -ForegroundColor Red $ErrorActionPreference = $prevEAP_t5 exit 1 } @@ -3061,7 +3383,7 @@ foreach ($pkg in @("transformers==5.5.0", "huggingface_hub==1.8.0", "hf_xet==1.4 } if ($t5PkgExit -ne 0) { Write-Host "[FAIL] Could not install $pkg into .venv_t5_550/" -ForegroundColor Red - Write-Host $output -ForegroundColor Red + Write-Host (Redact-InstallOutput $output) -ForegroundColor Red $ErrorActionPreference = $prevEAP_t5 exit 1 } @@ -3096,7 +3418,7 @@ foreach ($pkg in @("transformers==5.10.2", "huggingface_hub==1.8.0", "hf_xet==1. } if ($t5PkgExit -ne 0) { Write-Host "[FAIL] Could not install $pkg into .venv_t5_510/" -ForegroundColor Red - Write-Host $output -ForegroundColor Red + Write-Host (Redact-InstallOutput $output) -ForegroundColor Red $ErrorActionPreference = $prevEAP_t5 exit 1 } diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 666ea7ce10..b3a9b99c55 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -10,6 +10,8 @@ import pytest REPO_ROOT = Path(__file__).resolve().parents[2] INSTALL_SH = REPO_ROOT / "install.sh" INSTALL_PS1 = REPO_ROOT / "install.ps1" +SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1" +STACK_PY = REPO_ROOT / "studio" / "install_python_stack.py" class TestNoTorchBackendAutoInInstallSh: @@ -180,3 +182,607 @@ class TestUvBytecodeCompileTimeout: assert ( '$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text ), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT" + + +class TestTorchIndexOverrideParity: + """Every installer must honor UNSLOTH_TORCH_INDEX_URL / _FAMILY so a pinned wheel + index wins over GPU probing on all platforms (no asymmetric, per-OS coverage).""" + + @pytest.mark.parametrize( + "path", + [INSTALL_SH, INSTALL_PS1, SETUP_PS1, STACK_PY], + ids = ["install.sh", "install.ps1", "setup.ps1", "install_python_stack.py"], + ) + def test_installer_reads_override_env(self, path): + text = path.read_text(encoding = "utf-8") + for var in ("UNSLOTH_TORCH_INDEX_URL", "UNSLOTH_TORCH_INDEX_FAMILY"): + assert var in text, f"{path.name} does not honor {var}" + + @pytest.mark.parametrize( + "path", + [INSTALL_PS1, SETUP_PS1], + ids = ["install.ps1", "setup.ps1"], + ) + def test_amd_reroute_guarded_when_pinned(self, path): + # The AMD ROCm reroute must be skipped when the index is explicitly pinned, + # so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten. + text = path.read_text(encoding = "utf-8") + assert ( + "TorchIndexPinned" in text + ), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag" + + def test_cuda_pin_overrides_cvd_hide_gate(self): + # A pinned cu* index skips ALL host-GPU probing, so the CUDA repair must clear the + # CUDA_VISIBLE_DEVICES hide gate too (else the GPU-less CI case bails). + text = STACK_PY.read_text(encoding = "utf-8") + m = re.search(r"def _ensure_cuda_torch\(\).*?(?=\ndef )", text, re.DOTALL) + assert m, "could not locate _ensure_cuda_torch" + body = m.group(0) + assert "_cuda_pinned" in body, ( + "_ensure_cuda_torch should compute a CUDA-pin flag so the pin can " + "override the CVD hide gate" + ) + assert re.search( + r"if not _cuda_pinned and _cvd is not None", body + ), "the CVD hide gate must be bypassed when a CUDA index is pinned" + + def test_cpu_repair_pins_supported_torch_range(self): + # The explicit-CPU repair must use the bounded CPU/CUDA spec, not a bare trio (the + # /cpu index serves torch 2.11+, so a bare install could resolve out of range). + text = STACK_PY.read_text(encoding = "utf-8") + m = re.search(r"def _ensure_cpu_torch\(\).*?(?=\ndef )", text, re.DOTALL) + assert m, "could not locate _ensure_cpu_torch" + body = m.group(0) + assert "_CPU_TORCH_PKG_SPEC" in body, ( + "_ensure_cpu_torch should install the bounded _CPU_TORCH_PKG_SPEC, " + "not a bare torch/torchvision/torchaudio trio" + ) + + def test_setup_ps1_stale_check_gates_rocm_on_supported_arch(self): + # The stale check must expect ROCm torch only for arches the install path maps to a + # repo.amd.com index; expecting "rocm" for an unmapped arch marks a good CPU venv stale. + text = SETUP_PS1.read_text(encoding = "utf-8") + assert "_rocmWheelArches" in text, ( + "setup.ps1 stale check should restrict the ROCm expected-tag to the " + "supported gfx wheel arches" + ) + + +class TestGfx211AllowlistParity: + """The gfx per-arch 2.11-floor leaves (gfx120X-all / gfx1151 / gfx1150) must be the + SAME set in every installer and its stale/mismatch check. When they diverged, a + pinned gfx110X-all / gfx90a / gfx908 wheel (<2.11) was force-reinstalled every update.""" + + EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150"} + + def test_install_sh_allowlist(self): + text = INSTALL_SH.read_text(encoding = "utf-8").lower() + # install.sh: the TORCH_CONSTRAINT case (rocm7.2|gfx120x-all|gfx1151|gfx1150). + m = re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150", text) + assert m, "install.sh gfx-2.11 allowlist case not found / changed" + + def test_install_ps1_allowlist(self): + text = INSTALL_PS1.read_text(encoding = "utf-8").lower() + m = re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text) + assert m, "install.ps1 $_pinGfx211 allowlist not found / changed" + + def test_setup_ps1_defines_single_allowlist_helper(self): + # setup.ps1 must define the allowlist once (Test-RocmGfx211Leaf) and reuse it, so + # the stale check and install spec can't disagree. + text = SETUP_PS1.read_text(encoding = "utf-8") + assert ( + "function Test-RocmGfx211Leaf" in text + ), "setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper" + assert re.search( + r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower() + ), "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist" + assert "$_pinGfx211 = Test-RocmGfx211Leaf" in text, ( + "setup.ps1 install-spec path should reuse Test-RocmGfx211Leaf, not " + "re-hardcode the allowlist (they must not diverge)" + ) + + def test_stack_py_allowlist(self): + text = STACK_PY.read_text(encoding = "utf-8").lower() + assert ( + '"gfx120x-all", "gfx1151", "gfx1150"' in text + ), "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed" + + +class TestCudaLeafDigitParity: + """A wheel-family leaf is CUDA only when it is "cu" + digits (cu118/cu128/...). + A bare cu* glob wrongly catches mirror leaves like /custom or /current; when + that happened the venv was marked stale and rebuilt on every run. Every + installer must require a digit after "cu" in its family/CUDA classification.""" + + def test_stack_py_requires_cu_digit(self): + text = STACK_PY.read_text(encoding = "utf-8") + # EXACT cu+digits: a custom leaf like cu128-private must route to the + # verbatim/unknown path, not be compared against the installed +cu128 tag. + assert re.search( + r'r"cu\[0-9\]\+"', text + ), "install_python_stack.py _is_cuda_family_leaf must fullmatch cu[0-9]+" + + def test_setup_ps1_requires_cu_digit(self): + text = SETUP_PS1.read_text(encoding = "utf-8") + # EXACT cu+digits: cu128-private must not classify as CUDA (it would become + # the expected tag and rebuild the venv on every update). + assert re.search( + r"'\^cu\[0-9\]\+\$'", text + ), "setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9]+$, not a cu* prefix" + # The stale-venv branch must go through the digit-guarded helper. + assert ( + "Test-CudaFamilyLeaf $_pinLeaf" in text + ), "setup.ps1 stale check should classify CUDA via Test-CudaFamilyLeaf" + + def test_install_ps1_requires_cu_digit_in_gpu_branch(self): + text = INSTALL_PS1.read_text(encoding = "utf-8") + assert re.search( + r"'\^cu\[0-9\]'", text + ), "install.ps1 Get-TauriGpuBranch must require a digit after cu" + + def test_install_sh_requires_cu_digit_in_gpu_branch(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + # The _tauri_gpu_branch cuda case must be cu[0-9]*, not a bare cu*. + assert re.search( + r"cu\[0-9\]\*\)\s*echo \"cuda\"", text + ), "install.sh _tauri_gpu_branch cuda case must be cu[0-9]*, not cu*" + + def test_install_sh_backend_export_requires_cu_digit(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + # Brand CUDA only on cu[0-9]*; a bare catch-all *) -> cuda would mis-brand + # /current, /custom pins and skip ROCm repair on AMD hosts. + assert re.search( + r'cu\[0-9\]\*\)\s*export UNSLOTH_TORCH_BACKEND="cuda"', text + ), "install.sh backend export must brand cuda only on cu[0-9]*" + # An unknown leaf must NOT commit a cuda backend (it unsets instead). + assert re.search( + r"\*\)\s*unset UNSLOTH_TORCH_BACKEND", text + ), "install.sh backend export must unset (not force cuda) on an unknown leaf" + + def test_install_sh_lowercases_backend_leaf(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + # The leaf feeding both the backend case and the 2.11 floor case must be + # lowercased so the canonical gfx120X-all (capital X) matches. + assert re.search( + r"_torch_index_leaf=\$\(printf '%s' \"\$_torch_index_leaf\" \| tr '\[:upper:\]' '\[:lower:\]'\)", + text, + ), "install.sh must lowercase _torch_index_leaf before the gfx/rocm/cu case matches" + + +class TestKnown211SetParity: + """The KNOWN-2.11 rocm/gfx set must be identical across all four installers: + exactly {rocm7.2} plus the gfx allowlist {gfx120x-all, gfx1151, gfx1150}. + rocm7.3 / torch 2.12 do not exist, so no side may floor them speculatively.""" + + def test_install_sh_known_211_leaf_is_rocm72_and_gfx_allowlist(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + # The 2.11 floor case matches exactly rocm7.2 + the three gfx leaves. + assert re.search( + r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text + ), "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150" + # No speculative rocm7.3 anywhere. + assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3" + + def test_python_known_211_versions_is_only_rocm72(self): + text = STACK_PY.read_text(encoding = "utf-8") + assert "_ROCM_KNOWN_TORCH211_VERSIONS" in text + # The frozenset literal is exactly {(7, 2)}. + m = re.search(r"_ROCM_KNOWN_TORCH211_VERSIONS[^=]*=\s*frozenset\(\{([^}]*)\}\)", text) + assert m is not None, "install_python_stack.py must define _ROCM_KNOWN_TORCH211_VERSIONS" + assert "(7, 2)" in m.group(1) + assert "7, 3" not in m.group(1) and "7, 1" not in m.group(1) + + def test_setup_ps1_known_211_helper_is_only_rocm72(self): + text = SETUP_PS1.read_text(encoding = "utf-8") + assert "Test-RocmKnown211Version" in text + # The predicate is Major -eq 7 -and Minor -eq 2 (only rocm7.2). + assert re.search( + r"Test-RocmKnown211Version[\s\S]{0,400}\$Major -eq 7 -and \$Minor -eq 2", text + ), "setup.ps1 Test-RocmKnown211Version must accept only rocm7.2" + + def test_install_ps1_pin_floor_is_only_rocm72(self): + text = INSTALL_PS1.read_text(encoding = "utf-8") + # The pinned-ROCm install-spec floor must be Major -eq 7 -and Minor -eq 2, + # not the speculative >= 2 that would floor a non-existent rocm7.3. + assert re.search( + r"\$_pinRocm211 = \(\[int\]\$Matches\[1\] -eq 7 -and \[int\]\$Matches\[2\] -eq 2\)", + text, + ), "install.ps1 pinned-ROCm floor must be rocm7.2 only (no speculative >= 2)" + + def test_ps1_pin_floor_gate_is_anchored(self): + """The floor-selection gate that reads $_pinRocm211 from the raw leaf must anchor + the rocm match ($), or a suffixed custom leaf (rocm7.2-private) matches the rocm7.2 + prefix, takes the 2.11-floor branch, and is force-routed through the ROCm path + before the exact-match elseif can send it to the verbatim install (Codex P2).""" + for path, label in ((INSTALL_PS1, "install.ps1"), (SETUP_PS1, "setup.ps1")): + text = path.read_text(encoding = "utf-8") + assert "-match '^rocm(\\d+)\\.(\\d+)$'" in text, ( + f"{label} floor gate must anchor the rocm match (^rocm(\\d+)\\.(\\d+)$) so a " + "suffixed custom leaf is not floored/routed as rocm7.2" + ) + assert ( + "-match '^rocm(\\d+)\\.(\\d+)'\n" not in text + ), f"{label} floor gate must not use the unanchored ^rocm(\\d+)\\.(\\d+) prefix" + + def test_install_ps1_bounds_unknown_leaf_pinned_torch(self): + """install.ps1's pinned-torch install must bound BOTH companions on EVERY + index, cu families included: torchaudio 2.11 dropped its exact torch + pin from the wheel metadata, so a bare companion beside torch<2.11 can + resolve a mismatched 2.11.0 build (Codex P2, then unconditional per the + torchaudio 2.11 unpinning).""" + text = INSTALL_PS1.read_text(encoding = "utf-8") + assert ( + '$_pinVisionSpec = "torchvision>=0.19,<0.26.0"' in text + ), "install.ps1 custom-pin install must bound torchvision (>=0.19,<0.26.0)" + assert ( + '$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"' in text + ), "install.ps1 custom-pin install must bound torchaudio (>=2.4,<2.11.0)" + # No cu-family exemption: the bounds apply unconditionally. + assert ( + "$_pinCuLeaf" not in text + ), "install.ps1 must bound companions on every index (no cu-family exemption)" + # The bounded companions must actually be passed to the install command. + assert re.search( + r'"torch>=2\.4,<2\.11\.0" \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl', + text, + ), "install.ps1 custom-pin install must pass the bounded companion specs to uv" + + def test_gfx_allowlist_matches_across_installers(self): + # The gfx 2.11 allowlist {gfx120x-all, gfx1151, gfx1150} must appear in each. + gfx = ("gfx120x-all", "gfx1151", "gfx1150") + for path, label in ( + (INSTALL_SH, "install.sh"), + (INSTALL_PS1, "install.ps1"), + (SETUP_PS1, "setup.ps1"), + (STACK_PY, "install_python_stack.py"), + ): + low = path.read_text(encoding = "utf-8").lower() + for g in gfx: + assert g in low, f"{label} missing gfx 2.11 allowlist member {g}" + + +class TestPinnedRocmLeafDigitParity: + """A pinned index is a pip ROCm --default-index family only when its leaf is an + EXACT rocm+digits (rocm7 / rocm7.2) or gfx*. A ^rocm[0-9] PREFIX (or a bare rocm* + glob) wrongly catches a custom mirror / find-links leaf (rocm-current / + rocm-rel-7.2.1) AND a suffixed private-mirror leaf (rocm7.2-private / rocm7-current), + routing it through the ROCm install path (which silently falls back to CPU on + failure) or skipping the custom-index companion bounds, instead of the verbatim + --default-index install. All installers must match the family EXACTLY: Python and + install.sh via a shared _is_pip_rocm_family_leaf, setup.ps1 via Test-PipRocmFamilyLeaf, + install.ps1 via an anchored ^rocm[0-9]+(\\.[0-9]+)?$ reroute.""" + + def test_install_ps1_pinned_reroute_requires_rocm_digit(self): + text = INSTALL_PS1.read_text(encoding = "utf-8") + # The pinned gfx*/rocm reroute must match rocm EXACTLY (anchored), so a suffixed + # rocm7.2-private / rocm-current falls through to the verbatim --default-index path. + assert "-match '^rocm[0-9]+(\\.[0-9]+)?$'" in text, ( + "install.ps1 pinned-index reroute must anchor the rocm match " + "(^rocm[0-9]+(\\.[0-9]+)?$), not a bare -like 'rocm*' or an unanchored ^rocm\\d" + ) + # Neither the broad glob nor the unanchored prefix may drive that reroute. + assert ( + "-like 'rocm*'" not in text + ), "install.ps1 must not route a pinned index on a bare -like 'rocm*' glob" + assert ( + "-match '^rocm\\d'" not in text + ), "install.ps1 must not route a pinned index on an unanchored -match '^rocm\\d'" + + def test_setup_ps1_pinned_reroute_requires_rocm_digit(self): + text = SETUP_PS1.read_text(encoding = "utf-8") + # setup.ps1 routes every family decision through Test-PipRocmFamilyLeaf, which + # anchors the rocm match so a suffixed custom leaf stays on the verbatim path. + assert ( + "function Test-PipRocmFamilyLeaf" in text + ), "setup.ps1 must define Test-PipRocmFamilyLeaf (the exact rocm/gfx family gate)" + assert "'^rocm[0-9]+(\\.[0-9]+)?$'" in text, ( + "setup.ps1 Test-PipRocmFamilyLeaf must anchor the rocm match " + "(^rocm[0-9]+(\\.[0-9]+)?$) so rocm7.2-private / rocm-current stay verbatim" + ) + pinned_block = text[text.find("$_pinGfx211 = Test-RocmGfx211Leaf") :][:2000] + assert ( + "-like 'rocm*'" not in pinned_block + ), "setup.ps1 pinned reroute must not route on a bare -like 'rocm*' glob" + + def test_install_sh_repairable_requires_rocm_digit(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + # _torch_index_repairable routes rocm/gfx through the exact-match helper. + assert ( + "_is_pip_rocm_family_leaf" in text + ), "install.sh must define/use _is_pip_rocm_family_leaf for the exact rocm gate" + # gfx needs a following digit: gfx-private / gfxfoo are custom verbatim pins. + assert re.search( + r'case "\$1" in\n\s*gfx\[0-9\]\*\) return 0', text + ), "install.sh _is_pip_rocm_family_leaf must treat only gfx* as a family" + assert not re.search( + r'case "\$1" in\n\s*gfx\*\) return 0', text + ), "install.sh _is_pip_rocm_family_leaf must not family-match a bare gfx* glob" + + def test_stack_py_pip_rocm_family_requires_digit(self): + text = STACK_PY.read_text(encoding = "utf-8") + assert re.search( + r'fullmatch\(r"rocm\\d\+\(\?:\\\.\\d\+\)\?", leaf\)', text + ), "install_python_stack.py _is_pip_rocm_family_leaf must fullmatch rocm\\d+(?:\\.\\d+)?" + # The unanchored prefix must be gone from the family/flavor gates. + assert ( + 're.match(r"^rocm\\d"' not in text + ), "install_python_stack.py must not gate a family on an unanchored re.match(^rocm\\d)" + + def test_install_sh_rocm_side_effects_digit_gated(self): + """The AMD bitsandbytes + 'repair ROCm torch' side effects must fire only on + an EXACT ROCm family (rocm7.2/gfx*), not a bare */rocm* whole-URL glob nor a + ^rocm[0-9] prefix that catches a custom CPU/CUDA index like /rocm-current or a + suffixed /rocm7.2-private and force-repairs it from the wrong --default-index.""" + text = INSTALL_SH.read_text(encoding = "utf-8") + assert ( + 'if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then\n _torch_index_is_rocm_family=true' + in text + ), "install.sh must set _torch_index_is_rocm_family from the exact-match helper" + assert ( + '[ "$_torch_index_is_rocm_family" = true ]' in text + ), "install.sh ROCm bnb/repair hooks must gate on _torch_index_is_rocm_family" + assert ( + "*/rocm*|*/gfx*)\n _install_bnb_rocm" not in text + ), "install.sh must not gate _install_bnb_rocm on a bare */rocm* whole-URL glob" + + +class TestPinnedIndexClearsUvEnvParity: + """Every installer must neutralise the uv index env vars for a pinned torch + install (#6898). uv treats the default index (--index-url / --default-index) as + lowest priority, so an inherited UV_INDEX / UV_EXTRA_INDEX_URL mirror would win + under uv's first-index strategy and pull torch from the wrong index -- after + which the pinned wheel index is silently never used.""" + + UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL") + + def test_install_sh_clears_uv_index_vars(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + assert ( + "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in text + ), "install.sh run_install_cmd must clear the uv index vars for --default-index installs" + + def test_install_ps1_clears_uv_index_vars(self): + text = INSTALL_PS1.read_text(encoding = "utf-8") + for var in self.UV_VARS: + assert var in text, f"install.ps1 must clear {var} for pinned installs" + + def test_setup_ps1_clears_uv_index_vars(self): + text = SETUP_PS1.read_text(encoding = "utf-8") + for var in self.UV_VARS: + assert var in text, f"setup.ps1 must clear {var} for pinned installs" + + def test_stack_py_clears_uv_index_vars(self): + text = STACK_PY.read_text(encoding = "utf-8") + assert "_install_env_for_cmd" in text, ( + "install_python_stack.py must scrub inherited uv index vars for pinned " + "installs via _install_env_for_cmd (parity with install.sh #6898)" + ) + for var in self.UV_VARS: + assert var in text, f"install_python_stack.py must clear {var} for pinned installs" + + def test_all_installers_clear_uv_torch_backend(self): + """uv's torch backend redirects torch resolution to its own per-backend + index even against an explicit pin, so every installer's pinned-install + scrub must clear UV_TORCH_BACKEND too.""" + sh = INSTALL_SH.read_text(encoding = "utf-8") + assert "-u UV_TORCH_BACKEND" in sh, "install.sh pinned scrub must clear UV_TORCH_BACKEND" + for path in (INSTALL_PS1, SETUP_PS1): + text = path.read_text(encoding = "utf-8") + assert ( + "'UV_TORCH_BACKEND'" in text + ), f"{path.name} pinned scrub must clear UV_TORCH_BACKEND" + stack = STACK_PY.read_text(encoding = "utf-8") + assert ( + '"UV_TORCH_BACKEND",' in stack + ), "install_python_stack.py strip tuple must include UV_TORCH_BACKEND" + + def test_stack_py_strips_pip_extra_index_for_pip_fallback(self): + """The pip fallback honours PIP_EXTRA_INDEX_URL (pip adds it IN ADDITION + to --index-url), so the pinned-command scrub must strip it.""" + stack = STACK_PY.read_text(encoding = "utf-8") + assert ( + '"PIP_EXTRA_INDEX_URL",' in stack + ), "install_python_stack.py strip tuple must include PIP_EXTRA_INDEX_URL" + + def test_all_installers_scrub_find_links(self): + """uv's --find-links (env UV_FIND_LINKS) adds candidate locations that can + satisfy torch off a pinned index; every pinned-install scrub must clear it.""" + sh = INSTALL_SH.read_text(encoding = "utf-8") + assert "-u UV_FIND_LINKS" in sh + for path in (INSTALL_PS1, SETUP_PS1): + assert "'UV_FIND_LINKS'" in path.read_text(encoding = "utf-8"), path.name + stack = STACK_PY.read_text(encoding = "utf-8") + assert '"UV_FIND_LINKS",' in stack and '"PIP_FIND_LINKS",' in stack + + def test_setup_ps1_scrub_covers_pip_fallback(self): + """setup.ps1's Fast-Install must keep the scrub active through the pip + fallback (pip honours PIP_EXTRA_INDEX_URL / PIP_FIND_LINKS in addition to + --index-url); restoring the vars before the fallback reopens the hole.""" + text = SETUP_PS1.read_text(encoding = "utf-8") + fi = text[text.find("function Fast-Install") :][:2500] + assert "'PIP_EXTRA_INDEX_URL'" in fi and "'PIP_FIND_LINKS'" in fi + # the pip fallback must sit INSIDE the try whose finally restores the vars + assert fi.find("python -m pip install") < fi.find( + "finally" + ), "pip fallback must run before the scrub is restored" + + def test_all_installers_disable_uv_config_for_pinned_installs(self): + """A DISCOVERED uv.toml / pyproject [tool.uv] outranks the CLI pin + (verified with uv 0.10: [pip] torch-backend = "cpu" and a non-default + [[index]] both resolve torch+cpu against an explicit --index-url / + --default-index cu126 pin; UV_NO_CONFIG=1 restores the pin). Every + installer's pinned scrub must set UV_NO_CONFIG=1 and drop UV_CONFIG_FILE.""" + sh = INSTALL_SH.read_text(encoding = "utf-8") + assert "-u UV_CONFIG_FILE UV_NO_CONFIG=1" in sh, ( + "install.sh run_install_cmd must set UV_NO_CONFIG=1 and drop " + "UV_CONFIG_FILE for --default-index installs" + ) + for path in (INSTALL_PS1, SETUP_PS1): + text = path.read_text(encoding = "utf-8") + assert "'UV_CONFIG_FILE'" in text, f"{path.name} must drop UV_CONFIG_FILE" + assert ( + "$env:UV_NO_CONFIG = '1'" in text + ), f"{path.name} must set UV_NO_CONFIG=1 for pinned installs" + stack = STACK_PY.read_text(encoding = "utf-8") + assert ( + '"UV_CONFIG_FILE",' in stack + ), "install_python_stack.py strip tuple must include UV_CONFIG_FILE" + assert ( + 'env["UV_NO_CONFIG"] = "1"' in stack + ), "_install_env_for_cmd must set UV_NO_CONFIG=1 for pinned installs" + + def test_pip_fallbacks_disable_pip_config_files(self): + """The pip FALLBACK (uv missing/failed) honours user/site pip config files + even with the PIP_* env vars stripped: `pip config set + global.extra-index-url` still adds indexes to a pinned install. pip loads + NO configuration files when PIP_CONFIG_FILE is the platform devnull, so + the two installers that HAVE a pip fallback (install_python_stack.py and + setup.ps1's Fast-Install) must set it in their pinned scrub. install.sh + and install.ps1 are uv-only (no python -m pip fallback) and need no + equivalent.""" + stack = STACK_PY.read_text(encoding = "utf-8") + assert 'env["PIP_CONFIG_FILE"] = os.devnull' in stack, ( + "_install_env_for_cmd must point PIP_CONFIG_FILE at os.devnull for " + "pinned installs (pip fallback isolation)" + ) + setup = SETUP_PS1.read_text(encoding = "utf-8") + assert "$env:PIP_CONFIG_FILE = 'nul'" in setup, ( + "setup.ps1 Fast-Install pinned scrub must point PIP_CONFIG_FILE at nul " + "(Windows devnull) so the pip fallback ignores user/site pip config" + ) + assert ( + "'PIP_CONFIG_FILE'" in setup + ), "setup.ps1 must save/restore PIP_CONFIG_FILE around the pinned scrub" + + def test_setup_ps1_bounds_unknown_leaf_pinned_torch(self): + """A first-time/changed unknown-leaf custom pin routes through setup.ps1's + CUDA branch; install.ps1's fresh pinned install, install.sh, and the Python + verbatim path bound the WHOLE trio, so the Windows update path must too -- a + private mirror serving newer torch OR newer companions must not lift the venv + above the supported range under the pin.""" + text = SETUP_PS1.read_text(encoding = "utf-8") + # The custom-leaf branch bounds torch AND both companions (parity with the + # other installers' custom-pin trio bounds), gated on a non-cu-family leaf. + for spec in ( + '$cudaTorchSpec = "torch>=2.4,<2.11.0"', + '$cudaVisionSpec = "torchvision>=0.19,<0.26.0"', + '$cudaAudioSpec = "torchaudio>=2.4,<2.11.0"', + ): + assert spec in text, f"setup.ps1 must bound the custom-leaf trio: {spec}" + assert ( + "if ($TorchIndexPinned -and -not (Test-CudaFamilyLeaf $CuTag)) {" in text + ), "the custom-leaf trio bounds must be gated on a pinned non-cu-family leaf" + assert ( + "Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec" in text + ), "setup.ps1's CUDA branch must install via the bounded spec variables" + + def test_setup_ps1_bounds_pinned_cpu_torch(self): + """setup.ps1's CPU branch must bound the trio under an explicit pin (parity with + _CPU_TORCH_PKG_SPEC): the /cpu index serves newer torch, and _ensure_cpu_torch + keeps any CPU build, so a bare pinned trio could land an unsupported version. + An unpinned CPU host keeps the bare trio (pre-pin behavior unchanged).""" + text = SETUP_PS1.read_text(encoding = "utf-8") + for spec in ( + '$cpuTorchSpec = "torch>=2.4,<2.12.0"', + '$cpuVisionSpec = "torchvision>=0.19,<0.27.0"', + '$cpuAudioSpec = "torchaudio>=2.4,<2.12.0"', + ): + assert spec in text, f"setup.ps1 must bound the pinned CPU trio: {spec}" + assert ( + "if ($TorchIndexPinned) {" in text + ), "the CPU trio bounds must be gated on an explicit pin" + assert ( + "Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce" in text + ), "setup.ps1's CPU branch must install via the spec variables" + # The ceilings mirror the Python repair spec exactly. + stack = STACK_PY.read_text(encoding = "utf-8") + spec_block = re.search(r"_CUDA_TORCH_PKG_SPEC[^(]*\(\s*(.*?)\)", stack, re.DOTALL) + assert spec_block and '"torch>=2.4,<2.12.0"' in spec_block.group(1), ( + "_CPU_TORCH_PKG_SPEC (via _CUDA_TORCH_PKG_SPEC) must keep the torch<2.12 " + "ceiling the setup.ps1 pinned CPU branch mirrors" + ) + + def test_setup_ps1_stale_check_requires_rocm_digit(self): + """The stale-venv check must use the same EXACT rocm/gfx gate as the install + selection (Test-PipRocmFamilyLeaf), or a custom rocm-* / suffixed rocm7.2-private + leaf is stale-compared as a family and force-reinstalls on every studio update.""" + text = SETUP_PS1.read_text(encoding = "utf-8") + anchor = text.find("$_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx") + assert anchor >= 0, "setup.ps1 stale check must classify the pinned leaf" + stale = text[anchor:][:2500] + assert ( + "Test-PipRocmFamilyLeaf" in stale + ), "setup.ps1 stale check must gate rocm leaves via the exact Test-PipRocmFamilyLeaf" + assert ( + stale.count("-like 'rocm*'") == 0 + ), "setup.ps1 stale check must not use a bare -like 'rocm*' glob" + assert ( + "-match '^rocm\\d'" not in stale + ), "setup.ps1 stale check must not use an unanchored -match '^rocm\\d'" + + +class TestIndexPathSlashTrimParity: + """Every installer must trim trailing PATH slashes only on the verbatim + UNSLOTH_TORCH_INDEX_URL override, preserving a ?query/#fragment token: a whole-URL + strip corrupts a base64 token ending in "/", a single strip leaves a double-slash leaf + empty. The helper must be DEFINED and WIRED into the override return in all four.""" + + def test_helper_defined_in_all_installers(self): + assert "def _trim_index_path_slashes(" in STACK_PY.read_text(encoding = "utf-8") + assert "_trim_index_path_slashes()" in INSTALL_SH.read_text(encoding = "utf-8") + assert "function Trim-IndexPathSlashes" in INSTALL_PS1.read_text(encoding = "utf-8") + assert "function Trim-IndexPathSlashes" in SETUP_PS1.read_text(encoding = "utf-8") + + def test_helper_wired_into_override_in_all_installers(self): + assert "_trim_index_path_slashes(url)" in STACK_PY.read_text(encoding = "utf-8") + assert '_url=$(_trim_index_path_slashes "$_url")' in INSTALL_SH.read_text(encoding = "utf-8") + assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in INSTALL_PS1.read_text( + encoding = "utf-8" + ) + assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in SETUP_PS1.read_text( + encoding = "utf-8" + ) + + +class TestInstallOutputRedactionParity: + """uv/pip failure text embeds the failing --index-url verbatim, so a captured install + log dumped on error can leak a user:token@ or ?token= secret. Every installer must + DEFINE a redaction helper and WIRE it into the captured-output print path.""" + + def test_helper_defined_in_all_installers(self): + assert "def _redact_install_output(" in STACK_PY.read_text(encoding = "utf-8") + assert "_redact_install_output()" in INSTALL_SH.read_text(encoding = "utf-8") + assert "function Redact-InstallOutput" in INSTALL_PS1.read_text(encoding = "utf-8") + assert "function Redact-InstallOutput" in SETUP_PS1.read_text(encoding = "utf-8") + + def test_helper_wired_into_failure_print(self): + # install.sh dumps the captured log through the redactor on failure. + assert '_redact_install_output "$_log"' in INSTALL_SH.read_text(encoding = "utf-8") + # Both ps1 installers redact the captured $output before Write-Host on non-zero exit. + assert ( + "Write-Host (Redact-InstallOutput $output) -ForegroundColor Red" + in INSTALL_PS1.read_text(encoding = "utf-8") + ) + assert ( + "Write-Host (Redact-InstallOutput $output) -ForegroundColor Red" + in SETUP_PS1.read_text(encoding = "utf-8") + ) + # Python redacts the captured stdout before printing. + assert "_redact_install_output(" in STACK_PY.read_text(encoding = "utf-8") + + +class TestPipNoIndexScrubParity: + """The plain-pip fallback honours PIP_*: PIP_NO_INDEX=1 makes it ignore ALL indexes + (defeating the pinned --index-url) and PIP_INDEX_URL replaces the pin. The two installers + that HAVE a plain-pip fallback (Python + setup.ps1) must scrub both for a pinned install. + install.sh / install.ps1 are uv-only (--default-index), which ignores pip config/env.""" + + def test_python_scrubs_pip_no_index_and_pip_index_url(self): + text = STACK_PY.read_text(encoding = "utf-8") + assert '"PIP_NO_INDEX"' in text + assert '"PIP_INDEX_URL"' in text + + def test_setup_ps1_scrubs_pip_no_index_and_pip_index_url(self): + text = SETUP_PS1.read_text(encoding = "utf-8") + assert "'PIP_NO_INDEX'" in text + assert "'PIP_INDEX_URL'" in text diff --git a/tests/python/test_install_python_stack.py b/tests/python/test_install_python_stack.py index 9015ff8c9d..3a12e53f95 100644 --- a/tests/python/test_install_python_stack.py +++ b/tests/python/test_install_python_stack.py @@ -54,6 +54,24 @@ class TestBuildUvCmdTorchBackend: a.startswith("--torch-backend") for a in cmd ), f"Empty UV_TORCH_BACKEND should not add flag, got: {cmd}" + def test_uv_torch_backend_skipped_for_pinned_index(self): + """A pinned-index command must NOT get --torch-backend: uv's torch backend + redirects torch resolution to its own per-backend index even when + --index-url is given (verified: cu128 pin + backend cpu installs + torch+cpu), defeating the pin.""" + for pin_flag in ("--index-url", "--default-index"): + with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}): + cmd = self._call(("torch", pin_flag, "https://download.pytorch.org/whl/cu128")) + assert not any( + a.startswith("--torch-backend") for a in cmd + ), f"{pin_flag} command must not carry --torch-backend, got: {cmd}" + + def test_uv_torch_backend_kept_for_unpinned(self): + """Non-pinned commands still honour UV_TORCH_BACKEND.""" + with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}): + cmd = self._call(("somepackage",)) + assert "--torch-backend=cpu" in cmd + class TestUvSafePath: """_uv_safe_path hands uv a space-free `-c`/`-r` path (issue #6503).""" @@ -148,3 +166,119 @@ class TestUvSafePathHardening: assert " " not in value assert Path(value).read_text() == "transformers>=4.57.6\n" + + +class TestPinnedIndexClearsUvEnv: + """A pinned torch install (--index-url / --default-index) must neutralise an + inherited UV_INDEX / UV_EXTRA_INDEX_URL so the pinned wheel index wins. + + uv treats the default index (--index-url / --default-index) as LOWEST priority, + so an inherited UV_INDEX / UV_EXTRA_INDEX_URL (a corporate/CPU mirror) would be + searched first and, under uv's default first-index strategy, resolve torch from + the wrong mirror -- after which the marker records a wheel index that was never + used. install.sh (#6898), install.ps1 and setup.ps1 already clear these for + pinned installs; install_python_stack must match (parity across all installers). + """ + + UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL") + + def test_pinned_index_url_strips_uv_index_vars(self): + cmd = [ + "uv", + "pip", + "install", + "--force-reinstall", + "torch", + "torchvision", + "torchaudio", + "--index-url", + "https://download.pytorch.org/whl/cu128", + ] + with mock.patch.dict( + os.environ, + { + "UV_INDEX": "https://mirror.corp/simple", + "UV_EXTRA_INDEX_URL": "https://mirror.corp/extra", + "UV_INDEX_URL": "https://mirror.corp/root", + "UV_DEFAULT_INDEX": "https://mirror.corp/default", + }, + ): + env = ips._install_env_for_cmd(cmd) + assert env is not None, "a --index-url install must run with a scrubbed env" + for var in self.UV_VARS: + assert var not in env, f"{var} must be cleared for a pinned-index install" + + def test_pinned_default_index_strips_uv_index_vars(self): + # --default-index must be gated too (matches install.sh / install.ps1). + cmd = ["uv", "pip", "install", "torch", "--default-index", "https://x/cu126"] + with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}): + env = ips._install_env_for_cmd(cmd) + assert env is not None + assert "UV_INDEX" not in env + + def test_non_pinned_install_keeps_user_mirror(self): + # A plain install (no --index-url) must NOT scrub the env, so a user's mirror + # still applies to base packages. + cmd = ["uv", "pip", "install", "unsloth", "unsloth-zoo"] + with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}): + env = ips._install_env_for_cmd(cmd) + assert env is None, "non-pinned installs must inherit the caller env unchanged" + + def test_scrubbed_env_preserves_other_vars(self): + cmd = ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"] + with mock.patch.dict( + os.environ, + {"UV_INDEX": "https://mirror.corp/simple", "PATH_SENTINEL_XYZ": "keepme"}, + ): + env = ips._install_env_for_cmd(cmd) + assert env is not None + assert env.get("PATH_SENTINEL_XYZ") == "keepme", "only uv index vars are removed" + + def test_pinned_cmd_strips_pip_extra_index_url(self): + """PIP_EXTRA_INDEX_URL is stripped for pinned commands so the pip + fallback cannot satisfy torch from an inherited extra index.""" + with mock.patch.dict(os.environ, {"PIP_EXTRA_INDEX_URL": "https://mirror/simple"}): + env = ips._install_env_for_cmd( + ["pip", "install", "torch", "--index-url", "https://x/cu128"] + ) + assert env is not None and "PIP_EXTRA_INDEX_URL" not in env + + def test_pinned_cmd_strips_uv_torch_backend(self): + """UV_TORCH_BACKEND is stripped for pinned commands so uv cannot read it + from the environment and reroute torch off the pinned index.""" + with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}): + env = ips._install_env_for_cmd( + ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"] + ) + assert env is not None and "UV_TORCH_BACKEND" not in env + + def test_pinned_cmd_disables_uv_config_discovery(self): + """A DISCOVERED uv.toml / pyproject [tool.uv] outranks the CLI pin too + (verified with uv 0.10: [pip] torch-backend = "cpu" and a non-default + [[index]] both resolve torch+cpu against an explicit --index-url / + --default-index cu126 pin). Pinned commands must run with UV_NO_CONFIG=1 + and without an inherited UV_CONFIG_FILE.""" + with mock.patch.dict(os.environ, {"UV_CONFIG_FILE": "/etc/uv/uv.toml"}): + env = ips._install_env_for_cmd( + ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"] + ) + assert env is not None + assert env.get("UV_NO_CONFIG") == "1" + assert "UV_CONFIG_FILE" not in env + + def test_pinned_cmd_disables_pip_config_files(self): + """The pip FALLBACK honours user/site pip config files (pip config set + global.extra-index-url) even with the PIP_* env vars stripped; pip loads + NO configuration files when PIP_CONFIG_FILE is os.devnull. Harmless for + uv, decisive for the fallback.""" + env = ips._install_env_for_cmd( + ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"] + ) + assert env is not None + assert env.get("PIP_CONFIG_FILE") == os.devnull + + def test_non_pinned_cmd_keeps_uv_config_discovery(self): + """Non-pinned installs inherit the caller env unchanged, so a user's uv + configuration still applies to base packages.""" + env = ips._install_env_for_cmd(["uv", "pip", "install", "unsloth"]) + assert env is None diff --git a/tests/run_all.sh b/tests/run_all.sh index d03f4c4d4f..a31103a85b 100755 --- a/tests/run_all.sh +++ b/tests/run_all.sh @@ -15,6 +15,7 @@ sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh" sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh" sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh" sh "$TESTS_DIR/sh/test_torch_flavor.sh" +sh "$TESTS_DIR/sh/test_redact_install_output.sh" sh "$TESTS_DIR/sh/test_install_uv_override_space.sh" echo "" diff --git a/tests/sh/test_get_torch_index_url.sh b/tests/sh/test_get_torch_index_url.sh index 6656142625..23902097ef 100755 --- a/tests/sh/test_get_torch_index_url.sh +++ b/tests/sh/test_get_torch_index_url.sh @@ -23,6 +23,8 @@ _FAKE_SMI_DIR=$(mktemp -d) echo "" sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH" echo "" + sed -n '/^_trim_index_path_slashes()/,/^}/p' "$INSTALL_SH" + echo "" sed -n '/^get_torch_index_url()/,/^}/p' "$INSTALL_SH" } | sed "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \ > "$_FUNC_FILE" @@ -379,6 +381,61 @@ _result=$(run_func "$_dir" " -1 ") assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result" rm -rf "$_dir" +# --- explicit overrides (headless / container / CI; no GPU probing) ---------- +# 39) UNSLOTH_TORCH_INDEX_FAMILY pins the family with no GPU present (not the cpu fallback). +_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none") +assert_eq "family override (no GPU) -> cu128" "https://download.pytorch.org/whl/cu128" "$_result" + +# 40) Family override beats real detection: an nvidia-smi 12.6 host still gets cu128 +# (the Docker-build case -- builder sees the host driver but publishes a cu128 image). +_dir=$(make_mock_smi "12.6") +_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "$_dir") +assert_eq "family override beats detected 12.6 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result" +rm -rf "$_dir" + +# 41) UNSLOTH_TORCH_INDEX_URL is used verbatim and wins over detection. +_dir=$(make_mock_smi "12.6") +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu999" run_func "$_dir") +assert_eq "url override beats detection -> verbatim" "https://mirror.example.com/whl/cu999" "$_result" +rm -rf "$_dir" + +# 42) Family override is appended to UNSLOTH_PYTORCH_MIRROR (mirror still honoured). +_result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none") +assert_eq "mirror + family override -> mirror/cu128" "https://mirror.example.com/whl/cu128" "$_result" + +# 43) Trailing slash in UNSLOTH_TORCH_INDEX_URL is stripped. +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128/" run_func "none") +assert_eq "url override trailing slash stripped" "https://mirror.example.com/whl/cu128" "$_result" + +# 44) URL override takes precedence over family override. +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu130" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none") +assert_eq "url override beats family override -> url" "https://mirror.example.com/whl/cu130" "$_result" + +# 45) An empty override is ignored (falls through to normal detection). +_result=$(UNSLOTH_TORCH_INDEX_FAMILY="" UNSLOTH_TORCH_INDEX_URL="" run_func "none") +assert_eq "empty overrides ignored -> detected cpu" "https://download.pytorch.org/whl/cpu" "$_result" + +# 46) ALL trailing slashes are stripped from a URL override (not just one). +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128///" run_func "none") +assert_eq "url override double slash stripped" "https://mirror.example.com/whl/cu128" "$_result" + +# 47) Leading and trailing slashes stripped from a family override. +_result=$(UNSLOTH_TORCH_INDEX_FAMILY="//cu128//" run_func "none") +assert_eq "family override slashes stripped" "https://download.pytorch.org/whl/cu128" "$_result" + +# 48) A ?query token that ends in "/" is PRESERVED: only PATH slashes are trimmed, so a +# base64 token ending in "/" is not corrupted (path-only trim, not whole-URL rstrip). +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128?token=ab12cd/" run_func "none") +assert_eq "url override preserves query token slash" "https://mirror.example.com/whl/cu128?token=ab12cd/" "$_result" + +# 49) Double PATH slash before a query is collapsed while the query survives intact. +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128//?token=ab12cd/" run_func "none") +assert_eq "url override path slash trimmed, query kept" "https://mirror.example.com/whl/cu128?token=ab12cd/" "$_result" + +# 50) A #fragment ending in "/" is likewise preserved. +_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128#anchor/" run_func "none") +assert_eq "url override preserves fragment slash" "https://mirror.example.com/whl/cu128#anchor/" "$_result" + rm -f "$_FUNC_FILE" rm -rf "$_FAKE_SMI_DIR" rm -rf "$_TOOLS_DIR" diff --git a/tests/sh/test_redact_install_output.sh b/tests/sh/test_redact_install_output.sh new file mode 100755 index 0000000000..0f10122aea --- /dev/null +++ b/tests/sh/test_redact_install_output.sh @@ -0,0 +1,89 @@ +#!/bin/bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Unit tests for install.sh's _redact_install_output helper. uv/pip failure text embeds the +# failing --index-url verbatim, so a captured install log dumped on error can leak a +# user:token@ or ?token= secret. The helper redacts both before printing. Mirrors +# _redact_install_output (install_python_stack.py) / Redact-InstallOutput (install.ps1 / +# setup.ps1). +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +PASS=0 +FAIL=0 + +_FUNC_FILE=$(mktemp) +sed -n '/^_redact_install_output()/,/^}/p' "$INSTALL_SH" > "$_FUNC_FILE" +# shellcheck disable=SC1090 +. "$_FUNC_FILE" +rm -f "$_FUNC_FILE" + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label"; PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1)) + fi +} + +# Redact from a file (the actual call site passes a captured-log tempfile). +redact_str() { + _rs_tmp=$(mktemp) + printf '%s\n' "$1" > "$_rs_tmp" + _rs_out=$(_redact_install_output "$_rs_tmp") + rm -f "$_rs_tmp" + printf '%s' "$_rs_out" +} + +echo "=== _redact_install_output ===" +assert_eq "userinfo user:token@ redacted" \ + "ERROR: failed https://@download.pytorch.org/whl/cu128" \ + "$(redact_str 'ERROR: failed https://alice:s3cr3t@download.pytorch.org/whl/cu128')" + +assert_eq "bare-token@ userinfo redacted" \ + "fetch https://@host/whl/cu128 failed" \ + "$(redact_str 'fetch https://ghp_deadbeef@host/whl/cu128 failed')" + +assert_eq "single ?token= query redacted" \ + "url https://host/whl/cu128?token= unreachable" \ + "$(redact_str 'url https://host/whl/cu128?token=abcd1234 unreachable')" + +assert_eq "multiple query values redacted" \ + "https://host/whl/cu128?token=&channel=" \ + "$(redact_str 'https://host/whl/cu128?token=abcd1234&channel=beta')" + +assert_eq "http (not https) userinfo redacted" \ + "http://@host/simple" \ + "$(redact_str 'http://u:p@host/simple')" + +assert_eq "fragment token redacted" \ + "ERROR: could not fetch https://mirror.local/whl/cu128# (403)" \ + "$(redact_str 'ERROR: could not fetch https://mirror.local/whl/cu128#token=SECRET123 (403)')" + +assert_eq "query and fragment both redacted" \ + "https://host/whl/cu128?token=# done" \ + "$(redact_str 'https://host/whl/cu128?token=abc#sig=xyz done')" + +# Non-secret text is untouched (no false positives on ordinary log lines). +assert_eq "plain line untouched" \ + "Resolved 42 packages in 1.2s" \ + "$(redact_str 'Resolved 42 packages in 1.2s')" +assert_eq "plain url without creds untouched" \ + "downloading https://download.pytorch.org/whl/cu128/torch-2.8.0.whl" \ + "$(redact_str 'downloading https://download.pytorch.org/whl/cu128/torch-2.8.0.whl')" +assert_eq "bare hash comment untouched" \ + "# retrying with --no-cache-dir" \ + "$(redact_str '# retrying with --no-cache-dir')" + +# Regression guard: no secret substring survives. +_leak=$(redact_str 'https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET') +case "$_leak" in + *s3cr3t*|*SUPERSECRET*|*ALSOSECRET*) assert_eq "no secret leak" "clean" "leaked:$_leak" ;; + *) assert_eq "no secret leak" "clean" "clean" ;; +esac + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh index d60dfc9f90..bfafbd161b 100644 --- a/tests/sh/test_torch_constraint.sh +++ b/tests/sh/test_torch_constraint.sh @@ -108,6 +108,25 @@ assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var" _hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true) assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded" +# Companions must be bounded to torch's window everywhere: the <2.11 bound appears +# twice (default assignments + the pinned custom-leaf block), never bare. torchaudio +# 2.11 dropped its exact torch pin, so a bare companion next to a <2.11-capped torch +# resolves a mismatched 2.11 build. +_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"' "$INSTALL_SH" || true) +assert_eq "torchvision bounded (<0.26) at default + custom-leaf" "2" "$_count" +_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"' "$INSTALL_SH" || true) +assert_eq "torchaudio bounded (<2.11) at default + custom-leaf" "2" "$_count" +_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision"$' "$INSTALL_SH" || true) +assert_eq "no bare torchvision companion remains" "0" "$_count" +_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio"$' "$INSTALL_SH" || true) +assert_eq "no bare torchaudio companion remains" "0" "$_count" +# The cu* widen must carry the companions with it (torch <2.12 with torchaudio <2.11 +# would cap a mismatched pair the other way). +assert_eq "cu widen pairs torchaudio (<2.12)" "1" "$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0"' "$INSTALL_SH" || true)" +_gated=$(grep -c '_expected_torch_flavor_tag "$TORCH_INDEX_URL"' "$INSTALL_SH" || true) +_has_gate=$([ "$_gated" -ge 1 ] && echo "yes" || echo "no") +assert_eq "custom-companion bound gated on empty flavor tag" "yes" "$_has_gate" + # A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x land torch # 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC). _cuda_widen=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' "$INSTALL_SH" || true) @@ -285,6 +304,61 @@ bash -c " _uv_got2=$(cat "$_UV_LOG2" 2>/dev/null || echo "") assert_contains "mock uv arm64+py312 receives torch>=2.4" "$_uv_got2" "torch>=2.4,<2.11.0" +# ====================================================================== +# ROCm 2.11 floor: leaf is lowercased before the gfx*/rocm* allowlist match +# ====================================================================== +echo "" +echo "=== ROCm 2.11 floor case (leaf normalization) ===" + +# Structural: install.sh lowercases _torch_index_leaf before the floor case, so the +# canonical gfx120X-all (capital X) matches gfx120x-all. +_has_lc=$(grep -c '_torch_index_leaf=$(printf .* | tr .\[:upper:\]. .\[:lower:\].)' "$INSTALL_SH" || true) +_has_lc_ok=$([ "$_has_lc" -ge 1 ] && echo "yes" || echo "no") +assert_eq "install.sh lowercases _torch_index_leaf" "yes" "$_has_lc_ok" + +# Runtime: replicate install.sh's normalization + floor case and assert both gfx120X-all +# and gfx120x-all get the floor, while non-2.11 leaves keep the default. +run_floor_case() { + _url="$1" + bash -c ' + TORCH_CONSTRAINT="torch>=2.4,<2.11.0" + TORCHVISION_CONSTRAINT="torchvision" + TORCHAUDIO_CONSTRAINT="torchaudio" + _torch_index_leaf="${1%/}" + _torch_index_leaf="${_torch_index_leaf##*/}" + _torch_index_leaf=$(printf "%s" "$_torch_index_leaf" | tr "[:upper:]" "[:lower:]") + case "$_torch_index_leaf" in + rocm7.2|gfx120x-all|gfx1151|gfx1150) + TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" + TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" + TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" + ;; + esac + echo "$TORCH_CONSTRAINT" + ' _ "$_url" +} + +assert_eq "gfx120X-all (capital) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all')" +assert_eq "gfx120X-all trailing slash -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all/')" +assert_eq "gfx120x-all (lowercase) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120x-all')" +assert_eq "gfx1151 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1151')" +assert_eq "gfx1150 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1150')" +assert_eq "rocm7.2 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \ + "$(run_floor_case 'https://download.pytorch.org/whl/rocm7.2')" +assert_eq "gfx110X-all -> default (no floor)" "torch>=2.4,<2.11.0" \ + "$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx110X-all')" +assert_eq "rocm6.4 -> default (no floor)" "torch>=2.4,<2.11.0" \ + "$(run_floor_case 'https://download.pytorch.org/whl/rocm6.4')" +assert_eq "cu128 -> default (no floor)" "torch>=2.4,<2.11.0" \ + "$(run_floor_case 'https://download.pytorch.org/whl/cu128')" +assert_eq "cpu -> default (no floor)" "torch>=2.4,<2.11.0" \ + "$(run_floor_case 'https://download.pytorch.org/whl/cpu')" + # ====================================================================== # Summary # ====================================================================== diff --git a/tests/sh/test_torch_flavor.sh b/tests/sh/test_torch_flavor.sh index ead2c4164f..55da2c0f07 100755 --- a/tests/sh/test_torch_flavor.sh +++ b/tests/sh/test_torch_flavor.sh @@ -11,14 +11,21 @@ INSTALL_SH="$SCRIPT_DIR/../../install.sh" PASS=0 FAIL=0 -# Extract the three helper functions from install.sh and source them. +# Extract the helper functions from install.sh and source them +# (_torch_index_url_leaf is the shared leaf extractor the classifiers call). _FUNC_FILE=$(mktemp) { sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH" echo "" + sed -n '/^_torch_index_url_leaf()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_is_pip_rocm_family_leaf()/,/^}/p' "$INSTALL_SH" + echo "" sed -n '/^_expected_torch_flavor_tag()/,/^}/p' "$INSTALL_SH" echo "" sed -n '/^_torch_index_repairable()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_tauri_torch_index_family()/,/^}/p' "$INSTALL_SH" } > "$_FUNC_FILE" # shellcheck disable=SC1090 . "$_FUNC_FILE" @@ -56,6 +63,26 @@ assert_eq "amd gfx index" "rocm" "$(_expected_torch_flavor_tag 'https://re assert_eq "mirror cu130 leaf" "cu130" "$(_expected_torch_flavor_tag 'https://my.mirror/pytorch/whl/cu130')" assert_eq "unrecognized leaf" "" "$(_expected_torch_flavor_tag 'https://my.mirror/whl/simple')" assert_eq "empty url" "" "$(_expected_torch_flavor_tag '')" +# Query/fragment dropped before classification: .../cu128?token=x classifies as cu128, +# not an opaque leaf that reinstalls every run. +assert_eq "query-bearing cu128" "cu128" "$(_expected_torch_flavor_tag 'https://m/whl/cu128?token=x')" +assert_eq "fragment-bearing cpu" "cpu" "$(_expected_torch_flavor_tag 'https://m/whl/cpu#frag')" +# A cu-suffixed CUSTOM leaf (cu128-private, cu128x) is NOT the cu128 family (exact +# cu+digits only). Mirrors Python re.fullmatch(cu[0-9]+) / PowerShell. +assert_eq "cu-suffix custom leaf" "" "$(_expected_torch_flavor_tag 'https://m/whl/cu128-private')" +assert_eq "cu-alnum custom leaf" "" "$(_expected_torch_flavor_tag 'https://m/whl/cu128x')" +assert_eq "bare cu digits stays" "cu126" "$(_expected_torch_flavor_tag 'https://m/whl/cu126')" +# A custom leaf merely STARTING with rocm (rocm-current, rocm-rel-7.2.1) is NOT a pip +# rocm family -> "" (custom); real families (rocm7.2) and gfx indexes stay "rocm". +assert_eq "custom rocm-current" "" "$(_expected_torch_flavor_tag 'https://mirror/whl/rocm-current')" +assert_eq "radeon rocm-rel leaf" "" "$(_expected_torch_flavor_tag 'https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1')" +assert_eq "real rocm7.2 stays" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7.2')" +# A rocm-SUFFIX private mirror (rocm7.2-private, rocm7-current) is a custom pin -> +# "" (custom); match the family exactly, not the prefix. +assert_eq "suffixed rocm7.2-private" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7.2-private')" +assert_eq "suffixed rocm7-current" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7-current')" +assert_eq "two-dot rocm7.2.1" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7.2.1')" +assert_eq "bare rocm7 stays" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7')" echo "=== _torch_index_repairable ===" assert_eq "cu130 repairable" "yes" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cu130')" @@ -64,6 +91,66 @@ assert_eq "gfx repairable" "yes" "$(_torch_index_repairable 'https://repo. assert_eq "gfx1151 repairable" "yes" "$(_torch_index_repairable 'https://repo.amd.com/rocm/whl/gfx1151/')" assert_eq "cpu NOT repairable" "no" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cpu')" assert_eq "unknown NOT repair" "no" "$(_torch_index_repairable 'https://my.mirror/whl/simple')" +# A suffixed rocm leaf is a verbatim pin, not a --default-index repairable family. +assert_eq "rocm-private NOT repair" "no" "$(_torch_index_repairable 'https://co.internal/whl/rocm7.2-private')" + +echo "=== _is_pip_rocm_family_leaf ===" +assert_family() { + _label="$1"; _expected="$2"; _leaf="$3" + if _is_pip_rocm_family_leaf "$_leaf"; then _actual="yes"; else _actual="no"; fi + assert_eq "$_label" "$_expected" "$_actual" +} +assert_family "rocm7.2 family" "yes" "rocm7.2" +assert_family "rocm6.4 family" "yes" "rocm6.4" +assert_family "bare rocm7 family" "yes" "rocm7" +assert_family "gfx120x-all family" "yes" "gfx120x-all" +assert_family "gfx1151 family" "yes" "gfx1151" +assert_family "rocm7.2-private custom" "no" "rocm7.2-private" +assert_family "rocm7-current custom" "no" "rocm7-current" +assert_family "rocm-current custom" "no" "rocm-current" +assert_family "rocm-rel-7.2.1 custom" "no" "rocm-rel-7.2.1" +assert_family "rocm7.2.1 custom" "no" "rocm7.2.1" +# A trailing dot (rocm7.) or leading/double dot is NOT a family: both major and minor must +# be non-empty all-digits, matching Python re.fullmatch(rocm\d+(?:\.\d+)?). Bash previously +# accepted rocm7. via a bare %/-style trim while Python rejected it (validator asymmetry). +assert_family "rocm7. trailing-dot custom" "no" "rocm7." +assert_family "rocm.7 leading-dot custom" "no" "rocm.7" +assert_family "rocm7..2 double-dot custom" "no" "rocm7..2" +assert_family "cpu not rocm" "no" "cpu" +assert_family "cu128 not rocm" "no" "cu128" +assert_family "simple not rocm" "no" "simple" + +echo "=== _torch_index_url_leaf (ALL trailing slashes stripped -> non-empty leaf) ===" +# A double (or triple) trailing slash must yield the real leaf, not an empty string that +# fails every classifier arm. Python .rstrip("/") drops them all; bash must match (a bare +# %/ left .../cu128// classifying as ""). +assert_eq "double slash cu128 leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128//')" +assert_eq "triple slash rocm7.2 leaf" "rocm7.2" "$(_torch_index_url_leaf 'https://m/whl/rocm7.2///')" +assert_eq "double slash + token leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128//?token=x')" +assert_eq "single slash cu128 leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128/')" +# The classifier that consumes the leaf must therefore still tag a double-slash index. +assert_eq "double-slash cu128 tag" "cu128" "$(_expected_torch_flavor_tag 'https://m/whl/cu128//')" +assert_eq "double-slash rocm7.2 tag" "rocm" "$(_expected_torch_flavor_tag 'https://m/whl/rocm7.2//')" + +echo "=== _tauri_torch_index_family (credential redaction) ===" +# A token/fragment must be stripped BEFORE classification so it never reaches the +# [TAURI:DIAG] line (the family is the last path segment, which else carries the query). +SKIP_TORCH=false +assert_eq "token stripped from rocm" "rocm7.2" "$(_tauri_torch_index_family 'https://mirror/whl/rocm7.2?token=SECRET')" +assert_eq "token-bearing cu classifies" "cu128" "$(_tauri_torch_index_family 'https://m/whl/cu128?token=x')" +assert_eq "fragment stripped cpu" "cpu" "$(_tauri_torch_index_family 'https://m/whl/cpu#frag')" +assert_eq "plain rocm7.2 unchanged" "rocm7.2" "$(_tauri_torch_index_family 'https://download.pytorch.org/whl/rocm7.2')" +# A trailing slash must be stripped too, or the */cu128 and */cpu arms miss .../cu128/ +# and it falls through to "auto". +assert_eq "trailing slash cu128" "cu128" "$(_tauri_torch_index_family 'https://download.pytorch.org/whl/cu128/')" +assert_eq "slash + token cu128" "cu128" "$(_tauri_torch_index_family 'https://m/whl/cu128/?token=x')" +assert_eq "trailing slash cpu" "cpu" "$(_tauri_torch_index_family 'https://m/whl/cpu/')" +# Regression guard: no secret token substring may survive in any classification. +_leak=$(_tauri_torch_index_family 'https://mirror/whl/rocm7.2?token=SECRET') +case "$_leak" in + *SECRET*|*token*) assert_eq "no token leak in family" "clean" "leaked:$_leak" ;; + *) assert_eq "no token leak in family" "clean" "clean" ;; +esac echo "" echo "Results: $PASS passed, $FAIL failed" diff --git a/tests/studio/install/test_cuda_repair.py b/tests/studio/install/test_cuda_repair.py index cea4383268..c6d2b95316 100644 --- a/tests/studio/install/test_cuda_repair.py +++ b/tests/studio/install/test_cuda_repair.py @@ -64,15 +64,23 @@ def _run_cuda_repair( rocm_marker = False, smi_path = "/usr/bin/nvidia-smi", cvd = None, + index_family = None, + index_url = None, ): """Invoke _ensure_cuda_torch under a fully mocked host; return the pip mock. - cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it.""" + cvd controls CUDA_VISIBLE_DEVICES: None removes it from the env, any string sets it. + index_family sets UNSLOTH_TORCH_INDEX_FAMILY (the explicit wheel-index pin). + index_url sets UNSLOTH_TORCH_INDEX_URL (the full-URL pin form).""" env = {} if rocm_marker: env["UNSLOTH_ROCM_TORCH_INSTALLED"] = "1" if cvd is not None: env["CUDA_VISIBLE_DEVICES"] = cvd + if index_family is not None: + env["UNSLOTH_TORCH_INDEX_FAMILY"] = index_family + if index_url is not None: + env["UNSLOTH_TORCH_INDEX_URL"] = index_url def _which(name, *a, **k): if name == "nvidia-smi": @@ -99,6 +107,10 @@ def _run_cuda_repair( stack_mod.os.environ.pop("UNSLOTH_ROCM_TORCH_INSTALLED", None) if cvd is None: stack_mod.os.environ.pop("CUDA_VISIBLE_DEVICES", None) + if index_family is None: + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + if index_url is None: + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None) _ensure_cuda_torch() return mock_pip @@ -123,11 +135,74 @@ class TestCudaRepairFires: assert mock_pip.call_args.kwargs["constrain"] is False def test_rocm_in_version_string_triggers_repair(self): - # AMD SDK / Radeon wheels may encode rocm in __version__ without - # torch.version.hip; the probe prints "hip" for both. + # AMD SDK / Radeon wheels may encode rocm in __version__ without torch.version.hip; + # the probe prints "hip" for both. mock_pip = _run_cuda_repair(torch_state = "hip") assert mock_pip.call_count == 1 + def test_no_gpu_but_explicit_cuda_pin_repairs(self): + # Headless / CI cross-install: an explicit cu* pin commits to CUDA wheels with no + # NVIDIA GPU visible, so a ROCm-poisoned venv is still repaired to the pinned family. + mock_pip = _run_cuda_repair( + nvidia = False, + backend = "cuda", + index_family = "cu128", + torch_state = "hip", + ) + assert mock_pip.call_count == 1 + assert "cu128" in _index_url(mock_pip) + + def test_cvd_hidden_but_explicit_cuda_pin_repairs(self): + # CVD=-1/"" hides the GPU, but an explicit cu* pin skips ALL host-GPU probing, so the + # CVD hide gate must not suppress the repair (GPU-less CI: CVD=-1, FAMILY=cu128). + for _cvd in ("-1", ""): + mock_pip = _run_cuda_repair( + nvidia = False, + backend = "cuda", + cvd = _cvd, + index_family = "cu128", + torch_state = "hip", + ) + assert mock_pip.call_count == 1 + assert "cu128" in _index_url(mock_pip) + + def test_tagged_cuda_mismatch_repairs(self): + # A healthy CUDA torch whose +cuXXX differs from the pin is repaired. + mock_pip = _run_cuda_repair( + index_family = "cu128", + torch_state = "cuda|cu126", + cuda_version = "12.8", + ) + assert mock_pip.call_count == 1 + assert "cu128" in _index_url(mock_pip) + + def test_untagged_cuda_build_under_pin_repairs(self): + # An untagged CUDA build (no +cuXXX tag -> empty installed cu) can't be confirmed + # to match the pin, so the pin is enforced with a reinstall. + mock_pip = _run_cuda_repair( + index_family = "cu128", + torch_state = "cuda", # marker cuda, empty installed cu + cuda_version = "12.8", + ) + assert mock_pip.call_count == 1 + assert "cu128" in _index_url(mock_pip) + + def test_broken_probe_with_cuda_pin_repairs(self): + # torch present but unimportable under a CUDA pin: the base update won't repair a + # broken already-installed torch, so reinstall from the pin instead of stranding it. + mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1, index_family = "cu128") + assert mock_pip.call_count == 1 + assert "cu128" in _index_url(mock_pip) + + def test_broken_probe_with_cuda_url_pin_repairs(self): + mock_pip = _run_cuda_repair( + torch_state = "cpu", + torch_rc = 1, + index_url = "https://mirror.local/cu128", + ) + assert mock_pip.call_count == 1 + assert "https://mirror.local/cu128" in _index_url(mock_pip) + # No-op cases. @@ -157,8 +232,9 @@ class TestCudaRepairSkips: mock_pip = _run_cuda_repair(nvidia = False, torch_state = "hip") mock_pip.assert_not_called() - def test_torch_missing_skips(self): - # Non-zero probe exit = torch missing / un-importable. + def test_torch_missing_no_pin_skips(self): + # Non-zero probe exit = torch missing/un-importable. With NO CUDA pin the base + # install owns it, so leave it alone (a pinned build reinstalls). mock_pip = _run_cuda_repair(torch_state = "hip", torch_rc = 1) mock_pip.assert_not_called() @@ -191,6 +267,109 @@ class TestCudaRepairSkips: mock_pip = _run_cuda_repair(cvd = "0", torch_state = "hip") assert mock_pip.call_count == 1 + def test_matching_tagged_cuda_pin_no_repair(self): + # Healthy CUDA torch whose +cuXXX already matches the pin: no reinstall. + mock_pip = _run_cuda_repair( + index_family = "cu128", + torch_state = "cuda|cu128", + cuda_version = "12.8", + ) + mock_pip.assert_not_called() + + def test_custom_mirror_leaf_not_treated_as_cuda_pin(self): + # A mirror leaf starting with "cu" but not cuXXX (.../custom, .../current) must + # NOT be treated as a CUDA pin, so it can't bypass the NVIDIA gate. + for _leaf in ("custom", "current"): + mock_pip = _run_cuda_repair( + nvidia = False, + backend = "cuda", + index_url = f"https://mymirror.example/{_leaf}", + torch_state = "hip", + ) + mock_pip.assert_not_called() + + def test_explicit_cuda_family_leaf_helper(self): + # _explicit_cuda_torch_index_url matches cuXXX narrowly, not any cu* leaf. + import contextlib + + def _with(url): + with patch.dict(stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + return stack_mod._explicit_cuda_torch_index_url() + + assert _with("https://download.pytorch.org/whl/cu128") is not None + assert _with("https://download.pytorch.org/whl/cu126") is not None + assert _with("https://mymirror.example/custom") is None + assert _with("https://mymirror.example/current") is None + assert _with("https://download.pytorch.org/whl/cpu") is None + with contextlib.suppress(Exception): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None) + + +class TestTorchBackendDerivationFromPin: + """The module-level _TORCH_BACKEND derivation (standalone `studio update` + with no install.sh-set UNSLOTH_TORCH_BACKEND) must classify the pinned index + leaf via _is_cuda_family_leaf (^cu[0-9]), NOT a bare startswith("cu"). A + full-override URL ending in /current or /custom must fall through to backend + "" (probe the GPU) so _ensure_rocm_torch() still repairs a wrong/CPU torch on + AMD hosts, instead of being wrongly branded "cuda" and returning early.""" + + @staticmethod + def _derive(env): + # Re-run the module's import-time derivation, using its own _is_cuda_family_leaf + # so this stays in lockstep. + idx_override = ( + env.get("UNSLOTH_TORCH_INDEX_URL", "").strip() + or env.get("UNSLOTH_TORCH_INDEX_FAMILY", "").strip() + ) + backend = env.get("UNSLOTH_TORCH_BACKEND", "").lower() + if not backend: + leaf = idx_override.rstrip("/").rsplit("/", 1)[-1].lower() + if leaf.startswith(("rocm", "gfx")): + backend = "rocm" + elif leaf == "cpu": + backend = "cpu" + elif stack_mod._is_cuda_family_leaf(leaf): + backend = "cuda" + return backend + + def test_cu128_pin_is_cuda(self): + assert ( + self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://download.pytorch.org/whl/cu128"}) + == "cuda" + ) + + def test_cu128_family_is_cuda(self): + assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "cu128"}) == "cuda" + + def test_current_leaf_not_cuda(self): + # ^cu[0-9] rejects /current -> backend stays "" (probe GPU), so an AMD host still + # repairs a CPU/wrong torch instead of short-circuiting. + assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/current"}) == "" + + def test_custom_leaf_not_cuda(self): + assert self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://mymirror.example/custom"}) == "" + + def test_rocm_and_gfx_pins_are_rocm(self): + assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"}) == "rocm" + assert ( + self._derive({"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx120X-all"}) + == "rocm" + ) + + def test_cpu_pin_is_cpu(self): + assert self._derive({"UNSLOTH_TORCH_INDEX_FAMILY": "cpu"}) == "cpu" + + def test_source_uses_helper_not_bare_startswith(self): + # Guard against a regression back to elif _idx_leaf.startswith("cu"). + src = _STACK_PATH.read_text(encoding = "utf-8") + assert ( + "elif _is_cuda_family_leaf(_idx_leaf):" in src + ), "_TORCH_BACKEND derivation must classify CUDA via _is_cuda_family_leaf" + assert ( + 'elif _idx_leaf.startswith("cu"):' not in src + ), "_TORCH_BACKEND derivation must not use a bare startswith('cu')" + # CUDA index ladder. diff --git a/tests/studio/install/test_gpu_detection_followups.py b/tests/studio/install/test_gpu_detection_followups.py index f5ad9566d7..d2fd7ae8db 100644 --- a/tests/studio/install/test_gpu_detection_followups.py +++ b/tests/studio/install/test_gpu_detection_followups.py @@ -284,7 +284,9 @@ class TestBackendExportLeafClassification: def test_export_block_uses_leaf(self, install_src): anchor = install_src.find("_torch_index_leaf=") assert anchor >= 0, "backend export must classify on the final path segment" - window = install_src[anchor : anchor + 500] + # Window spans the leaf-normalization prelude (query/frag drop + all-slash trim loop) + # through the export case arms. + window = install_src[anchor : anchor + 900] assert 'export UNSLOTH_TORCH_BACKEND="rocm"' in window assert 'export UNSLOTH_TORCH_BACKEND="cpu"' in window assert 'export UNSLOTH_TORCH_BACKEND="cuda"' in window @@ -459,3 +461,130 @@ class TestHiddenCvdNotUsable: cvd, ) assert out == expected + + +class TestRedactInstallOutput: + """_redact_install_output scrubs index-URL credentials from a captured install log + before it is printed on failure (uv/pip embeds the failing --index-url verbatim).""" + + def test_userinfo_redacted(self): + out = stack_mod._redact_install_output( + "ERROR: failed https://alice:s3cr3t@download.pytorch.org/whl/cu128" + ) + assert out == "ERROR: failed https://@download.pytorch.org/whl/cu128" + + def test_bytes_input_decoded_and_redacted(self): + out = stack_mod._redact_install_output(b"fetch https://ghp_deadbeef@host/whl/cu128 failed") + assert out == "fetch https://@host/whl/cu128 failed" + + def test_query_values_redacted(self): + out = stack_mod._redact_install_output( + "url https://host/whl/cu128?token=abcd1234&channel=beta unreachable" + ) + assert out == "url https://host/whl/cu128?token=&channel= unreachable" + + def test_fragment_redacted(self): + out = stack_mod._redact_install_output( + "ERROR: could not fetch https://mirror.local/whl/cu128#token=SECRET123 (403)" + ) + assert out == "ERROR: could not fetch https://mirror.local/whl/cu128# (403)" + + def test_query_and_fragment_both_redacted(self): + out = stack_mod._redact_install_output("https://host/whl/cu128?token=abc#sig=xyz done") + assert out == "https://host/whl/cu128?token=# done" + + def test_bare_hash_comment_untouched(self): + # The fragment redaction is URL-anchored: a shell comment in tool output survives. + assert ( + stack_mod._redact_install_output("# retrying with --no-cache-dir") + == "# retrying with --no-cache-dir" + ) + + def test_plain_line_untouched(self): + assert ( + stack_mod._redact_install_output("Resolved 42 packages in 1.2s") + == "Resolved 42 packages in 1.2s" + ) + + def test_no_secret_substring_survives(self): + out = stack_mod._redact_install_output( + "https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET" + ) + assert "s3cr3t" not in out and "SUPERSECRET" not in out and "ALSOSECRET" not in out + + +class TestTrimIndexPathSlashes: + """_trim_index_path_slashes strips trailing PATH slashes only; a ?query/#fragment token + ending in "/" must survive (a whole-URL rstrip would corrupt a base64 token).""" + + def test_double_path_slash_collapsed(self): + assert stack_mod._trim_index_path_slashes("https://h/whl/cu128//") == "https://h/whl/cu128" + + def test_query_token_slash_preserved(self): + assert ( + stack_mod._trim_index_path_slashes("https://h/whl/cu128?token=ab12cd/") + == "https://h/whl/cu128?token=ab12cd/" + ) + + def test_path_slash_trimmed_query_kept(self): + assert ( + stack_mod._trim_index_path_slashes("https://h/whl/cu128//?token=ab12cd/") + == "https://h/whl/cu128?token=ab12cd/" + ) + + def test_fragment_slash_preserved(self): + assert ( + stack_mod._trim_index_path_slashes("https://h/whl/cu128#anchor/") + == "https://h/whl/cu128#anchor/" + ) + + +class TestRocmFamilyLeafParity: + """_is_pip_rocm_family_leaf must match re.fullmatch(rocm\\d+(?:\\.\\d+)?): a trailing dot + (rocm7.) is a CUSTOM pin, not a family (the historical bash/py validator asymmetry).""" + + @pytest.mark.parametrize( + "leaf, expected", + [ + ("rocm7", True), + ("rocm7.2", True), + ("gfx1151", True), + ("rocm7.", False), + ("rocm.7", False), + ("rocm7..2", False), + ("rocm7.2.1", False), + ("rocm7.2-private", False), + ("cpu", False), + ("cu128", False), + ], + ) + def test_family_classification(self, leaf, expected): + assert stack_mod._is_pip_rocm_family_leaf(leaf) is expected + + +class TestTorchIndexLeafAllSlashes: + """_torch_index_leaf drops query/fragment then strips ALL trailing slashes, so a + double-slash index still yields the real leaf (not an empty string).""" + + @pytest.mark.parametrize( + "url, expected", + [ + ("https://m/whl/cu128//", "cu128"), + ("https://m/whl/rocm7.2///", "rocm7.2"), + ("https://m/whl/cu128//?token=x", "cu128"), + ("https://m/whl/cu128/", "cu128"), + ], + ) + def test_leaf_never_empty_on_double_slash(self, url, expected): + assert stack_mod._torch_index_leaf(url) == expected + + +class TestUvIndexEnvVarsScrub: + """The pinned-install env scrub must drop PIP_NO_INDEX (which makes the pip fallback + ignore ALL indexes, defeating the pin) and PIP_INDEX_URL (replaces the pinned index).""" + + def test_pip_no_index_scrubbed(self): + assert "PIP_NO_INDEX" in stack_mod._UV_INDEX_ENV_VARS + + def test_pip_index_url_scrubbed(self): + assert "PIP_INDEX_URL" in stack_mod._UV_INDEX_ENV_VARS diff --git a/tests/studio/install/test_pr5940_followups.py b/tests/studio/install/test_pr5940_followups.py index d6dc8b2f8e..dd2a7ec487 100644 --- a/tests/studio/install/test_pr5940_followups.py +++ b/tests/studio/install/test_pr5940_followups.py @@ -784,7 +784,7 @@ def test_install_python_stack_windows_rocm_repair_pins_and_is_nonfatal(): assert re.search( r'"' + gfx + r'":\s*_ROCM_TORCH_PKG_SPECS\["rocm7\.2"\]', text ), f"{gfx} must pin to the rocm7.2 trio like install.ps1/setup.ps1" - i = text.find('f"ROCm torch (Windows, {gfx_arch})"') + i = text.find("f\"ROCm torch (Windows, {gfx_arch or 'pinned'})\"") assert i != -1, "Windows ROCm repair pip call not found" # The nearest preceding call must be the nonfatal pip_install_try, not pip_install. j = text.rfind("pip_install_try(", 0, i) diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index e7ac0ec82d..b94578369d 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -569,19 +569,27 @@ class TestEnsureRocmTorch: _ensure_rocm_torch() mock_pip.assert_not_called() + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) - def test_torch_already_has_cuda_skips(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): - """If torch already has CUDA, should skip ROCm reinstall.""" + def test_cuda_torch_on_amd_host_reinstalls( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """A CUDA-only torch build is unusable on an AMD-only host, so it must be + reinstalled to ROCm (has_hip_torch is driven by the empty HIP marker, not + by treating the CUDA version string as a HIP marker).""" mock_probe = MagicMock() mock_probe.returncode = 0 - mock_probe.stdout = b"12.6\n" # CUDA version + # Single-line probe: empty HIP marker before "|" for a CUDA build. + mock_probe.stdout = b"|2.10.0+cu126\n" with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() - mock_pip.assert_not_called() + assert mock_pip.call_count == 1 + assert "rocm7.1" in str(mock_pip.call_args_list[0]) @patch.object(stack_mod, "pip_install") @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) @@ -591,12 +599,31 @@ class TestEnsureRocmTorch: """If torch already has HIP, should skip ROCm reinstall.""" mock_probe = MagicMock() mock_probe.returncode = 0 - mock_probe.stdout = b"7.1.12345\n" # HIP version + mock_probe.stdout = b"7.1.12345|2.10.0+rocm7.1\n" # HIP marker + version with patch("os.path.isdir", return_value = True): with patch("subprocess.run", return_value = mock_probe): _ensure_rocm_torch() mock_pip.assert_not_called() + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 1)) + def test_cpu_torch_probe_line_not_read_as_hip(self, mock_ver, mock_gpu, mock_nvidia, mock_pip): + """A CPU build's probe line ("|2.10.0+cpu") must not read as HIP: the version + after the "|" separator is data, not a HIP marker, so has_hip_torch stays False + and the reinstall fires.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"|2.10.0+cpu\n" + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + with patch.object(stack_mod, "pip_install_try", return_value = True): + _ensure_rocm_torch() + assert mock_pip.call_count == 1 + assert "rocm7.1" in str(mock_pip.call_args_list[0]) + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @@ -680,6 +707,295 @@ class TestEnsureRocmTorch: torch_call = mock_pip.call_args_list[0] assert "rocm7.2" in str(torch_call) + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install_try", return_value = True) + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4)) + def test_explicit_gfx_index_honored_and_skips_strix_reroute( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """An explicit gfx wheel-index pin is authoritative: install from it verbatim + with torch 2.11, and never re-probe gfx codes to second-guess it (host ROCm 6.4 + would otherwise pick the rocm6.4 wheel / trigger the Strix re-route).""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"\n" # cpu torch -> reinstall + env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"} + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + # Would raise if the Strix block ran (it is skipped on an explicit pin). + with patch.object( + stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError + ): + _ensure_rocm_torch() + assert mock_pip.call_count == 1 + torch_call = str(mock_pip.call_args_list[0]) + assert "gfx1151" in torch_call + assert "torch>=2.11.0,<2.12.0" in torch_call + + def test_rocm_pin_family_mismatch_helper(self): + """_rocm_pin_family_mismatch: exact rocm compare, else the 2.11 line.""" + f = stack_mod._rocm_pin_family_mismatch + base = "https://download.pytorch.org/whl" + amd = "https://repo.amd.com/rocm/whl" + # Exact rocm version comparison. + assert f(f"{base}/rocm7.2", "2.11.0+rocm7.2") is False + assert f(f"{base}/rocm7.2", "2.10.0+rocm6.4") is True + assert f(f"{base}/rocm6.4", "2.10.0+rocm6.4") is False + # rocm7.2 is KNOWN-2.11. A +rocm7.2 wheel whose RELEASE drifted off 2.11 shares the + # tag but violates the spec -> mismatch (a plain version compare would accept it). + assert f(f"{base}/rocm7.2", "2.12.0+rocm7.2") is True + assert f(f"{base}/rocm7.2", "2.13.0+rocm7.2") is True + assert f(f"{base}/rocm7.2", "2.11.5+rocm7.2") is False # patch on 2.11 is in-spec + # An UNKNOWN newer rocm (not on the 2.11 allowlist) is not floored to 2.11, so a + # matching rocm version at any release line is NOT a mismatch on this branch. + assert f(f"{base}/rocm8.0", "2.12.0+rocm8.0") is False + # gfx pin (2.11 line) vs installed release line. + assert f(f"{amd}/gfx1151", "2.10.0+rocm6.4") is True + assert f(f"{amd}/gfx1151", "2.11.0+rocm7.13.0") is False + # rocm7.2 pin vs an untagged (no +rocm) wheel: a CPU/CUDA build never + # satisfies a ROCm pin, regardless of its release line -> always a mismatch. + assert f(f"{base}/rocm7.2", "2.10.0") is True + assert f(f"{base}/rocm7.2", "2.11.0") is True + assert f(f"{base}/rocm6.4", "2.10.0") is True + # A 2.11-allowlist gfx pin over a GENERIC (two-part +rocm7.2) 2.11 wheel mismatches: + # the user wants AMD's per-arch (three-part) wheel, not the generic one. + assert f(f"{amd}/gfx1151", "2.11.0+rocm7.2") is True + assert f(f"{amd}/gfx120X-all", "2.11.0+rocm7.2") is True + # ...but an already-installed per-arch (three-part) wheel is NOT re-flagged + # (no reinstall loop once the correct gfx wheel is present). + assert f(f"{amd}/gfx120X-all", "2.11.0+rocm7.13.0") is False + assert f(f"{amd}/gfx1150", "2.11.0+rocm7.13.0") is False + # A NON-2.11 gfx pin (gfx110X-all/gfx90a/gfx908) tracks the default <2.11 spec: a + # correct 2.10+rocm wheel is NOT a mismatch, a 2.11 build is. + assert f(f"{amd}/gfx110X-all", "2.10.0+rocm6.4") is False + assert f(f"{amd}/gfx90a", "2.10.0+rocm6.3") is False + assert f(f"{amd}/gfx908", "2.10.0+rocm7.0") is False + assert f(f"{amd}/gfx110X-all", "2.11.0+rocm7.2") is True + # A non-2.11 gfx pin over an untagged (no +rocm) wheel is a mismatch even + # when torch is already <2.11: a CPU/CUDA build never satisfies the ROCm pin. + assert f(f"{amd}/gfx110X-all", "2.10.0") is True + assert f(f"{amd}/gfx90a", "2.10.0") is True + # A major-only rocm pin (rocm7) compares on the major alone: rocm6.x mismatches, + # any rocm7.x satisfies it, an untagged wheel never does, a bare +rocm is lenient. + assert f(f"{base}/rocm7", "2.10.0+rocm6.4") is True + assert f(f"{base}/rocm7", "2.11.0+rocm7.2") is False + assert f(f"{base}/rocm7", "2.11.0+rocm7.13.0") is False + assert f(f"{base}/rocm7", "2.10.0") is True + assert f(f"{base}/rocm7", "2.10.0+rocm") is False + + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install_try", return_value = True) + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + def test_rocm_pin_mismatch_over_installed_rocm_reinstalls( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """A rocm7.2 pin over an already-installed OLDER +rocm6.4 build must reinstall, + even though has_hip_torch is True (the ROCm analogue of the CUDA cuXXX mismatch).""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + # HIP marker present (has_hip_torch=True) + installed +rocm6.4 wheel. + mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n" + env = {"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"} + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None) + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + torch_call = str(mock_pip.call_args_list[0]) + assert "rocm7.2" in torch_call + assert "torch>=2.11.0,<2.12.0" in torch_call + + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install_try", return_value = True) + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4)) + def test_gfx_pin_over_installed_pre211_rocm_reinstalls( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """A gfx* pin (2.11 line) over an installed pre-2.11 +rocm6.4 build reinstalls.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n" + env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"} + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + with patch.object( + stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError + ): + _ensure_rocm_torch() + torch_call = str(mock_pip.call_args_list[0]) + assert "gfx1151" in torch_call + assert "torch>=2.11.0,<2.12.0" in torch_call + + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install_try", return_value = True) + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + def test_rocm_pin_matches_installed_no_torch_reinstall( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """A rocm7.2 pin over an already-matching +rocm7.2 build must NOT reinstall torch + (no false reinstall of a correct ROCm venv).""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"7.2.12345|2.11.0+rocm7.2\n" + env = {"UNSLOTH_TORCH_INDEX_FAMILY": "rocm7.2"} + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_URL", None) + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + # No torch reinstall: any pip_install call must not target a torch index. + for _call in mock_pip.call_args_list: + _args = [str(a) for a in _call.args] + if "--index-url" in _args: + _url = _args[_args.index("--index-url") + 1] + assert "rocm7.2" not in _url or "torch" not in " ".join( + _args + ), "torch must not be reinstalled when the pin already matches" + # A torch reinstall would pass torch>=... as a positional; assert none did. + assert not any( + any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list + ) + + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install_try", return_value = True) + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (6, 4)) + def test_non211_gfx_pin_over_210_rocm_no_reinstall( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """A gfx110X-all pin (NOT in the 2.11 allowlist) over a correct 2.10+rocm + wheel must NOT be flagged stale -- the install path uses the default <2.11 + specs for that arch, so re-flagging would reinstall-loop on every update.""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"6.4.12345|2.10.0+rocm6.4\n" + env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx110X-all"} + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + _ensure_rocm_torch() + # has_hip_torch True + no mismatch -> torch must NOT be reinstalled. + assert not any( + any(str(a).startswith("torch") for a in _c.args) for _c in mock_pip.call_args_list + ) + + @patch.object(stack_mod, "IS_WINDOWS", False) + @patch.object(stack_mod, "pip_install_try", return_value = True) + @patch.object(stack_mod, "pip_install") + @patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = False) + @patch.object(stack_mod, "_has_rocm_gpu", return_value = True) + @patch.object(stack_mod, "_detect_rocm_version", return_value = (7, 2)) + def test_gfx_pin_over_generic_rocm211_reinstalls( + self, mock_ver, mock_gpu, mock_nvidia, mock_pip, mock_pip_try + ): + """A gfx1151 pin over a GENERIC (two-part +rocm7.2) 2.11 wheel must reinstall + the AMD per-arch wheel -- even though both are torch 2.11, the generic wheel + is not the per-arch build the user pinned (Strix stays off the generic wheel).""" + mock_probe = MagicMock() + mock_probe.returncode = 0 + mock_probe.stdout = b"7.2.12345|2.11.0+rocm7.2\n" + env = {"UNSLOTH_TORCH_INDEX_URL": "https://repo.amd.com/rocm/whl/gfx1151"} + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + with patch("os.path.isdir", return_value = True): + with patch("subprocess.run", return_value = mock_probe): + with patch.object( + stack_mod, "_detect_amd_gfx_codes", side_effect = AssertionError + ): + _ensure_rocm_torch() + torch_call = str(mock_pip.call_args_list[0]) + assert "gfx1151" in torch_call + assert "torch>=2.11.0,<2.12.0" in torch_call + + def test_radeon_url_not_classified_as_pip_rocm_family(self): + """A repo.radeon.com find-links dir (leaf rocm-rel-7.2.1) starts with "rocm" but is + NOT a pip --index-url ROCm family: it must route to the verbatim path, not a + --index-url reinstall that fails against a find-links listing.""" + leaf_f = stack_mod._is_pip_rocm_family_leaf + # Real pip ROCm families (download.pytorch.org/whl/rocmX.Y, repo.amd.com gfx). + assert leaf_f("rocm7.2") is True + assert leaf_f("rocm6.4") is True + assert leaf_f("gfx120x-all") is True + assert leaf_f("gfx1151") is True + # A bare rocm (no minor) is still an exact family. + assert leaf_f("rocm7") is True + # A Radeon find-links dir leaf, a custom mirror, cpu and cuda are NOT pip rocm. + assert leaf_f("rocm-rel-7.2.1") is False + assert leaf_f("simple") is False + assert leaf_f("current") is False + assert leaf_f("cpu") is False + assert leaf_f("cu128") is False + # A rocm-SUFFIX private mirror shares the family prefix but is a custom pin + # the verbatim path owns: a ^rocm\d PREFIX match would wrongly treat it as a + # --index-url family. Match EXACTLY. + assert leaf_f("rocm7.2-private") is False + assert leaf_f("rocm7-current") is False + assert leaf_f("rocm7.2.1") is False # two-part local suffix -> custom, not rocm7.2 + + radeon = "https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1" + pip_rocm = "https://download.pytorch.org/whl/rocm7.2" + amd_gfx = "https://repo.amd.com/rocm/whl/gfx120X-all" + + def _classify(url, fn): + with patch.dict(stack_mod.os.environ, {"UNSLOTH_TORCH_INDEX_URL": url}, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + return fn() + + rocm_fn = stack_mod._explicit_rocm_torch_index_url + unk_fn = stack_mod._explicit_unknown_family_torch_index_url + # Real pip rocm/gfx pins ARE a ROCm family (reinstallable via --index-url) and + # are NOT "unknown". + assert _classify(pip_rocm, rocm_fn) == pip_rocm + assert _classify(amd_gfx, rocm_fn) == amd_gfx + assert _classify(pip_rocm, unk_fn) is None + assert _classify(amd_gfx, unk_fn) is None + # The Radeon find-links URL is NOT a pip ROCm family (so _ensure_rocm_torch skips + # it) and IS unknown, so the family repair helpers leave it alone. + assert _classify(radeon, rocm_fn) is None + assert _classify(radeon, unk_fn) == radeon + + # A rocm-suffix private mirror routes the same way: NOT a pip rocm family, + # IS an unknown-family (verbatim) pin. + suffixed = "https://co.internal/whl/rocm7.2-private" + assert _classify(suffixed, rocm_fn) is None + assert _classify(suffixed, unk_fn) == suffixed + + @patch.object(stack_mod, "pip_install") + def test_ensure_cpu_torch_broken_probe_reinstalls(self, mock_pip): + """_ensure_cpu_torch: torch present but unimportable (probe exit != 0) under an + explicit CPU pin must reinstall from the pin, not return -- the base update does + not repair a broken installed torch, so returning would strand it (Codex P2).""" + mock_probe = MagicMock() + mock_probe.returncode = 1 # torch present but cannot import + mock_probe.stdout = b"" + env = {"UNSLOTH_TORCH_INDEX_URL": "https://mirror.local/cpu"} + with patch.dict(stack_mod.os.environ, env, clear = False): + stack_mod.os.environ.pop("UNSLOTH_TORCH_INDEX_FAMILY", None) + with patch("subprocess.run", return_value = mock_probe): + with patch.object(stack_mod, "NO_TORCH", False): + stack_mod._ensure_cpu_torch() + assert mock_pip.call_count == 1 + assert "https://mirror.local/cpu" in str(mock_pip.call_args) + @patch.object(stack_mod, "IS_WINDOWS", False) @patch.object(stack_mod, "pip_install_try", return_value = True) @patch.object(stack_mod, "pip_install") @@ -732,7 +1048,7 @@ class TestEnsureRocmTorch: mock_pip.assert_not_called() -# TEST: install_python_stack.py -- _has_rocm_gpu KFD sysfs vendor_id guard +# TEST: install_python_stack.py -- torch-index MARKER mechanism (PR #6692) class TestHasRocmGpuKfdVendorGuard: @@ -1716,9 +2032,8 @@ class TestDetectWindowsGfxArch: assert result == "gfx1200" def test_returns_arch_on_crash_with_gcnarchname_in_output(self): - # Regression #6043: hipinfo may crash (0xC0000005 on RDNA 4) after - # printing gcnArchName. Accept the arch whenever gcnArchName is in - # stdout, regardless of exit code (previously a CPU fallback). + # Regression #6043: hipinfo may crash (0xC0000005 on RDNA 4) after printing + # gcnArchName. Accept the arch whenever gcnArchName is in stdout, any exit code. mock_result = MagicMock() mock_result.returncode = -1073741819 # 0xC0000005 STATUS_ACCESS_VIOLATION mock_result.stdout = b"gcnArchName : gfx1200\nsome other line\n" @@ -2360,6 +2675,49 @@ class TestWindowsRocmTorchaoGuard: assert not any("torchao" in arg for arg in installed_specs) +class TestProgressStepCountMatchesTotal: + """The progress bar must reach exactly _TOTAL: every _progress() step is counted in + base_total. Regression for a repair step added without incrementing base_total, + which pushed _STEP past _TOTAL (Codex P2).""" + + def _run_stack(self, tmp_path, *, is_windows, is_macos, is_mac_arm): + unstructured_plugin = tmp_path / "unstructured" + github_plugin = tmp_path / "github" + unstructured_plugin.mkdir() + github_plugin.mkdir() + sub = MagicMock() + sub.returncode = 0 + sub.stdout = "" + with ( + patch.dict(os.environ, {"SKIP_STUDIO_BASE": "1"}), + patch.object(stack_mod, "IS_WINDOWS", is_windows), + patch.object(stack_mod, "IS_MACOS", is_macos), + patch.object(stack_mod, "IS_MAC_ARM", is_mac_arm), + patch.object(stack_mod, "NO_TORCH", False), + patch.object(stack_mod, "_rocm_windows_torch_installed", False), + patch.object(stack_mod, "_bootstrap_uv", return_value = False), + patch.object(stack_mod, "_installed_torch_is_windows_rocm", return_value = False), + patch.object(stack_mod, "_has_usable_nvidia_gpu", return_value = True), + patch.object(stack_mod, "_repair_bad_anyio"), + patch.object(stack_mod, "_ensure_cuda_torch"), + patch.object(stack_mod, "_ensure_rocm_torch"), + patch.object(stack_mod, "_ensure_cpu_torch"), + patch.object(stack_mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin), + patch.object(stack_mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin), + patch.object(stack_mod.subprocess, "run", return_value = sub), + ): + assert stack_mod.install_python_stack() == 0 + return stack_mod._STEP, stack_mod._TOTAL + + def test_windows_progress_reaches_total(self, tmp_path): + step, total = self._run_stack(tmp_path, is_windows = True, is_macos = False, is_mac_arm = False) + assert step == total, f"Windows progress {step} != total {total} (final step uncounted)" + + def test_linux_progress_reaches_total(self, tmp_path): + step, total = self._run_stack(tmp_path, is_windows = False, is_macos = False, is_mac_arm = False) + assert step == total, f"Linux progress {step} != total {total}" + + # TEST: worker.py -- Windows ROCm patches (source-level checks) @@ -2846,6 +3204,24 @@ class TestStrixRocm71Override: source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") assert "TORCH_CONSTRAINT" in source and "2.11" in source + def test_torch_constraint_211_matches_leaf_not_whole_url(self): + """The 2.11 constraint case must match the index LEAF, not the whole URL. + + A custom UNSLOTH_PYTORCH_MIRROR whose base path contains a gfx/rocm7.2 + segment (e.g. https://mirror.local/gfx-cache) with a cu*/cpu family must + not be pushed to the torch 2.11 line -- same leaf-only reasoning the + UNSLOTH_TORCH_BACKEND classification uses. + """ + source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") + # The 2.11 constraint block must switch on $_torch_index_leaf, not the full + # $TORCH_INDEX_URL (a */gfx* match false-positives on a mirror base path). Only the + # _grouped_mm-bug gfx families (gfx120X-all / gfx1151 / gfx1150) are pushed to 2.11; + # a bare gfx* would also floor gfx110X-all/gfx90a/gfx908, left bare on purpose. + assert 'case "$_torch_index_leaf" in\n rocm7.2|gfx120x-all|gfx1151|gfx1150)' in source, ( + "the torch>=2.11 constraint must match the specific gfx leaves that need " + "it (rocm7.2|gfx120x-all|gfx1151|gfx1150), not a bare gfx* or the whole URL" + ) + def test_amd_rocm_mirror_env_var_respected(self): """install.sh must honour UNSLOTH_AMD_ROCM_MIRROR for air-gapped installs.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") @@ -2934,9 +3310,9 @@ class TestServerStartupRocmFixes: assert '"BNB_ROCM_VERSION" not in os.environ' in source # ── hipInfo.exe PATH prepend (bitsandbytes arch-probe fix) ──────────────── - # bnb's get_rocm_gpu_arch() runs hipinfo.exe via PATH at import; the AMD - # wheel ships it in venv Scripts (on PATH only for activated venvs), so - # without the prepend bnb logs "[WinError 2]" when launched directly. + # bnb's get_rocm_gpu_arch() runs hipinfo.exe via PATH at import; the AMD wheel ships it + # in venv Scripts (on PATH only for activated venvs), so without the prepend bnb logs + # "[WinError 2]" when launched directly. def test_main_py_prepends_hipinfo_dir_to_path(self): """main.py must make hipInfo.exe resolvable before bnb imports.""" @@ -3178,11 +3554,10 @@ class TestRocmGfxForwarding: assert '$HelperReleaseRepo = "unslothai/llama.cpp"' in source assert "$HelperReleaseRepo = if (" not in source - # The text pins above guard the literal. The tests below *execute* the real - # routing line from setup.sh / setup.ps1 and assert the resolved release repo, - # so a refactor that reintroduces a conditional (or a ggml-org branch) is still - # caught. Inputs are varied -- CPU-only, inferred/forwarded gfx, usable NVIDIA -- - # to prove no host slips back onto ggml-org. No GPU, no tooling, no network. + # The text pins above guard the literal. The tests below execute the real routing line + # from setup.sh / setup.ps1 and assert the resolved release repo, so a refactor that + # reintroduces a conditional (or a ggml-org branch) is still caught. Inputs vary + # (CPU-only, inferred/forwarded gfx, usable NVIDIA) to prove no host hits ggml-org. @staticmethod def _resolve_setup_sh_repo( @@ -3277,8 +3652,8 @@ class TestRocmGfxForwarding: # TEST: _pick_rocm_gfx_target -- visible-device selection from rocminfo output. -# Honours CUDA/HIP_VISIBLE_DEVICES so a mixed-arch host installs the prebuilt -# for the selected GPU, not GPU 0. +# Honours CUDA/HIP_VISIBLE_DEVICES so a mixed-arch host installs the prebuilt for the +# selected GPU, not GPU 0. _pick_rocm_gfx_target = prebuilt_mod._pick_rocm_gfx_target @@ -3454,7 +3829,11 @@ class TestWslRerouteNvidiaGuard: source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") start = source.find("_maybe_reroute_strixhalo_to_2404()") assert start != -1 - body = source[start : start + 1200] + # Slice the WHOLE function body (to its closing brace at column 0), not a + # fixed-length window: preamble growth must not push the signals out of view. + end = source.find("\n}", start) + assert end != -1 + body = source[start:end] nv = body.find("_has_usable_nvidia_gpu") wmi = body.find("_wsl_amd_gpu_name") assert nv != -1, "reroute must consult _has_usable_nvidia_gpu before deciding to reroute" diff --git a/tests/studio/test_setup_pin_stale.ps1 b/tests/studio/test_setup_pin_stale.ps1 new file mode 100644 index 0000000000..2c92ae317f --- /dev/null +++ b/tests/studio/test_setup_pin_stale.ps1 @@ -0,0 +1,114 @@ +#!/usr/bin/env pwsh +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Unit test for studio/setup.ps1's pinned-torch-index stale-venv helpers +# (Test-RocmGfx211Leaf, Test-CudaFamilyLeaf, Get-RocmPinStaleTags). Pure helpers, +# AST-extracted and run in-process. Mirrors the Python _rocm_pin_family_mismatch / +# _is_cuda_family_leaf tests. +# Run: pwsh -NoProfile -File tests/studio/test_setup_pin_stale.ps1 + +$ErrorActionPreference = "Stop" +$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1") +$setupPath = (Resolve-Path $setupPath).Path + +# --- Parse setup.ps1 (also serves as a syntax gate) and extract the helpers --- +$tokens = $null; $errors = $null +$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors) +if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" } + +foreach ($name in @("Test-RocmGfx211Leaf", "Test-RocmKnown211Version", "Test-CudaFamilyLeaf", "Get-RocmPinStaleTags")) { + $fn = $ast.FindAll({ param($n) + $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name + }, $true) + if ($fn.Count -ne 1) { throw "expected exactly one $name in setup.ps1, found $($fn.Count)" } + # Pure helpers (no exit / external calls) -- safe to define in this scope. + Invoke-Expression $fn[0].Extent.Text +} + +$failures = 0 +function Check($name, $cond) { + if ($cond) { Write-Host " PASS $name" } + else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ } +} + +# A pinned gfx/rocm index is stale when Expected != Installed. +function IsStale($leaf, $ver) { + $t = Get-RocmPinStaleTags -PinLeaf $leaf -TorchVersion $ver + return $t.Expected -ne $t.Installed +} + +Write-Host "Test-RocmGfx211Leaf (the 2.11 gfx allowlist)" +Check "gfx1151 -> true" (Test-RocmGfx211Leaf "gfx1151") +Check "gfx1150 -> true" (Test-RocmGfx211Leaf "gfx1150") +Check "gfx120x-all -> true" (Test-RocmGfx211Leaf "gfx120x-all") +Check "gfx110x-all -> false" (-not (Test-RocmGfx211Leaf "gfx110x-all")) +Check "gfx90a -> false" (-not (Test-RocmGfx211Leaf "gfx90a")) +Check "gfx908 -> false" (-not (Test-RocmGfx211Leaf "gfx908")) + +Write-Host "Test-CudaFamilyLeaf (^cu[0-9])" +Check "cu118 -> true" (Test-CudaFamilyLeaf "cu118") +Check "cu128 -> true" (Test-CudaFamilyLeaf "cu128") +Check "cu130 -> true" (Test-CudaFamilyLeaf "cu130") +Check "custom -> false" (-not (Test-CudaFamilyLeaf "custom")) +Check "current -> false" (-not (Test-CudaFamilyLeaf "current")) +Check "cpu -> false" (-not (Test-CudaFamilyLeaf "cpu")) +Check "empty -> false" (-not (Test-CudaFamilyLeaf "")) + +Write-Host "Get-RocmPinStaleTags (mirror of _rocm_pin_family_mismatch)" +# Exact rocm version comparison. +Check "rocm7.2 pin + 2.11.0+rocm7.2 -> not stale" (-not (IsStale "rocm7.2" "2.11.0+rocm7.2")) +Check "rocm7.2 pin + 2.10.0+rocm6.4 -> stale" (IsStale "rocm7.2" "2.10.0+rocm6.4") +Check "rocm6.4 pin + 2.10.0+rocm6.4 -> not stale" (-not (IsStale "rocm6.4" "2.10.0+rocm6.4")) +# rocm7.2 is a KNOWN-2.11 index. A +rocm7.2 wheel whose RELEASE drifted off 2.11 shares +# the tag but violates the spec -> stale (mirror of _rocm_pin_family_mismatch). +Check "rocm7.2 pin + 2.12.0+rocm7.2 -> stale" (IsStale "rocm7.2" "2.12.0+rocm7.2") +Check "rocm7.2 pin + 2.13.0+rocm7.2 -> stale" (IsStale "rocm7.2" "2.13.0+rocm7.2") +Check "rocm7.2 pin + 2.11.5+rocm7.2 -> not stale" (-not (IsStale "rocm7.2" "2.11.5+rocm7.2")) +# An UNKNOWN newer rocm (off the 2.11 allowlist) isn't floored, so a matching version at +# any release line is NOT stale on this exact-compare branch. +Check "rocm8.0 pin + 2.12.0+rocm8.0 -> not stale" (-not (IsStale "rocm8.0" "2.12.0+rocm8.0")) +# An untagged (no +rocm) wheel never satisfies a ROCm pin -> always stale. +Check "rocm7.2 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm7.2" "2.10.0") +Check "rocm7.2 pin + 2.11.0 (untagged) -> stale" (IsStale "rocm7.2" "2.11.0") +Check "rocm6.4 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm6.4" "2.10.0") +# 2.11-allowlist gfx pin: per-arch (three-part) wheel is satisfied, generic is stale. +Check "gfx1151 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx1151" "2.11.0+rocm7.13.0")) +Check "gfx1150 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx1150" "2.11.0+rocm7.13.0")) +Check "gfx120x-all pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "gfx120x-all" "2.11.0+rocm7.13.0")) +Check "gfx1151 pin + 2.11.0+rocm7.2 (generic) -> stale" (IsStale "gfx1151" "2.11.0+rocm7.2") +Check "gfx1151 pin + 2.10.0+rocm6.4 -> stale" (IsStale "gfx1151" "2.10.0+rocm6.4") +# Non-2.11 gfx pin (gfx110X-all/gfx90a/gfx908): a valid <2.11 wheel is NOT stale. +Check "gfx110x-all pin + 2.10.0+rocm6.4 -> not stale" (-not (IsStale "gfx110x-all" "2.10.0+rocm6.4")) +Check "gfx90a pin + 2.10.0+rocm6.3 -> not stale" (-not (IsStale "gfx90a" "2.10.0+rocm6.3")) +Check "gfx908 pin + 2.10.0+rocm7.0 -> not stale" (-not (IsStale "gfx908" "2.10.0+rocm7.0")) +Check "gfx110x-all pin + 2.11.0+rocm7.2 -> stale" (IsStale "gfx110x-all" "2.11.0+rocm7.2") +# Non-2.11 gfx pin over an untagged wheel: never satisfies the pin -> stale, so the +# explicit ROCm index is applied even when torch is already <2.11. +Check "gfx110x-all pin + 2.10.0 (untagged) -> stale" (IsStale "gfx110x-all" "2.10.0") +Check "gfx90a pin + 2.10.0 (untagged) -> stale" (IsStale "gfx90a" "2.10.0") +# Capital gfx120X-all is lowercased by Get-TorchIndexLeaf before this helper, so the +# 2.11-allowlist branch fires and a generic/untagged wheel is stale. +Check "gfx120x-all pin + 2.11.0+rocm7.2 (generic) -> stale" (IsStale "gfx120x-all" "2.11.0+rocm7.2") +Check "gfx120x-all pin + 2.10.0 (untagged) -> stale" (IsStale "gfx120x-all" "2.10.0") + +# Major-only rocm pin (rocm7): majors compared alone; mirrors _rocm_pin_family_mismatch. +Check "rocm7 pin + 2.10.0+rocm6.4 -> stale" (IsStale "rocm7" "2.10.0+rocm6.4") +Check "rocm7 pin + 2.11.0+rocm7.2 -> not stale" (-not (IsStale "rocm7" "2.11.0+rocm7.2")) +Check "rocm7 pin + 2.11.0+rocm7.13.0 -> not stale" (-not (IsStale "rocm7" "2.11.0+rocm7.13.0")) +Check "rocm7 pin + 2.10.0 (untagged) -> stale" (IsStale "rocm7" "2.10.0") +Check "rocm7 pin + 2.10.0+rocm (unreadable) -> not stale" (-not (IsStale "rocm7" "2.10.0+rocm")) + +Write-Host "Test-RocmKnown211Version + KNOWN-2.11 fallback (rocm7.2 only; no speculative rocm7.3)" +Check "rocm7.2 -> known 2.11" (Test-RocmKnown211Version -Major 7 -Minor 2) +Check "rocm7.1 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 1)) +Check "rocm7.3 -> not known" (-not (Test-RocmKnown211Version -Major 7 -Minor 3)) +Check "rocm8.0 -> not known" (-not (Test-RocmKnown211Version -Major 8 -Minor 0)) +# Unreadable-installed fallback: a rocm7.3 pin (unknown -> <2.11 line) over a <2.11 +rocm +# wheel with an unreadable version is NOT stale; rocm7.2 (KNOWN-2.11) over the same wheel +# IS stale (#2534 alignment). +Check "rocm7.3 pin + 2.10.0+rocm (unreadable ver) -> not stale" (-not (IsStale "rocm7.3" "2.10.0+rocm")) +Check "rocm7.2 pin + 2.10.0+rocm (unreadable ver) -> stale" (IsStale "rocm7.2" "2.10.0+rocm") + +Write-Host "" +if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 } +Write-Host "All checks passed" -ForegroundColor Green diff --git a/tests/studio/test_torch_flavor.ps1 b/tests/studio/test_torch_flavor.ps1 index 50be4814b0..f2cc55c21c 100644 --- a/tests/studio/test_torch_flavor.ps1 +++ b/tests/studio/test_torch_flavor.ps1 @@ -15,7 +15,7 @@ $tokens = $null; $errors = $null $ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors) if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" } -foreach ($name in @("ConvertTo-TorchFlavorTag", "Get-ExpectedTorchFlavorTag")) { +foreach ($name in @("ConvertTo-TorchFlavorTag", "Get-ExpectedTorchFlavorTag", "Trim-IndexPathSlashes", "Redact-InstallOutput")) { $fn = $ast.FindAll({ param($n) $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name }, $true) @@ -49,6 +49,19 @@ Check "mirror cu130 leaf -> cu130" ((Get-ExpectedTorchFlavorTag -TorchIndexUrl Check "unrecognized leaf -> null" ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl "https://my.mirror/whl/simple")) Check "empty url -> null" ($null -eq (Get-ExpectedTorchFlavorTag -TorchIndexUrl "")) +Write-Host "Trim-IndexPathSlashes (install.ps1 parity: path-only, token-preserving)" +Check "double path slash collapsed" ((Trim-IndexPathSlashes "https://h/whl/cu128//") -eq "https://h/whl/cu128") +Check "single trailing slash trimmed" ((Trim-IndexPathSlashes "https://h/whl/cu128/") -eq "https://h/whl/cu128") +Check "query token slash preserved" ((Trim-IndexPathSlashes "https://h/whl/cu128?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/") +Check "path slash trimmed, query kept" ((Trim-IndexPathSlashes "https://h/whl/cu128//?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/") + +Write-Host "Redact-InstallOutput (install.ps1 parity: credential redaction)" +Check "userinfo redacted" ((Redact-InstallOutput "ERROR https://alice:s3cr3t@download.pytorch.org/whl/cu128") -eq "ERROR https://@download.pytorch.org/whl/cu128") +Check "query value redacted" ((Redact-InstallOutput "https://host/whl/cu128?token=abcd1234&channel=beta") -eq "https://host/whl/cu128?token=&channel=") +Check "fragment token redacted" ((Redact-InstallOutput "ERROR https://mirror.local/whl/cu128#token=SECRET123 (403)") -eq "ERROR https://mirror.local/whl/cu128# (403)") +Check "bare hash comment untouched" ((Redact-InstallOutput "# retrying with --no-cache-dir") -eq "# retrying with --no-cache-dir") +Check "plain line untouched" ((Redact-InstallOutput "Resolved 42 packages in 1.2s") -eq "Resolved 42 packages in 1.2s") + Write-Host "" if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 } Write-Host "All checks passed" -ForegroundColor Green diff --git a/tests/studio/test_torch_index_pin_hardening.ps1 b/tests/studio/test_torch_index_pin_hardening.ps1 new file mode 100644 index 0000000000..b8edf3da25 --- /dev/null +++ b/tests/studio/test_torch_index_pin_hardening.ps1 @@ -0,0 +1,78 @@ +#!/usr/bin/env pwsh +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Unit tests for setup.ps1's torch-index pin-hardening helpers: Trim-IndexPathSlashes +# (path-only slash trim, token-preserving), Redact-InstallOutput (credential redaction of +# captured install logs), Get-TorchIndexLeaf (ALL trailing slashes stripped) and +# Test-PipRocmFamilyLeaf (rocm7. is a custom pin, not a family). Pure helpers, AST-extracted +# and run in-process. Run: pwsh -NoProfile -File tests/studio/test_torch_index_pin_hardening.ps1 + +$ErrorActionPreference = "Stop" +$setupPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "studio", "setup.ps1") +$setupPath = (Resolve-Path $setupPath).Path +$setupText = Get-Content -Raw $setupPath + +# --- Parse setup.ps1 (also a syntax gate) and extract the pure helpers --- +$tokens = $null; $errors = $null +$ast = [System.Management.Automation.Language.Parser]::ParseFile($setupPath, [ref]$tokens, [ref]$errors) +if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "setup.ps1 has parse errors" } + +foreach ($name in @("Trim-IndexPathSlashes", "Redact-InstallOutput", "Get-TorchIndexLeaf", "Test-PipRocmFamilyLeaf")) { + $fn = $ast.FindAll({ param($n) + $n -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $n.Name -eq $name + }, $true) + if ($fn.Count -ne 1) { throw "expected exactly one $name in setup.ps1, found $($fn.Count)" } + Invoke-Expression $fn[0].Extent.Text +} + +$failures = 0 +function Check($name, $cond) { + if ($cond) { Write-Host " PASS $name" } + else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ } +} + +Write-Host "Trim-IndexPathSlashes (path-only, token-preserving)" +Check "double path slash collapsed" ((Trim-IndexPathSlashes "https://h/whl/cu128//") -eq "https://h/whl/cu128") +Check "single trailing slash trimmed" ((Trim-IndexPathSlashes "https://h/whl/cu128/") -eq "https://h/whl/cu128") +Check "no slash unchanged" ((Trim-IndexPathSlashes "https://h/whl/cu128") -eq "https://h/whl/cu128") +Check "query token slash preserved" ((Trim-IndexPathSlashes "https://h/whl/cu128?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/") +Check "path slash trimmed, query kept" ((Trim-IndexPathSlashes "https://h/whl/cu128//?token=ab12cd/") -eq "https://h/whl/cu128?token=ab12cd/") +Check "fragment slash preserved" ((Trim-IndexPathSlashes "https://h/whl/cu128#anchor/") -eq "https://h/whl/cu128#anchor/") + +Write-Host "Redact-InstallOutput (credential redaction)" +Check "userinfo redacted" ((Redact-InstallOutput "ERROR https://alice:s3cr3t@download.pytorch.org/whl/cu128") -eq "ERROR https://@download.pytorch.org/whl/cu128") +Check "bare-token@ redacted" ((Redact-InstallOutput "fetch https://ghp_deadbeef@host/whl/cu128 failed") -eq "fetch https://@host/whl/cu128 failed") +Check "single query value redacted" ((Redact-InstallOutput "url https://host/whl/cu128?token=abcd1234 unreachable") -eq "url https://host/whl/cu128?token= unreachable") +Check "multiple query values redacted" ((Redact-InstallOutput "https://host/whl/cu128?token=abcd1234&channel=beta") -eq "https://host/whl/cu128?token=&channel=") +Check "fragment token redacted" ((Redact-InstallOutput "ERROR https://mirror.local/whl/cu128#token=SECRET123 (403)") -eq "ERROR https://mirror.local/whl/cu128# (403)") +Check "query and fragment both redacted" ((Redact-InstallOutput "https://host/whl/cu128?token=abc#sig=xyz done") -eq "https://host/whl/cu128?token=# done") +Check "bare hash comment untouched" ((Redact-InstallOutput "# retrying with --no-cache-dir") -eq "# retrying with --no-cache-dir") +Check "plain line untouched" ((Redact-InstallOutput "Resolved 42 packages in 1.2s") -eq "Resolved 42 packages in 1.2s") +$leak = Redact-InstallOutput "https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET" +Check "no secret substring survives" (($leak -notmatch "s3cr3t") -and ($leak -notmatch "SUPERSECRET") -and ($leak -notmatch "ALSOSECRET")) + +Write-Host "Get-TorchIndexLeaf (ALL trailing slashes stripped)" +Check "double slash cu128 -> cu128" ((Get-TorchIndexLeaf "https://m/whl/cu128//") -eq "cu128") +Check "triple slash rocm7.2 -> rocm7.2" ((Get-TorchIndexLeaf "https://m/whl/rocm7.2///") -eq "rocm7.2") +Check "double slash + token -> cu128" ((Get-TorchIndexLeaf "https://m/whl/cu128//?token=x") -eq "cu128") +Check "single slash cu128 -> cu128" ((Get-TorchIndexLeaf "https://m/whl/cu128/") -eq "cu128") + +Write-Host "Test-PipRocmFamilyLeaf (rocm7. is a custom pin, not a family)" +Check "rocm7 family" (Test-PipRocmFamilyLeaf "rocm7") +Check "rocm7.2 family" (Test-PipRocmFamilyLeaf "rocm7.2") +Check "gfx1151 family" (Test-PipRocmFamilyLeaf "gfx1151") +Check "rocm7. trailing-dot NOT family" (-not (Test-PipRocmFamilyLeaf "rocm7.")) +Check "rocm.7 leading-dot NOT family" (-not (Test-PipRocmFamilyLeaf "rocm.7")) +Check "rocm7.2.1 two-dot NOT family" (-not (Test-PipRocmFamilyLeaf "rocm7.2.1")) +Check "rocm7.2-private NOT family" (-not (Test-PipRocmFamilyLeaf "rocm7.2-private")) +Check "cu128 NOT family" (-not (Test-PipRocmFamilyLeaf "cu128")) + +Write-Host "Fast-Install pinned-install env scrub (source assertion)" +# The pip fallback honours PIP_*; PIP_NO_INDEX=1 would make it ignore the pinned --index-url +# and PIP_INDEX_URL would replace it, so both must be scrubbed for a pinned install. +Check "PIP_NO_INDEX scrubbed" ($setupText -match "'PIP_NO_INDEX'") +Check "PIP_INDEX_URL scrubbed" ($setupText -match "'PIP_INDEX_URL'") + +Write-Host "" +if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 } +Write-Host "All checks passed" -ForegroundColor Green