Merge pull request #399 from PolicyEngine/fix/deprecate-national-orchestrator

Remove national orchestrator job

Author: anth-volk

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

@@ -104,53 +104,3 @@ 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,10 +50,6 @@ 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):
     """
@@ -62,57 +58,18 @@ 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))
 
-    # 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"
+    logger.info(f"Routing {request.country}:{resolved_version} to app {app_name}")
 
     # Get function reference from the target app
-    sim_func = modal.Function.from_name(app_name, func_name)
+    sim_func = modal.Function.from_name(app_name, "run_simulation")
 
     # Spawn the job (returns immediately)
     call = sim_func.spawn(payload)

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

@@ -1,9 +1,5 @@
 """Test fixtures for the simulation API tests."""
 
-from tests.fixtures.orchestration import (
-    MockFunctionCall,
-    create_mock_run_simulation,
-)
 from tests.fixtures.endpoints import (
     MockDict,
     MockFunction,
@@ -12,8 +8,6 @@
 )
 
 __all__ = [
-    "MockFunctionCall",
-    "create_mock_run_simulation",
     "MockDict",
     "MockFunction",
     "MockModalFunctionCall",