mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-30 18:14:02 +00:00
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.
This commit is contained in:
danielhanchen
parent
257ca41753
commit
5c93ffd450
3 changed files with 67 additions and 8 deletions
13
install.ps1
13
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
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue