ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

install: apply an explicit custom torch-index pin on the first update

Browse source 
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.
This commit is contained in:
Daniel Han 2026-07-12 12:39:29 +00:00
commit d671d8fb2f
4 changed files with 100 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

31
studio/install_python_stack.py
View file

@ -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",