From d671d8fb2f7e5d8f358022f2074bc08fd6fc86e3 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 12 Jul 2026 12:39:29 +0000 Subject: [PATCH] install: apply an explicit custom torch-index pin on the first update Round 6: an explicitly-set custom (unknown-family) UNSLOTH_TORCH_INDEX_URL was silently ignored on the first `studio update` of a venv that predates the marker feature, on both platforms, because the no-marker case was treated as "do nothing" and the version-tag heuristics cannot judge an unknown leaf. 1. install_python_stack.py _ensure_verbatim_torch_index now reinstalls verbatim when the marker is ABSENT (None), not only when it differs, and short-circuits only when the marker already records this exact pin. It then writes the marker, so every later update is a no-op. A user who did not set the override gets pin=None and is untouched, so an out-of-band torch install is never clobbered. 2. setup.ps1: for an unknown-family pin on a marker-less venv the stale-venv check now sets PinChangedForceReinstall so the torch block reinstalls in place from the pin. It deliberately does NOT set shouldRebuild, which would wipe the venv and strand a direct `studio update`. 3. setup.sh (the Linux `studio update` entry point) skipped install_python_stack.py entirely when unsloth was already current, so the marker-driven reinstall (both the verbatim custom pin and the cu/rocm flavor and family-change repair, e.g. gfx1151 to gfx120X-all) never ran. It now forces the dependency pass when a torch-index pin env var is set; the pass is idempotent and no-ops when the marker already matches. This mirrors setup.ps1's stale-venv pre-check. Tests: 3 new parity assertions. --- studio/install_python_stack.py | 31 +++++++----- studio/setup.ps1 | 11 +++++ studio/setup.sh | 15 ++++++ tests/python/test_cross_platform_parity.py | 55 ++++++++++++++++++++++ 4 files changed, 100 insertions(+), 12 deletions(-) diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 8c18d4753e..0798fbd634 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -1420,12 +1420,18 @@ def _ensure_verbatim_torch_index() -> None: pin (it is neither rocm/gfx nor cpu nor cuXXX), so without this the pin would be silently ignored and the GPU-probed default index used instead. - Fires ONLY when the marker exists and records a DIFFERENT index than the pin - (or after this reinstalls, the marker is rewritten to match). With NO marker it - is a no-op: an old venv (or torch installed out-of-band) must not be blindly - force-reinstalled from an unverified custom index -- backward compatibility. - macOS/no-torch: skipped (no torch to repair). The install uses the pinned URL - exclusively (--index-url) with bare specs so it "wins verbatim". + Fires when the marker differs from the pin (True) OR is ABSENT (None): a venv + predating the marker feature has no record, and the version-tag heuristics + cannot judge an unknown-family pin, so an explicitly-set URL would otherwise be + silently ignored on the first `studio update` -- the user asked for this index, + so apply it verbatim ONCE and record it. The write below makes every later + update a no-op (marker == pin -> False). Skips only when the marker already + records this exact pin (False). A user who did NOT set the override gets + pin=None and is never touched, so an out-of-band torch install is safe. macOS/ + no-torch: skipped (no torch to repair). The install uses the pinned URL + exclusively (--index-url) with bare specs so it "wins verbatim" -- an incomplete + mirror that cannot serve the trio fails loudly here, same as the marker-present + path (that is the cost of honouring an explicit pin). """ if NO_TORCH or IS_MACOS: return @@ -1433,13 +1439,14 @@ def _ensure_verbatim_torch_index() -> None: if pin is None: return _mismatch = _marker_pin_mismatch(pin) - if _mismatch is not True: - # None -> no marker (fall back / do nothing); False -> already this index. + if _mismatch is False: + # Marker already records this exact pin -> no reinstall (no per-update loop). + # True (marker differs) or None (no marker yet) both fall through to apply + # the explicit pin verbatim once, then _write_torch_index_marker below makes + # the next update a no-op. return - print( - f" explicit torch index pin ({pin}) differs from the recorded index -- " - f"reinstalling torch verbatim from it" - ) + _why = "differs from the recorded index" if _mismatch is True else "has no recorded index yet" + print(f" explicit torch index pin ({pin}) {_why} -- reinstalling torch verbatim from it") pip_install( "torch (pinned custom index)", "--force-reinstall", diff --git a/studio/setup.ps1 b/studio/setup.ps1 index f068230910..8ed74a6993 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2804,6 +2804,17 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode # not also rebuild on a bogus tag comparison here. $_expectedKnown = $false $expectedTorchTag = $installedTorchTag + # A venv predating the marker feature has $_markerMismatch = $null. + # The user has now explicitly set an unknown-family URL pin, but the + # flavor heuristic cannot judge it, so without the marker the pin would + # be silently ignored on this first update (the fast path keeps + # $SkipPythonDeps true and the --index-url install never runs). Apply + # it ONCE in place: PinChangedForceReinstall forces the dependency pass + # so the torch block reinstalls from $PinnedTorchIndexUrl and writes the + # marker, making later updates no-ops. Do NOT set $shouldRebuild here -- + # that wipes the venv and delegates to install.ps1, stranding a direct + # `studio update` at "Virtual environment not found". + if ($null -eq $_markerMismatch) { $script:PinChangedForceReinstall = $true } } } elseif ($HasNvidiaSmi) { $expectedTorchTag = Get-PytorchCudaTag diff --git a/studio/setup.sh b/studio/setup.sh index d244e3cdcf..c3935ae268 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -971,6 +971,21 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1) fi fi +# Honor an explicit torch-index pin on `studio update` even when unsloth is already +# current. install_python_stack.py owns the marker-driven torch reinstall (a verbatim +# custom-index pin via _ensure_verbatim_torch_index, and cu*/rocm flavor/family repair +# via _ensure_cuda/rocm_torch), and the fast "up to date" path above skips it -- so a +# newly-set or changed UNSLOTH_TORCH_INDEX_URL / _FAMILY (including the gfx1151 -> +# gfx120X-all marker case) would be silently ignored on Linux. Force the pass when a +# pin is set; it is idempotent -- the _ensure_* helpers no-op when the marker already +# records this index. Mirrors setup.ps1, whose stale-venv pre-check force-reinstalls on +# a pin change. Same shape as the anyio-repair override above. Unpinned updates keep +# the fast path. +if [ "$_SKIP_PYTHON_DEPS" = true ] && [ -n "${UNSLOTH_TORCH_INDEX_URL:-}${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]; then + substep "explicit torch-index pin set -- running dependency pass to apply/verify it..." + _SKIP_PYTHON_DEPS=false +fi + if [ "$_SKIP_PYTHON_DEPS" = false ]; then install_python_stack else diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index 1413a3f3bc..4bc457dc11 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -554,6 +554,61 @@ class TestPinnedRocmLeafDigitParity: ), "install.sh must not gate _install_bnb_rocm on a bare */rocm* whole-URL glob" +class TestFirstCustomPinAppliedWithoutMarker: + """An explicitly-set custom (unknown-family) UNSLOTH_TORCH_INDEX_URL must be + applied on the FIRST `studio update` of a venv that predates the marker feature. + Such a venv has no .unsloth-torch-index marker, so the marker compare returns + None/$null and the version-tag heuristics cannot judge an unknown leaf; without + treating "no marker" as "apply verbatim once", the explicit pin would be silently + ignored until a marker happened to exist. The verbatim reinstall writes the + marker, so every later update is a no-op (marker == pin).""" + + def test_stack_py_applies_pin_when_marker_absent(self): + text = STACK_PY.read_text(encoding = "utf-8") + body = text[text.find("def _ensure_verbatim_torch_index") :][:2200] + # The short-circuit must be "already this exact pin" (False), NOT the old + # "anything other than a definite mismatch" (is not True), which also bailed + # on a None (no-marker) result. + assert "if _mismatch is False:" in body, ( + "_ensure_verbatim_torch_index must reinstall on an absent marker (None), " + "returning early only when the marker already records this exact pin (False)" + ) + assert "if _mismatch is not True:" not in body, ( + "_ensure_verbatim_torch_index must no longer skip the reinstall when the " + "marker is absent (None)" + ) + + def test_setup_sh_forces_stack_pass_when_pin_set(self): + """On Linux, `studio update` runs setup.sh, which skips install_python_stack.py + (the only place the marker-driven torch reinstall lives) when unsloth is already + current. Without an override that fast path silently ignores an explicit torch + index pin, so setup.sh must force the dependency pass when a pin env var is set.""" + setup_sh = REPO_ROOT / "studio" / "setup.sh" + text = setup_sh.read_text(encoding = "utf-8") + assert ( + '[ -n "${UNSLOTH_TORCH_INDEX_URL:-}${UNSLOTH_TORCH_INDEX_FAMILY:-}" ]' in text + and "_SKIP_PYTHON_DEPS=false" in text + ), ( + "setup.sh must clear _SKIP_PYTHON_DEPS when UNSLOTH_TORCH_INDEX_URL/_FAMILY is " + "set so install_python_stack.py runs and applies the pin (marker no-ops if unchanged)" + ) + + def test_setup_ps1_forces_reinstall_when_marker_absent(self): + text = SETUP_PS1.read_text(encoding = "utf-8") + # The unknown-family else branch must promote a null marker to an in-place + # force-reinstall (PinChangedForceReinstall), NOT a wipe ($shouldRebuild). + assert 'if ($null -eq $_markerMismatch) { $script:PinChangedForceReinstall = $true }' in text, ( + "setup.ps1 must set PinChangedForceReinstall for an unknown-family pin on a " + "marker-less venv so the torch block reinstalls from $PinnedTorchIndexUrl in place" + ) + # Guard: the unknown-leaf branch must not trigger the venv wipe path. + else_branch = text[text.find("PEP 503 mirror ending in /simple") :][:900] + assert "$shouldRebuild = $true" not in else_branch, ( + "setup.ps1 unknown-family branch must repair in place (PinChangedForceReinstall), " + "not wipe the venv ($shouldRebuild), which would strand a direct studio update" + ) + + 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