PolicyEngine
diff --git a/‎.github/scripts/spawn_modal_pipeline.py‎
Lines changed: 15 additions & 0 deletions b/‎.github/scripts/spawn_modal_pipeline.py‎
Lines changed: 15 additions & 0 deletions
diff --git a/‎.github/workflows/pipeline.yaml‎
Lines changed: 20 additions & 0 deletions b/‎.github/workflows/pipeline.yaml‎
Lines changed: 20 additions & 0 deletions
diff --git a/‎changelog.d/818.changed.md‎
Lines changed: 5 additions & 0 deletions b/‎changelog.d/818.changed.md‎
Lines changed: 5 additions & 0 deletions
diff --git a/‎modal_app/data_build.py‎
Lines changed: 1 addition & 0 deletions b/‎modal_app/data_build.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎modal_app/local_area.py‎
Lines changed: 1 addition & 0 deletions b/‎modal_app/local_area.py‎
Lines changed: 1 addition & 0 deletions
diff --git a/‎modal_app/matrix_chunk_worker.py‎
Lines changed: 125 additions & 0 deletions b/‎modal_app/matrix_chunk_worker.py‎
Lines changed: 125 additions & 0 deletions
diff --git a/‎modal_app/pipeline.py‎
Lines changed: 63 additions & 6 deletions b/‎modal_app/pipeline.py‎
Lines changed: 63 additions & 6 deletions
@@ -15,6 +15,10 @@ def _as_bool(value: str) -> bool:
     return value.lower() == "true"
 
 
+def _env(name: str, default: str) -> str:
+    return os.environ.get(name, default)
+
+
 def _append_summary(function_call_id: str, context: RunContext) -> None:
     summary_path = os.environ.get("GITHUB_STEP_SUMMARY")
     if not summary_path:
@@ -36,6 +40,13 @@ def _append_summary(function_call_id: str, context: RunContext) -> None:
         handle.write(f"| HF staging | `{context.hf_staging_prefix}` |\n")
         if os.environ.get("SOURCE_SHA"):
             handle.write(f"| Source SHA | `{os.environ['SOURCE_SHA']}` |\n")
+        handle.write(
+            "| Matrix | "
+            f"`chunked={_env('CHUNKED_MATRIX', 'false')}, "
+            f"parallel={_env('PARALLEL_MATRIX', 'false')}, "
+            f"chunk_size={_env('CHUNK_SIZE', '25000')}, "
+            f"workers={_env('NUM_MATRIX_WORKERS', '50')}` |\n"
+        )
         handle.write(f"| Function call ID | `{function_call_id}` |\n\n")
         handle.write("**[Monitor on Modal Dashboard](https://modal.com/apps)**\n")
 
@@ -58,6 +69,10 @@ def main() -> None:
         "run_context": context.to_dict(),
         "modal_app_name": context.modal_app_name,
         "modal_environment": context.modal_environment,
+        "chunked_matrix": _as_bool(_env("CHUNKED_MATRIX", "false")),
+        "chunk_size": int(_env("CHUNK_SIZE", "25000")),
+        "parallel_matrix": _as_bool(_env("PARALLEL_MATRIX", "false")),
+        "num_matrix_workers": int(_env("NUM_MATRIX_WORKERS", "50")),
     }
     if environment_name:
         run_pipeline = modal.Function.from_name(
 
@@ -39,6 +39,22 @@ on:
         description: "Exact policyengine-us-data commit SHA to deploy"
         default: ""
         type: string
+      chunked_matrix:
+        description: "Build the calibration matrix in chunks (opt-in)"
+        default: false
+        type: boolean
+      chunk_size:
+        description: "Clone-household columns per chunk"
+        default: "25000"
+        type: string
+      parallel_matrix:
+        description: "Fan chunked matrix building across Modal workers"
+        default: false
+        type: boolean
+      num_matrix_workers:
+        description: "Number of Modal workers for parallel matrix build"
+        default: "50"
+        type: string
 
 concurrency:
   group: pipeline-${{ github.run_id }}-${{ github.run_attempt }}
@@ -80,6 +96,10 @@ jobs:
           RESUME_RUN_ID: ${{ inputs.resume_run_id || '' }}
           VERSION_OVERRIDE: ${{ inputs.version_override || '' }}
           SOURCE_SHA: ${{ inputs.source_sha || github.sha }}
+          CHUNKED_MATRIX: ${{ inputs.chunked_matrix || 'false' }}
+          CHUNK_SIZE: ${{ inputs.chunk_size || '25000' }}
+          PARALLEL_MATRIX: ${{ inputs.parallel_matrix || 'false' }}
+          NUM_MATRIX_WORKERS: ${{ inputs.num_matrix_workers || '50' }}
         run: |
           modal deploy --env="${MODAL_ENVIRONMENT}" --name="${US_DATA_MODAL_APP_NAME}" --tag="${US_DATA_RUN_ID}" modal_app/pipeline.py
           python .github/scripts/spawn_modal_pipeline.py
@@ -0,0 +1,5 @@
+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.
+
+Ensure Modal parallel matrix builds write coordinator shared state to the same volume-backed `matrix_build` directory used by workers, and release coordinator-side volume file handles before the final volume reload and CSR assembly.
@@ -45,6 +45,7 @@
 pipeline_volume = modal.Volume.from_name(
     os.environ.get("US_DATA_PIPELINE_VOLUME_NAME", "pipeline-artifacts"),
     create_if_missing=True,
+    version=2,
 )
 PIPELINE_MOUNT = "/pipeline"
 
 
@@ -55,6 +55,7 @@
 pipeline_volume = modal.Volume.from_name(
     os.environ.get("US_DATA_PIPELINE_VOLUME_NAME", "pipeline-artifacts"),
     create_if_missing=True,
+    version=2,
 )
 
 VOLUME_MOUNT = "/staging"
 
@@ -0,0 +1,125 @@
+"""Modal worker that materializes a batch of matrix chunks.
+
+The ``@app.function`` decorator here attaches to ``_calibration_app``
+(declared as ``policyengine-us-data-fit-weights`` in
+``remote_calibration_runner.py``), alongside ``build_package_remote``.
+At deploy time, ``modal_app/pipeline.py`` merges that app into the
+pipeline app via ``app.include(_calibration_app)``, so after
+``modal deploy modal_app/pipeline.py`` the function is registered
+under ``policyengine-us-data-pipeline`` in Modal's registry — that's
+the name ``dispatch_chunks_modal`` uses in its
+``modal.Function.from_name`` lookup.
+
+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],
+    resume_chunks: bool = False,
+) -> 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.
+        resume_chunks: Whether to trust matching pre-existing COO shards.
+            Fresh builds pass ``False`` so workers overwrite stale chunks.
+
+    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=resume_chunks,
+        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,
+    }
@@ -115,6 +115,7 @@
 pipeline_volume = modal.Volume.from_name(
     os.environ.get("US_DATA_PIPELINE_VOLUME_NAME", "pipeline-artifacts"),
     create_if_missing=True,
+    version=2,
 )
 staging_volume = modal.Volume.from_name(
     os.environ.get("US_DATA_STAGING_VOLUME_NAME", "local-area-staging"),
@@ -129,6 +130,31 @@ def _python_cmd(*args: str) -> list[str]:
     return [sys.executable, *args]
 
 
+def _calibration_package_parameters(
+    *,
+    workers: int,
+    n_clones: int,
+    target_config: str | None,
+    skip_county: bool,
+    chunked_matrix: bool,
+    chunk_size: int,
+    parallel_matrix: bool,
+    num_matrix_workers: int,
+) -> dict:
+    """Return manifest parameters that affect package construction."""
+    effective_parallel = bool(chunked_matrix and parallel_matrix)
+    return {
+        "workers": workers if not chunked_matrix else None,
+        "n_clones": n_clones,
+        "target_config": target_config,
+        "skip_county": skip_county,
+        "chunked_matrix": bool(chunked_matrix),
+        "chunk_size": chunk_size if chunked_matrix else None,
+        "parallel_matrix": effective_parallel,
+        "num_matrix_workers": num_matrix_workers if effective_parallel else None,
+    }
+
+
 def get_pinned_sha(branch: str) -> str:
     """Get the current tip SHA for a branch from GitHub."""
     result = subprocess.run(
@@ -207,6 +233,13 @@ def archive_diagnostics(
     PACKAGE_GPU_FUNCTIONS,
 )
 
+# Import registers ``build_matrix_chunk_worker`` on ``_calibration_app``
+# so a single ``modal deploy modal_app/pipeline.py`` also deploys the
+# worker via ``app.include(_calibration_app)`` below. Without this the
+# dispatch layer's ``modal.Function.from_name`` lookup would fail at
+# runtime.
+from modal_app.matrix_chunk_worker import build_matrix_chunk_worker  # noqa: E402, F401
+
 app.include(_calibration_app)
 
 from modal_app.local_area import app as _local_area_app  # noqa: E402
@@ -767,6 +800,10 @@ def run_pipeline(
     run_context: dict | None = None,
     modal_app_name: str = "",
     modal_environment: str = "",
+    chunked_matrix: bool = False,
+    chunk_size: int = 25_000,
+    parallel_matrix: bool = False,
+    num_matrix_workers: int = 50,
 ) -> str:
     """Run the full pipeline end-to-end.
 
@@ -792,6 +829,15 @@ def run_pipeline(
         run_context: Serialized run context from the launcher workflow.
         modal_app_name: Deployed Modal app name for this run.
         modal_environment: Modal environment used for this run.
+        chunked_matrix: Build the calibration matrix in clone-household
+            chunks instead of the non-chunked path. Opt-in; default off.
+        chunk_size: Clone-household columns per chunk when
+            ``chunked_matrix`` is True.
+        parallel_matrix: Fan chunked matrix building across Modal
+            workers via ``build_matrix_chunk_worker``. Only meaningful
+            when ``chunked_matrix`` is True; ignored otherwise.
+        num_matrix_workers: Number of Modal workers when
+            ``parallel_matrix`` is True.
 
     Returns:
         The run ID for use with promote.
@@ -1030,12 +1076,16 @@ def run_pipeline(
                 "database": _artifacts_dir(run_id) / "policy_data.db",
             }
         )
-        package_parameters = {
-            "workers": num_workers,
-            "n_clones": n_clones,
-            "target_config": None,
-            "skip_county": True,
-        }
+        package_parameters = _calibration_package_parameters(
+            workers=num_workers,
+            n_clones=n_clones,
+            target_config=None,
+            skip_county=True,
+            chunked_matrix=chunked_matrix,
+            chunk_size=chunk_size,
+            parallel_matrix=parallel_matrix,
+            num_matrix_workers=num_matrix_workers,
+        )
         package_reuse = _step_reusable(
             meta,
             BUILD_CALIBRATION_PACKAGE,
@@ -1068,6 +1118,13 @@ def run_pipeline(
                 workers=num_workers,
                 n_clones=n_clones,
                 run_id=run_id,
+                modal_app_name=current_run_context.modal_app_name,
+                modal_environment=current_run_context.modal_environment,
+                pipeline_volume_name=current_run_context.pipeline_volume_name,
+                chunked_matrix=chunked_matrix,
+                chunk_size=chunk_size,
+                parallel_matrix=parallel_matrix,
+                num_matrix_workers=num_matrix_workers,
             )
             print(f"  Package at: {pkg_path}")
Original file line number	Diff line number	Diff line change
`@@ -45,6 +45,7 @@`
`45`	`45`	`pipeline_volume = modal.Volume.from_name(`
`46`	`46`	`os.environ.get("US_DATA_PIPELINE_VOLUME_NAME", "pipeline-artifacts"),`
`47`	`47`	`create_if_missing=True,`
	`48`	`+ version=2,`
`48`	`49`	`)`
`49`	`50`	`PIPELINE_MOUNT = "/pipeline"`
`50`	`51`
Original file line number	Diff line number	Diff line change
`@@ -55,6 +55,7 @@`
`55`	`55`	`pipeline_volume = modal.Volume.from_name(`
`56`	`56`	`os.environ.get("US_DATA_PIPELINE_VOLUME_NAME", "pipeline-artifacts"),`
`57`	`57`	`create_if_missing=True,`
	`58`	`+ version=2,`
`58`	`59`	`)`
`59`	`60`
`60`	`61`	`VOLUME_MOUNT = "/staging"`