Add chunk_memory_mode: auto — container-memory-aware adaptive chunking

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.

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

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: