From 5736d5da9604838f476fa4a5eed67536ec941590 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 26 Jun 2026 08:33:48 +0000 Subject: [PATCH] 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. --- install.ps1 | 18 +++++++- studio/setup.ps1 | 49 ++++++++++++++++++---- tests/python/test_cross_platform_parity.py | 30 +++++++++++++ 3 files changed, 89 insertions(+), 8 deletions(-) diff --git a/install.ps1 b/install.ps1 index f7f9540970..2a929c81d5 100644 --- a/install.ps1 +++ b/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 diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 38398eb5f5..907367a69a 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -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) { diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 666ea7ce10..0fae0c3255 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,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" + )