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