diff --git a/CHANGELOG.md b/CHANGELOG.md
index fbc6a0b..15d062b 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/DESIGN.md b/DESIGN.md
index f89b189..974e76c 100644
--- a/DESIGN.md
+++ b/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.
diff --git a/bench/benchmark_plan.md b/bench/benchmark_plan.md
index d7e9f6d..859f825 100644
--- a/bench/benchmark_plan.md
+++ b/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").
diff --git a/bench/probe_ceiling.py b/bench/probe_ceiling.py
new file mode 100644
index 0000000..224ddca
--- /dev/null
+++ b/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()
diff --git a/bench/probe_decode.py b/bench/probe_decode.py
index b8331d5..94da5dd 100644
--- a/bench/probe_decode.py
+++ b/bench/probe_decode.py
@@ -6,20 +6,33 @@
Measurements over the first ``--max-chunks`` chunks (quick on a 25 GB store):
- 1. insitu throughput vs ``decode_threads`` (1,2,4,8,auto) — where decode parallelism
- saturates; flat past N cores => decode isn't the limit.
- 1b. **The V2 decoupling headline**: insitu throughput + peak residency vs
+ 1. **The V2 decoupling headline**: insitu throughput + peak residency vs
``max_inflight``, with ``block_chunks`` fixed small. V2 dials network
concurrency with ``max_inflight`` alone; throughput should rise to the network
knee and then stay *flat* (not fall, as v1's nested caps did when oversubscribed)
- while residency stays pinned at ``2*block_chunks`` — independent of concurrency.
- 2. raw obstore concurrent GET MB/s vs concurrency (1,4,8,16,32) — pure fetch, no
- decode.
+ while residency stays **pinned** — independent of concurrency. Pinned is the
+ claim, not a particular number: a plain single-variable pass sits at
+ ``2*block_chunks``, and a windowed one (``--window``) sits higher, because a
+ chunk feeds every block whose anchors reach it and is released only at the last.
+ 2. raw concurrent GET MB/s vs concurrency (1,4,8,16,32) — pure fetch, no decode.
3. (synthetic only) insitu compressed vs uncompressed — the codec's share.
-Diagnosis: if raw GET (2) far exceeds insitu (1/1b), we're decode/loop-limited and
+Diagnosis: if raw GET (2) far exceeds insitu (1), we're decode/loop-limited and
``max_inflight`` / the decode pool is the lever; if raw GET also caps near insitu,
it's the network/endpoint (more/bigger parallel streams, in-region S3, the gateway).
+
+**Comparing (1) against (2) assumes the data barely compresses.** insitu MB/s counts
+*decoded* bytes; raw GET counts the *stored* bytes it moved. On the synthetic store
+those agree, deliberately — ``make_dataset`` writes ``standard_normal`` f32, which zstd
+does not shrink (see ``bench/benchmark_plan.md``). On a real ``--url`` they do not: ERA5
+compresses ~2-4x, so raw GET is understated against insitu by that ratio, which biases
+the diagnosis *toward* "not decode-limited". Divide insitu's MB/s by the store's
+compression ratio before reading the gap.
+
+``decode_threads`` is a **process-wide** setting, not a sweep: the decode pool is built
+once, by the first dataset in the process, and a later different value is ignored. Set
+``--decode-threads`` and run the probe again to compare -- one value per process is the
+only way to measure it honestly.
"""
from __future__ import annotations
@@ -44,6 +57,7 @@
import zarr
from insitubatch import SplitManifest, SplitName, close_store, open_geometries, split_by_chunk
+from insitubatch.scheduler import decode_pool_workers
from insitubatch.source import InSituDataset
from ._profile import record_pyspy
@@ -213,6 +227,31 @@ def epoch_mb(epoch: int) -> float:
return cold, warm
+def stored_chunk_keys(store: Any, var: str, max_chunks: int) -> list[str]:
+ """Every stored-object key under the first ``max_chunks`` OUTER chunks of ``var``.
+
+ The grid comes from ``ChunkGrid.from_metadata`` and the key from the array's own
+ ``encode_chunk_key`` -- the same "one spelling of the shape of one stored object" the
+ engine plans reads with. Neither is cosmetic. ``metadata.chunks`` is the *inner* chunk
+ on a sharded array while a key addresses the *shard*, so building the grid from
+ ``chunks`` enumerates keys that do not exist (every dynamical.org archive is sharded);
+ and the ``c/i/j`` layout is zarr-v3's, while v2 writes ``i.j`` (WeatherBench2 ARCO is
+ v2). Either way a hand-built key 404s on exactly the stores the engine reads fine.
+ """
+ from zarr.core.chunk_grids import ChunkGrid
+
+ arr = zarr.open_array(store=store, path=var, mode="r")
+ meta = arr.metadata
+ stored = tuple(ChunkGrid.from_metadata(meta).chunk_shape) # the shard, if sharded
+ n_outer = min(max_chunks, math.ceil(arr.shape[0] / stored[0]))
+ inner_ranges = [range(math.ceil(s / c)) for s, c in zip(arr.shape[1:], stored[1:], strict=True)]
+ return [
+ f"{var}/" + meta.encode_chunk_key((oi, *inner))
+ for oi in range(n_outer)
+ for inner in itertools.product(*inner_ranges)
+ ]
+
+
def _raw_get_mb_s(
url: str, var: str, kw: dict[str, Any], concurrency: int, max_chunks: int, backend: str
) -> float:
@@ -221,21 +260,12 @@ def _raw_get_mb_s(
concurrent gcsfs ``cat_file``. Comparing the two is the cleanest "what does the
Python fsspec layer cost" number, isolated from decode/gather/loop.
- Enumerates the *real* stored-chunk grid (outer x inner) from the array's chunks, so
- it's correct for spatially-chunked arrays. max_chunks bounds the number of OUTER
- chunks; every inner chunk under them is fetched, concurrency threads at a time.
+ Reads the *real* stored-chunk grid via :func:`stored_chunk_keys`, so it is correct for
+ spatially-chunked, sharded and zarr-v2 arrays alike. max_chunks bounds the number of
+ OUTER chunks; every stored object under them is fetched, concurrency threads at a time.
"""
meta_store = build_store(backend, url, **kw)
- arr = zarr.open_array(store=meta_store, path=var, mode="r") # for the chunk grid
- n_outer = min(max_chunks, math.ceil(arr.shape[0] / arr.chunks[0]))
- inner_ranges = [
- range(math.ceil(s / c)) for s, c in zip(arr.shape[1:], arr.chunks[1:], strict=True)
- ]
- keys = [
- f"{var}/c/{oi}/" + "/".join(map(str, inner))
- for oi in range(n_outer)
- for inner in itertools.product(*inner_ranges)
- ]
+ keys = stored_chunk_keys(meta_store, var, max_chunks)
if backend == "fsspec":
import fsspec
@@ -275,18 +305,25 @@ def main() -> None:
"--backend",
default="obstore",
choices=["obstore", "fsspec"],
- help="store backend for the fetch-sensitive sections (1b, 2): obstore or fsspec/gcsfs",
+ help="store backend for the fetch-sensitive sections (1, 2): obstore or fsspec/gcsfs",
)
p.add_argument("--max-chunks", type=int, default=64, help="chunks to probe (bounds the cost)")
p.add_argument("--repeats", type=int, default=3, help="runs per point; report median (min-max)")
- p.add_argument("--decode-threads", default="1,2,4,8,0", help="sec 1 sweep (0=auto)")
+ p.add_argument(
+ "--decode-threads",
+ type=int,
+ default=0,
+ help="size the process-wide decode pool (0 = auto = min(32, cpu+4)). Applied to "
+ "the first dataset this process builds, which is the only one that can set it; "
+ "to compare thread counts, run the probe once per value.",
+ )
p.add_argument(
"--block-chunks",
type=int,
default=2,
- help="fixed shuffle window / residency for the max_inflight sweep (sec 1b)",
+ help="fixed shuffle window / residency for the max_inflight sweep (sec 1)",
)
- p.add_argument("--max-inflight", default="8,16,32,64", help="sec 1b sweep (the V2 dial)")
+ p.add_argument("--max-inflight", default="8,16,32,64", help="sec 1 sweep (the V2 dial)")
p.add_argument(
"--window",
type=int,
@@ -301,22 +338,16 @@ def main() -> None:
nargs="?",
const="probe-profile.svg",
default=None,
- help="record a py-spy --native flamegraph of the sec 1b sweep to PATH "
+ help="record a py-spy --native flamegraph of the sec 1 sweep to PATH "
"(default probe-profile.svg; .json => speedscope). Needs ptrace_scope=0 "
"or sudo; see bench/_profile.py.",
)
p.add_argument(
"--cache-dir",
default=None,
- help="run the cross-epoch cache test (sec 1c): cold vs cached epoch, mmap slots here",
+ help="run the cross-epoch cache test (sec 1b): cold vs cached epoch, mmap slots here",
)
p.add_argument("--no-raw", action="store_true", help="skip sec 2 (raw GET re-reads the data)")
- p.add_argument(
- "--no-decode-sweep",
- action="store_true",
- help="skip sec 1 (the decode_threads sweep) -- a one-time decode-saturation "
- "check; redundant when you only want the sec-1b max_inflight sawtooth",
- )
p.add_argument(
"--no-warm",
action="store_true",
@@ -348,17 +379,20 @@ def show(label: str, fn: Callable[[], float]) -> None:
# Prewarm: a cold S3 prefix is rate-limited and obstore's TLS pool is empty, so
# the first sweep point otherwise eats the ramp-up and reads at ~1 stream. One
- # throwaway burst at the top fixes it for every section below (and makes 1c's
+ # throwaway burst at the top fixes it for every section below (and makes 1b's
# "cold" epoch S3-warm/pool-cold, isolating decode+pool reuse from the S3 ramp).
if not a.no_warm:
mi_warm = max(int(x) for x in a.max_inflight.split(","))
print(f"0) prewarm: {a.max_chunks} chunks @ max_inflight={mi_warm} (discarded) ...")
try:
+ # Also where decode_threads lands: the pool is process-wide and the first
+ # dataset built sizes it, so passing it only to the sections below would be
+ # ignored -- with a warning, under a table that still looked like a result.
_insitu(
a.url,
a.var,
kw,
- decode_threads=0,
+ decode_threads=a.decode_threads,
max_chunks=a.max_chunks,
block_chunks=bc,
max_inflight=mi_warm,
@@ -372,25 +406,10 @@ def show(label: str, fn: Callable[[], float]) -> None:
if a.window:
print(f"(windowed: anchor input + {a.window} shifted target view(s); MB/s = input bytes)\n")
- if not a.no_decode_sweep:
- print(f"1) insitu MB/s vs decode_threads (median of {a.repeats}, block_chunks={bc}):")
- for dt in (int(x) for x in a.decode_threads.split(",")):
- show(
- f"decode_threads={dt or 'auto':>4}",
- partial(
- _insitu_mb,
- a.url,
- a.var,
- kw,
- decode_threads=dt,
- max_chunks=a.max_chunks,
- block_chunks=bc,
- window=a.window,
- backend=a.backend,
- ),
- )
-
- print(f"\n1b) insitu MB/s + peak residency vs max_inflight (block_chunks={bc}, decode auto):")
+ print(
+ f"\n1) insitu MB/s + peak residency vs max_inflight (block_chunks={bc}, "
+ f"decode pool={decode_pool_workers() or 'unbuilt'} threads):"
+ )
print(" V2 wants throughput flat (not falling) past the knee, residency pinned.")
profile = record_pyspy(a.profile) if a.profile else contextlib.nullcontext()
with profile: # profile scope = the V2 acceptance sweep (insitu only, no raw-GET)
@@ -400,7 +419,7 @@ def show(label: str, fn: Callable[[], float]) -> None:
a.url,
a.var,
kw,
- decode_threads=0,
+ decode_threads=a.decode_threads,
max_chunks=a.max_chunks,
block_chunks=bc,
max_inflight=mi,
@@ -417,7 +436,7 @@ def show(label: str, fn: Callable[[], float]) -> None:
)
if a.cache_dir:
- print(f"\n1c) cross-epoch cache ({a.max_chunks} chunks, budget holds all, mmap):")
+ print(f"\n1b) cross-epoch cache ({a.max_chunks} chunks, budget holds all, mmap):")
cold, warm = _insitu_cache(
a.url,
a.var,
@@ -449,11 +468,18 @@ def show(label: str, fn: Callable[[], float]) -> None:
variables=["t2m"],
compress=False,
)
- print("\n3) codec cost (synthetic, compressed vs uncompressed, auto decode_threads):")
+ print("\n3) codec cost (synthetic, compressed vs uncompressed):")
for label, u in (("compressed", a.url), ("uncompressed", none_url)):
show(
f"{label:12}",
- partial(_insitu_mb, u, "t2m", {}, decode_threads=0, max_chunks=a.max_chunks),
+ partial(
+ _insitu_mb,
+ u,
+ "t2m",
+ {},
+ decode_threads=a.decode_threads,
+ max_chunks=a.max_chunks,
+ ),
)
diff --git a/docs/architecture.md b/docs/architecture.md
index 4d0a0a9..9bd18d4 100644
--- a/docs/architecture.md
+++ b/docs/architecture.md
@@ -185,8 +185,14 @@ future via `call_soon_threadsafe`. What crosses the boundary is a Python object
```
insitu-prefetch zarr_io (ONE loop, shared with all zarr sync work) insitu-dec (N)
│ │ │
- │ _drive ──► _admit ──► pool.try_admit [_cv, brief] │
- │ _one, holding one `max_inflight` permit: │
+ │ _drive, per chunk — a block boundary ends the run, so │
+ │ each block claims what it reads on its own reference: │
+ │ ⓪ _take_ahead — ONE read-ahead permit, taken before │
+ │ we reference anything; back on consumer unpin_block │
+ │ ─► pin_if_ready — a hit is pinned and permitted alike, │
+ │ and costs no fetch │
+ │ ─► _admit ─► pool.try_admit [_cv, brief] │
+ │ _one, per tile, holding one `max_inflight` slot: │
│ │ │
│ ① await store.get(key) ──── obstore / gcsfs, inline ──┐ │
│ │ coroutine SUSPENDS, loop runs other tasks │ │
@@ -218,7 +224,13 @@ may sit on a loop shared with the rest of the process.
`max_inflight` is what bounds the executor: a `ThreadPoolExecutor` queue is unbounded, and
the semaphore is held across ①②③, so in-flight tiles, queued decode work and decoded-tile
-residency are all bounded by that one dial.
+residency are all bounded by that one dial. (It bounds *tiles*; read-ahead permits bound
+*chunks*. Two resources, two units — see the prefetch section.)
+
+A persisted slot stores its tiles **tile-major** — `(n_tiles, *tile_shape)`, each stored tile
+contiguous, in `inner_index` order — which is what keeps ③ to one code path: a revived chunk
+comes back as zero-copy views of the mapping, tiled exactly like a freshly fetched one, so
+`gather` never has to ask whether a slot was assembled or delivered.
#### Built on zarr's abstractions
@@ -293,6 +305,13 @@ Every physical axis of a zarr array has one **role**:
fetches each field axis's tiles and assembles them, decode-once, into the sample's
slot.
+`open_geometries(store)` takes the group's **data variables** and skips coordinates and
+grid mappings, detected from `dimension_names` (v3) or xarray's `_ARRAY_DIMENSIONS` (v2) —
+a CF/xarray-written group carries `time`/`latitude`/`longitude` and a 0-D `spatial_ref`
+alongside the data, and none of them has a sample axis to draw along. The inference is
+deliberately narrow: only a *self-named* 1-D array is a coordinate, so a station series over
+`('time',)` is still data, and naming an array in `variables=` bypasses the test entirely.
+
`ArrayGeometry.shape`/`chunks` are in **physical** (zarr) order — they mirror the
array metadata. The engine works in a *logical* view where the sample axis leads; the
single physical↔logical permutation is confined to the scheduler (address the store
@@ -400,6 +419,7 @@ flowchart LR
Q --> CONS["consumer __iter__
pops finished batches"]
CONS --> STEP[("train / infer step")]
Q -.->|"full ⇒ producer blocks
(backpressure)"| ASM
+ CONS -.->|"unpin_block ⇒ read-ahead permit returns
(bounds how far ahead)"| ASM
```
- **Producer** starts the scheduler over the epoch's chunks, then walks the epoch's
@@ -408,17 +428,50 @@ flowchart LR
`queue.Queue(maxsize=d)`.
- **Consumer** (`__iter__`) just pops finished batches → the train/infer step
overlaps with IO+decode+assembly of the next `d` batches.
-- **Backpressure / memory bound** — queue depth `d` + the pool's byte budget cap
- residency; a full queue pauses the consumer, a full budget pauses admission.
+- **Backpressure — two bounds, two questions.** *How far ahead* the producer may run
+ is bounded by **read-ahead permits**: it takes one per chunk and gets it back when the
+ consumer releases that chunk, so a pass runs at most its own working set ahead of its own
+ consumer. *How much is resident* is bounded by the pool's byte budget. A full queue pauses
+ the consumer, exhausted permits pause admission, and a full budget evicts unpinned-LRU to
+ admit a miss.
+- **The permit bound is computed, never configured** (`read_ahead_bound`): everything
+ still live at a block — first read at or before it, last read at or after it — plus the
+ whole of the next one, which the driver may work on while the consumer gathers this one.
+ It counts the read-unions the consumer actually waits on rather than `block_chunks`,
+ because a windowed variable's chunk feeds several blocks and is released only at the last
+ of them, and a variable chunked finer than the reference grid contributes several chunks
+ per anchor chunk. Below the bound a pass deadlocks against its own limit; above it one
+ iteration can claim the whole pool and starve another sharing it.
- **Continuous fetch** — the scheduler keeps `max_inflight` tiles in flight across
- block boundaries, and the budget (sized to ~two blocks) lets it admit the next
- block while the current one drains, so block-boundary IO overlaps compute. (At
- zero per-batch compute the loader is IO-throughput-bound, so the boundary is only
- smoothed, not removed — the network ceiling, not a scheduling gap.)
+ block boundaries, and the bound covers the next block as well as the live set, so it
+ admits the next block while the current one drains and block-boundary IO overlaps
+ compute. (At zero per-batch compute the loader is IO-throughput-bound, so the boundary is
+ only smoothed, not removed — the network ceiling, not a scheduling gap.)
- **Lifecycle** — early consumer exit sets a stop flag and drains the queue so a
producer parked on a full `put` can exit before the scheduler is closed.
- **Knobs:** `prefetch_depth` (queue depth `d`), `max_inflight`, `block_chunks`,
- `cache_budget_bytes`.
+ `cache_budget_bytes`. Read-ahead depth is **not** one — it is computed per pass.
+
+**One pool, many iterations.** One `InSituDataset` owns one `ChunkPool`, and every
+iteration open on it — `ds.train` and `ds.val` under `zip`, or two `DataLoader`s — is an
+independent **owner** inside it, minted when the iteration starts and dropped when it ends.
+Four things follow, and they are the same fact four times:
+
+- read-ahead permits belong to one owner and come back only from *that* owner's
+ `unpin_block`, so another pass being mid-wait says nothing about this one;
+- residency is the **sum** of the owners' working sets, not the maximum — the auto budget
+ covers one iteration deliberately (the engine cannot know how many you intend, and
+ guessing high costs memory in the single-iteration case that is almost every case), so a
+ second iteration the budget cannot hold raises when it starts rather than starving later;
+- a pass's report reads its **own** counters (`PassCounters`) while the pool keeps
+ cumulative totals, because a hit rate belongs to a pass and an eviction or a peak
+ co-residency belongs to the pool;
+- the driver may skip a chunk whose tiles are already in flight only when they are **its
+ own** tasks — an abandoned pass's tasks are cancelled, so anyone standing on them waits
+ on a delivery that is not coming.
+
+Sizing for concurrent iterations is on the
+[tuning page](tuning.md#several-iterations-at-once-multiply-the-budget).
Same shape as `torchdata.nodes.Prefetcher`, but async-native. This is what turns a
throughput win into a *GPU-fed* win.
@@ -480,10 +533,13 @@ The pipeline holds two guarantees, and they are **orthogonal** — batch size to
- **read-once** — a stored tile is fetched and decoded exactly once, however many samples
and batches reference it (the read plan dedups; the `ChunkPool` keeps it resident; gather
- reads from the slot). Across **epochs** it holds only while the tile stays resident: the
- default budget is the working set, which is read-once *per epoch*; raise
- `cache_budget_bytes` past it and an unevicted chunk is a cross-epoch hit too
- (see [The caching continuum](#the-caching-continuum)).
+ reads from the slot). It is bounded by residency, so it holds for as long as the tile stays
+ resident: the default budget is the working set, which is read-once *per epoch*; raise
+ `cache_budget_bytes` past it and an unevicted chunk is a cross-epoch hit too. One policy
+ trades it deliberately — a windowed, shuffled pass with `persist=True` releases each block's
+ chunks as it drains and re-admits them later **from the on-disk cache**, which is a re-read
+ but neither a re-fetch nor a re-decode (see
+ [The caching continuum](#the-caching-continuum)).
- **sample-once** — each valid sample lands in exactly one batch.
`order` is the ledger for sample-once: an `(N, 2)` array of `[chunk_id, within]`, one row per
@@ -559,16 +615,13 @@ the decode-once suite.
steps-per-epoch is `⌈N / bs⌉` and every batch is full except the last — the ordinary
data-loader contract. A batch may therefore span a shuffle-block boundary, which the producer
handles with two monotone frontiers over the block list: wait every block a batch draws from,
-and release a block once a batch has consumed its last row. Peak co-residency stays at **two
-blocks**, which is already the budget floor (the working set is sized at the current block plus
-one read-ahead block) — and when `bs` divides a block's row count no batch straddles at all, so
-the frontier costs nothing.
+and release a block once a batch has consumed its last row. On a non-windowed pass peak
+co-residency is **two blocks**, which is already the budget floor (the working set is sized at
+the current block plus one read-ahead block) — and when `bs` divides a block's row count no
+batch straddles at all, so the frontier costs nothing. A windowed view widens both terms: its
+chunks live from the first block that reads them to the last, which is what `working_set_bytes`
+sizes the floor from and what `read_ahead_bound` counts permits over.
-Batches used to be cut *within* each block's row range, which made the last batch of **every**
-block short whenever `block_chunks × 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.
There is deliberately no `drop_last`. Dropping the epoch's short tail is the caller's choice
and `len(batch)` is the whole implementation:
@@ -654,12 +707,10 @@ that halves `t2m` and no-ops on `u10` therefore makes the engine believe `u10` i
too, and `u10` is gathered as a truncated prefix of itself, its cache file permanently
unable to revive — with nothing raised. Right shape, right dtype, wrong numbers.
-The division of labour: "subtract 273.15" is a fact about Kelvin, "`2m_temperature` is in
-Kelvin" is a fact about your store. A transform may be *parameterized* by variable — a
-fitted `StandardScaler`'s statistics legitimately are, and it validates a declared scope
-against its own stats dict — but it must not decide whether it runs. Configuration inside,
-control flow outside. sklearn's `ColumnTransformer`, NVTabular's `cols >> Op()` and
-xarray's subset-then-`map` all place the selector the same way.
+The division of labour: a transform may be *parameterized* by variable — a fitted
+`StandardScaler`'s statistics legitimately are, and it validates a declared scope against its
+own stats dict — but it must not decide whether it runs. **Configuration inside, control flow
+outside.**
Scope also buys two smaller things: the engine skips the call entirely for arrays outside
it, and cache invalidation is per array — see
@@ -808,7 +859,7 @@ two train chunks shares its neighbours' weather). Set `contiguous=False` only wh
```python
from insitubatch import obstore_store, open_geometries, split_by_chunk
-store = obstore_store(url) # or fsspec_store / arraylake_store
+store = obstore_store(url) # or fsspec_store / icechunk_store / arraylake_store
geom = open_geometries(store)[var]
# time series (default): contiguous blocks, no cross-boundary leakage
manifest = split_by_chunk(geom, fractions=(0.8, 0.1, 0.1))
@@ -872,17 +923,27 @@ object, the `ChunkPool`, parameterized by a byte budget:
| layer | reuse scope | how |
|---|---|---|
| read-plan dedup | within a request | a chunk's tiles are fetched once, held by one slot |
-| residency | within an epoch | a small budget (the working set, ~2 blocks) |
+| residency | within an epoch | a small budget (the working set — `working_set_bytes`) |
| cache | across epochs | a large budget retains drained chunks |
-A chunk is **pinned** while the current epoch needs it; once its shuffle-block is
-drained it becomes **unpinned** — LRU-evictable but not dropped. The pool drops
-unpinned chunks only under budget pressure (evicting LRU to admit a miss). With a
-small budget that is prompt — the read-once buffer, where each chunk is still
-read+decoded **once per epoch** (a naive per-batch eviction would re-read chunks
-whose samples scatter across a shuffle block). Raise the budget past the working set
-and drained chunks linger, so a still-resident prepped chunk is a cross-epoch hit:
-the same machinery becomes the cache by *"don't evict."*
+A chunk is **pinned** until the last block that reads it has drained; then it becomes
+**unpinned** — LRU-evictable but not dropped. The pool drops unpinned chunks only under
+budget pressure (evicting LRU to admit a miss). With a small budget that is prompt — the
+read-once buffer, where each chunk is still read+decoded **once per epoch** (a naive
+per-batch eviction would re-read chunks whose samples scatter across a shuffle block). Raise
+the budget past the working set and drained chunks linger, so a still-resident prepped chunk
+is a cross-epoch hit: the same machinery becomes the cache by *"don't evict."*
+
+**Holding to the last use is the right trade only while the split fits.** A windowed view
+reads `anchor + offset` and shuffle permutes chunk order, so a chunk one block reads can be
+wanted again most of an epoch later — held first-use to last-use, residency is the *split*
+rather than two blocks. With `persist=True` a windowed, shuffled pass instead hands each
+block's chunks back as it drains and re-admits them later from the on-disk cache: on a
+256-chunk split the floor falls from 259 chunks to 12 for a three-chunk lead, and from 316 to
+18 for leads `{0, 24, 240}`, delivering the same samples. It is a rule rather than a knob, and
+it keys on `persist` rather than on `cache_dir` because only a persisted slot survives its own
+eviction as a revivable file — with `cache_dir` alone the backing is unlinked on eviction and
+the re-read goes back to the store, which is the opposite of the trade the policy assumes.
Backing is **heap or mmap** (`cache_dir` → mmap'd `.npy` on local NVMe): a heap slot
adopts the decoded tile, and the mmap tier writes it straight into its tile-major slice and
@@ -896,30 +957,6 @@ fat-chunk regime (one chunk → many batches); scoring/verification (reference c
reused across metrics, lead times, models); datasets that fit in RAM/NVMe
(effectively in-memory at GPU-fed speed after the first pass).
-### Two cache models: chunks vs batches
-
-insitubatch and xbatcher both cache — they cache *different things*, and each choice
-buys something real. [xbatcher](https://github.com/xarray-contrib/xbatcher) serializes
-**assembled batches** to a zarr store; insitubatch retains **decoded, chunk-transformed
-chunks** in the pool. Caching whole batches is a deliberate design choice, not a
-shortcut: it is exactly what lets xbatcher's cache survive process exit today.
-
-| dimension | insitubatch — chunk pool | xbatcher — batch cache |
-|---|---|---|
-| unit cached | decoded + chunk-transformed **chunk** (deduped) | **assembled batch**, in batch layout |
-| extra copy | none beyond the batch — `gather` copies rows from the slot straight into it | the batch, **plus** a separate materialized copy in the cache |
-| key | `(array, chunk_index)` | batch index → zarr store |
-| backing | heap or mmap'd `.npy` (reclaimable NVMe page cache) | zarr store (local dir or cloud) |
-| cross-epoch reuse | intrinsic (just don't evict) | yes |
-| cross-run persistence | **yes** — `persist=True` (see below) | **yes** — survives restarts |
-| shuffle | sits *before* shuffle: sample→batch membership re-drawn every epoch | batch composition frozen; only batch *order* reshuffles |
-| sweet spot | many samples per chunk, fat-chunk, multi-epoch, scoring reuse | one-sample-per-chunk, stable batch defs reused across runs |
-
-Both persist across runs now; the distinction is the **unit**. insitu caches deduped
-decoded chunks (stored once, with nothing staged between slot and batch) and keeps a stronger
-per-epoch shuffle because the cache is *upstream* of shuffling; xbatcher caches
-materialized batches with frozen composition. Pick by regime, not by a missing feature.
-
### Cross-run persistence
`persist=True` (with a `cache_dir`) turns the mmap tier into a cache that **survives
@@ -973,24 +1010,11 @@ lives in an array rather than in its source.
**The fingerprint is per *array*, over the transforms scoped to it.** `chunk_transforms`
is one ordered list, but a transform restricted with `applies` contributes only to the
-arrays it names, so each array's entries carry their own 16-hex `pipeline` hash. Editing a
-transform scoped to `t2m` marks `t2m` stale and leaves `u10`/`v10` — whose bytes did not
-change — valid. Staleness is therefore a *partial* answer, with three behaviours worth
-knowing:
-
-- the error **names which arrays** are stale, and `reset_stale_cache=True` deletes exactly
- those arrays' `.npy` files, not the directory;
-- an array with **no entries** is *cold*, not stale — adding a variable to a configuration
- is not a cache wipe;
-- an array **this run does not read** is left entirely alone, entries and files, so two
- configurations can share one `cache_dir`. Dropping a variable is not an edit to it: a
- feature-importance or ablation sweep can run variable subsets over one warm cache all day,
- and every array it omits is still there, valid, when a later run asks for it.
-
-Scope does *not* enter a transform's own token: narrowing `applies(["t2m", "u10"], f)` to
-`applies(["t2m"], f)` invalidates `u10` (it lost a transform) and leaves `t2m` valid (its
-pipeline is unchanged). The log is compacted — written to a temp name and renamed — when a
-load drops entries, which is safe because one writer holds the `cache_dir` lock.
+arrays it names, so each array's entries carry their own 16-hex `pipeline` hash. Staleness
+is therefore a *partial* answer: editing a transform scoped to `t2m` invalidates `t2m` and
+leaves `u10`/`v10` — whose bytes did not change — valid. Which edits invalidate what, the
+two cases that deliberately are *not* stale, and what narrowing a scope does are on the
+[tuning page](tuning.md#what-a-transform-edit-invalidates).
Only the manifest **format** is still all-or-nothing: a log we cannot parse tells us
nothing about any individual array, so an older format resets the whole cache.
@@ -1011,8 +1035,10 @@ isn't working is loud, not silent.
- **Crash-safe:** each *completed* chunk is appended to the log as it finishes (flushed to
the OS page cache, which survives process death), so a killed run keeps everything it
decoded. The log is self-deduplicating (a re-completed chunk across epochs/runs is not
- re-appended) and bounded to one line per cached chunk — no compaction needed. A `.npy`
- left half-written by a crash was never logged, so it's simply overwritten on re-decode.
+ re-appended) and holds one line per cached chunk. It is compacted — written to a temp name
+ and `rename`d — only when a load drops entries, so a reset is not re-read and re-rejected on
+ every later open; that is safe because one writer holds the `cache_dir` lock. A `.npy` left
+ half-written by a crash was never logged, so it's simply overwritten on re-decode.
**What you must ensure (the guarantee boundary):**
@@ -1023,12 +1049,35 @@ isn't working is loud, not silent.
that you bump when the **source data** changes — content/etag drift is deliberately
not detected.
- **`chunk_transforms` must be deterministic** (they are baked into the cached chunk).
- When you edit one, its fingerprint changes and the cache goes stale — construction then
- **raises** until you pass `reset_stale_cache=True` (wipe + rebuild) or delete the dir.
+ When you edit one, the fingerprint of every array it is scoped to changes and those arrays
+ go stale — construction then **raises**, naming them, until you pass
+ `reset_stale_cache=True` (which deletes just those arrays' files and rebuilds them cold) or
+ delete the dir.
For the fingerprint to *notice* a change, install the `cache` extra (cloudpickle) or set
an explicit `transform.cache_key`; the source-only fallback can miss closure/global
changes (and would then wrongly serve the old cache as a hit).
+**One writer per `cache_dir`.** The pool takes an advisory `flock` on the cache directory
+for its lifetime, so a second *writing* process fails fast — with the holder's PID and host —
+instead of quietly disagreeing with the first. It keys on `cache_dir` being set, not on
+`persist=True`, because ephemeral spill writes the same filenames. Chunk files are written
+under a temp name and `rename`d into place, so a mapping another process already holds keeps
+its inode and its data no matter who re-admits the chunk.
+
+`readonly_cache=True` is the workload this is shaped around — one job warms a cache, several
+score against it. It takes the lock **shared**, so any number coexist; it writes nothing; and
+a cache miss **raises**, naming the array and chunk. That is what makes it a contract — *this
+cache is complete for what I am about to read* — rather than a slow path that silently
+re-fetches whatever the warming run's split or transforms left out.
+
+Two consequences worth stating outright. There is no such thing as a stale `flock` — the
+kernel releases it when the process dies, `SIGKILL` and spot preemption included — so there
+is no cleanup procedure, and **deleting the lockfile is actively harmful**: it releases
+nothing and makes the next two processes lock different inodes. And a network `cache_dir`
+(where `flock` may be emulated per client) or a platform without POSIX locking is
+unarbitrated; both warn at construction, because that is the one configuration where two
+writers can still corrupt each other.
+
**Reshaping transforms.** A **reshaping** `chunk_transform` (e.g. `Regrid`, or a dtype
recast) is a first-class cacheable stage on every backing including the persistent mmap
tier. The transform declares its post-transform *inner* geometry via
@@ -1061,6 +1110,10 @@ adds over the store-swap is the **loader**:
- **Read-plan dedup across a request** — ensembles, lead times, and overlapping verification
windows touch the same chunks repeatedly; the plan collapses them to one read each.
- **Prefetch overlap** for sequential/autoregressive rollout, and split + shuffle for training.
+- **Device placement belongs to the adapter** — `to_torch(device=)` owns its H2D copy out of
+ page-locked host memory; `to_jax` places on `jax.devices()[0]`, where every other JAX
+ array-creation path puts one, and takes `device=` to choose another. On a CPU-only build
+ that put is a pass-through that still aliases the exported buffer, so it costs nothing.
A reference integration — an Earth2Studio `DataSource` backed by `InSituDataset` — is in
[emfdavid/earth2studio#1](https://github.com/emfdavid/earth2studio/pull/1).
diff --git a/src/insitubatch/pool.py b/src/insitubatch/pool.py
index f02dfda..6567632 100644
--- a/src/insitubatch/pool.py
+++ b/src/insitubatch/pool.py
@@ -1053,7 +1053,7 @@ def blocked_waiters(self, owner: int | None = None) -> list[tuple[str, int]]:
``owner`` narrows it to one pass, and which of the two a caller wants follows from
what it is waiting on. The byte budget is shared, so an admission stall is rightly
judged against every owner: another iteration's blocked consumer is holding budget
- we need. A fetch-ahead permit is not shared -- it belongs to one
+ we need. A read-ahead permit is not shared -- it belongs to one
:class:`~insitubatch.scheduler.Scheduler` and comes back only from that pass's own
``unpin_block`` -- so a permit stall must ask only about its own waiters
(:meth:`Scheduler._ahead_starvation`).
diff --git a/src/insitubatch/scheduler.py b/src/insitubatch/scheduler.py
index 5bc91ea..0cf2891 100644
--- a/src/insitubatch/scheduler.py
+++ b/src/insitubatch/scheduler.py
@@ -12,11 +12,20 @@
stitch the inner grid under a second cap is what couples them). See
[docs/architecture.md] for the full pipeline.
-Two bounded resources, deliberately distinct:
+Three bounded resources, deliberately distinct -- and they are bounded in different
+units, so no one of them can stand in for another:
-* **in-flight** (``max_inflight``, an ``asyncio.Semaphore``) -- tiles in flight; a
+* **in-flight** (``max_inflight``, an ``asyncio.Semaphore``) -- *tiles* in flight; a
slot is held from fetch-start to delivery, spanning fetch + decode + deliver. It is
also what bounds the decode pool's queue, which is otherwise unbounded.
+* **read-ahead** (``read_ahead_chunks``, an ``asyncio.Semaphore``) -- *chunks* this
+ iteration may hold ahead of its own consumer. One permit is taken per chunk before it
+ is referenced and handed back by that consumer's ``unpin_block``, so a pass runs at most
+ its own working set ahead. Computed by
+ :func:`~insitubatch.source.read_ahead_bound`, never configured. It is what keeps one
+ iteration from claiming the whole pool and starving another sharing it (#64); the byte
+ budget cannot do that job, because a driver far ahead of its own consumer is holding
+ chunks that are legitimately referenced.
* **residency** (the pool's byte budget) -- admission (``pool.try_admit``) evicts
ready-and-unreferenced LRU to make room and *references* (refcounted pin) the chunk,
so it stays resident from in-flight fetch through to the consumer's release; when the
@@ -323,7 +332,7 @@ def __init__(
self._loop = _get_loop()
self._inflight: asyncio.Semaphore | None = None
self._capacity: asyncio.Event | None = None # set on unpin -> wakes a parked admit
- self._ahead: asyncio.Semaphore | None = None # fetch-ahead permits (read_ahead_chunks)
+ self._ahead: asyncio.Semaphore | None = None # read-ahead permits (read_ahead_chunks)
# Outstanding tile tasks per chunk, this scheduler's only. A second block reading a
# chunk we are still fetching does not fetch it again.
self._fetching: dict[tuple[str, int], int] = {}
@@ -466,7 +475,7 @@ def unpin_block(self, keys: set[tuple[str, int]]) -> None:
self._loop.call_soon_threadsafe(self._release_ahead, len(keys))
def _release_ahead(self, n: int) -> None: # on the loop
- """Hand back `n` fetch-ahead permits, one per chunk the consumer released."""
+ """Hand back `n` read-ahead permits, one per chunk the consumer released."""
assert self._ahead is not None
for _ in range(n):
self._ahead.release()
@@ -592,7 +601,7 @@ def _fetch_done(self, key: tuple[str, int], _task: object) -> None:
self._fetching.pop(key, None)
async def _take_ahead(self, array: str, chunk_index: int) -> None:
- """Take a fetch-ahead permit, proving a stall terminal rather than waiting forever.
+ """Take a read-ahead permit, proving a stall terminal rather than waiting forever.
Permits come back from the consumer's ``unpin_block``, so a consumer that cannot
reach its next unpin never returns one -- the same shape as an admission stall, and