Architecture: loaders, parallelism, and prefetch¶

This doc contrasts the classic worker-based data loader with insitubatch's async-driven engine, then specifies the prefetch pipeline, the sample-geometry contract, and how downstream frameworks integrate. For the why behind the project see DESIGN.md.

The inversion in one line¶

Classic DataLoader: parallelism lives in num_workers OS processes, each running a synchronous __getitem__. insitubatch: parallelism lives in one async event loop; batch assembly is the consumer. That move is what unlocks read concurrency across a whole batch, a shared chunk cache, bounded memory, and prefetch overlap.

Classic worker-based loader¶

flowchart TB
    SRC[("cloud zarr")]:::cloud
    subgraph MAIN["main process"]
        SMP["sampler: shuffled indices"] --> IQ["index queue"]
        COL["collate + pin_memory"] --> GPU[("GPU step")]
    end
    subgraph WK["worker processes — num_workers, prefetch_factor"]
        W1["worker 1 · __getitem__ (sync)<br/>read chunk + decode"]
        W2["worker 2 · __getitem__ (sync)<br/>read chunk + decode"]
        WN["worker N · __getitem__ (sync)<br/>read chunk + decode"]
    end
    IQ -->|"fork + IPC"| W1 & W2 & WN
    SRC -->|"each worker re-reads<br/>the SAME chunks · no shared cache"| W1 & W2 & WN
    W1 & W2 & WN -->|"pickle / IPC<br/>(this queue is the prefetch buffer)"| COL
    classDef cloud fill:#cfe8ff,stroke:#379a4a;

Frictions against cloud ndim zarr:

