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

install: keep pinned torch repairs on the pinned index

Browse source 
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).
This commit is contained in:
Daniel Han 2026-07-12 10:17:01 +00:00
commit a98203b539
4 changed files with 175 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

9
install.ps1
View file

@ -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<digit> and gfx* are pip --default-index (PEP 503) families.
# A rocm-<nondigit> 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
}
}

33
studio/install_python_stack.py
View file

@ -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 <cuXXX/rocm/cpu>)
# 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:

72
tests/python/test_cross_platform_parity.py
View file

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

62
tests/python/test_install_python_stack.py
View file

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