Upload train/preflight.py with huggingface_hub

train/preflight.py

@@ -311,6 +311,21 @@ def torchrun(args, extra, timeout=7200):
               + extra, timeout=timeout, label="torchrun " + " ".join(extra[:6]))
 
 
+# The checkpoint rig is a timing fixture, not an artifact, so it is deleted on both ends of the run.
+_DEL = """
+import os, sys
+sys.path.insert(0, "/kaggle/working")
+import ounce100m_credentials
+ounce100m_credentials.install()
+from huggingface_hub import HfApi
+try:
+    HfApi().delete_repo(repo_id="%s", repo_type="dataset", token=os.environ["HF_TOKEN"])
+    print("hub repo deleted", flush=True)
+except Exception as e:
+    print("hub repo:", type(e).__name__, str(e)[:120], flush=True)
+"""
+
+
 def p2_throughput(args, R):
     """T1: 20L vs 22L at two sequence lengths, on the real data, for enough steps that the number is
     steady-state. The main run's whole schedule is division by this number. 22L/576 is the frozen shape;
@@ -358,41 +373,76 @@ def p2_throughput(args, R):
 
 
 def p3_cold_resume(args, R):
-    """Short run -> wipe every local trace -> resume with --resume auto, which must recover the exact
-    sample position from the Hub. Loss continuity is the pass condition: a restart that silently
-    re-initialised would jump the loss."""
+    """Three legs, each ending in one checkpoint, each started from the Hub with nothing local but the mix.
+
+    This is the test that matters most and the one that cannot be faked. Leg A trains and pushes. Legs B
+    and C each `--resume auto` after their local run directory has been deleted, so recovery is forced
+    through the Hub -- which is what every real interruption looks like. Two sequential resumes because §5
+    asks for more than one, and the cursor must advance monotonically and never re-read.
+    """
+    # The Hub side has to be wiped too: a latest.json left by an earlier attempt would make leg A a
+    # mid-run resume, and then the monotonic-cursor check below would pass for the wrong reason.
+    R["P3_hub_preclean"] = (sh(["python", "-c", _DEL % args.ckpt_repo],
+                                label="P3 hub preclean")["out"] or "").strip()[-200:]
+    # Same geometry as the main run (D-011): 4 x 1024 x 32 accum x 2 cards = 262,144 tokens/step, so P3
+    # exercises the real step, the real checkpoint size and the real cursor, not a cheaper stand-in.
     common = ["--seq-len", str(args.seq_len), "--tokens", str(args.tokens),
-              "--accum", str(args.accum), "--micro-batch", "2",
+              "--accum", str(args.accum), "--micro-batch", "4",
               "--hub-repo", args.ckpt_repo, "--prune", "--log-every", "5"]
-    a = torchrun(args, common + ["--max-steps", str(args.steps_a), "--out", WORK + "/p3"],
-                 timeout=9000)
-    # erase the instance's memory of the run: local checkpoints, the run dir, and the downloaded mix?
-    # No -- the mix is the dataset and a real interruption keeps it. Only the run state goes.
-    wipe = sh(["bash", "-c", f"rm -rf {WORK}/p3 {WORK}/run; df -h {WORK} | tail -1"],
-              label="P3 wipe local run state")
-    b = torchrun(args, common + ["--max-steps", str(args.steps_a + args.steps_b),
-                                 "--out", WORK + "/p3b", "--resume", "auto"], timeout=9000)
-    R["P3"] = {"first": {"rc": a["rc"], "tail": a["out"][-2500:]},
-               "wipe": wipe["out"][-400:],
-               "resumed": {"rc": b["rc"], "tail": b["out"][-2500:]}}
-    la = [l for l in a["out"].splitlines() if "'loss'" in l or "loss=" in l]
-    lb = [l for l in b["out"].splitlines() if "'loss'" in l or "loss=" in l]
-    R["P3_loss_last_before"] = la[-1][:200] if la else None
-    R["P3_loss_first_after"] = lb[0][:200] if lb else None
-    R["P3_pass"] = bool(a["rc"] == 0 and b["rc"] == 0 and "auto-resume: hub says step" in b["out"]
-                        and lb)
-    print("VERDICT P3_pass=", R["P3_pass"], flush=True)
+    legs, cursors, losses = [], [], []
+    for i in range(3):
+        if i:
+            # the wipe IS the test: forget everything the last leg left on this instance except the dataset
+            w = sh(["bash", "-c", "rm -rf " + WORK + "/p3* " + WORK + "/run && df -h " + WORK
+                     + " | tail -1"], label=f"P3 leg {chr(65 + i)} wipe")
+            R[f"P3_wipe_{chr(65 + i)}"] = (w["out"] or "").strip()[-200:]
+        steps = args.steps_a * (i + 1)
+        # One checkpoint per leg: pushing exactly at the leg's last step is what proves the cursor was
+        # written, verified on the Hub, and then read back cold by the next leg.
+        r = torchrun(args, common + ["--max-steps", str(steps), "--push-every-steps", str(steps),
+                                     "--resume", "auto",
+                                     "--out", WORK + f"/p3_{'abc'[i]}"], timeout=9000)
+        out = r["out"]
+        cur = None
+        for line in out.splitlines():
+            if line.startswith("resume from") and "cursor=" in line:
+                cur = line.split("cursor=", 1)[1][:220]
+            if line.startswith("CKPT "):
+                cursors.append(line[:260])
+            if "prior loss at resume:" in line:
+                losses.append(line[:120])
+        leg = {"rc": r["rc"], "seconds": r["seconds"], "max_steps": steps,
+               "auto_resume": ("auto-resume: hub says step" in out) or (i == 0),
+               "cursor_line": cur,
+               "tail": out[-1200:] if r["rc"] else None,
+               "err": (r["err"] or "")[-900:] if r["rc"] else None}
+        legs.append(leg)
+        print(f"VERDICT P3 leg {'abc'[i]} rc={r['rc']} auto_resume={leg['auto_resume']}", flush=True)
+        if r["rc"] != 0:
+            break
+    R["P3"] = {"legs": legs, "ckpt_lines": cursors, "prior_loss_lines": losses}
+    seen = [int(s.split("samples=")[1].split()[0].replace(",", "")) for s in cursors
+            if "samples=" in s]
+    R["P3_cursor_sequence"] = seen
+    ok = (len(legs) == 3 and all(l["rc"] == 0 for l in legs)
+          and all(l["auto_resume"] for l in legs[1:])
+          and len(cursors) >= 3 and len(seen) >= 3
+          and seen == sorted(seen) and len(set(seen)) == len(seen))
+    R["P3_pass"] = bool(ok)
+    print("VERDICT P3_pass=", R["P3_pass"], "cursors=", seen, flush=True)
+    R["P3_hub_postclean"] = (sh(["python", "-c", _DEL % args.ckpt_repo],
+                                label="P3 hub postclean")["out"] or "").strip()[-200:]
 
 
 def main():
     ap = argparse.ArgumentParser()
     ap.add_argument("--stage", required=True,
                     choices=["pargs", "p0", "p1", "p2", "p3", "summary", "all"])
-    ap.add_argument("--seq-len", type=int, default=2048)
+    ap.add_argument("--seq-len", type=int, default=1024)   # D-011: the config the main run will use
     ap.add_argument("--accum", type=int, default=32)
     ap.add_argument("--tokens", type=int, default=1_000_000_000)
-    ap.add_argument("--steps-a", type=int, default=60)
-    ap.add_argument("--steps-b", type=int, default=60)
+    ap.add_argument("--steps-a", type=int, default=20,
+                    help="P3 leg stride; legs run to 20, 40 and 60 cumulative steps")
     ap.add_argument("--ckpt-repo", default="Cion-lab/ounce100m-ckptbench-DELETEME")
     ap.add_argument("--out-json", default=WORK + "/preflight.json")
     args = ap.parse_args()