phase4_prep: free CPU rehearsal of everything a GPU session can fail -- pinned fetch, creds, pointer read, planner, full sha256 of the published mix, per-leg window arithmetic, resume-equivalence at real boundaries

kernels/phase4_prep.py

@@ -0,0 +1,176 @@
+# Phase 4 rehearsal on a CPU instance: every check a session can fail, run where failing costs no quota.
+#
+# The first GPU session of the run bills ~11 hours before it is provably healthy, and E-031 is exactly the
+# shape of bug that only shows up mid-run (a dataset-length guard that compared the wrong quantity would
+# have refused the run's third leg). So this kernel does, on free CPU: fetches the pinned code and asserts
+# its sha256s, resolves credentials, reads the checkpoint pointer, downloads the PUBLISHED mix from the
+# Hub, hashes every shard against the manifest, rehearses the session planner for the whole schedule, and
+# proves resume-equivalence of the reader at the real checkpoint boundaries with the real seed.
+#
+# It ends with the same PLAN_JSON the GPU session will print, so the numbers can be compared before and
+# after rather than trusting that the code did not change in between.
+import hashlib, json, os, subprocess, sys, time
+
+os.chdir("/kaggle/working")
+sys.path.insert(0, "/kaggle/working")
+REV = os.environ.get("OUNCE100M_REV") or "a6ae45de71482ecb6c810a25368db665f88e672d"
+WANT = {
+  "ounce100m_credentials.py": ("ounce100m_credentials.py",
+                               "6525f62f03f2d73650a1eb4f70fcb52d1194caad4ca88b2d8bd8fd54f88339b6"),
+  "shard_dataset.py": ("train/shard_dataset.py",
+                       "4111832e79445fb0ffce6a73a02da4c97255d6c7886762682e04a417f2695a78"),
+  "hubckpt.py": ("train/hubckpt.py",
+                 "c8c958417dc48db92f8be5c4b505581c746be9cc10110a358230ac47c480d5bf"),
+  "train_ounce100m.py": ("train/train_ounce100m.py",
+                         "487edc1be54f03b5f1a312901580cc3e1504c412322ed074ceff869febbfb8bc"),
+  "phase4_session.py": ("kernels/phase4_session.py",
+                        "a52c758ccca3f74d518d6fd1a765383c2851ea5893883429fba4c8f8b95bdfb2"),
+}
+BASE = "https://huggingface.co/Cion-lab/ounce100m-code/resolve/" + REV
+for p, (rp, want) in sorted(WANT.items()):
+    r = subprocess.run(["curl", "-sfL", BASE + "/" + rp, "-o", p], capture_output=True, text=True)
+    assert r.returncode == 0, ("fetch failed", rp, (r.stderr or "")[:200])
+    got = hashlib.sha256(open(p, "rb").read()).hexdigest()
+    assert got == want, ("SHA MISMATCH", rp, got, want)
+    print("OK", p, got[:12], flush=True)
+print("REV_OK", REV, flush=True)
+
+import ounce100m_credentials as C
+print("creds:", json.dumps(C.install(verify=True)), flush=True)
+T0 = time.time()
+
+
+def run(argv, label, timeout, env=None):
+    print("=== " + label, flush=True)
+    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 threading
+    killed = []
+    timer = threading.Timer(timeout, lambda: (killed.append(True), p.kill()))
+    timer.daemon = True
+    timer.start()
+    out = []
+    try:
+        for line in p.stdout:
+            out.append(line.rstrip("\n"))
+            print("   |", line[:300], flush=True)
+    finally:
+        timer.cancel()
+    rc = p.wait()
+    if killed:
+        print("   TIMEOUT after %d s, killed" % timeout, flush=True)
+        rc = -9
+    print("%s_RC %s seconds %.1f elapsed %.0f" % (label, rc, time.time() - t0, time.time() - T0),
+          flush=True)
+    return rc, "\n".join(out)
+
+
+# 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.
+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"})
+if rc1 != 0:
+    print("VERDICT PREP_STOP launcher rehearsal rc", rc1, flush=True)
+    raise SystemExit(2)
+
+# Stage 2 -- the published mix, byte for byte, and the reader at the boundaries the run will actually
+# cross. `snapshot_download` proves the repo is anonymously readable; this proves the bytes are the ones
+# the manifest certifies and that a resume at step X reads the same data an uninterrupted run would have.
+CHECK = """
+import hashlib, json, os, sys
+sys.path.insert(0, "/kaggle/working")
+ROOT = "/kaggle/working/mixroot"
+man = json.load(open(os.path.join(ROOT, "manifest.json")))
+tok_s = int(man["total_tokens"])
+bad, nb = [], 0
+for grp in ("shards", "val_shards"):
+    for rec in man.get(grp, []):
+        p = os.path.join(ROOT, rec["file"])
+        if not os.path.exists(p):
+            bad.append((rec["file"], "absent")); continue
+        h = hashlib.sha256()
+        with open(p, "rb") as fh:
+            for chunk in iter(lambda: fh.read(1 << 22), b""):
+                h.update(chunk)
+        nb += 1
+        if h.hexdigest() != rec["sha256"]:
+            bad.append((rec["file"], "sha " + h.hexdigest()[:12] + " != " + rec["sha256"][:12]))
+print("SHASHED", nb, "files", "total_tokens", format(tok_s, ","), "mismatches", len(bad), flush=True)
+for b in bad[:10]:
+    print("  BAD", b, flush=True)
+assert not bad, "the published mix does not match its own manifest"
+
+import shard_dataset as SD
+SEQ, MICRO, ACCUM, WORLD = 1024, 4, 32, 2
+seqs_per_step = MICRO * WORLD * ACCUM
+tps = seqs_per_step * SEQ
+steps = int(1_000_000_000 / tps)
+store = SD.PackedTokenStore(ROOT)
+n_samples = store.total_tokens // SEQ
+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)
+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
+    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
+
+# Resume-equivalence at three real boundaries, on the published bytes with the run's own seed.
+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)
+    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()
+    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)
+
+# 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
+# sibling module surfaced only after a GPU session had already been billed).
+rc3, _ = run([sys.executable, "-c",
+              "import sys; sys.path.insert(0, '/kaggle/working')\n"
+              "import shard_dataset as SD, json, os\n"
+              "os.makedirs('/tmp/cur', exist_ok=True)\n"
+              "c = SD.cursor_dict(1024, 97536, 20260919, dataset_files_sha='ab' * 32, step=381)\n"
+              "SD.save_cursor('/tmp/cur', c)\n"
+              "back = SD.load_cursor('/tmp/cur')\n"
+              "print('CURSOR', json.dumps(back), 'roundtrip', back == c)\n"
+              "assert back == c\n"
+              "import subprocess\n"
+              "p = subprocess.run([sys.executable, 'train_ounce100m.py', '--help'],\n"
+              "                   capture_output=True, text=True)\n"
+              "print('TRAINER_HELP_RC', p.returncode, 'lines', len((p.stdout or '').splitlines()))\n"
+              "assert p.returncode == 0, (p.stdout or '')[-1500:] + (p.stderr or '')[-1500:]\n"],
+             "CURSOR_AND_IMPORT", 900)
+print("VERDICT PHASE4_PREP rc_launcher", rc1, "rc_check", rc2, "rc_cursor", rc3,
+      "seconds", round(time.time() - T0, 1), flush=True)
+raise SystemExit(0 if (rc1 == rc2 == rc3 == 0) else 4)