diff --git a/scripts/modal_retry.sh b/scripts/modal_retry.sh new file mode 100755 index 0000000..24d0679 --- /dev/null +++ b/scripts/modal_retry.sh @@ -0,0 +1,18 @@ +#!/bin/bash +# Generic Modal retry wrapper — the resilience pattern this campaign +# re-implemented in /tmp seven times. Idempotent server-side functions +# make retries the resume mechanism. +# +# scripts/modal_retry.sh [args...] [retry_sleep_secs] +# +# Example: +# scripts/modal_retry.sh src/gpu/modal_export.py::parity --model khm-latn 120 +set -u +target="$1"; shift +sleep_secs="${!#}" # last arg if numeric +if [[ "$sleep_secs" =~ ^[0-9]+$ ]]; then set -- "${@:1:$#-1}"; else sleep_secs=60; fi +until modal run --detach "$target" "$@"; do + echo "[modal_retry] $target failed — retrying in ${sleep_secs}s ($(date))" + sleep "$sleep_secs" +done +echo "[modal_retry] $target completed" diff --git a/scripts/watch_marker.sh b/scripts/watch_marker.sh new file mode 100755 index 0000000..01d21b8 --- /dev/null +++ b/scripts/watch_marker.sh @@ -0,0 +1,19 @@ +#!/bin/bash +# Watch a Modal volume marker with debounce, then act (or just report). +# The pattern behind every poller this campaign needed. +# +# scripts/watch_marker.sh [check_interval_secs] [cmd...] +# +# Prints WAITING lines to stderr, exits 0 when the marker appears, then +# runs cmd if given. No debounce needed for a marker that only appears. +set -u +volume="$1"; marker="$2"; interval="${3:-600}"; shift 3 2>/dev/null || shift $# +while true; do + rm -f /tmp/.watch_marker_$$ + if modal volume get "$volume" "$marker" "/tmp/.watch_marker_$$" >/dev/null 2>&1; then + [ -f /tmp/.watch_marker_$$ ] && break + fi + sleep "$interval" +done +echo "[watch_marker] $volume:$marker present ($(date))" +if [ "$#" -gt 0 ]; then exec "$@"; fi diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 4229d84..a9ecbf5 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -57,6 +57,37 @@ CHECKPOINTS = modal.Volume.from_name("rababa-checkpoints") + +VOLUME_MOUNTS = { + "rababa": "/checkpoints", + "secryst": "/secryst-checkpoints", + "persian": "/persian-checkpoints", +} +DATA_MOUNTS = { + "secryst": "/secryst-datasets", + "persian": "/persian-datasets", +} + + +def resolve_spec(spec: dict) -> dict: + """Volume-relative paths for a spec — the single owner of "which + volume does this teacher/student/dataset live on" (previously four + pasted vol_map blocks).""" + teacher_vol = spec.get("teacher_volume", "rababa") + data_root = spec.get("data_volume", DATA_MOUNTS.get(teacher_vol, "/datasets")) + teacher = ( + spec["teacher"] if spec.get("teacher_is_hub") + else str(Path(VOLUME_MOUNTS[teacher_vol]) / spec["teacher"]) + ) + out_root = str(Path(VOLUME_MOUNTS[spec.get("out_volume", teacher_vol)]) / spec["out"]) + return { + "teacher_vol": teacher_vol, + "data_root": data_root, + "teacher": teacher, + "out_root": out_root, + "best": str(Path(out_root) / "best"), + } + def _ensure_src_path() -> None: # Modal copies the entry file to /root/.py while the repo image # sits at /root/interscript-ml — cover both layouts before importing @@ -398,19 +429,10 @@ def evaluate_per(spec_id: str, limit: int = 0) -> dict: from transformers import AutoModelForSeq2SeqLM, AutoTokenizer spec = SPECS[spec_id] - teacher_vol = spec.get("teacher_volume", "rababa") - vol_map = { - "rababa": "/checkpoints", - "secryst": "/secryst-checkpoints", - "persian": "/persian-checkpoints", - } - data_vol = {"secryst": "/secryst-datasets", - "persian": "/persian-datasets"}.get(teacher_vol, "/datasets") - data_vol = spec.get("data_volume", data_vol) - teacher_path = (spec["teacher"] if spec.get("teacher_is_hub") - else str(Path(vol_map[teacher_vol]) / spec["teacher"])) - student_vol = vol_map[spec.get("out_volume", teacher_vol)] - student_path = Path(student_vol) / spec["out"] / "best" + paths = resolve_spec(spec) + data_vol = paths["data_root"] + teacher_path = paths["teacher"] + student_path = Path(paths["best"]) test_rel = spec.get("eval_test") or spec.get("test") if not test_rel: raise RuntimeError(f"{spec_id}: no test path") @@ -510,20 +532,9 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict: spec = SPECS[spec_id] teacher_vol = spec.get("teacher_volume", "rababa") - vol_map = { - "rababa": "/checkpoints", - "secryst": "/secryst-checkpoints", - "persian": "/persian-checkpoints", - } - teacher_root = vol_map[teacher_vol] - out_root_vol = vol_map[spec.get("out_volume", teacher_vol)] - teacher_path = (spec["teacher"] if spec.get("teacher_is_hub") - else str(Path(teacher_root) / spec["teacher"])) - - data_vol = {"secryst": "/secryst-datasets", - "persian": "/persian-datasets"}.get(teacher_vol, "/datasets") - data_vol = spec.get("data_volume", data_vol) - train_path = Path(data_vol) / spec["train"] + paths = resolve_spec(spec) + teacher_path = paths["teacher"] + train_path = Path(paths["data_root"]) / spec["train"] # Teacher: use its OWN tokenizer (sentencepiece for umt5) teacher_tok = AutoTokenizer.from_pretrained(str(teacher_path)) @@ -644,7 +655,7 @@ def collate(batch): train_files = [(train_path, unit_limits[0] if unit_limits else 0)] for i, p in enumerate(spec.get("train_extra", [])): lim = unit_limits[i + 1] if i + 1 < len(unit_limits) else 0 - train_files.append((Path(data_vol) / p, lim)) + train_files.append((Path(paths["data_root"]) / p, lim)) train_ds = Pairs(train_files) print(f"[{spec_id}] train pairs: {len(train_ds)} from {len(train_files)} files", flush=True) label_beams = int(spec.get("label_beams", 4)) @@ -652,7 +663,7 @@ def collate(batch): # Step 1: teacher generates labels (beam-4) for the full corpus. # Resumable: evictions mid-labeling are routine on long jobs — # already-labeled srcs are skipped, the rest are appended. - out_root = Path(out_root_vol) / spec["out"] + out_root = Path(paths["out_root"]) out_root.mkdir(parents=True, exist_ok=True) labels_file = spec.get("labels_file", "teacher_labels.jsonl") if labels_file.endswith(".b64"): @@ -999,17 +1010,9 @@ def evaluate_der(spec_id: str, window: int = 1400, limit: int = 0) -> dict: ) spec = SPECS[spec_id] - vol_map = { - "rababa": "/checkpoints", - "secryst": "/secryst-checkpoints", - "persian": "/persian-checkpoints", - } - teacher_path = (spec["teacher"] if spec.get("teacher_is_hub") - else str(Path(vol_map[spec.get("teacher_volume", "rababa")]) / spec["teacher"])) - student_path = ( - Path(vol_map[spec.get("out_volume", spec.get("teacher_volume", "rababa"))]) - / spec["out"] / "best" - ) + paths = resolve_spec(spec) + teacher_path = paths["teacher"] + student_path = Path(paths["best"]) tok = AutoTokenizer.from_pretrained("google/byt5-small") teacher = AutoModelForSeq2SeqLM.from_pretrained(teacher_path).to("cuda").eval() @@ -1051,9 +1054,7 @@ def der_ce(model) -> dict: # what r7-style _init_choice probes read) import json - out_root = Path( - vol_map[spec.get("out_volume", spec.get("teacher_volume", "rababa"))] - ) / spec["out"] + out_root = Path(paths["out_root"]) out_root.mkdir(parents=True, exist_ok=True) (out_root / "final_eval.json").write_text( json.dumps(result, indent=2), encoding="utf-8" @@ -1103,20 +1104,10 @@ def distill_microkimi(spec_id: str, epochs: int = 3, calib_batches: int = 64, ) spec = SPECS[spec_id] - teacher_vol = spec.get("teacher_volume", "rababa") - vol_map = { - "rababa": "/checkpoints", - "secryst": "/secryst-checkpoints", - "persian": "/persian-checkpoints", - } - data_vol = {"secryst": "/secryst-datasets", - "persian": "/persian-datasets"}.get(teacher_vol, "/datasets") - data_vol = spec.get("data_volume", data_vol) - out_root_vol = vol_map[spec.get("out_volume", teacher_vol)] - teacher_path = (spec["teacher"] if spec.get("teacher_is_hub") - else str(Path(vol_map[teacher_vol]) / spec["teacher"])) - out_root = Path(out_root_vol) / spec["out"] + paths = resolve_spec(spec) + out_root = Path(paths["out_root"]) out_root.mkdir(parents=True, exist_ok=True) + teacher_path = paths["teacher"] student_tok = AutoTokenizer.from_pretrained("google/byt5-small") teacher = AutoModelForSeq2SeqLM.from_pretrained(teacher_path).to("cuda").eval() @@ -1404,7 +1395,6 @@ def qwen_next_chain() -> dict: modal run --detach src/gpu/modal_distill.py::qwen_chain """ - import json import time from pathlib import Path @@ -1414,38 +1404,29 @@ def qwen_next_chain() -> dict: ("ara-diac-small-muon", "rababa_arabic_distill_small/run-005-muon"), ("ara-diac-small-2", "rababa_arabic_distill_small/run-006-r7-muon"), ] - ROOT = Path("/checkpoints") - - def log(run: str, event: str) -> None: - # mkdir: a fresh arm's run dir does not exist until its training - # creates it — the first watch line must not crash on that - run_dir = ROOT / run - run_dir.mkdir(parents=True, exist_ok=True) - with (run_dir / "chain_log.jsonl").open("a", encoding="utf-8") as fh: - fh.write(json.dumps({"t": round(time.time()), "event": event}) + "\n") - CHECKPOINTS.commit() - def latest_step(run: str) -> int: - steps = [int(p.name.split("-")[1]) for p in (ROOT / run).glob("step-*")] - return max(steps) if steps else -1 + _ensure_src_path() + from gpu.runstate import RunState status = {} for spec_id, run in ARMS: - while not (ROOT / run / "best" / "config.json").exists(): + state = RunState(Path("/checkpoints") / run) + while not state.training_done(): CHECKPOINTS.reload() - before = latest_step(run) - log(run, f"watch step={before}") + before = state.latest_step() + state.log(f"watch step={before}", commit=CHECKPOINTS.commit) time.sleep(1200) CHECKPOINTS.reload() - after = latest_step(run) - if after == before and not (ROOT / run / "best" / "config.json").exists(): - log(run, f"stalled at step={after}; respawning {spec_id}") + after = state.latest_step() + if after == before and not state.training_done(): + state.log(f"stalled at step={after}; respawning {spec_id}", + commit=CHECKPOINTS.commit) distill_sequence.spawn(spec_id, epochs=3) - log(run, "training complete (best present)") - if not (ROOT / run / "final_eval.json").exists(): - log(run, "evaluating") + state.log("training complete (best present)", commit=CHECKPOINTS.commit) + if not state.eval_done(): + state.log("evaluating", commit=CHECKPOINTS.commit) evaluate_der.remote(spec_id=spec_id) - log(run, "eval done") + state.log("eval done", commit=CHECKPOINTS.commit) status[run] = "complete" return status diff --git a/src/gpu/runstate.py b/src/gpu/runstate.py new file mode 100644 index 0000000..25eb189 --- /dev/null +++ b/src/gpu/runstate.py @@ -0,0 +1,55 @@ +"""Run-state markers — the durable protocol every arm, orchestrator, and +supervisor agrees on. + +The contract (learned from incidents, now an interface): +- training is done when ``best/config.json`` exists (save_pretrained) +- evaluation is done when ``final_eval.json`` exists (evaluate_der) +- progress is observable as ``step-N`` checkpoint dirs; a stall is no + new step for STALL_SECS while training is incomplete +- ``chain_log.jsonl`` is the append-only audit trail + +The mkdir bug that killed the orchestrator's arm 4 (log() opened a file +inside a run dir only training creates) is unrepresentable through this +interface: every writer mkdirs. +""" + +from __future__ import annotations + +from pathlib import Path + +TRAINING_DONE = "best/config.json" +EVAL_DONE = "final_eval.json" +CHAIN_LOG = "chain_log.jsonl" + + +class RunState: + def __init__(self, root: Path) -> None: + self.root = Path(root) + + def training_done(self) -> bool: + return (self.root / TRAINING_DONE).exists() + + def eval_done(self) -> bool: + return (self.root / EVAL_DONE).exists() + + def latest_step(self) -> int: + steps = [int(p.name.split("-")[1]) for p in self.root.glob("step-*")] + return max(steps) if steps else -1 + + def log(self, event: str, commit=None) -> None: + import json + import time + + self.root.mkdir(parents=True, exist_ok=True) + with (self.root / CHAIN_LOG).open("a", encoding="utf-8") as fh: + fh.write(json.dumps({"t": round(time.time()), "event": event}) + "\n") + if commit is not None: + commit() + + def read_eval(self) -> dict | None: + import json + + path = self.root / EVAL_DONE + if not path.exists(): + return None + return json.loads(path.read_text(encoding="utf-8")) diff --git a/tests/test_gpu_runstate.py b/tests/test_gpu_runstate.py new file mode 100644 index 0000000..70e3879 --- /dev/null +++ b/tests/test_gpu_runstate.py @@ -0,0 +1,46 @@ +"""RunState marker protocol tests — pure CPU.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +from gpu.runstate import RunState # noqa: E402 + + +def test_fresh_run_state_reports_nothing_done(tmp_path: Path) -> None: + state = RunState(tmp_path / "run-x") + assert not state.training_done() + assert not state.eval_done() + assert state.latest_step() == -1 + assert state.read_eval() is None + + +def test_log_writes_into_missing_run_dir(tmp_path: Path) -> None: + # the arm-4 bug: first log must create the run dir, not crash + state = RunState(tmp_path / "brand-new-arm") + state.log("watch step=-1") + lines = (tmp_path / "brand-new-arm" / "chain_log.jsonl").read_text().splitlines() + assert len(lines) == 1 + assert "watch step=-1" in lines[0] + + +def test_step_and_marker_semantics(tmp_path: Path) -> None: + state = RunState(tmp_path / "run-y") + for n in (500, 2000, 1000): + (tmp_path / "run-y" / f"step-{n}").mkdir(parents=True) + assert state.latest_step() == 2000 + (tmp_path / "run-y" / "best").mkdir() + (tmp_path / "run-y" / "best" / "config.json").write_text("{}") + assert state.training_done() + assert not state.eval_done() + + +def test_read_eval_roundtrip(tmp_path: Path) -> None: + state = RunState(tmp_path / "run-z") + (tmp_path / "run-z").mkdir() + (tmp_path / "run-z" / "final_eval.json").write_text('{"gate_pass": true}') + assert state.eval_done() + assert state.read_eval()["gate_pass"] is True