diff --git a/activitysim/abm/models/location_choice.py b/activitysim/abm/models/location_choice.py index c955f4a60..2f298f17e 100644 --- a/activitysim/abm/models/location_choice.py +++ b/activitysim/abm/models/location_choice.py @@ -551,7 +551,7 @@ def run_location_sample( estimator, model_settings, chunk_size, - chunk_tag=f"{chunk_tag}.presample", + chunk_tag=f"{chunk_tag}.presample.{segment_name}", trace_label=trace_label, full_dest_size_terms=full_dest_size_terms, ) @@ -567,7 +567,7 @@ def run_location_sample( estimator, model_settings, chunk_size, - chunk_tag=f"{chunk_tag}.sample", + chunk_tag=f"{chunk_tag}.sample.{segment_name}", trace_label=trace_label, ) @@ -880,7 +880,7 @@ def run_location_choice( location_sample_df, model_settings, chunk_size, - chunk_tag=f"{chunk_tag}.logsums", + chunk_tag=f"{chunk_tag}.logsums.{segment_name}", trace_label=tracing.extend_trace_label( trace_label, "logsums.%s" % segment_name ), @@ -900,7 +900,7 @@ def run_location_choice( estimator, model_settings, chunk_size, - chunk_tag=f"{chunk_tag}.simulate", + chunk_tag=f"{chunk_tag}.simulate.{segment_name}", trace_label=tracing.extend_trace_label( trace_label, "simulate.%s" % segment_name ), diff --git a/activitysim/abm/models/trip_destination.py b/activitysim/abm/models/trip_destination.py index 7dc6154cc..a8bef5693 100644 --- a/activitysim/abm/models/trip_destination.py +++ b/activitysim/abm/models/trip_destination.py @@ -258,7 +258,7 @@ def destination_sample( chunk_size, trace_label, ): - chunk_tag = "trip_destination.sample" + chunk_tag = f"trip_destination.sample.{primary_purpose}" skims = skim_hotel.sample_skims(presample=False) alt_dest_col_name = model_settings.ALT_DEST_COL_NAME @@ -630,7 +630,7 @@ def destination_presample( trace_label, ): trace_label = tracing.extend_trace_label(trace_label, "presample") - chunk_tag = "trip_destination.presample" # distinguish from trip_destination.sample + chunk_tag = f"trip_destination.presample.{primary_purpose}" # distinguish from trip_destination.sample alt_dest_col_name = model_settings.ALT_DEST_COL_NAME @@ -876,7 +876,7 @@ def compute_logsums( logger.debug("Running %s with %d samples", trace_label, destination_sample.shape[0]) # chunk usage is uniform so better to combine - chunk_tag = "trip_destination.compute_logsums" + chunk_tag = f"trip_destination.compute_logsums.{primary_purpose}" # FIXME should pass this in? network_los = state.get_injectable("network_los") @@ -995,7 +995,7 @@ def trip_destination_simulate( destination alt chosen """ trace_label = tracing.extend_trace_label(trace_label, "trip_dest_simulate") - chunk_tag = "trip_destination.simulate" + chunk_tag = f"trip_destination.simulate.{primary_purpose}" spec = simulate.spec_for_segment( state, diff --git a/activitysim/abm/models/util/tour_destination.py b/activitysim/abm/models/util/tour_destination.py index 11bdb8ecd..ace7844cc 100644 --- a/activitysim/abm/models/util/tour_destination.py +++ b/activitysim/abm/models/util/tour_destination.py @@ -192,7 +192,7 @@ def destination_sample( chunk_size, trace_label, ): - chunk_tag = "tour_destination.sample" + chunk_tag = f"tour_destination.sample.{spec_segment_name}" # create wrapper with keys for this lookup # the skims will be available under the name "skims" for any @ expressions @@ -602,7 +602,7 @@ def destination_presample( trace_label, ): trace_label = tracing.extend_trace_label(trace_label, "presample") - chunk_tag = "tour_destination.presample" + chunk_tag = f"tour_destination.presample.{spec_segment_name}" logger.debug(f"{trace_label} location_presample") @@ -791,7 +791,7 @@ def run_destination_logsums( # if special person id is passed chooser_id_column = model_settings.CHOOSER_ID_COLUMN - chunk_tag = "tour_destination.logsums" + chunk_tag = f"tour_destination.logsums.{tour_purpose}" # merge persons into tours choosers = pd.merge( @@ -843,7 +843,7 @@ def run_destination_simulate( run destination_simulate on tour_destination_sample annotated with mode_choice logsum to select a destination from sample alternatives """ - chunk_tag = "tour_destination.simulate" + chunk_tag = f"tour_destination.simulate.{spec_segment_name}" model_spec = simulate.spec_for_segment( state, diff --git a/activitysim/abm/test/test_misc/test_trip_destination_sampling.py b/activitysim/abm/test/test_misc/test_trip_destination_sampling.py index 62bc4492b..f85768c4e 100644 --- a/activitysim/abm/test/test_misc/test_trip_destination_sampling.py +++ b/activitysim/abm/test/test_misc/test_trip_destination_sampling.py @@ -115,7 +115,7 @@ def fake_destination_sample( pd.Index([101, 102, 103], name="dest_taz"), ) assert captured["alt_dest_col_name"] == "dest_taz" - assert captured["chunk_tag"] == "trip_destination.sample" + assert captured["chunk_tag"] == "trip_destination.sample.eatout" assert captured["zone_layer"] is None assert captured["presample"] is False @@ -219,7 +219,7 @@ def fake_choose_maz_for_taz( pd.Index([1, 2, 3], name="zone_id"), ) assert captured["alt_dest_col_name"] == "dest_taz" - assert captured["chunk_tag"] == "trip_destination.presample" + assert captured["chunk_tag"] == "trip_destination.presample.eatout" assert captured["zone_layer"] == "taz" assert captured["presample"] is True assert captured["full_taz_index"] is None diff --git a/activitysim/core/chunk.py b/activitysim/core/chunk.py index f0074683d..8fb651cb7 100644 --- a/activitysim/core/chunk.py +++ b/activitysim/core/chunk.py @@ -22,6 +22,7 @@ logger = logging.getLogger(__name__) + # # CHUNK_METHODS and METRICS # @@ -88,6 +89,85 @@ MODE_PRODUCTION = "production" MODE_CHUNKLESS = "disabled" MODE_EXPLICIT = "explicit" + +# Auto-mode per-worker chunk budget (bytes) below which we WARN that memory is very tight. This is only +# a warning threshold — NOT a hard floor (flooring the per-worker budget would sum across workers and +# over-commit the shared ceiling). The actual budget is floored only to a positive value (>= 1) so that +# chunking stays active (chunk_size == 0 would disable chunking and OOM a large sample). +MIN_AUTO_BUDGET_WARN = 256_000_000 + +# Default for the chunk_peak_backoff_ratio setting: fraction of the (auto) per-worker budget a +# chunk's INCREMENTAL peak may reach before the adaptive sizer halves the next chunk (proactive +# back-off). Not a hard abort — just steers the ramp. +PEAK_BACKOFF_RATIO = 0.9 + +# Under chunk_size_mode=auto the first ("probe") chunk of a model with no cached row_size is the +# one allocation the budget cannot control — the per-row memory is unknown until it is measured, and +# in multiprocess ALL workers hit it simultaneously. Cap that probe small so measuring is cheap and +# a mis-guess cannot burst the ceiling (observed: a 15_000-row default probe drove per-worker peaks +# to ~2x the budget and OOM-killed a full-population run before any adaptation could kick in). +MAX_AUTO_PROBE_ROWS = 2000 + +# Under auto, cap post-probe growth by default: extrapolating a big chunk linearly from a tiny probe +# is the other unguarded leap. A user-set chunk_growth_cap still wins; 0 (legacy "no cap") only +# applies to fixed mode. +AUTO_DEFAULT_GROWTH_CAP = 2.0 + + +def _effective_growth_cap(settings) -> float: + """chunk_growth_cap, defaulting to AUTO_DEFAULT_GROWTH_CAP under chunk_size_mode=auto. + + A user-set cap (>= 1) wins. In fixed mode 0 (uncapped) remains the legacy default; auto mode + is expected to be safe out of the box, so growth is always capped there — to effectively + disable the cap under auto, set a large value (e.g. 100). + """ + cap = getattr(settings, "chunk_growth_cap", 0) or 0 + if not cap and getattr(settings, "chunk_size_mode", "fixed") == "auto": + return AUTO_DEFAULT_GROWTH_CAP + return cap + + +# one settings line per process (workers each log their own on first chunked model) +_CHUNK_SETTINGS_LOGGED = False + + +def log_chunking_settings(state: workflow.State) -> None: + """Log every effective chunking parameter once per process, for run audits. + + The values that govern chunking are spread across several settings (and two of them + have auto-mode defaults applied at runtime), so a run log otherwise never shows the + complete picture in one place. Emitted on the first chunked model of each process. + """ + global _CHUNK_SETTINGS_LOGGED + if _CHUNK_SETTINGS_LOGGED: + return + _CHUNK_SETTINGS_LOGGED = True + s = state.settings + num_processes = 1 + if getattr(s, "multiprocess", False): + try: + injected = state.get_injectable("num_processes", None) + except Exception: + injected = None + num_processes = injected or getattr(s, "num_processes", 1) or 1 + logger.info( + "chunking settings: " + f"chunk_size_mode={getattr(s, 'chunk_size_mode', 'fixed')} " + f"chunk_size={getattr(s, 'chunk_size', 0)} " + f"chunk_size_safety_factor={getattr(s, 'chunk_size_safety_factor', None)} " + f"chunk_growth_cap={getattr(s, 'chunk_growth_cap', 0)} " + f"(effective={_effective_growth_cap(s)}) " + f"chunk_row_size_margin={getattr(s, 'chunk_row_size_margin', 1.0)} " + f"chunk_training_mode={s.chunk_training_mode} " + f"chunk_method={getattr(s, 'chunk_method', None)} " + f"default_initial_rows_per_chunk={getattr(s, 'default_initial_rows_per_chunk', None)} " + f"(auto probe cap={MAX_AUTO_PROBE_ROWS} rows) " + f"num_processes={num_processes} " + f"chunk_peak_backoff_ratio={getattr(s, 'chunk_peak_backoff_ratio', PEAK_BACKOFF_RATIO)} " + f"min_auto_budget_warn={MIN_AUTO_BUDGET_WARN}" + ) + + TRAINING_MODES = [ MODE_RETRAIN, MODE_ADAPTIVE, @@ -185,6 +265,77 @@ 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_size_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_size_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_size_mode", "fixed") != "auto": + return state.settings.chunk_size + + limit = mem.get_memory_limit() + if not limit: + logger.warning( + "chunk_size_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_size_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) keeps the budget honest: it shrinks as more skim + # pages fault in, tracking the working set. `available == 0` legitimately means "no headroom" — only a + # None (couldn't read usage) falls back to the raw limit. + available = mem.get_available_memory() + if available is not None: + basis = max(0, min(limit, available)) + else: + basis = limit + # Multiprocess: the N workers share the single ceiling, so the budget is the shared free memory + # divided by the REAL per-step worker count — the 'num_processes' injectable 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 to the FULL budget → collective OOM). Fall back to the setting outside a worker. + num_processes = 1 + if getattr(state.settings, "multiprocess", False): + try: + injected = state.get_injectable("num_processes", None) + except Exception: + injected = None + num_processes = injected or getattr(state.settings, "num_processes", 1) or 1 + aggregate = int(basis * safety) + budget = aggregate // num_processes if num_processes > 1 else aggregate + # Keep chunking active: a budget of 0 makes chunk_size == 0, which disables chunking (process ALL + # choosers at once) and OOMs a large sample. Floor to a small positive value so rows-per-chunk clips + # to >= 1. Do NOT floor to a large PER-WORKER minimum — that sums across workers and over-commits + # (e.g. 8 workers x 1 GB on a 2 GB node, the very failure this feature prevents). A tiny budget is a + # genuine "memory is very tight" signal, so warn rather than silently inflate it. + if budget < MIN_AUTO_BUDGET_WARN and basis > 0: + logger.warning( + f"chunk_size_mode=auto: per-worker chunk budget {GB(budget)} is very small " + f"(limit={GB(limit)}, {num_processes} worker(s)); memory is tight — chunks will be minimal " + "and the run may be slow. Consider fewer workers or a higher memory limit." + ) + budget = max(budget, 1) + baseline, _ = mem.get_rss(force_garbage_collect=True) + # include the kernel's exact lifetime peak (getrusage ru_maxrss — never misses a transient + # spike, unlike a sampled high-water mark) so a completed run shows how close this process + # came to the ceiling: the number to look at when tuning chunk_size_safety_factor. + logger.info( + f"chunk_size_mode=auto: limit={GB(limit)} " + f"available={GB(available) if available is not None else 'n/a'} x safety_factor={safety}" + f"{f' / {num_processes} workers' if num_processes > 1 else ''} " + f"-> base_chunk_size={GB(budget)} " + f"(current rss={GB(baseline)}, peak rss={GB(mem.get_peak_rss())})" + ) + return budget + + def overhead_for_chunk_method(state: workflow.State, overhead, method=None): """ @@ -508,8 +659,17 @@ def audit( if not self.base_chunk_size: return + auto = getattr(state.settings, "chunk_size_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 +677,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}") @@ -706,8 +870,9 @@ def __init__( threading.Thread.__init__(self) def run(self): + tick = mem.MEM_SNOOP_TICK_LEN 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 +1002,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_size_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_size_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 +1061,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( @@ -882,10 +1078,19 @@ def initial_rows_per_chunk(self): else: # if no initial_row_size from cache, fall back to default_initial_rows_per_chunk self.initial_row_size = 0 - rows_per_chunk = min( - self.num_choosers, - self.state.settings.default_initial_rows_per_chunk, - ) + probe_rows = self.state.settings.default_initial_rows_per_chunk + if ( + getattr(self.state.settings, "chunk_size_mode", "fixed") == "auto" + and probe_rows > MAX_AUTO_PROBE_ROWS + ): + # auto: the probe's per-row memory is unknown, so the budget cannot bound it — + # keep the probe small; adaptation takes over from chunk 2 + logger.info( + f"{self.trace_label}.initial_rows_per_chunk - capping probe chunk " + f"{probe_rows} -> {MAX_AUTO_PROBE_ROWS} rows (chunk_size_mode=auto)" + ) + probe_rows = MAX_AUTO_PROBE_ROWS + rows_per_chunk = min(self.num_choosers, probe_rows) estimated_number_of_chunks = None if self.chunk_training_mode == MODE_PRODUCTION: @@ -965,11 +1170,57 @@ 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 = _effective_growth_cap(settings) + 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_size_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] + ) + backoff_ratio = ( + getattr(settings, "chunk_peak_backoff_ratio", PEAK_BACKOFF_RATIO) + or PEAK_BACKOFF_RATIO + ) + if peak_incremental > backoff_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"(> {backoff_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 = ( @@ -1215,6 +1466,7 @@ def adaptive_chunked_choosers( chunk_size: int | None = None, explicit_chunk_size: float = 0, ): + log_chunking_settings(state) # generator to iterate over choosers if state.settings.chunk_training_mode == MODE_CHUNKLESS or ( @@ -1244,8 +1496,14 @@ def adaptive_chunked_choosers( chunk_size = math.ceil(num_choosers * explicit_chunk_size) else: chunk_size = math.ceil(explicit_chunk_size / num_processes) - elif chunk_size is None: - chunk_size = state.settings.chunk_size + elif chunk_size is None or ( + getattr(state.settings, "chunk_size_mode", "fixed") == "auto" and chunk_size + ): + # auto replaces a POSITIVE statically passed chunk_size with the runtime budget + # (callers historically pass settings.chunk_size down explicitly; honoring it here + # would silently run parts of the pipeline on the static value). chunk_size == 0 is + # preserved in every mode: callers use it to run a component deliberately chunkless. + chunk_size = resolve_chunk_size(state) assert num_choosers > 0 assert chunk_size >= 0 @@ -1303,6 +1561,7 @@ def adaptive_chunked_choosers_and_alts( chunk_size: int | None = None, explicit_chunk_size: int = 0, ): + log_chunking_settings(state) """ generator to iterate over choosers and alternatives in chunk_size chunks @@ -1387,8 +1646,14 @@ def adaptive_chunked_choosers_and_alts( chunk_size = math.ceil(num_choosers * explicit_chunk_size) else: chunk_size = int(explicit_chunk_size / num_processes) - elif chunk_size is None: - chunk_size = state.settings.chunk_size + elif chunk_size is None or ( + getattr(state.settings, "chunk_size_mode", "fixed") == "auto" and chunk_size + ): + # auto replaces a POSITIVE statically passed chunk_size with the runtime budget + # (callers historically pass settings.chunk_size down explicitly; honoring it here + # would silently run parts of the pipeline on the static value). chunk_size == 0 is + # preserved in every mode: callers use it to run a component deliberately chunkless. + chunk_size = resolve_chunk_size(state) chunk_sizer = ChunkSizer( state, @@ -1468,6 +1733,7 @@ def adaptive_chunked_choosers_by_chunk_id( chunk_tag=None, explicit_chunk_size: int = 0, ): + log_chunking_settings(state) # generator to iterate over choosers in chunk_size chunks # like chunked_choosers but based on chunk_id field rather than dataframe length # (the presumption is that choosers has multiple rows with the same chunk_id that @@ -1496,7 +1762,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 6b1c8a6c5..8be7b1319 100644 --- a/activitysim/core/configuration/top.py +++ b/activitysim/core/configuration/top.py @@ -399,6 +399,82 @@ class Settings(PydanticBase, extra="allow", validate_assignment=True): minimum fraction of total chunk_size to reserve for adaptive chunking """ + chunk_size_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_size_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_size_safety_factor: float = 0.5 + """ + Fraction of ``(memory_limit - baseline)`` to use as the chunking budget under + ``chunk_size_mode: auto``. + + The default 0.5 tolerates a ~2x per-row-size mis-estimate — the realistic worst case when + one model tag spans segments with very different alternative sets (observed 58 KB -> 127 KB + per row between school_location segments), and in multiprocess ALL workers hit the + mis-estimated segment simultaneously. Sized so that even then the aggregate stays within the + ceiling. If a completed run's ``peak rss`` log lines show ample headroom, raise it (e.g. 0.75) + for bigger chunks and more throughput. + """ + + 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_peak_backoff_ratio: float = 0.9 + """ + Fraction of the per-worker budget a chunk's incremental memory peak may reach before the + sizer halves the next chunk (``chunk_size_mode: auto``, training/adaptive only). Lower + values back off earlier (more conservative); 1.0 backs off only when the budget is fully + consumed. Not a hard abort — it steers the ramp-up. + """ + + @model_validator(mode="after") + def _check_chunk_memory_settings(self): + if not (0 < self.chunk_size_safety_factor <= 1.0): + raise ValueError( + "chunk_size_safety_factor must be in (0, 1], " + f"got {self.chunk_size_safety_factor}" + ) + if self.chunk_growth_cap < 0 or 0 < self.chunk_growth_cap < 1.0: + raise ValueError( + "chunk_growth_cap must be 0 (off) or >= 1.0, " + f"got {self.chunk_growth_cap}" + ) + if self.chunk_row_size_margin < 1.0: + raise ValueError( + f"chunk_row_size_margin must be >= 1.0, got {self.chunk_row_size_margin}" + ) + if not (0 < self.chunk_peak_backoff_ratio <= 1.0): + raise ValueError( + "chunk_peak_backoff_ratio must be in (0, 1], " + f"got {self.chunk_peak_backoff_ratio}" + ) + return self + 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 fbfa11fbe..f1a8d223a 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,111 @@ 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_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/test/test_chunk_robust.py b/activitysim/core/test/test_chunk_robust.py new file mode 100644 index 000000000..366904619 --- /dev/null +++ b/activitysim/core/test/test_chunk_robust.py @@ -0,0 +1,245 @@ +# ActivitySim +# See full license in LICENSE.txt. +"""Tests for the cgroup-aware ``chunk_size_mode: auto`` budget and the adaptive sizing it feeds. +See also test_mem.py for the cgroup / worker-count helpers.""" +from __future__ import annotations + +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") + +GIB = 1024**3 + + +@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_size_mode = "auto" + state.settings.chunk_size_safety_factor = 0.75 + limit = mem.get_memory_limit() + budget = chunk.resolve_chunk_size(state) + # Budget = safety_factor * AVAILABLE memory (limit - current usage): strictly positive, and never + # above safety_factor * the real ceiling (available <= limit). + assert budget >= 1 + assert budget <= max(1, int(limit * 0.75)) + + +def test_resolve_chunk_size_auto_safety_factor_scales(state): + # a smaller safety_factor yields a smaller (or equal) budget + state.settings.chunk_size = 0 + state.settings.chunk_size_mode = "auto" + state.settings.chunk_size_safety_factor = 0.75 + hi = chunk.resolve_chunk_size(state) + state.settings.chunk_size_safety_factor = 0.25 + lo = chunk.resolve_chunk_size(state) + assert lo <= hi + + +def test_resolve_chunk_size_auto_zero_headroom_is_not_full_limit(state, monkeypatch): + # available == 0 legitimately means "no headroom" and must NOT be treated as "unknown" and replaced + # with the full limit (that would hand out a huge budget exactly when memory is exhausted). + state.settings.chunk_size = 0 + state.settings.chunk_size_mode = "auto" + monkeypatch.setattr(mem, "get_memory_limit", lambda *a, **k: 100 * GIB) + monkeypatch.setattr(mem, "get_available_memory", lambda *a, **k: 0) + budget = chunk.resolve_chunk_size(state) + assert budget < GIB # tiny (floored to keep chunking on), NOT ~75 GB + + +def test_resolve_chunk_size_auto_divides_by_workers(state, monkeypatch): + # multiprocess: the shared budget is divided by the per-step worker count (num_processes injectable) + state.settings.chunk_size = 0 + state.settings.chunk_size_mode = "auto" + state.settings.chunk_size_safety_factor = 1.0 + state.settings.multiprocess = True + monkeypatch.setattr(mem, "get_memory_limit", lambda *a, **k: 40 * GIB) + monkeypatch.setattr(mem, "get_available_memory", lambda *a, **k: 40 * GIB) + state.add_injectable("num_processes", 4) + budget = chunk.resolve_chunk_size(state) + assert budget == 10 * GIB # 40 GB / 4 workers + + +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_size_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_size_mode = "auto" + state.settings.chunk_training_mode = "training" + state.settings.default_initial_rows_per_chunk = 1 # tiny first (probe) chunk + monkeypatch.setattr(mem, "get_memory_limit", lambda *a, **k: 1000) + monkeypatch.setattr(mem, "get_available_memory", lambda *a, **k: 1) + + 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_chunk_memory_settings_validation(): + # the auto-mode knobs reject nonsensical values at configuration time + from pydantic import ValidationError + + from activitysim.core.configuration.top import Settings + + Settings(chunk_size_safety_factor=0.5) # in (0, 1] — ok + Settings(chunk_growth_cap=1.5) # off (0) or >= 1 — ok + Settings(chunk_row_size_margin=1.3) # >= 1 — ok + with pytest.raises(ValidationError): + Settings(chunk_size_safety_factor=0.0) + with pytest.raises(ValidationError): + Settings(chunk_size_safety_factor=1.5) + with pytest.raises(ValidationError): + Settings(chunk_growth_cap=0.5) # would shrink every chunk toward collapse + with pytest.raises(ValidationError): + Settings(chunk_row_size_margin=0.9) + Settings(chunk_peak_backoff_ratio=0.8) # in (0, 1] — ok + with pytest.raises(ValidationError): + Settings(chunk_peak_backoff_ratio=0.0) + with pytest.raises(ValidationError): + Settings(chunk_peak_backoff_ratio=1.5) + + +def test_auto_growth_cap_default(): + # auto mode caps growth by default; fixed keeps legacy uncapped; explicit setting wins + from activitysim.core.configuration.top import Settings + + assert chunk._effective_growth_cap(Settings(chunk_size_mode="fixed")) == 0 + assert ( + chunk._effective_growth_cap(Settings(chunk_size_mode="auto")) + == chunk.AUTO_DEFAULT_GROWTH_CAP + ) + assert ( + chunk._effective_growth_cap( + Settings(chunk_size_mode="auto", chunk_growth_cap=3.0) + ) + == 3.0 + ) + + +def test_auto_probe_chunk_is_capped(state, monkeypatch): + # with no cached row_size, the first (probe) chunk under auto is capped at + # MAX_AUTO_PROBE_ROWS even when default_initial_rows_per_chunk is huge + n = chunk.MAX_AUTO_PROBE_ROWS * 3 + data = pd.DataFrame({"x": range(n)}) + state.settings.chunk_size_mode = "auto" + state.settings.chunk_training_mode = "training" + state.settings.default_initial_rows_per_chunk = ( + 50_000 # deliberately oversized probe + ) + monkeypatch.setattr(mem, "get_memory_limit", lambda *a, **k: 100 * GIB) + monkeypatch.setattr(mem, "get_available_memory", lambda *a, **k: 100 * GIB) + + sizes = [ + len(chooser_chunk) + for _i, chooser_chunk, _label, _sizer in chunk.adaptive_chunked_choosers( + state, data, "test_auto_probe_cap" + ) + ] + assert sizes[0] <= chunk.MAX_AUTO_PROBE_ROWS + + +def test_chunking_settings_logged_once(state, caplog): + # the audit line contains every effective chunking parameter, once per process + import logging + + chunk._CHUNK_SETTINGS_LOGGED = False + state.settings.chunk_size_mode = "auto" + with caplog.at_level(logging.INFO, logger="activitysim.core.chunk"): + chunk.log_chunking_settings(state) + chunk.log_chunking_settings(state) # second call must be a no-op + msgs = [r.message for r in caplog.records if "chunking settings:" in r.message] + assert len(msgs) == 1 + for key in ( + "chunk_size_mode=auto", + "chunk_size=", + "chunk_size_safety_factor=", + "chunk_growth_cap=", + "(effective=", + "chunk_row_size_margin=", + "chunk_training_mode=", + "chunk_method=", + "default_initial_rows_per_chunk=", + "auto probe cap=", + "num_processes=", + "chunk_peak_backoff_ratio=", + ): + assert key in msgs[0], key + chunk._CHUNK_SETTINGS_LOGGED = False # don't leak state to other tests + + +def test_auto_overrides_passed_static_chunk_size(state, monkeypatch): + # callers historically pass settings.chunk_size down explicitly; under auto the + # runtime-resolved budget must win or parts of the pipeline silently run static + data = pd.DataFrame({"x": range(100)}) + state.settings.chunk_size_mode = "auto" + state.settings.chunk_training_mode = "training" + monkeypatch.setattr(mem, "get_memory_limit", lambda *a, **k: 10 * GIB) + monkeypatch.setattr(mem, "get_available_memory", lambda *a, **k: 10 * GIB) + budgets = [ + getattr(sizer, "base_chunk_size", None) or sizer.chunk_size + for _i, _c, _label, sizer in chunk.adaptive_chunked_choosers( + state, data, "test_auto_override", chunk_size=999 * GIB + ) + ] + assert budgets and all( + b <= 5 * GIB for b in budgets + ) # resolved from the 10 GiB limit, not the passed 999 GiB + + # chunk_size == 0 is the deliberate "run this component chunkless" signal — auto must + # NOT override it (tour scheduling logsums relies on this) + chunks = [ + c + for _i, c, _label, _sizer in chunk.adaptive_chunked_choosers( + state, data, "test_auto_keeps_chunkless", chunk_size=0 + ) + ] + assert len(chunks) == 1 and len(chunks[0]) == len(data) diff --git a/activitysim/core/test/test_mem.py b/activitysim/core/test/test_mem.py new file mode 100644 index 000000000..6074160b1 --- /dev/null +++ b/activitysim/core/test/test_mem.py @@ -0,0 +1,104 @@ +# ActivitySim +# See full license in LICENSE.txt. +"""Tests for the cgroup-aware memory-ceiling helpers used by adaptive chunking (chunk_size_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 + + +def test_get_peak_rss_without_resource_module(monkeypatch): + # Windows has no `resource` module (getrusage). The fallback must still return a + # positive, monotonic non-decreasing peak rather than raising on import or call. + monkeypatch.setattr(mem, "resource", None) + monkeypatch.setattr(mem, "_PEAK_RSS_FALLBACK", 0) + + p1 = mem.get_peak_rss() + assert isinstance(p1, int) and p1 > 0 + _ = [0] * 1_000_000 # allocate a little + p2 = mem.get_peak_rss() + assert p2 >= p1 # monotonic, like a real lifetime peak + + # a psutil failure degrades to the last known value instead of raising + class _Boom: + def __init__(self, *a, **k): + raise RuntimeError("no psutil here") + + monkeypatch.setattr(mem.psutil, "Process", _Boom) + assert mem.get_peak_rss() == p2 diff --git a/docs/core.rst b/docs/core.rst index fc695ddf1..39481d963 100644 --- a/docs/core.rst +++ b/docs/core.rst @@ -649,6 +649,42 @@ 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_size_mode: auto``) +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +By default (``chunk_size_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_size_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_size_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 (``num_processes``), so the N + workers that share the memory ceiling do not collectively exceed it. +* Each budget decision is logged together with the process's exact lifetime peak RSS + (``getrusage`` ``ru_maxrss`` via ``get_peak_rss``), so a completed run shows how close it came + to the ceiling — the number to look at when tuning ``chunk_size_safety_factor``. The + ``chunk_row_size_margin`` safety multiplier inflates the estimated per-row memory when sizing + chunks. +* Training-mode probe safety is built in: when a model has no cached per-row size, its first + ("probe") chunk is capped at 2000 rows (the smaller of that and ``default_initial_rows_per_chunk`` is used) — that chunk's memory cannot be bounded by the budget + because the per-row cost is unknown until measured, and in multiprocess all workers hit it at + once. Post-probe growth is capped at 2x per step by default (``chunk_growth_cap``, + user-overridable). A runtime budget alone cannot control these two bursts. + +This mode reuses the existing adaptive-chunking machinery; with the default ``fixed`` mode behavior +is unchanged. Settings: + +* chunk_size_mode: fixed - ``auto`` derives the chunk budget from the real memory ceiling; ``fixed`` (default) uses the static ``chunk_size`` +* chunk_size_safety_factor: 0.5 - fraction of the available memory ceiling to use as the budget (tolerates ~2x row-size mis-estimates across segments even when all workers hit them at once; raise if the ``peak rss`` log shows ample headroom) +* chunk_growth_cap: 0 - maximum multiplicative growth of rows-per-chunk between successive chunks (0 = off in ``fixed`` mode; under ``auto`` an unset value defaults to 2.0) +* chunk_row_size_margin: 1.0 - safety multiplier applied to the estimated per-row memory when sizing chunks +* chunk_peak_backoff_ratio: 0.9 - fraction of the per-worker budget a chunk's incremental peak may reach before the next chunk is halved + API ^^^ diff --git a/docs/dev-guide/changes.md b/docs/dev-guide/changes.md index cf4155ca3..2a05e90ae 100644 --- a/docs/dev-guide/changes.md +++ b/docs/dev-guide/changes.md @@ -14,6 +14,35 @@ branch (i.e., the main branch on GitHub), but not yet released in a stable versi of ActivitySim. See below under the various version headings for changes in released versions. +### Automatic Chunk Sizing from the Real Memory Limit (`chunk_size_mode: auto`) + +A new optional setting `chunk_size_mode` controls where adaptive chunking's memory +budget comes from. The default, `fixed`, preserves the existing behavior: the static +user-supplied `chunk_size` is used verbatim. Setting `chunk_size_mode: auto` ignores +`chunk_size` and derives the budget at runtime from the process's actual memory +ceiling — the Linux cgroup limit when running in a container (the limit that would +otherwise OOM-kill the run), or host RAM — minus current usage, scaled by the new +`chunk_size_safety_factor` setting (default 0.5), and divided across the multiprocess +worker count. The budget is recomputed at the start of every model component, so it +tracks memory actually in use. This removes the need to hand-tune `chunk_size` per +machine, and makes the same configuration portable across machines and containers of +different sizes. + +Auto mode also adds runtime safeguards that a static budget cannot provide: the first +("probe") chunk of any model with no cached row size is capped at 2000 rows, and +rows-per-chunk growth between successive chunks is capped (new `chunk_growth_cap` +setting; defaults to 2x under auto, off under fixed). A new `chunk_row_size_margin` +setting optionally inflates the estimated per-row memory when sizing chunks. + +Relatedly, the chunk cache tags for the `location_choice`, `tour_destination`, and +`trip_destination` components are now segmented by chooser segment (e.g. +`workplace_location.sample.work_high`), matching what `vectorize_tour_scheduling` +already does. Per-row memory can differ by more than 2x between segments of the same +component, so sizing one segment's chunks from another segment's cached row size could +badly overshoot. Existing `chunk_cache.csv` files will not match the new tags for +these components; a training-mode run rebuilds the cache (in the meantime the capped +probe keeps the first chunks safe). + ### Deprecated `SIMULATE_CHOOSER_COLUMNS` and `LOGSUM_CHOOSER_COLUMNS` The `SIMULATE_CHOOSER_COLUMNS` and `LOGSUM_CHOOSER_COLUMNS` settings were added