diff --git a/install.ps1 b/install.ps1 index 16416dd9ae..dd27a9c547 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', 'UV_TORCH_BACKEND') { + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND', 'UV_FIND_LINKS') { $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n) Remove-Item "Env:$n" -ErrorAction SilentlyContinue } diff --git a/install.sh b/install.sh index bbc7e52077..67ae1e4eba 100755 --- a/install.sh +++ b/install.sh @@ -165,7 +165,7 @@ run_install_cmd() { # 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 -u UV_TORCH_BACKEND "$@" ;; + *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL -u UV_TORCH_BACKEND -u UV_FIND_LINKS "$@" ;; esac if _is_verbose; then "$@" && return 0 diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index d88cdfab13..14bd465093 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -2594,7 +2594,9 @@ _UV_INDEX_ENV_VARS = ( "UV_INDEX", "UV_EXTRA_INDEX_URL", "UV_TORCH_BACKEND", + "UV_FIND_LINKS", "PIP_EXTRA_INDEX_URL", + "PIP_FIND_LINKS", ) diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 40c1c47397..a725b83c8c 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2938,23 +2938,30 @@ if (Get-Command uv -ErrorAction SilentlyContinue) { # Helper: install a package, preferring uv with pip fallback function Fast-Install { param([Parameter(ValueFromRemainingArguments=$true)]$Args_) - if ($UseUv) { - $VenvPy = (Get-Command python).Source - # An explicit --index-url must win. Inherited uv index env vars otherwise - # override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop - # 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', 'UV_TORCH_BACKEND') { - $saved[$n] = [Environment]::GetEnvironmentVariable($n) - Remove-Item "Env:$n" -ErrorAction SilentlyContinue - } + # An explicit --index-url must win. Inherited uv index env vars otherwise + # override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop + # them only for index-pinned installs; mirrors still apply elsewhere. The + # scrub must cover the WHOLE function, not just the uv attempt: the pip + # fallback honours PIP_EXTRA_INDEX_URL / PIP_FIND_LINKS in addition to the + # pinned --index-url, so restoring before the fallback reopens the hole. + # UV_TORCH_BACKEND / UV_FIND_LINKS likewise reroute a pinned uv resolve. + $saved = @{} + if (@($Args_) -contains '--index-url') { + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', + 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'PIP_EXTRA_INDEX_URL', 'PIP_FIND_LINKS') { + $saved[$n] = [Environment]::GetEnvironmentVariable($n) + Remove-Item "Env:$n" -ErrorAction SilentlyContinue } - try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 } - finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } } - if ($LASTEXITCODE -eq 0) { return } } - & python -m pip install @Args_ 2>&1 + try { + if ($UseUv) { + $VenvPy = (Get-Command python).Source + $result = & uv pip install --python $VenvPy @Args_ 2>&1 + if ($LASTEXITCODE -eq 0) { return } + } + & python -m pip install @Args_ 2>&1 + } + finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } } } # ── Check if Python deps need updating ── diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 291dca26c4..17bb39024e 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -567,6 +567,28 @@ class TestPinnedIndexClearsUvEnvParity: '"PIP_EXTRA_INDEX_URL",' in stack ), "install_python_stack.py strip tuple must include PIP_EXTRA_INDEX_URL" + def test_all_installers_scrub_find_links(self): + """uv's --find-links (env UV_FIND_LINKS) adds candidate locations that can + satisfy torch off a pinned index; every pinned-install scrub must clear it.""" + sh = INSTALL_SH.read_text(encoding = "utf-8") + assert "-u UV_FIND_LINKS" in sh + for path in (INSTALL_PS1, SETUP_PS1): + assert "'UV_FIND_LINKS'" in path.read_text(encoding = "utf-8"), path.name + stack = STACK_PY.read_text(encoding = "utf-8") + assert '"UV_FIND_LINKS",' in stack and '"PIP_FIND_LINKS",' in stack + + def test_setup_ps1_scrub_covers_pip_fallback(self): + """setup.ps1's Fast-Install must keep the scrub active through the pip + fallback (pip honours PIP_EXTRA_INDEX_URL / PIP_FIND_LINKS in addition to + --index-url); restoring the vars before the fallback reopens the hole.""" + text = SETUP_PS1.read_text(encoding = "utf-8") + fi = text[text.find("function Fast-Install") :][:2500] + assert "'PIP_EXTRA_INDEX_URL'" in fi and "'PIP_FIND_LINKS'" in fi + # the pip fallback must sit INSIDE the try whose finally restores the vars + assert fi.find("python -m pip install") < fi.find("finally"), ( + "pip fallback must run before the scrub is restored" + ) + 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