From a98203b5398cc508e868d65489835176e185d93d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 10:17:01 +0000 Subject: [PATCH] install: keep pinned torch repairs on the pinned index Two fixes for explicit index pins (UNSLOTH_TORCH_INDEX_FAMILY / _URL): 1. install_python_stack.py's repair paths ran uv without clearing the inherited uv index env vars. uv resolves the default index (--index-url or --default-index) at the LOWEST priority, so a UV_INDEX or UV_EXTRA_INDEX_URL mirror in the environment won for any package it served: a cu128-pinned repair could install torch from the mirror and then record the cu128 marker it never used. Verified empirically: with UV_EXTRA_INDEX_URL=.../cu126 exported, uv pip install torch --index-url .../cu128 resolves torch 2.13.0+cu126. Strip the four uv index env vars for pinned-index commands only, mirroring the gate install.sh, install.ps1 and setup.ps1 already have; non-pinned installs keep the user's mirror. 2. install.ps1 routed any pinned leaf matching rocm* through the ROCm --default-index path, so a custom find-links leaf like rocm-rel-7.2.1 was treated as a PEP 503 ROCm index and could silently fall back to CPU torch on resolution failure. Require a digit after rocm, matching install.sh's rocm[0-9]* and install_python_stack.py's ^rocm\d. Adds parity + unit tests for both (11 new tests). --- install.ps1 | 9 ++- studio/install_python_stack.py | 33 ++++++++++ tests/python/test_cross_platform_parity.py | 72 ++++++++++++++++++++++ tests/python/test_install_python_stack.py | 62 +++++++++++++++++++ 4 files changed, 175 insertions(+), 1 deletion(-) diff --git a/install.ps1 b/install.ps1 index 4618029111..71f120a86a 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2177,10 +2177,17 @@ exit 0 $PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0" $PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0" substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan" - } elseif ($_pinLeaf -like 'gfx*' -or $_pinLeaf -like 'rocm*') { + } elseif ($_pinLeaf -like 'gfx*' -or $_pinLeaf -match '^rocm\d') { # Other gfx per-arch indexes and older rocm (<=7.1) ship torch <2.11; # route via the ROCm path with bare specs so the pinned family resolves # its own wheels (matches the automatic path's bare floor for these). + # Require a DIGIT after rocm (rocm7.1, not rocm-rel-7.2.1 / rocm-current): + # only rocm and gfx* are pip --default-index (PEP 503) families. + # A rocm- leaf is a custom/find-links URL that must fall through + # to the verbatim --default-index install below, not this ROCm path (which + # silently falls back to CPU on failure). Mirrors install.sh's rocm[0-9]* + # (_torch_index_repairable) and install_python_stack.py's + # _is_pip_rocm_family_leaf (^rocm\d). $ROCmIndexUrl = $TorchIndexUrl } } diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 3146f6f6fe..6e8474176b 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -2368,6 +2368,7 @@ def run( cmd, stdout = subprocess.PIPE if quiet else None, stderr = subprocess.STDOUT if quiet else None, + env = _install_env_for_cmd(cmd), **_windows_hidden_subprocess_kwargs(), ) if result.returncode != 0: @@ -2566,6 +2567,36 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]: return cmd +# uv resolves the default index (passed as --index-url / --default-index) at the +# LOWEST priority: an inherited UV_INDEX / UV_EXTRA_INDEX_URL (a corporate or CPU +# mirror) is searched FIRST and, under uv's default first-index strategy, wins for +# any package it also serves. So a pinned torch repair (--index-url ) +# could silently resolve torch from that mirror instead of the pinned wheel index, +# and then _write_torch_index_marker() would record the pinned URL that was never +# actually used. install.sh (run_install_cmd, #6898), install.ps1 and setup.ps1 +# already neutralise these vars for their pinned installs; do the same here so the +# pin wins on every platform. Non-pinned installs (no --index-url) keep the mirror. +_UV_INDEX_ENV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL") + + +def _install_env_for_cmd(cmd: "list[str]") -> "dict[str, str] | None": + """Return an env with the uv index vars stripped for a pinned-index install. + + Returns None (inherit the caller's environment unchanged) when the command + does NOT pin an index, so ordinary installs still honour a user's UV_INDEX / + UV_EXTRA_INDEX_URL mirror. When the command passes --index-url / --default-index + (the torch repair paths), the four uv index env vars are removed so the pinned + index is not overridden by an inherited mirror (uv treats the default index as + lowest priority). Mirrors install.sh's run_install_cmd gate (#6898). + """ + if not any(arg in ("--index-url", "--default-index") for arg in cmd): + return None + env = os.environ.copy() + for name in _UV_INDEX_ENV_VARS: + env.pop(name, None) + return env + + def pip_install_try( label: str, *args: str, @@ -2592,6 +2623,7 @@ def pip_install_try( cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, + env = _install_env_for_cmd(cmd), ) if result.returncode == 0: return True @@ -2644,6 +2676,7 @@ def pip_install( uv_cmd, stdout = subprocess.PIPE, stderr = subprocess.STDOUT, + env = _install_env_for_cmd(uv_cmd), **_windows_hidden_subprocess_kwargs(), ) if result.returncode == 0: diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index d3c1f02709..1262b20af9 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -461,3 +461,75 @@ class TestKnown211SetParity: low = path.read_text(encoding = "utf-8").lower() for g in gfx: assert g in low, f"{label} missing gfx 2.11 allowlist member {g}" + + +class TestPinnedRocmLeafDigitParity: + """A pinned index is a pip ROCm --default-index family only when its leaf is + rocm+digit (rocm7.1 / rocm7.2) or gfx*. A bare `rocm*` glob wrongly catches a + custom mirror / Radeon find-links leaf (rocm-current / rocm-rel-7.2.1) and routes + it through the ROCm install path (which silently falls back to CPU on failure) + instead of the verbatim --default-index install. install.sh (_torch_index_repairable) + and install_python_stack.py (_is_pip_rocm_family_leaf) already require a digit + after rocm; install.ps1's pinned reroute must match.""" + + def test_install_ps1_pinned_reroute_requires_rocm_digit(self): + text = INSTALL_PS1.read_text(encoding = "utf-8") + # The pinned gfx*/rocm reroute must require a digit after rocm. + assert re.search( + r"\$_pinLeaf -like 'gfx\*' -or \$_pinLeaf -match '\^rocm\\d'", text + ), ( + "install.ps1 pinned-index reroute must use -match '^rocm\\d' (not a bare " + "-like 'rocm*'), so rocm-current / rocm-rel-* fall through to the verbatim " + "install instead of the ROCm --default-index path" + ) + # The broad glob must be gone from that reroute. + assert "-like 'rocm*'" not in text, ( + "install.ps1 must not route a pinned index on a bare -like 'rocm*' glob" + ) + + def test_install_sh_repairable_requires_rocm_digit(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + assert re.search( + r"cu\[0-9\]\*\|rocm\[0-9\]\*\|gfx\*", text + ), "install.sh _torch_index_repairable must require rocm[0-9]* (a digit after rocm)" + + def test_stack_py_pip_rocm_family_requires_digit(self): + text = STACK_PY.read_text(encoding = "utf-8") + assert re.search( + r'r"\^rocm\\d"', text + ), "install_python_stack.py _is_pip_rocm_family_leaf must match ^rocm\\d" + + +class TestPinnedIndexClearsUvEnvParity: + """Every installer must neutralise the uv index env vars for a pinned torch + install (#6898). uv treats the default index (--index-url / --default-index) as + lowest priority, so an inherited UV_INDEX / UV_EXTRA_INDEX_URL mirror would win + under uv's first-index strategy and pull torch from the wrong index -- after + which the torch-index marker records a wheel index that was never used.""" + + UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL") + + def test_install_sh_clears_uv_index_vars(self): + text = INSTALL_SH.read_text(encoding = "utf-8") + assert ( + "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in text + ), "install.sh run_install_cmd must clear the uv index vars for --default-index installs" + + def test_install_ps1_clears_uv_index_vars(self): + text = INSTALL_PS1.read_text(encoding = "utf-8") + for var in self.UV_VARS: + assert var in text, f"install.ps1 must clear {var} for pinned installs" + + def test_setup_ps1_clears_uv_index_vars(self): + text = SETUP_PS1.read_text(encoding = "utf-8") + for var in self.UV_VARS: + assert var in text, f"setup.ps1 must clear {var} for pinned installs" + + def test_stack_py_clears_uv_index_vars(self): + text = STACK_PY.read_text(encoding = "utf-8") + assert "_install_env_for_cmd" in text, ( + "install_python_stack.py must scrub inherited uv index vars for pinned " + "installs via _install_env_for_cmd (parity with install.sh #6898)" + ) + for var in self.UV_VARS: + assert var in text, f"install_python_stack.py must clear {var} for pinned installs" diff --git a/tests/python/test_install_python_stack.py b/tests/python/test_install_python_stack.py index 9015ff8c9d..d2ebf51eee 100644 --- a/tests/python/test_install_python_stack.py +++ b/tests/python/test_install_python_stack.py @@ -148,3 +148,65 @@ class TestUvSafePathHardening: assert " " not in value assert Path(value).read_text() == "transformers>=4.57.6\n" + + +class TestPinnedIndexClearsUvEnv: + """A pinned torch install (--index-url / --default-index) must neutralise an + inherited UV_INDEX / UV_EXTRA_INDEX_URL so the pinned wheel index wins. + + uv treats the default index (--index-url / --default-index) as LOWEST priority, + so an inherited UV_INDEX / UV_EXTRA_INDEX_URL (a corporate/CPU mirror) would be + searched first and, under uv's default first-index strategy, resolve torch from + the wrong mirror -- after which the marker records a wheel index that was never + used. install.sh (#6898), install.ps1 and setup.ps1 already clear these for + pinned installs; install_python_stack must match (parity across all installers). + """ + + UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL") + + def test_pinned_index_url_strips_uv_index_vars(self): + cmd = [ + "uv", "pip", "install", "--force-reinstall", + "torch", "torchvision", "torchaudio", + "--index-url", "https://download.pytorch.org/whl/cu128", + ] + with mock.patch.dict( + os.environ, + { + "UV_INDEX": "https://mirror.corp/simple", + "UV_EXTRA_INDEX_URL": "https://mirror.corp/extra", + "UV_INDEX_URL": "https://mirror.corp/root", + "UV_DEFAULT_INDEX": "https://mirror.corp/default", + }, + ): + env = ips._install_env_for_cmd(cmd) + assert env is not None, "a --index-url install must run with a scrubbed env" + for var in self.UV_VARS: + assert var not in env, f"{var} must be cleared for a pinned-index install" + + def test_pinned_default_index_strips_uv_index_vars(self): + # The uv-native spelling (--default-index) must be gated too, matching + # install.sh / install.ps1 which key off --default-index. + cmd = ["uv", "pip", "install", "torch", "--default-index", "https://x/cu126"] + with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}): + env = ips._install_env_for_cmd(cmd) + assert env is not None + assert "UV_INDEX" not in env + + def test_non_pinned_install_keeps_user_mirror(self): + # A plain install (no --index-url) must NOT scrub the env, so a user's + # UV_INDEX / UV_EXTRA_INDEX_URL mirror still applies to base packages. + cmd = ["uv", "pip", "install", "unsloth", "unsloth-zoo"] + with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}): + env = ips._install_env_for_cmd(cmd) + assert env is None, "non-pinned installs must inherit the caller env unchanged" + + def test_scrubbed_env_preserves_other_vars(self): + cmd = ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"] + with mock.patch.dict( + os.environ, + {"UV_INDEX": "https://mirror.corp/simple", "PATH_SENTINEL_XYZ": "keepme"}, + ): + env = ips._install_env_for_cmd(cmd) + assert env is not None + assert env.get("PATH_SENTINEL_XYZ") == "keepme", "only uv index vars are removed"