refactor: run on zarr's loop and make the pool's buffer unit the stored chunk

[emfdavid]

Summary

Closes #30. Two changes that had to land together, because both rewrite Scheduler._one:
sequencing them separately meant writing the concurrency core twice with a guaranteed
conflict.

1. Orchestration runs on zarr's event loop. There is no second loop, no per-read
bridge, and no per-pass thread churn — gcsfs and obstore are both a plain inline await.
Deleted: _io, _foreign_loop, _fsspec_io_loop, _run_loop, the insitu-sched thread,
and the per-pass ThreadPoolExecutor (three of four execution contexts were rebuilt every
epoch).

What made this possible is that zarr.core.chunk_utils.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 meant claiming that slot — which on a borrowed
loop retunes zarr's concurrency process-wide and then breaks it on close. Calling a sync
decode ourselves means we pass our executor explicitly and never touch the default.

2. The pool's buffer unit is the stored chunk. A slot was one assembled ndarray that
every decoded tile was memcpy'd into; it is now the decoded tiles themselves, adopted by
reference, with placement deferred to gather. Stored chunks are kept whole — one
buffer unit, one shape.

gather has no dispatch at all: a transformed chunk republishes as a one-tile slot
(output_geometry already sets post-transform chunks to the full inner shape), and a
persisted chunk is stored tile-major so a revived chunk is tiled exactly like a fresh one.
There is never an "assembled or tiled?" branch.

DESIGN.md's scatter-assemble was a documented design decision, so per CLAUDE.md
this went to #30 for discussion before code. That decision is now recorded as superseded
rather than quietly removed.

Built on zarr's abstractions, not a private dict of ndarrays

A deliberate goal, not incidental: the chunked pool speaks zarr's own vocabulary for
"a grid of inner chunks", so the structure stays legible to zarr and to the decode(out=)
conversation upstream.

ours	zarr	status
ArrayGeometry.tile_placement	zarr.core.indexing.ChunkProjection	adopted — it now returns one (chunk_coords, chunk_selection, out_selection, is_complete_chunk) instead of a private (dst, src) pair
our own sync codec whitelist	zarr.core.chunk_utils.ChunkTransform.decode_chunk	adopted — so the whitelist is work we no longer have to do; the design note had costed it, with GZip and zarr-v2 V2Codec traps
ArrayGeometry inner grid	DimensionGridLike	already conformant — proven by the pre-existing tests/test_zarr_indexing_parity.py, unchanged here

decode_and_scatter_chunk is not used, and that is the point: it is the decode+scatter
fusion the design note wanted hand-rolled, and a chunked slot removes the scatter entirely,
so there is nothing left to fuse.

Where we deliberately do not map. A _Slot is not a zarr shard, despite the
identical shape: a shard is a storage unit (one addressable object holding a grid of
chunks), a slot is a residency unit (the stored chunks sharing a sample-axis index,
never stored or addressed as one object). An earlier draft of this PR called it a "decoded
shard"; that was wrong and is corrected. Separately, what this codebase calls a tile is
what its own public API calls a stored chunk — two names for one concept. Both the
duplication and the broader question of aligning our "chunk" (a group, for which zarr has
no word) are split out to
#40, deliberately out of scope:
a rename that wide would hide the functional changes in this diff.

Two zarr facilities deliberately not adopted:

• FusedCodecPipeline.read_sync takes ByteGetters — it owns the IO. Adopting it
  surrenders our scheduler, max_inflight and back-pressure. decode_chunk is IO-free,
  which is why it is the one we want.
• The process-global codec_pipeline.path is not flipped. Same objection as mallopt:
  it retunes the substrate under user code. We construct our own ChunkTransform from the
  array's declared codecs instead.

Sold as simplicity, not throughput

Wall clock is a null. 4 vCPU (L3 33 MiB) and 16 vCPU / 62 GB, local zarr on NVMe,
file:// via obstore, zarr 3.3.0, batch_size=16, block_chunks=8, max_inflight=32,
decode_threads=8, 3 epochs, medians of 3 interleaved A/B/A/B reps, fingerprint-verified
per arm. main is the control, measured in the same session:

store	inner chunk	arm	wall (4 vCPU)	wall (16 vCPU)
arco_many (256×720×1440, spc 16)	45×90	warm	1.006×	1.004×
ragged (250×721×1440, spc 16)	180×360	warm	0.997×	1.007×
tiled512 (512×2048×2048, spc 1)	256×256	cold	1.040×	1.116×

The memcpy this deletes is 7.7–10.1% of process CPU on an inner-tiled store and never
surfaces at wall level, because it was not the bottleneck. No throughput claim is made.

The gather cost is not intrinsic — a controlled sweep varying only the inner tile grid
(same array, same config, warm so decode contention is zero):

tiles/chunk	1	4	16	256
inner chunk	720×1440	180×1440	180×360	45×90
gather, 4 vCPU	0.94×	0.83×	1.04×	1.84×
gather, 16 vCPU	1.02×	0.94×	0.97×	2.20×

Chunked gather is faster at coarse grids and loses only at the fine end, which restates
the chunking guidance already in docs/tuning.md rather than indicting the design.

#30's acceptance gate

a stress test that a consumer-stalled, back-pressured _drive sharing zarr's
process-global loop cannot starve or deadlock other zarr-sync work.

Passes. Consumer stalled 2 s/batch with prefetch_depth=1 so _drive parks in
_admit; a second thread hammers an unrelated zarr array through the plain sync API.
3 reps each, main as the control in the same session:

arm	victim reads during stall	p50	p95	errors
this branch (shared loop)	5670 / 5815 / 5617	1.69–1.77 ms	1.90–1.98 ms	none
main (private per-pass loop)	5590 / 5525 / 5687	1.74–1.80 ms	1.95–2.00 ms	none

But the gate as written asked the wrong question. Back-pressure never starved anything
— _admit parks on a real await, so it yields the loop. Every actual failure was in
close(), and all three were found only by testing teardown:

1. _shutdown cancelled every task on the loop (asyncio.all_tasks is the whole loop's
   set) → CancelledError inside unrelated zarr reads.
2. close() stopped and closed a loop it did not own → Cannot close a running event loop,
   and zarr's global loop dead for the rest of the process.
3. Shutting down the now-process-wide decode pool → the next scheduler raises
   cannot schedule new futures after shutdown.

They are races (the same arm produced them 2/3, 3/3 or 0/3 across runs), which is why
they are pinned by tests/test_loop_ownership.py rather than left to review. A scheduler
now cancels only the tasks it created and tears down nothing shared.

Verified against real GCS

The bridge existed for gcsfs, so file:// alone would not have proved this. Fingerprints
(SHA-256 over every batch byte + sample_indices) are identical to main on both
backends and both zarr formats:

store	format	obstore	gcsfs
gs://insitubatch-bench-insitubatch/adv_sweep_synth128_c64_i128.zarr	v3	adb1de656eef444b	adb1de656eef444b
gs://weatherbench2/.../1959-2022-6h-128x64_equiangular_conservative.zarr	v2	05cc29eee047e0e1	05cc29eee047e0e1

WeatherBench2 ARCO is zarr-v2, whose pipeline is a single V2Codec(filters, compressor)
rather than v3's codec chain — a v3-only sync-decode path would have silently missed our
main weather benchmark. Also verified byte-identical locally on a ragged grid and on
sample_axis=1.

For reviewers

Behavior is meant to be unchanged except for the two items under "User-visible changes"
below. Most valuable second look, in order:

1. Scheduler.close / _shutdown — the whole contract is "own nothing, tear down
   nothing shared". Worth checking the paths that release from the loop thread
   (tile_write's finally, on failure and cancellation) genuinely cannot reach
   _advance's assembly branch: they leave either pending > 0 or state FAILED.
2. ChunkPool._advance and pool.assembles. Delivery is inline on the tiled path (a
   dict write and a counter) but 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. An earlier draft delivered inline unconditionally and ran user code
   on zarr_io — pinned now by
   test_user_code_never_runs_on_the_shared_event_loop.
3. slot_charge_bytes is shared by the pool and the auto-sizer, deliberately. Sizing
   from the output shape while charging stored tiles under-provisions the budget and the
   pool starves mid-epoch — a hang-shaped failure, not a slow one. It caught me: raising
   the charge without raising the sizer broke 5 tests.
4. gather uses the source grid, not out_geom. output_geometry always sets chunks
   to the full inner shape, so out_geom describes a 1-tile grid — right for an assembling
   slot, wrong for a tiled one.

Two notes on scope. The starvation numbers are file:// on NVMe, 4 vCPU, a single dataset,
and the tiled path; they do not cover a remote store where the loop also carries network
waits. And decode_pool() sizes on first use, so the warning is the only signal a second
dataset's decode_threads was ignored.

User-visible changes

• Breaking: the persisted cache format is bumped to 3. A chunk's .npy now holds
  (n_tiles, *tile_shape) tile-major instead of one assembled array (file count per chunk
  unchanged). An existing cache_dir raises the usual stale-cache error and is rebuilt with
  reset_stale_cache=True. A version-2 file read as tile-major would be plausible-looking
  garbage, so it is refused rather than reinterpreted.
• decode_threads is process-wide, sized by the first dataset built in the process; a
  later different value is ignored with a warning. Thread count is a property of the machine,
  not of a dataset, and two datasets should not run two pools competing for the same cores.
• A ragged chunk grid is charged for the padding it holds (1.248× on ERA5 721×1440 at
  180×360; 1.997× on a short final outer chunk). The automatic budget accounts for it; a
  hand-set cache_budget_bytes on a ragged grid needs to be proportionally larger. Grids
  that divide evenly are unaffected. A chunk_transform still receives the logical
  clipped chunk and never sees padding.

Independent verification (examples + benchmarks, post-review)

The unit suite is green, but it does not run the examples — CI only lints and typechecks
them — and it does not touch a real object store or a GPU. So the examples and benchmarks
were run directly against main as a control, looking for what the tests could miss.

Setup. Two boxes: g2-standard-16 (16 vCPU, 62 GB, NVIDIA L4 — the box the tables
above were measured on) and an 8 vCPU / 31 GB box. Local zarr on NVMe, file:// via
obstore, zarr 3.3.0, torch 2.12.0+cu130, identical resolved envs on both arms. Control is
main @ 6282195, measured in the same session on the same stores. Probe stores rebuilt
with bench.make_dataset (arco_sq 180×360, arco_many 45×90, ragged 721×1440 @ 180×360
with a short final outer chunk, tiled 2048×2048 @ 256×256 spc=1).

Correctness: byte-identical to main everywhere

The failure mode this refactor can produce is plausible wrong data — a tile placed at the
wrong offset, an edge tile clipped wrong, a revived cache entry read with the wrong layout.
Shapes, dtypes, throughput and "it ran" all pass through that, so every arm is SHA-256 over
every delivered byte plus sample_indices, sorted into sample order so the digest is
shuffle-independent.

17 arms, run on both boxes, 34/34 digests identical to main: the four geometry
quadrants on the plain path; the windowed multi-geometry shape the Earth2Studio adapter
drives (two variables at one anchor + shift(), sample_range, shuffled 3-way split);
chunk_transform on three stores; two-epoch runs (cross-epoch reuse and unpinning);
and cache_dir cold/warm/transform pairs. Cold and warm digests match within each arm, so
the tile-major persist round-trips; and arco_sq and arco_many — same data, different
inner chunking — produce the same digest, so the chunk grid does not leak into delivered
bytes.

Real GCS, both backends. The deleted fsspec bridge is invisible to file://. On
WeatherBench2 ARCO (gs://weatherbench2/.../1959-2022-6h-128x64_equiangular_conservative.zarr,
zarr-v2, so the V2Codec path rather than a v3 codec chain), obstore and gcsfs × shuffled
and unshuffled all give 91df010e652bc478 on both arms.

Examples, including paths CI never executes. advection on torch, JAX, TF and CUDA:
persistence RMSE 0.884 in every run on both arms and both boxes — that number is
model-free, a property of the delivered val data, and is the only signal that has ever caught
a loader defect here. microscopy (sample_axis=2, per-variable chunking) on CPU and CUDA:
Otsu IoU 0.386 across all nine runs. Model skill varies run to run (0.46–0.63 on a
4-epoch CNN, on both arms) and is training noise, not loader signal. transforms (including
the reshaping Coarsen, which exercises the one-tile collapse) and fit_scaler produce
output identical to main. The bench smoke grid runs clean across all five engines.

The two user-visible claims above, checked directly. A version-2 cache_dir written by
main is refused by this branch with the documented stale-cache error naming
reset_stale_cache=True, not reinterpreted; passing that flag rebuilds it to the correct
digest. The ragged residency multiplier computes to 1.248× on 721×1440 @ 180×360,
matching the figure quoted above. slot_charge_bytes is confirmed to be the single function
used by both the auto-sizer (source.py:226) and the pool's own charge (pool.py:899) —
the coupling reviewer note 3 flags. And under-sizing is not hang-shaped in practice:
cache_budget_bytes is a floor (max(user, working_set), source.py:273), and genuine
exhaustion raises Scheduler._starvation's diagnostic rather than hanging.

The wall-clock table is load-bearing on compute_ms — please state it

The warm wall-clock nulls above reproduce, but only at a stated consumer load, and the
same stores on the same box give a very different ratio at zero consumer compute. Warm
medians on arco_many (45×90, 256 tiles/chunk — the worst quadrant), 16 vCPU,
batch_size=16 block_chunks=8 max_inflight=32, 4 epochs, median of the warm epochs, one
process per rep, main interleaved as the control:

per-batch consumer work	branch	main	ratio
none	0.870 s	0.315 s	2.76×
60 ms time.sleep (the harness's own proxy, bench/engines.py:63)	1.03 s	0.994 s	1.04×
60 ms single-threaded numpy	0.929 s	0.947 s	0.98×
60 ms multi-threaded numpy (BLAS takes all 16 cores)	1.949 s	1.113 s	1.75×

gather runs on the producer thread (source.py:454, inside produce()) behind a
prefetch_depth queue, so its cost overlaps consumer compute and disappears from wall clock
as soon as the consumer leaves the loader either cores or the GIL. The last row is core
starvation from a 16-thread BLAS consumer, not a gather defect — the same workload pinned to
one thread is at parity.

So the 1.004× reported above is right for a GPU-shaped consumer and wrong to read as the
loader's own throughput: the warm ceiling on a 256-tile grid is ~2.8× lower (and 1.88×
on 8 vCPU). That is the same effect the gather sweep already discloses at 1.84–2.20×, and it
only reaches wall clock when the loader is the bottleneck. It costs one line to say which
quantity the table reports; without it the table invites exactly the wrong reading.

The mainstream grids are a wash rather than a regression: arco_sq (16 tiles/chunk) and
ragged (20) came out 1.18–1.19× slower at zero consumer compute on 16 vCPU and 0.85×
faster on 8 vCPU, and the spc=1 tiled store — where each chunk is gathered once — is
0.82×, i.e. faster, on the branch.

Author attestation

• I have reviewed every change in this PR, I can explain why each one is correct, and I
  have verified the claims made in this description.

Left unchecked deliberately — this was drafted by Claude, and that box is the human
author's to tick.

Checklist

• Tests added or updated — tests/test_loop_ownership.py (6) and
  tests/test_chunked_slots.py (14) are new; the loop-ownership tests all fail on main,
  as do 4 of the chunked-slot tests, for the documented reasons
• uv run ruff check src tests bench examples, uv run mypy src bench examples and
  uv run pytest -q are green locally (318 passed, 21 skipped)
• Docstrings and API docs for any new or changed public surface — decode_pool,
  reset_decode_pool, slot_charge_bytes, ChunkPool.assembles,
  ArrayGeometry.tile_shape / inner_index, and DecodedChunk (which now states
  outright that user code never sees stored-chunk padding)
• User-facing behavior documented in docs/*.md — docs/architecture.md (one loop, the
  handoff diagram, teardown ownership), docs/tuning.md (process-wide decode_threads,
  ragged-grid residency), DESIGN.md (scatter-assemble superseded, M-GCS closed)
• A bullet added under ## Unreleased in CHANGELOG.md
• No load-bearing invariant is broken — Batch is still numpy, the hot path is still
  O(chunks) (gather is O(chunks × tiles-per-chunk), never per-sample), parallelism is
  still in the event loop, no dask, no reshard, no xr.DataArray
• Touches ChunkPool, the scheduler, and cross-thread readiness → free-threaded run
  passes: 3.13t with PYTHON_GIL=0 (asserted sys._is_gil_enabled() is False),
  99 passed / 2 skipped across pool, lifecycle, chunked-slots, loop-ownership, scheduler,
  residency, prefetch, pinned and window
• Performance claim? None made — the change is argued on simplicity, and the numbers
  above exist to show wall clock is a null. Setup and a main control are given for every
  table regardless.

[emfdavid]

Awesome simplification