Harden budget-window batch failure handling

Author: anth-volk

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

@@ -29,6 +29,12 @@
 )
 
 POLL_INTERVAL_SECONDS = 0.1
+REQUIRED_BUDGET_KEYS = (
+    "tax_revenue_impact",
+    "state_tax_revenue_impact",
+    "benefit_spending_impact",
+    "budgetary_impact",
+)
 
 
 def _extract_request_and_metadata(
@@ -87,6 +93,20 @@ def extract_budget_window_annual_impact(
     *, year: str, impact_data: dict[str, Any]
 ) -> BudgetWindowAnnualImpact:
     budget = impact_data.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.get("state_tax_revenue_impact", 0)
     tax_revenue_impact = budget.get("tax_revenue_impact", 0)
 
@@ -190,10 +210,18 @@ def run_budget_window_batch_impl(params: dict[str, Any]) -> dict[str, Any]:
                 put_batch_job_state(state)
                 return _failure_response(state)
 
-            annual_impact = extract_budget_window_annual_impact(
-                year=year,
-                impact_data=result,
-            )
+            try:
+                annual_impact = extract_budget_window_annual_impact(
+                    year=year,
+                    impact_data=result,
+                )
+            except Exception as exc:
+                error = str(exc)
+                mark_child_failed(state, year=year, error=error)
+                mark_batch_failed(state, error=error)
+                put_batch_job_state(state)
+                return _failure_response(state)
+
             mark_child_completed(state, year=year, annual_impact=annual_impact)
             put_batch_job_state(state)
             progress_made = True

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

@@ -187,8 +187,19 @@ def mark_batch_failed(
     *,
     error: str,
 ) -> BudgetWindowBatchState:
+    for year in list(state.running_years):
+        child = state.child_jobs.get(year)
+        if child is None:
+            continue
+        state.child_jobs[year] = BatchChildJobStatus(
+            job_id=child.job_id,
+            status="cancelled",
+            error=error,
+        )
+
     state.status = "failed"
     state.error = error
+    state.running_years = []
     return _touch(state)
 
 

projects/policyengine-api-simulation/tests/gateway/test_budget_window_state.py

@@ -6,6 +6,7 @@
     get_batch_job_state,
     mark_batch_complete,
     mark_batch_running,
+    mark_batch_failed,
     mark_child_completed,
     mark_child_failed,
     mark_child_started,
@@ -198,3 +199,33 @@ def test_state_transition_helpers_track_failed_child():
     assert state.status == "running"
     assert state.failed_years == ["2027"]
     assert state.child_jobs["2027"].status == "failed"
+
+
+def test_mark_batch_failed_cancels_any_remaining_running_children():
+    request = BudgetWindowBatchRequest(
+        country="us",
+        region="us",
+        start_year="2026",
+        window_size=2,
+    )
+
+    state = create_initial_batch_state(
+        batch_job_id="fc-parent-123",
+        request=request,
+        resolved_version="1.500.0",
+        resolved_app_name="policyengine-simulation-us1-500-0-uk2-66-0",
+        bundle=PolicyEngineBundle(model_version="1.500.0"),
+    )
+
+    mark_batch_running(state)
+    mark_child_started(state, year="2026", child_job_id="child-2026")
+    mark_child_started(state, year="2027", child_job_id="child-2027")
+    mark_child_failed(state, year="2026", error="boom")
+    mark_batch_failed(state, error="boom")
+
+    assert state.status == "failed"
+    assert state.running_years == []
+    assert state.failed_years == ["2026"]
+    assert state.child_jobs["2026"].status == "failed"
+    assert state.child_jobs["2027"].status == "cancelled"
+    assert state.child_jobs["2027"].error == "boom"

projects/policyengine-api-simulation/tests/test_budget_window_batch.py

@@ -226,6 +226,23 @@ def test_extract_budget_window_annual_impact_matches_v1_shape():
     )
 
 
+def test_extract_budget_window_annual_impact_rejects_malformed_child_result():
+    with pytest.raises(
+        ValueError,
+        match="Malformed budget-window child result: missing numeric budget.tax_revenue_impact",
+    ):
+        extract_budget_window_annual_impact(
+            year="2026",
+            impact_data={
+                "budget": {
+                    "state_tax_revenue_impact": 40,
+                    "benefit_spending_impact": 20,
+                    "budgetary_impact": 80,
+                }
+            },
+        )
+
+
 def test_build_budget_window_output_sums_totals():
     annual_impacts = [
         BudgetWindowAnnualImpact(
@@ -355,6 +372,51 @@ def test_run_budget_window_batch_impl_marks_failure(mock_batch_modal):
     assert result["status"] == "failed"
     assert result["failed_years"] == ["2026"]
     assert result["error"] == "child failed"
+    assert result["running_years"] == []
+    assert result["child_jobs"]["2027"]["status"] == "cancelled"
+    assert result["child_jobs"]["2027"]["error"] == "child failed"
     assert state is not None
     assert state.status == "failed"
     assert state.error == "child failed"
+    assert state.running_years == []
+    assert state.child_jobs["2027"].status == "cancelled"
+
+
+def test_run_budget_window_batch_impl_fails_on_malformed_child_result(
+    mock_batch_modal,
+):
+    request, payload = _build_parent_payload(window_size=1)
+    _seed_parent_batch(request, mock_batch_modal["parent_call_id"])
+
+    tracker = SpawnTracker()
+    run_simulation = MockRunSimulationFunction(
+        tracker=tracker,
+        results_by_year={
+            "2026": [
+                {
+                    "budget": {
+                        "state_tax_revenue_impact": 3,
+                        "benefit_spending_impact": 6,
+                        "budgetary_impact": 17,
+                    }
+                }
+            ],
+        },
+        call_registry=mock_batch_modal["call_registry"],
+    )
+    mock_batch_modal["functions"][
+        ("policyengine-simulation-us1-500-0-uk2-66-0", "run_simulation")
+    ] = run_simulation
+
+    result = run_budget_window_batch_impl(payload)
+    state = state_module.get_batch_job_state(mock_batch_modal["parent_call_id"])
+
+    assert result["status"] == "failed"
+    assert result["failed_years"] == ["2026"]
+    assert (
+        result["error"]
+        == "Malformed budget-window child result: missing numeric budget.tax_revenue_impact"
+    )
+    assert state is not None
+    assert state.status == "failed"
+    assert state.child_jobs["2026"].status == "failed"