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

[vincentgong7]

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.

4. 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.
5. 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_bytes → psutil 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.py — resolve_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.

[vincentgong7]

Thanks @janzill and all for the discussion at the 20th August 2026 meeting, and for writing up such detailed notes. The question of whether this need could be met with explicit chunking, conservative settings or deployment overlays is a fair one, and I would rather answer it plainly than argue for the PR. Where those alternatives work, I will say so.

The deployment

We run the Rotterdam (a city in the Netherlands) case on a small Kubernetes cluster that is shared with other services. Runs are submitted as Jobs from a web front end, so a run is not launched by a person sitting at a machine — it is scheduled onto whichever node has room at that moment. Nodes have 64 GB. The memory a run actually gets depends on which node it lands on and what else is running there at the time, and it is not the same from one run to the next. Skims are memory-mapped and are already resident when chunking starts.

So the number we would need to hand-tune, chunk_size, is a property of an environment that does not exist until the pod starts.

On the three alternatives

Explicit row-count chunking. This works well when the machine is known, and it is what we use wherever we can. Mode choice was the last core step we could not bound that way, which is why I contributed PR #1088 to plumb explicit_chunk through the mode-choice components; that is what let a full-sample run of 8.1M trips over 7787 zones fit on a 64 GB node. So I recognise Amir's experience directly: once row counts are calibrated for a machine they are easy to carry forward and quick to adjust.

Our difficulty is not that calibration is hard, it is that we do not know at authoring time which allocation the run will receive, and a row count that is right for the larger case is fatal for the smaller one. If our allocations were a fixed set of two or three known sizes, explicit chunking would cover us, with less code than this PR. I want to be clear about that.

Conservative settings for the smallest allocation. This also works, and it is what we did before. Two things pushed us off it. First, chunk_size is a byte budget for the whole run that each worker then takes in full, so a correct setting requires dividing by num_processes by hand and subtracting what is already resident; both are easy to get wrong, and the failure appears deep into a long multiprocess run rather than at startup. Second, sizing everything for the smallest allocation gives up the headroom on the runs that did get a large node. That second cost is real for us but I would not argue it is large in general — Amir's observation that runtime is not very sensitive to chunk size matches what we see, and I have not measured it carefully enough to claim otherwise.

Deployment-specific settings overlays. This is the closest alternative, and for deployments with an enumerable set of sizes I think it is the better answer. It stops working when the ceiling is continuous rather than one of a few known values.

What this PR does not address

Bo's case is a good check on scope. A workplace-location model that reaches 100 GB during preprocessing joins, before chunking takes effect, is not helped by this PR at all — the budget only governs the chunk loop, so a more accurate budget changes nothing upstream of it. The same is true of explicit chunking. I do not want the PR read as addressing that class of problem.

I only know that case from the notes, so I offer this tentatively: if the memory really is going to preprocessing rather than to the chunked work, then what bounds it is how far the chunk loop reaches rather than how the budget inside it is computed — which would be a different piece of work from either chunking mode.

On adaptive chunking not having been robust

I agree with that characterisation, and I think one part of why is measurable rather than a matter of opinion. Every backoff warning logs an exact (peak, budget) pair, so archived runs contain direct measurements of how far a chunk's transient peak overshoots the budget it was sized against. Across twelve archived full-population runs there are 802 such measurements. The overshoot ratio has a median of 1.08 and a maximum of 1.67.

That happens to explain the default. chunk_size_safety_factor: 0.7 tolerates an overshoot of 1.43, which sits between the 99th percentile (1.42) and the observed maximum (1.67) — so most runs survive and some do not, which is the behaviour we saw. 0.5 tolerates 2.0 and covers the observed maximum with margin.

The caveat matters as much as the number: this is one workload on one cluster, and all of it was measured with backoff already active, which truncates the tail. I offer it as a method that any deployment can run against its own logs, not as a constant anyone should trust from us. If it would be useful for the improve-or-deprecate discussion, I am happy to write up how the numbers were extracted.

A separate direction, if it is of interest

One thing that came out of the above is that the deeper fragility is not in any particular sizing strategy. It is that every strategy — adaptive, explicit, conservative — has to be right the first time, because being wrong means the run dies rather than slows down.

That property can be removed. On Linux, giving each worker an RLIMIT_DATA soft cap below the container limit turns an over-large allocation into a catchable MemoryError in the offending worker, instead of the cgroup group-kill that takes the whole container down with it. Memory-mapped skims are file-backed and exempt from that limit, so it caps the chunk working set and nothing else. The chunk loop can then halve the chunk and retry. On Windows the allocator already refuses over-commitment at allocation time, so the same retry works there without any cap.

The reason I mention it here rather than adding it to this PR is that it is independent of chunk_size_mode, and would protect explicit chunking just as well — a row count calibrated for a 256 GB machine would survive being run on a 128 GB one, slowly, instead of failing. If the consortium's decision goes towards deprecating adaptive chunking, this would still apply.

It is implemented and validated, single-process and multi-process, with the platform behaviour verified on both Linux and Windows. A retried chunk reproduces an unfailed run exactly: each row's random draws are seeded from its own index, and a retry rewinds the stream position of the rows it re-runs, so whether a MemoryError occurred does not change any result. On prototype_mtc with two worker processes, forcing a memory failure in both workers produced byte-identical output to an unforced run across every final table.

Where the cap is set below a component's baseline working set, the failure happens before there is any chunk to retry, and the run ends the way it does today — a catchable exception and a traceback, rather than the container being killed as a group. That degradation is the designed floor of the mechanism, not something it papers over.

I have not put it forward as a PR because the current one has not been reviewed, and I did not want to expand its scope. If it is worth having, I would open it separately.

Happy to answer anything that would help the discussion.