docker: harden rollback, publish, shim, and view-cleanup paths

Ten verified fixes from a 12-reviewer audit of the image tooling, each
reproduced before fixing:

1. install_llama_prebuilt.py move_install_dir_aside: the EXDEV fallback
   copied straight into the rollback path, so a copy that died halfway
   (ENOSPC, I/O error) left a partial tree that activation recovery would
   later restore over the intact install while deleting the good copy.
   Copy to a temp sibling and publish with one atomic rename; dst.exists()
   is now a truthful complete-tree signal.

2. unsloth_run.py --out truncated the existing output before nbconvert
   ran, so a timeout, missing kernel, or failed cell irreversibly
   destroyed the previous result. The input copy and executed result are
   staged as temp files next to the destination and published with
   os.replace only on exit code 0.

3. unsloth_nb_view.py cleanup treated every symlink in the view as its
   own: user-created links (and an operator's view-root routing symlink)
   were deleted on every rebuild. Cleanup now removes only links that
   resolve into the notebooks tree it links from, and builds inside a
   view-root symlink's target instead of unlinking it.

4. unsloth_llama_update.sh: the unconditional EXIT trap deleted the .old
   backup even when it was the only remaining copy (signal between the two
   renames, or a failed swap whose restore also failed). The handler now
   restores the backup first when the install dir is missing and removes
   it only after the new tree is verifiably active; HUP/INT/TERM route
   through the same handler.

5. unsloth_pip_shim.py: transitive dependencies could replace the baked
   torch stack (reproduced with a wheel requiring torch==99.0). Every
   forwarded install now carries a constraints file pinning the installed
   protected set, turning the swap into ResolutionImpossible.

6. unsloth_pip_shim.py: ${UPPER} env references in requirements files were
   classified before pip expanded them, bypassing the protected-package
   filter; the shim now expands with pip's exact regex first.

7. unsloth_pip_shim.py: a failure writing the filtered requirements copy
   returned the ORIGINAL file, forwarding exactly the protected pins it
   had detected; it now fails closed.

8. docker-publish.yml: workflow_dispatch defaulted unsloth_ref to 'main'
   while the stable-tag gates require '', so UI-default manual runs could
   never advance :core/:latest/:studio; the default is now empty.

9. entrypoint.sh: the sm_103/sm_121 branch rewrote libnvrtc.so.12 to the
   CUDA-13 build but the ordinary-GPU branch never restored it, so a
   container moved to an older GPU kept the stale link; it is now reversed
   when it points exactly at the .cu13 target.

Rejected after verification (no code change): timeout=0 semantics are
documented at the site with no zero callers, TORCHINDUCTOR_COMPILE_THREADS
override is deliberate, fetchNews is a string enum per JupyterLab's schema,
:base tag appears in no in-tree doc, install-cell digest exclusion is the
module's stated contract, transformers ceiling semantics are documented,
and the cloudflared download mirrors the pre-existing Studio downloader
(Cloudflare publishes no checksum asset). The UNSLOTH_ALLOW_CPU import
crash lives in unsloth_zoo (compiler.py / loss_utils.py capability probes),
not in this diff; the image consumes the zoo fix automatically once merged
there.

Tests: shim suite extended to 63 (constraints, env expansion, fail-closed),
jit-selector suite to 14 (NVRTC reversal transitions), plus staged-publish
and ownership repros; wider studio install suite green except failures
reproduced at the unmodified head.
This commit is contained in:
Daniel Han 2026-07-13 03:03:03 +00:00
commit 47d66ecb53
9 changed files with 349 additions and 63 deletions

View file

@ -80,10 +80,21 @@ def _run(shim, tool, args):
shim.main()
execd = None
except _Exec as exc:
# main() builds [REAL[tool]] + head + keep_args; head ends with the
# `install` verb, so everything after it is what we asserted on.
# main() builds [REAL[tool]] + head + keep_args + the protected
# constraints pair; head ends with the `install` verb, so everything
# after it is what we asserted on. The trailing
# `--constraint <unsloth-nb-protected-*.txt>` pair is injected on
# EVERY forwarded install (resolver-level protection); strip it here
# so each test asserts on its own arguments -- the dedicated
# constraint-injection tests below cover the pair itself.
i = exc.argv.index("install")
execd = exc.argv[i + 1 :]
if (
len(execd) >= 2
and execd[-2] == "--constraint"
and os.path.basename(execd[-1]).startswith("unsloth-nb-protected-")
):
execd = execd[:-2]
marker = shim._marker_path.read_text() if shim._marker_path.exists() else None
return execd, marker
@ -530,3 +541,104 @@ def test_upgrade_strategy_only_if_needed_also_dropped(shim):
# keeps the kept target installing normally.
execd, _ = _run(shim, "pip", ["--upgrade-strategy", "only-if-needed", "peft"])
assert execd == ["peft"], execd
# --------------------------------------------------------------------------
# Resolver-level protection: every forwarded install carries a constraints
# file pinning the installed protected packages, so a kept target's
# DEPENDENCY on an incompatible torch/transformers/etc. fails loudly instead
# of replacing the baked wheel.
# --------------------------------------------------------------------------
def _raw_execd(shim, tool, args):
"""Like _run but WITHOUT stripping the injected constraint pair."""
argv = ["uv", "pip", "install", *args] if tool == "uv" else ["pip", "install", *args]
with pytest.MonkeyPatch.context() as mp:
mp.setattr(shim.sys, "argv", argv)
try:
shim.main()
return None
except _Exec as exc:
return exc.argv[exc.argv.index("install") + 1 :]
def test_forwarded_install_carries_protected_constraints(shim):
execd = _raw_execd(shim, "pip", ["peft"])
assert execd is not None and execd[-2] == "--constraint", execd
pins = Path(execd[-1]).read_text(encoding = "utf-8").strip().splitlines()
assert pins, "constraints file must pin the installed protected packages"
assert all("==" in pin for pin in pins), pins
names = {pin.split("==", 1)[0].lower().replace("_", "-") for pin in pins}
protected = {"transformers"} | shim._KEEP | {"nvidia-"}
assert all(
n in shim._KEEP or n == "transformers" or n.startswith("nvidia-") for n in names
), names
def test_noop_install_gets_no_constraints(shim):
# A cell whose only target is protected still no-ops (no exec at all).
execd = _raw_execd(shim, "pip", ["torch"])
assert execd is None
# --------------------------------------------------------------------------
# pip expands ${UPPERCASE} in requirements files AFTER the shim classifies the
# literal text; classification must expand the same way or `${PKG}==...` with
# PKG=torch walks straight past _KEEP.
# --------------------------------------------------------------------------
def test_env_expanded_protected_requirement_dropped(shim, tmp_path, monkeypatch):
monkeypatch.setenv("PKG", "torch")
req = tmp_path / "reqs.txt"
req.write_text("${PKG}==2.11.0\nsnac==1.2.0\n", encoding = "utf-8")
execd, _ = _run(shim, "pip", ["-r", str(req)])
assert execd is not None and execd[0] == "-r", execd
filtered = Path(execd[1]).read_text(encoding = "utf-8")
assert "snac==1.2.0" in filtered
assert "${PKG}" not in filtered and "torch" not in filtered
def test_env_expanded_transformers_pin_recorded(shim, tmp_path, monkeypatch):
monkeypatch.setenv("TF_PKG", "transformers")
req = tmp_path / "reqs.txt"
req.write_text("${TF_PKG}==4.56.2\nsnac==1.2.0\n", encoding = "utf-8")
_, marker = _run(shim, "pip", ["-r", str(req)])
assert marker == "4.56.2"
def test_unset_env_reference_left_verbatim(shim, tmp_path, monkeypatch):
monkeypatch.delenv("NOT_SET_ANYWHERE", raising = False)
req = tmp_path / "reqs.txt"
req.write_text("${NOT_SET_ANYWHERE}==1.0\nsnac==1.2.0\n", encoding = "utf-8")
execd, _ = _run(shim, "pip", ["-r", str(req)])
# Nothing protected detected -> the original file is forwarded unchanged
# (pip forwards unset references verbatim too).
assert execd == ["-r", str(req)], execd
# --------------------------------------------------------------------------
# Filtered-copy write failures fail CLOSED: the original file pins protected
# packages, so forwarding it would hand pip exactly what must be filtered.
# --------------------------------------------------------------------------
def test_filter_write_failure_refuses_original_file(shim, tmp_path, monkeypatch):
req = tmp_path / "reqs.txt"
req.write_text("torch==2.11.0\nsnac==1.2.0\n", encoding = "utf-8")
def denied(*args, **kwargs):
raise OSError(30, "Read-only file system")
monkeypatch.setattr(shim.tempfile, "mkstemp", denied)
with pytest.raises(SystemExit, match = "refusing to forward"):
shim._filter_requirements_file(str(req))
def test_filter_write_failure_clean_file_passes_through(shim, tmp_path, monkeypatch):
# A file with nothing protected never needs the temp copy, so a broken
# TMPDIR must not block it.
req = tmp_path / "reqs.txt"
req.write_text("snac==1.2.0\n", encoding = "utf-8")
def denied(*args, **kwargs):
raise OSError(30, "Read-only file system")
monkeypatch.setattr(shim.tempfile, "mkstemp", denied)
path, recorded, dropped = shim._filter_requirements_file(str(req))
assert path == str(req) and recorded is None and dropped == []