From 75fd18b4bd41688413b0100d7a6333db80ce59d8 Mon Sep 17 00:00:00 2001 From: Vincent Gong Date: Wed, 12 Aug 2026 08:43:27 +0200 Subject: [PATCH] =?UTF-8?q?Add=20chunk=5Fmemory=5Fmode:=20auto=20=E2=80=94?= =?UTF-8?q?=20container-memory-aware=20adaptive=20chunking?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adaptive chunking sizes chunks against a static `chunk_size` byte budget that the user must hand-tune per machine, and it targets host RAM rather than the container's cgroup limit — so on memory-limited containers it can over-commit and get OOM-killed. This adds an opt-in `chunk_memory_mode: auto` that derives the chunk memory budget from the process's real memory ceiling at runtime and hardens the adaptive sizing. When `chunk_memory_mode: auto`: - Budget = (memory_limit - current usage) * chunk_memory_safety_factor, where memory_limit is read from the cgroup (v2 memory.max, then v1, then psutil) — the limit that actually OOM-kills the process. No machine-specific chunk_size to tune. - Multiprocess: the shared budget is divided by the real per-step worker count (read from the num_processes injectable); num_processes:0 auto-derives the worker count from the available non-reclaimable memory / chunk_worker_target_budget. - Chunks are sized from the budget vs the chunk's INCREMENTAL memory (not absolute rss), so a large memory-mapped shared skim buffer (reclaimable page cache) does not pollute the measurement and collapse chunks to a single row. - A budget floor keeps chunking active under pressure (a zero budget would disable chunking and process all choosers at once); a growth cap and an incremental peak-backoff bound overshoot; an optional circuit-breaker warns as memory nears the limit. All new behavior is gated on `chunk_memory_mode` (default `fixed` = current behavior), so existing runs are unaffected. Adds unit tests (core/test/test_mem.py, test_chunk_robust.py). --- activitysim/core/chunk.py | 268 +++++++++++++++++++-- activitysim/core/configuration/top.py | 64 +++++ activitysim/core/mem.py | 194 +++++++++++++++ activitysim/core/mp_tasks.py | 28 ++- activitysim/core/test/test_chunk_robust.py | 120 +++++++++ activitysim/core/test/test_mem.py | 83 +++++++ docs/core.rst | 31 +++ 7 files changed, 765 insertions(+), 23 deletions(-) create mode 100644 activitysim/core/test/test_chunk_robust.py create mode 100644 activitysim/core/test/test_mem.py diff --git a/activitysim/core/chunk.py b/activitysim/core/chunk.py index f0074683d8..4dc1f247a0 100644 --- a/activitysim/core/chunk.py +++ b/activitysim/core/chunk.py @@ -22,6 +22,13 @@ logger = logging.getLogger(__name__) + +class ChunkMemoryOverflow(Exception): + """Raised by the memory circuit breaker to abort an in-flight chunk that is approaching the + process memory ceiling, so the adaptive loop can retry it at a smaller size (see the + ``chunk_memory_circuit_breaker`` setting). Reserved for the Phase-2 circuit breaker.""" + + # # CHUNK_METHODS and METRICS # @@ -88,6 +95,11 @@ MODE_PRODUCTION = "production" MODE_CHUNKLESS = "disabled" MODE_EXPLICIT = "explicit" + +# Smallest auto-mode chunk budget (bytes). Keeps chunking active under memory pressure: a budget of 0 +# means chunk_size == 0, which disables chunking (process all choosers at once) and OOMs on big samples. +AUTO_BUDGET_FLOOR = 1_000_000_000 + TRAINING_MODES = [ MODE_RETRAIN, MODE_ADAPTIVE, @@ -185,6 +197,73 @@ def get_base_chunk_size(state: workflow.State): return state.chunk.CHUNK_SIZERS[0].chunk_size +def resolve_chunk_size(state: workflow.State) -> int: + """Memory budget (bytes) for adaptive chunking. + + Legacy ``chunk_memory_mode='fixed'`` returns the static ``chunk_size`` setting. ``'auto'`` derives + the budget from the process's REAL memory ceiling — the cgroup limit inside a container (what + actually OOM-kills us), else host RAM — scaled by ``chunk_memory_safety_factor``. The existing + ``available_headroom`` logic (budget − current rss) then automatically discounts memory already + resident (framework + shared skims), so no separate baseline subtraction is needed here. + """ + if getattr(state.settings, "chunk_memory_mode", "fixed") != "auto": + return state.settings.chunk_size + + limit = mem.get_memory_limit() + if not limit: + logger.warning( + "chunk_memory_mode=auto but the memory limit could not be determined; " + f"falling back to static chunk_size {GB(state.settings.chunk_size)}" + ) + return state.settings.chunk_size + + safety = state.settings.chunk_memory_safety_factor + # Base the budget on (limit − current usage). The shared skim set is memory-mapped from disk + # (reclaimable page cache), but the pages a chunk is actively reading are momentarily in-use, so the + # skim working set is a REAL transient cost that scales with chunk size. Counting currently-resident + # memory (which includes the resident skim cache) as used therefore keeps the budget honest: it + # shrinks as more skim pages fault in, tracking the working set — whereas basing it on non-reclaimable + # memory alone ignores the in-use skim and over-budgets (→ OOM on heavy models). Falls back to the raw + # limit if usage can't be read. (Conservative: also counts idle reclaimable cache as used, but the + # safety_factor margin + the kernel reclaiming under pressure keep that from starving throughput.) + available = mem.get_available_memory() + basis = max(0, min(limit, available)) if available else limit + budget = int(basis * safety) + # Multiprocess: N workers share the single memory ceiling, so each worker's budget is the shared free + # memory divided by the worker count — otherwise every worker would size chunks to the full budget + # and they'd collectively OOM. + num_processes = 1 + if getattr(state.settings, "multiprocess", False): + # In a worker subprocess the REAL per-step worker count is the 'num_processes' injectable that + # mp_tasks sets per step (state.settings.num_processes is 0 when the count is auto-derived, and a + # worker reading it would divide by 1 and size chunks to the FULL budget → the workers collectively + # overshoot the shared ceiling). Fall back to the setting outside a worker (parent/non-mp). + try: + injected = state.get_injectable("num_processes", None) + except Exception: + injected = None + num_processes = injected or getattr(state.settings, "num_processes", 1) or 1 + if num_processes > 1: + budget = int(budget / num_processes) + # Never let the auto budget collapse to ~0: chunk_size == 0 disables chunking entirely (activitysim + # then processes ALL choosers in one chunk), which OOMs on a large sample. Keep a small positive + # floor so chunking stays active with tiny-but-safe chunks even under severe memory pressure. + budget = max(budget, AUTO_BUDGET_FLOOR) + baseline, _ = mem.get_rss(force_garbage_collect=True) + logger.info( + f"chunk_memory_mode=auto: limit={GB(limit)} available={GB(available) if available else 'n/a'}" + f" x safety_factor={safety}" + f"{f' / {num_processes} workers' if num_processes > 1 else ''} " + f"-> base_chunk_size={GB(budget)} (current rss={GB(baseline)})" + ) + if budget <= baseline: + logger.warning( + f"auto chunk budget {GB(budget)} <= current rss {GB(baseline)}: memory is already tight; " + "chunking will proceed with minimal chunks." + ) + return budget + + def overhead_for_chunk_method(state: workflow.State, overhead, method=None): """ @@ -508,8 +587,17 @@ def audit( if not self.base_chunk_size: return + auto = getattr(state.settings, "chunk_memory_mode", "fixed") == "auto" mem_panic_threshold = self.base_chunk_size * (1 + MAX_OVERDRAFT) - bytes_panic_threshold = self.headroom + (self.base_chunk_size * MAX_OVERDRAFT) + # In auto mode the budget (base_chunk_size) is the available-memory allowance and chunk 'bytes' + # are tracked incrementally, so compare tracked bytes against the budget directly. The legacy + # headroom-based threshold subtracts absolute xss, which the shared skim buffer pollutes in + # multiprocess. + bytes_panic_threshold = ( + mem_panic_threshold + if auto + else self.headroom + (self.base_chunk_size * MAX_OVERDRAFT) + ) if bytes > bytes_panic_threshold: logger.warning( @@ -517,21 +605,25 @@ def audit( f"bytes: {bytes} headroom: {self.headroom} chunk_size: {self.base_chunk_size} {msg}" ) - if chunk_metric(state) == RSS and rss > mem_panic_threshold: - rss, _ = mem.get_rss(force_garbage_collect=True, uss=False) - if rss > mem_panic_threshold: - logger.warning( - f"out_of_chunk_memory: " - f"rss: {rss} chunk_size: {self.base_chunk_size} {msg}" - ) + # Absolute rss/uss include memory-mapped SHARED skim pages (tens of GB in multiprocess), so in + # auto mode they are not a meaningful per-chunk overflow signal and would spam false warnings — + # real OOM pressure is handled by the cgroup memory watchdog + the incremental peak-backoff. + if not auto: + if chunk_metric(state) == RSS and rss > mem_panic_threshold: + rss, _ = mem.get_rss(force_garbage_collect=True, uss=False) + if rss > mem_panic_threshold: + logger.warning( + f"out_of_chunk_memory: " + f"rss: {rss} chunk_size: {self.base_chunk_size} {msg}" + ) - if chunk_metric(state) == USS and uss > mem_panic_threshold: - _, uss = mem.get_rss(force_garbage_collect=True, uss=True) - if uss > mem_panic_threshold: - logger.warning( - f"out_of_chunk_memory: " - f"uss: {uss} chunk_size: {self.base_chunk_size} {msg}" - ) + if chunk_metric(state) == USS and uss > mem_panic_threshold: + _, uss = mem.get_rss(force_garbage_collect=True, uss=True) + if uss > mem_panic_threshold: + logger.warning( + f"out_of_chunk_memory: " + f"uss: {uss} chunk_size: {self.base_chunk_size} {msg}" + ) def close(self): logger.debug(f"ChunkLedger.close trace_label: {self.trace_label}") @@ -671,6 +763,37 @@ def get_hwm_bytes(self): return self.hwm_bytes["value"] +# per-process memory-watchdog state (per-process is correct: each MP worker watches its own rss) +_watchdog_breached = False + + +def memory_watchdog_check(state: workflow.State, rss: int, trace_label: str): + """When the circuit breaker is enabled, warn (loudly, once per breach) as rss approaches the hard + memory ceiling, so the danger is visible before the OS OOM-killer fires. The actual protection is + the proactive AIMD back-off in ``adaptive_rows_per_chunk``; a safe mid-flight abort-and-retry is + tracked as future work (it needs the generator/consumer contract reworked).""" + global _watchdog_breached + if not getattr(state.settings, "chunk_memory_circuit_breaker", False): + return + limit = mem.get_memory_limit() + if not limit: + return + # use the larger of the sampled rss and the kernel's EXACT lifetime peak (ru_maxrss): a short spike + # the periodic sampler missed still trips the warning before the OOM-killer would. + rss = max(rss, mem.get_peak_rss()) + abort_ratio = getattr(state.settings, "chunk_memory_abort_ratio", 0.9) + threshold = abort_ratio * limit + if rss >= threshold and not _watchdog_breached: + _watchdog_breached = True + logger.warning( + f"MEMORY WATCHDOG: rss {GB(rss)} reached {abort_ratio:.0%} of the {GB(limit)} memory " + f"limit at {trace_label}. Adaptive sizing will back off; if pressure persists, lower " + f"chunk_memory_safety_factor or set an explicit_chunk to avoid an OOM kill." + ) + elif rss < 0.8 * threshold and _watchdog_breached: + _watchdog_breached = False # recovered → re-arm the warning + + def log_rss(state: workflow.State, trace_label: str, force=False): if chunk_training_mode(state) == MODE_CHUNKLESS: # no memory tracing at all in chunkless mode @@ -690,6 +813,8 @@ def log_rss(state: workflow.State, trace_label: str, force=False): rss, uss = mem.trace_memory_info(hwm_trace_label, state=state) + memory_watchdog_check(state, rss, hwm_trace_label) + # check local hwm for all ledgers with state.chunk.ledger_lock: for c in state.chunk.CHUNK_LEDGERS: @@ -706,8 +831,13 @@ def __init__( threading.Thread.__init__(self) def run(self): + # Sample finer when the circuit breaker is armed so the watchdog catches a fast-rising spike + # well before the OOM-killer; otherwise keep the default (cheaper) cadence. + tick = mem.MEM_SNOOP_TICK_LEN + if getattr(self.state.settings, "chunk_memory_circuit_breaker", False): + tick = min(tick, 0.5) log_rss(self.state, self.trace_label) - while not self.stop_snooping.wait(timeout=mem.MEM_SNOOP_TICK_LEN): + while not self.stop_snooping.wait(timeout=tick): log_rss(self.state, self.trace_label) @@ -837,10 +967,38 @@ def available_headroom(self, xss): f"base_chunk_size: {util.INT(self.base_chunk_size)}" ) + # In 'auto' mode, prefer to SHRINK the chunk (down to the real headroom) rather than force + # a min_chunk_size chunk that can exceed real memory and OOM — since the budget already + # tracks the true ceiling. rows_per_chunk clips to >= 1 downstream, so progress continues + # with tiny chunks under pressure (slow but safe). Legacy 'fixed' mode keeps the old floor. + if getattr(self.state.settings, "chunk_memory_mode", "fixed") == "auto": + return max(headroom, 0) + headroom = self.min_chunk_size return headroom + def sizing_budget(self): + """Memory basis for CHOOSING rows-per-chunk (auto mode) — the deterministic 'auto-explicit' path. + + Legacy adaptive sizing divides ``available_headroom = base_chunk_size - xss`` by the row_size. + But ``xss`` (per-worker rss/uss) counts memory-mapped SHARED skim pages: in a multiprocess run + the ~38 GB skim buffer lives in /dev/shm and is mapped into every worker, so each worker's xss + is inflated by tens of GB. Subtracting it collapses the headroom to ~0 and forces 1-row chunks + (the crawl we observed), even though the chunk's real marginal cost is small. + + The correct basis is the BUDGET itself: ``base_chunk_size`` is already derived (in + resolve_chunk_size) from AVAILABLE memory = cgroup limit − current usage, so the resident shared + skims + framework are ALREADY excluded. Chunk overhead is accounted incrementally (hwm − prev), + so rows = budget / incremental_row_size is the right, shared-memory-immune sizing — deterministic + like explicit chunking, but with the size computed automatically from the real budget. The + growth-cap + incremental peak-backoff in adaptive_rows_per_chunk remain the safety net against an + under-estimated first row_size. Legacy 'fixed' mode keeps the original headroom-based sizing. + """ + if getattr(self.state.settings, "chunk_memory_mode", "fixed") == "auto": + return self.base_chunk_size + return self.headroom + def initial_rows_per_chunk(self): if self.chunk_training_mode == MODE_EXPLICIT: if self.rows_per_chunk: @@ -868,8 +1026,11 @@ def initial_rows_per_chunk(self): ), f"len(state.chunk.CHUNK_LEDGERS): {len(self.state.chunk.CHUNK_LEDGERS)}" if self.initial_row_size > 0: + margin = ( + getattr(self.state.settings, "chunk_row_size_margin", 1.0) or 1.0 + ) max_rows_per_chunk = np.maximum( - int(self.headroom / self.initial_row_size), 1 + int(self.sizing_budget() / (self.initial_row_size * margin)), 1 ) rows_per_chunk = np.clip(max_rows_per_chunk, 1, self.num_choosers) estimated_number_of_chunks = math.ceil( @@ -894,6 +1055,28 @@ def initial_rows_per_chunk(self): f"'production' but initial_row_size is zero in {self.trace_label}" ) + # safe-start (auto mode): the first chunk is the biggest gamble — it is sized from a CACHED + # row_size that may be stale/optimistic on a different machine or sample, and if it's too large + # it OOMs before any live measurement exists. Instead probe with a small first chunk (the + # default rows), measure its real memory, then let the adaptive sizing + growth-cap ramp up. + # Costs one small extra chunk; removes the first-chunk OOM gamble. + if ( + getattr(self.state.settings, "chunk_memory_mode", "fixed") == "auto" + and self.chunk_size > 0 + and self.initial_row_size > 0 + ): + probe = min( + self.num_choosers, self.state.settings.default_initial_rows_per_chunk + ) + if probe < rows_per_chunk: + logger.info( + f"{self.trace_label}: safe-start first chunk {int(rows_per_chunk)} -> {probe} rows " + f"(probe, then adapt)" + ) + rows_per_chunk = probe + self.initial_row_size = 0 # force re-measuring row_size from the probe + estimated_number_of_chunks = None + # cum_rows is out of phase with cum_overhead # since we won't know observed_chunk_size until AFTER yielding the chunk self.rows_per_chunk = rows_per_chunk @@ -965,11 +1148,54 @@ def adaptive_rows_per_chunk(self, i): # rows_per_chunk is closest number of chooser rows to achieve chunk_size without exceeding it if observed_row_size > 0: - self.rows_per_chunk = int(self.headroom / observed_row_size) + # inflate the (typically under-estimated) row size by the safety margin for a memory buffer + margin = getattr(self.state.settings, "chunk_row_size_margin", 1.0) or 1.0 + self.rows_per_chunk = int( + self.sizing_budget() / (observed_row_size * margin) + ) else: # they don't appear to have used any memory; increase cautiously in case small sample size was to blame self.rows_per_chunk = 2 * prev_rows_per_chunk + # --- robust adaptive sizing (AIMD): cap growth + back off if we neared the ceiling ---------- + # Adaptive chunking's classic OOM comes from (a) leaping to an over-large chunk by extrapolating + # a small, unrepresentative first chunk linearly, and (b) growing each chunk toward the budget + # even after a chunk already peaked dangerously close to it. Two guards address both, using only + # data already measured (no mid-flight interruption needed): + settings = self.state.settings + growth_cap = getattr(settings, "chunk_growth_cap", 0) or 0 + if growth_cap and prev_rows_per_chunk > 0: + capped = int(growth_cap * prev_rows_per_chunk) + if capped < self.rows_per_chunk: + logger.debug( + f"{self.trace_label}: growth-capped next chunk " + f"{self.rows_per_chunk} -> {capped} rows (<= {growth_cap}x prev)" + ) + self.rows_per_chunk = capped + + if ( + getattr(settings, "chunk_memory_mode", "fixed") == "auto" + and self.chunk_training_mode != MODE_PRODUCTION + and self.base_chunk_size > 0 + and prev_rows_per_chunk > 0 + ): + # Compare the chunk's INCREMENTAL peak (hwm − prev, this chunk's own marginal memory) to the + # budget — NOT absolute rss, which in multiprocess includes the shared skim pages and would + # trip the back-off on every chunk (collapsing to 1-row chunks). overhead[] is incremental. + peak_incremental = ( + overhead[USS] if chunk_metric(self.state) == USS else overhead[RSS] + ) + abort_ratio = getattr(settings, "chunk_memory_abort_ratio", 0.9) + if peak_incremental > abort_ratio * self.base_chunk_size: + backed_off = max(1, prev_rows_per_chunk // 2) + if backed_off < self.rows_per_chunk: + logger.warning( + f"{self.trace_label}: chunk peaked at {GB(peak_incremental)} (incremental) " + f"(> {abort_ratio:.0%} of budget {GB(self.base_chunk_size)}); " + f"backing off next chunk {self.rows_per_chunk} -> {backed_off} rows" + ) + self.rows_per_chunk = backed_off + self.rows_per_chunk = np.clip(self.rows_per_chunk, 1, rows_remaining) self.rows_processed += self.rows_per_chunk estimated_number_of_chunks = ( @@ -1245,7 +1471,7 @@ def adaptive_chunked_choosers( else: chunk_size = math.ceil(explicit_chunk_size / num_processes) elif chunk_size is None: - chunk_size = state.settings.chunk_size + chunk_size = resolve_chunk_size(state) assert num_choosers > 0 assert chunk_size >= 0 @@ -1388,7 +1614,7 @@ def adaptive_chunked_choosers_and_alts( else: chunk_size = int(explicit_chunk_size / num_processes) elif chunk_size is None: - chunk_size = state.settings.chunk_size + chunk_size = resolve_chunk_size(state) chunk_sizer = ChunkSizer( state, @@ -1496,7 +1722,7 @@ def adaptive_chunked_choosers_by_chunk_id( if state.settings.chunk_training_mode == MODE_EXPLICIT: chunk_size = explicit_chunk_size else: - chunk_size = state.settings.chunk_size + chunk_size = resolve_chunk_size(state) chunk_sizer = ChunkSizer( state, chunk_tag, diff --git a/activitysim/core/configuration/top.py b/activitysim/core/configuration/top.py index 6b1c8a6c54..a46f645f91 100644 --- a/activitysim/core/configuration/top.py +++ b/activitysim/core/configuration/top.py @@ -399,6 +399,70 @@ class Settings(PydanticBase, extra="allow", validate_assignment=True): minimum fraction of total chunk_size to reserve for adaptive chunking """ + chunk_memory_mode: Literal["fixed", "auto"] = "fixed" + """ + How the adaptive chunker derives its memory budget (``base_chunk_size``). + + * "fixed" (default, legacy behavior) + Use the static :ref:`chunk_size` setting as the memory budget. + * "auto" + Ignore the static ``chunk_size`` and derive the budget at runtime from the process's real + memory ceiling: ``(memory_limit - baseline) * chunk_memory_safety_factor``, where + ``memory_limit`` is the cgroup limit (a container/pod limit) or, if not containerized, host + RAM, and ``baseline`` is the resident memory already in use when chunking begins (framework + + shared skims). This targets the memory that will actually OOM-kill the process instead of a + hand-tuned number, and — crucially — respects a container memory limit that psutil can't see. + """ + + chunk_memory_safety_factor: float = 0.75 + """ + Fraction of ``(memory_limit - baseline)`` to use as the chunking budget under + ``chunk_memory_mode: auto``. Leaves headroom for measurement lag and transient peaks. + """ + + chunk_memory_circuit_breaker: bool = False + """ + (Reserved — adaptive circuit breaker.) When True, a memory watchdog aborts and retries a chunk at + a smaller size before the OS OOM-killer fires. Off by default. + """ + + chunk_memory_abort_ratio: float = 0.9 + """ + Fraction of the memory ceiling at which the circuit breaker aborts the in-flight chunk. + """ + + chunk_growth_cap: float = 0.0 + """ + Maximum multiplicative growth of rows-per-chunk from one chunk to the next under adaptive sizing + (e.g. 2.0 = at most double each step). 0 disables the cap (legacy behavior). Prevents a single + over-large jump when extrapolating a big chunk from a small, unrepresentative first chunk. + """ + + chunk_row_size_margin: float = 1.0 + """ + Safety multiplier applied to the estimated per-row memory when sizing chunks (>= 1.0, default 1.0 = + off). Because the row-size estimate is a sampled, linear extrapolation that tends to UNDER-estimate + a large chunk's true transient peak, inflating it (e.g. 1.3) sizes chunks ~30% smaller for a memory + safety buffer. This is the conservative-cache lever: it makes a re-used chunk_cache err toward + smaller, safe chunks. Bytes-per-row is machine-independent, and the auto budget already adapts the + ceiling to the actual machine/container at runtime, so the cache does not need machine-specific keys. + """ + + chunk_worker_target_budget: int = 0 + """ + Target per-worker chunk memory budget in bytes for MEMORY-AWARE AUTO WORKER COUNT (0 = off). When + ``num_processes`` is unset and ``chunk_memory_mode: auto``, the worker count is derived from the + non-reclaimable memory headroom instead of a cpu-only heuristic:: + + num_processes = clamp( (memory_limit − nonreclaimable) * chunk_memory_safety_factor + // chunk_worker_target_budget, 1, cpu_count ) + + i.e. the most workers whose per-worker share of the real (non-reclaimable) memory still meets this + budget. Prevents oversubscribing workers a memory-constrained node can't feed (which just OOMs). + Reclaimable skim page cache is intentionally not subtracted (the kernel evicts it under pressure). + Set it to the memory a single worker needs for the heaviest model's chunks (e.g. 10-12 GB). + """ + checkpoints: Union[bool, list] = True """ When to write checkpoint (intermediate table states) to disk. diff --git a/activitysim/core/mem.py b/activitysim/core/mem.py index fbfa11fbe5..db04045000 100644 --- a/activitysim/core/mem.py +++ b/activitysim/core/mem.py @@ -8,6 +8,7 @@ import logging import multiprocessing import os +import sys import threading import time @@ -17,6 +18,14 @@ from activitysim.core import config, util, workflow +try: + import resource # Unix-only (getrusage); not available on Windows +except ImportError: + resource = None + +# high-water mark for get_peak_rss's Windows fallback (kept monotonic) +_PEAK_RSS_FALLBACK = 0 + logger = logging.getLogger(__name__) USS = True @@ -282,6 +291,191 @@ def get_rss(force_garbage_collect=False, uss=False): return info.rss, 0 +# --- real memory-ceiling introspection (cgroup-aware) ---------------------------------------------- +# psutil reports the HOST's RAM, which is wrong inside a container: the process is bounded by its +# cgroup memory limit (e.g. a Kubernetes pod limit), not the node's total RAM. Chunk sizing that +# targets host RAM will overshoot the cgroup limit and get OOM-killed. These helpers read the real +# ceiling from the cgroup (v2, then v1), falling back to psutil when not containerized/unlimited. + +# cgroup "unlimited" is reported as a huge sentinel; treat anything at/above it as no-limit. +_CGROUP_UNLIMITED = 0x7FFFFFFFFFFFF000 # ~9.2e18 + + +def _read_cgroup_file(path): + try: + with open(path) as f: + return f.read().strip() + except OSError: + return None + + +def _finite_limit(raw): + """Parse a cgroup limit string; return an int only if it is a real finite limit.""" + if raw is None or raw == "max": + return None + try: + n = int(raw) + except (TypeError, ValueError): + return None + return n if 0 < n < _CGROUP_UNLIMITED else None + + +def get_memory_limit(cgroup_root: str = "/sys/fs/cgroup") -> int | None: + """This process's hard memory ceiling in bytes, or None if it can't be determined. + + Prefers the cgroup limit (what actually OOM-kills us in a container) over host RAM. Tries cgroup + v2 (``memory.max``), then cgroup v1 (``memory/memory.limit_in_bytes``), then psutil total RAM. + """ + limit = _finite_limit(_read_cgroup_file(os.path.join(cgroup_root, "memory.max"))) + if limit is not None: + return limit + for rel in ("memory/memory.limit_in_bytes", "memory.limit_in_bytes"): + limit = _finite_limit(_read_cgroup_file(os.path.join(cgroup_root, rel))) + if limit is not None: + return limit + try: + return int(psutil.virtual_memory().total) + except Exception: + return None + + +def get_available_memory(cgroup_root: str = "/sys/fs/cgroup") -> int | None: + """Best-effort bytes still available before this process hits its ceiling. + + Uses (cgroup limit - cgroup current usage) when containerized, else psutil available RAM. Note + cgroup ``memory.current`` counts reclaimable page cache as used, so this under-estimates the truly + available memory — which is the safe direction for chunk sizing (errs toward smaller chunks). + """ + limit = get_memory_limit(cgroup_root) + used = None + raw = _read_cgroup_file(os.path.join(cgroup_root, "memory.current")) + if raw is not None: + try: + used = int(raw) + except ValueError: + used = None + if used is None: + for rel in ("memory/memory.usage_in_bytes", "memory.usage_in_bytes"): + raw = _read_cgroup_file(os.path.join(cgroup_root, rel)) + if raw is not None: + try: + used = int(raw) + break + except ValueError: + used = None + if limit is not None and used is not None: + return max(0, limit - used) + try: + return int(psutil.virtual_memory().available) + except Exception: + return limit + + +def get_nonreclaimable_used(cgroup_root: str = "/sys/fs/cgroup") -> int | None: + """Bytes of NON-reclaimable memory currently charged to this cgroup, or None. + + = anonymous memory (worker private allocations — chunk DataFrames/arrays) + pinned shared memory + (tmpfs/shm). This is the memory the kernel CANNOT reclaim, so it is the true hard constraint for + the chunk budget. It deliberately EXCLUDES file-backed page cache: an ActivitySim run memory-maps + the skim set from on-disk .mmap files, so the (tens of GB) resident skims are reclaimable page + cache that the kernel evicts under pressure — counting them as 'used' (as limit − memory.current + does) needlessly starves the budget. Basing the budget on (limit − nonreclaimable) reflects the + real room for new chunk allocations, and it still correctly accounts for genuinely-pinned shm. + Reads cgroup v2 ``memory.stat`` (``anon`` + ``shmem``); falls back to v1 field names. + """ + fields = {} + for rel in ("memory.stat", "memory/memory.stat"): + raw = _read_cgroup_file(os.path.join(cgroup_root, rel)) + if raw: + for line in raw.splitlines(): + parts = line.split() + if len(parts) == 2: + try: + fields[parts[0]] = int(parts[1]) + except ValueError: + pass + break + if not fields: + return None + anon = fields.get("anon", fields.get("rss", 0)) + shmem = fields.get("shmem", fields.get("total_shmem", 0)) + return anon + shmem + + +def recommend_num_processes( + limit: int | None, + nonreclaimable: int | None, + safety: float, + target_per_worker: int, + cpu_count: int, +) -> int | None: + """Memory-aware multiprocess worker count, or None if it can't/shouldn't be computed. + + The most workers whose per-worker share of the non-reclaimable memory headroom still meets the + target per-worker chunk budget:: + + N = clamp( (limit − nonreclaimable) * safety // target_per_worker, 1, cpu_count ) + + Evaluated once at startup (before workers fork), so ``nonreclaimable`` ≈ the parent framework and + ``limit − nonreclaimable`` ≈ the full headroom the workers will share. Reclaimable skim page cache + is (correctly) NOT subtracted. Returns None when target_per_worker or limit is unset, so the caller + keeps its legacy cpu-based default. + """ + if not (limit and target_per_worker and cpu_count): + return None + available = max(0, limit - (nonreclaimable or 0)) + n = int((available * safety) // target_per_worker) + return max(1, min(n, cpu_count)) + + +def get_resident_shm(cgroup_root: str = "/sys/fs/cgroup") -> int | None: + """Bytes of shared/tmpfs memory (``/dev/shm`` etc.) currently charged to this cgroup, or None. + + In a multiprocess run this is the portion of the shared skim buffer that has already paged in. + Read from cgroup ``memory.stat`` (the ``shmem`` field). Used to reserve only the NOT-yet-resident + remainder of the known shared-skim footprint from the chunk budget (subtracting the already-resident + part would double-count, since get_available_memory already excludes it). NOTE: do NOT use + os.statvfs('/dev/shm') for this — inside a container it reports the node-level tmpfs size, not the + pod's emptyDir sizeLimit, which massively over-reserves and collapses the budget. + """ + for rel in ("memory.stat", "memory/memory.stat"): + raw = _read_cgroup_file(os.path.join(cgroup_root, rel)) + if raw: + for line in raw.splitlines(): + parts = line.split() + if len(parts) == 2 and parts[0] in ("shmem", "total_shmem"): + try: + return int(parts[1]) + except ValueError: + pass + return None + + +def get_peak_rss() -> int: + """Exact lifetime peak RSS of this process in bytes, from the kernel (``getrusage`` ru_maxrss). + + Unlike the MemMonitor's periodically-sampled high-water mark, this never misses a short-lived + transient allocation spike (a common cause of adaptive chunking under-estimating a chunk's true + peak). Linux reports ru_maxrss in kilobytes; macOS/BSD report bytes. On Windows (no ``resource`` + module) there is no getrusage peak, so this tracks a sampled high-water mark of the current RSS, + which keeps it monotonic non-decreasing like the real peak.""" + if resource is None: + # Windows: no getrusage; track a sampled high-water mark of current RSS so the result stays + # monotonic non-decreasing. + global _PEAK_RSS_FALLBACK + try: + rss = int(psutil.Process().memory_info().rss) + except Exception: + return _PEAK_RSS_FALLBACK + _PEAK_RSS_FALLBACK = max(_PEAK_RSS_FALLBACK, rss) + return _PEAK_RSS_FALLBACK + try: + maxrss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + except (ValueError, OSError): + return 0 + return int(maxrss) * 1024 if sys.platform.startswith("linux") else int(maxrss) + + def shared_memory_size(data_buffers): """ return total size of the multiprocessing shared memory block in data_buffers diff --git a/activitysim/core/mp_tasks.py b/activitysim/core/mp_tasks.py index c8507138f4..662a0c6c47 100644 --- a/activitysim/core/mp_tasks.py +++ b/activitysim/core/mp_tasks.py @@ -1925,8 +1925,32 @@ def get_run_list(state: workflow.State): # default settings that can be overridden by settings in individual steps global_chunk_size = state.settings.chunk_size - default_mp_processes = state.settings.num_processes or int( - 1 + multiprocessing.cpu_count() / 2.0 + # Memory-aware auto worker count: when num_processes is unset and chunk_memory_mode is 'auto' with a + # per-worker budget target, derive the worker count from the non-reclaimable memory headroom instead + # of the cpu-only heuristic (more workers than the RAM can feed just OOMs). Falls through to the + # legacy cpu-based default when not configured or not computable. + auto_mp_processes = None + if ( + not state.settings.num_processes + and getattr(state.settings, "chunk_memory_mode", "fixed") == "auto" + and getattr(state.settings, "chunk_worker_target_budget", 0) + ): + auto_mp_processes = mem.recommend_num_processes( + mem.get_memory_limit(), + mem.get_nonreclaimable_used(), + getattr(state.settings, "chunk_memory_safety_factor", 0.75) or 0.75, + state.settings.chunk_worker_target_budget, + multiprocessing.cpu_count(), + ) + if auto_mp_processes: + logger.info( + f"auto num_processes = {auto_mp_processes} " + f"(from non-reclaimable memory headroom / chunk_worker_target_budget)" + ) + default_mp_processes = ( + state.settings.num_processes + or auto_mp_processes + or int(1 + multiprocessing.cpu_count() / 2.0) ) if multiprocess and multiprocessing.cpu_count() == 1: diff --git a/activitysim/core/test/test_chunk_robust.py b/activitysim/core/test/test_chunk_robust.py new file mode 100644 index 0000000000..0f8faef543 --- /dev/null +++ b/activitysim/core/test/test_chunk_robust.py @@ -0,0 +1,120 @@ +# ActivitySim +# See full license in LICENSE.txt. +"""Tests for robust adaptive chunking: the real-memory-ceiling budget (chunk_memory_mode=auto), the +AIMD growth cap / back-off, and the memory watchdog. See also test_mem.py for the cgroup helpers.""" +from __future__ import annotations + +import logging +import os + +import pandas as pd +import pandas.testing as pdt +import pytest + +from activitysim.core import chunk, mem, simulate, workflow + +TESTDIR = os.path.dirname(__file__) +DATADIR = os.path.join(TESTDIR, "data") + + +@pytest.fixture +def state() -> workflow.State: + st = workflow.State() + st.initialize_filesystem( + working_dir=TESTDIR, data_dir=(DATADIR,) + ).default_settings() + st.settings.check_for_variability = False + return st + + +@pytest.fixture +def spec(state): + return state.filesystem.read_model_spec(file_name="sample_spec.csv") + + +@pytest.fixture +def data(): + return pd.read_csv(os.path.join(DATADIR, "data.csv")) + + +EXPECTED = pd.Series([1, 1, 1]) + + +def test_resolve_chunk_size_fixed_is_legacy(state): + # default (fixed) mode must return the static chunk_size verbatim -> no behavior change + state.settings.chunk_size = 123456 + assert chunk.resolve_chunk_size(state) == 123456 + + +def test_resolve_chunk_size_auto(state): + state.settings.chunk_size = 0 + state.settings.chunk_memory_mode = "auto" + state.settings.chunk_memory_safety_factor = 0.75 + limit = mem.get_memory_limit() + budget = chunk.resolve_chunk_size(state) + # Budget = safety_factor * AVAILABLE memory (limit - current usage), kept strictly positive by the + # floor and never above safety_factor * the real ceiling. (Available <= limit, so budget <= 0.75*limit.) + assert budget >= chunk.AUTO_BUDGET_FLOOR + assert 0 < budget <= max(chunk.AUTO_BUDGET_FLOOR, int(limit * 0.75)) + + +def test_resolve_chunk_size_auto_safety_factor_scales(state): + # a smaller safety_factor yields a smaller (or equal, if floored) budget + state.settings.chunk_size = 0 + state.settings.chunk_memory_mode = "auto" + state.settings.chunk_memory_safety_factor = 0.75 + hi = chunk.resolve_chunk_size(state) + state.settings.chunk_memory_safety_factor = 0.25 + lo = chunk.resolve_chunk_size(state) + assert lo <= hi + + +def test_auto_mode_simple_simulate_matches_fixed(state, data, spec): + # auto mode must produce the same choices as the legacy path + state.settings.chunk_size = 0 + state.settings.chunk_memory_mode = "auto" + state.settings.chunk_growth_cap = 2.0 + choices = simulate.simple_simulate(state, choosers=data, spec=spec, nest_spec=None) + pdt.assert_series_equal(choices.reset_index(drop=True), EXPECTED, check_dtype=False) + + +def test_auto_mode_splits_into_multiple_chunks(state, data, monkeypatch): + # Force a tiny auto budget so the choosers are split into MULTIPLE chunks (not a single-chunk run), + # exercising the auto chunk-sizing loop. Assert the chunks partition the choosers exactly. + state.settings.chunk_size = 0 + state.settings.chunk_memory_mode = "auto" + state.settings.chunk_training_mode = "training" + state.settings.default_initial_rows_per_chunk = 1 # tiny first (probe) chunk + monkeypatch.setattr(chunk, "AUTO_BUDGET_FLOOR", 1) + monkeypatch.setattr(mem, "get_memory_limit", lambda *a, **k: 1) + monkeypatch.setattr(mem, "get_nonreclaimable_used", lambda *a, **k: 0) + + chunks = [ + chooser_chunk.copy() + for _i, chooser_chunk, _label, _sizer in chunk.adaptive_chunked_choosers( + state, data, "test_auto_multichunk" + ) + ] + assert len(chunks) > 1 # the tiny budget forced more than one chunk + # chunks partition the original choosers exactly (rows + order preserved, none lost/duplicated) + pdt.assert_frame_equal(pd.concat(chunks), data) + + +def test_watchdog_warns_only_on_breach(state, caplog): + state.settings.chunk_memory_circuit_breaker = True + state.settings.chunk_memory_abort_ratio = 0.9 + limit = mem.get_memory_limit() + chunk._watchdog_breached = False + with caplog.at_level(logging.WARNING, logger="activitysim.core.chunk"): + chunk.memory_watchdog_check(state, int(0.5 * limit), "below") + assert not any("WATCHDOG" in r.message for r in caplog.records) + chunk.memory_watchdog_check(state, int(0.95 * limit), "over") + assert any("WATCHDOG" in r.message for r in caplog.records) + + +def test_watchdog_disabled_by_default(state, caplog): + # circuit breaker off (default) -> never warns even above the ratio + chunk._watchdog_breached = False + with caplog.at_level(logging.WARNING, logger="activitysim.core.chunk"): + chunk.memory_watchdog_check(state, int(0.99 * mem.get_memory_limit()), "over") + assert not any("WATCHDOG" in r.message for r in caplog.records) diff --git a/activitysim/core/test/test_mem.py b/activitysim/core/test/test_mem.py new file mode 100644 index 0000000000..d3c5a656c6 --- /dev/null +++ b/activitysim/core/test/test_mem.py @@ -0,0 +1,83 @@ +# ActivitySim +# See full license in LICENSE.txt. +"""Tests for the cgroup-aware memory-ceiling helpers used by adaptive chunking (chunk_memory_mode=auto).""" +from __future__ import annotations + +import os + +import psutil + +from activitysim.core import mem + +GIB = 1024**3 + + +def _write(root, rel, text): + path = os.path.join(root, rel) + os.makedirs(os.path.dirname(path), exist_ok=True) + with open(path, "w") as f: + f.write(text) + + +def test_finite_limit_parsing(): + assert mem._finite_limit("62000000000") == 62000000000 + assert mem._finite_limit("max") is None + assert mem._finite_limit(None) is None + assert mem._finite_limit("garbage") is None + assert mem._finite_limit("0") is None # non-positive is not a real limit + assert mem._finite_limit(str(mem._CGROUP_UNLIMITED)) is None # unlimited sentinel + + +def test_memory_limit_cgroup_v2(tmp_path): + root = str(tmp_path) + _write(root, "memory.max", "60000000000\n") + assert mem.get_memory_limit(cgroup_root=root) == 60000000000 + + +def test_memory_limit_cgroup_v2_max_falls_back(tmp_path): + # cgroup v2 present but unlimited ("max") -> fall through to host RAM (a positive int) + root = str(tmp_path) + _write(root, "memory.max", "max\n") + limit = mem.get_memory_limit(cgroup_root=root) + assert limit == int(psutil.virtual_memory().total) + assert limit > 0 + + +def test_memory_limit_cgroup_v1(tmp_path): + root = str(tmp_path) # no memory.max -> v1 path + _write(root, "memory/memory.limit_in_bytes", "48000000000\n") + assert mem.get_memory_limit(cgroup_root=root) == 48000000000 + + +def test_memory_limit_cgroup_v1_unlimited_falls_back(tmp_path): + root = str(tmp_path) + _write(root, "memory/memory.limit_in_bytes", str(mem._CGROUP_UNLIMITED)) + assert mem.get_memory_limit(cgroup_root=root) == int(psutil.virtual_memory().total) + + +def test_memory_limit_fallback_to_host(tmp_path): + # empty cgroup root -> psutil host total + assert mem.get_memory_limit(cgroup_root=str(tmp_path)) == int( + psutil.virtual_memory().total + ) + + +def test_available_memory_cgroup(tmp_path): + root = str(tmp_path) + _write(root, "memory.max", str(50 * GIB)) + _write(root, "memory.current", str(20 * GIB)) + assert mem.get_available_memory(cgroup_root=root) == 30 * GIB + + +def test_available_memory_fallback(tmp_path): + # no usage file -> psutil available (a non-negative int) + avail = mem.get_available_memory(cgroup_root=str(tmp_path)) + assert isinstance(avail, int) and avail >= 0 + + +def test_get_peak_rss(): + # exact lifetime peak RSS (getrusage ru_maxrss) — positive, and monotonic non-decreasing + p1 = mem.get_peak_rss() + assert isinstance(p1, int) and p1 > 0 + _ = [0] * 1_000_000 # allocate a little + assert mem.get_peak_rss() >= p1 diff --git a/docs/core.rst b/docs/core.rst index fc695ddf14..8fa41bdf52 100644 --- a/docs/core.rst +++ b/docs/core.rst @@ -649,6 +649,37 @@ Additional chunking settings: * keep_chunk_logs: True - whether to preserve or delete subprocess chunk logs when they are consolidated at end of multiprocess run * keep_mem_logs: True - whether to preserve or delete subprocess mem logs when they are consolidated at end of multiprocess run +Automatic memory-aware chunking (``chunk_memory_mode: auto``) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +By default (``chunk_memory_mode: fixed``) adaptive chunking sizes chunks against the static +``chunk_size`` byte budget, which must be hand-tuned per machine and targets host RAM. Setting +``chunk_memory_mode: auto`` instead derives the budget from the process's real memory ceiling at +runtime, so ``chunk_size`` need not be set and the run adapts to the actual machine or container +(this is especially useful inside memory-limited containers, where targeting host RAM can OOM-kill +the process): + +* The budget is ``(memory_limit - current usage) * chunk_memory_safety_factor``, where + ``memory_limit`` is read from the Linux cgroup (v2 ``memory.max``, then v1 + ``memory.limit_in_bytes``, then ``psutil`` host RAM) — the limit that actually OOM-kills the + process inside a container. +* In multiprocess mode the budget is divided by the number of workers. Set ``num_processes: 0`` to + additionally derive the worker count automatically from the available (non-reclaimable) memory and + the ``chunk_worker_target_budget`` per-worker target. +* Chunks are sized from the budget against each chunk's incremental memory growth, so a large + memory-mapped shared skim buffer (reclaimable page cache) does not distort the sizing. + +This mode reuses the existing adaptive-chunking machinery; with the default ``fixed`` mode behavior +is unchanged. Settings: + +* chunk_memory_mode: fixed - ``auto`` derives the chunk budget from the real memory ceiling; ``fixed`` (default) uses the static ``chunk_size`` +* chunk_memory_safety_factor: 0.75 - fraction of the available memory ceiling to use as the budget +* chunk_worker_target_budget: 0 - per-worker budget in bytes for the automatic worker count when ``num_processes: 0`` (0 = off) +* chunk_growth_cap: 0 - maximum multiplicative growth of rows-per-chunk between successive chunks (0 = off) +* chunk_row_size_margin: 1.0 - safety multiplier applied to the estimated per-row memory when sizing chunks +* chunk_memory_circuit_breaker: false - warn as memory approaches the ceiling +* chunk_memory_abort_ratio: 0.9 - fraction of the ceiling at which the adaptive back-off / warning triggers + API ^^^