From 47d66ecb53ccb0908a81a094994550c4a95483fa Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 13 Jul 2026 03:03:03 +0000 Subject: [PATCH] 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. --- .github/workflows/docker-publish.yml | 21 +++-- docker/entrypoint.sh | 25 ++++-- docker/unsloth_llama_update.sh | 29 ++++++- docker/unsloth_nb_view.py | 61 ++++++++----- docker/unsloth_pip_shim.py | 76 ++++++++++++++-- docker/unsloth_run.py | 49 ++++++++--- studio/install_llama_prebuilt.py | 21 ++++- tests/python/test_unsloth_pip_shim.py | 116 ++++++++++++++++++++++++- tests/sh/test_select_cuda_jit_tools.sh | 14 ++- 9 files changed, 349 insertions(+), 63 deletions(-) diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index e2cce67c52..9371d8e0d2 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -37,9 +37,14 @@ on: workflow_dispatch: inputs: unsloth_ref: - description: 'unsloth git ref to bake in' + # Blank means "the dispatched branch" (the resolver below falls back to + # the triggering sha, then main). The stable-tag gates (:core/:latest/ + # :studio) require this input to be EMPTY -- stable tags only when the + # operator did not override the source ref -- so a non-blank default + # would make every UI-default dispatch publish SHA tags only. + description: 'unsloth git ref override (blank = dispatched branch + stable tags)' required: false - default: 'main' + default: '' unsloth_zoo_ref: description: 'unsloth-zoo git ref to bake in' required: false @@ -114,10 +119,10 @@ jobs: # Freeze the requested unsloth ref to ONE concrete sha before the matrix # fans out, so both base arch legs AND the Studio build bake the identical - # unsloth commit even when the requested ref is a mutable branch (the - # workflow_dispatch default is unsloth_ref=main) that advances during the - # ~4h base + Studio run. Same requested-ref precedence the inline build-arg - # used: the dispatch input wins (default main), else the pushed tag, else + # unsloth commit even when the requested ref is a mutable branch that + # advances during the ~4h base + Studio run. Same requested-ref precedence + # the inline build-arg used: the dispatch input wins (blank by default, + # so stable tags stay enabled), else the pushed tag, else # the triggering commit sha, else main. A 40-char sha (branch/schedule # push) is already frozen; a branch/tag is resolved via ls-remote, exactly # like the zoo and notebooks steps, falling back to the bare ref on a @@ -276,8 +281,8 @@ jobs: # passed as a bogus --build-arg. Explanations live here instead: # UNSLOTH_REF (from the prepare job): resolved to ONE sha before the # matrix fans out, so both arch legs and the Studio build bake the - # identical unsloth commit even if a mutable branch (dispatch's - # unsloth_ref=main default) advances mid-run. Same requested-ref + # identical unsloth commit even if a mutable branch (an explicit + # dispatch unsloth_ref) advances mid-run. Same requested-ref # precedence as before: dispatch input, else the pushed tag, else # the triggering commit sha, else main. # UNSLOTH_ZOO_REF (from the prepare job): explicit dispatch input, diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index 5639103147..cae8e8c7fd 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -56,12 +56,25 @@ select_cuda_jit_tools() { 10.3|12.1) need_cu13=1 ;; esac done <<< "${caps}" - # Non-datacenter / undetectable / CPU host: nothing to do. cu12.8 is the - # immutable baked default (libnvrtc.so.12 -> .cu128.orig, Triton on its - # bundled cu12.8 ptxas), loadable on every supported 570+ driver, and needs - # NO write -- so a non-root `docker run --user` container is never left on a - # cu13 NVRTC a 570-579 driver cannot load. - [[ "${need_cu13}" -eq 1 ]] || return 0 + # Non-datacenter / undetectable / CPU host: cu12.8 is the immutable baked + # default (libnvrtc.so.12 -> .cu128.orig, Triton on its bundled cu12.8 + # ptxas), loadable on every supported 570+ driver, and needs NO write -- so + # a non-root `docker run --user` container is never left on a cu13 NVRTC a + # 570-579 driver cannot load. One exception needs a write: an earlier boot + # of this SAME container on sm_103/sm_121 left libnvrtc.so.12 -> .cu13 in + # the writable layer, and the container now runs on a GPU whose 570-579 + # driver cannot load cu13 output -- deterministically reverse exactly that + # selection (best-effort, same non-root caveat as the forward switch). + if [[ "${need_cu13}" -ne 1 ]]; then + for nvrtc_dir in \ + /opt/unsloth-venv/lib/python*/site-packages/nvidia/cuda_nvrtc/lib \ + "${UNSLOTH_STUDIO_HOME:-/opt/unsloth-studio}"/unsloth_studio/lib/python*/site-packages/nvidia/cuda_nvrtc/lib; do + [[ -e "${nvrtc_dir}/libnvrtc.so.12.cu128.orig" ]] || continue + [[ "$(readlink "${nvrtc_dir}/libnvrtc.so.12" 2>/dev/null)" == "libnvrtc.so.12.cu13" ]] || continue + ln -sf libnvrtc.so.12.cu128.orig "${nvrtc_dir}/libnvrtc.so.12" 2>/dev/null || true + done + return 0 + fi # Blackwell datacenter present: cu12.8 cannot emit compute_103/121, so point # Triton at cu13 ptxas and retarget each venv's libnvrtc.so.12 -> the staged # cu13 alias. -z guard leaves an explicit `docker run -e TRITON_PTXAS_PATH` diff --git a/docker/unsloth_llama_update.sh b/docker/unsloth_llama_update.sh index 3a796ed3c0..3d768bda34 100755 --- a/docker/unsloth_llama_update.sh +++ b/docker/unsloth_llama_update.sh @@ -98,7 +98,28 @@ fi # an atomic rename), then swap. On any failure the existing install is untouched. parent="$(dirname "$INSTALL_DIR")" work="$(mktemp -d "$parent/.llamaupd.XXXXXX")" -trap 'rm -rf "$work" "${INSTALL_DIR}.old.$$" 2>/dev/null || true' EXIT +backup="${INSTALL_DIR}.old.$$" +swap_done=0 +# The exit handler must never delete $backup while it is the ONLY copy of the +# install (signal between the two renames, or a failed swap whose restore also +# failed): put the old tree back first, and remove it only after the new tree +# is verifiably active. The signal traps make bash run the EXIT trap on +# HUP/INT/TERM too. +cleanup() { + if [ "$swap_done" -ne 1 ] && [ ! -e "$INSTALL_DIR" ] && [ -e "$backup" ]; then + if ! mv "$backup" "$INSTALL_DIR" 2>/dev/null; then + echo "[llama-update] CRITICAL: restore failed; previous install preserved at $backup" >&2 + fi + fi + rm -rf "$work" 2>/dev/null || true + if [ "$swap_done" = "1" ]; then + rm -rf "$backup" 2>/dev/null || true + fi +} +trap cleanup EXIT +trap 'exit 129' HUP +trap 'exit 130' INT +trap 'exit 143' TERM new="$work/llama.cpp" echo "[llama-update] fetching llama.cpp '$TAG' ($ARCH portable) ..." @@ -108,12 +129,12 @@ echo "[llama-update] fetching llama.cpp '$TAG' ($ARCH portable) ..." [ -e "$INSTALL_DIR/.unsloth-studio-owned" ] && touch "$new/.unsloth-studio-owned" echo "[llama-update] swapping into place ..." -mv "$INSTALL_DIR" "${INSTALL_DIR}.old.$$" +mv "$INSTALL_DIR" "$backup" if mv "$new" "$INSTALL_DIR"; then - rm -rf "${INSTALL_DIR}.old.$$" + swap_done=1 else echo "[llama-update] swap failed; restoring previous install" >&2 - mv "${INSTALL_DIR}.old.$$" "$INSTALL_DIR" + mv "$backup" "$INSTALL_DIR" exit 1 fi diff --git a/docker/unsloth_nb_view.py b/docker/unsloth_nb_view.py index 0d6bb5a4f3..2bbbacd11b 100644 --- a/docker/unsloth_nb_view.py +++ b/docker/unsloth_nb_view.py @@ -121,6 +121,14 @@ def build_view( if not os.path.isdir(nb_dir): raise SystemExit(f"no nb/ dir under {dest}") + # An operator may route the VIEW through a symlink to persistent/mounted + # storage. Build inside its target instead of unlinking the routing. + if os.path.islink(view): + resolved = os.path.realpath(view) + if not os.path.isdir(resolved): + raise SystemExit(f"view symlink has no directory target: {view} -> {resolved}") + view = resolved + rows = parse_readme(readme) if os.path.isfile(readme) else [] def allowed(fname): @@ -152,7 +160,7 @@ def build_view( # Rebuild VIEW: drop the symlinks/empty folders we made last boot, but never # the user's own files (VIEW is also JupyterLab's landing dir, so a user may # have saved real notebooks here). - _clear_view(view) + _clear_view(view, os.path.realpath(dest)) os.makedirs(view, exist_ok = True) n_links = 0 @@ -164,9 +172,9 @@ def build_view( target = os.path.join(nb_dir, fname) rel = os.path.relpath(target, folder) # ../../unsloth-notebooks/nb/ try: - if os.path.islink(link): + if os.path.islink(link) and _points_into(link, os.path.realpath(dest)): os.remove(link) # replace our own stale symlink - elif os.path.exists(link): + elif os.path.islink(link) or os.path.exists(link): # a real user file/dir already occupies this name -- never # clobber it; leave it and skip linking this notebook. print(f"[unsloth-nb] view: keep user file, skip link {fname}", file = sys.stderr) @@ -178,41 +186,54 @@ def build_view( return len(order), n_links -def _clear_view(path): +def _points_into(link, dest_real): + """True when a symlink resolves into the notebooks tree we link from. + + Every link this tool creates points at DEST/nb/, so this is the + ownership test for cleanup: a user's own symlink (to a dataset, project, + mounted dir, ...) resolves elsewhere and must survive a rebuild. realpath + resolves a broken link's path string too, so stale links to since-removed + notebooks are still recognised as ours. + """ + try: + target = os.path.realpath(link) + except OSError: + return False + return target == dest_real or target.startswith(dest_real + os.sep) + + +def _clear_view(path, dest_real): # Tear down a previously built VIEW in place. VIEW is also JupyterLab's - # landing directory, so a user may have saved real notebooks here -- those - # MUST survive a rebuild. We therefore unlink only symlinks (the notebooks we - # link) and rmdir only folders that end up empty; any regular file is left - # untouched, and a non-empty folder simply stays. + # landing directory, so a user may have saved real notebooks (or their own + # symlinks) here -- those MUST survive a rebuild. We therefore unlink only + # the symlinks we own (they resolve into DEST, see _points_into) and rmdir + # only folders that end up empty; any regular file and any user symlink is + # left untouched, and a non-empty folder simply stays. # - # islink is tested BEFORE isdir on the root: os.path.isdir() follows a - # symlink-to-directory, so without this a VIEW that is itself a symlink (e.g. - # pointed at the real nb/ tree) would be walked into and its target wiped. - if os.path.islink(path): - os.remove(path) - return - if not os.path.isdir(path): + # The VIEW root itself is never unlinked: build_view already resolved a + # symlinked root to its target, and an operator's routing symlink must + # survive. isdir on a non-link root is safe to walk. + if os.path.islink(path) or not os.path.isdir(path): return for root, dirs, files in os.walk(path, topdown = False): for name in files: p = os.path.join(root, name) - if os.path.islink(p): # our notebook symlinks only + if os.path.islink(p) and _points_into(p, dest_real): # our notebook symlinks only try: os.remove(p) except OSError: pass - # a regular file here is user-created -> keep it + # a regular file / user symlink here is user-created -> keep it for name in dirs: p = os.path.join(root, name) try: if os.path.islink(p): - os.remove(p) # symlinked dir: unlink, never recurse + if _points_into(p, dest_real): + os.remove(p) # our symlinked dir: unlink, never recurse else: os.rmdir(p) # succeeds only if we emptied it except OSError: pass # holds user files -> keep - # Leave the VIEW root itself in place: it may still hold user files, and - # build_view recreates it right after anyway. def main(argv): diff --git a/docker/unsloth_pip_shim.py b/docker/unsloth_pip_shim.py index afe51e6c8b..af3b8ce74c 100644 --- a/docker/unsloth_pip_shim.py +++ b/docker/unsloth_pip_shim.py @@ -245,6 +245,17 @@ def _version_pin(token): return m.group(1) if m else None +# pip expands ${UPPERCASE_NAME} in requirements files AFTER we classify the +# literal text (pip's ENV_VAR_RE; uv matches it), so `${PKG}==...` with +# PKG=torch would slip a protected package past _KEEP. Expand with the same +# syntax for CLASSIFICATION only; kept lines are forwarded verbatim. +_ENV_REF_RE = re.compile(r"\$\{([A-Z0-9_]+)\}") + + +def _expand_env_refs(text): + return _ENV_REF_RE.sub(lambda m: os.environ.get(m.group(1), m.group(0)), text) + + def _classify_flag_target(spec): """Classify the value that rides on -e/--editable or -P/--upgrade-package. @@ -319,9 +330,12 @@ def _rewrite_include(line, stripped, src_dir, depth): parent at that filtered copy. URLs and unreadable/absolute-unfiltered files fall back to an absolutised path so they still resolve. Returns (new_line, changed, recorded, dropped).""" - flag, target, comment = _parse_include(stripped) - if not target: + flag, raw_target, comment = _parse_include(stripped) + if not raw_target: return line, False, None, [] + # Resolve pip's ${VAR} references so the include we read/filter is the file + # pip would actually read (a literal `${DIR}/reqs.txt` never resolves here). + target = _expand_env_refs(raw_target) newline_char = "\n" if line.endswith("\n") else "" def _emit(new_target): @@ -336,7 +350,7 @@ def _rewrite_include(line, stripped, src_dir, depth): # (mirrors the top-level remote `-r`/`-c` refusal in main). new_line=None # tells the caller to remove the line entirely. if "://" in target: - return None, True, None, [flag + " " + target] + return None, True, None, [flag + " " + raw_target] abs_target = target if os.path.isabs(target) else os.path.join(src_dir, target) # Recursively filter the included file. Guard against cyclic / deep includes. if depth < 8: @@ -390,7 +404,7 @@ def _filter_requirements_file(path, _depth = 0): # the target is protected; a transformers pin is still recorded. e_flag, e_target, _e_comment = _parse_editable(stripped) if e_target is not None: - _action, _ver = _classify_flag_target(e_target) + _action, _ver = _classify_flag_target(_expand_env_refs(e_target)) if _action == "drop": if _ver and not recorded: recorded = _ver @@ -412,12 +426,13 @@ def _filter_requirements_file(path, _depth = 0): dropped.extend(inc_drp) continue spec = stripped.split(" #", 1)[0].strip() # drop any inline comment - name = _canon(spec) + classified = _expand_env_refs(spec) # classify what pip will SEE + name = _canon(classified) if name is None: out.append(line) # url / path / vcs / unparseable -> keep continue if name == "transformers": - v = _version_pin(spec) + v = _version_pin(classified) if v and not recorded: recorded = v dropped.append(spec) @@ -434,11 +449,50 @@ def _filter_requirements_file(path, _depth = 0): fd, tmp = tempfile.mkstemp(prefix = "unsloth-nb-req-", suffix = ".txt") with os.fdopen(fd, "w", encoding = "utf-8") as f: f.writelines(out) - except OSError: - return path, None, [] # can't write temp -> pass the file through unchanged + except OSError as exc: + # Fail CLOSED: protected requirements were detected in this file, so + # forwarding the original would hand pip exactly the specs we must + # filter. Abort the install with a clear error instead. + raise SystemExit( + f"[unsloth-nb] could not write a filtered copy of {path} ({exc}); " + "refusing to forward a requirements file that pins protected packages." + ) return tmp, recorded, dropped +def _protected_constraints_file(): + """Write `name==version` pins for every INSTALLED protected package to a + temp constraints file and return its path (None when nothing is pinned or + the file cannot be written). + + Argument filtering alone does not constrain pip/uv's RESOLVER: a kept + package may declare e.g. `torch==99.0` as a dependency and the tool would + replace the baked torch to satisfy it. Pinning the protected set on every + forwarded install makes such an install fail loudly instead. This is + belt-and-braces on top of the argument filtering, so a failure here keeps + the install usable rather than aborting it. + """ + try: + from importlib.metadata import distributions + + pins = {} + for dist in distributions(): + raw = (dist.metadata["Name"] or "").strip() + name = raw.lower().replace("_", "-") + if not name or name in pins: + continue + if name == "transformers" or name in _KEEP or name.startswith(_KEEP_PREFIX): + pins[name] = f"{raw}=={dist.version}" + if not pins: + return None + fd, tmp = tempfile.mkstemp(prefix = "unsloth-nb-protected-", suffix = ".txt") + with os.fdopen(fd, "w", encoding = "utf-8") as f: + f.write("\n".join(pins[name] for name in sorted(pins)) + "\n") + return tmp + except Exception: + return None + + def main(): tool = "uv" if os.path.basename(sys.argv[0]).startswith("uv") else "pip" argv = sys.argv[1:] @@ -667,6 +721,12 @@ def main(): print("[unsloth-nb] nothing to install after keeping the baked stack; ok.") return cmd = [REAL[tool]] + head + keep_args + # Constrain the resolver too: without this an allowed target could pull an + # incompatible torch/transformers/etc. in as a DEPENDENCY and replace the + # baked wheel even though the argument filter kept it off the command line. + constraints = _protected_constraints_file() + if constraints: + cmd += ["--constraint", constraints] sys.stdout.flush() os.execv(REAL[tool], cmd) diff --git a/docker/unsloth_run.py b/docker/unsloth_run.py index 5bd0be7793..a0f645aa02 100644 --- a/docker/unsloth_run.py +++ b/docker/unsloth_run.py @@ -68,19 +68,35 @@ def main(): want = args.tf or pin or (compat.tier_for_model(model) if compat else None) sidecar = compat.sidecar_for(want) if (compat and want) else None - # Materialise the notebook locally for nbconvert. + # Materialise the notebook locally for nbconvert. With --out, stage both the + # input copy and the executed result as temp files NEXT TO the destination + # (same dir, so the kernel cwd matches and the publish is one atomic + # os.replace) and only publish over an existing --out file when execution + # succeeded -- a timeout / failed cell / missing kernel must not destroy the + # previous output. tmp_dir = None - if args.notebook.startswith(("http://", "https://")) or args.out: - if args.out: - src_path = args.out - else: - tmp_dir = tempfile.mkdtemp() - src_path = os.path.join(tmp_dir, os.path.basename(args.notebook.split("?")[0])) + tmp_files = [] + publish_from = None + if args.out: + out_path = os.path.abspath(args.out) + out_dir = os.path.dirname(out_path) or "." + os.makedirs(out_dir, exist_ok = True) + fd, src_path = tempfile.mkstemp(prefix = ".unsloth-run-in-", suffix = ".ipynb", dir = out_dir) + with os.fdopen(fd, "w") as f: + json.dump(nb, f) + tmp_files.append(src_path) + fd, publish_from = tempfile.mkstemp(prefix = ".unsloth-run-out-", suffix = ".ipynb", dir = out_dir) + os.close(fd) + tmp_files.append(publish_from) + elif args.notebook.startswith(("http://", "https://")): + tmp_dir = tempfile.mkdtemp() + src_path = os.path.join(tmp_dir, os.path.basename(args.notebook.split("?")[0])) with open(src_path, "w") as f: json.dump(nb, f) + out_path = src_path else: src_path = args.notebook - out_path = args.out or src_path + out_path = src_path env = dict(os.environ) env["UNSLOTH_NB_SHIM"] = "1" # enable safe-install for the notebook's cells @@ -97,6 +113,7 @@ def main(): else: print("[unsloth-run] no transformers pin/model tier detected; using base venv") + nbconvert_out = publish_from if publish_from is not None else out_path cmd = [ "/opt/unsloth-venv/bin/jupyter", "nbconvert", @@ -107,17 +124,25 @@ def main(): "--ExecutePreprocessor.kernel_name=python3", src_path, "--output", - os.path.basename(out_path), + os.path.basename(nbconvert_out), "--output-dir", - os.path.dirname(os.path.abspath(out_path)) or ".", + os.path.dirname(os.path.abspath(nbconvert_out)) or ".", ] - print("[unsloth-run] executing:", os.path.basename(src_path)) + print("[unsloth-run] executing:", os.path.basename(args.notebook.split("?")[0]) if args.out else os.path.basename(src_path)) try: rc = subprocess.call(cmd, env = env) + if rc == 0 and publish_from is not None: + os.replace(publish_from, out_path) finally: - # Clean up the temp dir we materialised a downloaded notebook into. + # Clean up the temp dir we materialised a downloaded notebook into and + # any staging files left next to --out (already gone when published). if tmp_dir is not None: shutil.rmtree(tmp_dir, ignore_errors = True) + for p in tmp_files: + try: + os.remove(p) + except OSError: + pass sys.exit(rc) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 06a8ab9ca8..ba2860a21e 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -4798,14 +4798,31 @@ def move_install_dir_aside(src: Path, dst: Path) -> None: path is on a different overlay) fall back to copy + remove. A busy/in-use failure is deliberately NOT copy-faked here: the source is a live install and a partial copy + rmtree would be worse than failing, so it re-raises. + + The copy never writes into ``dst`` directly: callers treat ``dst.exists()`` + as proof of a complete tree (activation recovery restores a rollback dir + whenever it exists), so a copy that dies halfway (ENOSPC, I/O error) must + not leave a partial tree at ``dst``. Copy to a temp sibling and publish it + with one atomic rename; on failure remove the temp copy and leave ``src`` + untouched. """ try: os.replace(src, dst) except OSError as exc: if not is_cross_device_error(exc): raise - log(f"os.replace cross-device ({exc!r}); copy+remove {src} -> {dst}") - shutil.copytree(src, dst, dirs_exist_ok = True) + copy_tmp = dst.with_name(dst.name + ".copying") + counter = 0 + while copy_tmp.exists(): + counter += 1 + copy_tmp = dst.with_name(f"{dst.name}.copying-{counter}") + log(f"os.replace cross-device ({exc!r}); copy+publish {src} -> {dst}") + try: + shutil.copytree(src, copy_tmp) + os.replace(copy_tmp, dst) + except BaseException: + remove_tree(copy_tmp) + raise remove_tree(src) diff --git a/tests/python/test_unsloth_pip_shim.py b/tests/python/test_unsloth_pip_shim.py index dcb617255d..11e3fb89c2 100644 --- a/tests/python/test_unsloth_pip_shim.py +++ b/tests/python/test_unsloth_pip_shim.py @@ -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 ` 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 == [] diff --git a/tests/sh/test_select_cuda_jit_tools.sh b/tests/sh/test_select_cuda_jit_tools.sh index 578e0c2cb5..7838b835af 100755 --- a/tests/sh/test_select_cuda_jit_tools.sh +++ b/tests/sh/test_select_cuda_jit_tools.sh @@ -41,6 +41,9 @@ assert_eq() { # $1 = compute_cap(s) the mock nvidia-smi reports, ONE PER LINE ("none" -> no # nvidia-smi on PATH). A multi-line value models a mixed-GPU host so we can check # that every visible cap is scanned, not just the first. +# $2 (optional) = the target libnvrtc.so.12 starts on; defaults to the baked +# cu12.8 default, and "libnvrtc.so.12.cu13" models the stale link an earlier +# sm_103/sm_121 boot left in the same container's writable layer. # Builds a fake Studio venv NVRTC dir exactly as the build stages it: the real # cu12.8 lib as .cu128.orig, libnvrtc.so.12 -> it (the immutable default), and a # .cu13 alias pointing at a stand-in cu13 lib. Runs the function against it via @@ -48,6 +51,7 @@ assert_eq() { # host, so its glob is skipped. Prints " ". run_select() { _cap="$1" + _init="${2:-libnvrtc.so.12.cu128.orig}" _tmp=$(mktemp -d) mkdir -p "$_tmp/bin" if [ "$_cap" != "none" ]; then @@ -62,7 +66,7 @@ run_select() { : > "$_nvrtc/libnvrtc.so.12.cu128.orig" # real cu12.8 lib : > "$_nvrtc/libnvrtc.so.13.stub" # stand-in cu13 lib ln -sf libnvrtc.so.13.stub "$_nvrtc/libnvrtc.so.12.cu13" # staged cu13 alias - ln -sf libnvrtc.so.12.cu128.orig "$_nvrtc/libnvrtc.so.12" # immutable cu12.8 default + ln -sf "$_init" "$_nvrtc/libnvrtc.so.12" # cu12.8 default (or stale cu13) bash -c ' set -euo pipefail export PATH="'"$_tmp"'/bin:/usr/bin:/bin" @@ -101,6 +105,14 @@ assert_eq "B200 then GB10 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" assert_eq "B300 then H100 -> cu13 NVRTC selected" "UNSET libnvrtc.so.12.cu13" "$(run_select "$(printf '10.3\n9.0')")" assert_eq "H100 then A100 -> cu128 default kept" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select "$(printf '9.0\n8.0')")" +# Stateful transition: a cu13 selection left in the same container's writable +# layer by an earlier sm_103/sm_121 boot must be reversed when the container +# later starts on an ordinary GPU (or none) -- a 570-579 driver cannot load +# cu13-produced cubins -- and kept when the datacenter Blackwell is still there. +assert_eq "A100 after B300 -> cu128 restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select 8.0 libnvrtc.so.12.cu13)" +assert_eq "no GPU after B300 -> cu128 restored" "UNSET libnvrtc.so.12.cu128.orig" "$(run_select none libnvrtc.so.12.cu13)" +assert_eq "B300 after B300 -> cu13 kept" "UNSET libnvrtc.so.12.cu13" "$(run_select 10.3 libnvrtc.so.12.cu13)" + rm -f "$_FUNC_FILE" echo ""