docker: fix the notebook sync race and widen the Colab intro strip
Two bugs that compound. The sync backgrounds a GitHub refresh child and the parent exits immediately, firing its `trap finalize EXIT` (Colab intro strip plus categorized view rebuild) while the child is concurrently cp -a'ing refreshed notebooks into the same tree and rewriting the same state file. Six identical fresh-container boots reported cleaned 337/311/316/277/297/360 notebooks, and one of them published a categorized view holding 176 of 359 notebooks because both processes tore down and rebuilt the symlink farm at once. The lost writes are permanent: 222 to 309 recorded hashes no longer matched the file on disk, so those notebooks were treated as user-edited and skipped by every later strip, which is where 10 of the 23 notebooks still carrying the Colab intro came from. Keep the refresh detached, which is the whole point of it, and fix the ordering instead. One exclusive flock covers a whole invocation so the child cannot start until the parent has exited, the parent runs the finalize explicitly before it forks so the order holds even where flock is missing, the finalize is run-once, and the child re-arms it only when the refresh actually copied something. The strip itself only inspected cells[0], which 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 friends), and 2 (NeMo-Gym-*) wrap the sentence in a single-line HTML comment. Scan the leading markdown block instead, stopping at the first code cell so it can never reach prose between code cells, and match the closed single-line comment form. The strip stays idempotent and leaves the content signature of all 433 notebooks unchanged, so the boot refresh does not re-copy and re-strip them forever. Measured on the rebuilt image: ten consecutive fresh-container boots all report cleaned 536 notebook(s) and view 359 notebooks in 26 folders, 0 of 433 notebooks retain the Colab intro (was 23), 0 recorded hashes mismatch (was 222 to 309), and a second boot on the same volume is a no-op.
This commit is contained in:
parent
6162d4d87d
commit
9211c30cbc
4 changed files with 407 additions and 11 deletions
152
tests/python/test_docker_nb_strip_colab_scope.py
Normal file
152
tests/python/test_docker_nb_strip_colab_scope.py
Normal file
|
|
@ -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 `<a href="https://colab.research.google.com/...">` 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 = '<a href="https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/X.ipynb">badge</a>\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 = "<!-- " + INTRO.rstrip("\n") + " -->\n"
|
||||
doc = nb(md(commented, '<div class="align-center">\n'), code("print(1)"))
|
||||
assert strip._strip_intro(doc) is True
|
||||
assert not has_intro(doc)
|
||||
assert '<div class="align-center">\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 `<!--` swallows the rest of the cell when rendered.
|
||||
doc = nb(md("<!-- " + INTRO, "still inside the comment\n", "-->\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
|
||||
136
tests/python/test_docker_nb_sync_race.py
Normal file
136
tests/python/test_docker_nb_sync_race.py
Normal file
|
|
@ -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"
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue