E-037: torchrun's first positional IS the script -- passing sys.executable made it compile the python binary and die with 'source code cannot contain null bytes' in 8.6s. Affects the Phase 4 launcher identically at session 1. Probe also split into smoke (20 steps) and soak (180) modes

kernels/p4_stop_probe.py

@@ -50,11 +50,21 @@ 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"
-# 180 steps at the run's real step shape (262,144 tokens/step) = 47,185,920 tokens of the published mix:
-# 139,842,880 at the mid-run stop, 61,118,720 after the cold resume. Pushes every 60 steps, which is also
-# what a no-checkpointing main run would need as its recovery interval.
-TOKENS, PUSH_EVERY, STOP1, STEPS = 47185920, 60, 120, 180
-TPS = 262144
+# 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()
 
 
@@ -62,6 +72,7 @@ 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 = []
@@ -105,22 +116,22 @@ rc, out = run([sys.executable, "-c",
                "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)
+               "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 = [sys.executable, "-u", "train_ounce100m.py", "--root", ROOT, "--out", RUN,
+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", "2000000", "--log-every", "10",
+          "--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", 4200)
+              "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"]))
@@ -134,7 +145,7 @@ print("LEG1 pushed:", pushed, "pointer:", {k: ptr.get(k) for k in ("step", "path
 
 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)
+              "LEG2_TO_HORIZON", T_LEG2)
 try:
     listed2 = sorted(hubckpt.hub_listing(PROBE, "dataset", token=os.environ["HF_TOKEN"]))
 except Exception as e:
@@ -188,13 +199,16 @@ res["checks"] = {
     "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".
-    "peak_memory_below_13_6_gb": (rj1.get("peak_gpu_gb") or 99) <= 13.6
-    and (rj2.get("peak_gpu_gb") or 99) <= 13.6,
+    # 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")