train/preflight.py

# Gate 3 preflight driver: proves the pipeline, not the model, and writes the evidence itself.

# Every test here answers a question §5 Phase 3 lists, and each one emits a machine-readable verdict
# line so a future session can read the outcome without re-reading a log. GPU stages are few and short:
# the whole file is designed to fit inside the 6 GPU-hour lifetime test cap (memory/QUOTA.md), and it
# prints what it spent so the ledger can be updated from the log.

# Stages, in the cheap-first order they should run:
#   P0 (CPU)  reader: cursor math, resume slicing, determinism, val split really held out
#   P1 (CPU)  hub cycle: push_and_prune with a real-sized 100M checkpoint + optimizer state
#   P2 (GPU)  throughput: 20L vs 22L at seq 1024/2048 -- test T1, the number the main-run ETA is built on
#   P3 (GPU)  short train -> kill -> cold resume from the Hub on an empty disk -> loss continuity
#   P4 (GPU)  resume twice in sequence; the cursor must advance monotonically, never re-read
#
# P3/P4 are the tests that matter most and the ones that cannot be faked: a run that resumes from what it
# left on disk proves nothing, because every real interruption takes the disk with it (§3.13).

import argparse
import json
import os
import subprocess
import sys
import time

WORK = "/kaggle/working"
REV_DEFAULT = ""  # filled from the launcher; only used for reporting


def sh(argv, timeout=None, env=None, label=""):
    print(f"=== {label or ' '.join(argv[:3])}", flush=True)
    t0 = time.time()
    try:
        p = subprocess.run(argv, cwd=WORK, capture_output=True, text=True, timeout=timeout,
                           env=dict(os.environ, **(env or {})))
    except subprocess.TimeoutExpired as ex:
        # A timeout is a result, not an exception to propagate: letting it raise discarded every cell
        # that had already completed and orphaned both T4s (review finding).
        print(f"   TIMEOUT after {time.time() - t0:.0f}s", flush=True)
        return {"rc": -9, "out": (ex.stdout or b"").decode()[-200000:],
                "err": (ex.stderr or b"").decode()[-3000:],
                "seconds": round(time.time() - t0, 1), "timeout": True}
    for line in (p.stdout or "").splitlines():
        print("   |", line[:240], flush=True)
    if p.returncode != 0:
        print("   STDERR:", (p.stderr or "")[-3000:], flush=True)
    return {"rc": p.returncode, "out": (p.stdout or "")[-200000:],
            "err": (p.stderr or "")[-3000:], "seconds": round(time.time() - t0, 1)}


def last_json(text, begin, end):
    if begin not in text or end not in text:
        return None
    body = text.rsplit(begin, 1)[1].split(end, 1)[0]
    try:
        return json.loads(body)
    except Exception:
        return None


# ------------------------------------------------------------------------ Pargs (CPU, free)
def p_args(args, R):
    """Every keyword in train_ounce100m.py's TrainingArguments is a promise about a library version, and
    E-008/E-010 already taught us that transformers 5 drops things without warning. This constructs the
    real objects on the Kaggle image -- TrainingArguments, the config, the model, the optimizer and the
    trapezoid schedule -- with no data and no GPU, so a kwargs skew costs a free CPU minute instead of
    killing a billed 2xT4 session at import time."""
    src = r'''
import json, os, sys, math, time
sys.path.insert(0, "/kaggle/working")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
import torch
import transformers
from torch.optim import AdamW
from transformers import TrainingArguments
import train_ounce100m as T
R = {"transformers": transformers.__version__, "torch": torch.__version__}
cfg = T.build_config(2048)
R["config"] = {k: getattr(cfg, k, "<absent>") for k in
               ("hidden_size", "num_hidden_layers", "num_attention_heads", "num_key_value_heads",
                "intermediate_size", "vocab_size", "tie_word_embeddings", "max_position_embeddings",
                # rope_theta moved into rope_parameters in newer transformers; if the flat attribute is
                # gone the value we think we froze may not be the value the model got (verify, not assume)
                "rope_theta", "rope_parameters", "rms_norm_eps", "attention_dropout", "mlp_bias",
                "hidden_act")}
t0 = time.time()
m = T.LlamaForCausalLM(cfg)
R["params"] = T.count_params(m)
R["init_seconds"] = round(time.time() - t0, 1)
R["shapes_20L"] = T.count_params(T.LlamaForCausalLM(T.build_config(1024, layers=20)))
c20 = T.build_config(1024, layers=20)
R["shape20_ffn"] = c20.intermediate_size
R["shape20_kv"] = c20.num_key_value_heads
# Regression test for the defect that would have made T1 meaningless: passing --hidden 576 explicitly is
# the frozen width, so it must NOT fall through to the variant branch and become MHA-9 with a doubled FFN.
cv = T.build_config(2048, hidden=576)
R["hidden576_still_frozen"] = (cv.num_key_value_heads == 3 and cv.intermediate_size == 1536
                               and cv.num_hidden_layers == 22)
# The exact kwargs the trainer passes. If the image rejects one, this is where we learn it.
kw = dict(output_dir="/kaggle/working/argcheck", per_device_train_batch_size=2,
          gradient_accumulation_steps=32, learning_rate=6e-4, weight_decay=0.1, adam_beta1=0.9,
          adam_beta2=0.95, adam_epsilon=1e-8, max_grad_norm=1.0, lr_scheduler_type="constant",
          warmup_ratio=0.0, num_train_epochs=1, max_steps=10, fp16=True, bf16=False,
          gradient_checkpointing=False, ddp_find_unused_parameters=False, dataloader_num_workers=2,
          # the three knobs the tokens/sec sweep turns, checked here on CPU for free before any GPU cell
          # spends billed seconds discovering transformers 5 renamed or rejected one
          gradient_checkpointing_kwargs={"use_reentrant": False}, optim="adamw_torch_fused",
          torch_compile=False, per_device_eval_batch_size=4,
          dataloader_pin_memory=False, remove_unused_columns=False, ignore_data_skip=True,
          save_strategy="steps", save_steps=5, save_total_limit=1, logging_steps=5, report_to=[],
          seed=1, data_seed=1, accelerator_config={"dispatch_batches": False},
          average_tokens_across_devices=False)
try:
    ta = TrainingArguments(**kw); R["kwargs"] = "accepted"
except TypeError as e:
    R["kwargs"] = f"REJECTED: {str(e)[:300]}"
    bad = [k for k in kw if k in str(e)]
    R["suspect_keys"] = bad
    for k in bad:
        kw.pop(k, None)
    ta = TrainingArguments(**kw)
R["fp16_bf16"] = [ta.fp16, ta.bf16]
# schedule shape, standalone: warm 2%, flat to 80%, linear to zero
steps = 3815
opt = AdamW(m.parameters(), lr=6e-4, betas=(0.9, 0.95), weight_decay=0.1)
tr = T.TrapezoidTrainer.__new__(T.TrapezoidTrainer)   # no Trainer.__init__: we want the schedule only
tr.lr_shape = {"warmup_steps": max(100, int(steps * 0.02)), "decay_start": int(steps * 0.8),
               "total_steps": steps}
tr.optimizer, tr.lr_scheduler = opt, None
sch = tr.create_scheduler(steps, opt)
mults = []
for i in range(steps):
    mults.append(sch.lr_lambdas[0](i))
    sch.optimizer.param_groups[0]["lr"] = 6e-4 * mults[-1]
R["lr_shape"] = {"peak": max(mults), "at_step76": mults[76], "at_80pct": mults[int(steps*0.8)],
                 "at_90pct": mults[int(steps*0.9)], "final": mults[-1],
                 "hits_zero_only_past_last_step": mults[-1] > 0.0,
                 "plateau_fraction_steps": round(sum(1 for x in mults if x > 0.999) / steps, 3)}
# one real forward/backward on CPU with a tiny model to prove the collator output feeds the model.
# The report is printed BEFORE this optional probe: v3 of pargs lost every finding it had made because
# the last block raised and the print at the end never ran (E-012, reproduced in code written an hour
# earlier -- the guard is not enough, the ordering is the fix).
def _emit(R, tag="ARGS_JSON"):
    print(tag + "_BEGIN")
    print(json.dumps(R, default=str))
    print(tag + "_END")


_emit(R)
try:
    cfg2 = T.build_config(64, layers=2, hidden=128)
    m2 = T.LlamaForCausalLM(cfg2)
    ids = torch.randint(1, 1000, (3, 65))
    o = m2(input_ids=ids[:, :-1], labels=ids[:, 1:], attention_mask=torch.ones_like(ids[:, :-1]))
    o.loss.backward()
    fb = {"kv_heads": cfg2.num_key_value_heads, "loss": round(float(o.loss), 4),
          "finite": math.isfinite(float(o.loss)),
          "grads": sum(1 for p in m2.parameters() if p.grad is not None)}
except Exception as e:
    fb = {"error": f"{type(e).__name__}: {str(e)[:220]}"}
R["forward_backward"] = fb
R["PARGS_DONE"] = True
_emit(R)
'''
    r = sh([sys.executable, "-c", src], timeout=3600, label="Pargs construct on the Kaggle image")
    blocks = r["out"].split("ARGS_JSON_BEGIN")
    j = last_json(r["out"], "ARGS_JSON_BEGIN", "ARGS_JSON_END")
    R["PARGS"] = j
    R["PARGS_n_blocks"] = max(0, len(blocks) - 1)
    R["PARGS_rc"] = r["rc"]
    j = j or {}
    lr = j.get("lr_shape") or {}
    # Test the SHAPE of the schedule, not one value: flat plateau, then linear to ~0 at the final
    # update. (v4 asserted final == 0.0 and failed on 0.0013 -- LambdaLR reaches 0 at step T, which is
    # one past the last update at T-1. Every cosine implementation shares that off-by-one; the honest
    # assertion is "effectively zero at the last step", not "exactly zero".)
    shape_ok = (lr.get("peak") == 1.0 and lr.get("at_80pct") == 1.0
                and 0.45 < lr.get("at_90pct", -1) < 0.55
                and 0.0 <= lr.get("final", 1.0) < 0.005
                and 0.70 < lr.get("plateau_fraction_steps", 0) < 0.85)
    R["PARGS_pass"] = bool(j.get("PARGS_DONE") and j.get("kwargs") == "accepted"
                           and j.get("params", {}).get("sum_numel") == 106194240
                           # 20L must be 99,114,048 per docs/01-plan.md §2.3's closed-form table, and its
                           # FFN must stay 1536; a doubled FFN would make test T1 meaningless.
                           and j.get("shapes_20L", {}).get("sum_numel") == 99114048
                           and j.get("shape20_ffn") == 1536 and shape_ok
                           and j.get("hidden576_still_frozen") is True
                           and (j.get("forward_backward") or {}).get("finite"))
    print("VERDICT PARGS_pass=", R["PARGS_pass"], "blocks=", R["PARGS_n_blocks"], "shape_ok=", shape_ok,
          json.dumps({k: j.get(k) for k in ("transformers", "torch", "kwargs", "params", "shapes_20L",
                                            "lr_shape", "forward_backward")}, default=str)[:700],
          flush=True)


# ------------------------------------------------------------------------ P0 reader (CPU, free)
def p0_reader(args, R):
    """The reader is where an unrecoverable main run would hide: if two runs at the same cursor read
    different tokens, every resume silently trains on a subset. Checked on bytes, not on feelings."""
    src = r'''
import json, os, sys, numpy as np
sys.path.insert(0, "/kaggle/working")
import shard_dataset as SD
man = json.load(open("/kaggle/working/mixroot/manifest.json"))
store = SD.PackedTokenStore("/kaggle/working/mixroot", man)
L = 256
full = SD.make_dataset(store, L, 1234)
n = len(full)
part = SD.make_dataset(store, L, 1234, start_sample=n - 40)
same = all(bool((full[n - 40 + j]["input_ids"] == part[j]["input_ids"]).all()) for j in range(40))
labels_ok = bool((full[0]["labels"][:-1] == full[0]["input_ids"][1:]).all()
                 and int(full[0]["labels"][-1]) == int(full[0]["input_ids"][0]))
det2 = SD.make_dataset(store, L, 1234)
deterministic = bool((full[7]["input_ids"] == det2[7]["input_ids"]).all())
diffseed = SD.make_dataset(store, L, 999)
seed_matters = not bool((full[7]["input_ids"] == diffseed[7]["input_ids"]).all())
vstore = SD.PackedTokenStore("/kaggle/working/mixroot", man,
                             files=[s["file"] for s in man["val_shards"]])
tv = set(); vv = set()
vds = SD.make_dataset(vstore, L, 0, shuffle=False, count=min(400, vstore.total_tokens // L))
for i in range(min(len(full), 4000)):
    tv.add(full[i]["input_ids"][:64].numpy().tobytes())
for i in range(len(vds)):
    vv.add(vds[i]["input_ids"][:64].numpy().tobytes())
print("READER_JSON_BEGIN")
print(json.dumps({
  "total_tokens": store.total_tokens, "samples_at_L256": n,
  "resume_matches_uninterrupted": same, "labels_are_next_token": labels_ok,
  "same_seed_same_order": deterministic, "different_seed_different_order": seed_matters,
  "val_tokens": vstore.total_tokens, "val_windows_checked": len(vds),
  "train_val_prefix_collision": len(tv & vv),
  "tokens_per_shard_min": min(s["tokens"] for s in man["shards"]),
  "n_shards": len(man["shards"]),
}))
print("READER_JSON_END")
'''
    r = sh([sys.executable, "-c", src], label="P0 reader")
    R["P0"] = last_json(r["out"], "READER_JSON_BEGIN", "READER_JSON_END")
    R["P0_rc"] = r["rc"]
    j = R["P0"] or {}
    R["P0_pass"] = bool(r["rc"] == 0 and j.get("resume_matches_uninterrupted")
                        and j.get("labels_are_next_token") and j.get("same_seed_same_order")
                        and j.get("different_seed_different_order")
                        and j.get("train_val_prefix_collision") == 0
                        and j.get("val_tokens", 0) > 0)
    print("VERDICT P0_pass=", R["P0_pass"], flush=True)


# ------------------------------------------------------------------------ P1 hub cycle (CPU, free)
def p1_hub(args, R):
    """§3.13 at real size: a 106M-parameter checkpoint with optimizer state, pushed, verified from the
    Hub by re-listing and re-hashing, then pruned. Cheap because the weights are random; the BYTES are
    what is being timed."""
    src = r'''
import json, os, sys, time, numpy as np
sys.path.insert(0, "/kaggle/working")
import ounce100m_credentials, hubckpt
ounce100m_credentials.install()
from huggingface_hub import HfApi
api = HfApi(); tok = os.environ["HF_TOKEN"]
repo = "Cion-lab/ounce100m-ckptbench-DELETEME"
# Start from an empty repo: LFS dedupes identical bytes, so pushing into a repo that already holds this
# mock would time a no-op and report it as upload speed. Same lesson as the rehearsal's preclean.
try:
    api.delete_repo(repo_id=repo, repo_type="dataset", token=tok)
    print("preclean: deleted stale ckptbench repo", flush=True)
except Exception as e:
    print("preclean:", type(e).__name__, str(e)[:100], flush=True)
print("ensure_repo:", json.dumps(hubckpt.ensure_repo(repo, api, tok)), flush=True)
d = "/kaggle/working/ckptbench/checkpoint-1"
os.makedirs(d, exist_ok=True)
# 106,194,240 params x 4 B x 4 arrays (fp16-saved weights + fp32 master + Adam m and v) is the real
# checkpoint footprint, and it is the size D-007 measured at 1.6 GB. Timing a 425 MB mock would flatter
# the main run.
n = 106194240
t0 = time.time()
for name in ("weights.fp32", "master.fp32", "adam_m.fp32", "adam_v.fp32"):
    a = np.lib.format.open_memmap(os.path.join(d, "state." + name.replace(".", "_") + ".npy"),
                                  dtype=np.float32, mode="w+", shape=(n,))
    a[:] = np.float32(0.0)
    a.flush()
    del a
json.dump({"step": 1, "samples_consumed": 381500}, open(os.path.join(d, "cursor.json"), "w"))
up_t0 = time.time()
try:
    res = hubckpt.push_and_prune(repo, d, "ckpt/checkpoint-1", api, token=tok, prune=True)
    ver = res["verify"]
    # second, independent proof: pull it back into a clean directory and compare hashes
    back = "/kaggle/working/ckptbench/restored"
    got = hubckpt.download_checkpoint(repo, "ckpt/checkpoint-1", back, api, token=tok)
    same = json.load(open(os.path.join(got["dir"], "cursor.json"))) == {
        "step": 1, "samples_consumed": 381500}
    print("download landed in", got["dir"], "files", got["files"], flush=True)
    # The repo is a timing rig, not an artifact: leaving a 1.7 GB public blob invites a future session
    # to mistake it for a checkpoint. Everything measurable is in the log by this point.
    api.delete_repo(repo_id=repo, repo_type="dataset", token=tok)
except Exception as e:
    res, ver, got, same = {"error": f"{type(e).__name__}: {str(e)[:600]}"}, {}, {}, False
    print("   P1 raised:", res["error"], flush=True)
print("HUB_JSON_BEGIN")
print(json.dumps({"error": res.get("error"),
                  "push_verify_seconds": round(time.time()-up_t0,1),
                  "cycle": {k: res.get(k) for k in ("verify","pruned","free_before_gb","free_after_gb")},
                  "downloaded": got, "readback_ok": bool(ver.get("ok")) and same,
                  "gen_seconds": round(time.time()-t0,1)}, default=str))
print("HUB_JSON_END")
'''
    r = sh([sys.executable, "-c", src], timeout=5400,
           label="P1 hub cycle (1.7 GB mock checkpoint: weights+master+Adam m,v)")
    R["P1"] = last_json(r["out"], "HUB_JSON_BEGIN", "HUB_JSON_END")
    R["P1_rc"] = r["rc"]
    j = R["P1"] or {}
    cyc = ((j.get("cycle") or {}).get("verify") or {})
    R["P1_pass"] = bool(r["rc"] == 0 and cyc.get("ok") and j.get("readback_ok")
                        and (j.get("cycle") or {}).get("pruned"))
    print("VERDICT P1_pass=", R["P1_pass"], json.dumps(cyc)[:300], flush=True)
    R["P1_cleanup"] = "delete Cion-lab/ounce100m-ckptbench-DELETEME when done reading it"


# ------------------------------------------------------------------------ P2..P4 (GPU)
def torchrun(args, extra, timeout=7200):
    # --seq-len is deliberately NOT set here: each caller passes its own, and relying on argparse's
    # last-wins to undo a duplicate is a defect waiting for a reordering.
    return sh(["torchrun", "--nproc_per_node=2", "train_ounce100m.py", "--root", WORK + "/mixroot"]
              + 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;
    20L is the H-column alternative from docs/01-plan.md §2.3."""
    out = {}
    for layers in (20, 22):
        for seq in (1024, 2048):
            micro = 4 if seq == 1024 else 2
            r = torchrun(args, ["--layers", str(layers), "--seq-len", str(seq),
                                "--micro-batch", str(micro), "--accum", "1", "--max-steps", "30",
                                "--out", WORK + f"/t1_{layers}_{seq}", "--log-every", "10"],
                         timeout=5400)
            cell = {"rc": r["rc"], "seconds": r["seconds"], "micro_batch": micro}
            for line in r["out"].splitlines():
                if line.startswith("RUN_JSON "):
                    try:
                        cell.update(json.loads(line[len("RUN_JSON "):]))
                    except Exception as e:
                        cell["run_json_unparseable"] = str(e)[:120]
                elif line.startswith("params:"):
                    cell["params_line"] = line[:200]
                elif line.startswith("precision:"):
                    cell["precision"] = line[:200]
            if "tok_per_s" not in cell:
                cell["tail"] = r["out"][-1500:]
                cell["err"] = r["err"][-800:]
            out[f"L{layers}_s{seq}"] = cell
            if r["rc"] != 0:
                # Fast-fail: a systematic defect (import, OOM, batch-shape) repeats identically in all
                # four cells, and each repetition bills a GPU session. Stop at the first one, report it,
                # fix, then re-run. Cheapest possible way to be wrong.
                R["P2_aborted_after"] = f"L{layers}_s{seq}"
                break
        if R.get("P2_aborted_after"):
            break
    R["P2"] = out
    got = [c for c in out.values() if isinstance(c.get("tok_per_s"), (int, float))]
    # Pass means all four cells ran AND the frozen cell (22L/seq2048) is reported, because that is the
    # configuration the main run uses; a run where only the cheap cells succeeded proves nothing.
    R["P2_pass"] = bool(len(out) == 4 and all(c.get("rc") == 0 for c in out.values())
                        and len(got) == 4 and out["L22_s2048"].get("params") == 106194240)
    R["P2_eta_hours_at_1B"] = ({k: round(1e9 / c["tok_per_s"] / 3600, 2)
                                for k, c in out.items() if c.get("tok_per_s")})
    print("VERDICT P2_pass=", R["P2_pass"], "eta_hours:", json.dumps(R["P2_eta_hours_at_1B"]), flush=True)


def p_opts(args, R):
    """Tokens/sec, again -- because P3 came in at 5.99 GB peak on a 15.4 GB card and every number used to
    freeze D-011 was measured at accum 1, not the run's accum 32.

    Gradient checkpointing exists to buy memory, and it costs a full extra forward pass per layer. At
    5.99 GB of ~15.4 GB it is buying headroom nobody is spending, so the honest question is whether the
    frozen config is over-paying for it. Cells below change only *how* the same maths is executed --
    micro-batch/accum multiply to the same 262,144 tokens/step in every row, and the optimiser, LR,
    schedule, seed and data are identical -- so a winner can be adopted at freeze time without touching
    what the run is a test of.

    A baseline cell runs FIRST and the driver reports gains against it, because comparing across sessions
    on shared T4s is how the 4,071-vs-11,062 confusion happened in the first place.
    """
    cells = [
        ("base_m4_gc",         20, ["--micro-batch", "4", "--accum", "32", "--grad-ckpt"],
         {"micro_batch": 4, "accum": 32, "grad_ckpt": True}),
        ("m4_nogc",            20, ["--micro-batch", "4", "--accum", "32", "--no-grad-ckpt"],
         {"micro_batch": 4, "accum": 32, "grad_ckpt": False}),
        ("m8_gc_acc16",        20, ["--micro-batch", "8", "--accum", "16", "--grad-ckpt"],
         {"micro_batch": 8, "accum": 16, "grad_ckpt": True}),
        ("m8_nogc_acc16",      20, ["--micro-batch", "8", "--accum", "16", "--no-grad-ckpt"],
         {"micro_batch": 8, "accum": 16, "grad_ckpt": False}),
        ("m16_gc_acc8",        20, ["--micro-batch", "16", "--accum", "8", "--grad-ckpt"],
         {"micro_batch": 16, "accum": 8, "grad_ckpt": True}),
        ("m4_gc_nonreentrant", 20, ["--micro-batch", "4", "--accum", "32", "--grad-ckpt",
                                    "--gc-nonreentrant"],
         {"micro_batch": 4, "accum": 32, "grad_ckpt": True, "gc_nonreentrant": True}),
        ("m4_gc_fusedoptim",   20, ["--micro-batch", "4", "--accum", "32", "--grad-ckpt",
                                    "--optim", "adamw_torch_fused"],
         {"micro_batch": 4, "accum": 32, "grad_ckpt": True, "optim": "adamw_torch_fused"}),
    ]
    common = ["--seq-len", str(args.seq_len), "--tokens", str(args.tokens), "--attn", args.attn,
              "--hub-repo", "", "--log-every", "10", "--val-tokens", "200000"]
    out = {}
    base_rate = None
    for label, steps, extra, expect in cells:
        # Every row keeps the frozen 262,144 tokens/step: micro_batch x accum x 2 cards is constant, so a
        # winner changes how the same maths is executed, never what is being trained.
        if expect["micro_batch"] * expect["accum"] * 2 * args.seq_len != 262144:
            raise SystemExit(f"cell {label} does not preserve the frozen batch size in tokens")
        # An OOM in one row is a result, not a systematic defect: unlike P2, keep going and let the table
        # show which configurations this hardware can hold. --push-every-steps = the whole cell means one
        # checkpoint write per cell instead of ten, which would otherwise dominate elapsed and fill the
        # 19.5 GB working directory (review finding).
        r = torchrun(args, common + extra + ["--max-steps", str(steps),
                                             "--push-every-steps", str(steps),
                                             "--out", WORK + "/opt_" + label], timeout=1500)
        cell = {"rc": r["rc"], "seconds": r["seconds"], "argv": " ".join(extra),
                "expect": expect, "tokens_per_step": expect["micro_batch"] * expect["accum"] * 2
                * args.seq_len}
        for line in r["out"].splitlines():
            if line.startswith("RUN_JSON "):
                try:
                    cell.update(json.loads(line[len("RUN_JSON "):]))
                except Exception as e:
                    cell["run_json_unparseable"] = str(e)[:120]
            elif line.startswith("precision:"):
                cell["precision"] = line[:240]
            elif "OutOfMemory" in line or "CUDA out of memory" in line:
                cell["oom"] = True
        # Trust but verify: the cell is only the configuration it claims to be if the trainer says so.
        cell["expect_mismatch"] = {k: [v, cell.get(k)] for k, v in expect.items()
                                   if cell.get(k) is not None and cell.get(k) != v}
        if "tok_per_s" not in cell:
            cell["tail"] = (r["out"] or "")[-1200:]
            cell["err_first_error"] = ((r["err"] or "")[-1500:])
        if label == "base_m4_gc":
            base_rate = cell.get("tok_per_s")
        if base_rate and cell.get("tok_per_s"):
            cell["gain_vs_base_pct"] = round(100.0 * (cell["tok_per_s"] / base_rate - 1.0), 1)
        out[label] = cell
        print(f"VERDICT opt {label} rc={cell['rc']} tok_per_s={cell.get('tok_per_s')} "
              f"peak_gb={cell.get('peak_gpu_gb')} gain={cell.get('gain_vs_base_pct')} "
              f"mismatch={cell['expect_mismatch'] or 'none'}", flush=True)
        # Nothing here needs to survive the cell, and 8 cells x ~3 GB of export would fill the disk and
        # take the whole session down with it (§3.13's free-space discipline).
        sh(["bash", "-c", "rm -rf " + WORK + "/opt_" + label + " && df -h " + WORK + " | tail -1"],
           label="cleanup " + label)
    R["OPT"] = out
    got = {k: v["tok_per_s"] for k, v in out.items() if v.get("tok_per_s")}
    best = max(got, key=got.get) if got else None
    R["OPT_best"] = best
    R["OPT_best_tok_per_s"] = got.get(best)
    R["OPT_best_eta_hours_at_1B"] = (round(1e9 / got[best] / 3600, 2) if best else None)
    # "Passing" here means the table is trustworthy, not that anything improved: every cell must have
    # either run or OOMed (a real answer), and none may have died in a way that hides a defect.
    R["OPT_pass"] = bool(len(out) == len(cells) and base_rate and
                         all(("tok_per_s" in v) or v.get("oom") for v in out.values())
                         and not any(v["expect_mismatch"] for v in out.values()))
    print("VERDICT OPT_pass=", R["OPT_pass"], "best=", best, got, flush=True)


def p3_cold_resume(args, R):
    """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", "4",
              "--attn", args.attn, "--grad-ckpt",
              "--hub-repo", args.ckpt_repo, "--prune", "--log-every", "5"]
    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}
        for line in out.splitlines():
            if line.startswith("RUN_JSON "):
                try:
                    leg.update(json.loads(line[len("RUN_JSON "):]))
                except Exception as e:
                    leg["run_json_unparseable"] = str(e)[:120]
        legs.append(leg)
        print(f"VERDICT P3 leg {'abc'[i]} rc={r['rc']} auto_resume={leg['auto_resume']} "
              f"tok_per_s={leg.get('tok_per_s')} peak_gb={leg.get('peak_gpu_gb')}", flush=True)
        if r["rc"] != 0:
            break
    R["P3"] = {"legs": legs, "ckpt_lines": cursors, "prior_loss_lines": losses}
    # The last leg is the one that has run 60 steps of the frozen geometry, so it is also the throughput
    # cell D-012's tripwire is read from. Neither earlier 1024 A/B cell used micro 4 *with* grad-ckpt.
    frozen = legs[-1] if legs else {}
    R["P3_frozen_rate_tok_per_s"] = frozen.get("tok_per_s")
    R["P3_frozen_peak_gpu_gb"] = frozen.get("peak_gpu_gb")
    R["P3_frozen_params"] = frozen.get("params")
    R["P3_eta_hours_at_1B"] = (round(1e9 / frozen["tok_per_s"] / 3600, 2)
                               if frozen.get("tok_per_s") else None)
    R["P3_D012_fallback_fires"] = (R["P3_frozen_rate_tok_per_s"] < 5150
                                  if R["P3_frozen_rate_tok_per_s"] else None)
    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)
          # ...and the thing that resumed was the frozen model, not a smaller stand-in.
          and R["P3_frozen_params"] == 106194240 and bool(R["P3_frozen_rate_tok_per_s"]))
    R["P3_pass"] = bool(ok)
    print("VERDICT P3_pass=", R["P3_pass"], "cursors=", seen,
          "tok_per_s=", R["P3_frozen_rate_tok_per_s"], 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", "opts", "summary", "all"],
                    help="opts is the tokens/sec re-audit at the run's real accumulation; it is not a "
                         "Gate 3 stage and does not enter GATE_3_READY")
    ap.add_argument("--seq-len", type=int, default=1024)   # D-011: the config the main run will use
    ap.add_argument("--attn", default="eager", choices=["eager", "sdpa"],
                    help="passed through to the trainer by P3; named so a default change cannot silently"
                         " alter what the frozen-geometry measurement tested")
    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=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()

    R = {}
    if os.path.exists(args.out_json):
        try:
            R = json.load(open(args.out_json))
        except Exception:
            print("existing preflight.json unreadable; starting fresh", flush=True)
    todo = ["pargs", "p0", "p1", "p2", "p3"] if args.stage == "all" else [args.stage]
    t0 = time.time()
    if "pargs" in todo:
        p_args(args, R)
    if "p0" in todo:
        p0_reader(args, R)
    if "p1" in todo:
        p1_hub(args, R)
    if "p2" in todo:
        p2_throughput(args, R)
    if "p3" in todo:
        p3_cold_resume(args, R)
    if "opts" in todo:
        p_opts(args, R)
    R["seconds_this_invocation"] = round(time.time() - t0, 1)
    R["gpu_hours_this_invocation"] = round(R["seconds_this_invocation"] / 3600.0, 3)
    R["PASSES"] = {k: R.get(k + "_pass") for k in ("PARGS", "P0", "P1", "P2", "P3")}
    R["GATE_3_READY"] = all(v is True for v in R["PASSES"].values())
    with open(args.out_json, "w") as f:
        json.dump(R, f, indent=1, default=str)
    print("PREFLIGHT_JSON")
    print(json.dumps({"PASSES": R["PASSES"], "GATE_3_READY": R["GATE_3_READY"],
                      "gpu_hours_this_invocation": R["gpu_hours_this_invocation"]}))
    print("/PREFLIGHT_JSON")


if __name__ == "__main__":
    main()