phase4_prep: pin f6f10076; drive the leg check from the real planner, verify the manifest against itself and the mapped bytes, resume through load_cursor

kernels/phase4_prep.py

@@ -13,18 +13,18 @@ import hashlib, json, os, subprocess, sys, time
 
 os.chdir("/kaggle/working")
 sys.path.insert(0, "/kaggle/working")
-REV = os.environ.get("OUNCE100M_REV") or "eb1b3c0033289c9c73996555d8fa6c117e8d45c0"
+REV = os.environ.get("OUNCE100M_REV") or "f6f100763eec8782851f1b418fe869df6a4dcb92"
 WANT = {
   "ounce100m_credentials.py": ("ounce100m_credentials.py",
                                "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
   "shard_dataset.py": ("train/shard_dataset.py",
                        "adbcd96fbab9505a5e8dec4d2e83d00366a73f702786bffbde430195cf6f0471"),
   "hubckpt.py": ("train/hubckpt.py",
-                 "c8c958417dc48db92f8be5c4b505581c746be9cc10110a358230ac47c480d5bf"),
+                 "568c31b300906cb8d78a59ef890baa9731b18e064e3d69b0eaea4b45c546f23e"),
   "train_ounce100m.py": ("train/train_ounce100m.py",
-                         "487edc1be54f03b5f1a312901580cc3e1504c412322ed074ceff869febbfb8bc"),
+                         "c4c178c0ac17c68de1521d85612a0dc604e6d4748dfdc8bd3e374512bc90ce96"),
   "phase4_session.py": ("kernels/phase4_session.py",
-                        "a8cdac340a32d8a2f179d819bf526a048af706eef0cd94081519575f5645c072"),
+                        "6e55edab4e9eedd78363c572ebdd409558e0c3f76e13f4ebfbe54e6c4dd4964d"),
 }
 BASE = "https://huggingface.co/Cion-lab/ounce100m-code/resolve/" + REV
 for p, (rp, want) in sorted(WANT.items()):
@@ -45,11 +45,22 @@ def run(argv, label, timeout, env=None):
     t0 = time.time()
     e = dict(os.environ); e["PYTHONPATH"] = "/kaggle/working"
     e.update(env or {})
-    p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
-                         env=e, bufsize=1)
+    import signal as _sig
     import threading
+    p = subprocess.Popen(argv, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
+                         env=e, bufsize=1, start_new_session=True)
     killed = []
-    timer = threading.Timer(timeout, lambda: (killed.append(True), p.kill()))
+
+    def _kill():
+        # Kill the group, not the child: a staged `python -c snapshot_download` orphaned by a timeout
+        # keeps writing into mixroot while the next stage reads it.
+        killed.append(True)
+        try:
+            os.killpg(os.getpgid(p.pid), _sig.SIGTERM)
+        except Exception:
+            p.kill()
+
+    timer = threading.Timer(timeout, _kill)
     timer.daemon = True
     timer.start()
     out = []
@@ -71,12 +82,12 @@ def run(argv, label, timeout, env=None):
 # Stage 1 -- the launcher itself, in rehearsal mode. On a CPU box torchrun is never reached: PREP_ONLY
 # stops right after the pointer read, the mix download and the planner. Everything upstream of the first
 # optimiser step is exercised with the real files and the real Hub.
+PLANENV = {"PLANNING_TOK_PER_S": os.environ.get("PLANNING_TOK_PER_S", "9358"),
+           "SESSION_GPU_HOURS": os.environ.get("SESSION_GPU_HOURS", "6.9"),
+           "QUOTA_LEFT_HOURS": os.environ.get("QUOTA_LEFT_HOURS", "27.4"),
+           "MIN_FREE_GB": "6"}
 rc1, _ = run([sys.executable, "phase4_session.py", REV], "PREP_LAUNCHER", 3000,
-             env={"PHASE4_PREP_ONLY": "1",
-                  "PLANNING_TOK_PER_S": os.environ.get("PLANNING_TOK_PER_S", "9696"),
-                  "SESSION_GPU_HOURS": os.environ.get("SESSION_GPU_HOURS", "11.0"),
-                  "QUOTA_LEFT_HOURS": os.environ.get("QUOTA_LEFT_HOURS", "27.5"),
-                  "MIN_FREE_GB": "6"})
+             env=dict(PLANENV, PHASE4_PREP_ONLY="1"))
 if rc1 != 0:
     print("VERDICT PREP_STOP launcher rehearsal rc", rc1, flush=True)
     raise SystemExit(2)
@@ -108,50 +119,83 @@ for b in bad[:10]:
     print("  BAD", b, flush=True)
 assert not bad, "the published mix does not match its own manifest"
 
+import phase4_session as S          # constants and plan() only -- main() is behind __name__
 import shard_dataset as SD
-SEQ, MICRO, ACCUM, WORLD = 1024, 4, 32, 2
-seqs_per_step = MICRO * WORLD * ACCUM
+SEQ, steps = S.SEQ_LEN, S.HORIZON_STEPS
+seqs_per_step = S.MICRO_BATCH * S.WORLD * S.ACCUM       # the launcher's own geometry, not retyped
 tps = seqs_per_step * SEQ
-steps = int(1_000_000_000 / tps)
+assert tps == S.TOKENS_PER_STEP, f"prep disagrees with the launcher about tokens/step: {tps}"
+assert steps == int(S.TOKENS / tps), f"horizon disagrees: {steps}"
 store = SD.PackedTokenStore(ROOT)
 n_samples = store.total_tokens // SEQ
+fingerprint = SD.mix_fingerprint(ROOT, man)
+# The manifest is the corpus definition, and three different things have to agree about it: the record
+# list, the counts it claims, and the bytes the reader actually mapped. A rebuild that truncated the list
+# while leaving total_tokens alone would otherwise certify 1.11 B tokens of a smaller corpus (review).
+if len(man["shards"]) != int(man.get("n_shards") or -1) or len(store.shards) != len(man["shards"]):
+    raise AssertionError(f"shard records {len(man['shards'])} vs n_shards {man.get('n_shards')} vs "
+                         f"files opened {len(store.shards)}")
+summed = sum(int(s["tokens"]) for s in man["shards"])
+if summed != tok_s or store.total_tokens != summed:
+    raise AssertionError(f"manifest does not add up: records sum {summed}, total_tokens says {tok_s}, "
+                         f"bytes the reader mapped say {store.total_tokens}")
 print("READER", format(store.total_tokens, ","), "tokens in", len(store.shards), "shards ->",
       format(n_samples, ","), "windows of", SEQ, "|", steps, "steps x", tps, "tok =",
-      format(steps * tps, ","), "tokens", flush=True)
+      format(steps * tps, ","), "tokens | fingerprint", fingerprint, flush=True)
 assert n_samples >= steps * seqs_per_step, "not enough windows for the horizon"
 
-# Every leg the planner can produce, from every checkpoint boundary to the next.
-fails = []
-for start in range(0, steps + 1, 381):
-    stop = min(start + 381, steps)
-    need = (steps - start) * seqs_per_step          # the trainer's own guard
+# The schedule the run will actually live through: drive the real planner from step 0 until it reports the
+# horizon, and require every leg to start where the previous one stopped. This is the check that catches a
+# planner which stops one interval short of the end -- the last four steps of the run, unreachable, while
+# every later session fails its own budget test for want of a whole interval.
+legs, start = [], 0
+while start < steps and len(legs) < 40:
+    pl = S.plan(start, 15.0)
+    if not pl["planned_steps"]:
+        raise AssertionError(f"the planner refuses to start at step {start} with "
+                             f"{pl['usable_hours']} usable hours -- the run can never finish")
+    end = steps if pl["stop_after_steps"] == 0 else pl["stop_after_steps"]
+    assert end > start, (start, end)
+    need = (steps - start) * seqs_per_step              # the trainer's own guard, at this start step
     avail = n_samples - start * seqs_per_step
-    ok = avail >= need and (stop - start) * seqs_per_step <= avail
-    if not ok:
-        fails.append((start, stop, need, avail))
-    print("LEG start %5d stop %5d need %9d avail %9d %s" % (
-        start, stop, need, avail, "ok" if ok else "FAIL"), flush=True)
-print("LEGS", 1 + steps // 381, "checked,", len(fails), "failing", flush=True)
-assert not fails, fails
+    assert avail >= need, (start, need, avail)
+    store.tokens((end * seqs_per_step - 1) * SEQ, SEQ)  # last window this leg reads must be addressable
+    legs.append((start, end))
+    print("LEG %5d -> %5d  pct %5.1f -> %5.1f  need %9d avail %9d" % (
+        start, end, 100.0 * start / steps, 100.0 * end / steps, need, avail), flush=True)
+    start = end
+print("SESSIONS", len(legs), "ends_at", legs[-1][1], "horizon", steps, "contiguous",
+      [l[0] for l in legs[1:]] == [l[1] for l in legs[:-1]], flush=True)
+assert legs[-1][1] == steps and [l[0] for l in legs[1:]] == [l[1] for l in legs[:-1]], legs
 
-# Resume-equivalence at three real boundaries, on the published bytes with the run's own seed.
+# Resume-equivalence through the cursor the trainer would actually write, at boundaries the run crosses:
+# the first checkpoint, a mid-session one, and the last interval before the horizon. Building the resumed
+# dataset from load_cursor() rather than a hand-computed offset is the point -- the derivation from
+# step to samples_consumed is part of what is being tested.
 full = SD.make_dataset(store, SEQ, 20260919, start_sample=0)
-for boundary in (381, 1905, 3429):
-    s0 = boundary * seqs_per_step
-    resumed = SD.make_dataset(store, SEQ, 20260919, start_sample=s0)
+for boundary in sorted({S.CKPT_EVERY, legs[len(legs) // 2][0], steps - S.CKPT_EVERY}):
+    d = "/tmp/cur_%d" % boundary
+    os.makedirs(d, exist_ok=True)
+    SD.save_cursor(d, SD.cursor_dict(SEQ, boundary * seqs_per_step, 20260919, fingerprint,
+                                     step=boundary))
+    c = SD.load_cursor(d)
+    assert c["step"] * seqs_per_step == c["samples_consumed"], c
+    assert c["dataset_files_sha"] == fingerprint, c
+    s0 = c["samples_consumed"]
+    resumed = SD.make_dataset(store, SEQ, c["shuffle_seed"], start_sample=s0)
     assert len(resumed) == len(full) - s0, (len(resumed), len(full), s0)
     same = True
     for j in (0, 1, len(resumed) // 2, len(resumed) - 1):
-        a = full[s0 + j]["input_ids"].tolist()
-        b = resumed[j]["input_ids"].tolist()
-        same = same and a == b and full[s0 + j]["labels"].tolist() == resumed[j]["labels"].tolist()
+        a, b = full[s0 + j], resumed[j]
+        same = same and a["input_ids"].tolist() == b["input_ids"].tolist() \
+            and a["labels"].tolist() == b["labels"].tolist()
     print("RESUME_EQUIV step", boundary, "start_sample", s0, "len", len(resumed), "identical", same,
           flush=True)
     assert same, boundary
 print("CHECK_OK windows", format(n_samples, ","), "shards", len(store.shards), flush=True)
 """
 open("prepcheck.py", "w").write(CHECK)
-rc2, _ = run([sys.executable, "prepcheck.py"], "CHECK_PUBLISHED_MIX", 2400)
+rc2, _ = run([sys.executable, "prepcheck.py"], "CHECK_PUBLISHED_MIX", 2400, env=dict(PLANENV))
 
 # Stage 3 -- the reader's cursor format round-trips, and the trainer imports and prints --help on this
 # image. Both are cheap and both have bitten before (E-028: a build script that could not import its