pin 1dd45f80 (E-037 launcher fix)

kernels/p4_stop_probe.py

@@ -1,217 +1,217 @@
-# 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.
-#
-# It carries a second question, the one the user pushed back on (D-017): gradient checkpointing costs 32 %
-# of the throughput (9,696 -> 12,792 tok/s, 29.7 h -> 22.5 h) and gives up the memory headroom, sitting at
-# 12.25 GB of ~14.56. So this cell runs the *real* geometry at --accum 32, micro 4, WITHOUT checkpointing,
-# for 180 steps, through three push/verify/prune cycles and one forced cold resume and the end-of-run
-# validation pass -- the three things a 30-step throughput cell cannot show: allocator drift across a few
-# hundred steps, the save path's host/GPU copies while the card is 84 % full, and the eval forward pass.
-# Peak memory is asserted, not eyeballed. If it holds, the run adopts it as D-018 with a finer push cadence
-# as the bounded blast radius; if it does not, D-017 stands and this is the measurement that says so.
-import hashlib, json, os, shutil, signal, subprocess, sys, threading, time
-
-os.chdir("/kaggle/working")
-sys.path.insert(0, "/kaggle/working")
-REV = "91fb319faee801a74fb3199e63111db84673715f"
-WANT = {
-    "ounce100m_credentials.py": ("ounce100m_credentials.py",
-                                 "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
-    "shard_dataset.py": ("train/shard_dataset.py",
-                         "f35653bf4c8f2cfe7bb0c2c7835e308505fb5a84b4eb002d0f82f2cada768ca6"),
-    "hubckpt.py": ("train/hubckpt.py",
-                   "568c31b300906cb8d78a59ef890baa9731b18e064e3d69b0eaea4b45c546f23e"),
-    "train_ounce100m.py": ("train/train_ounce100m.py",
-                           "280f7ddc4632de30565426f1417f59c528a9ac0a59cfcfe18bc209216ede7c5e"),
-}
-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"
-# Two modes, one file, so the assertions are literally the same code in both. `smoke` is the user's
-# suggestion and it is the right order: 20 steps costs ~12 minutes and answers "does the training code run
-# at all, and does one checkpoint survive the push/verify/pointer/prune cycle" -- which is exactly what the
-# first run of this probe failed at, in 8.6 seconds, on a malformed torchrun command line (E-037). `soak`
-# is the 180-step memory question, and it is only worth 1.6 GPU-hours once the mechanics are known to work.
-MODE = os.environ.get("P4_PROBE_MODE", "soak")
-TPS = 262144                                        # the run's real step shape, in tokens
-if MODE == "smoke":
-    STEPS, PUSH_EVERY, STOP1, VAL = 20, 10, 10, 200000
-    T_LEG1, T_LEG2, T_FRESH = 1500, 1200, 2400
-else:
-    STEPS, PUSH_EVERY, STOP1, VAL = 180, 60, 120, 2000000
-    T_LEG1, T_LEG2, T_FRESH = 4200, 3000, 2400
-TOKENS = STEPS * TPS
-GATE_PEAK = (MODE != "smoke")                       # 20 steps says nothing about allocator drift
-T0 = time.time()
-
-
-def run(argv, label, timeout):
-    print("=== " + label, flush=True)
-    t0 = time.time()
-    e = dict(os.environ); e["PYTHONPATH"] = "/kaggle/working"
-    e["PYTHONUNBUFFERED"] = "1"   # the -u that torchrun cannot carry
-    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", T_FRESH)
-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 = ["train_ounce100m.py", "--root", ROOT, "--out", RUN,
-          "--hub-repo", PROBE, "--prune", "--seq-len", "1024", "--attn", "eager", "--grad-ckpt",
-          "--micro-batch", "4", "--accum", "32", "--no-grad-ckpt",
-          "--tokens", str(TOKENS), "--lr", "6e-4",
-          "--push-every-steps", str(PUSH_EVERY), "--val-tokens", str(VAL), "--log-every", "5",
-          "--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", T_LEG1)
-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", T_LEG2)
-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"),
-             "tok_per_s": rj1.get("tok_per_s"), "final_loss": rj1.get("final_loss"),
-             "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"),
-             "peak_gpu_gb": rj2.get("peak_gpu_gb"), "tok_per_s": rj2.get("tok_per_s"),
-             "final_loss": rj2.get("final_loss"),
-             "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": STEPS, "push_every": PUSH_EVERY, "stop1": STOP1,
-               "tokens_at_stop1": STOP1 * TPS, "tokens_at_horizon": STEPS * TPS, "grad_ckpt": False},
-}
-res["checks"] = {
-    "leg1_stopped_at_the_stop_step": rj1.get("final_step") == STOP1,
-    "leg1_reported_a_segment": rj1.get("segment_stop") is True,
-    # Integer compare: sorted() over names would order "checkpoint-120" before "checkpoint-60".
-    "leg1_pushed_every_interval": sorted(int(x.split("-")[1]) for x in pushed) ==
-                                 list(range(PUSH_EVERY, STOP1 + 1, PUSH_EVERY)),
-    "leg1_pointer_at_stop": ptr.get("step") == STOP1,
-    "leg2_resumed_and_finished": rj2.get("final_step") == STEPS,
-    "leg2_was_not_a_segment": rj2.get("segment_stop") is False,
-    "leg2_pushed_final": any(k.startswith("final/") for k in listed2),
-    "all_steps_present": sorted(int(x.split("-")[1]) for x in
-                                {k.split("/")[1] for k in listed2 if k.startswith("ckpt/checkpoint-")})
-    == list(range(PUSH_EVERY, STEPS + 1, PUSH_EVERY)),
-    "leg2_pointer_terminal": ptr2.get("step") == STEPS and ptr2.get("path_in_repo") == "final",
-    "tokens_match_the_arithmetic": rj1.get("tokens_consumed") == STOP1 * TPS
-    and rj2.get("tokens_consumed") == STEPS * TPS,
-    "params_are_the_frozen_model": rj1.get("params") == 106194240,
-    "checkpointing_really_off": rj1.get("grad_ckpt") is False,
-    # 13.6 GB of ~14.56 usable leaves the run room for a fragmentation spike and for the two CUDA contexts;
-    # above that the answer to the user's question is "not on this hardware", not "keep going". Reported in
-    # both modes, gated only in the soak: a 20-step peak is not evidence about 3,814 steps.
-    "peak_memory_below_13_6_gb": (not GATE_PEAK) or ((rj1.get("peak_gpu_gb") or 99) <= 13.6
-                                                     and (rj2.get("peak_gpu_gb") or 99) <= 13.6),
-    "no_nan_and_loss_moved": (rj1.get("final_loss") or 1e9) < 11.0
-    and (rj2.get("final_loss") or 1e9) < 11.0,
-}
-res["PROBE_PASSED"] = all(res["checks"].values()) and rc1 == 0 and rc2 == 0
-print("PROBE_MODE", MODE, "steps", STEPS, "push_every", PUSH_EVERY,
-      "stop1", STOP1, "peak_gate", GATE_PEAK, flush=True)
-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)
+# 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.
+#
+# It carries a second question, the one the user pushed back on (D-017): gradient checkpointing costs 32 %
+# of the throughput (9,696 -> 12,792 tok/s, 29.7 h -> 22.5 h) and gives up the memory headroom, sitting at
+# 12.25 GB of ~14.56. So this cell runs the *real* geometry at --accum 32, micro 4, WITHOUT checkpointing,
+# for 180 steps, through three push/verify/prune cycles and one forced cold resume and the end-of-run
+# validation pass -- the three things a 30-step throughput cell cannot show: allocator drift across a few
+# hundred steps, the save path's host/GPU copies while the card is 84 % full, and the eval forward pass.
+# Peak memory is asserted, not eyeballed. If it holds, the run adopts it as D-018 with a finer push cadence
+# as the bounded blast radius; if it does not, D-017 stands and this is the measurement that says so.
+import hashlib, json, os, shutil, signal, subprocess, sys, threading, time
+
+os.chdir("/kaggle/working")
+sys.path.insert(0, "/kaggle/working")
+REV = "1dd45f80384750d5e1509fdcb4ca8eea6145dcf1"
+WANT = {
+    "ounce100m_credentials.py": ("ounce100m_credentials.py",
+                                 "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
+    "shard_dataset.py": ("train/shard_dataset.py",
+                         "f35653bf4c8f2cfe7bb0c2c7835e308505fb5a84b4eb002d0f82f2cada768ca6"),
+    "hubckpt.py": ("train/hubckpt.py",
+                   "568c31b300906cb8d78a59ef890baa9731b18e064e3d69b0eaea4b45c546f23e"),
+    "train_ounce100m.py": ("train/train_ounce100m.py",
+                           "280f7ddc4632de30565426f1417f59c528a9ac0a59cfcfe18bc209216ede7c5e"),
+}
+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"
+# Two modes, one file, so the assertions are literally the same code in both. `smoke` is the user's
+# suggestion and it is the right order: 20 steps costs ~12 minutes and answers "does the training code run
+# at all, and does one checkpoint survive the push/verify/pointer/prune cycle" -- which is exactly what the
+# first run of this probe failed at, in 8.6 seconds, on a malformed torchrun command line (E-037). `soak`
+# is the 180-step memory question, and it is only worth 1.6 GPU-hours once the mechanics are known to work.
+MODE = os.environ.get("P4_PROBE_MODE", "soak")
+TPS = 262144                                        # the run's real step shape, in tokens
+if MODE == "smoke":
+    STEPS, PUSH_EVERY, STOP1, VAL = 20, 10, 10, 200000
+    T_LEG1, T_LEG2, T_FRESH = 1500, 1200, 2400
+else:
+    STEPS, PUSH_EVERY, STOP1, VAL = 180, 60, 120, 2000000
+    T_LEG1, T_LEG2, T_FRESH = 4200, 3000, 2400
+TOKENS = STEPS * TPS
+GATE_PEAK = (MODE != "smoke")                       # 20 steps says nothing about allocator drift
+T0 = time.time()
+
+
+def run(argv, label, timeout):
+    print("=== " + label, flush=True)
+    t0 = time.time()
+    e = dict(os.environ); e["PYTHONPATH"] = "/kaggle/working"
+    e["PYTHONUNBUFFERED"] = "1"   # the -u that torchrun cannot carry
+    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", T_FRESH)
+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 = ["train_ounce100m.py", "--root", ROOT, "--out", RUN,
+          "--hub-repo", PROBE, "--prune", "--seq-len", "1024", "--attn", "eager", "--grad-ckpt",
+          "--micro-batch", "4", "--accum", "32", "--no-grad-ckpt",
+          "--tokens", str(TOKENS), "--lr", "6e-4",
+          "--push-every-steps", str(PUSH_EVERY), "--val-tokens", str(VAL), "--log-every", "5",
+          "--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", T_LEG1)
+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", T_LEG2)
+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"),
+             "tok_per_s": rj1.get("tok_per_s"), "final_loss": rj1.get("final_loss"),
+             "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"),
+             "peak_gpu_gb": rj2.get("peak_gpu_gb"), "tok_per_s": rj2.get("tok_per_s"),
+             "final_loss": rj2.get("final_loss"),
+             "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": STEPS, "push_every": PUSH_EVERY, "stop1": STOP1,
+               "tokens_at_stop1": STOP1 * TPS, "tokens_at_horizon": STEPS * TPS, "grad_ckpt": False},
+}
+res["checks"] = {
+    "leg1_stopped_at_the_stop_step": rj1.get("final_step") == STOP1,
+    "leg1_reported_a_segment": rj1.get("segment_stop") is True,
+    # Integer compare: sorted() over names would order "checkpoint-120" before "checkpoint-60".
+    "leg1_pushed_every_interval": sorted(int(x.split("-")[1]) for x in pushed) ==
+                                 list(range(PUSH_EVERY, STOP1 + 1, PUSH_EVERY)),
+    "leg1_pointer_at_stop": ptr.get("step") == STOP1,
+    "leg2_resumed_and_finished": rj2.get("final_step") == STEPS,
+    "leg2_was_not_a_segment": rj2.get("segment_stop") is False,
+    "leg2_pushed_final": any(k.startswith("final/") for k in listed2),
+    "all_steps_present": sorted(int(x.split("-")[1]) for x in
+                                {k.split("/")[1] for k in listed2 if k.startswith("ckpt/checkpoint-")})
+    == list(range(PUSH_EVERY, STEPS + 1, PUSH_EVERY)),
+    "leg2_pointer_terminal": ptr2.get("step") == STEPS and ptr2.get("path_in_repo") == "final",
+    "tokens_match_the_arithmetic": rj1.get("tokens_consumed") == STOP1 * TPS
+    and rj2.get("tokens_consumed") == STEPS * TPS,
+    "params_are_the_frozen_model": rj1.get("params") == 106194240,
+    "checkpointing_really_off": rj1.get("grad_ckpt") is False,
+    # 13.6 GB of ~14.56 usable leaves the run room for a fragmentation spike and for the two CUDA contexts;
+    # above that the answer to the user's question is "not on this hardware", not "keep going". Reported in
+    # both modes, gated only in the soak: a 20-step peak is not evidence about 3,814 steps.
+    "peak_memory_below_13_6_gb": (not GATE_PEAK) or ((rj1.get("peak_gpu_gb") or 99) <= 13.6
+                                                     and (rj2.get("peak_gpu_gb") or 99) <= 13.6),
+    "no_nan_and_loss_moved": (rj1.get("final_loss") or 1e9) < 11.0
+    and (rj2.get("final_loss") or 1e9) < 11.0,
+}
+res["PROBE_PASSED"] = all(res["checks"].values()) and rc1 == 0 and rc2 == 0
+print("PROBE_MODE", MODE, "steps", STEPS, "push_every", PUSH_EVERY,
+      "stop1", STOP1, "peak_gate", GATE_PEAK, flush=True)
+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)