diff --git a/docker/unsloth_nb_strip_colab.py b/docker/unsloth_nb_strip_colab.py index 95b89a72f0..4ffa67a0a1 100644 --- a/docker/unsloth_nb_strip_colab.py +++ b/docker/unsloth_nb_strip_colab.py @@ -33,11 +33,31 @@ _INTRO_PREFIX = "to run this, press" _WIDGET_VIEW_MIME = "application/vnd.jupyter.widget-view+json" +def _is_intro_line(line): + """True for the Colab run announcement in either shipped spelling. + + Most notebooks open the line with the sentence itself, but two (NeMo-Gym-*) + ship it inside a single-line HTML comment: + + + + Only a comment that OPENS AND CLOSES on the same line is matched, so + dropping it can never leave a dangling `"): + return stripped[4:-3].strip().lower().startswith(_INTRO_PREFIX) + return False + + def _strip_lines(lines): """Drop the intro line (and an immediately-following blank). Return new list or None if there was nothing to strip.""" for i, line in enumerate(lines): - if line.lstrip().lower().startswith(_INTRO_PREFIX): + if _is_intro_line(line): out = lines[:i] + lines[i + 1 :] if i < len(out) and out[i].strip() == "": out = out[:i] + out[i + 1 :] @@ -45,14 +65,8 @@ def _strip_lines(lines): return None -def _strip_intro(nb): - """Strip the Colab intro sentence from cells[0]. Return True if changed.""" - cells = nb.get("cells") - if not isinstance(cells, list) or not cells: - return False - cell = cells[0] - if not isinstance(cell, dict) or cell.get("cell_type") != "markdown": - return False +def _strip_cell(cell): + """Strip the intro line out of ONE markdown cell. Return True if changed.""" src = cell.get("source") if isinstance(src, str): lines = src.splitlines(keepends = True) @@ -69,6 +83,29 @@ def _strip_intro(nb): return True +def _strip_intro(nb): + """Strip the Colab intro sentence from the LEADING markdown block. + + Scanning cells[0] alone missed 23 of the 433 shipped notebooks: 21 put the + Colab badge `` in cells[0] and the sentence in cells[1] + (Advanced_Llama3_2_(3B)_GRPO_LoRA, Falcon_H1-Alpaca, gpt-oss-(20B)-GRPO, + ...), and 2 (NeMo-Gym-*) wrap it in an HTML comment cells[0]-only matching + never saw. The scan stops at the first non-markdown cell, so it only ever + touches the header block a notebook opens with (at most 5 cells across the + shipped set) and can never reach explanatory prose between code cells. + Return True if any cell changed.""" + cells = nb.get("cells") + if not isinstance(cells, list): + return False + changed = False + for cell in cells: + if not isinstance(cell, dict) or cell.get("cell_type") != "markdown": + break # the first code cell ends the header block + if _strip_cell(cell): + changed = True + return changed + + def _clean_widgets(nb): """Drop baked ipywidget outputs + the orphan widget-state metadata that otherwise render as "Loading widget...". Return True if changed.""" diff --git a/docker/unsloth_sync_notebooks.sh b/docker/unsloth_sync_notebooks.sh index a164effcf1..522299d3a6 100644 --- a/docker/unsloth_sync_notebooks.sh +++ b/docker/unsloth_sync_notebooks.sh @@ -31,7 +31,9 @@ DEST="${UNSLOTH_NOTEBOOKS_DIR:-/workspace/unsloth-notebooks}" REMOTE="${UNSLOTH_NOTEBOOKS_REPO:-https://github.com/unslothai/notebooks}" STATE="$DEST/.unsloth_sync_state" # "sha256 relpath" of what we last wrote SYNCED="$DEST/.unsloth_sync_commit" # upstream commit we last synced to +LOCK="$DEST/.unsloth_sync.lock" # serialises this script against itself TIMEOUT="${UNSLOTH_NOTEBOOK_FETCH_TIMEOUT:-60}" +LOCK_WAIT="${UNSLOTH_NOTEBOOK_LOCK_TIMEOUT:-600}" # Resolve a helper script ($1 override, $2 PATH command, $3 sibling filename), # echoing the path or nothing. Used for SIG, VIEW and STRIP helpers. @@ -64,6 +66,40 @@ mkdir -p "$DEST" 2>/dev/null || exit 0 hash_of() { sha256sum "$1" 2>/dev/null | cut -d' ' -f1; } +# --- mutual exclusion -------------------------------------------------------- +# Every phase below mutates $DEST and rewrites $STATE, and the GitHub refresh +# runs in a DETACHED child of this same script, so two copies are live at once by +# design. Without a lock the parent's strip/view pass interleaved with the child's +# `cp -a` + state rewrite: six identical boots reported "cleaned" 279/289/293/297/ +# 300/306/307/315/330 notebooks, and every notebook the child copied while the +# parent was hashing it ended up permanently marked user-edited (its recorded +# hash no longer matched), so it was skipped by every later strip. +# +# One exclusive lock covers a whole invocation. The child therefore cannot start +# until the parent has finished and exited, which also fixes the ORDER: strip and +# view rebuild always run over a quiesced tree. flock is best-effort -- when it is +# unavailable, or $DEST cannot hold the lock file, we fall back to running +# unlocked (the parent still finalizes before forking, see below). +_LOCK_HELD=0 +lock_acquire() { + [ "$_LOCK_HELD" = "1" ] && return 0 + command -v flock >/dev/null 2>&1 || return 0 + # Group-redirect, not `exec ... 2>/dev/null`: bash reports a failed exec + # redirection before the redirection it was given applies, so a read-only + # $DEST would print "Permission denied" into the container log. + { exec 9>>"$LOCK"; } 2>/dev/null || return 0 + flock -w "$LOCK_WAIT" 9 2>/dev/null || return 0 + _LOCK_HELD=1 + return 0 +} +lock_release() { + [ "$_LOCK_HELD" = "1" ] || return 0 + _LOCK_HELD=0 + flock -u 9 2>/dev/null || true + exec 9>&- 2>/dev/null || true + return 0 +} + # --- categorized folder view + Docker-only Colab cleanups -------------------- # AMD/HIP detection: AMD-*.ipynb are shown only on an AMD GPU. UNSLOTH_NB_GPU # forces it (amd|cuda); otherwise probe nvidia-smi then the ROCm tools. @@ -107,8 +143,22 @@ strip_colab_intros() { # Apply both on EVERY exit after the basic guards, so the view + cleanups also # run on the common "nothing to refresh" / offline paths. Both are idempotent. -finalize() { strip_colab_intros; build_categorized_view; } -trap finalize EXIT +# Run-once: the parent calls this explicitly BEFORE it forks the refresh child +# (so the strip can never overlap the child's copy even where flock is missing), +# and the EXIT trap then has nothing left to do. +_FINALIZED=0 +finalize() { + [ "$_FINALIZED" = "1" ] && return 0 + _FINALIZED=1 + strip_colab_intros + build_categorized_view + return 0 +} +trap 'finalize; lock_release' EXIT + +# Everything past this point mutates $DEST / $STATE, so hold the lock for the +# whole run. A detached refresh child blocks here until its parent has exited. +lock_acquire # Record " " for every file currently under DEST (skip metadata). record_state() { @@ -117,6 +167,7 @@ record_state() { rel="${rel#./}" case "$rel" in .unsloth_sync_state|.unsloth_sync_state.tmp|.unsloth_sync_commit) continue ;; + .unsloth_sync.lock) continue ;; esac printf '%s %s\n' "$(hash_of "$DEST/$rel")" "$rel" >> "$STATE.tmp" done @@ -180,10 +231,23 @@ fi command -v git >/dev/null 2>&1 || exit 0 command -v sha256sum >/dev/null 2>&1 || exit 0 if [ "${UNSLOTH_NB_REFRESH_CHILD:-0}" != "1" ]; then + # Finalize BEFORE the fork, not from the EXIT trap after it: the trap used to + # fire while the child was already copying refreshed notebooks in, which is + # what made "cleaned N" differ on every boot. Doing it here also keeps the + # ordering deterministic on hosts without flock. Container startup is not + # delayed any further -- the trap ran exactly this work in the parent before. + finalize + lock_release UNSLOTH_NB_REFRESH_CHILD=1 "$0" >/dev/null 2>&1 & exit 0 fi +# --- refresh child ----------------------------------------------------------- +# The parent has already stripped + built the view for the tree as it stands, so +# suppress the EXIT-trap finalize; it is re-armed below only if this refresh +# actually rewrites notebooks, which keeps an up-to-date boot a true no-op. +_FINALIZED=1 + last="$(cat "$SYNCED" 2>/dev/null || true)" remote="$(timeout "$TIMEOUT" git ls-remote "$REMOTE" HEAD 2>/dev/null | cut -f1)" [ -z "$remote" ] && exit 0 # offline / unreachable -> keep what we have @@ -246,4 +310,11 @@ mv "$TMPSTATE" "$STATE" 2>/dev/null || rm -f "$TMPSTATE" echo "$remote" > "$SYNCED" 2>/dev/null || true rm -rf "$TMP" echo "[unsloth-nb] notebooks refreshed from GitHub: $updated updated, $kept kept (your edits), $unchanged kept (only header/footer changed upstream)" +# Freshly copied notebooks arrive with the upstream Colab intro, and new files +# have to enter the view, so re-arm the finalize -- but only when something was +# actually copied. Still under the lock, so nothing else is touching the tree. +if [ "$updated" -gt 0 ]; then + _FINALIZED=0 + finalize +fi exit 0 diff --git a/tests/python/test_docker_nb_strip_colab_scope.py b/tests/python/test_docker_nb_strip_colab_scope.py new file mode 100644 index 0000000000..a536a04bb5 --- /dev/null +++ b/tests/python/test_docker_nb_strip_colab_scope.py @@ -0,0 +1,152 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Regression guard for the Colab-intro strip in the Unsloth Docker image. + +Every generated Unsloth notebook opens with a Colab-only instruction ("To run +this, press Runtime > Run all ...") that is wrong inside Docker, so the image +strips it at sync time. The strip only ever inspected cells[0], and that missed +23 of the 433 shipped notebooks: + + * 21 put the Colab badge `` in + cells[0] and the sentence in cells[1] -- Advanced_Llama3_2_(3B)_GRPO_LoRA, + Falcon_H1-Alpaca, FunctionGemma_(270M)-LMStudio, gpt-oss-(20B)-GRPO, ... + * 2 (NeMo-Gym-Multi-Environment, NeMo-Gym-Sudoku) wrap the sentence in a + single-line HTML comment, so a "line starts with the sentence" match never + fired even though the sentence IS in cells[0]. + +Measured against the pristine baked template: a cells[0]-only strip left 23 of +433 notebooks carrying the line, a leading-markdown-block strip leaves 0, and +neither changes unsloth_nb_content_sig's middle digest for any of the 433 (which +matters, because a changed digest makes the boot refresh re-copy and re-strip the +notebook forever). + +The widening also has to stay narrow: the scan stops at the first non-markdown +cell so it can never reach explanatory prose between code cells, and it stays +idempotent so a second boot is a no-op. + +Static: imports the helper and feeds it in-memory notebooks. No docker, no GPU, +no network. +""" + +from __future__ import annotations + +import copy +import importlib.util +import json +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +STRIP_PATH = REPO_ROOT / "docker" / "unsloth_nb_strip_colab.py" + +INTRO = 'To run this, press "*Runtime*" and press "*Run all*" on a **free** Tesla T4 Google Colab instance!\n' +BADGE = 'badge\n' + + +@pytest.fixture(scope = "module") +def strip(): + assert STRIP_PATH.is_file(), f"missing {STRIP_PATH}" + spec = importlib.util.spec_from_file_location("unsloth_nb_strip_under_test", STRIP_PATH) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def md(*lines): + return {"cell_type": "markdown", "metadata": {}, "source": list(lines)} + + +def code(src): + return {"cell_type": "code", "metadata": {}, "execution_count": None, + "outputs": [], "source": [src]} + + +def nb(*cells): + return {"cells": list(cells), "metadata": {}, "nbformat": 4, "nbformat_minor": 5} + + +def text(cell): + src = cell.get("source", "") + return "".join(src) if isinstance(src, list) else src + + +def has_intro(notebook): + return any("to run this, press" in text(c).lower() for c in notebook["cells"]) + + +def test_intro_in_cell_zero_is_still_stripped(strip): + # The 386-notebook majority case must not regress. + doc = nb(md(INTRO, "\n", BADGE), code("print(1)")) + assert strip._strip_intro(doc) is True + assert not has_intro(doc) + assert BADGE in text(doc["cells"][0]), "the badge row must survive the strip" + + +def test_intro_in_cell_one_behind_the_badge_is_stripped(strip): + # 21 shipped notebooks; a cells[0]-only scan left every one of them. + doc = nb(md(BADGE), md(INTRO, "\n", "You will learn how to do data prep.\n"), code("print(1)")) + assert strip._strip_intro(doc) is True + assert not has_intro(doc) + assert "You will learn how to do data prep.\n" in text(doc["cells"][1]) + + +def test_intro_inside_a_single_line_html_comment_is_stripped(strip): + # NeMo-Gym-Multi-Environment / NeMo-Gym-Sudoku ship exactly this shape. + commented = "\n" + doc = nb(md(commented, '
\n'), code("print(1)")) + assert strip._strip_intro(doc) is True + assert not has_intro(doc) + assert '
\n' in text(doc["cells"][0]) + + +def test_multi_line_html_comment_is_left_alone(strip): + # A comment that does NOT close on the same line must not be half-removed, + # or the surviving `\n"), code("print(1)")) + assert strip._strip_intro(doc) is False + assert has_intro(doc) + + +def test_strip_stops_at_the_first_code_cell(strip): + # A markdown cell AFTER code is prose, not the header block: never touched. + later = md("Explanation.\n", INTRO) + doc = nb(md(BADGE), code("print(1)"), later) + assert strip._strip_intro(doc) is False + assert text(doc["cells"][2]) == "Explanation.\n" + INTRO + + +def test_strip_is_idempotent(strip): + doc = nb(md(BADGE), md(INTRO, "\n", "rest\n"), code("print(1)")) + assert strip._strip_intro(doc) is True + once = copy.deepcopy(doc) + assert strip._strip_intro(doc) is False, "a second boot must be a no-op" + assert doc == once + + +def test_a_notebook_without_the_intro_is_untouched(strip): + doc = nb(md(BADGE, "# Title\n"), code("print(1)")) + before = copy.deepcopy(doc) + assert strip._strip_intro(doc) is False + assert doc == before + + +def test_source_given_as_a_string_is_handled(strip): + doc = nb( + {"cell_type": "markdown", "metadata": {}, "source": BADGE}, + {"cell_type": "markdown", "metadata": {}, "source": INTRO + "\nrest\n"}, + code("print(1)"), + ) + assert strip._strip_intro(doc) is True + assert not has_intro(doc) + assert isinstance(doc["cells"][1]["source"], str) + + +def test_end_to_end_write_back_is_valid_json(strip, tmp_path): + p = tmp_path / "N.ipynb" + p.write_text(json.dumps(nb(md(BADGE), md(INTRO, "\n", "rest\n"), code("print(1)")))) + assert strip.strip_notebook(str(p)) is True + reloaded = json.loads(p.read_text()) + assert not has_intro(reloaded) + assert strip.strip_notebook(str(p)) is False diff --git a/tests/python/test_docker_nb_sync_race.py b/tests/python/test_docker_nb_sync_race.py new file mode 100644 index 0000000000..e93011f22d --- /dev/null +++ b/tests/python/test_docker_nb_sync_race.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-Present the Unsloth team. See /studio/LICENSE.AGPL-3.0 + +"""Regression guard for the notebook-sync race in the Unsloth Docker image. + +unsloth_sync_notebooks.sh populates /workspace/unsloth-notebooks on boot and then +refreshes from GitHub in a DETACHED child, so container start is never blocked on +a network fetch. The parent forked that child and exited immediately, which fired +its `trap finalize EXIT` -- the Colab-intro strip plus the categorized-view +rebuild -- while the child was concurrently `cp -a`-ing refreshed notebooks into +the same tree and rewriting the same state file. Both processes also ran +build_categorized_view, which tears down and rebuilds the symlink farm. + +Six identical fresh-container boots reported "cleaned" 279 / 289 / 293 / 297 / +300 / 306 / 307 / 315 / 330 notebooks; two consecutive `docker run`s of the same +image printed 378 and 372. Worse than the noise, the lost writes were permanent: +a notebook the child copied while the parent was hashing it ended up with a +recorded hash that no longer matched the file, so the strip treated it as +user-edited and skipped it on every later boot. That is where 10 of the 23 +notebooks still carrying the Colab intro came from. Setting +UNSLOTH_SKIP_NOTEBOOK_REFRESH=1 -- i.e. never forking the child -- made the +result stable and correctly idempotent, which is what pinned the cause. + +The fix keeps the refresh detached and fixes the ORDERING instead: one exclusive +lock covers a whole invocation so the child cannot start work until the parent +has exited, the parent runs the finalize explicitly BEFORE it forks (so the order +holds even on a host without flock), the finalize is run-once, and the child +re-arms it only when the refresh actually copied something. + +Static: parses the shell script. No docker, no GPU, no network. +""" + +from __future__ import annotations + +import re +from pathlib import Path + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[2] +SYNC = REPO_ROOT / "docker" / "unsloth_sync_notebooks.sh" + + +@pytest.fixture(scope = "module") +def sync() -> str: + assert SYNC.is_file(), f"missing {SYNC}" + return SYNC.read_text() + + +def test_the_refresh_is_still_detached(sync: str): + # The whole point of the child is that a 60s ls-remote + clone must not delay + # container startup. A fix that simply made the refresh synchronous would + # pass every other test here and regress boot time. + assert re.search(r'UNSLOTH_NB_REFRESH_CHILD=1 "\$0" >/dev/null 2>&1 &', sync), ( + "the GitHub refresh must stay a detached child" + ) + + +def test_an_exclusive_lock_serialises_the_two_processes(sync: str): + assert "lock_acquire()" in sync and "lock_release()" in sync + assert re.search(r"flock -w \"\$LOCK_WAIT\" 9", sync), ( + "the lock must be a real exclusive flock, and must not block forever" + ) + + +def test_the_lock_is_taken_before_anything_mutates_the_tree(sync: str): + lock = sync.index("\nlock_acquire\n") + populate = sync.index("# 1) First-boot populate") + assert lock < populate, ( + "populate / restore / refresh all rewrite the state file; the lock has to " + "cover them, not just the strip" + ) + + +def test_a_missing_flock_degrades_instead_of_hanging(sync: str): + block = sync[sync.index("lock_acquire()") : sync.index("lock_release()")] + assert "command -v flock" in block and "return 0" in block, ( + "a host without flock, or a $DEST that cannot hold the lock file, must " + "fall back to running unlocked rather than failing the boot" + ) + + +def test_the_parent_finalizes_before_it_forks(sync: str): + fork = sync.index('UNSLOTH_NB_REFRESH_CHILD=1 "$0"') + block = sync[sync.index('if [ "${UNSLOTH_NB_REFRESH_CHILD:-0}" != "1" ]; then') : fork] + assert re.search(r"^\s*finalize\s*$", block, re.M), ( + "the strip and view rebuild must be done BEFORE the child exists; running " + "them from the EXIT trap after the fork is the race itself" + ) + + +def test_finalize_runs_at_most_once(sync: str): + block = sync[sync.index("finalize() {") : sync.index("trap 'finalize; lock_release' EXIT")] + assert '[ "$_FINALIZED" = "1" ] && return 0' in block, ( + "the explicit pre-fork call and the EXIT trap must not strip twice" + ) + assert "_FINALIZED=1" in block + + +def test_the_exit_trap_still_covers_the_early_exits(sync: str): + # Offline / no-git / UNSLOTH_SKIP_NOTEBOOK_REFRESH all exit before the fork + # site, and still need the view built. + assert "trap 'finalize; lock_release' EXIT" in sync + + +def test_the_child_does_not_repeat_the_parents_finalize(sync: str): + tail = sync[sync.index("# --- refresh child ---") :] + assert re.search(r"^_FINALIZED=1\s*$", tail, re.M), ( + "the parent already stripped and built the view for the tree as it " + "stands; an unconditional second pass makes an up-to-date boot noisy" + ) + + +def test_the_child_re_arms_the_finalize_only_after_it_copies(sync: str): + tail = sync[sync.index("refreshed from GitHub") :] + assert re.search(r'if \[ "\$updated" -gt 0 \]; then\s*\n\s*_FINALIZED=0\s*\n\s*finalize', + tail), ( + "freshly copied notebooks arrive with the upstream Colab intro and have " + "to be stripped, but only when something was actually copied" + ) + + +def test_the_lock_file_is_not_recorded_as_a_notebook(sync: str): + block = sync[sync.index("record_state() {") :] + block = block[: block.index("\n}")] + assert ".unsloth_sync.lock) continue" in block, ( + "the lock file lives in $DEST next to the state file and must be excluded " + "from the managed-file state like the other metadata" + ) + + +def test_the_lock_lives_beside_the_state_it_protects(sync: str): + assert re.search(r'^LOCK="\$DEST/\.unsloth_sync\.lock"', sync, re.M), ( + "keeping the lock in $DEST also serialises two containers sharing the " + "notebooks volume, which /tmp would not" + )