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