Add chunk_size_mode: auto — derive the chunk budget from the real memory limit - #1106
Open
vincentgong7 wants to merge 1 commit into
Open
Add chunk_size_mode: auto — derive the chunk budget from the real memory limit#1106vincentgong7 wants to merge 1 commit into
vincentgong7 wants to merge 1 commit into
Conversation
…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
force-pushed
the
feature/chunk-size-mode-auto
branch
from
August 20, 2026 17:40
318a735 to
d769c36
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
ActivitySim sizes each chunk from two inputs: a memory budget, and an estimate of what one
chooser row costs.
The second input is already adaptive. In training mode ActivitySim measures the real per-row
cost of every component and caches it;
adaptiveandproductionreuse and refine thosecached values. The first input is not adaptive: in training, adaptive and production mode alike
the budget is the static
chunk_sizesetting, 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: autothe budget is derived at runtime from the process's actual memoryceiling 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: fixedreturns thestatic
chunk_sizeunchanged.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.
chunk_training_modeselects how theper-row costis obtained:
trainingmeasures the real cost of each component while it runs and writes it tochunk_cache.csv.adaptivestarts from the cached cost and keeps measuring, refining it as the run proceeds.productiontrusts the cached cost and does not measure, which is what makes it fast.This machinery estimates the divisor -
per-row cost- well. The dividend — thebudget— is the staticchunk_sizesetting in all three modes.Three problems with a static budget.
the cgroup limit, not at host RAM. A value chosen for the host over-commits in a container.
not subtracted, so real headroom is smaller than the number implies.
num_processes: 6andchunk_size: 14_000_000_000, the aggregate promise is 84 GB. The user must do that divisionby 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.
yet, so its size comes from
default_initial_rows_per_chunkregardless of the budget. Thatsetting 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.
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 secondfrom 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_sizetuned for the exactcontainer 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 = ((memory_limit − memory_in_use) × chunk_size_safety_factor) / num_processes.memory_limitcomes from thecgroup (v2
memory.max→ v1memory.limit_in_bytes→psutilhost RAM), so it is the limitthat would actually kill the process.
memory_in_useis the cgroup's current usage, somemory already held is subtracted before sizing. No machine-specific
chunk_sizeis needed.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.worker count. No per-worker minimum is added, because a minimum multiplies across workers.
to 2000 rows. This is a ceiling on the existing
default_initial_rows_per_chunksetting, not anew 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_chunkas before. The cap itself is a constantrather 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.
chunk_growth_capper step (default 2.0 in auto mode), bounding how far one small measurement is extrapolated.
chunk_peak_backoff_ratioof the budget (default 0.9), the next chunk is halved.segment is sized from its own measurement.
per-worker budget, current RSS and exact lifetime peak RSS (
getrusage, which does not missshort-lived spikes). Each process also logs one
chunking settings:line listing everyeffective chunking parameter, so a run log is self-describing when reviewed later.
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_processesreplaces problem 3 with arithmetic the code performs. Because the derivation isrepeated 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_schedulingalready keys its chunk cache per segment
(
segment_chunk_tag = extend_trace_label(tour_chunk_tag, tour_segment_name)). This PR appliesthe same pattern to
location_choice,tour_destinationandtrip_destination, so a segmentis 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, …) stayout 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 1makes 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), andget_peak_rss()(lifetime peakfrom
getrusage; theresourceimport is guarded and falls back to a monotonic psutilhigh-water mark on Windows).
activitysim/core/chunk.py—resolve_chunk_size()implements the budget, and the sizerimplements the probe cap (a module constant), growth cap and peak backoff. Under auto the budget also replaces a
positive
chunk_sizepassed down by a caller;chunk_size=0is preserved, since callers use itto 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 configurationis loaded.
docs/core.rst,docs/dev-guide/changes.md— documentation and an Upcoming Changes entry,including the cache migration note.
New settings
chunk_size_modefixedfixeduses the staticchunk_size;autoderives the budget from the real memory limitchunk_size_safety_factorchunk_growth_capchunk_peak_backoff_ratiochunk_row_size_marginThe default for
chunk_size_safety_factoris 0.5 because a full-size chunk's transient peak canbe 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.pyandtest_chunk_robust.py: cgroup v2/v1/host limitparsing, available memory, exact peak;
fixedreturnschunk_sizeunchanged; budget bounds,safety scaling, division across workers, and zero-headroom behavior; chunks partition the
choosers exactly; auto and fixed produce the same
simple_simulateresult; the probe cap and the autogrowth-cap default; auto replaces a passed static
chunk_sizebut preserveschunk_size=0;settings validation; the
get_peak_rssfallback used when the Unix-onlyresourcemodule is absent (Windows); and the audit line, logged once per process.
blackclean.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.
fixed, 10 GB per worker (tuned for this container)auto, safety 0.5auto, safety 0.5,default_initial_rows_per_chunk: 1000configuration 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.
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.
same configuration completed once and was OOM-killed once.
modes, two container sizes, and therefore different chunk boundaries — produced identical row
counts and a byte-identical households table (same MD5 for all eight).
(228 min vs 222 min).
Compatibility
chunk_size_modedefaults tofixed, which returns the staticchunk_sizeunchanged; a testasserts this. None of the auto defaults apply in fixed mode.
One change applies to every mode: the per-segment cache tags. Existing
chunk_cache.csventriesfor
location_choice,tour_destinationandtrip_destinationwill not match the new tags. Onetraining 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_filefromthe usage figure would address it and is left for follow-up work.
Automatic selection of
num_processesis deliberately out of scope. Where the skims arememory-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.