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() 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