Merge pull request #490 from PolicyEngine/rollback-budget-window-to-477-runtime-upgrades

Restore compact budget-window results

Author: anth-volk

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    changed:
+      - Bumped policyengine-core minimum version to 3.23.5 for pandas 3.0 compatibility

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

@@ -2,8 +2,6 @@
 
 import pytest
 
-from tests.fixtures.budget_window_outputs import make_single_year_macro_output
-
 
 class MockDict:
     """Mock for Modal.Dict to simulate version registry."""
@@ -36,12 +34,7 @@ class MockFunctionCall:
 
     def __init__(self, object_id: str = "mock-job-id-123"):
         self.object_id = object_id
-        self.result = make_single_year_macro_output(
-            tax_revenue_impact=1000000,
-            state_tax_revenue_impact=0,
-            benefit_spending_impact=0,
-            budgetary_impact=1000000,
-        )
+        self.result = {"budget": {"total": 1000000}}
         self.error = None
         self.running = False
         self.__class__.registry[object_id] = self

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

@@ -1,14 +1,14 @@
-"""Budget-window result validation and aggregation helpers."""
+"""Budget-window annual result extraction and aggregation helpers."""
 
 from __future__ import annotations
 
 from decimal import Decimal
 from typing import Any
 
 from src.modal.gateway.models import (
+    BudgetWindowAnnualImpact,
     BudgetWindowResult,
     BudgetWindowTotals,
-    SingleYearMacroOutput,
 )
 
 # The UK microsimulation has no state/province fiscal layer, so worker child
@@ -20,10 +20,7 @@
     "benefit_spending_impact",
     "budgetary_impact",
 )
-
-
-def _is_number(value: Any) -> bool:
-    return isinstance(value, int | float) and not isinstance(value, bool)
+OPTIONAL_BUDGET_KEYS = ("state_tax_revenue_impact",)
 
 
 def _as_decimal(value: float | int) -> Decimal:
@@ -35,60 +32,46 @@ def _as_decimal(value: float | int) -> Decimal:
     return Decimal(str(value))
 
 
-def validate_single_year_output(
+def extract_annual_impact(
     *,
     simulation_year: str,
     child_result: dict[str, Any],
-) -> SingleYearMacroOutput:
-    """Validate and normalize a child macro result.
-
-    UK worker results can omit ``state_tax_revenue_impact`` because there is
-    no state/province fiscal layer. The canonical output still includes that
-    field, defaulted to zero, so downstream clients receive one stable shape.
-    """
-
-    if not isinstance(child_result, dict):
-        raise ValueError(
-            "Malformed budget-window child result: expected object for "
-            f"{simulation_year}"
-        )
-
+) -> 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 _is_number(budget.get(key))
+        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}"
         )
 
-    normalized = dict(child_result)
-    normalized_budget = dict(budget)
-    if "state_tax_revenue_impact" not in normalized_budget:
-        normalized_budget["state_tax_revenue_impact"] = 0.0
-    elif not _is_number(normalized_budget["state_tax_revenue_impact"]):
-        raise ValueError(
-            "Malformed budget-window child result: missing numeric "
-            "budget.state_tax_revenue_impact"
-        )
-    normalized["budget"] = normalized_budget
-
-    try:
-        return SingleYearMacroOutput.model_validate(normalized)
-    except Exception as exc:
-        raise ValueError(
-            f"Malformed budget-window child result for {simulation_year}: {exc}"
-        ) from exc
+    tax_revenue_impact = budget["tax_revenue_impact"]
+    # UK worker results omit the state fiscal layer entirely; coerce to 0.0
+    # so the parent aggregator can still report federal/state splits with a
+    # uniform shape across countries.
+    state_tax_revenue_impact = budget.get("state_tax_revenue_impact")
+    if not isinstance(state_tax_revenue_impact, int | float):
+        state_tax_revenue_impact = 0.0
+
+    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_single_year_outputs(
-    *,
-    outputs_by_year: dict[str, SingleYearMacroOutput],
-    years: list[str],
+def sum_annual_impacts(
+    annual_impacts: list[BudgetWindowAnnualImpact],
 ) -> BudgetWindowTotals:
     """Sum per-year impacts using Decimal accumulators.
 
@@ -110,21 +93,18 @@ def sum_single_year_outputs(
         "budgetaryImpact": Decimal(0),
     }
 
-    for year in years:
-        output = outputs_by_year[year]
-        budget = output.model_dump(mode="json")["budget"]
-        tax_revenue_impact = budget["tax_revenue_impact"]
-        state_tax_revenue_impact = budget.get("state_tax_revenue_impact")
-
-        totals["taxRevenueImpact"] += _as_decimal(tax_revenue_impact)
+    for annual_impact in annual_impacts:
+        totals["taxRevenueImpact"] += _as_decimal(annual_impact.taxRevenueImpact)
         totals["federalTaxRevenueImpact"] += _as_decimal(
-            tax_revenue_impact - state_tax_revenue_impact
+            annual_impact.federalTaxRevenueImpact
+        )
+        totals["stateTaxRevenueImpact"] += _as_decimal(
+            annual_impact.stateTaxRevenueImpact
         )
-        totals["stateTaxRevenueImpact"] += _as_decimal(state_tax_revenue_impact)
         totals["benefitSpendingImpact"] += _as_decimal(
-            budget["benefit_spending_impact"]
+            annual_impact.benefitSpendingImpact
         )
-        totals["budgetaryImpact"] += _as_decimal(budget["budgetary_impact"])
+        totals["budgetaryImpact"] += _as_decimal(annual_impact.budgetaryImpact)
 
     return BudgetWindowTotals(**{key: float(value) for key, value in totals.items()})
 
@@ -133,25 +113,12 @@ def build_budget_window_result(
     *,
     start_year: str,
     window_size: int,
-    outputs_by_year: dict[str, SingleYearMacroOutput],
+    annual_impacts: list[BudgetWindowAnnualImpact],
 ) -> BudgetWindowResult:
-    years = [str(int(start_year) + offset) for offset in range(window_size)]
-    missing_years = [year for year in years if year not in outputs_by_year]
-    if missing_years:
-        raise ValueError(
-            "Cannot build budget-window result: missing outputs for "
-            + ", ".join(missing_years)
-        )
-
-    ordered_outputs = {year: outputs_by_year[year] for year in years}
     return BudgetWindowResult(
         startYear=start_year,
         endYear=str(int(start_year) + window_size - 1),
         windowSize=window_size,
-        years=years,
-        outputsByYear=ordered_outputs,
-        totals=sum_single_year_outputs(
-            outputs_by_year=ordered_outputs,
-            years=years,
-        ),
+        annualImpacts=annual_impacts,
+        totals=sum_annual_impacts(annual_impacts),
     )

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

@@ -14,7 +14,7 @@
 )
 from src.modal.budget_window_results import (
     build_budget_window_result,
-    validate_single_year_output,
+    extract_annual_impact,
 )
 from src.modal.budget_window_state import (
     build_batch_status_response,
@@ -167,7 +167,7 @@ def poll_running_children_once(self) -> bool:
                 return False
 
             try:
-                single_year_output = validate_single_year_output(
+                annual_impact = extract_annual_impact(
                     simulation_year=simulation_year,
                     child_result=child_result,
                 )
@@ -189,7 +189,7 @@ def poll_running_children_once(self) -> bool:
             mark_child_completed(
                 self.state,
                 year=simulation_year,
-                single_year_output=single_year_output,
+                annual_impact=annual_impact,
             )
             put_batch_job_state(self.state)
             progress_made = True
@@ -222,15 +222,15 @@ def fail_batch_for_child_error(
         put_batch_job_state(self.state)
 
     def complete_batch(self) -> dict[str, Any]:
-        outputs_by_year = {
-            simulation_year: self.state.partial_outputs_by_year[simulation_year]
+        annual_impacts = [
+            self.state.partial_annual_impacts[simulation_year]
             for simulation_year in self.state.years
-            if simulation_year in self.state.partial_outputs_by_year
-        }
+            if simulation_year in self.state.partial_annual_impacts
+        ]
         result = build_budget_window_result(
             start_year=self.state.start_year,
             window_size=self.state.window_size,
-            outputs_by_year=outputs_by_year,
+            annual_impacts=annual_impacts,
         )
         mark_batch_complete(self.state, result=result)
         put_batch_job_state(self.state)

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

@@ -9,12 +9,12 @@
 
 from src.modal.gateway.models import (
     BatchChildJobStatus,
+    BudgetWindowAnnualImpact,
     BudgetWindowBatchRequest,
     BudgetWindowBatchState,
     BudgetWindowBatchStatusResponse,
     BudgetWindowResult,
     PolicyEngineBundle,
-    SingleYearMacroOutput,
 )
 
 logger = logging.getLogger(__name__)
@@ -79,7 +79,7 @@ def create_initial_batch_state(
         completed_years=[],
         failed_years=[],
         child_jobs={},
-        partial_outputs_by_year={},
+        partial_annual_impacts={},
         result=None,
         error=None,
         created_at=now,
@@ -166,7 +166,7 @@ def mark_child_completed(
     state: BudgetWindowBatchState,
     *,
     year: str,
-    single_year_output: SingleYearMacroOutput,
+    annual_impact: BudgetWindowAnnualImpact,
 ) -> BudgetWindowBatchState:
     if year in state.running_years:
         state.running_years.remove(year)
@@ -178,7 +178,7 @@ def mark_child_completed(
         job_id=child.job_id,
         status="complete",
     )
-    state.partial_outputs_by_year[year] = single_year_output
+    state.partial_annual_impacts[year] = annual_impact
     return _touch(state)