Parallelize chunked matrix building approach (#818)

[juaristi22]

Fix #818.

Summary

Parallelizes UnifiedMatrixBuilder.build_matrix_chunked() across Modal workers behind an opt-in flag. Extracts the ~450-line inline chunked-matrix logic into a ChunkedMatrixAssembler coordinator class, replaces the list-then-concat final assembly with a two-pass streaming CSR build, and adds a Modal dispatch layer that fans per-chunk work across up to 50 workers.

The parallel path is opt-in (--parallel in combination with --chunked-matrix). Existing serial and non-chunked paths remain the default and are unchanged.

What changed

Phase 1 — refactor behind a coordinator (7a24547)

• New policyengine_us_data/calibration/chunked_matrix_assembler.py:
  • SharedBuildState, ChunkPlan, ChunkResult dataclasses.
  • partition_chunks() pure helper.
  • stream_csr_from_shards() — two-pass streaming CSR assembly: pass 1 counts per-row nonzeros to compute indptr; pass 2 scatters entries into preallocated data/indices arrays. No COO intermediate, no scipy-internal doubling.
  • ChunkedMatrixAssembler class with run_chunks, run_single_chunk, assemble_final.
• UnifiedMatrixBuilder.build_matrix_chunked() shrinks from ~450 lines to an ~80-line facade; public signature unchanged. Per-chunk loop and list-then-concat final assembly removed.
• _build_entity_relationship lifted to module scope so SharedBuildState stays pickle-clean for cross-process dispatch.

Phase 2 — Modal dispatch (0846a54, 1607259, 9429aaf, 81268f3)

• New policyengine_us_data/calibration/chunked_matrix_modal.py — coordinator-side dispatch. Pickles SharedBuildState to the pipeline volume, partitions chunk ids into contiguous batches, spawns workers via modal.Function.from_name, aggregates per-chunk errors, streams the final CSR after workers finish.
• New modal_app/matrix_chunk_worker.py — build_matrix_chunk_worker @app.function registered on _calibration_app. pipeline.py imports the worker and calls app.include(_calibration_app), so a single modal deploy modal_app/pipeline.py registers it under policyengine-us-data-pipeline.
• UnifiedMatrixBuilder.build_matrix_chunked() gains parallel, num_matrix_workers, run_id kwargs — all default off/50/empty.
• unified_calibration CLI gains --parallel / --no-parallel and --num-matrix-workers. run_id threads through the POLICYENGINE_US_DATA_RUN_ID env var set by build_package_remote.
• build_package_remote and _build_package_impl accept and forward chunked_matrix, chunk_size, parallel_matrix, num_matrix_workers to the unified_calibration subprocess.
• modal_app/pipeline.py: run_pipeline gains matching kwargs and forwards them to build_package_remote.remote().
• .github/workflows/pipeline.yaml: four new workflow_dispatch inputs — chunked_matrix, chunk_size, parallel_matrix, num_matrix_workers — all defaulting to off/25000/off/50.

Memory improvement

Streaming CSR assembly replaces the three-list-concat that previously held every shard in memory before building the CSR. Measured on a 5M-nnz synthetic fixture: peak RSS during assembly is 1.22× final CSR size (46.6 MiB RSS delta vs. 38.2 MiB final CSR bytes), compared with an estimated 2-3× for the old path.

Behavior by flag combination

chunked_matrix	parallel_matrix	Path
false	false	Non-chunked build_matrix (unchanged)
true	false	In-process serial chunked (identical to pre-PR --chunked-matrix)
false	true	Non-chunked; info log that --parallel was ignored
true	true	Modal dispatch, up to num_matrix_workers workers

Test coverage

Layer	File	Tests
Pure helpers + streaming CSR + assembler orchestration	tests/unit/calibration/test_chunked_matrix_assembler.py	11
Dispatch layer with fake worker + fake volume	tests/unit/calibration/test_chunked_matrix_modal.py	9
Fixture-scale byte-equivalence (chunked CSR ≡ non-chunked CSR)	tests/integration/test_chunked_matrix_builder.py	6 (unchanged, pass against the facade)
Env-gated Modal lookup smoke	tests/integration/test_matrix_chunk_worker_modal.py	2 (skipped unless MODAL_TOKEN_ID + MODAL_TOKEN_SECRET + POLICYENGINE_US_DATA_MODAL_SMOKE=1)

Full calibration unit suite: 392 passed, 7 skipped, 0 failed (up from 382 before phase 2).

Post-merge validation

The next automatic pipeline run (triggered by push.yaml's auto-generated Update package version commit) uses the default flags (off) and behaves identically to main today. To exercise the parallel path as a one-off validation run:

gh workflow run pipeline.yaml --ref main \
  -f chunked_matrix=true \
  -f parallel_matrix=true \
  -f num_matrix_workers=50

IMPORTANT: Do not trigger the pipeline run while a production job is running without specifying an alternative Modal environment as it might affect its processes.

Outstanding

• Pre-merge Modal deploy smoke not yet run — blocked on an in-flight production run on the same Modal environment.
• First real end-to-end parallel run will be the post-merge manual dispatch above.
• Full-CPS wall-time measurement will be added here after that run.

Related

• Depends on Pipeline environment fixes #819 (Modal venv-activation fixes) — merged, now on main.
• Implementation plan: Chunked Matrix Parallelization Plan.md at repo root.
• Long-term refactor this aligns with: US Data Pipeline Refactor.md Phase 4 (MatrixAssembler extraction).

Test plan

• pytest tests/unit/calibration/ tests/integration/test_chunked_matrix_builder.py — 398 pass (392 unit + 6 integration), 7 skipped, 0 failed.
• ruff check and ruff format --check clean on all changed files.
• modal deploy modal_app/pipeline.py succeeds and registers build_matrix_chunk_worker under policyengine-us-data-pipeline.
• POLICYENGINE_US_DATA_MODAL_SMOKE=1 pytest tests/integration/test_matrix_chunk_worker_modal.py — both tests pass.
• One-off gh workflow run pipeline.yaml dispatch with the parallel flags completes end-to-end; output CSR matches a serial baseline; wall time documented above.

🤖 Generated with Claude Code

[juaristi22]

Ready for review but not merge until smoke testing modal deployment.