Parallelize chunked matrix build across Modal workers (#818)

Phase 2 of issue #818. Builds on the phase-1 coordinator extraction
to fan chunked matrix building out across Modal workers, cutting the
full-CPS build's wall time from the ~14 h serial ceiling into a
wall-time window proportional to one chunk.

Changes:
- `policyengine_us_data/calibration/chunked_matrix_modal.py`: new
  `dispatch_chunks_modal` coordinator. Pickles `SharedBuildState` to
  `{chunk_root}/chunk_build_state.pkl` on the pipeline volume,
  partitions `range(n_chunks)` into contiguous batches, spawns one
  `build_matrix_chunk_worker` per batch via `modal.Function.from_name`,
  collects results, aggregates per-chunk errors, and streams the
  final CSR from shards on the volume. Contiguous batching (not
  round-robin) keeps resume-friendly prefixes on disk.
- `modal_app/matrix_chunk_worker.py`: new `build_matrix_chunk_worker`
  `@app.function` registered on the existing
  `policyengine-us-data-fit-weights` app (same app as
  `build_package_remote`). Mirrors `build_areas_worker` resources
  (memory=16384, cpu=1.0, timeout=28800, max_containers=50,
  nonpreemptible). Worker reads the pickled shared state, constructs
  a `ChunkedMatrixAssembler` with `resume=True`, calls `run_chunks`
  on its batch, commits the volume, returns nnz + errors.
- `SharedBuildState` grows a `chunk_size` field so a worker can
  reconstruct the assembler from the pickle alone (no extra args).
- `build_matrix_chunked` accepts `parallel`, `num_matrix_workers`,
  `run_id`; routes to `dispatch_chunks_modal` when `parallel=True`,
  preserves the in-process serial path otherwise.
- `unified_calibration` CLI gains `--parallel` (default off) and
  `--num-matrix-workers` (default 50). `--parallel` without
  `--chunked-matrix` logs an info message and runs the non-chunked
  path unchanged. `run_id` flows from env
  `POLICYENGINE_US_DATA_RUN_ID`, set by `build_package_remote`.
- `build_package_remote` / `_build_package_impl` take
  `chunked_matrix`, `chunk_size`, `parallel_matrix`,
  `num_matrix_workers` kwargs and forward them to the
  `unified_calibration` subprocess.

Tests:
- 9 unit tests in `tests/unit/calibration/test_chunked_matrix_modal.py`
  cover `partition_chunk_ids_contiguous` (exact division, remainder,
  more workers than chunks, zero chunks, invalid num_workers) and
  `dispatch_chunks_modal` with injected `worker_function` + `volume`
  fakes (spawn/assemble happy path, zero-chunks short-circuit, error
  aggregation, shared-state pickle side effect).
- New `test_shared_build_state_roundtrips_pickle` in the assembler
  unit tests guards the phase-2 boundary: if `SharedBuildState`
  stops pickling cleanly, we catch it here rather than in a Modal
  worker. 11 assembler unit tests total.
- `tests/integration/test_matrix_chunk_worker_modal.py` is an
  env-gated smoke test (MODAL_TOKEN_ID + MODAL_TOKEN_SECRET +
  POLICYENGINE_US_DATA_MODAL_SMOKE=1) that validates the deployed
  worker is lookupable via `modal.Function.from_name` and that
  production-scale batching is sane. Full end-to-end validation
  (write shared state, spawn, verify shards on the volume) is a
  pre-merge manual step documented in the PR.

Verified:
- 392 calibration unit tests pass (up from 382; 10 new phase-2
  tests; 0 regressions), plus the 6 existing chunked-matrix
  integration tests still green against the phase-1 facade.
- `ruff check` and `ruff format --check` clean on all changed files.

Validation still pending on real Modal (deploy + one-chunk
benchmark + full-CPS `workflow_dispatch` run); gated on the Modal
venv-activation PR landing.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Author: juaristi22

changelog.d/818.changed.md

@@ -1 +1,3 @@
 Extract `ChunkedMatrixAssembler` from `UnifiedMatrixBuilder.build_matrix_chunked` and replace the list-then-concat final assembly with a two-pass streaming CSR build. The facade signature and all existing chunked-matrix behaviour are unchanged.
+
+Parallelize chunked matrix building across Modal workers. Adds `dispatch_chunks_modal` (`policyengine_us_data/calibration/chunked_matrix_modal.py`) that pickles `SharedBuildState` to the pipeline volume, fans contiguous chunk-id batches to `build_matrix_chunk_worker` (`modal_app/matrix_chunk_worker.py`, registered on the `policyengine-us-data-fit-weights` app), and streams the final CSR from shards on the volume. New CLI flags on `unified_calibration`: `--parallel` (default off) and `--num-matrix-workers` (default 50). `build_package_remote` threads `run_id` to the subprocess via the `POLICYENGINE_US_DATA_RUN_ID` env var and forwards `--parallel` / `--num-matrix-workers` when `parallel_matrix=True`. `--parallel` without `--chunked-matrix` logs an info message and runs the non-chunked path unchanged.

modal_app/matrix_chunk_worker.py

@@ -0,0 +1,112 @@
+"""Modal worker that materializes a batch of matrix chunks.
+
+Registered on the same ``policyengine-us-data-fit-weights`` app as
+``build_package_remote`` so the coordinator can spawn workers via
+``modal.Function.from_name`` from inside the package build's
+subprocess. Each worker reads the shared ``ChunkedMatrixAssembler``
+state from ``pipeline_volume``, materializes its assigned chunks to
+COO shard files on the volume, and commits. The coordinator reads the
+shards back after all workers finish and streams them into the final
+CSR matrix.
+"""
+
+from __future__ import annotations
+
+import pickle
+import sys
+import traceback
+from pathlib import Path
+from typing import Dict, List
+
+_baked = "/root/policyengine-us-data"
+_local = str(Path(__file__).resolve().parent.parent)
+for _p in (_baked, _local):
+    if _p not in sys.path:
+        sys.path.insert(0, _p)
+
+from modal_app.images import cpu_image  # noqa: E402
+from modal_app.remote_calibration_runner import (  # noqa: E402
+    PIPELINE_MOUNT,
+    app,
+    hf_secret,
+    pipeline_vol,
+)
+
+
+def _chunk_root(run_id: str) -> str:
+    return f"{PIPELINE_MOUNT}/artifacts/{run_id}/matrix_build"
+
+
+@app.function(
+    image=cpu_image,
+    secrets=[hf_secret],
+    volumes={PIPELINE_MOUNT: pipeline_vol},
+    memory=16384,
+    cpu=1.0,
+    timeout=28800,
+    max_containers=50,
+    nonpreemptible=True,
+)
+def build_matrix_chunk_worker(run_id: str, chunk_ids: List[int]) -> Dict:
+    """Materialize ``chunk_ids`` from the pickled ``SharedBuildState``.
+
+    Args:
+        run_id: Pipeline run identifier; selects the volume path for
+            this worker's shared state and shard output directory.
+        chunk_ids: Chunk indices this worker is responsible for.
+
+    Returns:
+        Dict with ``chunk_ids``, ``nnz_per_chunk``, and ``errors``
+        lists suitable for the coordinator to aggregate.
+    """
+    from policyengine_us_data.calibration.chunked_matrix_assembler import (
+        ChunkedMatrixAssembler,
+    )
+
+    pipeline_vol.reload()
+    chunk_root = Path(_chunk_root(run_id))
+    state_path = chunk_root / "chunk_build_state.pkl"
+    if not state_path.exists():
+        return {
+            "chunk_ids": list(chunk_ids),
+            "nnz_per_chunk": [],
+            "errors": [
+                {
+                    "chunk_ids": list(chunk_ids),
+                    "error": f"Missing shared state at {state_path}",
+                }
+            ],
+        }
+
+    with open(state_path, "rb") as f:
+        shared_state = pickle.load(f)
+
+    assembler = ChunkedMatrixAssembler(
+        shared_state=shared_state,
+        chunk_root=chunk_root,
+        chunk_size=shared_state.chunk_size,
+        resume=True,
+        keep_chunks=False,
+    )
+
+    errors: List[Dict] = []
+    nnz_per_chunk: List[int] = []
+    for chunk_id in chunk_ids:
+        try:
+            result = assembler.run_single_chunk(chunk_id)
+            nnz_per_chunk.append(result.nnz)
+        except Exception as exc:
+            errors.append(
+                {
+                    "chunk_id": chunk_id,
+                    "error": str(exc),
+                    "traceback": traceback.format_exc(),
+                }
+            )
+
+    pipeline_vol.commit()
+    return {
+        "chunk_ids": list(chunk_ids),
+        "nnz_per_chunk": nnz_per_chunk,
+        "errors": errors,
+    }

modal_app/remote_calibration_runner.py

@@ -366,6 +366,10 @@ def _build_package_impl(
     workers: int = 8,
     n_clones: int = 430,
     run_id: str = "",
+    chunked_matrix: bool = False,
+    chunk_size: int = 25_000,
+    parallel_matrix: bool = False,
+    num_matrix_workers: int = 50,
 ) -> str:
     """Read data from pipeline volume, build X matrix, save package."""
     _setup_repo()
@@ -408,10 +412,26 @@ def _build_package_impl(
     if workers > 1:
         cmd.extend(["--workers", str(workers)])
     cmd.extend(["--n-clones", str(n_clones)])
+    if chunked_matrix:
+        cmd.extend(["--chunked-matrix", "--chunk-size", str(chunk_size)])
+        if parallel_matrix:
+            cmd.extend(
+                [
+                    "--parallel",
+                    "--num-matrix-workers",
+                    str(num_matrix_workers),
+                ]
+            )
 
+    build_env = os.environ.copy()
+    if run_id:
+        # ``unified_calibration.py`` reads this env var so workers can
+        # locate their shared state at {pipeline-artifacts}/{run_id}/
+        # matrix_build/chunk_build_state.pkl on the pipeline volume.
+        build_env["POLICYENGINE_US_DATA_RUN_ID"] = run_id
     build_rc, build_lines = _run_streaming(
         cmd,
-        env=os.environ.copy(),
+        env=build_env,
         label="build",
     )
     if build_rc != 0:
@@ -450,6 +470,10 @@ def build_package_remote(
     workers: int = 8,
     n_clones: int = 430,
     run_id: str = "",
+    chunked_matrix: bool = False,
+    chunk_size: int = 25_000,
+    parallel_matrix: bool = False,
+    num_matrix_workers: int = 50,
 ) -> str:
     return _build_package_impl(
         branch,
@@ -458,6 +482,10 @@ def build_package_remote(
         workers=workers,
         n_clones=n_clones,
         run_id=run_id,
+        chunked_matrix=chunked_matrix,
+        chunk_size=chunk_size,
+        parallel_matrix=parallel_matrix,
+        num_matrix_workers=num_matrix_workers,
     )
 
 

policyengine_us_data/calibration/chunked_matrix_assembler.py

@@ -67,6 +67,7 @@ class SharedBuildState:
     n_records: int
     n_clones: int
     n_targets: int
+    chunk_size: int
     target_variables: List[str]
     target_reform_ids: List[int]
     target_geo_info: List[Tuple[str, str]]

policyengine_us_data/calibration/chunked_matrix_modal.py

@@ -0,0 +1,208 @@
+"""Coordinator-side Modal dispatch for chunked matrix building.
+
+Writes shared ``SharedBuildState`` to a pipeline-volume path,
+spawns ``build_matrix_chunk_worker`` per contiguous batch of
+chunk ids, collects per-worker results, then streams the final CSR
+from all shards on the volume.
+
+Kept separate from ``unified_matrix_builder`` so the core matrix
+builder doesn't import Modal; only the dispatch path does.
+"""
+
+from __future__ import annotations
+
+import logging
+import math
+import pickle
+import time
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from scipy import sparse
+
+from policyengine_us_data.calibration.chunked_matrix_assembler import (
+    ChunkedMatrixAssembler,
+    SharedBuildState,
+)
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_NUM_MATRIX_WORKERS = 50
+MODAL_APP_NAME = "policyengine-us-data-fit-weights"
+WORKER_FUNCTION_NAME = "build_matrix_chunk_worker"
+
+
+def partition_chunk_ids_contiguous(n_chunks: int, num_workers: int) -> List[List[int]]:
+    """Split ``range(n_chunks)`` into ``num_workers`` contiguous batches.
+
+    Contiguous (not round-robin) so that a partial run leaves complete
+    prefixes on disk that future `--resume-chunks` invocations can
+    skip cleanly. Returns at most ``num_workers`` non-empty batches.
+    """
+    if n_chunks <= 0:
+        return []
+    if num_workers <= 0:
+        raise ValueError("num_workers must be positive")
+    batch_size = math.ceil(n_chunks / num_workers)
+    batches: List[List[int]] = []
+    for start in range(0, n_chunks, batch_size):
+        end = min(start + batch_size, n_chunks)
+        batches.append(list(range(start, end)))
+    return batches
+
+
+def _lookup_worker_function():
+    """Resolve the deployed Modal worker function.
+
+    Using ``Function.from_name`` avoids importing the worker module
+    here (Modal imports are heavy and would pull into every caller).
+    It also means unit tests can monkeypatch this function without
+    touching the worker module.
+    """
+    import modal
+
+    return modal.Function.from_name(MODAL_APP_NAME, WORKER_FUNCTION_NAME)
+
+
+def dispatch_chunks_modal(
+    *,
+    shared_state: SharedBuildState,
+    chunk_root: Path,
+    run_id: str,
+    num_workers: int = DEFAULT_NUM_MATRIX_WORKERS,
+    worker_function: Optional[Any] = None,
+    volume: Optional[Any] = None,
+) -> sparse.csr_matrix:
+    """Fan chunk materialization across Modal workers, then assemble.
+
+    Args:
+        shared_state: Read-only per-build state; pickled once to the
+            volume so workers can reconstruct the assembler without
+            receiving the arrays through ``.spawn()`` args.
+        chunk_root: Directory on the pipeline volume where shards land
+            (``{chunk_root}/coo/chunk_XXXXXX.npz``) and where the
+            shared state pickle lives
+            (``{chunk_root}/chunk_build_state.pkl``).
+        run_id: Forwarded to each worker so its volume paths align
+            with the coordinator's.
+        num_workers: Upper bound on workers; actual count equals
+            ``min(num_workers, n_chunks)``.
+        worker_function: Override for the Modal function (tests only).
+        volume: Override for the pipeline volume (tests only). When
+            omitted, resolves ``modal.Volume.from_name("pipeline-artifacts")``.
+
+    Raises:
+        RuntimeError: if any worker reports one or more chunk errors
+            after all workers finish. Raised after aggregating so no
+            errors are silently dropped.
+    """
+    chunk_root = Path(chunk_root)
+    chunk_root.mkdir(parents=True, exist_ok=True)
+    state_path = chunk_root / "chunk_build_state.pkl"
+    with open(state_path, "wb") as f:
+        pickle.dump(shared_state, f)
+
+    if volume is None:
+        import modal
+
+        volume = modal.Volume.from_name("pipeline-artifacts", create_if_missing=True)
+    # Make the shared-state pickle visible to workers.
+    volume.commit()
+
+    n_chunks = math.ceil(shared_state.n_total / shared_state.chunk_size)
+    batches = partition_chunk_ids_contiguous(n_chunks, num_workers)
+
+    if not batches:
+        # Nothing to materialize; fall through to assembly (which will
+        # return an empty CSR).
+        volume.reload()
+        assembler = ChunkedMatrixAssembler(
+            shared_state=shared_state,
+            chunk_root=chunk_root,
+            chunk_size=shared_state.chunk_size,
+            resume=True,
+            keep_chunks=False,
+        )
+        return assembler.assemble_final()
+
+    if worker_function is None:
+        worker_function = _lookup_worker_function()
+
+    logger.info(
+        "Dispatching %d chunks across %d workers (batch sizes: %s)",
+        n_chunks,
+        len(batches),
+        [len(b) for b in batches[:5]] + (["..."] if len(batches) > 5 else []),
+    )
+    t_dispatch = time.time()
+    handles = []
+    for batch_idx, chunk_ids in enumerate(batches):
+        handle = worker_function.spawn(run_id=run_id, chunk_ids=chunk_ids)
+        logger.info(
+            "Worker %d/%d: %d chunks (%d-%d), fc=%s",
+            batch_idx + 1,
+            len(batches),
+            len(chunk_ids),
+            chunk_ids[0],
+            chunk_ids[-1],
+            getattr(handle, "object_id", "unknown"),
+        )
+        handles.append((batch_idx, handle))
+
+    aggregated_errors: List[Dict] = []
+    for batch_idx, handle in handles:
+        try:
+            result = handle.get()
+        except Exception as exc:
+            aggregated_errors.append(
+                {
+                    "batch": batch_idx,
+                    "error": f"Worker crashed: {exc}",
+                }
+            )
+            logger.error("Worker %d crashed: %s", batch_idx, exc)
+            continue
+        if result is None:
+            aggregated_errors.append(
+                {"batch": batch_idx, "error": "Worker returned None"}
+            )
+            continue
+        errors = result.get("errors", [])
+        if errors:
+            for err in errors:
+                err_copy = dict(err)
+                err_copy["batch"] = batch_idx
+                aggregated_errors.append(err_copy)
+        logger.info(
+            "Worker %d done: %d chunks completed, %d errors",
+            batch_idx,
+            len(result.get("chunk_ids", [])) - len(errors),
+            len(errors),
+        )
+
+    logger.info(
+        "All workers finished in %.1fs; %d errors total",
+        time.time() - t_dispatch,
+        len(aggregated_errors),
+    )
+
+    if aggregated_errors:
+        preview = "; ".join(
+            f"batch {e['batch']}: {e.get('error', 'unknown')[:120]}"
+            for e in aggregated_errors[:3]
+        )
+        raise RuntimeError(
+            f"Parallel chunked matrix build failed with "
+            f"{len(aggregated_errors)} error(s). First: {preview}"
+        )
+
+    # All shards present on the volume; reload and assemble.
+    volume.reload()
+    assembler = ChunkedMatrixAssembler(
+        shared_state=shared_state,
+        chunk_root=chunk_root,
+        chunk_size=shared_state.chunk_size,
+        resume=True,
+        keep_chunks=False,
+    )
+    return assembler.assemble_final()