Phase 0: add ingest-throughput probe p0d

probes/p0d_ingest_throughput.py

@@ -0,0 +1,197 @@
+# Phase 0 CPU probe D -- measures the three numbers that decide whether Phase 2 is schedulable.
+# Zero GPU quota (§3.7: data work is a CPU job by rule).
+#
+#   (a) sustained Hub download MB/s from inside a Kaggle session,
+#   (b) `datasets` 5.0 streaming rows/s on a real multi-GB corpus,
+#   (c) BPE tokenization tokens/s on 1 core and on 2 and 4 cores.
+#
+# Why (c) matters more than it looks: the mix is ~1B tokens and the CPU shape has 4 vCPUs and no
+# swap. If tokenizing runs at a few thousand tok/s/core, Phase 2 costs days of wall clock before any
+# training starts and the shard layout has to be designed around that. Measure, don't hope.
+#
+# Readable canonical copy: Cion-lab/ounce100m-code, probes/p0d_ingest_throughput.py
+# Contains NO credentials.
+
+import json
+import multiprocessing as mp
+import os
+import socket
+import subprocess
+import sys
+import time
+
+R = {"host": socket.gethostname(), "python": sys.version.split()[0]}
+BUDGET_S = 1500  # self-imposed; this must never turn into a data-engineering run
+CORPUS_MB = 24
+TOK_CHUNK = 6 * 1024**2  # bytes of text used for each tokenize timing
+START = time.monotonic()
+
+
+def over():
+    return time.monotonic() - START > BUDGET_S
+
+
+# --------------------------------------------------------------------------- tokenizer
+def load_tokenizer():
+    from huggingface_hub import hf_hub_download
+
+    path = hf_hub_download("gpt2", "tokenizer.json")
+    os.environ["TOK_JSON"] = path
+    return path
+
+
+_g = {}
+
+
+def _init():
+    from tokenizers import PreTrainedTokenizerFast
+
+    _g["tk"] = PreTrainedTokenizerFast.from_file(os.environ["TOK_JSON"])
+
+
+def _encode_slice(text):
+    return len(_g["tk"].encode(text, add_special_tokens=False).ids)
+
+
+def tokenize(text, workers):
+    """tokens/s via `tokenizers`, sharded across `workers` processes.
+
+    Slices go through memory, not /kaggle/working: the working volume is 19.5 GB and the point of
+    this probe is not to fill it.
+    """
+    step = max(1, len(text) // workers)
+    slices = [text[i * step:(i + 1) * step] for i in range(workers)]
+    slices[-1] = text[(workers - 1) * step:]
+    if workers == 1:
+        _init()
+        return sum(_encode_slice(s) for s in slices)
+    # fork, not spawn: this script is exec()'d inside a kernel, so a spawned child would try to
+    # re-import a __main__ that is not the probe and would recurse instead of tokenizing.
+    ctx = mp.get_context("fork")
+    with ctx.Pool(workers, initializer=_init) as pool:
+        return int(sum(pool.map(_encode_slice, slices)))
+
+
+# --------------------------------------------------------------------------- corpus
+def grab_corpus(target_bytes):
+    """Stream real English off the Hub into memory, returning text plus streaming stats."""
+    import datasets
+
+    keep, total, t0, rows = [], 0, time.monotonic(), 0
+    ds = datasets.load_dataset("wikimedia/wikipedia", "20231101.en", split="train",
+                              streaming=True)
+    for row in ds:
+        txt = row["text"]
+        keep.append(txt)
+        total += len(txt)
+        rows += 1
+        if total >= target_bytes or time.monotonic() - t0 > 420:
+            break
+    el = time.monotonic() - t0
+    stats = {"rows": rows, "MB_text": round(total / 1024**2, 1), "seconds": round(el, 1),
+             "rows_per_s": round(rows / el, 1) if el else None,
+             "KB_text_per_s": round(total / el / 1024, 0) if el else None,
+             "hit_row_cap": total >= target_bytes}
+    return "\n".join(keep), stats
+
+
+# --------------------------------------------------------------------------- raw HTTP
+def list_parquet(repo_id, limit=8):
+    from huggingface_hub import HfApi
+
+    out = []
+    for f in HfApi().list_files_info(repo_id=repo_id, repo_type="dataset", revision="main"):
+        if f.path.endswith(".parquet"):
+            out.append({"path": f.path, "size_MB": round((f.size or 0) / 1024**2, 1)})
+            if len(out) >= limit:
+                break
+    return out
+
+
+def http_throughput(url, max_bytes=300 * 1024**2, chunk=4 * 1024**2):
+    import urllib.request
+
+    req = urllib.request.Request(url, headers={"User-Agent": "ounce100m-probe"})
+    got, t0, ttfb = 0, time.monotonic(), None
+    with urllib.request.urlopen(req, timeout=120) as r:
+        while got < max_bytes:
+            b = r.read(chunk)
+            if not b:
+                break
+            if ttfb is None:
+                ttfb = time.monotonic() - t0
+            got += len(b)
+    el = time.monotonic() - t0
+    return {"bytes_read_MB": round(got / 1024**2, 1), "seconds": round(el, 1),
+            "MB_per_s": round(got / el / 1024**2, 1) if el else None,
+            "ttfb_s": round(ttfb, 2) if ttfb else None}
+
+
+# --------------------------------------------------------------------------- run
+try:
+    R["tokenizer"] = {"path": load_tokenizer(), "id": "gpt2 tokenizer.json (stand-in)"}
+except Exception as e:
+    R["tokenizer"] = {"error": f"{type(e).__name__}: {e}"[:220]}
+
+try:
+    R["listing_wikipedia"] = list_parquet("wikimedia/wikipedia")[:3]
+except Exception as e:
+    R["listing_wikipedia"] = {"error": f"{type(e).__name__}: {e}"[:220]}
+
+try:
+    R["listing_fineweb_edu"] = list_parquet("HuggingFaceFW/fineweb-edu")[:3]
+except Exception as e:
+    R["listing_fineweb_edu"] = {"error": f"{type(e).__name__}: {e}"[:220]}
+
+text = ""
+if not over():
+    try:
+        text, R["streaming"] = grab_corpus(CORPUS_MB * 1024**2)
+        R["corpus_MB"] = round(len(text) / 1024**2, 1)
+    except Exception as e:
+        R["streaming"] = {"error": f"{type(e).__name__}: {e}"[:250]}
+
+tok_sample = text[:TOK_CHUNK]
+if tok_sample and "error" not in R["tokenizer"]:
+    for w in (1, 2, 4):
+        if over():
+            R["tokenize_stopped_at"] = w
+            break
+        try:
+            t0 = time.monotonic()
+            n = tokenize(tok_sample, w)
+            el = time.monotonic() - t0
+            tps = n / el
+            R[f"tokenize_{w}core"] = {
+                "tokens": n, "seconds": round(el, 1), "tokens_per_s": round(tps),
+                "hours_to_tokenize_1B_tokens": round(1e9 / tps / 3600, 1),
+                "chars_per_token": round(len(tok_sample) / n, 2) if n else None,
+                "GB_per_hour_at_1B_target": round(tps * 3600 * 4 / 1024**3, 2),
+            }
+        except Exception as e:
+            R[f"tokenize_{w}core"] = {"error": f"{type(e).__name__}: {e}"[:220]}
+        _g.clear()
+
+# raw HTTP last: it is the noisiest measurement and the least likely to be cache-warm
+for repo in ("wikimedia/wikipedia", "HuggingFaceFW/fineweb-edu"):
+    if over():
+        break
+    key = f"http_{repo}"
+    try:
+        big = [f for f in list_parquet(repo, limit=10) if f["size_MB"] > 150]
+        if not big:
+            R[key] = {"skipped": "no parquet >150 MB in the first 10"}
+            continue
+        url = f"https://huggingface.co/datasets/{repo}/resolve/main/{big[0]['path']}"
+        R[key] = {"file_MB": big[0]["size_MB"]}
+        R[key].update(http_throughput(url))
+    except Exception as e:
+        R[key] = {"error": f"{type(e).__name__}: {e}"[:220]}
+
+R["_budget_exhausted"] = over()
+R["_seconds_used"] = round(time.monotonic() - START, 1)
+R["disk_after"] = subprocess.run("df -h /kaggle/working /dev/shm", shell=True,
+                                 capture_output=True, text=True).stdout.strip()
+print("PROBE_JSON_BEGIN")
+print(json.dumps(R, indent=1, default=str))
+print("PROBE_JSON_END")