Recalculate an incomplete Stage 12 cache file outside the output-plan check's except block

changelog.d/1035.fixed.md

@@ -0,0 +1 @@
+An error raised while Stage 12 recalculates an incomplete cache file, including a Modal preemption, is no longer chained to the output-plan check that found the file incomplete, so the telemetry error report groups it by where it was raised.

projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/cached_simulation.py

@@ -141,13 +141,20 @@ def ensure(self) -> None:
             raise RuntimeError(msg)
         data = getattr(getattr(self, "output_dataset", None), "data", None)
         entity_data = getattr(data, "entity_data", None)
+        # Recalculate after leaving the except block. An error raised inside
+        # it would carry the output-plan error as its context, and the
+        # telemetry error report groups an error by its first traceback.
         try:
             validate_output_frames(entity_data or {}, plan)
         except (TypeError, ValueError):
             logger.warning(
                 "Stage 12 cache %s is incomplete; recalculating",
                 self.id,
             )
+            output_is_incomplete = True
+        else:
+            output_is_incomplete = False
+        if output_is_incomplete:
             self._record_outcome(CACHE_INCOMPLETE)
             self.run()
             self.save()

projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py

@@ -380,3 +380,73 @@ def load_incomplete_file_and_cache_it(simulation: Simulation) -> None:
         assert _cache.get(model.storage_id) is None
     finally:
         _cache.clear()
+
+
+class _RecalculationFailedError(Exception):
+    """Stands in for an error raised while Stage 12 recalculates a cache file."""
+
+
+def _load_incomplete_file(simulation: Simulation) -> None:
+    simulation.output_dataset = _output(complete=False)
+
+
+@pytest.mark.parametrize("failing_step", ["run", "save"])
+def test_recalculation_error_carries_no_output_plan_error(
+    monkeypatch, failing_step: str
+) -> None:
+    """An error from recalculating an incomplete file isn't chained to the check.
+
+    The telemetry error report groups an error by the first traceback in its
+    stack, so a chained output-plan error would group every recalculation
+    failure under the check that found the file incomplete.
+    """
+    recalculation_failure = _RecalculationFailedError(f"{failing_step} failed")
+
+    def raise_recalculation_failure(simulation: Simulation) -> None:
+        raise recalculation_failure
+
+    def recalculate_complete_output(simulation: Simulation) -> None:
+        simulation.output_dataset = _output()
+
+    model = _cached_model()
+    monkeypatch.setattr(Simulation, "ensure", _load_incomplete_file)
+    monkeypatch.setattr(
+        Simulation,
+        "run",
+        raise_recalculation_failure
+        if failing_step == "run"
+        else recalculate_complete_output,
+    )
+    monkeypatch.setattr(
+        Simulation,
+        "save",
+        raise_recalculation_failure if failing_step == "save" else lambda self: None,
+    )
+
+    with pytest.raises(_RecalculationFailedError) as raised:
+        model.ensure()
+
+    assert raised.value is recalculation_failure
+    assert raised.value.__context__ is None
+    assert raised.value.__cause__ is None
+    assert model.stage12_cache_outcome == CACHE_INCOMPLETE
+
+
+def test_preemption_during_recalculation_carries_no_output_plan_error(
+    monkeypatch,
+) -> None:
+    """A Modal preemption during the recalculation isn't chained to the check."""
+
+    def preempted_run(simulation: Simulation) -> None:
+        raise KeyboardInterrupt
+
+    model = _cached_model()
+    monkeypatch.setattr(Simulation, "ensure", _load_incomplete_file)
+    monkeypatch.setattr(Simulation, "run", preempted_run)
+    monkeypatch.setattr(Simulation, "save", lambda self: None)
+
+    with pytest.raises(KeyboardInterrupt) as raised:
+        model.ensure()
+
+    assert raised.value.__context__ is None
+    assert model.stage12_cache_outcome == CACHE_INCOMPLETE