From 26782a5a67f4506900949a5b6c9cd2055c76e498 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:13:50 +0400 Subject: [PATCH 1/2] Test that Stage 12's recalculation errors carry no output-plan error Three tests fail on main: an error from run() or save(), and a KeyboardInterrupt from run(), raised while Stage12CachedSimulation.ensure() recalculates an incomplete cache file carry the output-plan ValueError as __context__. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_cache_runtime.py | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py index 16b1b2468..ece37bc94 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py +++ b/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 From 79e6f891375a2f2171cd19f4b38d439e085ddbd8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 10 Nov 2026 18:15:30 +0400 Subject: [PATCH 2/2] Recalculate an incomplete Stage 12 cache file outside the except block Stage12CachedSimulation.ensure() now only notes in its except block that the loaded file failed the output plan, and runs and saves after the try statement, so a recalculation error or a preemption carries no chained output-plan error. Fixes #1035. Co-Authored-By: Claude Opus 5.5 --- changelog.d/1035.fixed.md | 1 + .../stage12_runtime/cached_simulation.py | 7 +++++++ 2 files changed, 8 insertions(+) create mode 100644 changelog.d/1035.fixed.md diff --git a/changelog.d/1035.fixed.md b/changelog.d/1035.fixed.md new file mode 100644 index 000000000..549f145a8 --- /dev/null +++ b/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. diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/cached_simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/cached_simulation.py index 9029edb1f..6e779eb52 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/cached_simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/cached_simulation.py @@ -141,6 +141,9 @@ 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): @@ -148,6 +151,10 @@ def ensure(self) -> None: "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()