Skip to content

Add chunk_size_mode: auto — derive the chunk budget from the real memory limit - #1106

Open
vincentgong7 wants to merge 1 commit into
ActivitySim:mainfrom
vincentgong7:feature/chunk-size-mode-auto
Open

Add chunk_size_mode: auto — derive the chunk budget from the real memory limit#1106
vincentgong7 wants to merge 1 commit into
ActivitySim:mainfrom
vincentgong7:feature/chunk-size-mode-auto

Conversation

@vincentgong7

Copy link
Copy Markdown
Contributor

Summary

ActivitySim sizes each chunk from two inputs: a memory budget, and an estimate of what one
chooser row costs.

rows_per_chunk = budget / per_row_cost

The second input is already adaptive. In training mode ActivitySim measures the real per-row
cost of every component and caches it; adaptive and production reuse and refine those
cached values. The first input is not adaptive: in training, adaptive and production mode alike
the budget is the static chunk_size setting, a number the user must hand-tune to the machine.
Inside a container the process is killed at the cgroup limit, which is usually well below host
RAM, so a value tuned for the host over-commits and the run dies — often deep into a long
multiprocess run.

This PR makes the first input adaptive as well, and closes two gaps in the second. With
chunk_size_mode: auto the budget is derived at runtime from the process's actual memory
ceiling and divided across the worker count. The chunk cache is keyed per chooser segment rather
than per component, so segments with different row costs are sized from their own measurements.
And the chunks for which no measurement exists yet — the first chunk of a component, and every
chunk in production mode — are bounded explicitly. It is an enhancement of the existing
adaptive-chunking machinery, not a replacement. The default chunk_size_mode: fixed returns the
static chunk_size unchanged.

Motivation

How chunk sizing works today. For each component, ActivitySim splits the choosers into
batches ("chunks") sized to fit in memory. The number of rows per chunk is the memory budget
divided by an estimated per-row memory cost.

rows_per_chunk = budget / per_row_cost

chunk_training_mode selects how the per-row cost
is obtained:

  • training measures the real cost of each component while it runs and writes it to
    chunk_cache.csv.
  • adaptive starts from the cached cost and keeps measuring, refining it as the run proceeds.
  • production trusts the cached cost and does not measure, which is what makes it fast.

This machinery estimates the divisor - per-row cost - well. The dividend — the budget — is the static
chunk_size setting in all three modes.

Three problems with a static budget.

  1. It does not see the container memory limit. In Kubernetes the kernel kills the process at
    the cgroup limit, not at host RAM. A value chosen for the host over-commits in a container.
  2. It does not see memory already in use. Skims, framework state and charged page cache are
    not subtracted, so real headroom is smaller than the number implies.
  3. Each worker takes the full value. With num_processes: 6 and
    chunk_size: 14_000_000_000, the aggregate promise is 84 GB. The user must do that division
    by hand.

Two problems on the adaptive side. A correct budget is not sufficient on its own, because
the per-row estimate it is divided by is missing or imprecise in specific places.

  1. Some chunks have no estimate at all. The first chunk of a component has not been measured
    yet, so its size comes from default_initial_rows_per_chunk regardless of the budget. That
    setting defaults to 100 rows, which is safe, but it is a throughput knob: on a large machine
    a small first chunk wastes time, so it is often raised. Our own production configuration set
    it to 15,000. In multiprocessing every worker runs that chunk at the same moment, so a value
    raised for throughput is multiplied by the worker count at the least informed point of the
    run.
  2. The cached estimate is an average, recorded per component rather than per segment. The
    location and destination components run their chooser segments one after another under a
    single chunk tag, so the cost measured for one segment sizes the first chunk of the next. We
    measured 58 KB/row and 127 KB/row for two segments of school_location; sizing the second
    from the first overshoots by that factor on every worker simultaneously. Separately, a
    full-size chunk's transient peak exceeds the cached average it was sized from, and production
    mode applies that average at full size with no measurement and no ramp.

These are not hypothetical. In the benchmark below, a chunk_size tuned for the exact
container completed in training and adaptive mode but was OOM-killed in production mode: the
budget was right, and the run still died on problems 4 and 5.

