From 74c3dd6caef9c450e5070ce39a152243f8c6575f Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Wed, 16 Sep 2026 00:25:25 +0000 Subject: [PATCH] docs: describe the 0.2.0 engine, and fix the probes that measure it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `docs/architecture.md` last moved substantively on 2026-09-05; 26 `src/` commits landed after it, including most of what 0.2.0 changed. The page was not stale but wrong — eight claims a reader sizing a budget would have acted on. Admission is bounded by read-ahead permits, not by the byte budget; peak co-residency is two blocks only on a plain pass; read-once is a per-pass default that release-and-re-read deliberately trades; the cache log does need compaction and a stale cache is not a whole-directory wipe. Both mermaid diagrams were further behind than the prose, which is the part worth noting — nobody greps a diagram. Adds what 0.2.0 shipped and the page had no home for: the permit protocol, one pool with many accountable iteration owners, per-block release, the one-writer `cache_dir` contract, and `readonly_cache`. Unifies "read-ahead permits" across `docs/`, `DESIGN.md`, `scheduler.py` and `pool.py`, where one word had meant both that bound and `max_inflight`. Paid for by deletion: two rationale blocks move to `DESIGN.md` per the docs rule in `CLAUDE.md`, two duplications go. Net +53 lines and one heading fewer. The benchmark probes that produce the page's numbers had two defects and a false premise: - `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, and the probe's own prewarm built it at auto before the sweep began, so every point ran on one pool. On an 8-vCPU box it returned 756 -> 808 -> 921 MB/s across `decode_threads=1,2,8`: noise arranged in the most convincing shape it could take. The documented diagnosis for a flat sweep is "decode isn't the limit", so the failure pointed away from itself. Sweep removed; the header now reports the pool's effective thread count. - The raw-GET ceiling could not be measured on a sharded or a zarr-v2 store — every key 404'd. Keys were hand-built from `metadata.chunks` in v3's `c/i/j/k` layout; on a sharded array that names the inner chunk, not the shard, and on v2 the separator is `.`. Now `ChunkGrid.from_metadata` plus the array's own `encode_chunk_key`: the spelling the engine plans reads with. - Every published "% of raw-GET ceiling" figure mixes units. insitu counts decoded bytes, a raw GET moves stored bytes, and `benchmark_plan.md` licensed dividing one by the other by calling the synthetic stores incompressible. They compress ~8% (1.081x, measured by HEAD), so the figures run ~6 points high, one-directional and flattering. `bench/probe_ceiling.py` measures the ratio per store instead of assuming it, reports "% kept" in one unit, and flags insitu points the permit bound capped below their `max_inflight`. Behavior is unchanged: the only `src/` edits are a module docstring and one comment. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_019waAAauPpmW4QX1VaN4vqP --- CHANGELOG.md | 84 +++++++++++++ DESIGN.md | 21 +++- bench/benchmark_plan.md | 16 ++- bench/probe_ceiling.py | 169 ++++++++++++++++++++++++++ bench/probe_decode.py | 140 +++++++++++++--------- docs/architecture.md | 223 ++++++++++++++++++++++------------- src/insitubatch/pool.py | 2 +- src/insitubatch/scheduler.py | 19 ++- 8 files changed, 520 insertions(+), 154 deletions(-) create mode 100644 bench/probe_ceiling.py 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