diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 6e8474176b..599f71c9da 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -2560,9 +2560,13 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]: cmd.extend(_translate_pip_args_for_uv(args)) # Torch is pre-installed by install.sh/setup.ps1. Do not add # --torch-backend by default -- it can cause solver dead-ends on CPU-only - # machines. Callers that need it can set UV_TORCH_BACKEND. + # machines. Callers that need it can set UV_TORCH_BACKEND. Never add it to a + # pinned-index command: uv's torch backend redirects torch resolution to its + # own per-backend index (verified: --index-url .../cu128 with UV_TORCH_BACKEND + # =cpu resolves torch+cpu), which would defeat the explicit pin exactly like + # the index env vars _install_env_for_cmd() strips. _tb = os.environ.get("UV_TORCH_BACKEND", "") - if _tb: + if _tb and not _is_pinned_index_cmd(cmd): cmd.append(f"--torch-backend={_tb}") return cmd @@ -2576,7 +2580,22 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]: # 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") +# UV_TORCH_BACKEND is stripped too: uv's torch backend redirects torch resolution +# 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). +_UV_INDEX_ENV_VARS = ( + "UV_DEFAULT_INDEX", + "UV_INDEX_URL", + "UV_INDEX", + "UV_EXTRA_INDEX_URL", + "UV_TORCH_BACKEND", +) + + +def _is_pinned_index_cmd(cmd: "list[str] | tuple[str, ...]") -> bool: + """True when the command pins an index via --index-url / --default-index.""" + return any(arg in ("--index-url", "--default-index") for arg in cmd) def _install_env_for_cmd(cmd: "list[str]") -> "dict[str, str] | None": @@ -2585,11 +2604,12 @@ def _install_env_for_cmd(cmd: "list[str]") -> "dict[str, str] | None": 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). + (the torch repair paths), the uv index env vars (and UV_TORCH_BACKEND) are + removed so the pinned index is not overridden by an inherited mirror or torch + backend (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): + if not _is_pinned_index_cmd(cmd): return None env = os.environ.copy() for name in _UV_INDEX_ENV_VARS: diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 26219ad45a..5b782e3fc6 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -3168,10 +3168,15 @@ if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) { $ROCmVisionSpec = "torchvision>=0.26.0,<0.27.0" $ROCmAudioSpec = "torchaudio>=2.11.0,<2.12.0" substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchSpec" "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 (matches the automatic path's # bare floor for these arches). + # Require a DIGIT after rocm (rocm7.1, not rocm-rel-7.2.1 / rocm-current): + # only rocm and gfx* are pip --index-url (PEP 503) families; a + # rocm- leaf is a custom/find-links URL that must stay on the + # verbatim unknown-pin path. Mirrors install.ps1 and + # install_python_stack.py's _is_pip_rocm_family_leaf (^rocm\d). $ROCmIndexUrl = $PinnedTorchIndexUrl $ROCmTorchSpec = "torch" $ROCmVisionSpec = "torchvision" diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 384dc858bc..e7f5a2441f 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -485,6 +485,17 @@ class TestPinnedRocmLeafDigitParity: "-like 'rocm*'" not in text ), "install.ps1 must not route a pinned index on a bare -like 'rocm*' glob" + def test_setup_ps1_pinned_reroute_requires_rocm_digit(self): + text = SETUP_PS1.read_text(encoding = "utf-8") + assert "-match '^rocm\\d'" in text, ( + "setup.ps1 pinned-index reroute must use -match '^rocm\\d' (not a bare " + "-like 'rocm*' glob) so custom find-links leaves stay on the verbatim path" + ) + pinned_block = text[text.find("$_pinGfx211 = Test-RocmGfx211Leaf") :][:2000] + assert "-like 'rocm*'" not in pinned_block, ( + "setup.ps1 pinned reroute must not route 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( diff --git a/tests/python/test_install_python_stack.py b/tests/python/test_install_python_stack.py index a991c09169..6a6db4abb3 100644 --- a/tests/python/test_install_python_stack.py +++ b/tests/python/test_install_python_stack.py @@ -54,6 +54,24 @@ class TestBuildUvCmdTorchBackend: a.startswith("--torch-backend") for a in cmd ), f"Empty UV_TORCH_BACKEND should not add flag, got: {cmd}" + def test_uv_torch_backend_skipped_for_pinned_index(self): + """A pinned-index command must NOT get --torch-backend: uv's torch backend + redirects torch resolution to its own per-backend index even when + --index-url is given (verified: cu128 pin + backend cpu installs + torch+cpu), defeating the pin.""" + for pin_flag in ("--index-url", "--default-index"): + with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}): + cmd = self._call(("torch", pin_flag, "https://download.pytorch.org/whl/cu128")) + assert not any( + a.startswith("--torch-backend") for a in cmd + ), f"{pin_flag} command must not carry --torch-backend, got: {cmd}" + + def test_uv_torch_backend_kept_for_unpinned(self): + """Non-pinned commands still honour UV_TORCH_BACKEND.""" + with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}): + cmd = self._call(("somepackage",)) + assert "--torch-backend=cpu" in cmd + class TestUvSafePath: """_uv_safe_path hands uv a space-free `-c`/`-r` path (issue #6503).""" @@ -216,3 +234,10 @@ class TestPinnedIndexClearsUvEnv: 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" + + 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.""" + with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}): + env = ips._install_env_for_cmd(["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]) + assert env is not None and "UV_TORCH_BACKEND" not in env diff --git a/tests/studio/test_torch_index_marker.ps1 b/tests/studio/test_torch_index_marker.ps1 index 516f3f6f90..896fc618f9 100644 --- a/tests/studio/test_torch_index_marker.ps1 +++ b/tests/studio/test_torch_index_marker.ps1 @@ -39,12 +39,14 @@ function Check($name, $cond) { } Write-Host "Get-NormalizedIndexUrl (trim / strip trailing slash / lowercase leaf)" -Check "trailing slashes + leaf lowered" ` - ((Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all///") -eq "https://repo.amd.com/rocm/whl/gfx120x-all") +# -ceq: PowerShell -eq is case-INsensitive for strings, which would make these +# case-normalization checks vacuous (any casing would pass). +Check "trailing slashes + known family leaf lowered" ` + ((Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all///") -ceq "https://repo.amd.com/rocm/whl/gfx120x-all") Check "whitespace trimmed" ` - ((Get-NormalizedIndexUrl " https://download.pytorch.org/whl/cu128 ") -eq "https://download.pytorch.org/whl/cu128") -Check "host case preserved, leaf lowered" ` - ((Get-NormalizedIndexUrl "https://Mirror.Local/Simple/") -eq "https://Mirror.Local/simple") + ((Get-NormalizedIndexUrl " https://download.pytorch.org/whl/cu128 ") -ceq "https://download.pytorch.org/whl/cu128") +Check "host case preserved, unknown custom leaf keeps case" ` + ((Get-NormalizedIndexUrl "https://Mirror.Local/Simple/") -ceq "https://Mirror.Local/Simple") Check "gfx120X-all == gfx120x-all after normalize" ` ((Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120X-all") -eq (Get-NormalizedIndexUrl "https://repo.amd.com/rocm/whl/gfx120x-all")) Check "empty -> null" ($null -eq (Get-NormalizedIndexUrl " "))