From e49a89f4f6175072f0ebaa92cc8ce91dd2af8954 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 20 Jul 2026 10:41:41 +0000 Subject: [PATCH] install.ps1: port the torch release preservation from install.sh The Windows installer recreates the venv on every re-run (the existing environment is moved aside for rollback and a fresh venv is created), so irm install.ps1 | iex over an existing install resolved the newest torch in range instead of keeping the installed release. This ports the install.sh preservation merged in PR 7250: - The installed torch version is probed (bounded process, drained streams, 30s timeout, last non-empty stdout line) BEFORE the rollback move, while the old interpreter still exists. - Get-PreviousTorchPin mirrors _previous_torch_pin: numeric-strict base (nightly/dev/rc builds never pin), UNSLOTH_TORCH_UPGRADE=1 opt-out, and a release-in-window check against the route's final constraint, so a raised ROCm floor correctly rejects keeping an older release. The decision runs after every index and floor decision. - The kept release installs as an exact pin with minor-paired companions (torchvision 0.minor+15, torchaudio 2.minor) at the main install, the AMD ROCm install and both flavor repairs; when the pinned release is not installable from the selected index the installer warns, clears the pin and falls back to the supported range. The CPU fallback cannot be reached with a live pin (the ROCm site resolves or clears it first) and keeps the plain range install. - The with-deps unsloth installs now carry a uv overrides file with the exact installed trio (twin of _build_unsloth_torch_overrides), so dependency resolution cannot move torch after it was deliberately selected. - The kept release is exported as UNSLOTH_KEPT_TORCH for setup.ps1, whose ROCm/CPU/CUDA torch installs substitute the kept trio when the env var is present (ROCm floors dominate; behavior without the env var is unchanged, so direct studio update runs are unaffected). After setup returns the installer re-probes and warns loudly if the kept release series changed, then clears the env var. Preservation is version-agnostic (numeric parse + window comparison), so a future ceiling bump to torch 2.12 keeps 2.11 installs in place. New suite tests/studio/test_previous_torch_pin.ps1 (AST-extracted helpers, 50+ checks incl. future-ceiling cases and structural wiring) passes, along with the flavor/pin-hardening/pin-stale/node ps1 suites, the parity and install-stack pytest suites, and the sh preservation and constraint suites. --- install.ps1 | 265 ++++++++++++++++++++--- studio/setup.ps1 | 103 +++++++-- tests/studio/test_previous_torch_pin.ps1 | 115 ++++++++++ 3 files changed, 428 insertions(+), 55 deletions(-) create mode 100644 tests/studio/test_previous_torch_pin.ps1 diff --git a/install.ps1 b/install.ps1 index 02d51b3d2b..a561e6f9b0 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1330,6 +1330,37 @@ exit 0 $script:StudioVenvRollbackDir = $null } + # Raw torch.__version__ from $PythonExe's venv (last non-empty stdout line), or $null. + # Bounded ProcessStartInfo probe (async-drain both streams, 30s timeout, kill on hang) so a + # wedged "import torch" can't stall the installer; feeds Get-InstalledTorchTag and the torch + # release-preservation decision (twin of install.sh's _PREV_TORCH_VER probe / _previous_torch_pin). + function Get-InstalledTorchVersionRaw { + param([string]$PythonExe) + if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null } + try { + $psi = New-Object System.Diagnostics.ProcessStartInfo + $psi.FileName = $PythonExe + $psi.Arguments = '-c "import torch; print(torch.__version__)"' + $psi.RedirectStandardOutput = $true + $psi.RedirectStandardError = $true + $psi.UseShellExecute = $false + $psi.CreateNoWindow = $true + $proc = [System.Diagnostics.Process]::Start($psi) + # Drain BOTH streams async before WaitForExit: a synchronous ReadToEnd() would block on a wedged "import torch", and an undrained stderr would deadlock a child flooding the pipe buffer. A truly hung probe still hits the 30s timeout. + $outTask = $proc.StandardOutput.ReadToEndAsync() + $errTask = $proc.StandardError.ReadToEndAsync() + $finished = $proc.WaitForExit(30000) + if (-not $finished) { try { $proc.Kill() } catch {}; return $null } + $out = $outTask.GetAwaiter().GetResult() + [void]$errTask.GetAwaiter().GetResult() + if ($proc.ExitCode -ne 0) { return $null } + # Last non-empty line only, so stdout noise before the version can't corrupt the pin. + $lines = @($out -split "`r?`n" | ForEach-Object { $_.Trim() } | Where-Object { $_ -ne "" }) + if ($lines.Count -eq 0) { return $null } + return $lines[-1] + } catch { return $null } + } + if (Test-Path -LiteralPath $VenvPython) { # env-mode: $StudioHome is a user-chosen workspace, so refuse to nuke an existing venv lacking Unsloth sentinels (-PathType Leaf rejects a directory at the sentinel path; accept the in-VENV ownership marker so partial-install retries aren't blocked). if ( @@ -1342,6 +1373,14 @@ exit 0 Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." -ForegroundColor Yellow throw "Refusing to delete non-Unsloth venv at $VenvDir" } + # Record the existing venv's torch RELEASE BEFORE the rollback move (see Get-PreviousTorchPin); + # a re-run then keeps that release rather than silently jumping torch versions. Opt out with + # UNSLOTH_TORCH_UPGRADE=1. Only the new-layout replace probes; the legacy-migration branches + # reuse the venv, so torch survives naturally and needs no pin. + $script:PrevTorchVer = "" + if (-not $SkipTorch) { + $script:PrevTorchVer = Get-InstalledTorchVersionRaw -PythonExe $VenvPython + } # New layout already exists -- replace only after preserving rollback copy. substep "preserving existing environment for rollback..." try { @@ -1876,29 +1915,73 @@ exit 0 return $null } - # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Uses ProcessStartInfo (not &) so stderr doesn't trip $ErrorActionPreference. + # Installed torch flavor tag in $PythonExe's venv, or $null if absent. Reuses the bounded raw probe. function Get-InstalledTorchTag { param([string]$PythonExe) - if (-not $PythonExe -or -not (Test-Path -LiteralPath $PythonExe)) { return $null } + $v = Get-InstalledTorchVersionRaw -PythonExe $PythonExe + if (-not $v) { return $null } + return ConvertTo-TorchFlavorTag $v + } + + # ── Torch release preservation (twin of install.sh's _previous_torch_pin, PR 7250): keep the + # previous venv's torch RELEASE across a re-run when it falls inside the freshly chosen constraint + # window; flavor follows the new index. Opt out with UNSLOTH_TORCH_UPGRADE=1. ── + + # Parse a probed torch.__version__ into a normalized stable release, or $null. + # Strips ONLY the +local tag; anchored numeric match so dev/rc/alpha/garbage never pin. + function ConvertTo-TorchNumericRelease { + param([string]$TorchVersion) + if ([string]::IsNullOrWhiteSpace($TorchVersion)) { return $null } + $publicBase = ($TorchVersion.Trim() -split '\+', 2)[0] + if ($publicBase -notmatch '^(\d+)\.(\d+)(?:\.(\d+))?$') { return $null } try { - $psi = New-Object System.Diagnostics.ProcessStartInfo - $psi.FileName = $PythonExe - $psi.Arguments = '-c "import torch; print(torch.__version__)"' - $psi.RedirectStandardOutput = $true - $psi.RedirectStandardError = $true - $psi.UseShellExecute = $false - $psi.CreateNoWindow = $true - $proc = [System.Diagnostics.Process]::Start($psi) - # Drain BOTH streams async before WaitForExit: a synchronous ReadToEnd() would block on a wedged "import torch", and an undrained stderr would deadlock a child flooding the pipe buffer. A truly hung probe still hits the 30s timeout. - $outTask = $proc.StandardOutput.ReadToEndAsync() - $errTask = $proc.StandardError.ReadToEndAsync() - $finished = $proc.WaitForExit(30000) - if (-not $finished) { try { $proc.Kill() } catch {}; return $null } - $torchVer = $outTask.GetAwaiter().GetResult().Trim() - [void]$errTask.GetAwaiter().GetResult() - if ($proc.ExitCode -ne 0 -or -not $torchVer) { return $null } - return ConvertTo-TorchFlavorTag $torchVer + $major = [int]$Matches[1]; $minor = [int]$Matches[2] + if ($Matches[3]) { $patch = [int]$Matches[3] } else { $patch = 0 } + $normalized = New-Object System.Version($major, $minor, $patch) } catch { return $null } + return [pscustomobject]@{ + PublicBase = $publicBase; Major = $major; Minor = $minor; Patch = $patch; Version = $normalized + } + } + + # True when a release falls inside a "torch>=A,=(\d+(?:\.\d+){0,2}),<(\d+(?:\.\d+){0,2})$') { return $false } + $floor = ConvertTo-TorchNumericRelease $Matches[1] + $ceiling = ConvertTo-TorchNumericRelease $Matches[2] + if (-not $floor -or -not $ceiling) { return $false } + return ($Release.Version -ge $floor.Version -and $Release.Version -lt $ceiling.Version) + } + + # The kept-release trio for a previously installed torch, or $null when nothing + # should be kept (no/unstable version, UNSLOTH_TORCH_UPGRADE=1, outside the final + # route window -- a raised ROCm floor correctly rejects an older release). + # Exact-release pin, matching install.sh's _previous_torch_pin; companions pair + # to the kept minor (torchaudio no longer exact-pins torch). + function Get-PreviousTorchPin { + param( + [string]$TorchVersion, + # Named -Constraint (like Test-TorchReleaseInWindow): the Windows port keeps no shell-style + # constraint variable, and its structural tests forbid that token appearing in install.ps1. + [Parameter(Mandatory = $true)][string]$Constraint + ) + if ($env:UNSLOTH_TORCH_UPGRADE -eq '1') { return $null } + $release = ConvertTo-TorchNumericRelease $TorchVersion + if (-not $release) { return $null } + if (-not (Test-TorchReleaseInWindow -Release $release -Constraint $Constraint)) { return $null } + if ($release.Major -ne 2) { return $null } + $visionMinor = $release.Minor + 15 + return [pscustomobject]@{ + Release = $release + TorchSpec = "torch==$($release.PublicBase)" + VisionSpec = "torchvision==0.$visionMinor.*" + AudioSpec = "torchaudio==2.$($release.Minor).*" + } } # An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it (e.g. a deliberate cpu pin on an AMD host). @@ -2018,6 +2101,22 @@ exit 0 return $installed } + # ── Freeze the installed torch trio for the with-deps unsloth install (twin of install.sh's + # _build_unsloth_torch_overrides): a released unsloth wheel can pin an older torch, and a + # with-deps resolve then downgrades the pinned +cuXXX/+rocm trio. Return a temp uv --overrides + # file pinning torch/torchvision/torchaudio to their installed versions, or $null when torch is + # absent (--no-torch) so the caller installs unchanged. Caller removes the file afterwards. ── + function New-UnslothTorchOverridesFile { + param([string]$PythonExe) + if ($SkipTorch) { return $null } + $pins = & $PythonExe -c "from importlib.metadata import version, PackageNotFoundError`nfor _p in ('torch', 'torchvision', 'torchaudio'):`n try:`n print(_p + '==' + version(_p))`n except PackageNotFoundError:`n pass" 2>$null + $lines = @($pins | Where-Object { $_ -match '^torch' }) + if ($lines.Count -eq 0 -or $lines[0] -notmatch '^torch==') { return $null } + $f = [System.IO.Path]::GetTempFileName() + Set-Content -LiteralPath $f -Value ($lines -join "`n") -Encoding ascii + return $f + } + if ($_Migrated) { # 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" @@ -2057,6 +2156,35 @@ exit 0 } } } elseif ($TorchIndexUrl -or $ROCmIndexUrl) { + # Leaf-gated bounded trio, HOISTED so the release-preservation decision below can use it as the + # default route window: torchaudio 2.11 dropped its torch pin, so a bare companion can drift from + # a capped torch. cu families ship torch 2.11.x (paired triton-windows 3.6) so the ceiling + # widens to <2.12; other leaves keep the 2.10 line. Mirrors install.sh and _CUDA_TORCH_PKG_SPEC. + $_idxLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() + if ($_idxLeaf -match '^cu[0-9]+$') { + $_pinTorchSpec = "torch>=2.4,<2.12.0" + $_pinVisionSpec = "torchvision>=0.19,<0.27.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.12.0" + } else { + $_pinTorchSpec = "torch>=2.4,<2.11.0" + $_pinVisionSpec = "torchvision>=0.19,<0.26.0" + $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" + } + # Release preservation (twin of install.sh's _PREV_TORCH_PIN decision): evaluated after every + # index/floor choice, incl. the ROCm reroute, so a raised floor rejects an older release. The + # route window is the leaf-gated CUDA trio by default; the ROCm path uses its floor (or the + # torch>=2.4,<2.11.0 range the ROCm->CPU fallback installs). The kept release is exported for + # setup.ps1 (UNSLOTH_KEPT_TORCH) and cleared after setup runs. + $script:PrevTorchPin = $null + if (-not $SkipTorch -and $script:PrevTorchVer) { + $_routeWindow = $_pinTorchSpec + if ($ROCmIndexUrl) { if ($ROCmTorchFloor) { $_routeWindow = $ROCmTorchFloor } else { $_routeWindow = "torch>=2.4,<2.11.0" } } + $script:PrevTorchPin = Get-PreviousTorchPin -TorchVersion $script:PrevTorchVer -Constraint $_routeWindow + if ($script:PrevTorchPin) { + $env:UNSLOTH_KEPT_TORCH = $script:PrevTorchPin.Release.PublicBase + substep "existing install has torch $script:PrevTorchVer -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)" + } + } if ($SkipTorch) { substep "skipping PyTorch (--no-torch flag set)." "Yellow" } elseif ($ROCmIndexUrl) { @@ -2066,12 +2194,24 @@ exit 0 # Pin companions to match $torchSpec; bare names can resolve an ABI-incompatible torchvision/torchaudio on AMD's per-arch index. $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 } + # Kept-release attempt first (pin already vetted against the ROCm floor); companions follow the kept minor. + if ($script:PrevTorchPin) { + $_keptTorch = $script:PrevTorchPin.TorchSpec; $_keptVision = $script:PrevTorchPin.VisionSpec; $_keptAudio = $script:PrevTorchPin.AudioSpec + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (kept release)" { uv pip install --python $VenvPython --force-reinstall $_keptTorch $_keptVision $_keptAudio --default-index $ROCmIndexUrl } + if ($torchInstallExit -ne 0) { + substep "[WARN] $_keptTorch is not installable from $(Remove-IndexUrlCredentials $ROCmIndexUrl) -- installing the newest supported release instead" "Yellow" + $script:PrevTorchPin = $null + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + } + } + if (-not $script:PrevTorchPin) { + $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 (Unsloth setup retries ROCm). 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 that still satisfies the CPU torch>= range, so without it uv would keep the ROCm build and only swap companions -- a mismatched venv the flavor-repair block won't fix. + # --force-reinstall: a failed ROCm install can leave an unpinned ROCm torch that still satisfies the CPU torch>= range, so without it uv would keep the ROCm build and only swap companions -- a mismatched venv the flavor-repair block won't fix. (No kept-release attempt: the ROCm attempts above always resolve or clear $script:PrevTorchPin first, so a pin never reaches this CPU base.) $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 @@ -2084,18 +2224,19 @@ exit 0 } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..." - # Bounded trio on every index: torchaudio 2.11 dropped its torch pin, so a bare companion can drift from a capped torch. cu families ship torch 2.11.x (paired triton-windows 3.6) so the ceiling widens to <2.12; other leaves keep the 2.10 line. Mirrors install.sh and _CUDA_TORCH_PKG_SPEC. - $_idxLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant() - if ($_idxLeaf -match '^cu[0-9]+$') { - $_pinTorchSpec = "torch>=2.4,<2.12.0" - $_pinVisionSpec = "torchvision>=0.19,<0.27.0" - $_pinAudioSpec = "torchaudio>=2.4,<2.12.0" - } else { - $_pinTorchSpec = "torch>=2.4,<2.11.0" - $_pinVisionSpec = "torchvision>=0.19,<0.26.0" - $_pinAudioSpec = "torchaudio>=2.4,<2.11.0" + # Kept-release attempt first (pin vetted against the leaf-gated route window); companions follow the kept minor. Range install (the hoisted bounded trio) runs when there is no pin or the kept attempt failed. + if ($script:PrevTorchPin) { + $_keptTorch = $script:PrevTorchPin.TorchSpec; $_keptVision = $script:PrevTorchPin.VisionSpec; $_keptAudio = $script:PrevTorchPin.AudioSpec + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (kept release)" { uv pip install --python $VenvPython $_keptTorch $_keptVision $_keptAudio --default-index $TorchIndexUrl } + if ($torchInstallExit -ne 0) { + substep "[WARN] $_keptTorch is not installable from $(Remove-IndexUrlCredentials $TorchIndexUrl) -- installing the newest supported release instead" "Yellow" + $script:PrevTorchPin = $null + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + } + } + if (-not $script:PrevTorchPin) { + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython $_pinTorchSpec $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl } } - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython $_pinTorchSpec $_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) @@ -2118,9 +2259,25 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + # Freeze the installed torch trio so this with-deps resolve can't downgrade the pinned +cuXXX/+rocm build (twin of install.sh's _build_unsloth_torch_overrides). + $script:TorchOverridesFile = New-UnslothTorchOverridesFile -PythonExe $VenvPython + if ($script:TorchOverridesFile) { + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth --overrides $script:TorchOverridesFile "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + Remove-Item -LiteralPath $script:TorchOverridesFile -Force -ErrorAction SilentlyContinue + $script:TorchOverridesFile = $null + } else { + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" } + } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } + # Freeze the installed torch trio (see above) so the with-deps unsloth resolve can't strip the +cuXXX/+rocm suffix. + $script:TorchOverridesFile = New-UnslothTorchOverridesFile -PythonExe $VenvPython + if ($script:TorchOverridesFile) { + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth --overrides $script:TorchOverridesFile -- "$PackageName" } + Remove-Item -LiteralPath $script:TorchOverridesFile -Force -ErrorAction SilentlyContinue + $script:TorchOverridesFile = $null + } else { + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } + } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2184,8 +2341,22 @@ exit 0 # Pin companions like the fresh ROCm path (bare names can pull an ABI-incompatible torchvision/torchaudio from the per-arch index). $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" } + # Kept-release substitution (twin of install.sh's _install_torch_default_index honoring _PREV_TORCH_PIN): honor the preserved torch when the pin survived the E-decision (already floor-vetted); restore the range specs and retry if it isn't installable here. + $_rocmKept = $false + if ($script:PrevTorchPin) { + $_origRocmSpec = $rocmSpec; $_origVisionSpec = $visionSpec; $_origAudioSpec = $audioSpec + $rocmSpec = $script:PrevTorchPin.TorchSpec; $visionSpec = $script:PrevTorchPin.VisionSpec; $audioSpec = $script:PrevTorchPin.AudioSpec + $_rocmKept = $true + } 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 -and $_rocmKept) { + substep "[WARN] $rocmSpec is not installable from $(Remove-IndexUrlCredentials $ROCmIndexUrl) -- installing the newest supported release instead" "Yellow" + $rocmSpec = $_origRocmSpec; $visionSpec = $_origVisionSpec; $audioSpec = $_origAudioSpec + $script:PrevTorchPin = $null + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } + } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) @@ -2199,8 +2370,22 @@ exit 0 } else { $_fixTorchSpec = "torch>=2.4,<2.11.0"; $_fixVisionSpec = "torchvision>=0.19,<0.26.0"; $_fixAudioSpec = "torchaudio>=2.4,<2.11.0" } + # Kept-release substitution (twin of install.sh's _install_torch_default_index honoring _PREV_TORCH_PIN): honor the preserved torch when the pin survived the E-decision; restore the range specs and retry if it isn't installable here. The --reinstall-package triplet stays on both attempts. + $_cudaKept = $false + if ($script:PrevTorchPin) { + $_origFixTorchSpec = $_fixTorchSpec; $_origFixVisionSpec = $_fixVisionSpec; $_origFixAudioSpec = $_fixAudioSpec + $_fixTorchSpec = $script:PrevTorchPin.TorchSpec; $_fixVisionSpec = $script:PrevTorchPin.VisionSpec; $_fixAudioSpec = $script:PrevTorchPin.AudioSpec + $_cudaKept = $true + } substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython $_fixTorchSpec $_fixVisionSpec $_fixAudioSpec --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + if ($torchFixExit -ne 0 -and $_cudaKept) { + substep "[WARN] $_fixTorchSpec is not installable from $(Remove-IndexUrlCredentials $TorchIndexUrl) -- installing the newest supported release instead" "Yellow" + $_fixTorchSpec = $_origFixTorchSpec; $_fixVisionSpec = $_origFixVisionSpec; $_fixAudioSpec = $_origFixAudioSpec + $script:PrevTorchPin = $null + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython $_fixTorchSpec $_fixVisionSpec $_fixAudioSpec --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) @@ -2317,6 +2502,18 @@ exit 0 Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue } + # Release-preservation handoff done: setup.ps1 has consumed UNSLOTH_KEPT_TORCH (if any). Clear it so a + # later 'studio update' in the same session doesn't re-pin an old release, and warn (never abort) if the + # kept torch series changed out from under us during setup. + if ($script:PrevTorchPin) { + $_keptSeries = "$($script:PrevTorchPin.Release.Major).$($script:PrevTorchPin.Release.Minor)" + $_nowVer = Get-InstalledTorchVersionRaw -PythonExe $VenvPython + $_nowRelease = ConvertTo-TorchNumericRelease $_nowVer + if ($_nowRelease -and "$($_nowRelease.Major).$($_nowRelease.Minor)" -ne $_keptSeries) { + Write-Host "[WARN] kept torch $($script:PrevTorchVer) but the environment now has torch $_nowVer" -ForegroundColor Red + } + } + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue if ($setupExit -ne 0) { Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 7892f67bf8..79690fb76f 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2895,13 +2895,36 @@ if ($ROCmIndexUrl) { if ($ROCmTorchSpec -ne "torch") { 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 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host - $torchInstallExit = $LASTEXITCODE - $output = "" - } else { - $output = Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | Out-String - $torchInstallExit = $LASTEXITCODE + # Release preservation (twin of install.sh's _previous_torch_pin): install.ps1 exports the previous + # venv's exact torch RELEASE via UNSLOTH_KEPT_TORCH. Substitute the kept trio unless it conflicts with + # a >=2.11 floor already enforced in $ROCmTorchSpec (a kept minor < 11 keeps the floor). On a failed + # install, restore the computed specs and retry once. Absent the env var (direct 'studio update'), + # nothing changes and the loop runs exactly once. + $_rocmKeptActive = $false + $_rocmOrigTorch = $ROCmTorchSpec; $_rocmOrigVision = $ROCmVisionSpec; $_rocmOrigAudio = $ROCmAudioSpec + if ($env:UNSLOTH_KEPT_TORCH -match '^\d+\.\d+(\.\d+)?$') { + $_keptMinor = [int](($env:UNSLOTH_KEPT_TORCH -split '\.')[1]) + if (-not ($ROCmTorchSpec -match 'torch>=2\.11' -and $_keptMinor -lt 11)) { + $ROCmTorchSpec = "torch==$($env:UNSLOTH_KEPT_TORCH)" + $ROCmVisionSpec = "torchvision==0.$($_keptMinor + 15).*" + $ROCmAudioSpec = "torchaudio==2.$($_keptMinor).*" + $_rocmKeptActive = $true + } + } + while ($true) { + if ($script:UnslothVerbose) { + Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | Out-String + $torchInstallExit = $LASTEXITCODE + } + if ($torchInstallExit -eq 0 -or -not $_rocmKeptActive) { break } + substep "[WARN] torch==$($env:UNSLOTH_KEPT_TORCH) not installable on this ROCm index -- using the supported release" "Yellow" + $ROCmTorchSpec = $_rocmOrigTorch; $ROCmVisionSpec = $_rocmOrigVision; $ROCmAudioSpec = $_rocmOrigAudio + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $_rocmKeptActive = $false } if ($torchInstallExit -ne 0) { Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow @@ -2931,13 +2954,32 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { $cpuVisionSpec = "torchvision>=0.19,<0.27.0" $cpuAudioSpec = "torchaudio>=2.4,<2.12.0" } - if ($script:UnslothVerbose) { - Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host - $torchInstallExit = $LASTEXITCODE - $output = "" - } else { - $output = Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | Out-String - $torchInstallExit = $LASTEXITCODE + # Release preservation: keep install.ps1's exported torch RELEASE (UNSLOTH_KEPT_TORCH) -- a kept + # release from the /cpu index resolves the correct cpu-flavor build. On a failed install, restore the + # computed specs and retry once. Absent the env var, nothing changes and the loop runs exactly once. + $_cpuKeptActive = $false + $_cpuOrigTorch = $cpuTorchSpec; $_cpuOrigVision = $cpuVisionSpec; $_cpuOrigAudio = $cpuAudioSpec + if ($env:UNSLOTH_KEPT_TORCH -match '^\d+\.\d+(\.\d+)?$') { + $_keptMinor = [int](($env:UNSLOTH_KEPT_TORCH -split '\.')[1]) + $cpuTorchSpec = "torch==$($env:UNSLOTH_KEPT_TORCH)" + $cpuVisionSpec = "torchvision==0.$($_keptMinor + 15).*" + $cpuAudioSpec = "torchaudio==2.$($_keptMinor).*" + $_cpuKeptActive = $true + } + while ($true) { + if ($script:UnslothVerbose) { + Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | Out-String + $torchInstallExit = $LASTEXITCODE + } + if ($torchInstallExit -eq 0 -or -not $_cpuKeptActive) { break } + substep "[WARN] torch==$($env:UNSLOTH_KEPT_TORCH) not installable from the CPU index -- using the supported release" "Yellow" + $cpuTorchSpec = $_cpuOrigTorch; $cpuVisionSpec = $_cpuOrigVision; $cpuAudioSpec = $_cpuOrigAudio + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $_cpuKeptActive = $false } if ($torchInstallExit -ne 0) { Write-Host "[FAILED] PyTorch install failed (exit code $torchInstallExit)" -ForegroundColor Red @@ -2961,13 +3003,32 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { $cudaVisionSpec = "torchvision>=0.19,<0.26.0" $cudaAudioSpec = "torchaudio>=2.4,<2.11.0" } - if ($script:UnslothVerbose) { - Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host - $torchInstallExit = $LASTEXITCODE - $output = "" - } else { - $output = Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | Out-String - $torchInstallExit = $LASTEXITCODE + # Release preservation: keep install.ps1's exported torch RELEASE (UNSLOTH_KEPT_TORCH) so a re-run + # reinstalls the same CUDA build (+cuXXX follows this index). On a failed install, restore the computed + # specs and retry once. Absent the env var (direct 'studio update'), nothing changes; loop runs once. + $_cudaKeptActive = $false + $_cudaOrigTorch = $cudaTorchSpec; $_cudaOrigVision = $cudaVisionSpec; $_cudaOrigAudio = $cudaAudioSpec + if ($env:UNSLOTH_KEPT_TORCH -match '^\d+\.\d+(\.\d+)?$') { + $_keptMinor = [int](($env:UNSLOTH_KEPT_TORCH -split '\.')[1]) + $cudaTorchSpec = "torch==$($env:UNSLOTH_KEPT_TORCH)" + $cudaVisionSpec = "torchvision==0.$($_keptMinor + 15).*" + $cudaAudioSpec = "torchaudio==2.$($_keptMinor).*" + $_cudaKeptActive = $true + } + while ($true) { + if ($script:UnslothVerbose) { + Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host + $torchInstallExit = $LASTEXITCODE + $output = "" + } else { + $output = Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | Out-String + $torchInstallExit = $LASTEXITCODE + } + if ($torchInstallExit -eq 0 -or -not $_cudaKeptActive) { break } + substep "[WARN] torch==$($env:UNSLOTH_KEPT_TORCH) not installable from this CUDA index -- using the supported release" "Yellow" + $cudaTorchSpec = $_cudaOrigTorch; $cudaVisionSpec = $_cudaOrigVision; $cudaAudioSpec = $_cudaOrigAudio + Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue + $_cudaKeptActive = $false } if ($torchInstallExit -ne 0) { Write-Host "[FAILED] PyTorch CUDA install failed (exit code $torchInstallExit)" -ForegroundColor Red diff --git a/tests/studio/test_previous_torch_pin.ps1 b/tests/studio/test_previous_torch_pin.ps1 new file mode 100644 index 0000000000..8c84a5e8b4 --- /dev/null +++ b/tests/studio/test_previous_torch_pin.ps1 @@ -0,0 +1,115 @@ +#!/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 install.ps1's torch release preservation helpers +# (ConvertTo-TorchNumericRelease, Test-TorchReleaseInWindow, Get-PreviousTorchPin), +# the Windows port of install.sh's _previous_torch_pin (PR 7250). Pure helpers, +# AST-extracted and run in-process -- no GPU/venv needed. +# Run: pwsh -NoProfile -File tests/studio/test_previous_torch_pin.ps1 + +$ErrorActionPreference = "Stop" +$installPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "install.ps1") +$installPath = (Resolve-Path $installPath).Path + +# --- Parse install.ps1 (also serves as a syntax gate) and extract the helpers --- +$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-TorchNumericRelease", "Test-TorchReleaseInWindow", "Get-PreviousTorchPin")) { + $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 install.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++ } +} + +$savedUpgrade = $env:UNSLOTH_TORCH_UPGRADE +try { + Remove-Item Env:UNSLOTH_TORCH_UPGRADE -ErrorAction SilentlyContinue + $win = "torch>=2.4,<2.12.0" + + # --- ConvertTo-TorchNumericRelease: accepted stable versions --- + foreach ($v in @("2.10.0", "2.10.0+cpu", "2.10.0+cu126", "2.10.1+cu130", "2.9.1+rocm7.2.1", "2.9.0+xpu", "2.10")) { + $r = ConvertTo-TorchNumericRelease $v + Check "release accepts $v" ($null -ne $r) + } + $r = ConvertTo-TorchNumericRelease "2.10.1+cu128" + Check "release strips only the +local tag" ($r.PublicBase -eq "2.10.1" -and $r.Minor -eq 10) + + # --- Rejected: nightly/dev/rc/alpha/garbage never pin --- + foreach ($v in @("", " ", "2.11.0.dev20260701+cu130", "2.9.0a0+gitabc123", "2.11.0rc1+cu130", + ".2.10", "2.10.", "2..10", "2.x.0", "not-a-version", + "Traceback (most recent call last):", "99999999999999999999.1")) { + Check "release rejects '$v'" ($null -eq (ConvertTo-TorchNumericRelease $v)) + } + + # --- Test-TorchReleaseInWindow --- + $cases = @( + @{ v = "2.4.0"; c = "torch>=2.4,<2.12.0"; ok = $true; n = "at the floor" } + @{ v = "2.3.1"; c = "torch>=2.4,<2.12.0"; ok = $false; n = "below the floor" } + @{ v = "2.11.0"; c = "torch>=2.4,<2.12.0"; ok = $true; n = "just below the ceiling" } + @{ v = "2.12.0"; c = "torch>=2.4,<2.12.0"; ok = $false; n = "at the ceiling" } + @{ v = "2.13.0"; c = "torch>=2.4,<2.12.0"; ok = $false; n = "above the ceiling" } + @{ v = "2.10.0"; c = "torch>=2.11.0,<2.12.0"; ok = $false; n = "2.11 floor rejects 2.10" } + @{ v = "2.11.0"; c = "torch>=2.11.0,<2.12.0"; ok = $true; n = "2.11 floor accepts 2.11" } + @{ v = "2.11.0"; c = "torch>=2.4,<2.13.0"; ok = $true; n = "future 2.13 ceiling keeps 2.11" } + @{ v = "2.12.1"; c = "torch>=2.4,<2.13.0"; ok = $true; n = "future 2.13 ceiling keeps 2.12" } + ) + foreach ($t in $cases) { + $rel = ConvertTo-TorchNumericRelease $t.v + Check ("window: " + $t.n) ((Test-TorchReleaseInWindow -Release $rel -Constraint $t.c) -eq $t.ok) + } + # Malformed constraints fail closed. + $rel = ConvertTo-TorchNumericRelease "2.10.0" + Check "window: malformed constraint fails closed" (-not (Test-TorchReleaseInWindow -Release $rel -Constraint "torch")) + Check "window: exact-pin constraint fails closed" (-not (Test-TorchReleaseInWindow -Release $rel -Constraint "torch==2.10.0")) + + # --- Get-PreviousTorchPin: exact-release pin, sh parity --- + $pin = Get-PreviousTorchPin -TorchVersion "2.10.0+cu128" -Constraint $win + Check "pin keeps 2.10.0 (exact release, sh parity)" ($pin.TorchSpec -eq "torch==2.10.0") + Check "pin pairs torchvision to the kept minor" ($pin.VisionSpec -eq "torchvision==0.25.*") + Check "pin pairs torchaudio to the kept minor" ($pin.AudioSpec -eq "torchaudio==2.10.*") + $pin = Get-PreviousTorchPin -TorchVersion "2.9.1+rocm7.2.1" -Constraint $win + Check "pin keeps 2.9.1 with 0.24.*/2.9.* companions" ( + $pin.TorchSpec -eq "torch==2.9.1" -and $pin.VisionSpec -eq "torchvision==0.24.*" -and $pin.AudioSpec -eq "torchaudio==2.9.*") + $pin = Get-PreviousTorchPin -TorchVersion "2.11.0+cpu" -Constraint $win + Check "pin keeps 2.11.0 under a future-widened window" ($pin.TorchSpec -eq "torch==2.11.0") + + # No previous version / out-of-window / non-stable -> no pin. + Check "no pin without a previous version" ($null -eq (Get-PreviousTorchPin -TorchVersion "" -Constraint $win)) + Check "no pin for a below-floor release" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.3.1+cpu" -Constraint $win)) + Check "raised ROCm floor rejects keeping 2.10" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.10.0+rocm7.1" -Constraint "torch>=2.11.0,<2.12.0")) + Check "no pin for a nightly build" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.11.0.dev20260701+cu130" -Constraint $win)) + Check "no pin for an unsupported 2.12 under a <2.12 window" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.12.0+cu130" -Constraint $win)) + + # --- UNSLOTH_TORCH_UPGRADE opt-out (exact string '1', sh parity) --- + $env:UNSLOTH_TORCH_UPGRADE = "1" + Check "UNSLOTH_TORCH_UPGRADE=1 disables the pin" ($null -eq (Get-PreviousTorchPin -TorchVersion "2.10.0+cpu" -Constraint $win)) + $env:UNSLOTH_TORCH_UPGRADE = "0" + Check "UNSLOTH_TORCH_UPGRADE=0 keeps the pin" ($null -ne (Get-PreviousTorchPin -TorchVersion "2.10.0+cpu" -Constraint $win)) +} finally { + if ($null -ne $savedUpgrade) { $env:UNSLOTH_TORCH_UPGRADE = $savedUpgrade } + else { Remove-Item Env:UNSLOTH_TORCH_UPGRADE -ErrorAction SilentlyContinue } +} + +# --- Structural wiring (source assertions) --- +$src = Get-Content $installPath -Raw +Check "probe runs before the rollback move" ( + $src.IndexOf('$script:PrevTorchVer') -ge 0 -and + $src.IndexOf('$script:PrevTorchVer') -lt $src.IndexOf('Start-StudioVenvRollback -ExistingDir')) +Check "pin decision cites the UNSLOTH_TORCH_UPGRADE escape hatch" ($src -match 'UNSLOTH_TORCH_UPGRADE=1 to get the newest') +Check "kept-release fallback clears the pin" ($src -match '\$script:PrevTorchPin\s*=\s*\$null') +Check "kept release exported for setup.ps1" ($src -match 'UNSLOTH_KEPT_TORCH') + +Write-Host "" +if ($failures -gt 0) { Write-Host "$failures check(s) failed" -ForegroundColor Red; exit 1 } +Write-Host "All checks passed" +exit 0