feat(runner): batch simulation job execution to prevent OOM kills

[daniel-klein]

Summary

• Replace all-at-once task submission with configurable batched execution in run_simulation_job()
• Batch size configurable via MODELOPS_JOB_BATCH_SIZE env var (default: 200)
• Detect OOM kill signals in task failures and log guidance to reduce batch size

Problem

The runner held references to ALL simulation futures for the entire job duration (to support a second gather() pass for model_outputs Parquet writing). This prevented Dask from freeing any SimReturn from distributed memory, causing:

1. Worker OOM kills — distributed memory accumulates across all workers
2. Progress going backwards — when a worker dies, all results stored on it are lost and must be recomputed
3. Scheduler overload — 1,000 params × 6 targets = 7,000 tasks in the graph simultaneously; at 10K params this becomes 70,000

Solution

Submit tasks in batches. Per batch:

1. Submit sims + aggs (e.g., 200 params = 1,400 tasks)
2. Gather all results for the batch
3. Release batch references — Dask frees distributed memory
4. Next batch starts with clean worker memory

	Before	After (batch=200)
Peak task graph	7,000 (all at once)	1,400
Peak distributed memory	1,000 SimReturns	200 SimReturns
Worker death blast radius	Entire job	1 batch
Scales to 10K params	No (70K tasks)	Yes (1,400 tasks)

Test plan

• Run make submit-tiny with small batch (MODELOPS_JOB_BATCH_SIZE=2) to verify batching works
• Run make submit with 1,000 params and verify no OOM kills
• Verify model_outputs Parquet files are written correctly
• Verify target loss Parquet files match previous results

Closes #25

[daniel-klein]

Why this PR is still needed (despite #29/#30 worker-saturation fix)

Root cause: two separate memory problems

Problem 1 — Task over-saturation (fixed by #29/#30): Without worker-saturation=1.0, the scheduler piles hundreds of queued tasks onto workers, each holding result data in memory. The saturation fix limits each worker to 1 active task per thread.

Problem 2 — Cumulative memory growth in long-lived worker processes (NOT fixed by #29/#30): Each simulation allocates millions of small Python objects (numpy arrays, starsim states, etc.). When a task completes, Python frees the objects but CPython's pymalloc allocator cannot return fragmented memory pages to the OS. Worker RSS grows monotonically across sequential tasks — even with only 1 task at a time.

Why worker-saturation alone doesn't prevent OOM

Current worker config: 6 nworkers per pod × --memory-limit 4.0GiB = exactly 24 GiB pod limit. Zero headroom. Even with saturation=1.0, workers running tasks sequentially accumulate fragmented memory until they exceed 4 GiB. Since all 6 workers grow together, the combined RSS hits the 24 GiB K8s pod limit before Dask's nanny (which monitors individual workers at the 95% threshold = 3.8 GiB) can intervene. Result: K8s OOMKills the entire pod, taking all 6 workers down at once.

Observed on job-e56c635a (2026-02-19): 9 out of 21 running worker pods OOMKilled within 30 minutes, despite running the latest image with worker-saturation fix.

Why the blast radius is so large

The current gather() pattern uses wait(all_futures) — it waits for ALL tasks to complete before collecting any results. This means every worker holds results for every task it has ever run, for the entire job duration. When a worker dies (OOMKill or nanny restart), ALL its previously completed results are lost — not just the 1 in-flight task, potentially hundreds. This is what causes the dramatic "progress going backwards" effect.

With a pod OOMKill: 6 workers × N completed results each = massive result loss and recomputation.

Why as_completed doesn't fit this architecture

An incremental as_completed() gather would reduce result retention on workers, but doesn't work well here because:

• Aggregation tasks depend on grouped sim results: each agg task needs ALL replicates for a param_id as *args. Can't process individual sim results in isolation.
• Would require architectural change: either move aggregation to the runner (losing worker-side parallelism) or dynamically track replicate group completion and submit agg tasks on the fly (complex, error-prone).

Why batching works

Batching naturally fits the grouped sim→agg architecture:

1. Submit a batch of param_ids (e.g., 200) with their replicates + agg tasks
2. Dask handles within-batch sim→agg dependencies normally
3. Gather all batch results to the runner, persist to blob storage
4. Release batch references → Dask frees distributed memory
5. Next batch starts with bounded worker memory

This fixes all three aspects:

• Bounded memory: workers hold at most batch_size × replicates results, not the entire job
• Fault containment: a worker death loses at most 1 batch worth of results
• Memory recycling: between batches, fragmented worker memory is bounded because old results are released (and optionally workers can be restarted)

Production validation

This PR has been tested at scale (2,500 params × 5 replicates × 6 targets, batch_size=200) with zero OOM kills on the same cluster configuration that consistently OOM-kills without batching.