No shared chunk cache — a chunk is fetched + decompressed once per worker whose samples land in it.
Fan-out is one sample deep — a worker can drive async obstore (zarr's sync API runs a process-wide event loop), but __getitem__ returns one sample before the next starts, so concurrency never spans the samples a batch needs, and the budget is per-process, multiplied N times against the store.
dask thread pool nested in each worker — procs × threads oversubscription, slow fork startup, fat memory.
The fork-safety tax — a modern object store (obstore) runs a Rust tokio runtime. fork (the Linux default for workers) copies only the calling thread and leaves that runtime's threads dead with their locks held, so the first read in a forked worker deadlocks. Every escape is a cost the process model imposes: spawn (relaunch the interpreter per worker), forkserver (keep a pristine pre-fork server around), or a stack that rebuilds its loop on a PID change (s3fs/gcsfs) — and even that only if the store is reopened in the worker, never inherited across the fork. An obstore handle opened pre-fork deadlocks; a gcsfs one raises Future attached to a different loop. We hit both: the obstore workers baseline hung under fork, and the gcsfs xbatcher example raised — so fork is off the table for async cloud stores.
Note: the worker model does prefetch (via prefetch_factor — workers run ahead into the IPC result queue). Prefetch is not the differentiator; how we prefetch is.

Why this is the argument for the single loop. Every friction above — no shared cache, per-sample fan-out, thread oversubscription, the fork-safety tax — follows from putting parallelism in OS processes. insitubatch drives one in-process event loop (num_workers=0): there is no fork, so there is no fork-safety tax, no per-worker runtime to relaunch, and the chunk cache and the obstore runtime are simply shared. The deadlock we hit benchmarking the baseline is a symptom of the very thing we replace.

Startup latency — the inference angle¶

Training amortizes worker spin-up over many epochs, so the start-method tax above mostly disappears. Inference does not: you typically make a single pass from a cold loader, and a long-lived server holding a DataLoader open (pinned workers, held file handles) is rare. There, time-to-first-batch is dominated by process startup. The worker model's best case is forkserver with set_forkserver_preload([...]) — heavy imports paid once in the server, forked workers skip them. insitubatch's first batch is just the first read: no processes to start. The two runnable examples make this concrete and measurable: examples/wb2_dataloader.py (insitubatch) and examples/wb2_xbatcher.py (--compare prints TTFB across spawn / forkserver / forkserver-preload), both on the public WeatherBench2 ERA5.

insitubatch async-driven pipeline¶

flowchart TB
    SRC[("cloud zarr · gs:// / s3://")]:::cloud
    subgraph PLAN["planner — per epoch"]
        PERM["chunk permutation + shuffle-block order"] --> RP["read plan:<br/>samples → deduped STORED-chunk (tile) reads"]
    end
    subgraph LOOP["scheduler · single async event loop · one thread"]
        PROD["prefetch producer<br/>looks ahead d batches"] --> SEM["one budget · max_inflight tiles"]
        SEM --> SKIP{"pool already<br/>holds it?"}
        SKIP -->|"hit"| POOL
        SKIP -->|"miss"| OBS["obstore get (stored chunk)"]
        OBS --> DEC["decode · GIL released"]
        DEC --> CT["chunk_transform (vectorized)"]
        CT --> DLV["deliver: adopt the tile (no copy)"]
    end
    DLV --> POOL["ChunkPool — byte budget + pin/LRU<br/>residency AND cache · heap or mmap NVMe"]
    POOL --> GTH["gather + batch_transform"]
    GTH --> PQ["prefetch queue · depth d"]
    PQ --> GPU[("device + device_transform<br/>→ model step")]
    SRC -->|"each stored chunk read ONCE"| OBS
    RP --> PROD
    GPU -.->|"pulls"| PQ
    PROD -.->|"refills ahead of consumption"| PQ
    classDef cloud fill:#cfe8ff,stroke:#379a4a;

The unit of work is the stored chunk (an (outer, inner) tile), and one max_inflight budget spans every fetch — so read concurrency is one dial, with no nested inner/outer caps. Each decoded tile is adopted by reference into its outer chunk's slot in the ChunkPool — the tile is the resident buffer, so the fill path has no memcpy — and the pool is the residency tier and the cache in one: a byte budget with pin/unpin + LRU, backed by heap or an mmap'd .npy on NVMe. Placement is deferred to gather, which writes each stored chunk straight into its rectangle of the batch. Before fetching, the scheduler asks the pool whether it already holds a chunk — a hit (cross-epoch, since the pool persists) skips fetch + decode + transform entirely.

Properties: parallelism in the loop (not processes); each stored chunk read once and amortized across every sample that touches it; read concurrency (max_inflight) and residency (the pool's byte budget) are independent dials; total memory is the budget + the prefetch queue (depth d) + the in-flight tiles — every term a tunable cap, none scaling with batch size or epoch length.

The batch buffers are pooled too, on different machinery: gather takes one from a reuse pool that allocates on a miss and reclaims a buffer once nothing outside the pool references it. Lifetime tracking is what makes that safe against the DLPack export, which aliases the buffer into the consumer's tensor — so retaining a batch retains its buffer automatically and the pool just allocates elsewhere. There is no depth parameter (it converges on however many buffers are genuinely in flight) and a short final batch is a prefix view, costing neither an allocation nor a copy. INSITUBATCH_NO_BUFFER_REUSE is the escape hatch back to one fresh array per batch. Host memory is pageable by default; as_torch(view, device=...) swaps in a page-locked allocator, which is what makes the H2D copy genuinely asynchronous — pinning is coupled to device on purpose, since it is only safe when whoever issues the copy also knows when it landed. See Known limitations in DESIGN.md.

One event loop, one decode pool¶

A genuinely-async fsspec backend (gcsfs, s3fs) binds its aiohttp session to the first event loop that awaits it, permanently — no constructor knob pins it. For a zarr store that first loop is zarr's process-wide store-IO loop (zarr.core.sync._get_loop()), because any zarr sync call (open_geometries, xr.open_zarr, user code) touches the store first. That is the correct owner: the session living on zarr's loop is what keeps the store drivable by any zarr code.

insitubatch runs its orchestration on that loop too. There is no second loop, no per-read bridge, and no per-pass thread churn: gcsfs and obstore are both a plain inline await. What made this possible is that ChunkTransform.decode_chunk (zarr 3.3.0) is synchronous. The async codec pipeline dispatches to the loop's default executor, so keeping our own decode pool used to mean claiming that slot — which on a loop we share retunes zarr's concurrency process-wide, and breaks it when we shut the pool down. Calling a sync decode ourselves means we pass our executor explicitly and never touch the default.

flowchart TD
    subgraph proc["insitubatch process"]
      OG["any zarr sync call<br/>open_geometries / xr.open_zarr"]
      ZL["zarr store-IO loop (zarr_io)<br/>zarr.core.sync._get_loop()<br/><b>scheduler orchestration runs here</b>"]
      DP["insitu-dec pool<br/>process-wide, N threads<br/>decode (GIL released)"]
      SESS[("gcsfs aiohttp session")]
      OBS["obstore ObjectStore<br/>Rust tokio runtime, loop-agnostic"]
    end
    OG -->|"first await creates the session here"| ZL
    ZL -->|owns| SESS
    ZL -->|"fsspec read: await inline (no bridge)"| SESS
    ZL -->|"obstore read: await inline"| OBS
    ZL -->|"run_in_executor(decode_pool, decode_chunk)"| DP
    DP -->|"NDBuffer handed back by reference"| ZL

The handoff, per stored tile¶

loop.run_in_executor(pool, fn, *args) submits to the pool and wraps the resulting future; the coroutine awaits it and suspends, so the loop is free to service other tiles, other in-flight reads, and everyone else's zarr work. The worker resolves the future via call_soon_threadsafe. What crosses the boundary is a Python object reference — no copy, no pickling; it is a thread boundary, not a process one.

insitu-prefetch          zarr_io  (ONE loop, shared with all zarr sync work)   insitu-dec (N)
     │                       │                                                      │
     │                  _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    │      │
     │                       │◄───────────────────────── Buffer ─────────────┘      │
     │                       │                                                      │
     │                  ② await run_in_executor(pool, decode_chunk, buf, spec)       │
     │                       │──── submit ─────────────────────────────────────────►│
     │                       │   coroutine SUSPENDS                   decode_chunk   │
     │                       │                                   (GIL RELEASED in    │
     │                       │                                    numcodecs)         │
     │                       │◄── call_soon_threadsafe ── NDBuffer ─────────────────│
     │                       │                                                      │
     │                  np.moveaxis(tile, ax, 0)   [a VIEW, no copy]                 │
     │                       │                                                      │
     │                  ③ deliver:                                                   │
     │                       │  tiled     → w.deliver(tile) INLINE  [dict + _cv]     │
     │                       │  assembling→ await run_in_executor(...) ─────────────►│
     │                       │                              _advance → assemble →    │
     │                       │                              chunk_transform → persist│
     │                       │◄────────────────── done ─────────────────────────────│
     │◄── _cv.notify_all() wakes wait_ready ──┘
   gather

Delivery is inline on the tiled path because it is a dict write and a counter. It becomes an executor hop when the pool assembles, because delivering the last tile also runs the assembly memcpy, the user chunk_transform and the mmap write-back — and none of that 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. (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¶

The pool speaks zarr's own vocabulary for "a grid of inner chunks", so the structure stays legible to zarr rather than being a private dict of ndarrays:

ours	zarr
`ArrayGeometry.tile_placement`	returns a `zarr.core.indexing.ChunkProjection` — `chunk_coords`, `chunk_selection`, `out_selection`, `is_complete_chunk`
the sync decode	`zarr.core.chunk_utils.ChunkTransform.decode_chunk` — so we do not maintain a codec whitelist
`ArrayGeometry` inner grid	conforms to `DimensionGridLike` (proven in `tests/test_zarr_indexing_parity.py`)

A _Slot deliberately has no zarr counterpart. It is a residency unit — the stored chunks sharing one sample-axis index — whereas a zarr shard is a storage unit, one addressable object holding a grid of chunks. They have the same shape, which makes the analogy tempting and wrong: a slot is never stored or addressed as one object. What this codebase calls a tile is what its public API calls a stored chunk (one zarr chunk, or one zarr shard on a sharded array — whatever a single store.get returns); that duplication is tracked in #40, not fixed here.

On a sharded array the stored chunk is the shard, and the distinction is load-bearing: zarr reports chunks as the inner chunk (read granularity inside a shard) and chunk_grid.chunk_shape as the shard itself, which is what a chunk key addresses. The geometry and the read plan use the latter, so a sharded array is read one whole shard at a time. That is a good fit where a run consumes the samples a shard holds and a poor one where it wants a small spatial subset of a large shard — fetching only the inner chunks a run needs is #57.

Two zarr facilities are deliberately not adopted. FusedCodecPipeline.read_sync takes ByteGetters and therefore owns the IO, which would surrender the scheduler, max_inflight and back-pressure; decode_chunk is IO-free, which is why it is the one we use. And the process-global codec_pipeline.path is never flipped — that retunes the substrate under user code — so we build our own ChunkTransform from the array's declared codecs, which is also what keeps zarr-v2 (V2Codec(filters, compressor)) working.

A scheduler tears down only what it created. It cancels the tasks it started (never asyncio.all_tasks, which is the whole loop's set), and it does not stop the loop, close it, or shut down the decode pool. Each of those is correct on a private loop and destructive on a shared one — cancelling the task set raises CancelledError inside unrelated zarr reads, stopping the loop hangs every later zarr.core.sync.sync() call, and shutting the shared pool down makes the next scheduler raise cannot schedule new futures after shutdown. All three were observed while designing this, are pinned by tests/test_loop_ownership.py, and are why the acceptance gate for sharing a loop is a close() test rather than a steady-state one.

Measured: with a consumer-stalled, back-pressured _drive, unrelated zarr-sync work completes ~5,700 reads during the stall at p50 1.69–1.77 ms, against 1.74–1.80 ms for a private per-pass loop — no starvation, no deadlock, and the victim keeps running through teardown. Teardown also closes the fsspec session on its own loop (close_store), so gcsfs's finalizer — which wrongly targets fsspec.asyn.get_loop() — becomes a no-op.

Sample geometry — the axis-role contract¶

This is the stable contract the engine commits to, and the extension points reserved for it. (For why it evolved this way — and why the tempting generalizations are deferred — see the "Sample geometry" entry in DESIGN.md.)

Every physical axis of a zarr array has one role:

sample axis — the axis the engine draws / shuffles / splits / windows along. Exactly one per variable. ArrayGeometry.sample_axis names it as a physical axis index (default 0: time for ERA5/HRRR; 2 for the Z of an OME-NGFF (T,C,Z,Y,X) microscopy stack sampled slice-by-slice — open_geometries(store, sample_axis=2)). A sample is a slice along this axis that does not cross a chunk boundary.
field axis — every other axis, carried whole as the per-sample payload (inner_shape). Field axes may be inner-chunked (the ARCO/ERA5 norm); the engine 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 key via physical_chunk_coord; moveaxis the decoded tile sample-first). With sample_axis=0 both are identities, so the common path is untouched.

Windowing is orthogonal to axis choice: a variable is a (label, path, offset) view, and g.shift(k) reads array[anchor + offset] along the sample axis. Several views of one array (same path, different offset) decode once and share slots — the forecasting {"x": g, "y": g.shift(horizon)} pattern, which composes with any sample_axis.

The frozen promises (an API a v0.1 user can build against):

Batch.arrays[label] is always (batch, *sample_inner) — a leading batch axis plus the per-sample payload. Future capabilities change what sample_inner resolves to, never this shape.
The (label, path, view) variable model, decode-once dedup on shared path, and the window helpers (Batch.sample_indices / offsets / read_indices / stack).
Framework adapters infer tensor shape from view.geometries (post-transform inner_shape/dtype).

Reserved extension points (committed as additive, not yet built):

Multiple sample axes (a product index — HCS Well×Field×Time, Year×Day). sample_axis: int widens to int | tuple[int, ...] non-breakingly (a scalar still means one axis). Deferred: a tuple trades away the "sample is a contiguous range in one axis" invariant and pulls in the full product-draw-space refactor — see DESIGN.md.
Patch / sliding-window sampling (crop a giant field into (64,64) tiles). A field axis gains an extent (declarative geometry); the crop origins are a sampler policy (a strided grid, or a random draw — random crop already works today as a batch_transform over the whole field). Memory-optimal patching (read only the tiles a crop touches) is a pool/residency change, not a Batch-contract change.
Cross-node concatenation (many well-arrays into one stream) is a catalog layer above the core, not part of this contract.

Use-case support¶

Concrete cross-domain use cases against the contract above — what works today, what is a reserved (additive) extension, and what is deliberately out of scope.

Supported now:

Use case	Example	How
Train over a time/sample axis	ERA5/HRRR reanalysis; astronomy light curves	`sample_axis=0` (default)
Sample over any single axis	OME-NGFF 2D segmentation over `Z`; `T`-frame or volumetric stacks	`open_geometries(store, sample_axis=2)`
Paired inputs+targets with different sample-axis chunking	microscopy raw image + label mask (Z-chunk 1 vs 30)	two variables, one manifest
Multiple co-registered variables at one anchor (data + mask + weight)	radio astronomy MSv4 visibilities: `VISIBILITY` + `FLAG` + `WEIGHT`, sampled over `time`	variables share one manifest
Forecasting input/target windows	weather: input@`t`, target@`t+h`	`g.shift(h)` views (decode-once)
Whole inner field per sample, spatial grid fetched decode-once	ARCO/ERA5 `721×1440` fields; microscopy planes	inner (field) chunking
Cross-variable derived fields & per-sample random augmentation	`windspeed=√(u²+v²)`; random crop of the whole field	`batch_transform`
Chunk-aligned splits, block shuffle, cross-run cache	all domains	shipped

Reserved (committed as additive, not yet built):

Use case	Example	Extension
Multiple sample axes (product index)	HCS `Well×Field×Time`; `Year×Day`	`sample_axis` widens `int → int \\| tuple`
Native spatial patch / sliding window (memory-optimal crops)	ERA5 `64×64` patches for a ViT; radio-astronomy cubes	patch extent (geometry) + origin sampler + partial-field residency
GPU `device_transform` stage	any	milestone M2

Out of scope (by design):

Not supported	Example	Why
Cross-chunk / cross-sample-boundary stencils	finite differences across a time-chunk seam	a sample cannot span a chunk boundary
Cross-node concatenation into one stream	many HCS well-arrays as one dataset	catalog layer above the core
General compute graph / cross-chunk reductions	lazy dask-style evaluation	dask is off the hot path by design; reductions run over the loader
Resharding into a sample-per-file format	MDS/tar/WebDataset ETL	we train in place, no reshard
Scattered / boolean sample selection for splits	arbitrary index masks	splits are chunk-contiguous (leakage-safe)

Mapping a new domain onto the contract — worked example: radio astronomy (MSv4). The xradio MeasurementSet-v4 visibility dataset stores VISIBILITY (time, baseline_id, frequency, polarization) (complex) alongside FLAG (same shape, boolean) and WEIGHT. This is the same shape as the shipped microscopy case: pick the sample axis (time for per-integration samples, or frequency/baseline_id), carry the rest as the inner field, and gather VISIBILITY + FLAG + WEIGHT at one anchor from one manifest — no reshard. The natural task, RFI-flag detection (predict FLAG from |VISIBILITY|), is the segmentation pattern with an amplitude-threshold baseline in place of Otsu. Complex data needs no special support: take np.abs() in a batch_transform (or at the model boundary), so the numpy Batch and the DLPack adapters only ever see real arrays. This is a schema-level mapping, not yet a shipped example — but every capability it needs is validated today.

Prefetch¶

source.InSituDataset.__iter__ runs a background producer thread that assembles batches ahead of the consumer:

✅ Intra-batch concurrency — a batch's missing chunks are fetched concurrently via the async loop (stored-chunk fan-out under max_inflight).
✅ Inter-batch overlap — the producer assembles batches N+1..N+depth while the caller works on batch N; the consumer just drains a bounded queue. (A demand-driven loop would leave the event loop idle during the compute step.)

Design (producer/consumer pipeline)¶

flowchart LR
    subgraph PRODUCER["producer thread"]
        WALK["walk draw order"] --> ASM["assemble batch<br/>(async read + gather + transforms)"]
    end
    ASM --> Q["bounded queue · maxsize d"]
    Q --> CONS["consumer __iter__<br/>pops finished batches"]
    CONS --> STEP[("train / infer step")]
    Q -.->|"full ⇒ producer blocks<br/>(backpressure)"| ASM
    CONS -.->|"unpin_block ⇒ read-ahead permit returns<br/>(bounds how far ahead)"| ASM

Producer starts the scheduler over the epoch's chunks, then walks the epoch's batches — waiting each block a batch draws from, gathering, and unpinning a block once a batch has consumed its last row — pushing batches to a bounded 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 — 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 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. 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 DataLoaders — 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.

Same shape as torchdata.nodes.Prefetcher, but async-native. This is what turns a throughput win into a GPU-fed win.

Trade-offs: chunk size, shuffle window, concurrency, batch size¶

Four dials shape throughput and memory, and the engine keeps them as independent as possible so you can move one without paying on the others. This is the model behind them; for the values to actually set, see Tuning.

dial	what it trades	bounded by
stored-chunk size (`inner_chunks`, write time)	fetch granularity ↔ per-request overhead	a tile of ~10–50 MB
read concurrency (`max_inflight`)	network saturation ↔ in-flight memory	the store's raw-GET knee
shuffle window (`block_chunks`)	shuffle quality ↔ resident memory	RAM / cache budget
batch size (`batch_size`)	the model's step size	the window's sample pool

Chunk size is the amortization lever. insitu reads each stored chunk once and gathers every sample inside it, so the work saved versus a per-sample __getitem__ grows with samples-per-chunk. Fat chunks amortize more; the one-sample-per-chunk (GRIB) end has nothing to amortize. Chunk size also sets the memory unit — residency is counted in whole outer chunks.

Stored-chunk size decouples concurrency cost from chunk size. When a chunk is split into an inner grid of tiles, a read fetches a tile, not the whole chunk, so raising max_inflight costs tile-sized memory, not chunk-sized. That is why "fat, spatial" is the sweet spot and "fat, single inner" is not: with one tile per chunk the two collapse and concurrency costs full chunks. The decoupled scheduler is what makes read concurrency and residency independent dials in the first place.

Batch size is largely orthogonal to IO. A batch is a vectorized gather from the resident window, so batch size sets the step the model sees, not the read pattern — as long as the window's pool (block_chunks × samples-per-chunk) stays well above it.

Why the block-local shuffle is enough¶

The shuffle is approximate, not global: chunks are permuted each epoch, and within a window of block_chunks chunks all samples are shuffled together (the scope boundary on exact shuffle explains why a global shuffle is incompatible with chunk-aligned, low-copy reads). Two things make that converge to a full shuffle in practice:

Within an epoch, each batch is a uniform draw from the window's pool of block_chunks × samples-per-chunk samples. Keep that pool well above batch_size and a single batch is already well-mixed locally.
Across epochs, the per-epoch chunk permutation re-randomizes which chunks share a window, so any two samples' chunks eventually co-occur. Over a run the set of samples a given sample is ever batched with approaches the whole dataset.

So even a modest window asymptotes quickly toward a global shuffle over the many epochs training actually runs — at memory cost O(block_chunks), not O(dataset). shuffle_quality scores an emitted order against a perfect global shuffle if you want to see it on your own data. (shuffle=False swaps in a sequential order — chunks and samples in order — for eval / inference / reconstruction.)

Read-once and sample-once¶

The pipeline holds two guarantees, and they are orthogonal — batch size touches neither:

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). 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).
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 drawn sample. It is a permutation of every valid anchor, so each sample appears once. The sample-level fancy index is never stored — gather recomputes it per batch from the rows (anchor = chunk_id·spc + within; sample = anchor + offset; then, per unique read chunk, one coalesced write per stored tile into its sub-rectangle of the batch), Python work O(chunks-in-batch × tiles-per-chunk), never O(samples).

The figure traces one epoch, archive → batch. Colour tracks the bytes — a chunk keeps its colour from the archive ① through the resident pool ③; the read plan ② in between is metadata (tile keys in shuffle-block order, not bytes), so the plan (dashed) stays distinct from the movement (thick). In stage ④ a batch is a row of colour-coded slices drawn across several resident chunks: one chunk feeds many batches yet is never broken up — it stays whole and reused, which is what read-once buys. (Loaders that pre-shuffle to disk instead dissolve the chunk into scattered rows; see the annbatch contrast in DESIGN.md.)

flowchart TB
    subgraph S1["① immutable archive — chunks along the sample axis"]
        direction LR
        A0["chunk 0"]:::cA
        A1["chunk 1"]:::cB
        A2["chunk 2"]:::cC
        A3["chunk 3"]:::cD
    end
    subgraph S2["② read plan — deduped, shuffle-block order"]
        direction LR
        P["what to fetch, in what order · read 2 · 0 · 3 · 1 · … · <i>tile keys — not bytes</i>"]:::plan
    end
    subgraph S3["③ ChunkPool — block_chunks tiles resident &amp; intact"]
        direction LR
        R2["chunk 2"]:::cC
        R0["chunk 0"]:::cA
        R3["chunk 3"]:::cD
    end
    subgraph S4["④ windowed anchor draw — each batch interleaves slices from several resident chunks (which stay whole, reused)"]
        direction TB
        subgraph B1["batch #1"]
            direction LR
            b1a["c2"]:::cC
            b1b["c0"]:::cA
            b1c["c3"]:::cD
        end
        subgraph B2["batch #2"]
            direction LR
            b2a["c0"]:::cA
            b2b["c3"]:::cD
            b2c["c2"]:::cC
        end
        Bn["… batch #n"]:::batch
    end
    S1 -.->|"plan the reads<br/>dedup + draw order"| S2
    S2 ==>|"scheduler fetches + decodes<br/>one per tile — the movement"| S3
    S3 -->|"gather by (chunk_id, within)<br/>draw rows"| S4
    classDef cA fill:#cfe0f7,stroke:#5b7aa6;
    classDef cB fill:#f7cfd8,stroke:#a65b6e;
    classDef cC fill:#f7e0cf,stroke:#a6805b;
    classDef cD fill:#d3f0d8,stroke:#5ba66e;
    classDef plan fill:#ffffff,stroke:#666,stroke-dasharray:5 3,color:#333;
    classDef batch fill:#ededed,stroke:#888;

"Exactly once" means every valid anchor. order is built to handle the edges: a short final chunk emits within only up to its real length, and windowed sampling drops anchors whose anchor+offset would read off the array (incomplete windows — correct to skip for training; the inference path validates the range and raises instead). Guaranteed by test_order_covers_every_sample_exactly_once, test_order_handles_partial_final_chunk, and the decode-once suite.

The tail is ragged once per epoch. Batches are cut over the whole epoch order, so 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. 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.

There is deliberately no drop_last. Dropping the epoch's short tail is the caller's choice and len(batch) is the whole implementation:

for batch in ds.train:
    if len(batch) < ds.batch_size:
        continue

Transforms — three stages, placed by cost¶

Models need preprocessing (at minimum scaling; often regridding). The interesting question is where a transform runs, because placement is a performance lever tied to the core principle (Python work scales with chunks, not samples).

read → decode ─►[chunk_transform]─► buffer → gather ─►[batch_transform]─► DLPack ─►[device_transform]─► model
                   O(chunks)                              O(batches)                    O(batches), on GPU
                   amortized over every                   needs the                     cheap-on-device
                   sample in the chunk                    assembled batch               ops

chunk_transform(DecodedChunk) -> DecodedChunk — per-chunk, on the decode thread pool, before shuffle/gather. Amortized over every sample that draws from the chunk. Home for per-element, sample-order-independent ops: scaling / normalization, unit conversion, dtype cast, chunk-local regrid. Sees one variable, one chunk. Restrict it to particular variables with applies — see Scoping a transform to variables.
batch_transform(Batch) -> Batch — per-batch, after gather. For ops that need the assembled batch: cross-variable derived fields, channel stacking, per-sample random augmentation/crops, collation to model layout.
device_transform — in the framework adapter, after DLPack, on-GPU, overlapping compute. For ops cheap on device (GPU normalization, batched interpolation, FFTs).

Placement principle: push each transform as early and as shared (per-chunk) as possible; move later only when it needs the batch, per-sample randomness, or is cheaper on-device. A per-sample transform in __getitem__ (the torch way) redoes work for every reused sample — we refuse that by default.

Free advantage: parallelism is in one event loop, not worker processes, so transforms need not be picklable — stateful normalizers, closures, GPU objects all work. torch's DataLoader forces picklable transforms across fork.

Runnable side-by-side example: examples/transforms.py — a Kelvin→Celsius chunk_transform (one variable, cached) and a cross-variable windspeed batch_transform (needs the assembled batch, uncached).

Scoping a transform to variables¶

chunk_transforms is one ordered list and a bare transform runs on every variable. To restrict one, wrap it at the call site with applies:

from insitubatch import applies

InSituDataset(..., chunk_transforms=[
    applies(["2m_temperature"], kelvin_to_celsius),   # this variable only
    Coarsen(2),                                       # every variable
])

The names are zarr array paths (what chunk.read.array carries), not dict labels — two labels may alias one array (t2m_now / t2m_next), and the cache is keyed by path. A name matching no array raises at construction.

Limitation: aliased labels share one chunk pipeline. Because scope is per array path, two labels backed by the same array cannot be given different chunk_transforms — there is no way to express "scale t2m_now but not t2m_next" at the chunk stage. This is the de-duplication the pool exists for: aliases are one decoded, transformed chunk read at two sample offsets, and two pipelines over it would mean two transformed copies of the same bytes and two cache entries under one key. Per-label work belongs in a batch_transform, which sees the assembled batch and each label separately (uncached, which is the honest price). The same holds for any two labels pointing at one path, aliased or not.

Do not gate inside the transform. An if chunk.read.array == "t2m" in the body is invisible to the engine, and that is not a style point — it is a correctness bug for any reshaping transform. output_geometry folds every transform's declared output_inner into every variable's geometry; it cannot see a name test in a function body. A transform that halves t2m and no-ops on u10 therefore makes the engine believe u10 is halved 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: 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 the cache fingerprint.

Standard scaler — pre-fit GLOBAL stats (not per-chunk)¶

@dataclass
class StandardScaler:
    """Global per-variable (optionally per-level) standardization with PRE-FIT,
    FIXED stats. Applied identically to every chunk — never recomputed per chunk."""
    mean: dict[str, np.ndarray]   # per var, shaped to broadcast: surface (1,1); per-level (level,1,1)
    std:  dict[str, np.ndarray]
    eps: float = 1e-8
    def __call__(self, chunk: DecodedChunk) -> DecodedChunk:
        m, s = self.mean[chunk.read.array], self.std[chunk.read.array]
        chunk.data = (chunk.data - m) / (s + self.eps)
        return chunk

Pre-fit the stats however you like and pass them in. The recommended way is to fit over the loader itself with scikit-learn's incremental StandardScaler.partial_fit — covered next — which also warms the cache. StandardScaler above is then the chunk-stage applier, for when you want the normalization cached with the decoded chunk; the fit pass and the apply stage are independent.

A re-fitted scaler needs a cache_key (or cloudpickle)

Stats large enough for numpy to summarize in repr — anything over 1000 elements, so per-gridpoint mean/std — are invisible to the best-effort source fingerprint, and a re-fit reopens a persisted cache as a hit carrying the old normalization. Install insitubatch[cache], or pass StandardScaler(..., cache_key=...) with a stats version you bump on every fit. See the fingerprint. Fitting at the batch stage (below) sidesteps this entirely: the cache then holds raw chunks.

Alternative: fit at the batch stage with community tooling, warming the cache. Standardization is elementwise, so per-chunk and per-batch are identical — which means you can also fit it over the loader itself: iterate once with no scaler (decoding + caching the raw chunks) while a sklearn.preprocessing.StandardScaler.partial_fit (or dask_ml) accumulates per-variable stats, then attach the fitted scaler as a batch_transform. The cache then holds raw chunks — normalization-agnostic and reusable across experiments — and the fit pass is the warm-up; training reads decode-once. It also composes cleanly with a preceding chunk_transform (a regrid), since the fit sees the chunk stage's output. Runnable: examples/fit_scaler.py.

Regrid — precomputed weights, placement by regime¶

@dataclass
class Regrid:
    """Bilinear lat/lon → target grid. Chunk-local (spatial dims whole per chunk).
    Weights computed ONCE; apply is a vectorized sparse gather. inner_shape changes
    consistently across chunks."""
    src_lat, src_lon, dst_lat, dst_lon: np.ndarray
    def __post_init__(self):
        self._idx, self._w = _bilinear_weights(self.src_lat, self.src_lon,
                                               self.dst_lat, self.dst_lon)
    def __call__(self, chunk: DecodedChunk) -> DecodedChunk:
        chunk.data = _apply_weights(chunk.data, self._idx, self._w)
        return chunk
    def output_inner(self, geom: ArrayGeometry) -> tuple[tuple[int, ...], np.dtype]:
        return (len(self.dst_lat), len(self.dst_lon)), geom.dtype  # new inner shape

A reshaping chunk_transform declares output_inner so the cache can size its slot at the post-transform shape (see the caching section); a shape/dtype-preserving one (scaling) omits it. Only the inner dims and dtype may change — the sample axis is the engine's to keep.

Fat chunks → chunk_transform (amortized over the chunk's samples).
ARCO chunk-1 → reuse the same weights as a sparse tensor in a device_transform (batched on GPU), since per-chunk == per-sample there.

Developing a transform — `check_transform`¶

Both transform contracts (vectorized/GIL-releasing, and a correct output_inner) are easy to get wrong and otherwise only surface at training time. Run the bundled CLI against one chunk of your real store while you write the transform:

# a runnable example against the public WeatherBench2 ERA5, checking the bundled K->C transform:
insitubatch-check-transform \
    gs://weatherbench2/datasets/era5/1959-2022-6h-128x64_equiangular_with_poles_conservative.zarr \
    --var 2m_temperature --transform examples/transforms.py:kelvin_to_celsius --skip-signature
# your own:  insitubatch-check-transform s3://bucket/era5.zarr --var t2m --transform ./prep.py:Regrid
#            (target is module:attr or path.py:attr; a transform class is instantiated)

It prints the chunk geometry (the decoded MB the cache will hold), runs the transform on real data, validates a declared output_inner against the actual output (catching the mismatch ChunkPool._persist would later raise, or a reshape that forgot to declare one), and runs a thread-scaling probe that flags a transform holding the GIL — a pure-Python per-element transform that would serialize the decode pool. Non-zero exit on a failed check, so it can gate a pre-commit hook. --no-gil-probe does the geometry + cacheability checks only (fast, deterministic); the GIL probe wants a realistically-sized chunk (a toy array is dominated by per-call overhead). --skip-signature / --request-payer read public / Requester-Pays stores.

Bad / corrupt chunks¶

Real archives — especially GRIB-under-zarr (HRRR) — ship the occasional truncated or corrupt stored chunk. By default a decode failure fails fast (on_bad_chunk="raise"). Set on_bad_chunk="nan" and a bad (or missing) tile is filled with NaN (float dtypes) or the array's fill value instead of poisoning the epoch — the outer chunk assembles with a hole that you repair with an ordinary chunk_transform:

def fill_nan(chunk):                      # your policy: climatology, interpolate, ...
    np.nan_to_num(chunk.data, copy=False, nan=0.0)
    return chunk

ds = InSituDataset(store, manifest, on_bad_chunk="nan", chunk_transforms=[fill_nan])
for batch in ds.train:
    ...
print(ds.bad_chunks)   # the (array, chunk_index, inner_coord) reads that were bad this epoch

Granularity is the stored chunk (tile), so one corrupt inner tile NaNs only its region of the outer chunk, not the whole field. A failure during delivery (a genuine bug, not a bad chunk) still poisons — the policy only covers fetch/decode. Dropping NaN-containing samples is deliberately not automatic (it would break the fixed-shape vectorized gather); exclude known-bad chunks at the split/manifest level instead (ds.bad_chunks gives you the list to quarantine).

Splits — chunk-aligned and leakage-safe¶

Train/val/test are partitioned ahead of time, at chunk granularity along the sample axis — split_by_chunk assigns whole chunks to each split, never individual samples. Two reasons:

No leakage. A sample never straddles a split boundary, and the temporally adjacent, autocorrelated samples inside a chunk can't land on opposite sides of train/val.
Reads stay chunk-aligned. Every read serves exactly one split — no half-chunk waste, no sample shared between two splits.

The result is a SplitManifest (which chunk indices belong to each split), persisted as JSON for reproducibility; the dataset's views (ds.train / ds.val / ds.test) read from it.

contiguous is the decision to get right. By default (contiguous=True) each split is a contiguous block of chunks — the safe choice for time series, where a randomly interleaved split still leaks through autocorrelation across chunk boundaries (a val chunk wedged between two train chunks shares its neighbours' weather). Set contiguous=False only when samples are exchangeable (independent scenes); it shuffles chunks before partitioning.

from insitubatch import obstore_store, open_geometries, split_by_chunk

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))
# independent scenes: shuffle chunks before splitting
manifest = split_by_chunk(geom, fractions=(0.8, 0.1, 0.1), contiguous=False)

fractions are fractions of chunks, not samples (for many chunks the two converge). See split_by_chunk in the API reference and "Splits" in DESIGN.md for the rationale.

Subsetting to a window — define it with xarray¶

The SplitManifest records which sample-axis chunks belong to each split, so to train on a window of a long archive you just restrict the manifest. split_by_chunk takes a sample_range=(start, stop) of sample indices and keeps the chunks overlapping it before partitioning. And because you probably think in time, not indices, you can define the window with the xarray API you already know and translate it — xarray is used only for this off-hot-path planning step; the engine itself never touches xarray:

import xarray as xr
from insitubatch import InSituDataset, obstore_store, open_geometries, split_by_chunk

store = obstore_store(url)                 # one Store, reused for planning + the engine

# define the window in xarray
xds = xr.open_zarr(store)
sel = xds.sel(time=slice("2020-01-01", "2021-01-01"))
times = xds.indexes["time"]
i0 = times.get_loc(sel.time.values[0])
i1 = times.get_loc(sel.time.values[-1]) + 1  # half-open

# pure zarr/numpy from here
geom = open_geometries(store)[var]
manifest = split_by_chunk(geom, fractions=(0.8, 0.1, 0.1), sample_range=(i0, i1))
ds = InSituDataset(store, manifest, ...)

Limitation — chunk-aligned and contiguous. The selection snaps outward to chunk boundaries: a window that starts or ends mid-chunk pulls in that whole partial edge chunk, because splits are chunk-granular (you subset whole chunks, never individual samples). For day/hour chunks against a multi-month window that's effectively exact. It is only for a single contiguous window — scattered/boolean selections (e.g. "summers only") don't map cleanly, since straddling chunks would silently re-add the samples you meant to drop. Subsetting the inner (spatial) axes is a separate, later feature.

The caching continuum¶

The cache boundary IS the chunk-transform boundary. Every chunk is keyed (array, chunk_index), and chunk_transforms are deterministic and applied before shuffle — so what's worth keeping is the decoded + scaled + regridded array, and a hit skips fetch and decode and normalize and regrid (not just bytes). batch_transforms (per-sample / random, post-shuffle) run after and are never cached. So chunk_transforms are exactly the deterministic prefix safe to persist.

Dedup → buffer → cache is one continuum — and in the engine it is literally one 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 — `working_set_bytes`)
cache	across epochs	a large budget retains drained chunks

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 keeps the view — so a hit needs no copy out of a separate cache either way. mmap makes the footprint reclaimable kernel page cache, bounded on disk by bytes, so the working set stays bounded. Caching the prepped representation is strictly stronger than a raw-byte NVMe cache for an ML pipeline. The default budget is the working set (read-once); raise cache_budget_bytes to cache.

Heavy-reuse tasks unlocked: multi-epoch training (epoch 0 warms it); the fat-chunk regime (one chunk → many batches); scoring/verification (reference chunks reused across metrics, lead times, models); datasets that fit in RAM/NVMe (effectively in-memory at GPU-fed speed after the first pass).

Cross-run persistence¶

persist=True (with a cache_dir) turns the mmap tier into a cache that survives process exit: ready slot files are kept on close, an append-only insitu_cache.jsonl log records each completed entry the moment it lands, and a fresh dataset over the same cache_dir revives them as ready hits on first touch — the fetch driver's existing hit path skips fetch and decode and transform, so a warm run re-decodes zero chunks. Because the log is written incrementally (not at close), a killed process — spot preemption, OOM, SIGTERM — still leaves a usable cache: the next run re-decodes only the chunks that hadn't finished. Without persist, cache_dir is ephemeral spill (files unlinked on close).

ds = InSituDataset(store, manifest, cache_dir="/mnt/nvme/era5/v3",
                   persist=True, cache_budget_bytes=...)

Two guards keep a reopened cache honest:

a chunk-transform fingerprint stamped on every log entry, over the transforms scoped to that array (identity is carried on the entry, not in the header; the design record has the why) — change your chunk_transforms (or bump the log format) and the arrays they affect are stale. A stale cache is almost never what you intended, so by default it raises at construction, naming which arrays, rather than silently rebuilding or serving stale data. Pass reset_stale_cache=True to opt into deleting just those arrays' files and rebuilding them cold (or delete the cache_dir yourself); and
a per-entry shape/dtype check on revive — a chunk whose stored geometry no longer matches the current array is a miss (re-fetch + overwrite), never an error.

(Corruption or tampering — an unreadable header, a malformed interior entry, or a file that isn't a bare basename — always raises, regardless of reset_stale_cache; that flag governs an expected stale cache, not a damaged one.)

The fingerprint resolves each transform by, in order: an explicit transform.cache_key (authoritative), a cloudpickle hash (the optional cache extra — captures closures and referenced globals, so a changed closed-over constant invalidates), or a best-effort source hash (catches an edited body but not a changed closure/global — it warns once so the weaker guarantee is visible).

The source fallback is blind to large statistics. It folds a class-based transform's repr in so instance config still counts, but numpy summarizes an array over 1000 elements in repr — array([0., 1., 2., ..., 1997., 1998., 1999.]) — so two transforms carrying different per-gridpoint arrays repr identically and hash the same. Re-fit a StandardScaler over a (721, 1440) grid, reopen a persisted cache without cloudpickle, and every chunk revives as a hit, normalized with the old statistics: no error, no miss, wrong numbers. Install insitubatch[cache] (cloudpickle hashes the values), or set pass StandardScaler(..., cache_key="<stats version>"), whenever a transform's identity 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. 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.

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.

Observability. dataset.cache_hits / cache_misses give per-epoch counts; a plain miss is silent, but persist=True serving zero hits while the cache was consulted-and-rejected emits one WARNING per epoch — a cache you configured but that isn't working is loud, not silent.

Invariants the engine guarantees:

The cache key is (array_path, global chunk_index) — the absolute zarr chunk index, not relative to a split or sample_range. So overlapping subsets/splits and a later fuller run share entries: chunk 5 is always the same .npy.
A hit returns the prepped chunk (post-chunk_transform): no fetch, no decode, no transform. gather copies its rows straight into the batch, with nothing staged in between.
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 holds one line per cached chunk. It is compacted — written to a temp name and renamed — 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):

The cache_dir path is the dataset identity. The store URL is not in the key (an Icechunk/Arraylake session store has no round-trippable URL), so pointing one cache_dir at two sources that share array paths + geometry would serve the wrong chunks. Use a distinct dir per dataset and bury a version in the path (…/era5/v3/) 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, 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 renamed 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 output_inner(geom) -> (inner_shape, dtype); the engine sizes the cache slot at that output shape and assembles decoded tiles into a transient source-shaped scratch buffer, then writes the transformed result into the slot (the sample axis is spliced back from the source, so a chunk_transform can never move it). Shape/dtype-preserving transforms omit output_inner and keep the zero-copy fast path (the slot is buffer and cache in one).

Limitations (deferred). A raw-decoded tier (keyed by source only, for transform experimentation) and the kvikio/GDS NVMe→GPU feed off the persistent .npy tier remain on the roadmap (DESIGN.md).

Downstream integration — tensors, not xarray¶

insitubatch delivers tensor batches and never builds xr.DataArray. That decides how it plugs into an xarray-native inference framework such as Earth2Studio: around the models, not inside their xarray loop. A raw ObjectStore(obstore) store-swap in such a framework's ARCO source buys faster bytes — but that is an obstore win, not ours. insitubatch earns its place for batched training / fine-tuning / hindcast scoring: read the ARCO/zarr store → DLPack → (torch.Tensor, coords) and drive the model directly (the coords we supply is a light OrderedDict of coordinate arrays — metadata, not the xarray machinery), instead of routing every (time, variable) through the framework's per-request xarray DataSource. What that adds over the store-swap is the loader:

Bounded fan-out — one max_inflight budget vs an unbounded gather over thousands of timesteps; sustained throughput at bounded memory.
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.

What this does NOT do (scope boundaries)¶

These are deliberate current boundaries — the design is honest about them rather than pretending to be a general compute graph.

chunk_transform sees ONE variable and ONE chunk. It cannot combine variables. So windspeed = sqrt(U10² + V10²) is not a chunk transform.
It is a batch_transform — the Batch holds all variables aligned on the sample axis (batch.arrays["u10"], batch.arrays["v10"]), so derived cross-variable fields compute cleanly there. Caveat: batch transforms run after the cache, so a derived field is recomputed per batch/draw, not cached. A cached cross-variable derived variable (compute once from co-scheduled input chunks, store as a pseudo-chunk keyed like any other) is a deliberate future feature.
No cross-chunk / cross-sample-boundary ops. A sample is a slice of the sample axis (any single physical axis — see "the axis-role contract") that does not span a chunk boundary. So temporal stencils or windows that straddle two chunks (e.g. finite differences across the seam, or a 6-step window crossing chunk edges) are not supported. Windows spanning n chunks are a future opt-in that trades away zero-copy.
A field (inner) axis is carried whole. A sample is the entire payload on every non-sample axis; there is no spatial patching/cropping in the read yet, so a giant field is assembled whole (crop it in a batch_transform). Native patch geometry — and memory-optimal partial-field residency — is a reserved extension (see the axis-role contract). The shipped random-crop example (examples/wb2_dataloader.py) does exactly this — crops in a batch_transform over the whole-field cache, so the crop re-randomizes each epoch for free — and its crop is vectorized: one batched RNG draw for every sample's window origin, then a single sliding_window_view gather, so each window copies as a block rather than walking the axes element by element. What a whole-field read costs is bytes, not Python — the crop is cheap, the read is not.
Not a compute framework. No general task graph, no cross-chunk reductions on the hot path, no lazy dask-style evaluation — by design (dask on the hot path is the thing we route around). Reductions like fitting a scaler run over the loader (e.g. sklearn partial_fit), not as a graph.
Shuffle is approximate, not global — chunk permutation + shuffle-block (block_chunks is the quality↔memory knob). Exact global shuffle is incompatible with chunk-aligned, low-copy reads.
Variables must share the sample-axis length — an enforced invariant; InSituDataset raises ValueError otherwise (samples are paired row-for-row across variables). They may chunk that axis differently: the manifest defines a reference anchor grid and each variable maps global anchors onto its own chunk grid, so an OME-NGFF raw array (Z-chunk 1) pairs with its label mask (Z-chunk 30) with no reshard. Which physical axis is the sample axis is also free (the axis-role contract) — only the length along it must match.

Rule of thumb: per-variable, per-chunk, deterministic → chunk stage (cacheable). Cross-variable or per-sample-random → batch stage (not cached). Cross-chunk → not supported.