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"