Cut batches over the epoch, not per shuffle-block

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`).

docs/architecture.md

@@ -295,9 +295,10 @@ flowchart LR
     Q -.->|"full ⇒ producer blocks<br/>(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
 

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

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:

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``.
 

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: