From 85573ddae35cd2ef5e7f18d6ec784a682e17ce40 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Sun, 2 Aug 2026 03:12:20 +0000 Subject: [PATCH 1/3] fix: cut batches over the epoch, not per shuffle-block Batches were drawn within each block's row range (`range(rstart, rstop, bs)`) and restarted at the boundary, so every block whose row count was not a multiple of batch_size ended in a short batch -- one per block rather than one per epoch. Sample-once held either way, so nothing failed loudly. What broke was any fixed-shape consumer: torch.compile and jax.jit retrace per shape, and counting the epoch's own short final block that was up to THREE distinct shapes, so the retrace cache never settled. Measured over the real order/partition, 20 chunks: spc blkc bs rows/blk | short now shapes | short after shapes 64 16 32 1024 | 0 1 | 0 1 64 16 48 1024 | 2 2 | 1 2 73 16 32 1168 | 2 3 | 1 2 10 4 32 40 | 5 2 | 1 2 64 4 96 256 | 5 2 | 1 2 Note row 1: when bs divides a block's rows there was never a problem. This is a sharp edge users fall off, not a constant tax -- which is why it went unnoticed. The fix is smaller than DESIGN.md's entry assumed (it costed this as option (b), "extends that block's pin lifetime past its own batches"). Two things already existed: `order` is one flat array for the epoch with blocks as contiguous row ranges, so a straddling batch is just `order[start:start+bs]` not being truncated; and `last_use` already defers a chunk's release past its own block for windowed reads. So this is two monotone frontiers over the block list -- `ready` (waited) and `freed` (released) -- not new machinery. Each block passes through each exactly once, so the producer stays O(blocks) of Python. Residency is unchanged. A straddling batch needs blocks f and f+1 co-resident, which is already the budget floor (`source.py` sizes the working set at 2 * block_chunks: current block plus one read-ahead). Simulated across the configs above, peak co-pinned is 2 in every straddling case and 1 when bs divides -- never 3. The floor is now load-bearing for correctness rather than only for overlap, which scheduler.py's note now says. DESIGN.md's option (c) worried that a whole-epoch rebatch would "dilute the block-local residency guarantee". It does not: block-local residency is untouched, only the *batch* spans a boundary. Release also moves ahead of the queue put, which is a real change and not just code motion. `out_q.put` blocks when full, so releasing after it kept a whole block pinned for as long as the consumer was slow, and pinned chunks cannot be evicted -- consumer slowness propagated 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. Tests: batch-size multiset, sample-once, and a positive assertion that a batch really does straddle (so the first cannot pass by dropping the remainder), over three straddling configs. Plus one that a completed epoch leaves no block pinned -- release is now conditional, so "a block never released" is newly reachable and would not raise, it would starve the next epoch's admission. That one needs care: `_iterate` calls `pool.unpin_all()` at the START of each epoch to absorb pins an aborted epoch leaked, so a test that just runs two epochs passes whether the frontier released anything or not. It reads the pin table between epochs instead, and was verified to fail on all three configs when the release condition is weakened to `<`. test_buffers' epoch-summary test asserted 10 batches over 40 rows at bs=5, documenting the old raggedness; it is 8 now. Co-Authored-By: Claude Opus 5 --- DESIGN.md | 25 +++---- docs/architecture.md | 39 +++++++---- src/insitubatch/scheduler.py | 10 ++- src/insitubatch/source.py | 86 +++++++++++++++-------- src/insitubatch/types.py | 18 +++++ tests/test_buffers.py | 8 ++- tests/test_source.py | 131 +++++++++++++++++++++++++++++++++++ 7 files changed, 259 insertions(+), 58 deletions(-) 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..bbfe212 100644 --- a/tests/test_buffers.py +++ b/tests/test_buffers.py @@ -398,9 +398,11 @@ def counts(line: str) -> tuple[int, int]: return int(m[1]), int(m[2]) (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 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" From 9dc46f2af53703bebbad5df53819b0bbb88ff9ac Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Sun, 2 Aug 2026 03:16:58 +0000 Subject: [PATCH 2/3] test: cover early-break teardown under a straddling draw Answers a review question: does unpin_all still guard correctness when a partial epoch stops with two blocks pinned? Yes, and structurally it always did -- unpin_all does `self._pinned.clear()`, so the pinned count never mattered. But the existing regression used spc=8, block_chunks=2, batch_size=4: 16 rows per block, which divides evenly, so no batch in it ever straddles. Parametrized to add batch_size=12, which does. Two honest limits, stated in the docstring rather than left implied: - The straddling arm does NOT reach a pinned state the dividing arm cannot. Cutting batches over the epoch defers a block's release until a batch consumes its last row, but that extra hold spans only 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 this branch. So this is coverage that teardown works under a straddling draw, not a distinct leak. - Both arms now assert the precondition (`len(pool._pinned) > block_chunks`) before tearing down, so neither can pass on a single-block state that would not exercise the leak at all. Same trap as the epoch-boundary test earlier on this branch, which passed because unpin_all rescued it rather than because the release frontier drained. Co-Authored-By: Claude Opus 5 --- tests/test_prefetch.py | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) 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. From 1d5c083fa3f24b2f75dba138d76277de37f665a1 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Sun, 2 Aug 2026 04:11:02 +0000 Subject: [PATCH 3/3] test: hold the epoch-summary bound only where refcounts are exact CI's 3.13t job failed on `assert alloc1 <= 1` -- a warm pool allocating 2 more buffers in epoch 1. Not caused by this branch, and not corruption. Measured on a 3.13t build of this box, same scenario, 120 runs: buffers alloc(e1) runs 2 0 9 3 0 39 3 1 35 4 1 20 4 2 16 5 2 1 So the tight bound fails ~14% of the time under free-threading. Flake rate for the whole test, 75 runs per arm: main 5/75 (6.7%), this branch 8/75 (10.7%) -- a difference of 4% against a standard error near 4.6%, i.e. indistinguishable. It is pre-existing. The mechanism is the open question buffers.py already documents: sys.getrefcount may read stale under free-threading. Reading HIGH makes a free buffer look lent, so the pool allocates instead of reusing -- benign, and the direction the `refs < idle_refs` guard deliberately does not police because it is not the dangerous one. The dangerous direction (reading LOW -> reuse a live buffer) did not fire once: zero liveness RuntimeErrors across those 120 runs, and the pool settled at 2-5 buffers against 8 batches rather than climbing toward one each. So the assertion, not the pool, was wrong: it asserted a tightness the runtime never guaranteed off-GIL. Under the GIL it stays exact. Off-GIL it asserts what that 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 would make it a memory floor. Verified 0/30 failures on 3.13t after. `sys._is_gil_enabled` is 3.13+, so the check defaults to True on 3.12. This does NOT fix the underlying hazard, and the fix is not "impossible" as I first said. Polling `weakref.ref(lent_view)` alongside the refcount would veto reuse whenever the lent view is still held, even on a stale-low read, and the two are complementary rather than redundant (measured: a consumer holding only a crop kills the weakref -- numpy collapses base chains -- but keeps the owner refcount high). It is partial: a crop plus a stale-low read still slips through. Out of scope here, and it needs a reproducer first -- 0/120 means writing a fix for a hazard we cannot yet demonstrate. Tracked separately. Co-Authored-By: Claude Opus 5 --- tests/test_buffers.py | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/tests/test_buffers.py b/tests/test_buffers.py index bbfe212..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,6 +398,11 @@ 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]) # 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 @@ -408,7 +414,23 @@ def counts(line: str) -> tuple[int, int]: # 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: