Merge pull request #442 from PolicyEngine/feat/multi-year-sim-api

Add budget-window batch gateway contracts

Author: anth-volk

projects/policyengine-api-simulation/fixtures/gateway/test_endpoints.py

@@ -74,6 +74,7 @@ def from_name(cls, app_name: str, func_name: str):
 @pytest.fixture
 def mock_modal(monkeypatch):
     """Patch Modal calls in the gateway endpoints module."""
+    from src.modal import budget_window_state
     from src.modal.gateway import endpoints
 
     mock_func = MockFunction()
@@ -101,6 +102,7 @@ class MockModal:
         FunctionCall = MockFunctionCall
 
     monkeypatch.setattr(endpoints, "modal", MockModal)
+    monkeypatch.setattr(budget_window_state, "modal", MockModal)
 
     return {
         "func": mock_func,

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

@@ -0,0 +1,104 @@
+"""Helpers for budget-window batch job state."""
+
+from __future__ import annotations
+
+from datetime import UTC, datetime
+
+import modal
+
+from src.modal.gateway.models import (
+    BudgetWindowBatchRequest,
+    BudgetWindowBatchState,
+    BudgetWindowBatchStatusResponse,
+    PolicyEngineBundle,
+)
+
+BUDGET_WINDOW_JOB_DICT_NAME = "simulation-api-budget-window-jobs"
+
+
+def _budget_window_job_store():
+    return modal.Dict.from_name(BUDGET_WINDOW_JOB_DICT_NAME, create_if_missing=True)
+
+
+def _utc_now_iso() -> str:
+    return datetime.now(UTC).isoformat()
+
+
+def _build_years(start_year: str, window_size: int) -> list[str]:
+    base_year = int(start_year)
+    return [str(base_year + offset) for offset in range(window_size)]
+
+
+def create_initial_batch_state(
+    *,
+    batch_job_id: str,
+    request: BudgetWindowBatchRequest,
+    resolved_version: str,
+    resolved_app_name: str,
+    bundle: PolicyEngineBundle,
+) -> BudgetWindowBatchState:
+    years = _build_years(request.start_year, request.window_size)
+    now = _utc_now_iso()
+
+    return BudgetWindowBatchState(
+        batch_job_id=batch_job_id,
+        status="submitted",
+        country=request.country,
+        region=request.region,
+        version=resolved_version,
+        target=request.target,
+        resolved_app_name=resolved_app_name,
+        policyengine_bundle=bundle,
+        start_year=request.start_year,
+        window_size=request.window_size,
+        max_parallel=request.max_parallel,
+        request_payload=request.model_dump(exclude={"telemetry"}, mode="json"),
+        years=years,
+        queued_years=list(years),
+        running_years=[],
+        completed_years=[],
+        failed_years=[],
+        child_jobs={},
+        partial_annual_impacts={},
+        result=None,
+        error=None,
+        created_at=now,
+        updated_at=now,
+        run_id=request.telemetry.run_id if request.telemetry else None,
+    )
+
+
+def get_batch_job_state(batch_job_id: str) -> BudgetWindowBatchState | None:
+    payload = _budget_window_job_store().get(batch_job_id)
+    if payload is None:
+        return None
+    return BudgetWindowBatchState.model_validate(payload)
+
+
+def put_batch_job_state(state: BudgetWindowBatchState) -> None:
+    serialized = state.model_dump(mode="json")
+    _budget_window_job_store()[state.batch_job_id] = serialized
+
+
+def build_batch_status_response(
+    state: BudgetWindowBatchState,
+) -> BudgetWindowBatchStatusResponse:
+    total_years = len(state.years)
+    progress = (
+        0 if total_years == 0 else round(len(state.completed_years) / total_years * 100)
+    )
+
+    return BudgetWindowBatchStatusResponse(
+        status=state.status,
+        progress=progress,
+        completed_years=state.completed_years,
+        running_years=state.running_years,
+        queued_years=state.queued_years,
+        failed_years=state.failed_years,
+        child_jobs=state.child_jobs,
+        result=state.result,
+        error=state.error,
+        resolved_app_name=state.resolved_app_name,
+        policyengine_bundle=state.policyengine_bundle,
+        run_id=state.run_id,
+    )

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

@@ -10,6 +10,9 @@
 from fastapi.responses import JSONResponse
 
 from src.modal.gateway.models import (
+    BudgetWindowBatchRequest,
+    BudgetWindowBatchStatusResponse,
+    BudgetWindowBatchSubmitResponse,
     JobStatusResponse,
     JobSubmitResponse,
     PingRequest,
@@ -156,6 +159,25 @@ async def submit_simulation(request: SimulationRequest):
     )
 
 
+@router.post(
+    "/simulate/economy/budget-window",
+    response_model=BudgetWindowBatchSubmitResponse,
+    response_model_exclude_none=True,
+)
+async def submit_budget_window_batch(request: BudgetWindowBatchRequest):
+    """
+    Submit a budget-window batch job.
+
+    This contract-first endpoint is intentionally disabled until the
+    orchestration worker lands in the follow-up PR. That keeps the route
+    mergeable without falsely claiming that work has started.
+    """
+    raise HTTPException(
+        status_code=501,
+        detail="Budget-window batch orchestration is not implemented yet",
+    )
+
+
 @router.get(
     "/jobs/{job_id}",
     response_model=JobStatusResponse,
@@ -205,6 +227,24 @@ async def get_job_status(job_id: str):
         )
 
 
+@router.get(
+    "/budget-window-jobs/{batch_job_id}",
+    response_model=BudgetWindowBatchStatusResponse,
+    response_model_exclude_none=True,
+)
+async def get_budget_window_job_status(batch_job_id: str):
+    """
+    Poll for budget-window batch status.
+
+    This contract-first endpoint is intentionally disabled until the
+    orchestration worker lands in the follow-up PR.
+    """
+    raise HTTPException(
+        status_code=501,
+        detail="Budget-window batch orchestration is not implemented yet",
+    )
+
+
 @router.get("/versions")
 async def list_versions():
     """List all available versions for all countries."""

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

@@ -15,6 +15,9 @@
 from fastapi import FastAPI
 
 from src.modal.gateway.models import (
+    BudgetWindowBatchRequest,
+    BudgetWindowBatchStatusResponse,
+    BudgetWindowBatchSubmitResponse,
     JobStatusResponse,
     JobSubmitResponse,
     PingRequest,
@@ -48,6 +51,24 @@ async def submit_simulation(request: SimulationRequest) -> JobSubmitResponse:
         """
         raise NotImplementedError("Stub for OpenAPI generation")
 
+    @app.post(
+        "/simulate/economy/budget-window",
+        response_model=BudgetWindowBatchSubmitResponse,
+        responses={
+            200: {"description": "Budget-window batch submitted successfully"},
+            400: {"description": "Invalid request (unknown country/version/year)"},
+        },
+    )
+    async def submit_budget_window_batch(
+        request: BudgetWindowBatchRequest,
+    ) -> BudgetWindowBatchSubmitResponse:
+        """
+        Submit a budget-window batch job.
+
+        Returns immediately with a parent batch job ID for polling.
+        """
+        raise NotImplementedError("Stub for OpenAPI generation")
+
     @app.get(
         "/jobs/{job_id}",
         response_model=JobStatusResponse,
@@ -70,6 +91,27 @@ async def get_job_status(job_id: str) -> JobStatusResponse:
         """
         raise NotImplementedError("Stub for OpenAPI generation")
 
+    @app.get(
+        "/budget-window-jobs/{batch_job_id}",
+        response_model=BudgetWindowBatchStatusResponse,
+        responses={
+            200: {
+                "description": "Batch complete",
+                "model": BudgetWindowBatchStatusResponse,
+            },
+            202: {"description": "Batch submitted or running"},
+            404: {"description": "Batch job not found"},
+            500: {"description": "Batch failed"},
+        },
+    )
+    async def get_budget_window_job_status(
+        batch_job_id: str,
+    ) -> BudgetWindowBatchStatusResponse:
+        """
+        Poll for budget-window batch status.
+        """
+        raise NotImplementedError("Stub for OpenAPI generation")
+
     @app.get("/versions")
     async def list_versions() -> dict:
         """List all available versions for all countries."""