Files
clawhdf5/benchmarks/bench_worldmodel_h5py.py
osobh b2dce41532
CI / test (push) Failing after 3s
bench: world-model sample loading — clawhdf5 reads h5py files 7x faster
than h5py (5e)

stable-worldmodel (arXiv 2605.21800, LeCun/Balestriero) supports HDF5 as
one of three native formats and measures generic HDF5 at 1,416-1,474
samples/s for per-frame sample loading. This measures clawhdf5 against
that shape, hardware-controlled: clawhdf5 and h5py reading the SAME file
on the SAME machine.

worldmodel_sampling example: mmap an (N,H,W,C) uint8 observation dataset,
read each frame once per pass in shuffled (dataloader) order. The file is
written by h5py (benchmarks/gen_worldmodel_frames.py) — clawhdf5 parsing
an externally-produced HDF5 file is itself the interop result — and read
by both clawhdf5 and the h5py counterpart (benchmarks/bench_worldmodel_h5py.py,
opening exactly stable-worldmodel's HDF5Dataset: swmr + 256 MB cache).

Results (tank, Ryzen 7 7800X3D, 20000x64x64x3 = 246 MB, in page cache,
median of 3):

  clawhdf5 zero-copy view        593k samples/sec   8.1x
  clawhdf5 materialised copy     518k samples/sec   7.1x
  h5py (swmr, 256 MB cache)       73k samples/sec   1.0x

The materialised-copy row is the fair equal-work comparison (to_vec per
frame, matching h5py's numpy materialisation) and is still 7.1x faster;
that the copy costs almost nothing shows the gap is h5py's per-frame call
overhead, not data movement. Honest caveats in BENCHMARKS.md: absolute
numbers are NOT comparable to the paper's (different hardware, smaller
frames, no torch/transform), only the same-machine ratio is; this is an
in-page-cache measurement isolating read-path overhead, not disk
bandwidth.

Adds only an example, two benchmark scripts, and a BENCHMARKS.md section —
no library code. (Workspace clippy has pre-existing toolchain drift
unrelated to this change; tracked separately.)
2026-08-07 22:54:26 -07:00

39 lines
1.1 KiB
Python

#!/usr/bin/env python3
"""h5py counterpart to worldmodel_sampling.rs — same file, same shuffled
per-frame access, same minimal touch (sum the frame bytes). Reports
samples/sec so the two sit side by side on one machine."""
import sys, time, numpy as np, h5py
path = sys.argv[1]
passes = int(sys.argv[2]) if len(sys.argv) > 2 else 5
def shuffled(n):
v = list(range(n))
state = 0x9E3779B97F4A7C15
for i in range(n - 1, 0, -1):
state = (state * 6364136223846793005 + 1442695040888963407) & 0xFFFFFFFFFFFFFFFF
j = (state >> 33) % (i + 1)
v[i], v[j] = v[j], v[i]
return v
# swmr + a 256 MB chunk cache: exactly stable-worldmodel's HDF5Dataset._open_h5.
f = h5py.File(path, "r", swmr=True, rdcc_nbytes=256 * 1024 * 1024)
d = f["observation"]
n = d.shape[0]
order = shuffled(n)
# warm
sink = 0
for i in order:
sink += int(d[i].sum())
t0 = time.perf_counter()
sink = 0
for _ in range(passes):
for i in order:
sink += int(d[i].sum())
elapsed = time.perf_counter() - t0
total = n * passes
print(f"h5py: {n} frames x {passes} passes = {total} reads in {elapsed:.3f}s")
print(f"h5py: {total/elapsed:.0f} samples/sec")