install: honor torch-index override on the Windows installers too
The pinned-index work landed for install.sh and install_python_stack.py, but the Windows installers still picked the wheel index from GPU probing. Extend the same UNSLOTH_TORCH_INDEX_URL / _FAMILY contract so a pinned index wins on every platform: - install.ps1: Get-TorchIndexUrl returns the pinned URL/family before nvidia-smi probing; the AMD ROCm reroute is skipped when the index is pinned, so an explicit cpu/cu* pin on an AMD host is not overwritten. - studio/setup.ps1: add shared Get-PinnedTorchIndexUrl / Get-TorchIndexLeaf helpers; the stale-venv check, the install selection and the AMD reroute all honor the pin, and the CPU/CUDA install pulls from the resolved index URL. - tests: parity test that all four installers read both override vars and the two Windows installers gate the AMD reroute on the pinned flag.
This commit is contained in:
Daniel Han
parent
5d69193e1e
commit
5736d5da96
3 changed files with 89 additions and 8 deletions
18
install.ps1
18
install.ps1
|
|
@ -1939,6 +1939,17 @@ exit 0
|
|||
# Mirrors Get-PytorchCudaTag in setup.ps1.
|
||||
function Get-TorchIndexUrl {
|
||||
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
|
||||
# Explicit pin -- skip ALL GPU probing when the caller names the wheel index
|
||||
# (headless / CI / cross-install). Matches install.sh::get_torch_index_url and
|
||||
# install_python_stack.py: UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim);
|
||||
# UNSLOTH_TORCH_INDEX_FAMILY is the convenience leaf (cpu, cu128, rocm6.4, ...)
|
||||
# appended to the mirror base so UNSLOTH_PYTORCH_MIRROR is still honoured.
|
||||
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
|
||||
return $env:UNSLOTH_TORCH_INDEX_URL.Trim().TrimEnd('/')
|
||||
}
|
||||
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
|
||||
return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
|
||||
}
|
||||
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
|
||||
try {
|
||||
$output = Invoke-NvidiaSmiBounded $NvidiaSmiExe
|
||||
|
|
@ -2016,6 +2027,11 @@ exit 0
|
|||
} catch { return $null }
|
||||
}
|
||||
|
||||
# An explicit UNSLOTH_TORCH_INDEX_URL / _FAMILY pin is authoritative: the AMD
|
||||
# ROCm reroute below must not rewrite it (e.g. a deliberate cpu pin on an AMD
|
||||
# host, or a pinned ROCm family we already resolved in Get-TorchIndexUrl).
|
||||
$TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or `
|
||||
(-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY))
|
||||
$TorchIndexUrl = Get-TorchIndexUrl
|
||||
|
||||
# ── GPU arch → newest compatible Windows ROCm wheel release ──
|
||||
|
|
@ -2027,7 +2043,7 @@ 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) {
|
||||
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
|
||||
|
|
|
|||
|
|
@ -402,6 +402,29 @@ function Get-PytorchCudaTag {
|
|||
return "cu126"
|
||||
}
|
||||
|
||||
# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY), shared by the
|
||||
# stale-venv check and the install selection below so a pinned wheel index wins
|
||||
# over GPU probing -- matching install.sh, install.ps1 and install_python_stack.py.
|
||||
# UNSLOTH_TORCH_INDEX_URL is verbatim (full URL); _FAMILY is the leaf (cpu, cu128,
|
||||
# rocm6.4, ...) joined to the mirror base so UNSLOTH_PYTORCH_MIRROR is honoured.
|
||||
function Get-PinnedTorchIndexUrl {
|
||||
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
|
||||
return $env:UNSLOTH_TORCH_INDEX_URL.Trim().TrimEnd('/')
|
||||
}
|
||||
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
|
||||
$base = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
|
||||
return "$base/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
|
||||
}
|
||||
return $null
|
||||
}
|
||||
|
||||
# The last path segment of a wheel index URL (cu128 / cpu / rocm6.4 / gfx1151).
|
||||
function Get-TorchIndexLeaf {
|
||||
param([string]$Url)
|
||||
if ([string]::IsNullOrWhiteSpace($Url)) { return $null }
|
||||
return ($Url.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
|
||||
}
|
||||
|
||||
# VS generator -> MSBuild BuildCustomizations dir; toolset tracks the VS major
|
||||
# (18->v180, 17->v170), defaulting to v170 when unparseable.
|
||||
function Get-VcBuildCustomizationsDir {
|
||||
|
|
@ -2536,7 +2559,8 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
|
|||
}
|
||||
|
||||
if (-not $shouldRebuild) {
|
||||
$expectedTorchTag = if ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" }
|
||||
$_pinnedIdx = Get-PinnedTorchIndexUrl
|
||||
$expectedTorchTag = if ($_pinnedIdx) { Get-TorchIndexLeaf $_pinnedIdx } elseif ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" }
|
||||
if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
|
||||
$shouldRebuild = $true
|
||||
}
|
||||
|
|
@ -2712,7 +2736,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"
|
||||
|
|
@ -2733,7 +2763,7 @@ $ROCmIndexUrl = $null
|
|||
# SDK -- which flips Studio out of chat-only (CHAT_ONLY) and enables Train/Export.
|
||||
# Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed
|
||||
# ROCm install still falls back to CPU below, so this is safe.
|
||||
if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
|
||||
if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
|
||||
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
|
||||
$archFamilyMap = @{
|
||||
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
|
||||
|
|
@ -2785,6 +2815,11 @@ if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
|
|||
|
||||
$PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
|
||||
|
||||
# A full UNSLOTH_TORCH_INDEX_URL pin is used verbatim; a family pin already set
|
||||
# $CuTag, so $PyTorchWhlBase/$CuTag is the requested family index. The CPU/CUDA
|
||||
# install branches below pull from this instead of re-joining mirror + tag.
|
||||
$TorchInstallIndexUrl = if ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" }
|
||||
|
||||
$ROCmCpuFallback = $false
|
||||
if ($ROCmIndexUrl) {
|
||||
substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..."
|
||||
|
|
@ -2821,11 +2856,11 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
|
|||
$cpuForce = @()
|
||||
if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") }
|
||||
if ($script:UnslothVerbose) {
|
||||
Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu"
|
||||
Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl
|
||||
$torchInstallExit = $LASTEXITCODE
|
||||
$output = ""
|
||||
} else {
|
||||
$output = Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" | Out-String
|
||||
$output = Fast-Install torch torchvision torchaudio @cpuForce --index-url $TorchInstallIndexUrl | Out-String
|
||||
$torchInstallExit = $LASTEXITCODE
|
||||
}
|
||||
if ($torchInstallExit -ne 0) {
|
||||
|
|
@ -2837,11 +2872,11 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
|
|||
substep "installing PyTorch with CUDA support ($CuTag)..."
|
||||
substep "(This download is ~2.8 GB -- may take a few minutes)"
|
||||
if ($script:UnslothVerbose) {
|
||||
Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag"
|
||||
Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl
|
||||
$torchInstallExit = $LASTEXITCODE
|
||||
$output = ""
|
||||
} else {
|
||||
$output = Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" | Out-String
|
||||
$output = Fast-Install torch torchvision torchaudio --index-url $TorchInstallIndexUrl | Out-String
|
||||
$torchInstallExit = $LASTEXITCODE
|
||||
}
|
||||
if ($torchInstallExit -ne 0) {
|
||||
|
|
|
|||
|
|
@ -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,31 @@ 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"
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue