Speed up raw IO for many-small-read workloads (deep learning)

.gitignore

@@ -103,4 +103,7 @@ venv/
 .hypothesis/
 .ruff_cache/
 .ipynb_checkpoints/
-/.claude/
+/.claude/
+# IO benchmark synthetic data (regenerate via benchmarks/io_dl/generate_data.py)
+benchmarks/io_dl/data/
+benchmarks/io_dl/data_big/

benchmarks/io_dl/ab_test.py

@@ -0,0 +1,101 @@
+"""Clean interleaved A/B with explicit environment control.
+
+Arm INSTALLED: site-packages mne (unpatched release copy)
+Arm TREE:      working tree mne (patched)
+
+Also reports component breakdown inside each process.
+"""
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+TREE = "/Users/bruaristimunha/Projects/libraries/mne_python/mne_python_more_io_speed"
+HERE = Path(__file__).parent
+
+CHILD = r'''
+import gc, json, os, sys, time
+import numpy as np
+which = sys.argv[1]
+# `python -c` puts cwd ('') first on sys.path; remove anything that could
+# resolve to the working tree unless this is the TREE arm
+TREE = "{tree}"
+if which == "tree":
+    if TREE not in sys.path:
+        sys.path.insert(0, TREE)
+else:
+    sys.path[:] = [p for p in sys.path if p not in ("", ".", TREE)]
+    os.environ.pop("PYTHONPATH", None)
+import mne
+mne.set_log_level("ERROR")
+fmt, pl = sys.argv[2], sys.argv[3] == "1"
+READERS = {
+    "edf": lambda p_: mne.io.read_raw_edf("{here}/data/bench.edf", preload=p_),
+    "brainvision": lambda p_: mne.io.read_raw_brainvision("{here}/data/bench.vhdr", preload=p_),
+    "fif": lambda p_: mne.io.read_raw_fif("{here}/data/bench_raw.fif", preload=p_),
+}
+raw = READERS[fmt](pl)
+rng = np.random.default_rng(0)
+starts = rng.integers(0, raw.n_times - 513, size=4000).astype(int)
+stops = starts + 512
+
+def bench(fn, warmup=300):
+    for s in starts[:warmup]:
+        fn(int(s))
+    ts = []
+    gc.disable()
+    for s, e in zip(starts[warmup:], stops[warmup:]):
+        t0 = time.perf_counter_ns()
+        fn(s)
+        ts.append(time.perf_counter_ns() - t0)
+    gc.enable()
+    arr = np.asarray(ts) / 1e3
+    return dict(med=float(np.median(arr)), p10=float(np.percentile(arr, 10)))
+
+out = dict(which=which, file=mne.__file__)
+out["get_data_only"] = bench(lambda s: raw.get_data(start=s, stop=s + 512))
+out["get_data_sum"] = bench(lambda s: raw.get_data(start=s, stop=s + 512).sum())
+mm = np.memmap("{here}/data/bench.bin", dtype="<f4", mode="r", shape=(64, raw.n_times))
+out["memmap_copy"] = bench(lambda s: mm[:, s:s + 512].astype(np.float64).sum())
+print(json.dumps(out))
+'''.replace("{tree}", TREE).replace("{here}", str(HERE))
+
+
+FMT = os.environ.get("AB_FMT", "brainvision")
+PRELOAD = os.environ.get("AB_PRELOAD", "1")
+
+
+def run(which):
+    env = {k: v for k, v in os.environ.items()}
+    out = subprocess.run([sys.executable, "-c", CHILD, which, FMT, PRELOAD],
+                         capture_output=True, text=True, env=env)
+    if out.returncode != 0:
+        raise RuntimeError(out.stderr[-800:])
+    return json.loads(out.stdout.strip().splitlines()[-1])
+
+
+def main():
+    rounds = int(sys.argv[1]) if len(sys.argv) > 1 else 5
+    acc = {"INSTALLED": [], "TREE": []}
+    for i in range(rounds):
+        row = {}
+        for key, which in (("INSTALLED", "installed"), ("TREE", "tree")):
+            r = run(which)
+            acc[key].append(r)
+            row[key] = r
+            src = "TREE" if "more_io_speed" in r["file"] else "site-packages"
+            print(f"round {i+1} {key:<10} src={src:<13} "
+                  f"data={r['get_data_only']['med']:6.1f}us  "
+                  f"data+sum={r['get_data_sum']['med']:6.1f}us  "
+                  f"floor={r['memmap_copy']['med']:5.1f}us")
+    print("\n=== best-of medians ===")
+    for key in ("INSTALLED", "TREE"):
+        best = min(acc[key], key=lambda r: r["get_data_only"]["med"])
+        g, gs, f = best["get_data_only"]["med"], best["get_data_sum"]["med"], best["memmap_copy"]["med"]
+        print(f"{key:<10} get_data={g:6.1f}us  get_data+sum={gs:6.1f}us  memmap-floor={f:5.1f}us")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/io_dl/backends-data_big-20260824-161621.json

@@ -0,0 +1,74 @@
+{
+  "floor_memmap": {
+    "full_ms": 151.852959,
+    "win_ms": 42.519625,
+    "mbps": 3107.3428078540105,
+    "us_per_win": 141.73208333333332
+  },
+  "mne_edf": {
+    "full_ms": 453.770375,
+    "win_ms": 575.13625,
+    "mbps": 1039.8633890544309,
+    "us_per_win": 1917.1208333333334
+  },
+  "mne_bdf": {
+    "full_ms": 994.840167,
+    "win_ms": 954.506333,
+    "mbps": 474.3065425503673,
+    "us_per_win": 3181.6877766666667
+  },
+  "mne_bv": {
+    "full_ms": 842.084667,
+    "win_ms": 479.360625,
+    "mbps": 560.3465049197957,
+    "us_per_win": 1597.86875
+  },
+  "mne_fif": {
+    "full_ms": 534.52,
+    "win_ms": 418.168083,
+    "mbps": 882.7718326723041,
+    "us_per_win": 1393.89361
+  },
+  "edfio_edf_eager": {
+    "full_ms": 255.659167,
+    "win_ms": null,
+    "mbps": 922.9578065550061,
+    "us_per_win": null
+  },
+  "npz": {
+    "full_ms": 206.126375,
+    "win_ms": null,
+    "mbps": 2289.17429901923,
+    "us_per_win": null
+  },
+  "h5_full": {
+    "full_ms": 153.334458,
+    "win_ms": 16.073,
+    "mbps": 3077.320037222162,
+    "us_per_win": 53.57666666666667
+  },
+  "h5_t10s": {
+    "full_ms": 247.235,
+    "win_ms": 217.567209,
+    "mbps": 1908.5453111412219,
+    "us_per_win": 725.22403
+  },
+  "h5_win": {
+    "full_ms": 213.12475,
+    "win_ms": 38.790833,
+    "mbps": 2214.004708509922,
+    "us_per_win": 129.30277666666666
+  },
+  "zarr_t10s": {
+    "full_ms": 4688.113042,
+    "win_ms": 9798.711709,
+    "mbps": 100.65013274481532,
+    "us_per_win": 32662.37236333333
+  },
+  "zarr_win": {
+    "full_ms": 765.368083,
+    "win_ms": 760.479458,
+    "mbps": 616.5127740243122,
+    "us_per_win": 2534.9315266666667
+  }
+}

benchmarks/io_dl/backends-data_big-20260824-161927.json

@@ -0,0 +1,92 @@
+{
+  "floor_memmap": {
+    "full_ms": 155.608959,
+    "win_ms": 42.666833,
+    "mbps": 3032.339545437098,
+    "us_per_win": 142.22277666666668
+  },
+  "mne_edf": {
+    "full_ms": 497.573125,
+    "win_ms": 608.397875,
+    "mbps": 948.3213145806458,
+    "us_per_win": 2027.9929166666666
+  },
+  "mne_edf_preloaded": {
+    "full_ms": null,
+    "win_ms": 683.8465,
+    "mbps": null,
+    "us_per_win": 2279.4883333333332
+  },
+  "mne_bdf": {
+    "full_ms": 1003.320542,
+    "win_ms": 852.924292,
+    "mbps": 470.2975572087908,
+    "us_per_win": 2843.0809733333335
+  },
+  "mne_bv": {
+    "full_ms": 814.662583,
+    "win_ms": 389.382458,
+    "mbps": 579.2081407033272,
+    "us_per_win": 1297.9415266666665
+  },
+  "mne_fif": {
+    "full_ms": 551.795,
+    "win_ms": 508.93025,
+    "mbps": 855.1349686024702,
+    "us_per_win": 1696.4341666666667
+  },
+  "mne_fif_preloaded": {
+    "full_ms": null,
+    "win_ms": 704.338583,
+    "mbps": null,
+    "us_per_win": 2347.7952766666667
+  },
+  "edfio_edf_eager": {
+    "full_ms": 244.071083,
+    "win_ms": null,
+    "mbps": 966.7782889298688,
+    "us_per_win": null
+  },
+  "edfio_edf_lazy_wins": {
+    "full_ms": null,
+    "win_ms": 347.37475,
+    "mbps": null,
+    "us_per_win": 1157.9158333333332
+  },
+  "npz": {
+    "full_ms": 228.083208,
+    "win_ms": null,
+    "mbps": 2068.802890566148,
+    "us_per_win": null
+  },
+  "h5_full": {
+    "full_ms": 158.631,
+    "win_ms": 28.184583,
+    "mbps": 2974.571174612781,
+    "us_per_win": 93.94861
+  },
+  "h5_t10s": {
+    "full_ms": 263.300583,
+    "win_ms": 279.080709,
+    "mbps": 1792.093259436497,
+    "us_per_win": 930.2690300000002
+  },
+  "h5_win": {
+    "full_ms": 217.69475,
+    "win_ms": 53.925958,
+    "mbps": 2167.5267777472814,
+    "us_per_win": 179.75319333333334
+  },
+  "zarr_t10s": {
+    "full_ms": 5098.628375,
+    "win_ms": 9509.274958,
+    "mbps": 92.54630172962938,
+    "us_per_win": 31697.583193333336
+  },
+  "zarr_win": {
+    "full_ms": 632.534083,
+    "win_ms": 503.473208,
+    "mbps": 745.9822524693898,
+    "us_per_win": 1678.2440266666665
+  }
+}

benchmarks/io_dl/backends-data_big-20260824-174932.json

@@ -0,0 +1,104 @@
+{
+  "floor_memmap": {
+    "full_ms": 123.97925,
+    "win_ms": 34.394292,
+    "mbps": 3805.953012298429,
+    "us_per_win": 114.64764000000001
+  },
+  "mne_edf": {
+    "full_ms": 464.572083,
+    "win_ms": 463.496,
+    "mbps": 1015.6856541033267,
+    "us_per_win": 1544.9866666666667
+  },
+  "mne_edf_preloaded": {
+    "full_ms": null,
+    "win_ms": 561.771959,
+    "mbps": null,
+    "us_per_win": 1872.5731966666667
+  },
+  "mne_bdf": {
+    "full_ms": 844.267375,
+    "win_ms": 643.647791,
+    "mbps": 558.8978254667248,
+    "us_per_win": 2145.4926366666664
+  },
+  "mne_bv": {
+    "full_ms": 528.579208,
+    "win_ms": 298.229583,
+    "mbps": 892.6934560770691,
+    "us_per_win": 994.0986099999999
+  },
+  "mne_fif": {
+    "full_ms": 474.158625,
+    "win_ms": 327.07125,
+    "mbps": 995.150515294328,
+    "us_per_win": 1090.2375
+  },
+  "mne_fif_preloaded": {
+    "full_ms": null,
+    "win_ms": 584.748375,
+    "mbps": null,
+    "us_per_win": 1949.16125
+  },
+  "mne_set": {
+    "full_ms": 510.646625,
+    "win_ms": 295.776875,
+    "mbps": 924.0425313689285,
+    "us_per_win": 985.9229166666667
+  },
+  "nwb": {
+    "full_ms": 135.857542,
+    "win_ms": 29.9985,
+    "mbps": 3473.191057733107,
+    "us_per_win": 99.995
+  },
+  "edfio_edf_eager": {
+    "full_ms": 171.696875,
+    "win_ms": null,
+    "mbps": null,
+    "us_per_win": null
+  },
+  "edfio_edf_lazy_wins": {
+    "full_ms": null,
+    "win_ms": 314.790875,
+    "mbps": null,
+    "us_per_win": 1049.3029166666668
+  },
+  "npz": {
+    "full_ms": 152.86025,
+    "win_ms": null,
+    "mbps": 3086.8665987396985,
+    "us_per_win": null
+  },
+  "h5_full": {
+    "full_ms": 104.1435,
+    "win_ms": 15.96875,
+    "mbps": 4530.855982370479,
+    "us_per_win": 53.229166666666664
+  },
+  "h5_t10s": {
+    "full_ms": 177.1665,
+    "win_ms": 160.295667,
+    "mbps": 2663.3658169010505,
+    "us_per_win": 534.31889
+  },
+  "h5_win": {
+    "full_ms": 148.633666,
+    "win_ms": 27.039625,
+    "mbps": 3174.6455072971153,
+    "us_per_win": 90.13208333333333
+  },
+  "zarr_t10s": {
+    "full_ms": 4115.615792,
+    "win_ms": 8396.288958,
+    "mbps": 114.65093532715262,
+    "us_per_win": 27987.629859999997
+  },
+  "zarr_win": {
+    "full_ms": 540.063791,
+    "win_ms": 373.51525,
+    "mbps": 873.710120662394,
+    "us_per_win": 1245.0508333333332
+  }
+}