ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

install: extend the pinned-index guards to every remaining surface

Browse source 
Round 3 of the pinned-index hardening, closing the same holes on the
surfaces the earlier rounds missed:

1. install.sh's pinned-install env scrub now clears UV_TORCH_BACKEND (uv's
   torch backend redirects torch resolution to its own per-backend index
   even against --default-index), and both PowerShell wrappers clear it in
   their pinned-install scrubs, matching install_python_stack.py.

2. setup.ps1's marker stale check still classified any rocm* leaf as a
   PyTorch ROCm family while the install selection is digit-gated, so a
   custom rocm-current / rocm-rel-7.2.1 pin stale-compared as
   not-rocm vs rocm and force-reinstalled on every studio update. The
   stale check now uses the same ^rocm\d gate.

3. install_python_stack.py's pinned-command scrub also strips
   PIP_EXTRA_INDEX_URL for the pip fallback: pip adds the env extra index
   in addition to --index-url, so an inherited mirror could satisfy torch
   off the pin while the marker recorded the pinned URL. PIP_INDEX_URL
   needs no strip since the explicit --index-url flag overrides it.

Parity + unit tests extended (4 new tests).
This commit is contained in:
Daniel Han 2026-07-12 10:58:14 +00:00
commit 1509023d0c
6 changed files with 60 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

2
install.ps1
View file

@ -479,7 +479,7 @@ function Install-UnslothStudio {
$savedUvIndex = $null
if ($Command.ToString() -match '--default-index') {
$savedUvIndex = @{}
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND') {
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}

6
install.sh
View file

@ -161,9 +161,11 @@ run_install_cmd() {
shift
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when we pass --default-index, neutralize every uv index env var so
# the pinned index wins. Other installs keep the user's mirror.
# the pinned index wins. UV_TORCH_BACKEND is cleared too: uv's torch backend
# redirects torch resolution to its own per-backend index even against a
# --default-index pin. Other installs keep the user's mirror and backend.
case " $* " in
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL -u UV_TORCH_BACKEND "$@" ;;
esac
if _is_verbose; then
"$@" && return 0

5
studio/install_python_stack.py
View file

@ -2584,12 +2584,17 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
# to its own per-backend index even when a --index-url is given, so it defeats a
# pin the same way ( _build_uv_cmd also refuses to turn it into a --torch-backend
# flag for pinned commands; this covers uv reading the env var directly).
# PIP_EXTRA_INDEX_URL matters for the pip FALLBACK (uv missing/failed): pip's
# --extra-index-url env form adds indexes "in addition to --index-url", so an
# inherited mirror could still satisfy torch off the pin. (PIP_INDEX_URL needs no
# strip: the command's explicit --index-url flag overrides it.)
_UV_INDEX_ENV_VARS = (
"UV_DEFAULT_INDEX",
"UV_INDEX_URL",
"UV_INDEX",
"UV_EXTRA_INDEX_URL",
"UV_TORCH_BACKEND",
"PIP_EXTRA_INDEX_URL",
)

9
studio/setup.ps1
View file

@ -2770,7 +2770,12 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
$_markerMismatch = Test-MarkerPinMismatch -VenvDir $VenvDir -PinUrl $_pinnedIdx
if ($_markerMismatch -eq $true) { $shouldRebuild = $true }
# cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds.
if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') {
# Digit-gated like the install selection below: a custom rocm-* leaf
# (rocm-current / rocm-rel-7.2.1) is NOT a PyTorch ROCm family, stays on
# the verbatim/marker path, and must not be stale-compared as one (that
# reported not-rocm vs rocm(torch<2.11) and force-reinstalled on every
# studio update even though the pin never changed).
if ($_pinLeaf -like 'gfx*' -or $_pinLeaf -match '^rocm\d') {
# Do NOT collapse a pinned ROCm/gfx leaf to a generic "rocm": that
# would match any installed +rocm wheel and mask a pin change from
# one ROCm family to another (e.g. rocm6.4 -> gfx1151, or rocm6.4
@ -2940,7 +2945,7 @@ function Fast-Install {
# them only for index-pinned installs; mirrors still apply elsewhere.
$saved = @{}
if (@($Args_) -contains '--index-url') {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND') {
$saved[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}

36
tests/python/test_cross_platform_parity.py
View file

@ -542,3 +542,39 @@ class TestPinnedIndexClearsUvEnvParity:
)
for var in self.UV_VARS:
assert var in text, f"install_python_stack.py must clear {var} for pinned installs"
def test_all_installers_clear_uv_torch_backend(self):
"""uv's torch backend redirects torch resolution to its own per-backend
index even against an explicit pin, so every installer's pinned-install
scrub must clear UV_TORCH_BACKEND too."""
sh = INSTALL_SH.read_text(encoding = "utf-8")
assert "-u UV_TORCH_BACKEND" in sh, "install.sh pinned scrub must clear UV_TORCH_BACKEND"
for path in (INSTALL_PS1, SETUP_PS1):
text = path.read_text(encoding = "utf-8")
assert "'UV_TORCH_BACKEND'" in text, f"{path.name} pinned scrub must clear UV_TORCH_BACKEND"
stack = STACK_PY.read_text(encoding = "utf-8")
assert '"UV_TORCH_BACKEND",' in stack, (
"install_python_stack.py strip tuple must include UV_TORCH_BACKEND"
)
def test_stack_py_strips_pip_extra_index_for_pip_fallback(self):
"""The pip fallback honours PIP_EXTRA_INDEX_URL (pip adds it IN ADDITION
to --index-url), so the pinned-command scrub must strip it."""
stack = STACK_PY.read_text(encoding = "utf-8")
assert '"PIP_EXTRA_INDEX_URL",' in stack, (
"install_python_stack.py strip tuple must include PIP_EXTRA_INDEX_URL"
)
def test_setup_ps1_stale_check_requires_rocm_digit(self):
"""The marker stale check must use the same rocm+digit gate as the
install selection, or a custom rocm-* leaf force-reinstalls on every
studio update."""
text = SETUP_PS1.read_text(encoding = "utf-8")
stale = text[text.find("Get-RocmPinStaleTags -PinLeaf") - 2500 :][:2500]
assert "-match '^rocm\\d'" in stale.replace("\\", "\\"), (
"setup.ps1 stale check must digit-gate rocm leaves"
)
assert stale.count("-like 'rocm*'") == 0, (
"setup.ps1 stale check must not use a bare -like 'rocm*' glob"
)

7
tests/python/test_install_python_stack.py
View file

@ -235,6 +235,13 @@ class TestPinnedIndexClearsUvEnv:
assert env is not None
assert env.get("PATH_SENTINEL_XYZ") == "keepme", "only uv index vars are removed"
def test_pinned_cmd_strips_pip_extra_index_url(self):
"""PIP_EXTRA_INDEX_URL is stripped for pinned commands so the pip
fallback cannot satisfy torch from an inherited extra index."""
with mock.patch.dict(os.environ, {"PIP_EXTRA_INDEX_URL": "https://mirror/simple"}):
env = ips._install_env_for_cmd(["pip", "install", "torch", "--index-url", "https://x/cu128"])
assert env is not None and "PIP_EXTRA_INDEX_URL" not in env
def test_pinned_cmd_strips_uv_torch_backend(self):
"""UV_TORCH_BACKEND is stripped for pinned commands so uv cannot read it
from the environment and reroute torch off the pinned index."""