Merge pull request #388 from PolicyEngine/feat/district-breakdowns

Orchestrate state results into US national results with congressional district breakdowns

Author: anth-volk

.github/workflows/pr.yml

@@ -91,27 +91,27 @@ jobs:
     name: Test integration
     needs: docker-build  # Only run if docker builds succeed
     runs-on: ubuntu-latest
-    
+
     steps:
     - uses: actions/checkout@v4
-    
+
     - name: Set up Python
       uses: actions/setup-python@v5
       with:
         python-version: '3.13'
-    
+
     - name: Install uv
       uses: astral-sh/setup-uv@v3
       with:
         enable-cache: true
-    
+
     - name: Set up Docker Buildx
       uses: docker/setup-buildx-action@v3
-    
+
     - name: Generate API clients
       run: |
         ./scripts/generate-clients.sh
-    
+
     - name: Start services
       run: |
         docker compose -f deployment/docker-compose.yml up -d
@@ -126,22 +126,30 @@ jobs:
           echo "Waiting for services... (attempt $i/30)"
           sleep 2
         done
-    
-    - name: Run integration tests
+
+    - name: Run integration tests (local services only)
       run: |
         cd projects/policyengine-apis-integ
         uv sync --extra test
-        # Run tests that don't require GCP credentials
-        # Modal staging tests run against deployed staging environment
-        simulation_integ_test_base_url="https://policyengine-staging--policyengine-simulation-gateway-web-app.modal.run" \
-          uv run pytest tests/ -v -m "not requires_gcp"
-    
+        # Run tests against local Docker services only (not Modal staging)
+        uv run pytest tests/ -v -m "not requires_gcp and not beta_only"
+
     - name: Show service logs on failure
       if: failure()
       run: |
         docker compose -f deployment/docker-compose.yml logs
-    
+
     - name: Stop services
       if: always()
       run: |
-        docker compose -f deployment/docker-compose.yml down
+        docker compose -f deployment/docker-compose.yml down
+
+  # Deploy to Modal staging and run Modal-specific integration tests
+  test-modal-integration:
+    name: Test Modal integration
+    needs: docker-build
+    uses: ./.github/workflows/modal-deploy.reusable.yml
+    with:
+      environment: beta
+      modal_environment: staging
+    secrets: inherit

projects/policyengine-api-simulation/pyproject.toml

@@ -16,11 +16,12 @@ dependencies = [
     "pydantic-settings (>=2.7.1,<3.0.0)",
     "opentelemetry-instrumentation-fastapi (>=0.51b0,<0.52)",
     "policyengine-fastapi",
-    "policyengine==0.9.0",
+    "policyengine==0.10.1",
     "policyengine-uk>=2.22.8",
     "policyengine-us>=1.370.2",
     "tables>=3.10.2",
     "modal>=0.73.0",
+    "logfire>=3.0.0",
 ]
 
 [tool.hatch.build.targets.wheel]

projects/policyengine-api-simulation/src/modal/app.py

@@ -47,7 +47,7 @@ def get_app_name(us_version: str, uk_version: str) -> str:
     .pip_install(
         f"policyengine-us=={US_VERSION}",
         f"policyengine-uk=={UK_VERSION}",
-        "policyengine==0.8.1",
+        "policyengine==0.10.0",
         "tables>=3.10.2",
         "logfire",
     )
@@ -104,3 +104,53 @@ def run_simulation(params: dict) -> dict:
             return result
     finally:
         logfire.force_flush()
+
+
+@app.function(
+    image=simulation_image,
+    cpu=2.0,
+    memory=4096,
+    timeout=7200,  # 2 hours to wait for all 52 jobs
+    retries=0,
+    secrets=[gcp_secret, logfire_secret],
+)
+def run_national_with_breakdowns(params: dict) -> dict:
+    """
+    Orchestrate parallel simulations and aggregate results.
+
+    Spawns:
+    - 1 national ECPS simulation (region="us")
+    - State-level simulations (51 states or 10 test states if _test_mode=True)
+
+    Each spawned job runs in its own container via run_simulation.
+    Returns combined national results with congressional district breakdowns.
+
+    If _test_mode=True in params, runs only 10 test states instead of all 51.
+    """
+    import logfire
+
+    from src.modal.orchestration import run_national_orchestration
+    from src.modal.utils.state_codes import TEST_STATE_CODES
+
+    configure_logfire()
+
+    # Check for test mode
+    test_mode = params.pop("_test_mode", False)
+    state_codes = TEST_STATE_CODES if test_mode else None
+
+    try:
+        with logfire.span(
+            "run_national_with_breakdowns",
+            input_params=params,
+            test_mode=test_mode,
+        ) as span:
+            result = run_national_orchestration(params, run_simulation, state_codes)
+            span.set_attribute(
+                "total_districts",
+                len(
+                    result.get("congressional_district_impact", {}).get("districts", [])
+                ),
+            )
+            return result
+    finally:
+        logfire.force_flush()

projects/policyengine-api-simulation/src/modal/gateway/endpoints.py

@@ -50,6 +50,10 @@ def get_app_name(country: str, version: Optional[str]) -> tuple[str, str]:
     return app_name, resolved_version
 
 
+NATIONAL_WITH_BREAKDOWNS = "national-with-breakdowns"
+NATIONAL_WITH_BREAKDOWNS_TEST = "national-with-breakdowns-test"
+
+
 @router.post("/simulate/economy/comparison", response_model=JobSubmitResponse)
 async def submit_simulation(request: SimulationRequest):
     """
@@ -58,19 +62,59 @@ async def submit_simulation(request: SimulationRequest):
     Matches the existing Cloud Run API endpoint path.
     Routes to the appropriate app based on country and version params.
     Returns immediately with job_id for polling.
+
+    Special handling for data="national-with-breakdowns":
+    - Only supported for country="us"
+    - Spawns 52 parallel simulations (1 national + 51 states)
+    - Returns aggregated results with congressional district breakdowns
+
+    Special handling for data="national-with-breakdowns-test":
+    - Only supported for country="us"
+    - Spawns 11 parallel simulations (1 national + 10 test states)
+    - Returns aggregated results with congressional district breakdowns
     """
     try:
         app_name, resolved_version = get_app_name(request.country, request.version)
     except ValueError as e:
         raise HTTPException(status_code=400, detail=str(e))
 
-    logger.info(f"Routing {request.country}:{resolved_version} to app {app_name}")
+    # Check for national-with-breakdowns special cases
+    payload = request.model_dump(exclude={"version"})
+    data_value = payload.get("data")
+    is_national_breakdowns = data_value in (
+        NATIONAL_WITH_BREAKDOWNS,
+        NATIONAL_WITH_BREAKDOWNS_TEST,
+    )
+
+    if is_national_breakdowns:
+        if request.country.lower() != "us":
+            raise HTTPException(
+                status_code=400,
+                detail="national-with-breakdowns is only supported for country='us'",
+            )
+
+        # Add test_mode flag to payload for orchestration to use
+        if data_value == NATIONAL_WITH_BREAKDOWNS_TEST:
+            payload["_test_mode"] = True
+            logger.info(
+                f"Routing {request.country}:{resolved_version} to {app_name} "
+                f"(national-with-breakdowns-test orchestration - 10 states)"
+            )
+        else:
+            logger.info(
+                f"Routing {request.country}:{resolved_version} to {app_name} "
+                f"(national-with-breakdowns orchestration - all states)"
+            )
+
+        func_name = "run_national_with_breakdowns"
+    else:
+        logger.info(f"Routing {request.country}:{resolved_version} to app {app_name}")
+        func_name = "run_simulation"
 
     # Get function reference from the target app
-    sim_func = modal.Function.from_name(app_name, "run_simulation")
+    sim_func = modal.Function.from_name(app_name, func_name)
 
     # Spawn the job (returns immediately)
-    payload = request.model_dump(exclude={"version"})
     call = sim_func.spawn(payload)
 
     return JobSubmitResponse(

projects/policyengine-api-simulation/src/modal/orchestration.py

@@ -0,0 +1,138 @@
+"""
+Orchestration logic for national-with-breakdowns simulations.
+
+This module handles spawning 52 parallel simulations (1 national + 51 states)
+and aggregating the results into a single response with congressional district breakdowns.
+"""
+
+import logfire
+from typing import Any, Callable
+
+from src.modal.utils.state_codes import STATE_CODES, TEST_STATE_CODES
+
+# Re-export for backwards compatibility
+__all__ = ["STATE_CODES", "TEST_STATE_CODES", "run_national_orchestration"]
+
+
+def run_national_orchestration(
+    params: dict,
+    run_simulation: Callable,
+    state_codes: list[str] | None = None,
+) -> dict:
+    """
+    Orchestrate parallel simulations and aggregate results.
+
+    Spawns:
+    - 1 national ECPS simulation (region="us")
+    - State-level simulations for each state in state_codes (or all 51 if not specified)
+
+    Each spawned job runs in its own container via run_simulation.
+
+    Partial failure handling:
+    - If ALL states fail, the entire request fails
+    - If SOME states fail, the request succeeds with null values for failed states
+
+    Args:
+        params: Base simulation parameters (reform, baseline, time_period, etc.)
+        run_simulation: The Modal function to spawn for each simulation
+        state_codes: Optional list of state codes to run. If None, runs all 51.
+
+    Returns:
+        Aggregated result with national metrics + all congressional district breakdowns
+    """
+    states_to_run = state_codes if state_codes is not None else STATE_CODES
+
+    # Prepare base params (remove the special data flag)
+    base_params = {k: v for k, v in params.items() if k != "data"}
+
+    # 1. Spawn national ECPS simulation
+    logfire.info("Spawning national ECPS simulation")
+    national_params = {
+        **base_params,
+        "region": "us",
+        # data=None lets policyengine use default ECPS dataset
+    }
+    national_call = run_simulation.spawn(national_params)
+
+    # 2. Spawn state simulations (each gets its own container)
+    logfire.info("Spawning state-level simulations", state_count=len(states_to_run))
+    state_calls: dict[str, Any] = {}
+    for state_code in states_to_run:
+        state_params = {
+            **base_params,
+            "region": f"state/{state_code.lower()}",
+            # data=None lets get_default_dataset resolve to states/{CODE}.h5
+        }
+        state_calls[state_code] = run_simulation.spawn(state_params)
+
+    logfire.info(
+        "All simulations spawned, waiting for results",
+        total_jobs=len(states_to_run) + 1,
+    )
+
+    # 3. Wait for national result first
+    logfire.info("Waiting for national ECPS result")
+    national_result = national_call.get()
+    logfire.info("National ECPS simulation complete")
+
+    # 4. Wait for all state results and extract district data
+    all_districts: list[dict] = []
+    failed_states: list[str] = []
+    successful_states: list[str] = []
+
+    for state_code in states_to_run:
+        logfire.info("Waiting for state result", state_code=state_code)
+        call = state_calls[state_code]
+
+        try:
+            state_result = call.get()
+
+            # Extract congressional_district_impact.districts from state result
+            district_impact = state_result.get("congressional_district_impact", {})
+            districts = district_impact.get("districts", [])
+            logfire.info(
+                "State result received",
+                state_code=state_code,
+                districts_extracted=len(districts),
+            )
+            all_districts.extend(districts)
+            successful_states.append(state_code)
+
+        except Exception as e:
+            logfire.warn(
+                "State simulation failed",
+                state_code=state_code,
+                error=str(e)[:200],
+            )
+            failed_states.append(state_code)
+            # Add null placeholder for each district in this state
+            # We don't know how many districts, so we skip adding placeholders
+            # The response will simply be missing these districts
+
+    logfire.info(
+        "State simulations complete",
+        successful_count=len(successful_states),
+        failed_count=len(failed_states),
+    )
+
+    # 5. Check if ALL states failed
+    if len(failed_states) == len(states_to_run):
+        raise RuntimeError(
+            f"All {len(states_to_run)} state simulations failed. "
+            f"Failed states: {failed_states}"
+        )
+
+    if failed_states:
+        logfire.warn("Some states failed", failed_states=failed_states)
+
+    logfire.info("Total districts collected", total_districts=len(all_districts))
+
+    # 6. Merge: national result + aggregated districts + metadata
+    final_result = national_result.copy()
+    final_result["congressional_district_impact"] = {
+        "districts": all_districts,
+        "failed_states": failed_states if failed_states else None,
+        "successful_states": successful_states,
+    }
+
+    return final_result

projects/policyengine-api-simulation/src/modal/utils/__init__.py

@@ -1,3 +1,7 @@
 """
 Utility functions for Modal deployment.
 """
+
+from src.modal.utils.state_codes import STATE_CODES, TEST_STATE_CODES
+
+__all__ = ["STATE_CODES", "TEST_STATE_CODES"]