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

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: