From 5c93ffd4501b52314f151be640e831bf353a2164 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Mon, 6 Jul 2026 05:33:36 +0000 Subject: [PATCH] install: honor custom pins and repair pinned venvs in place Four follow-ups to the torch-index marker work: - install_python_stack.py: _ensure_cuda_torch/_ensure_rocm_torch now bail when an explicit custom-index pin names no known torch family, so a verbatim URL override (a private/simple mirror) is not clobbered by auto-detected CUDA/ROCm wheels before _ensure_verbatim_torch_index applies it. - install_python_stack.py: the ROCm marker is additive, not a substitute -- a matching marker still runs the family/version check so a wheel swapped after the marker was written is caught. Mirrors setup.ps1. - setup.ps1: a stale venv under an explicit pin, whose torch still imports, is repaired in place (force-reinstall torch from the pin in the dependency pass) instead of wiped. The wipe path only delegates to install.ps1, so on a direct update it stranded the user at "Virtual environment not found" instead of applying the new pin. A broken venv or unpinned drift still wipes/delegates. - install.ps1: when a pinned ROCm install fails over to a CPU base, the marker now records the CPU index actually used instead of the ROCm pin, so the next managed setup does not see CPU torch under a ROCm pin and abort as stale. --- install.ps1 | 13 +++++++++++-- studio/install_python_stack.py | 28 +++++++++++++++++++++++++--- studio/setup.ps1 | 34 +++++++++++++++++++++++++++++++--- 3 files changed, 67 insertions(+), 8 deletions(-) diff --git a/install.ps1 b/install.ps1 index fc005e2c9a..148672b7cc 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2080,6 +2080,11 @@ exit 0 # Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs. $ROCmIndexUrl = $null $ROCmTorchFloor = $null + # Set when a pinned ROCm install fails and a CPU base is installed instead; the + # marker must then record the CPU index actually used, not the ROCm pin left in + # $TorchIndexUrl (else the next managed setup sees CPU torch under a ROCm pin and + # aborts as stale rather than continuing with the intended CPU base/retry path). + $RocmCpuFallbackIndexUrl = $null $PinnedRocmVisionSpec = $null $PinnedRocmAudioSpec = $null if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) { @@ -2305,6 +2310,9 @@ exit 0 # reinstalls ROCm afterwards (recomputes its own index URL). $ROCmIndexUrl = $null $ROCmTorchFloor = $null + # Record the CPU index actually installed from, so the marker below + # reflects the CPU base rather than the ROCm pin still in $TorchIndexUrl. + $RocmCpuFallbackIndexUrl = $CpuFallbackIndexUrl } } else { Write-TauriLog "STEP" "Installing PyTorch" @@ -2439,10 +2447,11 @@ exit 0 # Torch is now resolved; write the exact --index-url used so setup.ps1 / # install_python_stack.py can detect a later pin change by an EXACT string # compare instead of the version-tag heuristic. Reflects the installed family: - # $ROCmIndexUrl when the ROCm path ran, else the CUDA/CPU/pinned $TorchIndexUrl. + # $ROCmIndexUrl when the ROCm path ran, the CPU fallback index when a pinned + # ROCm install failed over to a CPU base, else the CUDA/CPU/pinned $TorchIndexUrl. # Skipped for --no-torch (nothing installed). Matches install.sh / setup.ps1. if (-not $SkipTorch) { - $MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } else { $TorchIndexUrl } + $MarkerIndexUrl = if ($ROCmIndexUrl) { $ROCmIndexUrl } elseif ($RocmCpuFallbackIndexUrl) { $RocmCpuFallbackIndexUrl } else { $TorchIndexUrl } Write-TorchIndexMarker -VenvDir $VenvDir -IndexUrl $MarkerIndexUrl } diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 90aac9533b..d5b0aa15af 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1414,6 +1414,11 @@ def _ensure_cuda_torch() -> None: # (or any unrecognised value) are deliberate and must not be overridden. if _TORCH_BACKEND not in ("", "cuda"): return + # An explicit custom-index pin whose leaf names no known torch family wins + # VERBATIM (_ensure_verbatim_torch_index applies it); do not override it with + # auto-detected CUDA wheels here. + if _explicit_unknown_family_torch_index_url() is not None: + 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. if IS_MACOS or IS_WINDOWS or NO_TORCH: @@ -1611,6 +1616,14 @@ def _ensure_rocm_torch() -> None: # different environment (different PATH, CUDA_VISIBLE_DEVICES, etc.). if _TORCH_BACKEND in ("cuda", "cpu"): return + # An explicit custom-index pin whose leaf names no known torch family (a private + # PEP 503 mirror, /simple, /current, ...) wins VERBATIM -- _ensure_verbatim_torch_index() + # applies it. Never override such a pin with the auto-detected ROCm index here: the + # user chose that index (a valid cross-install / custom-mirror case), and installing + # ROCm wheels over it before the verbatim pass runs would leave the URL override + # unhonored (the marker then matches the pin, so verbatim skips the restore). + if _explicit_unknown_family_torch_index_url() is not None: + 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. @@ -1823,10 +1836,19 @@ def _ensure_rocm_torch() -> None: _rocm_pin_mismatch = False if has_hip_torch and _rocm_pin is not None: _marker_verdict = _marker_pin_mismatch(_rocm_pin) - if _marker_verdict is None: - _rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver) + if _marker_verdict is True: + # Marker records a DIFFERENT index than the pin -> reinstall. This is the + # only signal that catches a per-arch gfx switch (gfx1151 -> gfx120X-all, + # both +rocm7.13.0) the version-tag heuristic below cannot see. + _rocm_pin_mismatch = True else: - _rocm_pin_mismatch = _marker_verdict + # Marker matches OR is absent: the marker is an ADDITIONAL rebuild signal, + # not a substitute for validating the installed wheel. Still run the + # family/version check so a stale wheel (torch swapped after the marker was + # written -- e.g. marker records gfx1151 but the venv now carries generic + # +rocm7.2 or an older +rocm6.4) is caught. Mirrors setup.ps1, which keeps + # the flavor check alongside the marker for this reason. + _rocm_pin_mismatch = _rocm_pin_family_mismatch(_rocm_pin, _installed_torch_ver) rocm_torch_ready = has_hip_torch and not _rocm_pin_mismatch diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 7b0d09255a..daa045f498 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2700,6 +2700,9 @@ 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 an explicit torch-index pin is repaired in place + # (force-reinstall torch from the pin) instead of wiped -- see the rebuild block. + $script:PinChangedForceReinstall = $false if (Test-Path -LiteralPath $VenvPyExe) { try { @@ -2830,6 +2833,19 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode } } + # A stale venv under an explicit torch-index pin, whose torch still imports, is + # repaired IN PLACE: the dependency pass below force-reinstalls torch from the pin + # instead of wiping. setup.ps1's rebuild path only removes the venv and delegates + # to install.ps1, so on a direct `unsloth studio update` a wipe would strand the + # user at "Virtual environment not found" rather than applying the new pin. Only a + # broken venv (torch unimportable) or an unpinned GPU-detected drift falls through + # to the wipe/delegate path below. + 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) { @@ -2995,6 +3011,11 @@ 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 (the fast "up to date" path would +# otherwise skip it and leave the old wheel under the changed pin). +if ($script:PinChangedForceReinstall) { $SkipPythonDeps = $false } + if (-not $SkipPythonDeps) { if ($script:UnslothVerbose) { @@ -3184,7 +3205,9 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { # Build the array directly: an if-expression collapses @("x") to a scalar string, # which @splat would then enumerate char-by-char into broken single-letter args. $cpuForce = @() - if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") } + # --force-reinstall also 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 ($ROCmCpuFallback -or $script:PinChangedForceReinstall) { $cpuForce = @("--force-reinstall") } if ($script:UnslothVerbose) { Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl $torchInstallExit = $LASTEXITCODE @@ -3201,12 +3224,17 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { } 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 local tag), so without it uv would + # keep the old build and the changed CUDA pin (e.g. cu126 -> cu128) never applies. + $cudaForce = @() + if ($script:PinChangedForceReinstall) { $cudaForce = @("--force-reinstall") } if ($script:UnslothVerbose) { - Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl + Fast-Install torch torchvision torchaudio @cudaForce --index-url $TorchInstallIndexUrl $torchInstallExit = $LASTEXITCODE $output = "" } else { - $output = Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl | Out-String + $output = Fast-Install torch torchvision torchaudio @cudaForce --index-url $TorchInstallIndexUrl | Out-String $torchInstallExit = $LASTEXITCODE } if ($torchInstallExit -ne 0) {