What it does (when chunk_size_mode: auto)

  • Budget from the real ceiling.
    budget = ((memory_limit − memory_in_use) × chunk_size_safety_factor) / num_processes. memory_limit comes from the
    cgroup (v2 memory.max → v1 memory.limit_in_bytespsutil host RAM), so it is the limit
    that would actually kill the process. memory_in_use is the cgroup's current usage, so
    memory already held is subtracted before sizing. No machine-specific chunk_size is needed.
  • Recomputed per component. The budget is derived again at the start of every component — in all training modes, including production, which reads the per-row cost from cache but not the budget — so it follows memory that is actually held rather than a value fixed at startup.
  • Multiprocess-aware. The ceiling is shared by all workers, so the budget is divided by the
    worker count. No per-worker minimum is added, because a minimum multiplies across workers.
  • Probe chunks are capped. The first chunk of a component with no cached row size is limited
    to 2000 rows. This is a ceiling on the existing default_initial_rows_per_chunk setting, not a
    new setting: the smaller of the two is used, so a configuration that already asks for a small
    first chunk is unaffected, and one raised for throughput is brought back down. Users who want
    a smaller probe lower default_initial_rows_per_chunk as before. The cap itself is a constant
    rather than a setting, because a safety bound a user can raise is not a safety bound. The first
    chunk remains an unmeasured guess; capping it bounds what a wrong guess can cost.
  • Growth after the probe is capped. Rows per chunk may grow by at most chunk_growth_cap
    per step (default 2.0 in auto mode), bounding how far one small measurement is extrapolated.
  • Chunks back off near the ceiling. If a chunk's measured incremental peak exceeds
    chunk_peak_backoff_ratio of the budget (default 0.9), the next chunk is halved.
  • Per-segment chunk cache tags. The chunk cache is keyed per chooser segment, so each
    segment is sized from its own measurement.
  • Observability. Every budget decision is logged with the limit, available memory,
    per-worker budget, current RSS and exact lifetime peak RSS (getrusage, which does not miss
    short-lived spikes). Each process also logs one chunking settings: line listing every
    effective chunking parameter, so a run log is self-describing when reviewed later.
  • Guards. A suspiciously small budget produces a warning, not a silent floor. Zero available
    memory is treated as "no headroom, size down", never as "unknown, use the full limit".

Approach

The change has three layers, one for each group of problems above.

1. Derive the budget from the real limit, at runtime (problems 1-3). Reading the cgroup
limit and current usage replaces problems 1 and 2 with a measured quantity, and dividing by
num_processes replaces problem 3 with arithmetic the code performs. Because the derivation is
repeated per component, the budget tracks actual usage instead of a startup estimate.

2. Bound the chunks that have no estimate (problem 4). A budget only constrains a chunk
whose per-row cost is known. The probe cap, the growth cap and the peak backoff bound the
remaining cases: they keep the first measurement cheap, limit extrapolation from it, and shrink
the next chunk when a measured peak approaches the budget. These are what make production and
training mode safe, rather than the budget value itself.

3. Make the cached estimate granular enough to trust (problem 5). The guards above limit the
damage of a bad estimate; per-segment tags reduce how often the estimate is bad.
vectorize_tour_scheduling
already keys its chunk cache per segment
(segment_chunk_tag = extend_trace_label(tour_chunk_tag, tour_segment_name)). This PR applies
the same pattern to location_choice, tour_destination and trip_destination, so a segment
is sized from its own measured per-row cost instead of the previous segment's. No core change is
required, because the chunk historian keys by tag string. Iteration numbers (i1, i2, …) stay
out of the tag so shadow-pricing iterations continue to share history.

The three layers map onto the two inputs of rows_per_chunk = budget / per_row_cost: layer 1
makes the dividend correct for the machine, layer 3 makes the divisor more accurate, and layer 2
covers the chunks for which no divisor has been measured yet.

Changes

activitysim/core/mem.py — three helpers: get_memory_limit() (cgroup v2 → v1 → host RAM),
get_available_memory() (limit minus cgroup current usage), and get_peak_rss() (lifetime peak
from getrusage; the resource import is guarded and falls back to a monotonic psutil
high-water mark on Windows).

activitysim/core/chunk.pyresolve_chunk_size() implements the budget, and the sizer
implements the probe cap (a module constant), growth cap and peak backoff. Under auto the budget also replaces a
positive chunk_size passed down by a caller; chunk_size=0 is preserved, since callers use it
to run a component chunkless. Adds the budget log line and the chunking settings: audit line.

activitysim/abm/models/location_choice.py, util/tour_destination.py,
trip_destination.py
— chunk cache tags now include the chooser segment (12 one-line changes).

activitysim/core/configuration/top.py — the new settings, validated when the configuration
is loaded.

docs/core.rst, docs/dev-guide/changes.md — documentation and an Upcoming Changes entry,
including the cache migration note.

New settings

setting default meaning
chunk_size_mode fixed fixed uses the static chunk_size; auto derives the budget from the real memory limit
chunk_size_safety_factor 0.5 fraction of available memory used as the budget
chunk_growth_cap 0.0 maximum growth of rows-per-chunk between chunks (0 = off; auto uses 2.0 when unset)
chunk_peak_backoff_ratio 0.9 fraction of the budget a chunk's incremental peak may reach before the next chunk is halved
chunk_row_size_margin 1.0 multiplier applied to the estimated per-row memory when sizing chunks

The default for chunk_size_safety_factor is 0.5 because a full-size chunk's transient peak can
be roughly twice the cached average per-row cost, and all workers reach such a chunk at the same
time. The benchmark below tests this value directly.

Testing

Unit tests — 22 tests in test_mem.py and test_chunk_robust.py: cgroup v2/v1/host limit
parsing, available memory, exact peak; fixed returns chunk_size unchanged; budget bounds,
safety scaling, division across workers, and zero-headroom behavior; chunks partition the
choosers exactly; auto and fixed produce the same simple_simulate result; the probe cap and the auto
growth-cap default; auto replaces a passed static chunk_size but preserves chunk_size=0;
settings validation; the get_peak_rss fallback used when the Unix-only resource
module is absent (Windows); and the audit line, logged once per process. black clean.

Full-population benchmark — 1.22 M households (8.1 M trips), 4 workers, Kubernetes containers
of two sizes, every run on this code. Two budget settings crossed with all three training modes.
Production runs read the cache written by the training run of the same budget; adaptive runs
start with an empty cache.

budget container training production adaptive
fixed, 10 GB per worker (tuned for this container) 60 GiB ✅ 222 min ❌ OOM after ~20 min ✅ 228 min
auto, safety 0.5 60 GiB ✅ 228 min ✅ 192 min ✅ 227 min
auto, safety 0.5, default_initial_rows_per_chunk: 1000 40 GiB ✅ 271 min ✅ 223 min ✅ 257 min
  1. Auto completed all six runs, across two container sizes and all three training modes. The
    configuration was identical for both containers; only the container changed. The logged
    per-worker budget adapted from about 8.3 GB to about 3.7 GB.
  2. Production mode failed under the static budget and succeeded under auto. The cache it read was
    accurate — written by a training run on the same machine with the same settings — so cache
    accuracy is not what makes production mode safe. Under auto it was also the fastest
    configuration measured (192 min, 16% below training), because it skips measurement.
  3. The safety factor matters at the observed margin. At 0.5 all six runs completed. At 0.7 the
    same configuration completed once and was OOM-killed once.
  4. Chunking does not change results. All eight completed runs — both budgets, three training
    modes, two container sizes, and therefore different chunk boundaries — produced identical row
    counts and a byte-identical households table (same MD5 for all eight).
  5. Auto's cost against a well-tuned static baseline is about 3% of wall time in training mode
    (228 min vs 222 min).

Compatibility

chunk_size_mode defaults to fixed, which returns the static chunk_size unchanged; a test
asserts this. None of the auto defaults apply in fixed mode.

One change applies to every mode: the per-segment cache tags. Existing chunk_cache.csv entries
for location_choice, tour_destination and trip_destination will not match the new tags. One
training run rebuilds them. In the meantime production mode falls back to the capped probe chunk
for those components. This is noted in changes.md.

Known limitation

Late in a run the cgroup's charged page cache raises reported usage, which lowers the computed
budget. Auto therefore becomes more conservative as a long run proceeds: smaller chunks and more
"memory is tight" warnings. This is safe but noisy. Subtracting reclaimable inactive_file from
the usage figure would address it and is left for follow-up work.

Automatic selection of num_processes is deliberately out of scope. Where the skims are
memory-mapped, the shared buffer reported before workers fork does not reflect their real
footprint, so the worker count cannot be sized reliably at that point.

…ory limit

Adaptive chunking sizes chunks against a static `chunk_size` byte budget that the user must hand-tune
per machine, and it targets host RAM. Inside a container (k8s/cgroup) the process is OOM-killed at the
cgroup limit, not host RAM, so a chunk_size set from host RAM over-commits and the run dies. This adds
an opt-in `chunk_size_mode: auto` that derives the chunk budget at runtime from the process's real
memory ceiling (cgroup v2 memory.max -> v1 -> psutil), scaled by chunk_size_safety_factor and, in
multiprocess, divided by the per-step worker count. It reuses the existing adaptive machinery and its
measured, cached row_size; it only changes where the byte budget comes from.

Details:
- Budget = (memory_limit - current usage) * chunk_size_safety_factor. Computed per model at runtime,
  so it tracks the memory actually resident (framework + skims paged in). `available == 0` means no
  headroom, not "unknown" -> full limit. The budget is floored only to a positive value so chunking
  stays active; it is NOT floored to a large per-worker minimum (that would sum across workers and
  over-commit). A very small budget warns instead.
- Multiprocess: divide the budget by the per-step worker count (the num_processes injectable), so the
  N workers sharing the ceiling don't collectively exceed it.
- Accuracy: get_peak_rss() (exact getrusage ru_maxrss peak, Windows-safe fallback) and
  chunk_row_size_margin improve the measured row_size; chunk_growth_cap bounds chunk-to-chunk growth.

All new behavior is gated on chunk_size_mode (default `fixed` = current behavior unchanged). Adds
unit tests (core/test/test_mem.py, test_chunk_robust.py) and a docs/core.rst section.
@vincentgong7
vincentgong7 force-pushed the feature/chunk-size-mode-auto branch from 318a735 to d769c36 Compare August 20, 2026 17:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant