From 837b09122e3d2fcf3736f1d4fff073bcf0c42fc2 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 27 Jul 2026 14:02:16 +0000 Subject: [PATCH] docker: close five failure paths the review found Notebook sync, in-place publish. entrypoint.sh runs sync_notebooks and then execs the container command, so the detached refresh child is still copying while JupyterLab serves the same tree. cp -a writes through the destination inode, so a reader can catch half-written JSON and a save made after the recorded-hash check is destroyed and then recorded as pristine. Publish through a same-dir dot-prefixed temp plus an atomic rename, and re-read the hash once the staging copy is complete (the earlier check sits before middle_unchanged, a python subprocess, so the window was most of the loop). A single-file bind mount cannot be renamed over, so that path falls back to the previous copy. Notebook sync, first boot. A pre-existing file whose bytes already match the baked template fell through to cp -a, which is --preserve=all: as root that stamps root:root, the baked mode and the build mtime onto a bind-mounted host file and locks its owner out of editing it. Record it as managed instead. The hash is identical, so the state file is byte-for-byte what the copy wrote. unsloth-studio-update. The post-update import check only warned, then the default restart replaced a process that was serving fine with one known not to import. supervisord retries startretries times, lands in FATAL and never leaves it on its own, so the container serves nothing until someone execs in. Keep the running service and exit non-zero with the remedy. unsloth-llama-update --check. resolve_latest swallows every failure into an empty string, which fell into the "up to date" branch and exited 0, so the command reported a state it could not observe. Report UNKNOWN and fail. unsloth-llama-update rollback. The in-place restore iterates the backup's entries, so a file the new release introduced survives it and the restored tree is mixed-version; ggml dlopens every libggml-*.so next to the binaries. Clear the install dir before restoring, gated on the drain having completed, because before that an entry there can still be the only copy of an old file. docker-publish ref freeze. git ls-remote exits 0 whether or not a ref matched, so a non-zero exit means the remote was never reached. That exit was lost twice over: first element of a pipeline, and a run step with no explicit shell runs under bash -e without pipefail. The step exited 0 and published ref=main, which the amd64, arm64 and Studio builds each resolve again, so one multi-arch tag could carry different revisions. Fail the prepare job instead, keeping the passthrough for the reachable-but-no-match case it was written for. Jupyter output select. lastPointerOutput was only replaced by another pointer-down, but J/K/arrow cell navigation fires none, so Ctrl/Cmd+A on a later cell selected the previously clicked output and suppressed notebook:select-all; after a re-run the node is detached and the chord did nothing at all. Revalidate the remembered output (still in the document, still in the active cell) before using it as the fallback. Tests: four static guards in test_docker_nb_sync_race.py, a new behavioural test_docker_update_helpers.py driving both helpers with stub pip, supervisorctl and mv, a new test_docker_publish_ref_freeze.py that executes each resolver step under bash -e with a failing ls-remote, and a source check in validate_studio_features.py. Each fails against the code before this change; the interrupted-drain case also fails against the unconditional form of the rollback fix. --- .github/workflows/docker-publish.yml | 29 +- .../unsloth_labext/src/outputSelect.ts | 16 +- docker/unsloth_llama_update.sh | 27 +- docker/unsloth_studio_update.sh | 10 +- docker/unsloth_sync_notebooks.sh | 42 ++- tests/python/test_docker_nb_sync_race.py | 56 ++++ .../python/test_docker_publish_ref_freeze.py | 132 +++++++++ tests/python/test_docker_update_helpers.py | 250 ++++++++++++++++++ tests/validate_studio_features.py | 9 + 9 files changed, 558 insertions(+), 13 deletions(-) create mode 100644 tests/python/test_docker_publish_ref_freeze.py create mode 100644 tests/python/test_docker_update_helpers.py diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml index d6f9a1d905..0382f3dc51 100644 --- a/.github/workflows/docker-publish.yml +++ b/.github/workflows/docker-publish.yml @@ -105,7 +105,17 @@ jobs: if printf '%s' "$REF" | grep -Eq '^[0-9a-f]{40}$'; then SHA="$REF" else - SHA="$(git ls-remote https://github.com/unslothai/unsloth "$REF" | awk 'NR==1{print $1}')" + # ls-remote exits 0 whether or not a ref matched, so a non-zero exit + # means we never reached the remote. The pipe into awk would hide it + # (no pipefail under the default `bash -e` shell) and the fallback + # below would then hand a MUTABLE name to the amd64, arm64 and Studio + # builds, which each resolve it again -- the exact split this job + # exists to prevent. Fail the run instead. + if ! LS_OUT="$(git ls-remote https://github.com/unslothai/unsloth "$REF")"; then + echo "::error::unslothai/unsloth unreachable; cannot freeze ref '${REF}' to a sha" + exit 1 + fi + SHA="$(printf '%s\n' "$LS_OUT" | awk 'NR==1{print $1}')" [ -n "$SHA" ] || SHA="$REF" fi echo "ref=${SHA}" >> "$GITHUB_OUTPUT" @@ -129,7 +139,14 @@ jobs: if printf '%s' "$REF" | grep -Eq '^[0-9a-f]{40}$'; then SHA="$REF" else - SHA="$(git ls-remote https://github.com/unslothai/unsloth-zoo "$REF" | awk 'NR==1{print $1}')" + # Same rule as the unsloth ref above: a non-zero ls-remote is a + # transport failure, not "no such ref", and forwarding the branch + # name would let the three builds each pick a different commit. + if ! LS_OUT="$(git ls-remote https://github.com/unslothai/unsloth-zoo "$REF")"; then + echo "::error::unslothai/unsloth-zoo unreachable; cannot freeze ref '${REF}' to a sha" + exit 1 + fi + SHA="$(printf '%s\n' "$LS_OUT" | awk 'NR==1{print $1}')" [ -n "$SHA" ] || SHA="$REF" fi echo "ref=${SHA}" >> "$GITHUB_OUTPUT" @@ -146,7 +163,13 @@ jobs: if printf '%s' "$REF" | grep -Eq '^[0-9a-f]{40}$'; then SHA="$REF" else - SHA="$(git ls-remote https://github.com/unslothai/notebooks "$REF" | awk 'NR==1{print $1}')" + # Same rule as the two refs above: only a reachable remote with no + # matching ref may fall through to the literal "$REF". + if ! LS_OUT="$(git ls-remote https://github.com/unslothai/notebooks "$REF")"; then + echo "::error::unslothai/notebooks unreachable; cannot freeze ref '${REF}' to a sha" + exit 1 + fi + SHA="$(printf '%s\n' "$LS_OUT" | awk 'NR==1{print $1}')" [ -n "$SHA" ] || SHA="$REF" fi echo "commit=${SHA}" >> "$GITHUB_OUTPUT" diff --git a/docker/jupyter/unsloth_labext/src/outputSelect.ts b/docker/jupyter/unsloth_labext/src/outputSelect.ts index 54074c2d1a..1c31ade442 100644 --- a/docker/jupyter/unsloth_labext/src/outputSelect.ts +++ b/docker/jupyter/unsloth_labext/src/outputSelect.ts @@ -64,6 +64,20 @@ const outputSelectPlugin: JupyterFrontEndPlugin = { // Remember the last pointer-down: a click on an image/widget output leaves no // text selection, so the anchor alone can't tell which output is meant. let lastPointerOutput: HTMLElement | null = null; + // ...but only trust it while that output is still in the document AND still + // inside the ACTIVE cell. Keyboard cell navigation (J/K, arrows) fires no + // pointer event, so an unvalidated value would make the chord on a later cell + // select the previously clicked output and swallow `notebook:select-all`; and + // a re-executed cell replaces the node, leaving a detached range that selects + // nothing at all while still suppressing the shortcut. + const rememberedOutput = (): HTMLElement | null => { + const output = lastPointerOutput; + if (!output || !output.isConnected) { + return null; + } + const cell = output.closest('.jp-Cell'); + return cell && cell.classList.contains('jp-mod-active') ? output : null; + }; document.addEventListener( 'pointerdown', (event: PointerEvent): void => { @@ -85,7 +99,7 @@ const outputSelectPlugin: JupyterFrontEndPlugin = { // Own the chord only when in an output: the target, else the last click // (not the stale selection anchor; see the header). const output = - closestOutput(event.target as Node | null) ?? lastPointerOutput; + closestOutput(event.target as Node | null) ?? rememberedOutput(); if (!output) { return; } diff --git a/docker/unsloth_llama_update.sh b/docker/unsloth_llama_update.sh index 7b6dfc984e..8ab0e1f068 100755 --- a/docker/unsloth_llama_update.sh +++ b/docker/unsloth_llama_update.sh @@ -87,7 +87,16 @@ echo "[llama-update] installed: $CUR" if [ "$CHECK_ONLY" = "1" ]; then LATEST="$(resolve_latest)" echo "[llama-update] latest: ${LATEST:-unknown}" - if [ -n "$LATEST" ] && [ "$LATEST" != "$CUR" ]; then + # resolve_latest swallows every failure into "" (line 75), so an empty value + # means the lookup did not happen -- no network, proxy, GitHub down. Printing + # "up to date" there is the one answer --check must never give: it reports a + # state it could not observe. Say unknown and exit non-zero instead. + if [ -z "$LATEST" ]; then + echo "[llama-update] could not reach the release feed; update status UNKNOWN" >&2 + echo "[llama-update] (retry once the container has network access)" >&2 + exit 1 + fi + if [ "$LATEST" != "$CUR" ]; then echo "[llama-update] an update is available (run without --check to apply)" else echo "[llama-update] up to date" @@ -121,6 +130,7 @@ else backup="${INSTALL_DIR}.old.$$" fi swap_done=0 +drained=0 # The exit handler must never delete $backup while it's the ONLY copy: restore the # old tree first, remove it only after the new tree is active. Signal traps run # the EXIT trap on HUP/INT/TERM too. @@ -132,6 +142,18 @@ cleanup() { # half-moved NEW one: drop it, then move the old one back. if [ -d "$backup" ]; then _restore_fail=0 + # The per-name loop below only sees entries the OLD tree had, so a + # file the new release introduced survives it and the "restored" + # dir ends up mixed-version -- ggml dlopens every libggml-*.so it + # finds next to the binaries. Once the drain finished, every + # remaining entry is a half-moved NEW one, so clear them all. + # Gated on "drained": before the drain completes an entry here can + # still be the ONLY copy of an old one, and deleting it loses data. + if [ "$drained" = "1" ]; then + find "$INSTALL_DIR" -mindepth 1 -maxdepth 1 \ + ! -path "$work" ! -path "$backup" \ + -exec rm -rf {} + 2>/dev/null || true + fi for _e in "$backup"/* "$backup"/.[!.]* "$backup"/..?*; do { [ -e "$_e" ] || [ -L "$_e" ]; } || continue _b="$(basename "$_e")" @@ -176,6 +198,9 @@ if [ "$IN_PLACE" = "1" ]; then mkdir "$backup" find "$INSTALL_DIR" -mindepth 1 -maxdepth 1 \ ! -path "$work" ! -path "$backup" -exec mv -t "$backup" {} + + # Every old entry now lives in $backup, so from here the trap may clear the + # install dir before restoring. set -e means a failed drain never gets here. + drained=1 if find "$new" -mindepth 1 -maxdepth 1 -exec mv -t "$INSTALL_DIR" {} +; then swap_done=1 else diff --git a/docker/unsloth_studio_update.sh b/docker/unsloth_studio_update.sh index fbda364fe4..8855e7b19e 100755 --- a/docker/unsloth_studio_update.sh +++ b/docker/unsloth_studio_update.sh @@ -91,11 +91,17 @@ fi echo "[studio-update] after: unsloth $(version_of)" # Sanity: the backend must still import after the swap (a missing --no-deps -# transitive dep shows up here). Non-fatal: just warn with the remedy. +# transitive dep shows up here). Restarting into code that cannot import kills a +# process that is serving fine and leaves supervisord's studio program in FATAL +# after startretries, which it never leaves on its own. Keep the running service +# and fail instead, so the operator can add the dep or roll back with Studio up. if ! "$PY" -c "import studio.backend.main" >/dev/null 2>&1; then - echo "[studio-update] WARNING: 'import studio.backend.main' failed after update." >&2 + echo "[studio-update] ERROR: 'import studio.backend.main' failed after update." >&2 echo "[studio-update] A new dependency may be missing. Re-run with --with-deps:" >&2 echo "[studio-update] unsloth-studio-update --with-deps" >&2 + echo "[studio-update] NOT restarting Studio: the running process keeps serving." >&2 + echo "[studio-update] Once fixed: supervisorctl restart studio" >&2 + exit 1 fi if [ "$RESTART" = "1" ]; then diff --git a/docker/unsloth_sync_notebooks.sh b/docker/unsloth_sync_notebooks.sh index 522299d3a6..2c23daf984 100644 --- a/docker/unsloth_sync_notebooks.sh +++ b/docker/unsloth_sync_notebooks.sh @@ -184,9 +184,17 @@ if [ ! -f "$STATE" ]; then # A pre-existing file (bind-mounted or hand-created) is user data: keep it # and do NOT record it, else the refresh below would treat it as pristine # and overwrite it. Only files we lay down are recorded as managed. - if [ -e "$DEST/$rel" ] \ - && [ "$(hash_of "$DEST/$rel")" != "$(hash_of "$TEMPLATE/$rel")" ]; then - echo "[unsloth-nb] kept existing user file: $DEST/$rel" + if [ -e "$DEST/$rel" ]; then + if [ "$(hash_of "$DEST/$rel")" != "$(hash_of "$TEMPLATE/$rel")" ]; then + echo "[unsloth-nb] kept existing user file: $DEST/$rel" + continue + fi + # Same bytes already on disk (a bind-mounted checkout of the same + # notebooks). cp -a is --preserve=all, so copying would only stamp the + # baked root:root ownership, mode and build mtime onto the host user's + # own file and lock them out of editing it. Record it as managed -- the + # hash is identical, so the state is byte-for-byte what cp would write. + printf '%s %s\n' "$(hash_of "$DEST/$rel")" "$rel" >> "$STATE.tmp" continue fi if cp -a "$TEMPLATE/$rel" "$DEST/$rel" 2>/dev/null; then @@ -300,9 +308,31 @@ while IFS= read -r -d '' f; do continue fi mkdir -p "$(dirname "$dst")" 2>/dev/null || true - if cp -a "$f" "$dst" 2>/dev/null; then - printf '%s %s\n' "$(hash_of "$dst")" "$rel" >> "$TMPSTATE" - updated=$((updated + 1)) + # Publish through a same-dir temp + rename. This child is forked before the + # entrypoint execs the container command, so JupyterLab is already serving + # $DEST while this loop runs: cp -a writes in place (the inode is reused), so + # a reader can catch half-written JSON, and a save made between the recorded- + # hash check above and this write is destroyed and then recorded as pristine. + # rename(2) is atomic, and re-reading the hash once the temp is complete + # shrinks the check-to-write window to the rename itself. The staging name is + # dot-prefixed and per-PID so a killed refresh leaves nothing visible in the + # file browser; unsloth_nb_strip_colab.py already publishes these same files + # this way. + new="$(dirname "$dst")/.unsloth_nb_new.$$" + if cp -a "$f" "$new" 2>/dev/null; then + if [ -e "$dst" ] && [ "$(hash_of "$dst")" != "${LAST[$rel]:-}" ]; then + # Saved while we were copying -> their edit wins, keep the marker. + rm -f "$new" + printf '%s %s\n' "${LAST[$rel]:-}" "$rel" >> "$TMPSTATE" + kept=$((kept + 1)) + continue + fi + # A single-FILE bind mount cannot be renamed over (EBUSY); fall back to the + # previous in-place copy there so that setup keeps working as it does today. + if mv -f "$new" "$dst" 2>/dev/null || { rm -f "$new"; cp -a "$f" "$dst" 2>/dev/null; }; then + printf '%s %s\n' "$(hash_of "$dst")" "$rel" >> "$TMPSTATE" + updated=$((updated + 1)) + fi fi done < <(find "$TMP" -type f -print0) diff --git a/tests/python/test_docker_nb_sync_race.py b/tests/python/test_docker_nb_sync_race.py index de2c26dca6..1365cc1235 100644 --- a/tests/python/test_docker_nb_sync_race.py +++ b/tests/python/test_docker_nb_sync_race.py @@ -135,3 +135,59 @@ def test_the_lock_lives_beside_the_state_it_protects(sync: str): "keeping the lock in $DEST also serialises two containers sharing the " "notebooks volume, which /tmp would not" ) + + +# --- concurrent-publish safety ------------------------------------------------ +# The detach above is deliberate, but entrypoint.sh runs `sync_notebooks` and then +# `exec "$@"`, so the child is still copying while JupyterLab serves the same tree. +# `cp -a` writes THROUGH the destination inode, so it both exposes half-written +# JSON to a reader and destroys a save made after the recorded-hash check. The +# publish therefore has to go via a same-dir temp plus an atomic rename. + + +def test_the_refresh_publishes_each_notebook_atomically(sync: str): + block = sync[sync.index("while IFS= read -r -d '' f; do") :] + block = block[: block.index("done < <(find")] + assert re.search(r'cp -a "\$f" "\$new"', block), ( + "the refresh must copy into a staging file, not onto the live notebook" + ) + assert re.search(r'mv -f "\$new" "\$dst"', block), ( + "the staged copy must be published with an atomic rename" + ) + + +def test_the_staging_file_is_hidden_and_beside_the_destination(sync: str): + assert re.search(r'new="\$\(dirname "\$dst"\)/\.unsloth_nb_new\.\$\$"', sync), ( + "the staging file must be dot-prefixed (invisible in the file browser), " + "per-PID (two containers on one volume) and in the destination directory " + "(a rename cannot cross filesystems)" + ) + + +def test_the_recorded_hash_is_rechecked_immediately_before_publishing(sync: str): + block = sync[sync.index("while IFS= read -r -d '' f; do") :] + block = block[: block.index("done < <(find")] + recheck = block.index('cp -a "$f" "$new"') + assert re.search( + r'if \[ -e "\$dst" \] && \[ "\$\(hash_of "\$dst"\)" != "\$\{LAST\[\$rel\]:-\}" \]', + block[recheck:], + ), ( + "the earlier check sits before middle_unchanged (a python subprocess), so " + "the hash has to be re-read once the staging copy is complete or a save " + "made in between is silently overwritten" + ) + + +def test_a_pristine_pre_existing_file_is_not_rewritten_on_first_boot(sync: str): + block = sync[sync.index('if [ ! -f "$STATE" ]; then') :] + block = block[: block.index('mv "$STATE.tmp" "$STATE"')] + assert "kept existing user file" in block + # A bind-mounted file whose bytes already match the template used to fall + # through to `cp -a`, i.e. --preserve=all stamping root:root, the baked mode + # and the build mtime onto the host user's own file. Record, don't copy. + same = block.index("kept existing user file") + tail = block[same:] + assert tail.index("$STATE.tmp") < tail.index('cp -a "$TEMPLATE/$rel"'), ( + "an existing file with the template's exact bytes must be recorded as " + "managed without being copied over" + ) diff --git a/tests/python/test_docker_publish_ref_freeze.py b/tests/python/test_docker_publish_ref_freeze.py new file mode 100644 index 0000000000..b7133148b5 --- /dev/null +++ b/tests/python/test_docker_publish_ref_freeze.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""The docker publish workflow must never forward an unfrozen ref. + +`prepare` resolves unsloth, unsloth-zoo and notebooks to ONE commit each so the +amd64 leg, the arm64 leg and the Studio build all bake identical source; that is +the whole reason the job exists. Each resolver was + + SHA="$(git ls-remote "$REF" | awk 'NR==1{print $1}')" + [ -n "$SHA" ] || SHA="$REF" + +`git ls-remote` exits 0 whether or not a ref matched, so a non-zero exit means +the remote was never reached. That exit was lost twice over: it is the first +element of a pipeline, and a `run:` step with no explicit `shell:` runs under +`bash -e` WITHOUT pipefail, so the step exited 0 and published `ref=main`. Each +build then resolved `main` independently, and a branch advance between them +would ship one multi-arch tag containing different revisions. The stable-tag +gates key off the inputs, not off whether resolution worked, so `:latest` would +still be moved onto it. + +Static plus behavioural: the resolver `run:` blocks are executed under `bash -e` +with a `git` stub. No docker, no network. +""" + +from __future__ import annotations + +import os +import re +import shutil +import subprocess +from pathlib import Path + +import pytest +import yaml + +REPO_ROOT = Path(__file__).resolve().parents[2] +WORKFLOW = REPO_ROOT / ".github" / "workflows" / "docker-publish.yml" + +RESOLVER_STEPS = ("unsloth_ref", "zoo_ref", "notebooks") + +pytestmark = pytest.mark.skipif( + shutil.which("bash") is None, reason = "needs bash", +) + + +@pytest.fixture(scope = "module") +def steps() -> dict: + assert WORKFLOW.is_file(), f"missing {WORKFLOW}" + doc = yaml.safe_load(WORKFLOW.read_text(encoding = "utf-8")) + found = {} + for step in doc["jobs"]["prepare"]["steps"]: + if step.get("id") in RESOLVER_STEPS: + found[step["id"]] = step["run"] + missing = set(RESOLVER_STEPS) - set(found) + assert not missing, f"resolver steps missing from the prepare job: {missing}" + return found + + +def test_the_workflow_never_pins_a_shell_so_bash_e_has_no_pipefail(steps: dict): + # If someone later adds `shell: bash` the runner switches to + # `bash --noprofile --norc -eo pipefail`, which would make the guards below + # redundant rather than wrong -- but until then they are the only protection. + doc = yaml.safe_load(WORKFLOW.read_text(encoding = "utf-8")) + assert "shell" not in doc.get("defaults", {}).get("run", {}), ( + "this test models the default `bash -e` shell; update it if a default " + "shell with pipefail is introduced" + ) + + +@pytest.mark.parametrize("step_id", RESOLVER_STEPS) +def test_an_unreachable_remote_fails_the_step(steps: dict, step_id: str, tmp_path: Path): + script = _expand(steps[step_id]) + res = _run_with_failing_ls_remote(script, tmp_path) + assert res.returncode != 0, ( + "a transport failure must fail the prepare job, not fall through to the " + f"mutable ref:\nstdout={res.stdout}\nstderr={res.stderr}" + ) + + +@pytest.mark.parametrize("step_id", RESOLVER_STEPS) +def test_an_unreachable_remote_never_emits_a_mutable_ref( + steps: dict, step_id: str, tmp_path: Path, +): + script = _expand(steps[step_id]) + res = _run_with_failing_ls_remote(script, tmp_path) + emitted = (tmp_path / "github_output").read_text(encoding = "utf-8") \ + if (tmp_path / "github_output").exists() else "" + for line in emitted.splitlines(): + key, _, value = line.partition("=") + assert re.fullmatch(r"[0-9a-f]{40}", value), ( + f"{step_id} published {key}={value!r}, which the three builds each " + "resolve again, so they can bake different revisions" + ) + assert res.returncode != 0 + + +def _expand(run: str) -> str: + """Replace the `${{ ... }}` expressions with the empty string the default + (push to main, no dispatch inputs) trigger produces.""" + return re.sub(r"\$\{\{[^}]*\}\}", "", run) + + +def _run_with_failing_ls_remote(script: str, tmp_path: Path): + bin_dir = tmp_path / "bin" + bin_dir.mkdir(parents = True, exist_ok = True) + stub = bin_dir / "git" + stub.write_text( + "#!/usr/bin/env bash\n" + 'if [ "$1" = "ls-remote" ]; then\n' + ' echo "fatal: unable to access: Could not resolve host" >&2\n' + " exit 128\n" + "fi\n" + "exit 0\n", + encoding = "utf-8", + ) + stub.chmod(0o755) + out = tmp_path / "github_output" + out.write_text("", encoding = "utf-8") + env = dict(os.environ) + env["PATH"] = f"{bin_dir}{os.pathsep}" + env["PATH"] + env["GITHUB_OUTPUT"] = str(out) + # Whatever the expansions above blanked out; the resolvers default to "main". + for name in ("INPUT_REF", "TAG_REF", "PUSH_SHA"): + env[name] = "" + path = tmp_path / "step.sh" + path.write_text(script, encoding = "utf-8") + # Exactly how the runner invokes a `run:` step with no explicit `shell:`. + return subprocess.run( + ["bash", "-e", str(path)], + capture_output = True, text = True, env = env, timeout = 60, + ) diff --git a/tests/python/test_docker_update_helpers.py b/tests/python/test_docker_update_helpers.py new file mode 100644 index 0000000000..097930a497 --- /dev/null +++ b/tests/python/test_docker_update_helpers.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Behavioural guards for the two in-container update helpers of the Docker image. + +Both are `docker exec` entry points that mutate a running container, so a wrong +answer costs an outage or a mixed-version install: + +* `unsloth-studio-update` swaps the Studio Python packages and then restarts the + service. It verifies the new backend imports first, but only warned -- so a + release that pulls in a dependency `--no-deps` did not install got the healthy + old process killed and replaced by one that cannot start. supervisord retries + `startretries` times, lands in FATAL and never leaves it on its own, so the + container serves nothing until someone exec's in. +* `unsloth-llama-update --check` reported "up to date" when it could not reach + the release feed at all, and its in-place rollback only removed entries whose + names the OLD tree also had, leaving new-release-only shared objects beside + the restored files. ggml dlopen()s every `libggml-*.so` next to the binaries, + so that mix is loaded on the next GGUF run. + +These drive the real scripts with stub `pip` / `supervisorctl` / `python` / +`mv` on PATH. No docker, no GPU, no network. +""" + +from __future__ import annotations + +import os +import shutil +import subprocess +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STUDIO_UPDATE = REPO_ROOT / "docker" / "unsloth_studio_update.sh" +LLAMA_UPDATE = REPO_ROOT / "docker" / "unsloth_llama_update.sh" + +pytestmark = pytest.mark.skipif( + shutil.which("bash") is None, reason = "needs bash", +) + + +def _stub(directory: Path, name: str, body: str) -> None: + directory.mkdir(parents = True, exist_ok = True) + path = directory / name + path.write_text("#!/usr/bin/env bash\n" + body, encoding = "utf-8") + path.chmod(0o755) + + +def _run(script: Path, args, env, cwd = None): + return subprocess.run( + ["bash", str(script), *args], + capture_output = True, text = True, env = env, cwd = cwd, timeout = 120, + ) + + +# --- unsloth-studio-update ---------------------------------------------------- + +def _studio_env(tmp_path: Path, *, import_ok: bool) -> dict: + home = tmp_path / "studio" + venv_bin = home / "unsloth_studio" / "bin" + venv_bin.mkdir(parents = True) + _stub( + venv_bin, "python", + 'if [ "$1" = "-c" ]; then\n' + + (" exit 0\n" if import_ok else ' case "$2" in *studio.backend.main*) exit 1;; esac\n exit 0\n') + + 'fi\n' + 'if [ "$1" = "-m" ] && [ "$2" = "pip" ]; then\n' + ' if [ "$3" = "show" ]; then echo "Version: 2026.7.5"; exit 0; fi\n' + ' echo "STUB-PIP $*" >> "$STUB_LOG"; exit 0\n' + 'fi\n' + 'exit 0\n', + ) + bin_dir = tmp_path / "bin" + _stub(bin_dir, "supervisorctl", + 'echo "STUB-SUPERVISORCTL $*" >> "$STUB_LOG"\n' + 'if [ "$1" = "status" ]; then exit 0; fi\nexit 0\n') + env = dict(os.environ) + env["PATH"] = f"{bin_dir}{os.pathsep}" + env["PATH"] + env["UNSLOTH_STUDIO_HOME"] = str(home) + env["STUB_LOG"] = str(tmp_path / "calls.log") + return env + + +def test_studio_update_restarts_when_the_backend_imports(tmp_path: Path): + env = _studio_env(tmp_path, import_ok = True) + res = _run(STUDIO_UPDATE, [], env) + calls = Path(env["STUB_LOG"]).read_text() if Path(env["STUB_LOG"]).exists() else "" + assert res.returncode == 0, res.stderr + assert "STUB-SUPERVISORCTL restart studio" in calls, calls + + +def test_studio_update_does_not_restart_into_a_backend_that_cannot_import(tmp_path: Path): + env = _studio_env(tmp_path, import_ok = False) + res = _run(STUDIO_UPDATE, [], env) + calls = Path(env["STUB_LOG"]).read_text() if Path(env["STUB_LOG"]).exists() else "" + assert "STUB-SUPERVISORCTL restart studio" not in calls, ( + "restarting into code that cannot import kills a process that is serving " + "fine and parks supervisord's studio program in FATAL:\n" + calls + ) + assert res.returncode != 0, "a broken update must not report success" + assert "--with-deps" in res.stderr, "the remedy must still be printed" + + +# --- unsloth-llama-update ----------------------------------------------------- + +def _llama_env(tmp_path: Path, *, latest: str | None) -> dict: + install = tmp_path / "llama.cpp" + install.mkdir(parents = True) + (install / "UNSLOTH_PREBUILT_INFO.json").write_text( + '{"tag": "b1111-old"}\n', encoding = "utf-8", + ) + fetcher = tmp_path / "fetch_llama_prebuilt.py" + resolve = ( + " raise RuntimeError('unreachable')\n" if latest is None + else f" return {latest!r}\n" + ) + fetcher.write_text( + "def resolve_latest_tag(repo):\n" + resolve, encoding = "utf-8", + ) + env = dict(os.environ) + env["UNSLOTH_LLAMA_CPP_PATH"] = str(install) + env["UNSLOTH_LLAMA_FETCHER"] = str(fetcher) + return env + + +def _llama_check(tmp_path: Path, latest): + env = _llama_env(tmp_path, latest = latest) + return _run(LLAMA_UPDATE, ["--check"], env) + + +def test_llama_check_reports_an_available_update(tmp_path: Path): + res = _llama_check(tmp_path, "b2222-new") + assert res.returncode == 0, res.stderr + assert "an update is available" in res.stdout + + +def test_llama_check_reports_up_to_date(tmp_path: Path): + res = _llama_check(tmp_path, "b1111-old") + assert res.returncode == 0, res.stderr + assert "up to date" in res.stdout + + +def test_llama_check_does_not_claim_up_to_date_when_it_could_not_look(tmp_path: Path): + res = _llama_check(tmp_path, None) + assert "up to date" not in res.stdout, ( + "--check exists to report update status; saying 'up to date' for a lookup " + "that never happened is the one answer it must never give:\n" + res.stdout + ) + assert res.returncode != 0, "an unperformed check must not exit 0" + assert "UNKNOWN" in res.stdout + res.stderr + + +def _llama_inplace_env(tmp_path: Path, old: list[str], new: list[str]) -> dict: + """An in-place (volume-mounted) install whose activation fails part-way.""" + install = tmp_path / "llama.cpp" + install.mkdir(parents = True) + for name in old: + (install / name).write_text("OLD\n", encoding = "utf-8") + (install / "UNSLOTH_PREBUILT_INFO.json").write_text( + '{"tag": "b1111-old"}\n', encoding = "utf-8", + ) + fetcher = tmp_path / "fetch_llama_prebuilt.py" + fetcher.write_text( + "import os, sys\n" + "def resolve_latest_tag(repo):\n" + " return 'b2222-new'\n" + "if __name__ == '__main__':\n" + " dest = sys.argv[3]\n" + " os.makedirs(dest, exist_ok = True)\n" + f" for name in {new!r}:\n" + " open(os.path.join(dest, name), 'w').write('NEW\\n')\n" + " open(os.path.join(dest, 'UNSLOTH_PREBUILT_INFO.json'), 'w')" + ".write('{\"tag\": \"b2222-new\"}\\n')\n", + encoding = "utf-8", + ) + # Fail the ACTIVATION move (-t ) AFTER it has moved the files, so + # the install dir is populated with the new tree and `find` still reports the + # failure -- the mid-swap abort the rollback exists for. The drain + # (-t ) and the rollback's own per-file moves must keep working, so + # only that one invocation is broken. + bin_dir = tmp_path / "bin" + _stub( + bin_dir, "mv", + 'if [ "$1" = "-t" ] && [ "$2" = "$FAIL_MV_TARGET" ]; then\n' + ' shift 2\n' + ' for _s in "$@"; do /bin/mv "$_s" "$FAIL_MV_TARGET/"; done\n' + ' exit 1\n' + 'fi\n' + 'exec /bin/mv "$@"\n', + ) + env = dict(os.environ) + env["PATH"] = f"{bin_dir}{os.pathsep}" + env["PATH"] + env["UNSLOTH_LLAMA_CPP_PATH"] = str(install) + env["UNSLOTH_LLAMA_FETCHER"] = str(fetcher) + env["UNSLOTH_LLAMA_UPDATE_IN_PLACE"] = "1" + env["FAIL_MV_TARGET"] = str(install) + return env + + +def test_llama_rollback_leaves_no_new_release_files_behind(tmp_path: Path): + # "libggml-hexagon.so" exists only in the new release, so the rollback loop -- + # which iterates the BACKUP's entries -- cannot see it. ggml dlopen()s every + # libggml-*.so sitting next to the binaries, so a leftover is loaded against + # the restored older libggml-base.so. + old = ["libggml-base.so", "libggml-cpu-icelake.so", "llama-cli"] + new = ["libggml-base.so", "libggml-cpu-icelake.so", "llama-cli", + "libggml-hexagon.so", "llama-mtmd-cli"] + env = _llama_inplace_env(tmp_path, old, new) + res = _run(LLAMA_UPDATE, [], env) + assert res.returncode != 0, "a failed swap must not report success" + install = tmp_path / "llama.cpp" + present = sorted(p.name for p in install.iterdir()) + leftovers = [n for n in ("libggml-hexagon.so", "llama-mtmd-cli") if n in present] + assert not leftovers, ( + f"new-release-only files survived the rollback: {leftovers} in {present}" + ) + for name in old: + assert (install / name).read_text() == "OLD\n", ( + f"{name} was not restored from the backup: {present}" + ) + + +def test_llama_rollback_keeps_every_old_file_when_the_drain_is_interrupted(tmp_path: Path): + # The mirror image: abort while the OLD tree is still being moved into the + # backup. The entries left in the install dir are then the only copy of those + # old files, so clearing the directory before restoring would destroy them. + old = ["libggml-base.so", "libggml-cpu-icelake.so", "llama-cli", "llama-quantize"] + env = _llama_inplace_env(tmp_path, old, old) + install = tmp_path / "llama.cpp" + # Fail the DRAIN (-t /.old.) after moving only the first + # source, so half the old tree is still sitting in the install dir when the + # rollback runs. Those entries are then the only copy there is. + _stub( + tmp_path / "bin", "mv", + 'case "${1:-}:${2:-}" in\n' + ' -t:*/.old.*)\n' + ' _t="$2"; shift 2\n' + ' [ $# -gt 0 ] && /bin/mv "$1" "$_t/"\n' + ' exit 1;;\n' + 'esac\n' + 'exec /bin/mv "$@"\n', + ) + res = _run(LLAMA_UPDATE, [], env) + assert res.returncode != 0 + survivors = sorted(p.name for p in install.rglob("*") if p.is_file()) + for name in old: + assert name in survivors, ( + f"{name} was lost during an interrupted drain: {survivors}" + ) diff --git a/tests/validate_studio_features.py b/tests/validate_studio_features.py index 5c7f4338c7..73930f9017 100644 --- a/tests/validate_studio_features.py +++ b/tests/validate_studio_features.py @@ -226,6 +226,15 @@ def test_labext_and_branding() -> None: # uiChrome hides the right activity bar; CTRL+A output-select selects nodes. check("right activity bar hidden", "jp-mod-right" in all_src and "display: none" in all_src) check("ctrl+A output select", "selectNodeContents" in all_src) + # The remembered pointer-down is only replaced by another pointer-down, but + # J/K/arrow cell navigation fires none, so it has to be revalidated (still in + # the document, still in the ACTIVE cell) before it is used as the fallback -- + # otherwise Ctrl+A on a later cell selects the old output and swallows + # JupyterLab's notebook:select-all. + check( + "ctrl+A fallback revalidated", + "isConnected" in all_src and "jp-mod-active" in all_src, + ) # branding assets login = os.path.join(JUPYTER, "login.html") login_src = open(login, encoding = "utf-8").read() if os.path.isfile(login) else ""