diff --git a/DESIGN.md b/DESIGN.md index 2e4c6e8..79b7769 100644 --- a/DESIGN.md +++ b/DESIGN.md @@ -99,8 +99,9 @@ The plan buys two orthogonal guarantees — **read-once** (a stored tile fetched decoded once, however many samples touch it) and **sample-once** (each valid sample in exactly one batch, tracked by the `order` ledger, with the fancy index recomputed in `gather` and never stored) — both independent of batch size. The mechanism, the edge -cases (short final chunk, dropped windowed edges), and the ragged per-block tail are the -live contract: see [Read-once and sample-once](docs/architecture.md#read-once-and-sample-once). +cases (short final chunk, dropped windowed edges), and the single ragged tail per epoch are +the live contract: see +[Read-once and sample-once](docs/architecture.md#read-once-and-sample-once). ## Sample geometry — how the ladder evolved @@ -417,16 +418,16 @@ The shape above wasn't the first cut. The pivots that got here, and the roads no Things wrong or missing in *our* code today, with the reasoning that sets their priority. (Things simply *not built yet* are milestones — see Roadmap.) -- **Per-block ragged batches (possible wart).** Batches never span shuffle blocks, so a - block's sample count not dividing `batch_size` yields a short final batch *per block*, not - one per epoch — behaviour and step-count impact in - [Read-once and sample-once](docs/architecture.md#read-once-and-sample-once). Correct - (sample-once holds), but it can surprise consumers assuming a uniform batch size. No - `drop_last` today. Options if users want uniformity: (a) opt-in `drop_last`; (b) carry a - block's remainder into the next block's first batch (extends that block's pin lifetime past - its own batches); (c) a whole-epoch re-batch (cheap — `order` is already materialized — but - mixes chunks across block boundaries, diluting the block-local residency guarantee). - **Priority pending user feedback** — no correctness impact, so we hold until someone hits it. +- ~~**Per-block ragged batches.**~~ **FIXED.** Batches are now cut over the whole epoch + `order`, so steps-per-epoch is `⌈N / bs⌉` and only the epoch's last batch is short — + option (b), carrying a block's remainder into the next block's first batch. It cost less + than the entry assumed: `order` was already one flat array with blocks as contiguous row + ranges, and `last_use` already deferred a chunk's release past its own block for windowed + reads, so the change is two monotone frontiers (wait / release) in the producer, not new + machinery. Peak co-residency stays two blocks — already the budget floor — so option (c)'s + worry about "diluting the block-local residency guarantee" does not apply: block-local + residency is unchanged, only the *batch* spans a boundary. Still no `drop_last`; the epoch's + short tail is the ordinary loader contract and `len(batch)` is how a caller drops it. - ✅ **SHIPPED — batch buffer reuse + pinning** (`buffers.BatchBuffers`, `frameworks.as_torch(device=...)` / `pin_host_buffers`; branch `batch-buffer-ring`). diff --git a/docs/architecture.md b/docs/architecture.md index 4e59b99..db2319f 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -295,9 +295,10 @@ flowchart LR Q -.->|"full ⇒ producer blocks
(backpressure)"| ASM ``` -- **Producer** starts the scheduler over the epoch's chunks, then per shuffle-block - waits the block assembled, gathers its batches, and unpins it, pushing batches to - a bounded `queue.Queue(maxsize=d)`. +- **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 / memory bound** — queue depth `d` + the pool's byte budget cap @@ -443,15 +444,29 @@ training; the inference path validates the range and raises instead). Guaranteed `test_order_covers_every_sample_exactly_once`, `test_order_handles_partial_final_chunk`, and the decode-once suite. -**The tail is ragged per block, not per epoch.** Batches do not span shuffle blocks: the -producer batches `order` within each block's row range and restarts at the boundary. So when a -block's sample count (`≈ block_chunks × spc`, minus edges) is not a multiple of `batch_size` — -the common case — the **last batch of every block is short**. No sample is lost or duplicated -(sample-once holds), but an epoch yields *several* short batches, and steps-per-epoch is -`Σ ⌈block_samples / bs⌉`, not `⌈N / bs⌉` — worth knowing for BatchNorm on a small tail or -step-count math. There is deliberately no `drop_last` today; whether to add one is an open -question (see Known limitations in -[DESIGN.md](https://github.com/emfdavid/insitubatch/blob/main/DESIGN.md#known-limitations--defects)). +**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. 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. + +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: + +```python +for batch in ds.train: + if len(batch) < ds.batch_size: + continue +``` ## Transforms — three stages, placed by cost diff --git a/src/insitubatch/scheduler.py b/src/insitubatch/scheduler.py index 86cccec..3dbd184 100644 --- a/src/insitubatch/scheduler.py +++ b/src/insitubatch/scheduler.py @@ -33,9 +33,13 @@ hanging. Budget floor: a batch may draw from any chunk in its shuffle-block, so the whole -block must be co-resident to gather -- the budget must hold at least one block (the -producer sizes it to two: the current block plus one read-ahead block, so -block-boundary IO overlaps the current block's compute). +block must be co-resident to gather -- and since batches are cut over the whole epoch +order rather than per block, the batch that straddles a boundary needs *two* blocks +co-resident. The producer already sizes the budget to two (the current block plus one +read-ahead, so block-boundary IO overlaps the current block's compute), so the +straddling batch costs nothing beyond that floor -- but the floor is now load-bearing +for correctness, not only for overlap. Never three: a block is released as soon as a +batch has consumed its last row. """ from __future__ import annotations diff --git a/src/insitubatch/source.py b/src/insitubatch/source.py index 0696127..0cc76a9 100644 --- a/src/insitubatch/source.py +++ b/src/insitubatch/source.py @@ -87,8 +87,9 @@ class InSituDataset: for batch in ds.train: ... # one epoch; ds.set_epoch(e) reshuffles for batch in ds.val: ... - One epoch over a view = permute the split's chunks -> walk shuffle-blocks -> per - block, stream-fetch its stored chunks into the pool, gather coalesced batches, evict. + One epoch over a view = permute the split's chunks -> walk shuffle-blocks, stream-fetching + each block's stored chunks into the pool -> gather coalesced batches cut over the whole + epoch (so only the last is short, and one may span a block boundary) -> evict. Batches are numpy :class:`Batch`; convert to a framework with :mod:`insitubatch.frameworks` (``as_torch`` / ``to_jax`` / ``as_tf_dataset``). A different per-split configuration (e.g. train-only augmentation) is a separate dataset. @@ -337,13 +338,15 @@ def _iterate(self, split: SplitName | None, shuffle: bool) -> Iterator[Batch]: one :class:`ChunkPool` -- so a chunk a windowed read pulls across a split boundary is decoded once and reused by both splits. - A producer thread starts the scheduler over the split's chunks (in draw order), - then for each shuffle-block waits the block assembled, gathers its batches, and - unpins it; this consumer pops from a bounded queue (depth ``prefetch_depth``) that - provides backpressure and inter-batch overlap. The scheduler keeps ``max_inflight`` - tiles continuously in flight and fetches one block ahead, so block-boundary IO - overlaps the per-batch compute. Chunks the pool already holds (cross-epoch or - cross-split hits) cost no fetch. + A producer thread starts the scheduler over the split's chunks (in draw order), then + walks the epoch's batches: waiting each shuffle-block a batch draws from, gathering, + and unpinning a block once a batch has consumed its last row. Batches are cut over the + whole epoch order, so only the last one is short -- a batch may span a block boundary, + and the two frontiers in :func:`produce` are what makes that safe. This consumer pops + from a bounded queue (depth ``prefetch_depth``) that provides backpressure and + inter-batch overlap. The scheduler keeps ``max_inflight`` tiles continuously in flight + and fetches one block ahead, so block-boundary IO overlaps the per-batch compute. + Chunks the pool already holds (cross-epoch or cross-split hits) cost no fetch. """ spc = self._ref_spc # the manifest anchor grid, shared by every variable order = self._draw_order(split, shuffle) @@ -374,28 +377,55 @@ def _iterate(self, split: SplitName | None, shuffle: bool) -> Iterator[Batch]: release[bi].add(key) def produce(sched: Scheduler) -> None: + # Batches are cut over the WHOLE epoch order, not per block. Cutting them per + # block (`range(rstart, rstop, bs)`) made every block whose row count was not a + # multiple of batch_size end in a short batch -- one per block rather than one + # per epoch, and up to *three* distinct shapes once the epoch's own short final + # block is counted. Sample-once held either way, so nothing failed; what broke + # was any fixed-shape consumer (torch.compile / jax.jit retrace per shape, and + # with three shapes the retrace cache never settles). + # + # Blocks stay exactly what they were -- contiguous row ranges over disjoint + # chunks -- and `order` was always one flat array for the epoch, so a batch that + # crosses a boundary is just `order[start : start + bs]` not being truncated. + # What the boundary needs is bookkeeping, tracked as two monotone frontiers over + # the block list: `ready` (waited) and `freed` (released). Both only advance, and + # each block passes through each exactly once, so this stays O(blocks) of Python. bs = self.batch_size + ready = freed = 0 try: sched.start(ordered_chunks, spc) - for bi, (rstart, rstop, _cids) in enumerate(blocks): - # A block's batches draw across its whole read-union, so wait it all - # assembled (and claimed by the driver -- see ChunkPool.wait_ready) - # before gathering; each wait is cheap once ready. - for path, cid in block_keys[bi]: - sched.pool.wait_ready(path, cid) - for start in range(rstart, rstop, bs): - if stop.is_set(): - return - rows = order[start : min(start + bs, rstop)] - batch = sched.pool.gather(rows, self.variables, spc) - for transform in self.batch_transforms: - batch = transform(batch) - out_q.put(batch) # blocks when full -> backpressure - # Release the driver's reference on chunks whose *last* use is this - # block: now LRU-evictable (retained for reuse if budget allows), - # unblocking the read-ahead. Chunks read again later keep their - # reference until then. - sched.unpin_block(release[bi]) + for start in range(0, len(order), bs): + if stop.is_set(): + return + stop_row = min(start + bs, len(order)) + # Wait every block this batch draws from. A batch spans at most two + # (bs <= a block's rows in any sane configuration, and the loop is + # correct regardless): a block's batches draw across its whole + # read-union, so it must be assembled -- and claimed by the driver, see + # ChunkPool.wait_ready -- before gathering. Each wait is cheap once ready. + while ready < len(blocks) and blocks[ready][0] < stop_row: + for path, cid in block_keys[ready]: + sched.pool.wait_ready(path, cid) + ready += 1 + batch = sched.pool.gather(order[start:stop_row], self.variables, spc) + # Release the driver's reference on chunks whose *last* use is a block now + # fully behind the frontier -- a batch has consumed its last row, so it is + # done. Now LRU-evictable (retained for reuse if budget allows), unblocking + # the read-ahead. Chunks a later block reads again keep their reference. + # + # Here, not after the queue put: `out_q.put` blocks when full, so releasing + # after it keeps a whole block pinned for as long as the consumer is slow, + # and pinned chunks cannot be evicted -- consumer slowness would propagate + # into a read-ahead stall via `try_admit` parking on a full budget. Safe + # this early because `gather` COPIES out of the slots, so the batch never + # aliases chunk memory and the reference is already dead weight. + while freed < len(blocks) and blocks[freed][1] <= stop_row: + sched.unpin_block(release[freed]) + freed += 1 + for transform in self.batch_transforms: + batch = transform(batch) + out_q.put(batch) # blocks when full -> backpressure except Exception as exc: # noqa: BLE001 - forwarded to the consumer out_q.put(exc) finally: diff --git a/src/insitubatch/types.py b/src/insitubatch/types.py index 29da9f7..7af071e 100644 --- a/src/insitubatch/types.py +++ b/src/insitubatch/types.py @@ -243,6 +243,24 @@ class Batch: sample_indices: np.ndarray = field(default_factory=lambda: np.empty(0, dtype=np.int64)) offsets: dict[str, int] = field(default_factory=dict) # label -> sample-axis read offset + def __len__(self) -> int: + """Rows in the batch -- how a caller implements ``drop_last`` for itself. + + Every batch is full except the epoch's last, which is short whenever the split's + sample count does not divide ``batch_size``. That is the ordinary data-loader + contract, and dropping it is the caller's choice, so the check has to be a one-liner + that does not require naming a variable:: + + for batch in ds.train: + if len(batch) < ds.batch_size: + continue + + Counted off ``sample_indices`` rather than an array's leading axis: it is the anchor + row count the engine sets on every batch and the one :meth:`read_indices` already + builds on, so this adds a spelling rather than a second notion of how long a batch is. + """ + return len(self.sample_indices) + def read_indices(self, label: str) -> np.ndarray: """Global sample index each row of ``label`` was read from: ``anchor + offset``. diff --git a/tests/test_buffers.py b/tests/test_buffers.py index 479b833..44eab72 100644 --- a/tests/test_buffers.py +++ b/tests/test_buffers.py @@ -17,6 +17,7 @@ from __future__ import annotations import re +import sys import threading import time from typing import Any @@ -397,16 +398,39 @@ def counts(line: str) -> tuple[int, int]: assert m, line return int(m[1]), int(m[2]) + def held(line: str) -> int: + m = re.search(r"(\d+) x heap", line) + assert m, line + return int(m[1]) + (lent0, alloc0), (lent1, alloc1) = counts(lines[0]), counts(lines[1]) - # 10 batches, not 8: a batch never crosses a shuffle block, so 40 rows in blocks of - # 16/16/8 give ragged batches. One buffer lent per batch, every epoch. - assert lent0 == lent1 == 10 + # 8 batches: batches are cut over the whole epoch order, so 40 rows at batch_size 5 give + # 8 full ones and no remainder -- the blocks (16/16/8 rows) do not each end in a short + # batch. One buffer lent per batch, every epoch. This read 10 while batches were cut + # per block; see tests/test_source.py's draw-policy section. + assert lent0 == lent1 == 8 # The point of the line. Epoch 0 allocates however many are genuinely in flight -- 3 or 4 # here, decided by producer/consumer timing, so it is not a fixed number. What must hold is # that a warm pool stops allocating: were buffers failing to come back, this would climb # toward one per batch and the pool would be a growing memory floor. assert alloc0 >= 1 - assert alloc1 <= 1, f"warm epoch still allocating: {lines[1]}" + # `sys._is_gil_enabled` is 3.13+; on 3.12 the GIL is always on, so default to True. + if getattr(sys, "_is_gil_enabled", lambda: True)(): + assert alloc1 <= 1, f"warm epoch still allocating: {lines[1]}" + else: + # Free-threading cannot hold the tight bound, and the reason is the open question + # `buffers.py` documents: `sys.getrefcount` may read stale. Reading *high* makes a + # free buffer look lent, so the pool allocates another instead of reusing -- benign + # (over-allocation, never corruption), and the direction the liveness guard in + # `_Buffer.free` does not police because it is not the dangerous one. + # + # Measured over 120 runs of exactly this scenario on 3.13t: `alloc1` came back 0, 1 + # or 2 (so the tight bound fails ~14% of the time -- it flakes on main too, at a + # rate indistinguishable from any branch), the pool settled at 2-5 buffers, and the + # liveness guard raised zero times. So assert what the tight bound is a proxy for: + # the pool is not degenerating toward one buffer per batch, which is what an actual + # "buffers never come back" regression looks like and what makes it a memory floor. + assert held(lines[1]) < lent1, f"pool grew toward one buffer per batch: {lines[1]}" def test_exported_torch_tensor_holds_the_buffer() -> None: diff --git a/tests/test_prefetch.py b/tests/test_prefetch.py index 83d8240..90d3554 100644 --- a/tests/test_prefetch.py +++ b/tests/test_prefetch.py @@ -11,6 +11,7 @@ import time import numpy as np +import pytest import zarr from insitubatch import ( @@ -129,7 +130,8 @@ def test_partial_iteration_reaps_producer(tmp_path) -> None: ds.close() -def test_early_break_then_next_epoch_does_not_deadlock(tmp_path) -> None: +@pytest.mark.parametrize("batch_size", [4, 12], ids=["divides-a-block", "straddles-blocks"]) +def test_early_break_then_next_epoch_does_not_deadlock(tmp_path, batch_size: int) -> None: """A capped epoch (early break) must not poison the next epoch. Regression: ``try_admit`` pins each chunk the driver admits; the producer only @@ -140,6 +142,20 @@ def test_early_break_then_next_epoch_does_not_deadlock(tmp_path) -> None: free no room and the driver deadlocks on ``_capacity`` while the consumer hangs in ``wait_ready`` (observed on the bench: scheduler loop idle, prefetch thread parked, no in-flight work). Pins must not survive an epoch. + + Parametrized on whether a batch can straddle a block boundary: a block is 2 chunks x 8 + samples = 16 rows, so ``batch_size=4`` divides it and no batch ever crosses, while + ``batch_size=12`` makes the second batch span two blocks. + + Be precise about what that buys. Cutting batches over the epoch defers a block's release + until a batch consumes its last row, but that extra hold lasts only from the straddling + ``wait_ready`` to the release a few statements later, inside one producer step -- a + consumer ``break`` cannot land there. The *durable* two-block pin at a break comes from + read-ahead, which predates all of this. So the straddling arm does not reach a pinned + state the dividing arm cannot; it is coverage that the teardown path works under a + straddling draw, not a distinct leak. ``unpin_all`` clears the pin table wholesale, so + the count never mattered -- and the assertion below keeps *both* arms from passing on a + single-block state that would not exercise the leak at all. """ url = _write(tmp_path, n=160, spc=8) # 20 chunks; train split ~16 geom = open_geometries(obstore_store(url))["t2m"] @@ -147,11 +163,12 @@ def test_early_break_then_next_epoch_does_not_deadlock(tmp_path) -> None: # budget = 2 * block_chunks chunks; batch_size < one block so a block yields # several batches -> breaking after one batch leaves the current block pinned too. + block_chunks = 2 ds = InSituDataset( obstore_store(url), manifest, - batch_size=4, - block_chunks=2, + batch_size=batch_size, + block_chunks=block_chunks, prefetch_depth=2, ) @@ -160,6 +177,12 @@ def test_early_break_then_next_epoch_does_not_deadlock(tmp_path) -> None: it = iter(ds.train) next(it) time.sleep(0.3) # let read-ahead pin a block or two before we tear down + # The precondition the test rests on: the break really is leaving more than one + # block's worth of chunks pinned, so epoch 1 is inheriting the state described above. + assert len(ds._pool._pinned) > block_chunks, ( + f"only {len(ds._pool._pinned)} chunks pinned at the break -- " + "not the multi-block leak this test exists to clear" + ) it.close() # deterministic generator teardown (GeneratorExit -> __iter__ finally) # epoch 1: must run to completion, not deadlock on leaked pins. diff --git a/tests/test_source.py b/tests/test_source.py index 0002415..798e6b1 100644 --- a/tests/test_source.py +++ b/tests/test_source.py @@ -382,3 +382,134 @@ def test_val_view_is_deterministic_train_shuffles(write_zarr) -> None: val5 = np.concatenate([b.sample_indices for b in ds.val]) np.testing.assert_array_equal(val0, val5) # val ignores epoch (no shuffle) np.testing.assert_array_equal(val0, np.sort(val0)) # and is in order + + +# --- batch draw policy across shuffle-block boundaries ------------------------------------- +# +# Batches used to be drawn *within* each block (`for start in range(rstart, rstop, bs)`), so +# every block whose row count was not a multiple of batch_size ended in a short batch -- one +# per block, not one per epoch, and up to three distinct shapes once the epoch's own short +# final block is counted. That is a fixed-shape consumer's problem (torch.compile / jax.jit +# retrace per shape, and the retrace cache never settles) and a silent one: sample-once still +# held, so nothing failed. Batches are now drawn over the whole epoch order, with block +# readiness and release driven off a frontier. + +_STRADDLE = [ + # (n, spc, block_chunks, batch_size) -- rows/block deliberately not a multiple of bs + pytest.param(160, 4, 8, 12, id="32-rows-per-block-bs12"), + pytest.param(160, 4, 4, 10, id="16-rows-per-block-bs10"), + pytest.param(120, 8, 2, 7, id="16-rows-per-block-bs7"), +] + + +def _batch_sizes(url, *, n, spc, block_chunks, batch_size, shuffle=True): # type: ignore[no-untyped-def] + geom = open_geometries(obstore_store(url))["t2m"] + manifest = split_by_chunk(geom, fractions=(1.0, 0.0, 0.0)) + ds = InSituDataset( + obstore_store(url), + manifest, + batch_size=batch_size, + block_chunks=block_chunks, + shuffle=shuffle, + seed=0, + ) + ds.set_epoch(0) + batches = list(ds.train) + return ds, batches + + +@pytest.mark.parametrize(("n", "spc", "block_chunks", "batch_size"), _STRADDLE) +def test_only_the_epochs_final_batch_is_short( + write_zarr, n: int, spc: int, block_chunks: int, batch_size: int +) -> None: + """One ragged batch per *epoch*, not per block -- the normal data-loader contract.""" + url, _ = write_zarr(n=n, spc=spc) + _ds, batches = _batch_sizes(url, n=n, spc=spc, block_chunks=block_chunks, batch_size=batch_size) + sizes = [len(b.sample_indices) for b in batches] + + short = [i for i, s in enumerate(sizes) if s != batch_size] + assert short in ([], [len(sizes) - 1]), f"short batches at {short} of {len(sizes)}: {sizes}" + assert len(set(sizes)) <= 2, f"more than two shapes: {sorted(set(sizes))}" + + +@pytest.mark.parametrize(("n", "spc", "block_chunks", "batch_size"), _STRADDLE) +def test_every_sample_still_appears_exactly_once( + write_zarr, n: int, spc: int, block_chunks: int, batch_size: int +) -> None: + """Sample-once is the invariant the whole draw policy exists to preserve.""" + url, _ = write_zarr(n=n, spc=spc) + _ds, batches = _batch_sizes(url, n=n, spc=spc, block_chunks=block_chunks, batch_size=batch_size) + idx = np.concatenate([b.sample_indices for b in batches]) + assert sorted(idx.tolist()) == list(range(n)) + + +def test_a_batch_may_span_a_block_boundary(write_zarr) -> None: + """The positive statement, so the test above cannot pass by dropping the remainder. + + Sequential order makes blocks contiguous sample ranges: 4 samples/chunk x 8 chunks = 32 + rows per block, so with batch_size 12 the third batch is rows [24, 36) and must straddle. + """ + url, _ = write_zarr(n=160, spc=4) + _ds, batches = _batch_sizes(url, n=160, spc=4, block_chunks=8, batch_size=12, shuffle=False) + rows_per_block = 4 * 8 + spanning = [b for b in batches if len(set(b.sample_indices // rows_per_block)) > 1] + assert spanning, "no batch crossed a block boundary" + assert len(spanning[0]) == 12 + + +@pytest.mark.parametrize(("n", "spc", "block_chunks", "batch_size"), _STRADDLE) +def test_residency_is_still_bounded_by_two_blocks( + write_zarr, n: int, spc: int, block_chunks: int, batch_size: int +) -> None: + """A straddling batch needs two blocks co-resident -- which is already the budget floor. + + `scheduler.py`'s "Budget floor" note sizes the working set at two blocks (current plus + one read-ahead), so the frontier costs nothing here. A regression to three would mean a + block is being held past its last row. + """ + url, _ = write_zarr(n=n, spc=spc) + ds, _batches = _batch_sizes(url, n=n, spc=spc, block_chunks=block_chunks, batch_size=batch_size) + assert ds.resident_peak <= 2 * block_chunks + + +def test_batch_len_is_its_row_count(write_zarr) -> None: + """`len(batch)` is how a consumer implements drop_last -- it must not need a variable name.""" + url, _ = write_zarr(n=50, spc=8) + _ds, batches = _batch_sizes(url, n=50, spc=8, block_chunks=2, batch_size=8, shuffle=False) + assert [len(b) for b in batches] == [len(b.sample_indices) for b in batches] + assert len(batches[-1]) == 50 % 8 # the one short batch a caller may choose to drop + assert all(len(b) == b.arrays["t2m"].shape[0] for b in batches) + + +@pytest.mark.parametrize(("n", "spc", "block_chunks", "batch_size"), _STRADDLE) +def test_a_completed_epoch_leaves_no_block_pinned( + write_zarr, n: int, spc: int, block_chunks: int, batch_size: int +) -> None: + """The release frontier must drain on its own -- not be rescued by the epoch backstop. + + Release is now conditional (``blocks[freed].rstop <= stop_row``) where it used to be + unconditional per block, so "a block never released" is newly reachable: it would not + raise, it would starve the next epoch's admission, which parks on a full budget awaiting + a release that never comes. + + Asserting that is harder than it looks. ``_iterate`` calls ``pool.unpin_all()`` at the + *start* of every epoch precisely to absorb pins an aborted epoch leaked, so any test that + merely runs two epochs passes whether the frontier released anything or not -- it exercises + the backstop, not the frontier. So this reads the pin table between epochs, where the + backstop cannot mask it. Verified to fail if the release condition is weakened to ``<``. + """ + url, _ = write_zarr(n=n, spc=spc) + geom = open_geometries(obstore_store(url))["t2m"] + manifest = split_by_chunk(geom, fractions=(1.0, 0.0, 0.0)) + ds = InSituDataset( + obstore_store(url), + manifest, + batch_size=batch_size, + block_chunks=block_chunks, + shuffle=True, + seed=0, + ) + ds.set_epoch(0) + idx = np.concatenate([b.sample_indices for b in ds.train]) + assert sorted(idx.tolist()) == list(range(n)) + assert not ds._pool._pinned, f"{len(ds._pool._pinned)} chunks pinned after a full epoch"