hub sync for staged sources + cold-resume rehearsal script

build/build_mix.py

@@ -219,12 +219,34 @@ class State:
 
 # --------------------------------------------------------------------- stage
 def stage(args, st, tokenizer):
+    remote = HS.remote_stage_files(args.hub_repo) if (args.hub_repo and not args.merge_only) else set()
+    if args.hub_repo:
+        print(f"hub: {len(remote)} staged files already on {args.hub_repo}", flush=True)
     for src in SOURCES:
         key = src_key(src)
         rec = st.d["staged"].get(key, {})
         if rec.get("complete"):
             print(f"stage: skip {key} (complete)", flush=True)
             continue
+        if args.hub_repo:
+            # A source whose shards AND record are already on the Hub is done; pulling the record is
+            # cheaper than re-staging and keeps the merge's proportional weights honest.
+            have = [p for p in remote if p.startswith(f"stage/{key}/") and p.endswith(".bin")]
+            if have and f"stage/{key}/record.json" in remote:
+                try:
+                    HS.download_missing(st.root, [key], args.hub_repo, remote=remote)
+                    rec_json = json.load(open(os.path.join(st.stage_dir, key, "record.json")))
+                    st.d["staged"][key] = {"complete": True, "shards": rec_json["shards"],
+                                           "tokens": rec_json["tokens"],
+                                           "docs_seen": rec_json["docs_seen"],
+                                           "dropped": rec_json["dropped"],
+                                           "target": rec_json["tokens"], "restored_from_hub": True}
+                    st.commit([], [])
+                    print(f"stage: restored {key} from Hub ({rec_json['tokens']:,} tok), skipping",
+                          flush=True)
+                    continue
+                except Exception as e:
+                    print(f"!! restore failed for {key}: {type(e).__name__}: {str(e)[:140]}", flush=True)
         want = min(src["tokens"], args.cap_tokens_per_source or src["tokens"])
         if args.budget_tokens and st.staged_tokens_total() >= args.budget_tokens:
             print("stage: overall budget reached", flush=True)
@@ -328,6 +350,19 @@ def stage(args, st, tokenizer):
                                "docs_seen": docs_seen, "dropped": dict(drops),
                                "target": want, "seconds": round(time.monotonic() - t0, 1)}
         st.commit(new_h, new_k)
+        # /kaggle/working is destroyed when the session ends, so a source that took 40 minutes to stage
+        # must not exist only there. Published per source, because that is the unit a resume can skip.
+        if complete and args.hub_repo:
+            with open(os.path.join(st.stage_dir, key, "record.json"), "w") as f:
+                json.dump({"key": key, "src": {k: v for k, v in src.items()},
+                           "shards": shards, "tokens": got, "docs_seen": docs_seen,
+                           "dropped": dict(drops)}, f, indent=1)
+            try:
+                n = HS.upload_source(st.root, key, args.hub_repo)
+                print(f"  published {key} -> {args.hub_repo}/stage/{key} ({n} files)", flush=True)
+            except Exception as e:
+                print(f"!! publish FAILED for {key}: {type(e).__name__}: {str(e)[:160]} -- "
+                      f"this source will have to be re-staged if the session dies", flush=True)
         el = time.monotonic() - t0
         print(f"stage: {key} {got:,} tok of {want:,} from {docs_seen:,} docs; "
               f"dup={drops['exact_dup']} keydup={drops['key_dup']} lang={drops['non_english']} "
@@ -433,6 +468,9 @@ def main():
     ap.add_argument("--max-chars", type=int, default=100_000)
     ap.add_argument("--sec-per-source", type=int, default=0)
     ap.add_argument("--only", default="", help="comma list of config names to stage")
+    ap.add_argument("--hub-repo", default="Cion-lab/ounce100m-mix-stage",
+                    help="staged sources are published here so a session kill cannot lose them; "
+                         "empty string disables sync")
     ap.add_argument("--stage-only", action="store_true")
     ap.add_argument("--merge-only", action="store_true")
     args = ap.parse_args()
@@ -451,7 +489,41 @@ def main():
     print(f"tokenizer vocab={tk.get_vocab_size()} "
           f"sha256={hashlib.sha256(open(tok_path, 'rb').read()).hexdigest()[:12]}", flush=True)
 
+    global HS
+    HS = None
+    if args.hub_repo:
+        try:
+            import hubsync as HS_mod
+            HS = HS_mod
+            HS.ensure_repo(args.hub_repo)
+        except Exception as e:
+            print(f"!! hub sync unavailable: {type(e).__name__}: {str(e)[:160]}; "
+                  f"build proceeds but is NOT interruption-safe", flush=True)
+            HS = None
+
     st = State(args.root)
+    if args.hub_repo and not st.d["staged"]:
+        # The state file lives on /kaggle/working too, so in a fresh session it is empty even when the
+        # Hub holds a complete stage. Rebuild the stage bookkeeping from the Hub's record.json files --
+        # without this, --merge-only on a new instance would find nothing and quietly emit an empty mix.
+        try:
+            remote = HS.remote_stage_files(args.hub_repo)
+            keys = sorted({p.split("/")[1] for p in remote if p.startswith("stage/") and "/" in p})
+            pulled = HS.download_missing(args.root, keys, args.hub_repo, remote=remote)
+            for k in keys:
+                rp = os.path.join(args.root, "stage", k, "record.json")
+                if not os.path.exists(rp):
+                    continue
+                j = json.load(open(rp))
+                st.d["staged"][k] = {"complete": True, "shards": j["shards"], "tokens": j["tokens"],
+                                     "docs_seen": j["docs_seen"], "dropped": j["dropped"],
+                                     "target": j["tokens"], "restored_from_hub": True}
+            st.commit([], [])
+            print(f"hub: restored {len(st.d['staged'])} staged source(s) "
+                  f"({st.staged_tokens_total():,} tokens)", flush=True)
+        except Exception as e:
+            print(f"!! hub restore failed: {type(e).__name__}: {str(e)[:160]}", flush=True)
+
     t0 = time.monotonic()
     if not args.merge_only:
         stage(args, st, tk)

build/hubsync.py

@@ -0,0 +1,105 @@
+# Hub sync for the mix build: the instance is scratch, so nothing that took an hour to produce may live
+# only there (§3.13 "never let the instance be the only copy", applied to a CPU job rather than a
+# checkpoint).
+#
+# Granularity is one *source*: staging a source takes 5-40 minutes, so publishing its shards the moment
+# it finishes bounds the loss from any interruption to the one source in flight, and a resume can skip
+# every source already on the Hub without re-reading it.
+#
+# The stage repo is private-in-practice but public by project rule; it holds no credentials and no
+# benchmark material -- only token ids from the approved mix.
+
+import json
+import os
+
+
+def repo_id_from_env(default=None):
+    return os.environ.get("OUNCE100M_STAGE_REPO") or default
+
+
+def api():
+    import os
+
+    import ounce100m_credentials
+
+    ounce100m_credentials.install()
+    from huggingface_hub import HfApi
+
+    return HfApi(token=os.environ["HF_TOKEN"])
+
+
+def ensure_repo(repo):
+    from huggingface_hub import create_repo
+
+    # hub 1.11 (the image's version) takes the owner inside repo_id and has no `namespace` kwarg (E-013)
+    create_repo(repo_id=repo, repo_type="dataset", private=False, exist_ok=True)
+    return repo
+
+
+def remote_stage_files(repo):
+    """Set of relative paths under stage/ that already exist on the Hub."""
+    import urllib.request
+
+    url = f"https://huggingface.co/api/datasets/{repo}/tree/main/stage?recursive=true"
+    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
+
+
+def upload_source(root, key, repo):
+    """Publish stage/<key>/* in one commit. Returns the file names pushed."""
+    a = api()
+    local = os.path.join(root, "stage", key)
+    files = sorted(f for f in os.listdir(local) if f.endswith((".bin", ".json")))
+    if not files:
+        return []
+    a.upload_folder(folder_path=local, path_in_repo=f"stage/{key}", repo_id=repo,
+                    repo_type="dataset", commit_message=f"stage {key}: {len(files)} shard(s)",
+                    allow_patterns=["*.bin", "*.json"])
+    return files
+
+
+def download_missing(root, keys_wanted, repo, remote=None):
+    """Fetch stage/<key>/*.bin for every wanted key not already on disk. Returns {key: [filenames]}."""
+    from huggingface_hub import hf_hub_download
+
+    remote = remote if remote is not None else remote_stage_files(repo)
+    got = {}
+    for key in keys_wanted:
+        want = sorted(p for p in remote if p.startswith(f"stage/{key}/") and p.endswith(".bin"))
+        if not want:
+            got[key] = []
+            continue
+        local_dir = os.path.join(root, "stage", key)
+        os.makedirs(local_dir, exist_ok=True)
+        have = {f"stage/{key}/{f}" for f in os.listdir(local_dir)}
+        for rp in want:
+            if rp in have:
+                continue
+            hf_hub_download(repo_id=repo, repo_type="dataset", filename=rp,
+                            local_dir=root, force_download=False)
+        got[key] = sorted(os.listdir(local_dir))
+    return got
+
+
+def status(repo):
+    import urllib.request
+
+    url = f"https://huggingface.co/api/datasets/{repo}"
+    try:
+        with urllib.request.urlopen(url, timeout=60) as r:
+            d = json.loads(r.read().decode())
+        return {"exists": True, "gated": d.get("gated"), "private": d.get("private"),
+                "lastModified": d.get("lastModified"), "id": d.get("id")}
+    except Exception as e:
+        return {"exists": False, "error": f"{type(e).__name__}: {str(e)[:120]}"}

build/smoke_hub_resume.py

@@ -0,0 +1,100 @@
+# 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()