p0d v2: version-free tree API + subprocess rayon bench (fixes 3 API-skew faults)

probes/p0d_ingest_throughput.py

@@ -1,111 +1,52 @@
-# 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).
+# Phase 0 CPU probe D v2 -- ingest throughput for sizing Phase 2. Quota-free (CPU) 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.
+# v1 measured one thing and failed at three others, all for the same reason: it used APIs that exist
+# in this workspace's huggingface_hub 1.32 / tokenizers 0.22 and do not exist in the Kaggle image's
+# (hub 1.11 has no HfApi.list_files_info; tokenizers 0.22 does not export PreTrainedTokenizerFast at
+# top level), and it used multiprocessing.Pool, which cannot work in a kernel that was exec()'d
+# because the worker tries to pickle functions by qualified name out of a __main__ that is not an
+# importable module. See ERRORS.md E-007/E-009.
 #
-# 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.
+# v2 therefore: plain HTTP against the public /api/.../tree endpoint for file listings, the generic
+# tokenizers.Tokenizer for encoding, and re-exec of THIS file as a subprocess per thread setting so
+# rayon's pool size is chosen before the library is imported. No credentials in this file.
 #
-# Readable canonical copy: Cion-lab/ounce100m-code, probes/p0d_ingest_throughput.py
-# Contains NO credentials.
+# Measures: (a) sustained Hub download MB/s, (b) datasets streaming rows/s + KB text/s,
+#           (c) BPE tokens/s at 1, 2 and 4 cores -- which is what says whether Phase 2 is hours or days.
 
 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
+BUDGET_S = 1500
 START = time.monotonic()
+CORPUS_MB = 24
+TOK_MB = 6
+GPT2_CHARS_PER_TOKEN = 3.6
 
 
 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)))
+def http_json(url):
+    import urllib.request
 
+    req = urllib.request.Request(url, headers={"User-Agent": "ounce100m-probe"})
+    with urllib.request.urlopen(req, timeout=90) as r:
+        return json.loads(r.read().decode())
 
-# --------------------------------------------------------------------------- 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 tree_parquet(repo_id, repo_type="datasets", limit=10):
+    """Largest files in a repo, via the public tree API -- no client library, no version skew."""
+    url = f"https://huggingface.co/api/{repo_type}/{repo_id}/tree/main?recursive=true"
+    out = [e for e in http_json(url)
+           if e.get("type") == "file" and str(e.get("path", "")).endswith(".parquet")]
+    out.sort(key=lambda e: -(e.get("size") or 0))
+    return [{"path": e["path"], "size_MB": round((e.get("size") or 0) / 1024**2, 1)} for e in out[:limit]]
 
 
 def http_throughput(url, max_bytes=300 * 1024**2, chunk=4 * 1024**2):
@@ -127,62 +68,105 @@ def http_throughput(url, max_bytes=300 * 1024**2, chunk=4 * 1024**2):
             "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]}
+def grab_corpus(target_bytes):
+    import datasets
 
-try:
-    R["listing_wikipedia"] = list_parquet("wikimedia/wikipedia")[:3]
-except Exception as e:
-    R["listing_wikipedia"] = {"error": f"{type(e).__name__}: {e}"[:220]}
+    keep, total, t0, rows = [], 0, time.monotonic(), 0
+    for row in datasets.load_dataset("wikimedia/wikipedia", "20231101.en", split="train",
+                                     streaming=True):
+        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
+    return "\n".join(keep), {"rows": rows, "MB_text": round(total / 1024**2, 1),
+                             "seconds": round(el, 1), "rows_per_s": round(rows / el, 1),
+                             "KB_text_per_s": round(total / el / 1024, 0)}
+
+
+# ---------------------------------------------------------------- worker mode: one thread budget
+if len(sys.argv) > 1 and sys.argv[1] == "tokworker":
+    from tokenizers import Tokenizer
+
+    tk = Tokenizer.from_file(os.environ["TOK_JSON"])
+    tk.no_truncate()
+    with open("/dev/shm/" + sys.argv[2]) as f:
+        chunks = json.load(f)
+    t = time.monotonic()
+    enc = tk.encode_batch(chunks)          # rayon-parallel across cores, released GIL
+    el = time.monotonic() - t
+    n = sum(len(e.ids) for e in enc)
+    print("TOKRESULT " + json.dumps({"tokens": n, "seconds": round(el, 2),
+                                     "rayon": os.environ.get("RAYON_NUM_THREADS")}), flush=True)
+    sys.exit(0)
+
+# ---------------------------------------------------------------- orchestrator
+R = {"host": socket.gethostname()}
+text = ""
 
 try:
-    R["listing_fineweb_edu"] = list_parquet("HuggingFaceFW/fineweb-edu")[:3]
+    from huggingface_hub import hf_hub_download
+
+    os.environ["TOK_JSON"] = hf_hub_download("gpt2", "tokenizer.json")
+    R["tokenizer"] = {"id": "gpt2/tokenizer.json", "path": os.environ["TOK_JSON"]}
 except Exception as e:
-    R["listing_fineweb_edu"] = {"error": f"{type(e).__name__}: {e}"[:220]}
+    R["tokenizer"] = {"error": f"{type(e).__name__}: {e}"[:200]}
 
-text = ""
-if not over():
+for repo in ("wikimedia/wikipedia", "HuggingFaceFW/fineweb-edu"):
     try:
-        text, R["streaming"] = grab_corpus(CORPUS_MB * 1024**2)
-        R["corpus_MB"] = round(len(text) / 1024**2, 1)
+        R[f"tree_{repo}"] = tree_parquet(repo)
     except Exception as e:
-        R["streaming"] = {"error": f"{type(e).__name__}: {e}"[:250]}
+        R[f"tree_{repo}"] = {"error": f"{type(e).__name__}: {e}"[:200]}
 
-tok_sample = text[:TOK_CHUNK]
-if tok_sample and "error" not in R["tokenizer"]:
-    for w in (1, 2, 4):
+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]}
+
+if text and "error" not in R["tokenizer"]:
+    sample = text[:TOK_MB * 1024**2]
+    # Chunks sized for rayon: many small pieces keep all cores fed far better than one blob.
+    step = max(1, len(sample) // 64)
+    chunks = [sample[i:i + step] for i in range(0, len(sample), step)]
+    with open("/dev/shm/p0d_chunks.json", "w") as f:
+        json.dump(chunks, f)
+    # The launcher wrote this exact filename into cwd (/kaggle/working), so re-exec that rather than
+    # sys.argv[0], which points at the kernel's own wrapper script, not at this probe.
+    here = "p0d_ingest_throughput.py"
+    for threads 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),
+        env = {**os.environ, "RAYON_NUM_THREADS": str(threads), "OMP_NUM_THREADS": str(threads)}
+        p = subprocess.run([sys.executable, here, "tokworker"], capture_output=True, text=True,
+                           env=env, timeout=900)
+        line = next((l for l in p.stdout.splitlines() if l.startswith("TOKRESULT")), None)
+        if line:
+            w = json.loads(line[len("TOKRESULT"):])
+            tps = w["tokens"] / w["seconds"]
+            R[f"tokenize_{threads}core"] = {
+                **w, "tokens_per_s": round(tps),
+                "hours_to_tokenize_1B_tokens": round(1e9 / tps / 3600, 2),
+                "GB_text_per_hour": round(tps * 3600 * GPT2_CHARS_PER_TOKEN / 1024**3, 2),
+                "chars_per_token": round(len(sample) / w["tokens"], 2),
             }
-        except Exception as e:
-            R[f"tokenize_{w}core"] = {"error": f"{type(e).__name__}: {e}"[:220]}
-        _g.clear()
+        else:
+            R[f"tokenize_{threads}core"] = {"error": (p.stderr or p.stdout)[-300:]}
+    with open("/dev/shm/p0d_chunks.json", "w") as f:
+        f.write("{}")
 
-# 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}"
+    big = [f for f in R.get(key.replace("http", "tree"), []) if f["size_MB"] > 150]
+    if not big:
+        R[key] = {"skipped": "no parquet >150MB listed"}
+        continue
+    url = f"https://huggingface.co/datasets/{repo}/resolve/main/{big[0]['path']}"
     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: