p4_stop_probe: exercise --stop-after-steps, the push/verify/pointer/prune order and a cold Hub resume on 2xT4 before session 1 depends on them

kernels/p4_stop_probe.py

@@ -0,0 +1,177 @@
+# P4 probe: does --stop-after-steps really produce a resumable, Hub-verified segment boundary?
+#
+# Why this exists and what it buys with ~0.4 GPU-hours: `--stop-after-steps` is the mechanism every Phase 4
+# session ends on, and Gate 3 never exercised it -- P3's legs used `--max-steps`, so each leg *was* a whole
+# run. The launcher's plan, the trainer's stop, HubPush's push-verify-pointer-prune order, the resume scan
+# that refuses a stale pointer, and the final-push-plus-terminal-pointer path are therefore untested
+# together. E-035/4 is precisely a defect in that untested seam, and the rehearsal kernel cannot reach it
+# because it needs two T4s. This does, against the published mix and the frozen 22L geometry, in one
+# session: leg 1 trains to a mid-run stop step, leg 2 wipes the disk, resumes from the Hub and runs to the
+# horizon.
+#
+# The probe writes to its own checkpoint repo, never to Cion-lab/ounce100m-ckpt: the first real session has
+# to find that repo absent, which is the RepoMissing branch E-034 was about.
+import hashlib, json, os, shutil, signal, subprocess, sys, threading, time
+
+os.chdir("/kaggle/working")
+sys.path.insert(0, "/kaggle/working")
+REV = "f6f100763eec8782851f1b418fe869df6a4dcb92"
+WANT = {
+    "ounce100m_credentials.py": ("ounce100m_credentials.py",
+                                 "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
+    "shard_dataset.py": ("train/shard_dataset.py",
+                         "adbcd96fbab9505a5e8dec4d2e83d00366a73f702786bffbde430195cf6f0471"),
+    "hubckpt.py": ("train/hubckpt.py",
+                   "568c31b300906cb8d78a59ef890baa9731b18e064e3d69b0eaea4b45c546f23e"),
+    "train_ounce100m.py": ("train/train_ounce100m.py",
+                           "c4c178c0ac17c68de1521d85612a0dc604e6d4748dfdc8bd3e374512bc90ce96"),
+}
+BASE = "https://huggingface.co/Cion-lab/ounce100m-code/resolve/" + REV
+for p, (rp, want) in sorted(WANT.items()):
+    assert subprocess.run(["curl", "-sfL", f"{BASE}/{rp}", "-o", p]).returncode == 0, ("fetch", rp)
+    got = hashlib.sha256(open(p, "rb").read()).hexdigest()
+    assert got == want, ("SHA MISMATCH", rp, got[:16], want[:16])
+    print("OK", p, got[:12], flush=True)
+
+import ounce100m_credentials as C
+print("creds:", json.dumps(C.install(verify=True)), flush=True)
+import hubckpt
+from huggingface_hub import HfApi
+
+MIX = "Cion-lab/ounce100m-mix-v1"
+PROBE = "Cion-lab/ounce100m-ckpt-probe"          # never the real run's repo, see the header
+ROOT, RUN = "/kaggle/working/mixroot", "/kaggle/working/run"
+# 60 steps at accum 2 (16 seqs/step, 16,384 tokens/step) = 983,040 tokens of the real mix. accum is the
+# only thing deviating from the frozen recipe, and the probe tests checkpoint mechanics, not throughput.
+TOKENS, PUSH_EVERY, STOP1 = 983040, 10, 30
+T0 = time.time()
+
+
+def run(argv, label, timeout):
+    print("=== " + label, flush=True)
+    t0 = time.time()
+    e = dict(os.environ); e["PYTHONPATH"] = "/kaggle/working"
+    p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
+                         env=e, bufsize=1, start_new_session=True)
+    killed = []
+
+    def _kill():
+        killed.append(True)
+        try:
+            os.killpg(os.getpgid(p.pid), signal.SIGTERM)
+        except Exception:
+            p.kill()
+
+    timer = threading.Timer(timeout, _kill)
+    timer.daemon = True
+    timer.start()
+    keep, lines = [], []
+    try:
+        for line in p.stdout:
+            line = line.rstrip("\n")
+            lines.append(line)
+            if line.startswith(("CKPT ", "RUN_JSON ", "resume from", "auto-resume", "precision:",
+                                "params:", "mix:", "checkpoint hub target", "segment boundary",
+                                "latest.json", "TRAIN DONE", "validation loss", "Traceback", "Error")):
+                print("  KEY>", line[:300], flush=True)
+            keep.append(line)
+            del keep[:-40]
+    finally:
+        timer.cancel()
+    rc = p.wait()
+    if killed:
+        print("   TIMEOUT after %d s" % timeout, flush=True)
+        rc = -9
+    if rc != 0:
+        print("   TAIL:\n" + "\n".join(keep)[-2500:], flush=True)
+    print("%s_RC %s seconds %.1f elapsed %.0f" % (label, rc, time.time() - t0, time.time() - T0),
+          flush=True)
+    return rc, "\n".join(lines)
+
+
+rc, out = run([sys.executable, "-c",
+               "import sys; sys.path.insert(0, '/kaggle/working')\n"
+               "from huggingface_hub import snapshot_download\n"
+               "p = snapshot_download(repo_id='%s', repo_type='dataset',\n"
+               "                        local_dir='/kaggle/working/mixroot', max_workers=4)\n"
+               "print('mix at', p)\n" % MIX], "FETCH_MIX", 2400)
+if rc != 0:
+    raise SystemExit("VERDICT P4PROBE_STOP could not fetch the published mix")
+man = json.load(open(os.path.join(ROOT, "manifest.json")))
+print("mix", man["n_shards"], "shards", format(int(man["total_tokens"]), ","), "tokens", flush=True)
+
+common = [sys.executable, "-u", "train_ounce100m.py", "--root", ROOT, "--out", RUN,
+          "--hub-repo", PROBE, "--prune", "--seq-len", "1024", "--attn", "eager", "--grad-ckpt",
+          "--micro-batch", "4", "--accum", "2", "--tokens", str(TOKENS), "--lr", "6e-4",
+          "--push-every-steps", str(PUSH_EVERY), "--val-tokens", "100000", "--log-every", "10",
+          "--resume", "auto"]
+
+shutil.rmtree(RUN, ignore_errors=True)
+rc1, o1 = run(["torchrun", "--nproc_per_node=2"] + common + ["--stop-after-steps", str(STOP1)],
+              "LEG1_STOP_EARLY", 3000)
+api = HfApi(token=os.environ["HF_TOKEN"])
+try:
+    listed = sorted(hubckpt.hub_listing(PROBE, "dataset", token=os.environ["HF_TOKEN"]))
+except Exception as e:
+    listed = ["<listing failed: %s>" % type(e).__name__]
+pushed = sorted({k.split("/")[1] for k in listed if k.startswith("ckpt/checkpoint-")
+                 and len(k.split("/")) > 1})
+ptr = hubckpt.latest_pointer(PROBE, token=os.environ["HF_TOKEN"])
+print("LEG1 pushed:", pushed, "pointer:", {k: ptr.get(k) for k in ("step", "path_in_repo", "error")},
+      flush=True)
+
+shutil.rmtree(RUN, ignore_errors=True)                      # force the cold-resume path (E-029's shape)
+rc2, o2 = run(["torchrun", "--nproc_per_node=2"] + common + ["--stop-after-steps", "0"],
+              "LEG2_TO_HORIZON", 3000)
+try:
+    listed2 = sorted(hubckpt.hub_listing(PROBE, "dataset", token=os.environ["HF_TOKEN"]))
+except Exception as e:
+    listed2 = ["<listing failed: %s>" % type(e).__name__]
+ptr2 = hubckpt.latest_pointer(PROBE, token=os.environ["HF_TOKEN"])
+
+
+def harvest(text, tag):
+    for line in (text or "").splitlines():
+        if line.startswith(tag):
+            try:
+                return json.loads(line[len(tag):].strip())
+            except Exception:
+                return {"unparsed": line[:200]}
+    return {}
+
+
+rj1, rj2 = harvest(o1, "RUN_JSON "), harvest(o2, "RUN_JSON ")
+res = {
+    "leg1_rc": rc1, "leg2_rc": rc2,
+    "leg1": {"final_step": rj1.get("final_step"), "segment_stop": rj1.get("segment_stop"),
+             "tokens_consumed": rj1.get("tokens_consumed"), "tok_per_s": rj1.get("tok_per_s"),
+             "params": rj1.get("params"), "peak_gpu_gb": rj1.get("peak_gpu_gb"),
+             "pushed": pushed, "pointer": {k: ptr.get(k) for k in ("step", "path_in_repo")}},
+    "leg2": {"final_step": rj2.get("final_step"), "segment_stop": rj2.get("segment_stop"),
+             "tokens_consumed": rj2.get("tokens_consumed"), "val_ppl": rj2.get("val_ppl"),
+             "pointer_after": {k: ptr2.get(k) for k in ("step", "path_in_repo")},
+             "has_final": any(k.startswith("final/") for k in listed2)},
+    "expect": {"steps_planned": 60, "push_every": PUSH_EVERY, "stop1": STOP1,
+               "tokens_at_stop1": STOP1 * 16384, "tokens_at_horizon": 60 * 16384},
+}
+res["checks"] = {
+    "leg1_stopped_at_the_stop_step": rj1.get("final_step") == STOP1,
+    "leg1_reported_a_segment": rj1.get("segment_stop") is True,
+    "leg1_pushed_every_interval": [p for p in pushed if p.startswith("checkpoint-")] ==
+                                 [f"checkpoint-{s}" for s in range(PUSH_EVERY, STOP1 + 1, PUSH_EVERY)],
+    "leg1_pointer_at_stop": ptr.get("step") == STOP1,
+    "leg2_resumed_and_finished": rj2.get("final_step") == 60,
+    "leg2_was_not_a_segment": rj2.get("segment_stop") is False,
+    "leg2_pushed_final": any(k.startswith("final/") for k in listed2),
+    "leg2_pointer_terminal": ptr2.get("step") == 60 and ptr2.get("path_in_repo") == "final",
+    "tokens_match_the_arithmetic": rj1.get("tokens_consumed") == STOP1 * 16384
+    and rj2.get("tokens_consumed") == 60 * 16384,
+    "params_are_the_frozen_model": rj1.get("params") == 106194240,
+}
+res["PROBE_PASSED"] = all(res["checks"].values()) and rc1 == 0 and rc2 == 0
+print("PROBE_JSON_BEGIN")
+print(json.dumps(res, indent=1, default=str))
+print("PROBE_JSON_END")
+print("VERDICT P4PROBE", "PASS" if res["PROBE_PASSED"] else "FAIL",
+      [k for k, v in res["checks"].items() if not v], "seconds", round(time.time() - T0, 1), flush=True)
+raise SystemExit(0 if res["PROBE_PASSED"] else 5)