stage listing raises instead of returning {}; empty-repo case handled explicitly; verify refuses manifest-less PASS; +val holdout in merge

build/build_mix.py

@@ -45,6 +45,9 @@ import time
 MAGIC_FLAGS = 0
 TOKENS_PER_SHARD = 8_000_000    # 16 MB as uint16; ~140 shards for 1.1B tokens
 MERGE_BATCH_DOCS = 32           # granularity of the round-robin; ~20k tokens, so a 262k batch spans ~13
+VAL_STRIDE = 50                 # every 50th document is held out ~= 2% of tokens, sampled across
+                                # sources rather than whole documents, so validation PPL is honest
+VAL_TOKENS_PER_SHARD = 2_000_000
 ROW_TOKEN_CAP = 16_384          # one document must never dominate a shard
 SHUFFLE_SEED = 20260919         # recorded so a rebuild is reproducible
 HASHES_FILE = "seen_hashes.bin"
@@ -412,10 +415,36 @@ def merge(args, st, tokenizer):
 
     cap = args.merge_budget_tokens or int(total_w)
     buf, buf_src, buf_tokens = [], [], 0
+    val_buf, val_tokens, docs_seq = [], 0, 0
     out_dir = os.path.join(st.root, "mix")
     os.makedirs(out_dir, exist_ok=True)
     t0 = time.monotonic()
 
+    def flush_val():
+        nonlocal val_buf, val_tokens
+        if not val_buf:
+            return
+        idx = len(st.d.get("val_shards", []))
+        p = os.path.join(st.root, "val", f"shard-{idx:04d}.bin")
+        os.makedirs(os.path.dirname(p), exist_ok=True)
+        tmp = p + ".tmp"
+        offs = array.array("I", [0])
+        ids = array.array("H")
+        for d in val_buf:
+            ids.extend(d)
+            offs.append(len(ids))
+        with open(tmp, "wb") as f:
+            f.write(struct.pack("<IIII", vocab, len(val_buf), len(ids), MAGIC_FLAGS))
+            offs.tofile(f)
+            ids.tofile(f)
+        os.replace(tmp, p)
+        st.d.setdefault("val_shards", []).append({
+            "file": os.path.relpath(p, st.root), "docs": len(val_buf), "tokens": len(ids),
+            "sha256": hashlib.sha256(open(p, "rb").read()).hexdigest(), "bytes": os.path.getsize(p)})
+        st.commit([], [])
+        print(f"  val shard {idx:04d}: {len(ids):,} tok / {len(val_buf):,} docs", flush=True)
+        val_buf, val_tokens = [], 0
+
     def flush():
         nonlocal buf, buf_src, buf_tokens
         if not buf:
@@ -446,6 +475,17 @@ def merge(args, st, tokenizer):
             d = next_doc(pick)
             if d is None:
                 break
+            docs_seq += 1
+            if docs_seq % VAL_STRIDE == 0:
+                # held out of training entirely: routed to the validation split instead of the buffer
+                val_buf.append(d)
+                val_tokens += len(d)
+                if val_tokens >= VAL_TOKENS_PER_SHARD:
+                    flush_val()
+                pick["emitted_tokens"] += len(d)
+                pick["pending"] = None
+                took += 1
+                continue
             buf.append((d, pick["key"]))
             buf_tokens += len(d)
             pick["emitted_tokens"] += len(d)
@@ -457,6 +497,9 @@ def merge(args, st, tokenizer):
         if all(s["pending"] is None and s["exhausted"] for s in streams):
             break
     flush()
+    flush_val()
+    print(f"merge: {len(st.d['val_shards']) if 'val_shards' in st.d else 0} val shards, "
+          f"{sum(s['tokens'] for s in st.d.get('val_shards', [])):,} held-out tokens", flush=True)
     print(f"merge: {len(st.d['merged_shards'])} shards, "
           f"{sum(s['tokens'] for s in st.d['merged_shards']):,} tokens", flush=True)
 
@@ -555,7 +598,11 @@ def main():
                   for k, v in st.d["staged"].items()},
         "staged_tokens": st.staged_tokens_total(),
         "shards": st.d["merged_shards"],
+        "val_shards": st.d.get("val_shards", []),
         "total_tokens": sum(s["tokens"] for s in st.d["merged_shards"]),
+        "val_tokens": sum(s["tokens"] for s in st.d.get("val_shards", [])),
+        "val_holdout": {"stride": VAL_STRIDE, "note": "every 50th merged document, sampled across "
+                         "sources rather than whole documents, excluded from training shards"},
         "distinct_sources_per_shard_min": min((len(s["sources"]) for s in st.d["merged_shards"]),
                                               default=0),
         "seconds": round(time.monotonic() - t0, 1),

build/hubsync.py

@@ -44,24 +44,58 @@ def ensure_repo(repo):
     return repo
 
 
+class StageListingError(RuntimeError):
+    """Raised when the stage repo's contents cannot be determined.
+
+    Returning an empty set here would be the single most dangerous shortcut in this pipeline: an empty
+    listing is indistinguishable from "nothing staged yet", so a *broken* listing makes a resume
+    silently re-stage everything (wasting hours) or, worse, makes a merge believe there is nothing to
+    merge while still looking consistent. Callers must see the difference between empty and unknown."""
+
+
 def remote_stage_files(repo):
-    """Set of relative paths under stage/ that already exist on the Hub."""
+    """Relative paths under stage/ that exist on the Hub. Raises rather than returning {} on failure."""
+    import json
     import urllib.request
 
-    url = f"https://huggingface.co/api/datasets/{repo}/tree/main/stage?recursive=true"
+    errors = []
+    a = api()
     try:
-        with urllib.request.urlopen(url, timeout=90) as r:
-            entries = json.loads(r.read().decode())
-    except Exception:
-        return set()
-    out = set()
-    for e in entries:
-        if e.get("type") != "file":
-            continue
-        out.add(e["path"])
-        # a directory in the tree listing means we must recurse; recursive=true should cover it, but
-        # if the server ignored the flag we get dirs and no files, which the caller detects as empty
-    return out
+        files = a.list_repo_files(repo_id=repo, repo_type="dataset", revision="main")
+        return {f for f in files if f.startswith("stage/")}
+    except Exception as e:
+        errors.append(f"list_repo_files: {type(e).__name__}: {str(e)[:160]}")
+        # A freshly created repo has no commits, so "main" does not resolve and list_repo_files raises
+        # for a reason that genuinely means empty. Distinguish it by asking whether the repo exists at
+        # all -- NOT by swallowing the exception, which is what made the first version ambiguous.
+        try:
+            info = a.repo_info(repo_id=repo, repo_type="dataset")
+            if not getattr(info, "sha", None):
+                return set()
+        except Exception as e2:
+            errors.append(f"repo_info: {type(e2).__name__}: {str(e2)[:120]}")
+
+    # REST fallback. Note the path goes in the query string, not the URL body -- the first version put
+    # it in the path, got a 404, and the bare `except` turned that into a silent empty listing (E-018).
+    for url in (f"https://huggingface.co/api/datasets/{repo}/tree/main?recursive=true",
+                f"https://huggingface.co/api/datasets/{repo}/tree/main"):
+        try:
+            with urllib.request.urlopen(url, timeout=90) as r:
+                entries = json.loads(r.read().decode())
+            paths = set()
+            stack = [(e, "") for e in entries]
+            while stack:
+                e, prefix = stack.pop()
+                p = f"{prefix}/{e['path']}".lstrip("/") if prefix else e.get("path", "")
+                if e.get("type") == "dir":
+                    stack.extend((c, p) for c in (e.get("children") or []))
+                elif p.startswith("stage/"):
+                    paths.add(p)
+            return paths
+        except Exception as e:
+            errors.append(f"tree {url.split('/api/')[1]}: {type(e).__name__}: {str(e)[:140]}")
+
+    raise StageListingError(f"cannot list {repo}: " + " | ".join(errors))
 
 
 def upload_source(root, key, repo):

build/smoke_hub_resume.py

@@ -1,100 +1,106 @@
-# Cold-resume rehearsal for the mix builder, in one session.
-#
-# What it proves and why it matters: /kaggle/working is destroyed when a Kaggle session ends, so the
-# staged shards AND build_state.json are gone even though the Hub still holds the stage. The full build is
-# hours of streaming; if a session kill meant starting that over, the build would never finish. This
-# script therefore:
-#   1. stages two small sources and publishes them to a throwaway stage repo,
-#   2. WIPES the local tree -- a faithful simulation of the instance dying, not a soft restart,
-#   3. runs --merge-only against the Hub alone, which must restore the stage bookkeeping from the Hub's
-#      record.json files and emit a mixed shard containing both sources,
-#   4. runs verify_mix.py over the result, and
-#   5. deletes the throwaway repo so no test artifact is mistaken for the real stage later.
-#
-# Step 3 is the one that would catch a silent failure: with an empty state and no restore path, merge
-# would find nothing and emit an empty mix while still exiting 0.
-#
-# CPU session, zero GPU quota. No credentials in this file; they come from the private store (D-006).
-
-import json
-import os
-import shutil
-import subprocess
-import sys
-import time
-
-ROOT = "/kaggle/working/mixroot"
-REPO = "Cion-lab/ounce100m-mix-stagetest"
-TWO = "cosmopedia__khanacademy,SimpleStories__default"
-
-
-def sh(argv, label):
-    print(f"=== {label}", flush=True)
-    p = subprocess.run([sys.executable] + argv, cwd="/kaggle/working", capture_output=True, text=True,
-                       timeout=3600)
-    out = p.stdout
-    for line in out.splitlines():
-        print("   |", line[:220], flush=True)
-    if p.returncode != 0:
-        print("   STDERR:", (p.stderr or "")[-1500:], flush=True)
-    return {"rc": p.returncode, "out": out[-6000:], "err": (p.stderr or "")[-1500:]}
-
-
-def grab_json(text, begin, end):
-    if begin in text and end in text:
-        try:
-            return json.loads(text.split(begin, 1)[1].split(end, 1)[0])
-        except Exception:
-            return None
-    return None
-
-
-def main():
-    R = {"repo": REPO, "sources": TWO}
-    base = ["build_mix.py", "--root", ROOT, "--only", TWO, "--cap-tokens-per-source", "1200000",
-            "--hub-repo", REPO]
-
-    shutil.rmtree(ROOT, ignore_errors=True)
-    r1 = sh(base + ["--stage-only"], "1. stage two sources and publish")
-    R["stage"] = {"rc": r1["rc"], "json": grab_json(r1["out"], "BUILD_JSON_BEGIN", "BUILD_JSON_END")}
-
-    # the wipe IS the test: after this, nothing local survives except this script and the code
-    shutil.rmtree(ROOT, ignore_errors=True)
-    R["wiped"] = {"root_exists": os.path.exists(ROOT),
-                  "stage_dir_exists": os.path.exists(os.path.join(ROOT, "stage"))}
-
-    r2 = sh(base + ["--merge-only"], "2. cold resume: merge from the Hub with an empty local tree")
-    R["merge"] = {"rc": r2["rc"], "json": grab_json(r2["out"], "BUILD_JSON_BEGIN", "BUILD_JSON_END")}
-
-    r3 = sh(["verify_mix.py", "--root", ROOT, "--min-sources-per-shard", "2",
-             "--sample-decodes", "3"], "3. verify the restored mix")
-    R["verify"] = grab_json(r3["out"], "VERIFY_JSON_BEGIN", "VERIFY_JSON_END")
-
-    staged = (R["merge"].get("json") or {}).get("total_tokens", 0)
-    R["COLD_RESUME_PASSED"] = bool(
-        r1["rc"] == 0 and r2["rc"] == 0 and staged > 1_000_000
-        and (R["verify"] or {}).get("PASS") is True
-        and (R["verify"] or {}).get("distinct_sources_in_mix") == 2)
-
-    # cleanup: the throwaway stage repo must not outlive the test, or a later real build could restore
-    # capped 1.2M-token sources and believe they were complete
-    def cleanup():
-        import ounce100m_credentials
-        ounce100m_credentials.install()
-        from huggingface_hub import delete_repo
-        delete_repo(repo_id=REPO, repo_type="dataset", token=os.environ["HF_TOKEN"])
-        return "deleted"
-
-    try:
-        R["cleanup"] = cleanup()
-    except Exception as e:
-        R["cleanup"] = {"error": f"{type(e).__name__}: {str(e)[:200]}",
-                        "action_needed": f"delete the throwaway dataset {REPO} manually"}
-
-    print("SMOKE_JSON_BEGIN")
-    print(json.dumps(R, indent=1, default=str)[:6000])
-    print("SMOKE_JSON_END")
-
-
-if __name__ == "__main__":
-    main()
+# Cold-resume rehearsal for the mix builder, in one session.
+#
+# What it proves and why it matters: /kaggle/working is destroyed when a Kaggle session ends, so the
+# staged shards AND build_state.json are gone even though the Hub still holds the stage. The full build is
+# hours of streaming; if a session kill meant starting that over, the build would never finish. This
+# script therefore:
+#   1. stages two small sources and publishes them to a throwaway stage repo,
+#   2. WIPES the local tree -- a faithful simulation of the instance dying, not a soft restart,
+#   3. runs --merge-only against the Hub alone, which must restore the stage bookkeeping from the Hub's
+#      record.json files and emit a mixed shard containing both sources,
+#   4. runs verify_mix.py over the result, and
+#   5. deletes the throwaway repo so no test artifact is mistaken for the real stage later.
+#
+# Step 3 is the one that would catch a silent failure: with an empty state and no restore path, merge
+# would find nothing and emit an empty mix while still exiting 0.
+#
+# CPU session, zero GPU quota. No credentials in this file; they come from the private store (D-006).
+
+import json
+import os
+import shutil
+import subprocess
+import sys
+import time
+
+ROOT = "/kaggle/working/mixroot"
+REPO = "Cion-lab/ounce100m-mix-stagetest"
+TWO = "cosmopedia__khanacademy,SimpleStories__default"
+
+
+def sh(argv, label):
+    print(f"=== {label}", flush=True)
+    p = subprocess.run([sys.executable] + argv, cwd="/kaggle/working", capture_output=True, text=True,
+                       timeout=3600)
+    out = p.stdout
+    for line in out.splitlines():
+        print("   |", line[:220], flush=True)
+    if p.returncode != 0:
+        print("   STDERR:", (p.stderr or "")[-1500:], flush=True)
+    return {"rc": p.returncode, "out": out[-6000:], "err": (p.stderr or "")[-1500:]}
+
+
+def grab_json(text, begin, end):
+    if begin in text and end in text:
+        try:
+            return json.loads(text.split(begin, 1)[1].split(end, 1)[0])
+        except Exception:
+            return None
+    return None
+
+
+def main():
+    R = {"repo": REPO, "sources": TWO}
+    base = ["build_mix.py", "--root", ROOT, "--only", TWO, "--cap-tokens-per-source", "1200000",
+            "--hub-repo", REPO]
+
+    shutil.rmtree(ROOT, ignore_errors=True)
+    r1 = sh(base + ["--stage-only"], "1. stage two sources and publish")
+    R["stage"] = {"rc": r1["rc"], "json": grab_json(r1["out"], "BUILD_JSON_BEGIN", "BUILD_JSON_END")}
+
+    # the wipe IS the test: after this, nothing local survives except this script and the code
+    shutil.rmtree(ROOT, ignore_errors=True)
+    R["wiped"] = {"root_exists": os.path.exists(ROOT),
+                  "stage_dir_exists": os.path.exists(os.path.join(ROOT, "stage"))}
+
+    r2 = sh(base + ["--merge-only"], "2. cold resume: merge from the Hub with an empty local tree")
+    R["merge"] = {"rc": r2["rc"], "json": grab_json(r2["out"], "BUILD_JSON_BEGIN", "BUILD_JSON_END")}
+
+    r3 = sh(["verify_mix.py", "--root", ROOT, "--min-sources-per-shard", "2",
+             "--sample-decodes", "3"], "3. verify the restored mix")
+    R["verify"] = grab_json(r3["out"], "VERIFY_JSON_BEGIN", "VERIFY_JSON_END")
+
+    staged = (R["merge"].get("json") or {}).get("total_tokens", 0)
+    R["COLD_RESUME_PASSED"] = bool(
+        r1["rc"] == 0 and r2["rc"] == 0 and staged > 1_000_000
+        and (R["verify"] or {}).get("PASS") is True
+        and (R["verify"] or {}).get("distinct_sources_in_mix") == 2)
+
+    # cleanup: the throwaway stage repo must not outlive the test, or a later real build could restore
+    # capped 1.2M-token sources and believe they were complete
+    def cleanup():
+        import ounce100m_credentials
+        ounce100m_credentials.install()
+        from huggingface_hub import delete_repo
+        delete_repo(repo_id=REPO, repo_type="dataset", token=os.environ["HF_TOKEN"])
+        return "deleted"
+
+    try:
+        # Only clean up on success. If the rehearsal failed, the stage repo IS the evidence and deleting
+        # it destroys the thing the next session would need in order to diagnose the failure.
+        if R["COLD_RESUME_PASSED"]:
+            R["cleanup"] = cleanup()
+        else:
+            R["cleanup"] = ("KEPT for diagnosis - delete manually or it will be restored as a real "
+                            f"stage source: {REPO}")
+    except Exception as e:
+        R["cleanup"] = {"error": f"{type(e).__name__}: {str(e)[:200]}",
+                        "action_needed": f"delete the throwaway dataset {REPO} manually"}
+
+    print("SMOKE_JSON_BEGIN")
+    print(json.dumps(R, indent=1, default=str)[:6000])
+    print("SMOKE_JSON_END")
+
+
+if __name__ == "__main__":
+    main()

build/verify_mix.py

@@ -1,176 +1,188 @@
-# Gate 2 verification: recount and re-derive the mix from its shard files, independently of the builder.
-#
-# A number checked by the code that produced it proves nothing about a truncated write, a bad header, or a
-# source that silently contributed nothing. This re-derives everything from the artifacts:
-#   1. sha256 of every shard matches the manifest        (catches corrupt/partial uploads)
-#   2. header n_tokens == byte length, no trailing bytes (catches torn writes)
-#   3. offsets[0] == 0, strictly non-decreasing, offsets[-1] == n_tokens (catches a corrupt index)
-#   4. every id < vocab                                   (an out-of-range id is a train-time crash)
-#   5. total tokens recounted from the bytes, not from counters
-#   6. per-shard SOURCE MIX, and it is asserted as a gate, not reported as a diagnostic -- this is the
-#      check that caught E-016, where shard 0 was 100% fineweb-edu while every other check passed
-#   7. the (shard_index, token_offset) cursor resolves to stable ids, and document boundaries from the
-#      offset table actually reconstruct the document the builder tokenised
-#
-# Read-only. CPU session. No credentials.
-
-import argparse
-import array
-import hashlib
-import json
-import os
-import random
-import struct
-
-
-def sha256_file(path, chunk=8 * 1024**2):
-    h = hashlib.sha256()
-    with open(path, "rb") as f:
-        for b in iter(lambda: f.read(chunk), b""):
-            h.update(b)
-    return h.hexdigest()
-
-
-def read_shard(path):
-    with open(path, "rb") as f:
-        head = f.read(16)
-        if len(head) != 16:
-            raise ValueError(f"short header ({len(head)} bytes)")
-        vocab, n_docs, n_tokens, flags = struct.unpack("<IIII", head)
-        offs = array.array("I")
-        offs.fromfile(f, n_docs + 1)
-        ids = array.array("H")
-        ids.fromfile(f, n_tokens)
-        trailing = len(f.read())
-    return {"vocab": vocab, "n_docs": n_docs, "n_tokens": n_tokens, "flags": flags,
-            "offsets": offs, "ids": ids, "trailing": trailing}
-
-
-def main():
-    ap = argparse.ArgumentParser()
-    ap.add_argument("--root", default="/kaggle/working/mixroot")
-    ap.add_argument("--sample-decodes", type=int, default=6)
-    ap.add_argument("--cursor-probes", type=int, default=40)
-    ap.add_argument("--min-sources-per-shard", type=int, default=-1,
-                    help="-1 = half the distinct sources in the mix")
-    a = ap.parse_args()
-
-    man = json.load(open(os.path.join(a.root, "manifest.json")))
-    mix_vocab = man["tokenizer"]["vocab_size"]
-    out = {"shards_expected": len(man["shards"]), "sha_ok": 0, "sha_bad": [],
-           "header_ok": 0, "header_bad": [], "offset_bad": [], "out_of_range_ids": 0,
-           "total_tokens_recounted": 0, "total_docs_recounted": 0, "trailing_bytes": 0}
-
-    tok = None
-    try:
-        from tokenizers import Tokenizer
-        from huggingface_hub import hf_hub_download
-        tok = Tokenizer.from_file(hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json"))
-        tok.no_truncation()
-    except Exception as e:
-        out["tokenizer_note"] = f"decode checks disabled: {type(e).__name__}: {e}"[:160]
-
-    per_shard_sources = []
-    sizes = []
-    for rec in man["shards"]:
-        p = os.path.join(a.root, rec["file"])
-        if not os.path.exists(p):
-            out["header_bad"].append({"file": rec["file"], "error": "missing"})
-            continue
-        got = sha256_file(p)
-        if got != rec["sha256"]:
-            out["sha_bad"].append({"file": rec["file"], "want": rec["sha256"][:12], "got": got[:12]})
-            continue
-        out["sha_ok"] += 1
-        try:
-            sh = read_shard(p)
-        except Exception as e:
-            out["header_bad"].append({"file": rec["file"], "error": f"{type(e).__name__}: {e}"[:160]})
-            continue
-        sizes.append((rec["file"], sh["n_tokens"]))
-        o = sh["offsets"]
-        if (sh["vocab"] != mix_vocab or sh["trailing"] != 0 or sh["n_tokens"] != rec["tokens"]
-                or sh["n_docs"] != rec["docs"]):
-            out["header_bad"].append({"file": rec["file"], "vocab": sh["vocab"], "n": sh["n_tokens"],
-                                      "manifest_n": rec["tokens"], "trailing": sh["trailing"]})
-        else:
-            out["header_ok"] += 1
-        if len(o) != sh["n_docs"] + 1 or o[0] != 0 or o[-1] != sh["n_tokens"] or \
-                any(o[i] > o[i + 1] for i in range(len(o) - 1)):
-            out["offset_bad"].append(rec["file"])
-        out["out_of_range_ids"] += sum(1 for i in sh["ids"] if i >= mix_vocab)
-        out["total_tokens_recounted"] += sh["n_tokens"]
-        out["total_docs_recounted"] += sh["n_docs"]
-        out["trailing_bytes"] += sh["trailing"]
-        per_shard_sources.append((rec["file"], sorted(set(rec["sources"]))))
-
-    distinct = sorted({s for _f, sl in per_shard_sources for s in sl})
-    out["distinct_sources_in_mix"] = len(distinct)
-    out["sources_per_shard"] = {"min": min((len(s) for _f, s in per_shard_sources), default=0),
-                                "max": max((len(s) for _f, s in per_shard_sources), default=0),
-                                "median": sorted(len(s) for _f, s in per_shard_sources)[
-                                    len(per_shard_sources) // 2] if per_shard_sources else 0}
-    floor = a.min_sources_per_shard if a.min_sources_per_shard >= 0 else max(2, len(distinct) // 2)
-    out["min_sources_per_shard_required"] = floor
-    out["shards_below_source_floor"] = [f for f, sl in per_shard_sources if len(sl) < floor]
-
-    # tokens per source, recomputed from the shard headers rather than trusted from the manifest
-    src_tokens = {}
-    for rec in man["shards"]:
-        per = rec["tokens"] / max(1, len(rec["sources"]))
-        for s in rec["sources"]:
-            src_tokens[s] = src_tokens.get(s, 0) + per
-    out["approx_tokens_per_source_equal_share"] = {k: int(v) for k, v in sorted(src_tokens.items())}
-    out["note_per_source"] = ("equal split of a shard's tokens across its sources is an approximation; "
-                             "exact per-source totals live in the stage manifest")
-
-    if tok is not None and sizes:
-        rng = random.Random(11)
-        decs = []
-        for _ in range(a.sample_decodes):
-            fname, n = rng.choice(sizes)
-            sh = read_shard(os.path.join(a.root, fname))
-            i = rng.randrange(0, sh["n_docs"])
-            d = sh["ids"][sh["offsets"][i]:sh["offsets"][i + 1]]
-            decs.append({"shard": fname, "doc": i, "doc_tokens": len(d),
-                         "text": tok.decode(list(d))[:200]})
-        out["decoded_documents"] = decs
-
-        probes = mism = boundary_ok = 0
-        total = sum(n for _f, n in sizes)
-        for _ in range(a.cursor_probes):
-            g = rng.randrange(0, max(1, total - 64))
-            acc, si, lo = 0, 0, 0
-            for i, (_f, n) in enumerate(sizes):
-                if acc + n > g:
-                    si, lo = i, g - acc
-                    break
-                acc += n
-            a1 = read_shard(os.path.join(a.root, sizes[si][0]))
-            a2 = read_shard(os.path.join(a.root, sizes[si][0]))
-            if lo + 64 > a1["n_tokens"]:
-                continue
-            probes += 1
-            if list(a1["ids"][lo:lo + 64]) != list(a2["ids"][lo:lo + 64]):
-                mism += 1
-            # a document read via its offsets must equal the same document read twice
-            j = rng.randrange(0, a1["n_docs"])
-            d1 = list(a1["ids"][a1["offsets"][j]:a1["offsets"][j + 1]])
-            d2 = list(a2["ids"][a2["offsets"][j]:a2["offsets"][j + 1]])
-            if d1 == d2 and 0 < len(d1) <= 16_384:
-                boundary_ok += 1
-        out["cursor_probes"] = {"done": probes, "mismatches": mism, "boundaries_ok": boundary_ok}
-
-    out["manifest_claims_total_tokens"] = man["total_tokens"]
-    out["recount_matches_manifest"] = man["total_tokens"] == out["total_tokens_recounted"]
-    out["PASS"] = (not out["sha_bad"] and not out["header_bad"] and not out["offset_bad"]
-                   and not out["shards_below_source_floor"] and out["recount_matches_manifest"]
-                   and out["out_of_range_ids"] == 0 and out["trailing_bytes"] == 0
-                   and out.get("cursor_probes", {}).get("mismatches", 1) == 0)
-    print("VERIFY_JSON_BEGIN")
-    print(json.dumps(out, indent=1, default=str))
-    print("VERIFY_JSON_END")
-
-
-if __name__ == "__main__":
-    main()
+# Gate 2 verification: recount and re-derive the mix from its shard files, independently of the builder.
+#
+# A number checked by the code that produced it proves nothing about a truncated write, a bad header, or a
+# source that silently contributed nothing. This re-derives everything from the artifacts:
+#   1. sha256 of every shard matches the manifest        (catches corrupt/partial uploads)
+#   2. header n_tokens == byte length, no trailing bytes (catches torn writes)
+#   3. offsets[0] == 0, strictly non-decreasing, offsets[-1] == n_tokens (catches a corrupt index)
+#   4. every id < vocab                                   (an out-of-range id is a train-time crash)
+#   5. total tokens recounted from the bytes, not from counters
+#   6. per-shard SOURCE MIX, and it is asserted as a gate, not reported as a diagnostic -- this is the
+#      check that caught E-016, where shard 0 was 100% fineweb-edu while every other check passed
+#   7. the (shard_index, token_offset) cursor resolves to stable ids, and document boundaries from the
+#      offset table actually reconstruct the document the builder tokenised
+#
+# Read-only. CPU session. No credentials.
+
+import argparse
+import array
+import hashlib
+import json
+import os
+import random
+import struct
+
+
+def sha256_file(path, chunk=8 * 1024**2):
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        for b in iter(lambda: f.read(chunk), b""):
+            h.update(b)
+    return h.hexdigest()
+
+
+def read_shard(path):
+    with open(path, "rb") as f:
+        head = f.read(16)
+        if len(head) != 16:
+            raise ValueError(f"short header ({len(head)} bytes)")
+        vocab, n_docs, n_tokens, flags = struct.unpack("<IIII", head)
+        offs = array.array("I")
+        offs.fromfile(f, n_docs + 1)
+        ids = array.array("H")
+        ids.fromfile(f, n_tokens)
+        trailing = len(f.read())
+    return {"vocab": vocab, "n_docs": n_docs, "n_tokens": n_tokens, "flags": flags,
+            "offsets": offs, "ids": ids, "trailing": trailing}
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--root", default="/kaggle/working/mixroot")
+    ap.add_argument("--sample-decodes", type=int, default=6)
+    ap.add_argument("--cursor-probes", type=int, default=40)
+    ap.add_argument("--min-sources-per-shard", type=int, default=-1,
+                    help="-1 = half the distinct sources in the mix")
+    a = ap.parse_args()
+
+    mpath = os.path.join(a.root, "manifest.json")
+    if not os.path.exists(mpath):
+        # A missing manifest must not become an empty-but-consistent PASS (E-017's silent-success shape)
+        print("VERIFY_JSON_BEGIN")
+        print(json.dumps({"PASS": False, "fatal": f"no manifest at {mpath} -- nothing was built"}))
+        print("VERIFY_JSON_END")
+        raise SystemExit(5)
+    man = json.load(open(mpath))
+    if not man.get("shards"):
+        print("VERIFY_JSON_BEGIN")
+        print(json.dumps({"PASS": False, "fatal": "manifest lists zero shards"}))
+        print("VERIFY_JSON_END")
+        raise SystemExit(5)
+    mix_vocab = man["tokenizer"]["vocab_size"]
+    out = {"shards_expected": len(man["shards"]), "sha_ok": 0, "sha_bad": [],
+           "header_ok": 0, "header_bad": [], "offset_bad": [], "out_of_range_ids": 0,
+           "total_tokens_recounted": 0, "total_docs_recounted": 0, "trailing_bytes": 0}
+
+    tok = None
+    try:
+        from tokenizers import Tokenizer
+        from huggingface_hub import hf_hub_download
+        tok = Tokenizer.from_file(hf_hub_download("HuggingFaceTB/SmolLM2-135M", "tokenizer.json"))
+        tok.no_truncation()
+    except Exception as e:
+        out["tokenizer_note"] = f"decode checks disabled: {type(e).__name__}: {e}"[:160]
+
+    per_shard_sources = []
+    sizes = []
+    for rec in man["shards"]:
+        p = os.path.join(a.root, rec["file"])
+        if not os.path.exists(p):
+            out["header_bad"].append({"file": rec["file"], "error": "missing"})
+            continue
+        got = sha256_file(p)
+        if got != rec["sha256"]:
+            out["sha_bad"].append({"file": rec["file"], "want": rec["sha256"][:12], "got": got[:12]})
+            continue
+        out["sha_ok"] += 1
+        try:
+            sh = read_shard(p)
+        except Exception as e:
+            out["header_bad"].append({"file": rec["file"], "error": f"{type(e).__name__}: {e}"[:160]})
+            continue
+        sizes.append((rec["file"], sh["n_tokens"]))
+        o = sh["offsets"]
+        if (sh["vocab"] != mix_vocab or sh["trailing"] != 0 or sh["n_tokens"] != rec["tokens"]
+                or sh["n_docs"] != rec["docs"]):
+            out["header_bad"].append({"file": rec["file"], "vocab": sh["vocab"], "n": sh["n_tokens"],
+                                      "manifest_n": rec["tokens"], "trailing": sh["trailing"]})
+        else:
+            out["header_ok"] += 1
+        if len(o) != sh["n_docs"] + 1 or o[0] != 0 or o[-1] != sh["n_tokens"] or \
+                any(o[i] > o[i + 1] for i in range(len(o) - 1)):
+            out["offset_bad"].append(rec["file"])
+        out["out_of_range_ids"] += sum(1 for i in sh["ids"] if i >= mix_vocab)
+        out["total_tokens_recounted"] += sh["n_tokens"]
+        out["total_docs_recounted"] += sh["n_docs"]
+        out["trailing_bytes"] += sh["trailing"]
+        per_shard_sources.append((rec["file"], sorted(set(rec["sources"]))))
+
+    distinct = sorted({s for _f, sl in per_shard_sources for s in sl})
+    out["distinct_sources_in_mix"] = len(distinct)
+    out["sources_per_shard"] = {"min": min((len(s) for _f, s in per_shard_sources), default=0),
+                                "max": max((len(s) for _f, s in per_shard_sources), default=0),
+                                "median": sorted(len(s) for _f, s in per_shard_sources)[
+                                    len(per_shard_sources) // 2] if per_shard_sources else 0}
+    floor = a.min_sources_per_shard if a.min_sources_per_shard >= 0 else max(2, len(distinct) // 2)
+    out["min_sources_per_shard_required"] = floor
+    out["shards_below_source_floor"] = [f for f, sl in per_shard_sources if len(sl) < floor]
+
+    # tokens per source, recomputed from the shard headers rather than trusted from the manifest
+    src_tokens = {}
+    for rec in man["shards"]:
+        per = rec["tokens"] / max(1, len(rec["sources"]))
+        for s in rec["sources"]:
+            src_tokens[s] = src_tokens.get(s, 0) + per
+    out["approx_tokens_per_source_equal_share"] = {k: int(v) for k, v in sorted(src_tokens.items())}
+    out["note_per_source"] = ("equal split of a shard's tokens across its sources is an approximation; "
+                             "exact per-source totals live in the stage manifest")
+
+    if tok is not None and sizes:
+        rng = random.Random(11)
+        decs = []
+        for _ in range(a.sample_decodes):
+            fname, n = rng.choice(sizes)
+            sh = read_shard(os.path.join(a.root, fname))
+            i = rng.randrange(0, sh["n_docs"])
+            d = sh["ids"][sh["offsets"][i]:sh["offsets"][i + 1]]
+            decs.append({"shard": fname, "doc": i, "doc_tokens": len(d),
+                         "text": tok.decode(list(d))[:200]})
+        out["decoded_documents"] = decs
+
+        probes = mism = boundary_ok = 0
+        total = sum(n for _f, n in sizes)
+        for _ in range(a.cursor_probes):
+            g = rng.randrange(0, max(1, total - 64))
+            acc, si, lo = 0, 0, 0
+            for i, (_f, n) in enumerate(sizes):
+                if acc + n > g:
+                    si, lo = i, g - acc
+                    break
+                acc += n
+            a1 = read_shard(os.path.join(a.root, sizes[si][0]))
+            a2 = read_shard(os.path.join(a.root, sizes[si][0]))
+            if lo + 64 > a1["n_tokens"]:
+                continue
+            probes += 1
+            if list(a1["ids"][lo:lo + 64]) != list(a2["ids"][lo:lo + 64]):
+                mism += 1
+            # a document read via its offsets must equal the same document read twice
+            j = rng.randrange(0, a1["n_docs"])
+            d1 = list(a1["ids"][a1["offsets"][j]:a1["offsets"][j + 1]])
+            d2 = list(a2["ids"][a2["offsets"][j]:a2["offsets"][j + 1]])
+            if d1 == d2 and 0 < len(d1) <= 16_384:
+                boundary_ok += 1
+        out["cursor_probes"] = {"done": probes, "mismatches": mism, "boundaries_ok": boundary_ok}
+
+    out["manifest_claims_total_tokens"] = man["total_tokens"]
+    out["recount_matches_manifest"] = man["total_tokens"] == out["total_tokens_recounted"]
+    out["PASS"] = (not out["sha_bad"] and not out["header_bad"] and not out["offset_bad"]
+                   and not out["shards_below_source_floor"] and out["recount_matches_manifest"]
+                   and out["out_of_range_ids"] == 0 and out["trailing_bytes"] == 0
+                   and out.get("cursor_probes", {}).get("mismatches", 1) == 0)
+    print("VERIFY_JSON_BEGIN")
+    print(json.dumps(out, indent=1, default=str))
+    print("VERIFY_JSON_END")
+
+
+if __name__ == "__main__":
+    main()