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

Implement budget-window batch orchestration

Author: anth-volk

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

@@ -59,16 +59,29 @@ def __init__(self):
         self.last_payload = None
         self.last_from_name_call = None
         self.last_call = None
+        self.calls = []
 
-    def spawn(self, payload: dict) -> MockFunctionCall:
-        self.last_payload = payload
-        self.last_call = MockFunctionCall()
-        return self.last_call
+    def bind(self, app_name: str, func_name: str) -> "BoundMockFunction":
+        return BoundMockFunction(self, app_name, func_name)
 
-    @classmethod
-    def from_name(cls, app_name: str, func_name: str):
-        """Mock from_name that returns a MockFunction."""
-        raise NotImplementedError("Mock not configured")
+
+class BoundMockFunction:
+    """Function handle returned by Modal.Function.from_name."""
+
+    def __init__(self, recorder: MockFunction, app_name: str, func_name: str):
+        self.recorder = recorder
+        self.app_name = app_name
+        self.func_name = func_name
+
+    def spawn(self, payload: dict) -> MockFunctionCall:
+        self.recorder.last_payload = payload
+        is_batch = self.func_name == "run_budget_window_batch"
+        object_id = "mock-batch-job-id-123" if is_batch else "mock-job-id-123"
+        self.recorder.last_call = MockFunctionCall(object_id=object_id)
+        if is_batch:
+            self.recorder.last_call.running = True
+        self.recorder.calls.append((self.app_name, self.func_name, payload, object_id))
+        return self.recorder.last_call
 
 
 @pytest.fixture
@@ -94,7 +107,7 @@ class MockModalFunction:
         @staticmethod
         def from_name(app_name: str, func_name: str):
             mock_func.last_from_name_call = (app_name, func_name)
-            return mock_func
+            return mock_func.bind(app_name, func_name)
 
     class MockModal:
         Dict = MockModalDict

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

@@ -104,3 +104,32 @@ def run_simulation(params: dict) -> dict:
             return result
     finally:
         logfire.force_flush()
+
+
+@app.function(
+    image=simulation_image,
+    cpu=1.0,
+    memory=4096,
+    timeout=3600,
+    retries=0,
+    max_containers=100,
+    secrets=[gcp_secret, logfire_secret],
+)
+def run_budget_window_batch(params: dict) -> dict:
+    """Execute a multi-year budget-window batch orchestration."""
+    import logfire
+
+    from src.modal.budget_window_batch import run_budget_window_batch_impl
+
+    configure_logfire()
+
+    try:
+        with logfire.span(
+            "run_budget_window_batch",
+            input_params=params,
+        ) as span:
+            result = run_budget_window_batch_impl(params)
+            span.set_attribute("output_result", result)
+            return result
+    finally:
+        logfire.force_flush()

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

@@ -0,0 +1,19 @@
+"""Thin entrypoint for budget-window batch execution."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import modal
+
+from src.modal.budget_window_context import build_batch_context
+from src.modal.budget_window_scheduler import BudgetWindowBatchRunner
+
+
+def run_budget_window_batch_impl(params: dict[str, Any]) -> dict[str, Any]:
+    context = build_batch_context(
+        params,
+        batch_job_id=modal.current_function_call_id(),
+    )
+    runner = BudgetWindowBatchRunner(context)
+    return runner.run()

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

@@ -0,0 +1,100 @@
+"""Typed context and child request helpers for budget-window batches."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+from src.modal.gateway.models import BudgetWindowBatchRequest, PolicyEngineBundle
+
+BATCH_ONLY_FIELDS = {
+    "version",
+    "start_year",
+    "window_size",
+    "max_parallel",
+    "target",
+    "_metadata",
+    "_telemetry",
+}
+
+
+@dataclass(frozen=True)
+class BudgetWindowBatchContext:
+    """Resolved parent batch execution context."""
+
+    batch_job_id: str
+    request: BudgetWindowBatchRequest
+    resolved_version: str
+    resolved_app_name: str
+    bundle: PolicyEngineBundle
+    raw_params: dict[str, Any]
+
+
+@dataclass(frozen=True)
+class ChildSimulationRequest:
+    """Expanded single-year child simulation request."""
+
+    simulation_year: str
+    payload: dict[str, Any]
+
+
+@dataclass
+class ChildSimulationHandle:
+    """Tracked child job handle for a single simulation year."""
+
+    simulation_year: str
+    job_id: str
+    call: Any | None = None
+
+
+def build_batch_context(
+    params: dict[str, Any],
+    *,
+    batch_job_id: str,
+) -> BudgetWindowBatchContext:
+    request = BudgetWindowBatchRequest.model_validate(params)
+    metadata = params.get("_metadata")
+    if not isinstance(metadata, dict):
+        raise ValueError("Missing internal batch metadata")
+
+    resolved_app_name = metadata.get("resolved_app_name")
+    resolved_version = metadata.get("resolved_version")
+    bundle_payload = metadata.get("policyengine_bundle")
+
+    if not isinstance(resolved_app_name, str) or not resolved_app_name:
+        raise ValueError("Missing resolved_app_name in batch metadata")
+    if not isinstance(resolved_version, str) or not resolved_version:
+        raise ValueError("Missing resolved_version in batch metadata")
+    if not isinstance(bundle_payload, dict):
+        raise ValueError("Missing policyengine_bundle in batch metadata")
+
+    return BudgetWindowBatchContext(
+        batch_job_id=batch_job_id,
+        request=request,
+        resolved_version=resolved_version,
+        resolved_app_name=resolved_app_name,
+        bundle=PolicyEngineBundle.model_validate(bundle_payload),
+        raw_params=params,
+    )
+
+
+def build_child_simulation_request(
+    context: BudgetWindowBatchContext,
+    *,
+    simulation_year: str,
+) -> ChildSimulationRequest:
+    payload = {
+        key: value
+        for key, value in context.raw_params.items()
+        if key not in BATCH_ONLY_FIELDS
+    }
+    payload["time_period"] = simulation_year
+
+    telemetry = context.raw_params.get("_telemetry")
+    if isinstance(telemetry, dict):
+        payload["_telemetry"] = telemetry
+
+    return ChildSimulationRequest(
+        simulation_year=simulation_year,
+        payload=payload,
+    )

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

@@ -0,0 +1,87 @@
+"""Budget-window annual result extraction and aggregation helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from src.modal.gateway.models import (
+    BudgetWindowAnnualImpact,
+    BudgetWindowResult,
+    BudgetWindowTotals,
+)
+
+REQUIRED_BUDGET_KEYS = (
+    "tax_revenue_impact",
+    "state_tax_revenue_impact",
+    "benefit_spending_impact",
+    "budgetary_impact",
+)
+
+
+def extract_annual_impact(
+    *,
+    simulation_year: str,
+    child_result: dict[str, Any],
+) -> BudgetWindowAnnualImpact:
+    budget = child_result.get("budget", {})
+    if not isinstance(budget, dict):
+        raise ValueError("Malformed budget-window child result: missing budget object")
+
+    missing_keys = [
+        key
+        for key in REQUIRED_BUDGET_KEYS
+        if not isinstance(budget.get(key), int | float)
+    ]
+    if missing_keys:
+        missing = ", ".join(f"budget.{key}" for key in missing_keys)
+        raise ValueError(
+            f"Malformed budget-window child result: missing numeric {missing}"
+        )
+
+    state_tax_revenue_impact = budget["state_tax_revenue_impact"]
+    tax_revenue_impact = budget["tax_revenue_impact"]
+
+    return BudgetWindowAnnualImpact(
+        year=simulation_year,
+        taxRevenueImpact=tax_revenue_impact,
+        federalTaxRevenueImpact=tax_revenue_impact - state_tax_revenue_impact,
+        stateTaxRevenueImpact=state_tax_revenue_impact,
+        benefitSpendingImpact=budget["benefit_spending_impact"],
+        budgetaryImpact=budget["budgetary_impact"],
+    )
+
+
+def sum_annual_impacts(
+    annual_impacts: list[BudgetWindowAnnualImpact],
+) -> BudgetWindowTotals:
+    totals = {
+        "taxRevenueImpact": 0,
+        "federalTaxRevenueImpact": 0,
+        "stateTaxRevenueImpact": 0,
+        "benefitSpendingImpact": 0,
+        "budgetaryImpact": 0,
+    }
+
+    for annual_impact in annual_impacts:
+        totals["taxRevenueImpact"] += annual_impact.taxRevenueImpact
+        totals["federalTaxRevenueImpact"] += annual_impact.federalTaxRevenueImpact
+        totals["stateTaxRevenueImpact"] += annual_impact.stateTaxRevenueImpact
+        totals["benefitSpendingImpact"] += annual_impact.benefitSpendingImpact
+        totals["budgetaryImpact"] += annual_impact.budgetaryImpact
+
+    return BudgetWindowTotals(**totals)
+
+
+def build_budget_window_result(
+    *,
+    start_year: str,
+    window_size: int,
+    annual_impacts: list[BudgetWindowAnnualImpact],
+) -> BudgetWindowResult:
+    return BudgetWindowResult(
+        startYear=start_year,
+        endYear=str(int(start_year) + window_size - 1),
+        windowSize=window_size,
+        annualImpacts=annual_impacts,
+        totals=sum_annual_impacts(annual_impacts),
+    )