docs: describe the 0.2.0 engine, and fix the probes that measure it

CHANGELOG.md

@@ -1,5 +1,89 @@
 # Changelog
 
+## Unreleased
+
+- **`bench/probe_decode.py`'s decode-thread sweep had stopped measuring anything, and printed
+  a clean scaling curve anyway.** The decode pool went process-wide in 0.2.0 — built once by
+  the first dataset in the process, a later different value ignored with a warning — and the
+  probe's own prewarm built it at auto *before* the sweep began. So every point ran on the
+  same pool. The warnings went to stderr above the report; the table below still read as a
+  result, and on an 8-vCPU box it came back `756 -> 808 -> 921 MB/s` across
+  `decode_threads=1,2,8`: run-to-run noise on one 20-thread pool, arranged in the most
+  convincing shape it could have taken. The documented diagnosis for a flat sweep is "decode
+  isn't the limit", so the failure pointed *away* from itself. The sweep is gone.
+  `--decode-threads` stays as a single value, applied to the first dataset the probe builds —
+  the only one that can set it — and the section header now reports the pool's *effective*
+  thread count (`decode_pool_workers()`), so a number that never applied cannot be quoted.
+  Sweeping thread counts means one run per value, which the docstring says.
+
+- **The probe's raw-GET ceiling could not be measured on a sharded or a zarr-v2 store: every
+  key 404'd.** `_raw_get_mb_s` hand-built chunk keys from `metadata.chunks` in zarr-v3's
+  `c/i/j/k` layout. On a sharded array `chunks` is the *inner* chunk while a key addresses the
+  shard (the #56 distinction), so half the keys named nothing; on zarr-v2 the separator is `.`
+  with no `c/` prefix, so none of them did. It raised rather than returning a wrong number,
+  but it meant the insitu-vs-raw comparison — the probe's whole diagnosis — was unavailable on
+  every dynamical.org archive (all sharded) and on WeatherBench2 ARCO (v2), which are the
+  stores the module docstring offers as example invocations and the ones the engine has read
+  since #56. The grid now comes from `ChunkGrid.from_metadata` and the key from the array's
+  own `encode_chunk_key`: one spelling, the same one the engine plans reads with. Verified
+  against all three layouts, and the sharded case now correctly charges the shard.
+
+- **Every published "% of raw-GET ceiling" figure was computed in mixed units, and the premise
+  that licensed it is false.** insitu's MB/s counts *decoded* bytes; a raw GET moves *stored*
+  bytes. Dividing one by the other overstates "% kept" by exactly the **logical-to-stored
+  ratio**, and `bench/benchmark_plan.md` licensed the division by asserting the synthetic
+  stores are **incompressible** — `standard_normal` f32, so zstd "shrinks them ~0%" and
+  bytes-moved ≈ logical size. It is not incompressible: zstd level 0 takes **~7.6% off it
+  (1.082×)**. The error is small, one-directional and flattering, which is why it survived —
+  a ratio of two plausible MB/s numbers has nothing visibly wrong with it. `docs/benchmarks.md`
+  story 4's **81% / 84% are each inflated ~6 points** — the honest figures are ~75% / ~78% —
+  as is the 25 Gb/s scaling table, which used the same method. Those pages are **not**
+  corrected here.
+
+  The quantity `bench/probe_ceiling.py` now measures per store is deliberately **not** the
+  compression ratio but logical-to-stored, because that is the conversion "% kept" actually
+  needs — and the two differ. Zarr stores edge chunks padded to full size (the benchmark
+  geometry's 361 inner rows occupy 4×91 = 364 stored rows), so ~0.8% of every pass is padding
+  the user never sees. On that geometry the two effects nearly cancel and the ratio lands at
+  **1.081×**. Validated against exact known answers before being used to correct anything: a
+  store with no edge padding and no codec reports **exactly 1.000000**, and a padded
+  uncompressed one **exactly 361/364**, with `HEAD`-derived bytes equal to the filesystem's on
+  every store.
+
+- **The probe's residency claim was right about the mechanism and wrong about the number.**
+  Residency is *pinned* — the decoupling headline — but it is `2*block_chunks` only on a plain
+  pass. A windowed one sits higher, because a chunk feeds every block whose anchors reach it
+  and is released at the last (measured: 6 chunks at `block_chunks=2 --window 2`).
+
+- **The architecture page still described the pre-0.2.0 engine, and a reader sizing a budget
+  would have acted on it.** Eight claims were false rather than merely dated. Admission is
+  bounded by read-ahead permits and not by the byte budget, so "a full budget pauses
+  admission" and "the budget (sized to ~two blocks) lets it admit the next block" both named
+  the wrong mechanism. Peak co-residency was stated as two blocks unconditionally, which a
+  windowed view is not. Read-once was stated as a flat guarantee, which release-and-re-read
+  deliberately trades. And the persistence section said the cache log needs no compaction and
+  that a stale cache is a whole-directory wipe — on a page that 40 lines earlier described
+  per-array staleness correctly.
+
+  Both diagrams were worse off than the prose, because nobody greps a mermaid block: the
+  producer/consumer figure showed queue depth as the only backpressure path, omitting the
+  permit that returns on `unpin_block`, and the handoff figure showed `_drive -> _admit` with
+  neither the read-ahead permit nor the `pin_if_ready` hit check that now precede it, nor the
+  block boundary that ends an admission run.
+
+  The mechanisms 0.2.0 added are now documented where they were missing: the permit protocol,
+  one pool with many accountable iteration owners (which is the single parent of permit
+  ownership, budget-as-a-sum, per-pass counters and the own-tasks-only fetch skip), per-block
+  release, the one-writer `cache_dir` contract and `readonly_cache`. The vocabulary is single
+  again — "read-ahead permits" throughout, matching `read_ahead_bound` and the `read-ahead
+  permits exhausted` error, where the page had used one word for both that bound and
+  `max_inflight`. `scheduler.py` claimed "two bounded resources"; there are three.
+
+  Paid for by deletion rather than growth: the xbatcher cache comparison duplicated
+  `DESIGN.md`, the transform-fingerprint prose duplicated `docs/tuning.md`, and two rationale
+  blocks — the per-block batch-cutting history and the in-body-gating precedents — moved to
+  `DESIGN.md`, where a superseded design belongs. Net one heading fewer.
+
 ## 0.2.0 — 2026-09-11
 
 ### Upgrading

DESIGN.md

@@ -456,6 +456,25 @@ The shape above wasn't the first cut. The pivots that got here, and the roads no
   configurations are **slower than plain zarr-python in every row** (e.g. 1.93 s vs 1.10 s
   reading a 2 GB array) at **~2× the peak memory** (4.3 GB vs 2.4 GB) — and that is on local
   NVMe, without the forked-worker nesting that makes it worse for us.
+- **Batches cut over the epoch, not within each block.** Batches were once cut inside each
+  block's row range, which made the last batch of *every* block short whenever
+  `block_chunks x spc` was not a multiple of `batch_size`. Sample-once held either way, so
+  nothing failed loudly; what it broke was fixed-shape consumers — `torch.compile` and
+  `jax.jit` retrace per shape, and with the epoch's own short final block that was up to
+  *three* shapes, so the retrace cache never settled. Cutting over the whole epoch `order`
+  leaves exactly one short batch — the ordinary data-loader contract — and costs two monotone
+  frontiers over the block list.
+- **Transform scope at the call site, not inside the body.** The convention was for a
+  `chunk_transform` to test `chunk.read.array` and no-op on the rest. `output_geometry` cannot
+  see a test inside a function body, so it folded *every* transform's declared `output_inner`
+  into *every* variable's geometry: a regrid gated by an `if` to `t2m` also made the engine
+  believe `u10` was regridded, and `u10` was gathered as a truncated prefix of itself with
+  nothing raised. `applies(["t2m"], f)` puts the selector where the engine can read it.
+  "Subtract 273.15" is a fact about Kelvin; "`2m_temperature` is in Kelvin" is a fact about
+  your store — the old pattern welded them together, which is why one unit conversion had to
+  enumerate two archives' naming conventions and silently skipped a third. Configuration
+  inside, control flow outside: the same placement as sklearn's `ColumnTransformer`,
+  NVTabular's `cols >> Op()` and xarray's subset-then-`map`.
 - **No reshard.** The usual fix — rewrite the archive to one-sample-per-file — is a second
   copy that discards the chunk locality the store already has. We train in place and pay for
   it with an *approximate* shuffle instead.
@@ -962,7 +981,7 @@ measurements live in [docs/benchmarks.md](docs/benchmarks.md).
   removes all backpressure and the driver eagerly fetches the *whole* resident set,
   starving the current block and inflating cold TTFB (observed: ~7 s → ~50 s with
   `--cache-resident`). The stopgap — a low `--max-inflight` doubling as a throttle —
-  couples two dials. **Fix:** a fetch-ahead semaphore released on the consumer's `unpin`,
+  couples two dials. **Fix:** a read-ahead semaphore released on the consumer's `unpin`,
   bounding how far ahead of the consumer the driver admits, independent of retention. Then
   the cache retains consumed chunks (fast warm epochs) without eager whole-split fetching
   (fast cold TTFB). Applies to the V2 engine and the M-W residency path alike.

bench/benchmark_plan.md

@@ -134,11 +134,17 @@ stable median. Each family member is ~3.2 GB; the fat grids ~3.3 GB each.
 > **What the `make_dataset` print reports — and what we measure.** The
 > `~N MB uncompressed` line is the *logical* size (`n_samples × inner × 4`),
 > printed for a quick sanity check — not a `stat` of the written store. The values
-> are `standard_normal` f32, which is **incompressible**, so `compress=auto` (zstd)
-> shrinks them ~0% and the **on-disk footprint ≈ this number**. That is deliberate:
-> bytes-moved ≈ uncompressed keeps the MB/s and "% of raw-GET ceiling" math clean,
-> and the codec still runs on the decode path (zstd decompress per chunk) so the
-> decode-cost story is real. Caveat: real ERA5 compresses ~2–4×, so this synthetic
+> are `standard_normal` f32, which zstd still compresses by **~7.6%** (`compress=auto`),
+> and zarr pads edge chunks to full size, which adds **~0.8%** back whenever the inner
+> dims are not a multiple of `inner_chunks` (the `361 → 4×91` grids here). What matters
+> downstream is neither of those alone but their product, the **logical-to-stored ratio**
+> — ≈**1.08×** on these grids, measured per store, never assumed. It is load-bearing for
+> exactly one metric: insitu's MB/s counts *decoded* bytes while a raw GET moves *stored*
+> bytes, so any "% of raw-GET ceiling" must divide by that ratio before the two are
+> comparable. `bench/probe_ceiling.py` measures it via `HEAD` on the objects the raw sweep
+> fetches and reports the corrected figure; take it from there rather than from this line.
+> The codec runs on the decode path either way (zstd decompress
+> per chunk), so the decode-cost story is real. Caveat: real ERA5 compresses ~2–4×, so this synthetic
 > data moves *more* bytes/sample than production (conservative for the network
 > ceiling) — see the `make_dataset` docstring ("synthetic-but-realistic").
 

bench/probe_ceiling.py

@@ -0,0 +1,169 @@
+"""Probe: how much of the object store's raw-GET ceiling survives decode + gather?
+
+    uv run python -m bench.probe_ceiling --url gs://bucket/era5_fat_g16.zarr --var t2m
+    uv run python -m bench.probe_ceiling --url s3://bucket/era5_fat_g16.zarr --anon
+
+This is `docs/benchmarks.md` story 4 as a standalone, self-checking measurement. It exists
+because that story's published ratio is easy to get wrong in two specific ways, and both are
+silent -- the run completes and prints a plausible percentage either way.
+
+**1. The two sides are counted in different units.** insitu's MB/s is *decoded* bytes; a raw
+GET moves *stored* bytes. Dividing one by the other overstates "% kept" by exactly the
+**logical-to-stored** ratio -- which is what this probe measures, via ``HEAD`` on the same
+objects the raw sweep fetches, reporting the corrected number alongside the naive one.
+
+That ratio is not the compression ratio. `standard_normal` f32 is close to incompressible
+but not actually so (zstd level 0 takes ~7.6% off it), while zarr pads edge chunks to full
+size, which puts ~0.8% back on the wire. On the benchmark grids the two nearly cancel and
+the ratio is **~1.08x** -- ~6 points of "% kept". Assuming either effect instead of
+measuring both gets a plausible, flattering number: a store with no padding and no codec
+must report exactly 1.000, and a padded uncompressed one exactly `inner / padded_inner`.
+
+**2. Read-ahead permits can cap insitu below the concurrency being compared.** A pass may
+hold ``read_ahead_bound`` chunks ahead of its own consumer, so in-flight tiles are bounded by
+``bound x tiles-per-chunk`` however high ``max_inflight`` goes -- while the raw sweep, a flat
+thread pool over a key list, has no such bound. Comparing the two at a concurrency insitu
+cannot reach measures the bound, not the loader. The insitu sweep therefore reports
+``inflight_peak`` beside each point and flags any that never reached its ``max_inflight``.
+
+Both sides are measured in one session, interleaved rather than blocked, so provider drift
+hits every arm equally. The raw sweep is the control: it is the same bytes over the same
+network from the same box, with decode and gather removed.
+"""
+
+from __future__ import annotations
+
+import argparse
+import statistics
+import time
+import warnings
+from typing import Any
+
+import numpy as np
+import obstore
+
+from insitubatch import open_geometries, split_by_chunk
+from insitubatch.source import InSituDataset
+
+from .backend import build_store
+from .probe_decode import _raw_get_mb_s, _store_kwargs, stored_chunk_keys
+
+
+def stored_bytes(url: str, var: str, kw: dict[str, Any], max_chunks: int) -> int:
+    """Bytes actually on the wire for one pass -- ``HEAD``, not ``GET``, so it costs nothing.
+
+    Paired with the logical size this gives the **logical-to-stored** ratio the two sides
+    differ by -- not the compression ratio, which is a different number: zarr pads edge
+    chunks to full size, so the stored form carries padding the user never sees. Measured
+    per store, never assumed; ~8% of the wire is 6 points of "% kept".
+    """
+    store = build_store("obstore", url, **kw)
+    obs = obstore.store.from_url(url, **kw)
+    return sum(obstore.head(obs, k)["size"] for k in stored_chunk_keys(store, var, max_chunks))
+
+
+def insitu_pass(
+    url: str, var: str, kw: dict[str, Any], *, block_chunks: int, max_inflight: int
+) -> tuple[float, int, int]:
+    """One full pass; (decoded MB/s, peak in-flight tiles, peak resident chunks)."""
+    store = build_store("obstore", url, **kw)
+    geom = open_geometries(store, variables=[var])[var]
+    ds = InSituDataset(
+        store,
+        split_by_chunk(geom, fractions=(1.0, 0.0, 0.0)),
+        geometries={var: geom},
+        batch_size=32,
+        block_chunks=block_chunks,
+        max_inflight=max_inflight,
+        shuffle=False,
+    )
+    bps = int(np.prod(geom.inner_shape)) * geom.dtype.itemsize
+    n = 0
+    t = time.perf_counter()
+    for b in ds.train:
+        n += b.arrays[var].shape[0]
+    dt = time.perf_counter() - t
+    return n * bps / 1e6 / dt, ds.inflight_peak, ds.resident_peak
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    p.add_argument("--url", required=True, help="zarr URL (gs:// | s3:// | file://)")
+    p.add_argument("--var", default="t2m")
+    p.add_argument("--max-chunks", type=int, default=256, help="outer chunks per pass")
+    p.add_argument("--repeats", type=int, default=3, help="passes per point; median reported")
+    p.add_argument("--block-chunks", default="2,4,8", help="insitu sweep (sets the permit bound)")
+    p.add_argument("--max-inflight", default="32,64,128,256", help="insitu sweep")
+    p.add_argument("--concurrency", default="16,32,64,128", help="raw-GET sweep")
+    p.add_argument("--anon", action="store_true", help="anonymous (public gs:// / s3://)")
+    p.add_argument("--request-payer", action="store_true")
+    p.add_argument("--s3-express", action="store_true")
+    a = p.parse_args()
+
+    kw = _store_kwargs(a.url, "obstore", a.anon, a.request_payer, a.s3_express)
+    bcs = [int(x) for x in a.block_chunks.split(",")]
+    mis = [int(x) for x in a.max_inflight.split(",")]
+
+    store = build_store("obstore", a.url, **kw)
+    geom = open_geometries(store, variables=[a.var])[a.var]
+    tiles = len(stored_chunk_keys(store, a.var, a.max_chunks)) // min(a.max_chunks, geom.n_chunks)
+    n_chunks = min(a.max_chunks, geom.n_chunks)
+    logical = n_chunks * geom.sample_chunk_size * int(np.prod(geom.inner_shape))
+    logical *= geom.dtype.itemsize
+    on_wire = stored_bytes(a.url, a.var, kw, a.max_chunks)
+    ratio = logical / on_wire
+
+    print(f"probe {a.url}  var={a.var}")
+    print(f"  {n_chunks} outer chunks x {tiles} stored objects each")
+    print(f"  {logical / 1e6:.0f} MB logical / {on_wire / 1e6:.0f} MB on the wire = {ratio:.4f}x")
+
+    print("\nprewarm (discarded) ...")
+    insitu_pass(a.url, a.var, kw, block_chunks=bcs[0], max_inflight=mis[-1])
+
+    grid = [(bc, mi) for bc in bcs for mi in mis]
+    acc: dict[tuple[int, int], list[float]] = {k: [] for k in grid}
+    peak: dict[tuple[int, int], tuple[int, int]] = {}
+    for _ in range(a.repeats):  # interleaved, so drift hits every arm alike
+        for bc, mi in grid:
+            mb, ip, rp = insitu_pass(a.url, a.var, kw, block_chunks=bc, max_inflight=mi)
+            acc[(bc, mi)].append(mb)
+            peak[(bc, mi)] = (ip, rp)
+
+    print(f"\ninsitu decoded MB/s (median of {a.repeats}):")
+    print(
+        f"  {'block_chunks':>12} {'max_inflight':>12} {'inflight_pk':>11}"
+        f" {'resident':>8} {'MB/s':>9}"
+    )
+    for bc, mi in grid:
+        ip, rp = peak[(bc, mi)]
+        flag = "  <-- permit-bound" if ip < mi else ""
+        med = statistics.median(acc[(bc, mi)])
+        print(f"  {bc:>12} {mi:>12} {ip:>11} {rp:>8} {med:>9.1f}{flag}")
+
+    print(f"\nraw obstore GET, the control -- same bytes, no decode (median of {a.repeats}):")
+    raws: dict[int, float] = {}
+    for c in (int(x) for x in a.concurrency.split(",")):
+        raws[c] = statistics.median(
+            _raw_get_mb_s(a.url, a.var, kw, c, a.max_chunks, "obstore") for _ in range(a.repeats)
+        )
+        print(f"  concurrency={c:>4}: {raws[c]:8.1f} MB/s")
+
+    best_in = max(statistics.median(v) for v in acc.values())
+    best_raw = max(raws.values())
+    if any(ip < mi for (_, mi), (ip, _) in peak.items()):
+        warnings.warn(
+            "some insitu points never reached their max_inflight -- the permit bound capped "
+            "them. Raise --block-chunks before reading the ratio.",
+            stacklevel=1,
+        )
+    print(f"\nbest insitu (decoded) : {best_in:8.1f} MB/s")
+    print(f"best raw GET (stored) : {best_raw:8.1f} MB/s")
+    print(f"  naive ratio         : {100 * best_in / best_raw:5.1f}%   <-- mixed units, too high")
+    print(f"  raw as decoded-equiv: {best_raw * ratio:8.1f} MB/s")
+    print(f"  % of ceiling kept   : {100 * best_in / (best_raw * ratio):5.1f}%")
+
+
+if __name__ == "__main__":
+    main()