spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expH_capture_cost_frontier/implementation/benchmark_cold.py5#  Purpose   : Run #2 — cold-cache storage-format throughput (purge per repeat)6#  Author    : Simon-Pierre Boucher7#  Contact   : contact@spboucher.ai8#  Website   : https://modelmap.io9#  Created   : 2026-08-1210#  Modified  : 2026-08-1211#  Platform  : macOS / Apple Silicon (arm64)12#  License   : All rights reserved (research code)13# =============================================================================14"""expH run #2 (see hypothesis.md, registered before this run).1516Storage-only cold-cache variant: `sudo purge` empties the unified buffer17cache before EVERY timed read repetition. Designed to run standalone on a18MacLustr node (numpy + zarr + safetensors only, no torch/mlx). Results JSON19is self-contained (embeds a local hardware manifest) and is collected back20into results/ by the driver on the laptop.2122Usage: python3 benchmark_cold.py --workdir /path --out results.json \23       [--purge-cmd "sudo -n purge"]24"""2526from __future__ import annotations2728import argparse29import json30import platform31import shutil32import subprocess33import time34from pathlib import Path3536import numpy as np3738REPEATS = 339ROWS, DIM = 200_000, 4096          # ~1.6 GB fp16 per format40WRITE_CHUNK = 4_09641BATCH = 4_09642N_BATCHES = 2443SEED = 0444546def sysctl(key: str) -> str:47    try:48        return subprocess.run(["sysctl", "-n", key], capture_output=True,49                              text=True, check=True).stdout.strip()50    except subprocess.CalledProcessError:51        return ""525354def local_manifest() -> dict:55    return {56        "author": "Simon-Pierre Boucher",57        "chip": sysctl("machdep.cpu.brand_string"),58        "cores": int(sysctl("hw.ncpu") or 0),59        "unified_gb": round(int(sysctl("hw.memsize") or 0) / 2**30, 1),60        "os": platform.mac_ver()[0],61        "python": platform.python_version(),62        "numpy": np.__version__,63        "host": platform.node(),64    }656667def main() -> int:68    ap = argparse.ArgumentParser()69    ap.add_argument("--workdir", required=True)70    ap.add_argument("--out", required=True)71    ap.add_argument("--purge-cmd", default="")72    args = ap.parse_args()7374    workdir = Path(args.workdir)75    workdir.mkdir(parents=True, exist_ok=True)7677    def purge():78        if args.purge_cmd:79            subprocess.run(args.purge_cmd, shell=True, check=True,80                           capture_output=True)8182    def timed_cold(fn):83        out = []84        for _ in range(REPEATS):85            purge()86            t0 = time.perf_counter()87            fn()88            out.append(time.perf_counter() - t0)89        return out9091    rng = np.random.default_rng(SEED)92    data = rng.standard_normal((WRITE_CHUNK, DIM)).astype(np.float16)93    batches = [rng.integers(0, ROWS, BATCH) for _ in range(N_BATCHES)]94    total_bytes = ROWS * DIM * 295    batch_bytes = BATCH * DIM * 2 * N_BATCHES96    results = []9798    def record(fmt, op, times, nbytes):99        results.append({"format": fmt, "op": op, "bytes": nbytes, "seconds": times,100                        "gb_per_s_mean": nbytes / 2**30 / np.mean(times),101                        "cache": "cold" if op != "write" else "warm"})102        print(f"  {fmt:18s} {op:12s} {nbytes/2**30/np.mean(times):8.2f} GB/s", flush=True)103104    # ---- raw np.memmap105    p = workdir / "acts.raw"106    def write_raw():107        m = np.memmap(p, dtype=np.float16, mode="w+", shape=(ROWS, DIM))108        for i in range(0, ROWS, WRITE_CHUNK):109            end = min(i + WRITE_CHUNK, ROWS)110            m[i:end] = data[: end - i]111        m.flush(); del m112    t0 = time.perf_counter(); write_raw()113    record("raw-mmap", "write", [time.perf_counter() - t0], total_bytes)114    def seq_raw():115        m = np.memmap(p, dtype=np.float16, mode="r", shape=(ROWS, DIM))116        float(np.asarray(m).sum(dtype=np.float32)); del m117    def rnd_raw():118        m = np.memmap(p, dtype=np.float16, mode="r", shape=(ROWS, DIM))119        for b in batches:120            m[b].sum(dtype=np.float32)121        del m122    record("raw-mmap", "seq-scan", timed_cold(seq_raw), total_bytes)123    record("raw-mmap", "random-batch", timed_cold(rnd_raw), batch_bytes)124125    # ---- safetensors126    from safetensors import safe_open127    from safetensors.numpy import save_file128    ps = workdir / "acts.safetensors"129    full = np.memmap(p, dtype=np.float16, mode="r", shape=(ROWS, DIM))130    t0 = time.perf_counter(); save_file({"acts": np.asarray(full)}, str(ps))131    record("safetensors", "write", [time.perf_counter() - t0], total_bytes)132    del full133    def seq_st():134        f = safe_open(str(ps), framework="np")135        float(f.get_tensor("acts").sum(dtype=np.float32))136    def rnd_st():137        f = safe_open(str(ps), framework="np")138        t = f.get_tensor("acts")139        for b in batches:140            t[b].sum(dtype=np.float32)141    record("safetensors", "seq-scan", timed_cold(seq_st), total_bytes)142    record("safetensors", "random-batch", timed_cold(rnd_st), batch_bytes)143144    # ---- zarr variants145    import zarr146    for codec, name in ((None, "zarr-uncompressed"), ("default", "zarr-zstd")):147        pz = workdir / f"acts_{name}.zarr"148        kwargs = {} if codec == "default" else {"compressors": None}149        if pz.exists():150            shutil.rmtree(pz)151        t0 = time.perf_counter()152        z = zarr.create_array(store=str(pz), shape=(ROWS, DIM),153                              chunks=(WRITE_CHUNK, DIM), dtype=np.float16, **kwargs)154        for i in range(0, ROWS, WRITE_CHUNK):155            end = min(i + WRITE_CHUNK, ROWS)156            z[i:end] = data[: end - i]157        record(name, "write", [time.perf_counter() - t0], total_bytes)158        def seq_z(pz=pz):159            zz = zarr.open_array(store=str(pz), mode="r")160            float(zz[:].sum(dtype=np.float32))161        def rnd_z(pz=pz):162            zz = zarr.open_array(store=str(pz), mode="r")163            for b in batches:164                zz[np.sort(b)].sum(dtype=np.float32)165        record(name, "seq-scan", timed_cold(seq_z), total_bytes)166        record(name, "random-batch", timed_cold(rnd_z), batch_bytes)167168    doc = {169        "experiment": "expH_capture_cost_frontier",170        "run": 2,171        "scope": "storage formats, COLD cache (purge per repeat), second hardware",172        "config": {"rows": ROWS, "dim": DIM, "write_chunk": WRITE_CHUNK,173                   "batch": BATCH, "n_batches": N_BATCHES, "repeats": REPEATS,174                   "purge_cmd": args.purge_cmd or "(none — warm!)", "seed": SEED},175        "manifest": local_manifest(),176        "storage": results,177    }178    Path(args.out).write_text(json.dumps(doc, indent=2) + "\n")179    print(f"results -> {args.out}")180    return 0181182183if __name__ == "__main__":184    raise SystemExit(main())185