From 1509023d0ccbe21bc18ce04679efc06cf6e9d6d7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 10:58:14 +0000 Subject: [PATCH] install: extend the pinned-index guards to every remaining surface 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). --- install.ps1 | 2 +- install.sh | 6 ++-- studio/install_python_stack.py | 5 +++ studio/setup.ps1 | 9 ++++-- tests/python/test_cross_platform_parity.py | 36 ++++++++++++++++++++++ tests/python/test_install_python_stack.py | 7 +++++ 6 files changed, 60 insertions(+), 5 deletions(-) diff --git a/install.ps1 b/install.ps1 index 71f120a86a..16416dd9ae 100644 --- a/install.ps1 +++ b/install.ps1 @@ -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 } diff --git a/install.sh b/install.sh index fc571e4917..bbc7e52077 100755 --- a/install.sh +++ b/install.sh @@ -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 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 599f71c9da..d88cdfab13 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -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", ) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 5b782e3fc6..40c1c47397 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -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 } diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 0b6162dcff..8d0be19e64 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -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" + ) + diff --git a/tests/python/test_install_python_stack.py b/tests/python/test_install_python_stack.py index b47276e173..88febe37cf 100644 --- a/tests/python/test_install_python_stack.py +++ b/tests/python/test_install_python_stack.py @@ -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."""