p4_stop_probe: repair the soak branch's timeout tuple (an earlier edit dropped it, which would have been a NameError only in soak mode) and right-size the leg ceilings under the notebook timeout

kernels/p4_stop_probe.py

@@ -1,275 +1,279 @@
-# 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 = "850729d3ef34359c3b10826d898681634cf3274c"
-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",
-                           "98f1402ace0f76dfff89bc50bbdb8f49a270e165d711b8ceef71cd6030ba2665"),
-}
-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"
-# A scratch repo per attempt (dated), never the real run's repo: v1 left `final` at step 20 and v2 was
-# correctly refused by the stale-stop guard, so rather than clearing state between attempts each one gets an
-# empty repo. That also makes every attempt walk the RepoMissing resume branch (E-034) that session 1 hits.
-PROBE = os.environ.get("P4_PROBE_REPO") or (
-    "Cion-lab/ounce100m-ckpt-probe-" + time.strftime("%m%d-%H%M", time.gmtime()))
-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")
-if MODE not in ("smoke", "soak"):
-    # A typo here would otherwise run the 1.6-hour soak when a 13-minute smoke was asked for.
-    raise SystemExit("P4_PROBE_MODE must be smoke or soak, got %r" % MODE)
-TLOG = "/kaggle/working/tlogs"
-# torchrun takes its first positional as the SCRIPT, not a command: passing sys.executable made
-# it compile the Python binary (E-037). --redirects is a bitmask into the --log-dir per rank:
-# 2 = redirect only stderr, so rank 1's traceback lands in a file while stdout (the loss curve,
-# which §5 requires monitoring on every wake) still streams live to this log. 3 would file both.
-TORCHRUN = ["torchrun", "--nproc_per_node=2", "--redirects", "3", "--tee", "3", "--log-dir", TLOG]
-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 = 1200, 900, 1500
-else:
-    STEPS, PUSH_EVERY, STOP1, VAL = 180, 60, 120, 2000000
-    T_LEG1, T_LEG2, T_FRESH = 4800, 2700, 1500
-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
-    e["NCCL_DEBUG"] = "WARN"        # a rank that dies in a collective says so here and nowhere else
-    e["TORCH_CPP_LOG_LEVEL"] = "WARNING"
-    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()
-
-    def _hard():
-        killed.append(True)
-        try:
-            os.killpg(os.getpgid(p.pid), signal.SIGKILL)
-        except Exception:
-            p.kill()
-
-    timer = threading.Timer(timeout, _kill)
-    timer.daemon = True
-    timer.start()
-    hard = threading.Timer(timeout + 90, _hard)
-    hard.daemon = True
-    hard.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()
-        hard.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",
-      "| probe repo", PROBE, flush=True)
-
-common = ["train_ounce100m.py", "--root", ROOT, "--out", RUN,
-          "--hub-repo", PROBE, "--prune", "--seq-len", "1024", "--attn", "eager",
-          # The flag under test. Passing both --grad-ckpt and --no-grad-ckpt would leave it to argparse's
-          # last-wins ordering, which is not a thing to be ambiguous about in a probe of this recipe.
-          "--no-grad-ckpt",
-          "--micro-batch", "4", "--accum", "32",
-          "--tokens", str(TOKENS), "--lr", "6e-4",
-          "--push-every-steps", str(PUSH_EVERY), "--val-tokens", str(VAL), "--log-every", "5",
-          "--resume", "auto"]
-
-
-def diag(tag):
-    """torchrun's ChildFailedError said `error_file: <N/A>` and printed no traceback, which is useless for
-    a rank-1-only failure. --redirects writes each rank's own stdout/stderr into the log dir, so print
-    those on a failed leg: the exception is in there, and guessing at it costs a session."""
-    # v4's diag found no *.log at all, so list the tree as well as tailing it: if torchrun named the
-    # files something else, this says so instead of printing nothing and leaving me guessing.
-    rc, out = run(["bash", "-c",
-                   'ls -R "%s" 2>&1 | head -40; '
-                   'for f in $(find "%s" -type f 2>/dev/null | head -8); do echo "==== $f"; '
-                   'tail -70 "$f"; done' % (TLOG, TLOG)], "DIAG_" + tag, 180)
-    return out
-
-
-print("PROBE_REPO", PROBE, flush=True)
-shutil.rmtree(RUN, ignore_errors=True)
-rc1, o1 = run(TORCHRUN + common + ["--stop-after-steps", str(STOP1)], "LEG1_STOP_EARLY", T_LEG1)
-if rc1 != 0:
-    diag("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 + common + ["--stop-after-steps", "0"], "LEG2_TO_HORIZON", T_LEG2)
-if rc2 != 0:
-    diag("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"),
-             "peak_alloc_gb": rj1.get("peak_gpu_gb"), "peak_alloc_gb_max_rank":
-             rj1.get("peak_gpu_gb_max_rank"), "peak_reserved_gb_max_rank":
-             rj1.get("peak_reserved_gb_max_rank"), "val_error": rj1.get("val_error"),
-             "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"), "peak_gpu_gb_max_rank":
-             rj2.get("peak_gpu_gb_max_rank"), "peak_reserved_gb_max_rank":
-             rj2.get("peak_reserved_gb_max_rank"), "val_error": rj2.get("val_error"),
-             "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.
-    # rank 0's allocated peak is not the constraint: reserved bytes, the CUDA contexts and the *other*
-    # rank's peak are what an OOM is actually about, so the max across ranks is what gets gated.
-    "peak_memory_below_13_6_gb": (not GATE_PEAK) or (
-        (rj1.get("peak_reserved_gb_max_rank") or rj1.get("peak_gpu_gb_max_rank") or 99) <= 13.6
-        and (rj2.get("peak_reserved_gb_max_rank") or rj2.get("peak_gpu_gb_max_rank") or 99) <= 13.6),
-    "validation_actually_ran": (rj1.get("val_ppl") is not None and rj1.get("val_error") is None
-                                and rj2.get("val_ppl") is not None and rj2.get("val_error") is None),
-    "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 = "850729d3ef34359c3b10826d898681634cf3274c"
+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",
+                           "98f1402ace0f76dfff89bc50bbdb8f49a270e165d711b8ceef71cd6030ba2665"),
+}
+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"
+# A scratch repo per attempt (dated), never the real run's repo: v1 left `final` at step 20 and v2 was
+# correctly refused by the stale-stop guard, so rather than clearing state between attempts each one gets an
+# empty repo. That also makes every attempt walk the RepoMissing resume branch (E-034) that session 1 hits.
+PROBE = os.environ.get("P4_PROBE_REPO") or (
+    "Cion-lab/ounce100m-ckpt-probe-" + time.strftime("%m%d-%H%M", time.gmtime()))
+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")
+if MODE not in ("smoke", "soak"):
+    # A typo here would otherwise run the 1.6-hour soak when a 13-minute smoke was asked for.
+    raise SystemExit("P4_PROBE_MODE must be smoke or soak, got %r" % MODE)
+TLOG = "/kaggle/working/tlogs"
+# torchrun takes its first positional as the SCRIPT, not a command: passing sys.executable made
+# it compile the Python binary (E-037). --redirects is a per-rank bitmask into --log-dir (1=stderr,
+# 2=stdout, 3=both) and --tee repeats the same streams to this process, so the loss curve §5 requires
+# watching stays live *and* each rank's traceback is on disk for diag() below.
+TORCHRUN = ["torchrun", "--nproc_per_node=2", "--redirects", "3", "--tee", "3", "--log-dir", TLOG]
+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 = 1200, 900, 1500
+else:
+    STEPS, PUSH_EVERY, STOP1, VAL = 180, 60, 120, 2000000
+    # Sized from the model rather than guessed: 120 steps at ~20.5 s + build + three pushes + eval is
+    # ~3,000 s, and 60 steps + cold pull + push + eval is ~1,700 s. The three ceilings must also fit under
+    # the notebook's own session timeout (10,800 s for the soak) with room for the model build, because a
+    # stage that outlives the container reports nothing at all (review point B5).
+    T_LEG1, T_LEG2, T_FRESH = 3900, 2400, 1200
+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
+    e["NCCL_DEBUG"] = "WARN"        # a rank that dies in a collective says so here and nowhere else
+    e["TORCH_CPP_LOG_LEVEL"] = "WARNING"
+    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()
+
+    def _hard():
+        killed.append(True)
+        try:
+            os.killpg(os.getpgid(p.pid), signal.SIGKILL)
+        except Exception:
+            p.kill()
+
+    timer = threading.Timer(timeout, _kill)
+    timer.daemon = True
+    timer.start()
+    hard = threading.Timer(timeout + 90, _hard)
+    hard.daemon = True
+    hard.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()
+        hard.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",
+      "| probe repo", PROBE, flush=True)
+
+common = ["train_ounce100m.py", "--root", ROOT, "--out", RUN,
+          "--hub-repo", PROBE, "--prune", "--seq-len", "1024", "--attn", "eager",
+          # The flag under test. Passing both --grad-ckpt and --no-grad-ckpt would leave it to argparse's
+          # last-wins ordering, which is not a thing to be ambiguous about in a probe of this recipe.
+          "--no-grad-ckpt",
+          "--micro-batch", "4", "--accum", "32",
+          "--tokens", str(TOKENS), "--lr", "6e-4",
+          "--push-every-steps", str(PUSH_EVERY), "--val-tokens", str(VAL), "--log-every", "5",
+          "--resume", "auto"]
+
+
+def diag(tag):
+    """torchrun's ChildFailedError said `error_file: <N/A>` and printed no traceback, which is useless for
+    a rank-1-only failure. --redirects writes each rank's own stdout/stderr into the log dir, so print
+    those on a failed leg: the exception is in there, and guessing at it costs a session."""
+    # v4's diag found no *.log at all, so list the tree as well as tailing it: if torchrun named the
+    # files something else, this says so instead of printing nothing and leaving me guessing.
+    rc, out = run(["bash", "-c",
+                   'ls -R "%s" 2>&1 | head -40; '
+                   'for f in $(find "%s" -type f 2>/dev/null | head -8); do echo "==== $f"; '
+                   'tail -70 "$f"; done' % (TLOG, TLOG)], "DIAG_" + tag, 180)
+    return out
+
+
+print("PROBE_REPO", PROBE, flush=True)
+shutil.rmtree(RUN, ignore_errors=True)
+rc1, o1 = run(TORCHRUN + common + ["--stop-after-steps", str(STOP1)], "LEG1_STOP_EARLY", T_LEG1)
+if rc1 != 0:
+    diag("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 + common + ["--stop-after-steps", "0"], "LEG2_TO_HORIZON", T_LEG2)
+if rc2 != 0:
+    diag("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"),
+             "peak_alloc_gb": rj1.get("peak_gpu_gb"), "peak_alloc_gb_max_rank":
+             rj1.get("peak_gpu_gb_max_rank"), "peak_reserved_gb_max_rank":
+             rj1.get("peak_reserved_gb_max_rank"), "val_error": rj1.get("val_error"),
+             "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"), "peak_gpu_gb_max_rank":
+             rj2.get("peak_gpu_gb_max_rank"), "peak_reserved_gb_max_rank":
+             rj2.get("peak_reserved_gb_max_rank"), "val_error": rj2.get("val_error"),
+             "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.
+    # rank 0's allocated peak is not the constraint: reserved bytes, the CUDA contexts and the *other*
+    # rank's peak are what an OOM is actually about, so the max across ranks is what gets gated.
+    "peak_memory_below_13_6_gb": (not GATE_PEAK) or (
+        (rj1.get("peak_reserved_gb_max_rank") or rj1.get("peak_gpu_gb_max_rank") or 99) <= 13.6
+        and (rj2.get("peak_reserved_gb_max_rank") or rj2.get("peak_gpu_gb_max_rank") or 99) <= 13.6),
+    "validation_actually_ran": (rj1.get("val_ppl") is not None and rj1.get("val_error") is None
+                                and rj2.get("val_ppl") is not None and rj2.get("val_error") is None),
+    "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)