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 2026-06-26 08:33:48 +00:00
commit 5736d5da96
3 changed files with 89 additions and 8 deletions

View file

@ -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"
)