docs: name concurrent iterations as the cause when the residency budget starves

CHANGELOG.md

@@ -2,6 +2,20 @@
 
 ## Unreleased
 
+- **Under-sizing the budget for concurrent iterations now says so.** One `InSituDataset`
+  owns one chunk pool and every active iteration shares it — `zip(ds.train, ds.val)`, or
+  two `DataLoader`s — but each holds its *own* chunk references, so residency is the sum
+  of their working sets, not the maximum. The auto-sized default covers one iteration and
+  deliberately stays that way: the engine cannot know how many you intend to run, and
+  guessing high would cost memory in the single-iteration case that is almost every case.
+  What was missing was the diagnostic. Starvation previously advised "raise
+  cache_budget_bytes, or lower batch_size / block_chunks" — correct but not actionable
+  when *every* resident chunk is legitimately referenced and the caller has no way to see
+  why. It now names how many iterations are sharing the pool, and the pattern that
+  produces that. Owners count from mint to release rather than from their first pin,
+  because the iteration that starves before it can pin anything is exactly the one that
+  needs naming. `docs/tuning.md` gains the sizing rule.
+
 - **The chunk pool's "safe to take away" predicate is now true, not approximately true.**
   Eviction eligibility was spread across five loosely-coupled fields, and each one could
   lie. `fail()` had to set `ready = True` to wake a waiter — the only lever available —

docs/tuning.md

@@ -95,6 +95,37 @@ If you set `cache_budget_bytes` above the working set, residency rises to that b
 purpose — that extra memory *is* the cross-epoch cache. Point `cache_dir` at local NVMe to
 spill it to disk instead of RAM.
 
+### Several iterations at once multiply the budget
+
+The three bounds above describe **one** iteration. One `InSituDataset` owns one chunk pool,
+and every active iteration shares it — `zip(ds.train, ds.val)`, or two `DataLoader`s over the
+same dataset. That is supported, and chunks a windowed read pulls across a split boundary are
+decoded once and reused by both. But each iteration holds its **own** references to the chunks
+it is working on, so residency is the sum, not the maximum:
+
+```
+cache_budget_bytes  >=  n_concurrent_iterations x (block_chunks x outer_chunk_bytes)
+```
+
+**The auto-sized default is deliberately computed for one iteration, and stays that way.**
+The engine cannot know how many iterations you intend to run, so any automatic multiplier
+would be a guess that silently costs memory for the single-iteration case — which is almost
+every case. Running several is the explicit choice, so sizing for it is yours too.
+
+Run two without raising the budget and the loader stops with `residency budget exhausted: ...`,
+naming how many iterations are sharing the pool. It cannot free a slot, because every resident
+chunk is legitimately referenced by one of them. Raise `cache_budget_bytes` (or lower
+`block_chunks`, which shrinks each iteration's share) — not `max_inflight`, which is a
+concurrency dial and is not what is binding here.
+
+!!! note "This got stricter once pin accounting became owner-scoped"
+
+    Before that fix, starting a second iteration silently released the first one's
+    references. That freed budget by accident and hid the requirement — and the cost was
+    that the first iteration's in-use chunks could be evicted mid-gather, delivering
+    plausible wrong data. A budget that appeared to work for `zip(ds.train, ds.val)` was
+    relying on that bug. Sizing for the sum is the honest requirement.
+
 ### The batch queue is a high-water mark
 
 Batch buffers are pooled and held for the dataset's lifetime, so that third bound is the

src/insitubatch/buffers.py

@@ -246,6 +246,10 @@ class BatchBuffers:
     held only across that decision -- never across a gather -- so it costs one uncontended
     acquire per variable per batch, against a gather that copies megabytes.
 
+    Running two iterations also costs *residency*: each holds its own chunk references, so the
+    pool's budget must cover both working sets. The auto-sized default covers one -- see
+    docs/tuning.md, "Several iterations at once multiply the budget".
+
     Writing a lent buffer stays lock-free and always was: a buffer is lent to exactly one
     producer, which is the only thread that touches it.
     """

src/insitubatch/pool.py

@@ -393,6 +393,10 @@ def __init__(
         # single `claimed` bool lets one iteration's claim satisfy another's wait_ready.
         self._pinned: dict[tuple[str, int], dict[int, int]] = {}
         self._owner_seq = 0  # monotonic; owners are opaque tokens minted by new_owner()
+        # Owners live from mint to release_owner, NOT from their first pin: an iteration
+        # that starves before it can pin anything is exactly the case the starvation
+        # diagnostic has to name, and counting pin-holders would miss it.
+        self._owners: set[int] = set()
         self._cv = threading.Condition(threading.Lock())
         self._error: BaseException | None = None  # global poison (driver death)
         self.max_resident = 0  # peak distinct outer chunk positions held at once
@@ -414,6 +418,7 @@ def new_owner(self) -> int:
         """
         with self._cv:
             self._owner_seq += 1
+            self._owners.add(self._owner_seq)
             return self._owner_seq
 
     def _refs(self, key: tuple[str, int]) -> int:  # call under the lock
@@ -452,6 +457,22 @@ def budget_bytes(self) -> int | None:
         """The residency ceiling, or ``None`` for an unbounded pool."""
         return self._budget
 
+    @property
+    def active_owners(self) -> int:
+        """How many iterations are live on this pool right now.
+
+        Live means minted and not yet released -- *not* "currently holds a pin". An
+        iteration that starves before it can pin anything is precisely the case the
+        starvation diagnostic exists to name, and counting pin-holders would miss it.
+
+        More than one means several iterations share this pool (``zip(ds.train,
+        ds.val)``, two DataLoaders) and each needs its own working set resident at the
+        same time -- the most common reason a budget auto-sized for one cannot admit.
+        Read-only snapshot.
+        """
+        with self._cv:
+            return len(self._owners)
+
     def blocked_waiters(self) -> list[tuple[str, int]]:
         """``(path, chunk_index)`` keys some thread is currently blocked on in
         :meth:`wait_ready`.
@@ -663,6 +684,7 @@ def release_owner(self, owner: int) -> None:
                 if slot.state is not SlotState.READY and slot.quiescent and self._refs(k) == 0
             ]:
                 self._drop(key)  # an abandoned partial can never be a valid cache entry
+            self._owners.discard(owner)  # this iteration is done; it no longer counts
             self._cv.notify_all()  # freed budget may unpark an admission
 
     def reset_epoch_counters(self) -> None:

src/insitubatch/scheduler.py

@@ -380,14 +380,25 @@ def _starvation(self, array: str, chunk_index: int) -> RuntimeError | None:
         if not waiting:
             return None
         budget = self.pool.budget_bytes
+        # A budget sized for ONE iteration cannot serve several: each holds its own
+        # references, so the requirement multiplies. Name it rather than leave the caller
+        # to rediscover it -- it is the likeliest cause once more than one owner is live.
+        owners = self.pool.active_owners
+        concurrent = (
+            f" NOTE: {owners} iterations are sharing this pool (e.g. `zip(ds.train, "
+            f"ds.val)`, or two DataLoaders); each needs its own working set resident at "
+            f"once, so the budget must cover all {owners}."
+            if owners > 1
+            else ""
+        )
         return RuntimeError(
             f"residency budget exhausted: cannot admit chunk {chunk_index} of {array!r}. "
             f"The pool holds {self.pool.resident_chunks} chunk(s) "
             f"({self.pool.resident_bytes} of {budget} bytes), every one of them pinned or "
             f"in flight; no tile is in flight; and the consumer is blocked waiting on "
             f"{sorted(waiting)[:4]}. Nothing can free a slot, so this would hang. The "
             f"working set is larger than the budget it was sized for -- raise "
-            f"cache_budget_bytes, or lower batch_size / block_chunks."
+            f"cache_budget_bytes, or lower batch_size / block_chunks.{concurrent}"
         )
 
     async def _io(self, coro: Coroutine[Any, Any, _T]) -> _T:

src/insitubatch/source.py

@@ -257,6 +257,13 @@ def var_bytes(g: ArrayGeometry, o: ArrayGeometry, samples: int) -> int:
                 train_samples = len(self.manifest.chunks[SplitName.TRAIN.value]) * self._ref_spc
                 train_ws = sum(var_bytes(g, o, train_samples) for g, o in pairs)
                 working_set = max(working_set, train_ws)
+        # Sized for ONE iteration, deliberately. Every active iteration shares this pool and
+        # holds its own references, so N concurrent iterations need ~N x this -- but the engine
+        # cannot know N, and guessing high would cost memory in the single-iteration case that
+        # is almost every case. Running several is an explicit choice, so sizing for it is the
+        # caller's: pass `cache_budget_bytes`. Under-sizing is not silent -- admission raises
+        # and names how many iterations are sharing the pool (`Scheduler._starvation`).
+        # See docs/tuning.md, "Several iterations at once multiply the budget".
         self.cache_budget_bytes = max(int(cache_budget_bytes or 0), working_set)
         # persist turns the cache_dir mmap tier into a cross-run cache (files + manifest
         # survive close; reopen revives them as hits). It needs a dir to keep files in;

tests/test_residency.py

@@ -168,3 +168,49 @@ def drain(sched: Scheduler) -> int:
     with Scheduler(obstore_store(small_store), geoms, pool, SchedulerConfig()) as sched:
         sched.start(list(range(geom.n_chunks)), geom.sample_chunk_size)
         assert run_by(DEADLINE, lambda: drain(sched)) == geom.n_chunks
+
+
+def test_starvation_names_concurrent_iterations_as_the_cause(small_store, run_by) -> None:
+    """Under-sizing for concurrent iterations must say so, not just "raise the budget".
+
+    Sizing for one iteration is the deliberate default (the engine cannot know how many
+    you intend to run), so the *diagnostic* is what has to carry the cost. Every resident
+    chunk here is legitimately referenced -- by one of two iterations -- which is exactly
+    the case a caller cannot infer from "every one of them pinned or in flight".
+
+    Guards the number too: reporting a count means it has to be the count of distinct
+    owners, not of pins or of resident chunks.
+    """
+    geoms = open_geometries(obstore_store(small_store))
+    geom = geoms["t2m"]
+    manifest = split_by_chunk(geom, fractions=(1.0, 0.0, 0.0))
+    chunk_bytes = geom.sample_chunk_size * 2 * 2 * 4
+    ds = InSituDataset(
+        obstore_store(small_store),
+        manifest,
+        shuffle=False,
+        batch_size=geom.sample_chunk_size,
+        block_chunks=2,
+        cache_budget_bytes=2 * chunk_bytes,  # fits one iteration, not two
+    )
+    ds.set_epoch(0)
+
+    def interleave() -> None:
+        a, b = iter(ds.train), iter(ds.train)
+        try:
+            for _ in range(geom.n_chunks):
+                next(a)
+                next(b)
+        finally:
+            for it in (a, b):
+                it.close()
+
+    with pytest.raises(RuntimeError, match="residency budget exhausted") as exc:
+        run_by(DEADLINE, interleave)
+
+    msg = str(exc.value)
+    assert "2 iterations are sharing this pool" in msg, (
+        "the diagnostic must name concurrent iterations as the cause -- every chunk is "
+        f"legitimately referenced, so 'raise the budget' alone is not actionable: {msg}"
+    )
+    assert "zip(ds.train, ds.val)" in msg, "point at the pattern that produces this"