From 0ce254e10c6f22ca332db40d6af13e7364cab9b5 Mon Sep 17 00:00:00 2001 From: David Stuebe <84335963+emfdavid@users.noreply.github.com> Date: Mon, 31 Aug 2026 03:14:51 +0000 Subject: [PATCH] docs: name concurrent iterations as the cause when the residency budget starves One InSituDataset owns one ChunkPool and every active iteration shares it, but each holds its own chunk references -- so residency is the sum of their working sets, not the maximum. Sizing the default for ONE iteration is the right default and stays: the engine cannot know how many iterations a caller intends to run, and guessing high would cost memory in the single-iteration case that is almost every case. What was missing is the diagnostic. Starvation 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 cannot see why. It now names how many iterations share the pool and the pattern that produces that. Owners count from mint to release rather than from their first pin: the iteration that starves BEFORE it can pin anything is exactly the one that needs naming, and counting pin-holders misses it (caught by the test, which failed on the first implementation). docs/tuning.md gains the sizing rule and the reason the default is not multiplied. The auto-sizing site and the docstrings that advertise the shared-pool configuration (buffers.BatchBuffers) carry the caveat, so it is not only in the prose. Refs #38 Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01WxdT3e62pYT3EVMBF1sA7C --- CHANGELOG.md | 14 +++++++++++ docs/tuning.md | 31 ++++++++++++++++++++++++ src/insitubatch/buffers.py | 4 ++++ src/insitubatch/pool.py | 22 +++++++++++++++++ src/insitubatch/scheduler.py | 13 +++++++++- src/insitubatch/source.py | 7 ++++++ tests/test_residency.py | 46 ++++++++++++++++++++++++++++++++++++ 7 files changed, 136 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b3df517..5ad5018 100644 --- a/CHANGELOG.md +++ b/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 — diff --git a/docs/tuning.md b/docs/tuning.md index 902d1bb..601fffe 100644 --- a/docs/tuning.md +++ b/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 diff --git a/src/insitubatch/buffers.py b/src/insitubatch/buffers.py index aa99d07..4d68796 100644 --- a/src/insitubatch/buffers.py +++ b/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. """ diff --git a/src/insitubatch/pool.py b/src/insitubatch/pool.py index 6128bfe..c956400 100644 --- a/src/insitubatch/pool.py +++ b/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: diff --git a/src/insitubatch/scheduler.py b/src/insitubatch/scheduler.py index 72dcdfe..a6e75f3 100644 --- a/src/insitubatch/scheduler.py +++ b/src/insitubatch/scheduler.py @@ -380,6 +380,17 @@ 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) " @@ -387,7 +398,7 @@ def _starvation(self, array: str, chunk_index: int) -> RuntimeError | None: 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: diff --git a/src/insitubatch/source.py b/src/insitubatch/source.py index 340e29f..ab5437f 100644 --- a/src/insitubatch/source.py +++ b/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; diff --git a/tests/test_residency.py b/tests/test_residency.py index ef5cf30..c34bf88 100644 --- a/tests/test_residency.py +++ b/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"