From d94b18a1fccf43c1b35a6f6066e75e1b7a27e018 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:19:00 +0400 Subject: [PATCH 1/2] Share one Stage 12 child wait deadline and cancel unfinished children The coordinator waited up to 3,000 s for baseline and then up to a fresh 3,000 s for reform, so queued workers could push it past its own 4,500 s Modal timeout before it recorded the failure. Both waits now draw on one 3,000 s window, leaving 1,500 s for the 900 s production-result wait and 600 s of aggregation and failure persistence. When a child wait fails or times out, the coordinator now cancels every unfinished child, best effort, and records a cancelled sibling as failed, because Modal's InputCancellation is a BaseException the worker cannot record itself. Cancellation and state-update errors are logged and never replace the original child failure. This keeps the sequential baseline-then-reform wait; concurrent observation remains part of #712. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 1 + .../stage12_runtime/coordination.py | 81 +++++- .../stage12_runtime/dependencies.py | 2 + .../tests/test_stage12_modal_app.py | 27 ++ .../tests/test_stage12_runtime.py | 264 ++++++++++++++++++ 5 files changed, 373 insertions(+), 2 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 066e85eb1..59dd6602b 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -10,3 +10,4 @@ - Accepted a UK dataset without area codes only when the policyengine.py bundle routes constituency and local-authority regions to another dataset, and kept requiring area codes on every constituency and local-authority run - Raised a clear error, never an omission, when a UK constituency or local-authority breakdown is built over a dataset without area codes - Bound UK constituency and local-authority weight matrices to the certified data release of the selected dataset instead of the unversioned bucket copy, and stopped runs before they start when the matrix lacks the run year or household dimension + - Gave Stage 12 baseline and reform waits one shared deadline inside the coordinator timeout, and cancelled unfinished child simulations when a child fails or times out diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index cdb22d832..493c0153e 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -2,10 +2,12 @@ from __future__ import annotations -from collections.abc import Callable, Mapping +import logging +from collections.abc import Callable, Collection, Mapping, Sequence from contextlib import nullcontext from datetime import UTC, datetime from hashlib import sha256 +from time import monotonic from typing import Any from uuid import uuid4 @@ -58,6 +60,12 @@ ) from .simulation import descriptor_from_record, simulation_input_sha256 +logger = logging.getLogger(__name__) + +# One wait window shared by every child simulation, equal to the workers' +# own 3,000 s Modal timeout. The coordinator's 4,500 s Modal timeout then +# leaves 1,500 s: the 900 s production-result wait plus 600 s to aggregate, +# persist, and record any failure before Modal stops the coordinator. SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 @@ -187,6 +195,65 @@ def _claim_parent( return claimed, True +def _cancel_unfinished_children( + calls: Mapping[str, ChildCall], + *, + failed_role: str, + finished_roles: Collection[str], + simulations: Sequence[PlannedSimulationExecutionInput], + persistence: ComparisonStore, +) -> None: + """Stop children that can no longer contribute to a failed report. + + Best effort: a cancellation or state-update error must never replace the + child failure that the caller re-raises. + """ + for role, call in calls.items(): + if role in finished_roles: + continue + try: + call.cancel() + except Exception as cancel_error: # noqa: BLE001 + logger.warning( + "stage12_child_cancel_failed", + extra={ + "simulation_role": role, + "error_type": type(cancel_error).__name__, + }, + ) + continue + if role == failed_role: + # The caller records the outcome of the child it was waiting on. + continue + # A cancelled worker cannot record its own outcome, so record it here. + simulation = next(item for item in simulations if item.role.value == role) + try: + sibling = persistence.get_simulation(simulation.simulation_execution_id) + if sibling.status is ComparisonRunLifecycleStatus.RUNNING: + cancelled_at = datetime.now(UTC) + persistence.replace_simulation( + sibling.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_cancelled", + "error_summary": ( + f"Cancelled after the {failed_role} simulation failed" + ), + "updated_at": cancelled_at, + "completed_at": cancelled_at, + } + ) + ) + except Exception as persistence_error: # noqa: BLE001 + logger.warning( + "stage12_cancelled_child_state_update_failed", + extra={ + "simulation_role": role, + "error_type": type(persistence_error).__name__, + }, + ) + + def coordinate_report( payload: object, context_payload: object, @@ -392,6 +459,9 @@ def stage_scope( updated_at=datetime.now(UTC), ) # Every required call has been started before the coordinator waits. + # The calls share one deadline, so a slow baseline never gives reform + # a second full window beyond the coordinator's own timeout. + wait_deadline = monotonic() + SIMULATION_WAIT_TIMEOUT_SECONDS for role, call in calls.items(): try: wait_span = ( @@ -404,9 +474,16 @@ def stage_scope( ) with wait_span: descriptors[role] = SimulationArtifactDescriptor.model_validate( - call.get(timeout=SIMULATION_WAIT_TIMEOUT_SECONDS) + call.get(timeout=max(0.0, wait_deadline - monotonic())) ) except Exception as error: + _cancel_unfinished_children( + calls, + failed_role=role, + finished_roles=descriptors.keys(), + simulations=planned_simulations, + persistence=persistence, + ) simulation = next( item for item in simulations if item.role.value == role ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py index e5d1517be..efaf6e573 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py @@ -64,6 +64,8 @@ class ChildCall(Protocol): def get(self, *, timeout: float | None = None) -> object: ... + def cancel(self) -> None: ... + class ChildInvoker(Protocol): def spawn( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index ade681c12..d9242a545 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -207,6 +207,33 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( assert "asgi_app" not in vars(module) +def test_coordinator_timeout_leaves_time_to_record_after_the_child_window( + monkeypatch, +) -> None: + from policyengine_simulation_executor.stage12_runtime.comparison import ( + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, + ) + from policyengine_simulation_executor.stage12_runtime.coordination import ( + SIMULATION_WAIT_TIMEOUT_SECONDS, + ) + + module = _load(monkeypatch) + functions = {name: options for name, options in module.app.function_calls} + minimum_aggregation_and_persistence_seconds = 600 + + assert ( + SIMULATION_WAIT_TIMEOUT_SECONDS + == functions["run_single_simulation_us"]["timeout"] + == functions["run_single_simulation_uk"]["timeout"] + ) + assert ( + functions["coordinate_report"]["timeout"] + - SIMULATION_WAIT_TIMEOUT_SECONDS + - PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + >= minimum_aggregation_and_persistence_seconds + ) + + @pytest.mark.parametrize( "function_name", [ diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index f989c3a90..60c5c55ab 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -75,6 +75,9 @@ from policyengine_simulation_executor.stage12_runtime import ( coordinate_report as coordinate_report_impl, ) +from policyengine_simulation_executor.stage12_runtime import ( + coordination as coordination_module, +) from policyengine_simulation_executor.stage12_runtime.aggregation import ( validate_uk_local_authority_metadata, ) @@ -916,6 +919,9 @@ def __init__(self, object_id, future): def get(self, *, timeout=None): return self.future.result(timeout=timeout) + def cancel(self): + self.future.cancel() + class ImmediateCall: def __init__(self, object_id, result, observed_timeouts): @@ -1505,6 +1511,9 @@ def __init__(self, object_id): def get(self, *, timeout=None): raise TimeoutError("contains sensitive timeout details") + def cancel(self): + pass + class TimeoutInvoker: def __init__(self): @@ -1548,3 +1557,258 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: timed_out_child = store.children[BASELINE_ID] assert timed_out_child.status is ComparisonRunLifecycleStatus.FAILED assert timed_out_child.error_summary == "TimeoutError" + + +class SteppedClock: + def __init__(self, start_seconds=10_000.0): + self.current_seconds = start_seconds + + def monotonic(self): + return self.current_seconds + + def advance(self, elapsed_seconds): + self.current_seconds += elapsed_seconds + + +class ScriptedChildCall: + def __init__( + self, + object_id, + *, + result, + clock, + wait_error=None, + elapsed_seconds=0.0, + cancel_error=None, + ): + self.object_id = object_id + self.result = result + self.clock = clock + self.wait_error = wait_error + self.elapsed_seconds = elapsed_seconds + self.cancel_error = cancel_error + self.wait_timeouts = [] + self.cancel_count = 0 + + def get(self, *, timeout=None): + self.wait_timeouts.append(timeout) + self.clock.advance(self.elapsed_seconds) + if self.wait_error is not None: + raise self.wait_error + return self.result + + def cancel(self): + self.cancel_count += 1 + if self.cancel_error is not None: + raise self.cancel_error + + +class ScriptedInvoker: + def __init__( + self, + artifacts, + clock, + *, + wait_errors=None, + elapsed_seconds=None, + cancel_errors=None, + ): + self.artifacts = artifacts + self.clock = clock + self.wait_errors = wait_errors or {} + self.elapsed_seconds = elapsed_seconds or {} + self.cancel_errors = cancel_errors or {} + self.calls = {} + + def spawn(self, *, simulation, **_): + parsed = PlannedSimulationExecutionInput.model_validate(simulation) + role = parsed.role.value + descriptor = self.artifacts.add_simulation( + parsed, _frames(100.0 if role == "baseline" else 120.0) + ) + call = ScriptedChildCall( + f"call-{role}", + result=descriptor.model_dump(mode="json"), + clock=self.clock, + wait_error=self.wait_errors.get(role), + elapsed_seconds=self.elapsed_seconds.get(role, 0.0), + cancel_error=self.cancel_errors.get(role), + ) + self.calls[role] = call + return call + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") + + +@pytest.fixture +def stepped_clock(monkeypatch): + clock = SteppedClock() + monkeypatch.setattr(coordination_module, "monotonic", clock.monotonic) + return clock + + +def _coordinate_with(store, artifacts, invoker): + return coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: {"result": "complete"}, + ) + + +def test_children_share_one_wait_window(stepped_clock) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + elapsed_seconds={"baseline": 1_200.0}, + ) + + result = _coordinate_with(store, artifacts, invoker) + + assert result["artifact"]["uri"] == "gs://private/aggregate.json" + assert invoker.calls["baseline"].wait_timeouts == [ + coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + ] + assert invoker.calls["reform"].wait_timeouts == [ + coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS - 1_200.0 + ] + assert [call.cancel_count for call in invoker.calls.values()] == [0, 0] + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +def test_exhausted_wait_window_gives_the_next_child_no_extra_time( + stepped_clock, +) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + elapsed_seconds={ + "baseline": coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + 60.0 + }, + ) + + _coordinate_with(store, artifacts, invoker) + + assert invoker.calls["reform"].wait_timeouts == [0.0] + + +def test_baseline_failure_cancels_and_records_the_reform_child( + stepped_clock, +) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + wait_errors={"baseline": LookupError("baseline worker failed")}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(store, artifacts, invoker) + + reform_call = invoker.calls["reform"] + assert reform_call.cancel_count == 1 + assert reform_call.wait_timeouts == [] + assert artifacts.aggregate_writes == [] + failed_baseline = store.children[BASELINE_ID] + assert failed_baseline.status is ComparisonRunLifecycleStatus.FAILED + assert failed_baseline.error_code == "simulation_invocation_failed" + cancelled_reform = store.children[REFORM_ID] + assert cancelled_reform.status is ComparisonRunLifecycleStatus.FAILED + assert cancelled_reform.error_code == "simulation_cancelled" + assert cancelled_reform.error_summary == ( + "Cancelled after the baseline simulation failed" + ) + assert cancelled_reform.completed_at is not None + + +def test_baseline_timeout_cancels_both_unfinished_children(stepped_clock) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + wait_errors={"baseline": TimeoutError("baseline still running")}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(store, artifacts, invoker) + + # The timed-out baseline is still running on Modal, so it is stopped too. + assert invoker.calls["baseline"].cancel_count == 1 + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + assert store.children[REFORM_ID].error_code == "simulation_cancelled" + + +def test_reform_failure_leaves_the_finished_baseline_alone(stepped_clock) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + wait_errors={"reform": LookupError("reform worker failed")}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(store, artifacts, invoker) + + assert invoker.calls["baseline"].cancel_count == 0 + assert store.children[BASELINE_ID].error_code is None + assert store.children[REFORM_ID].error_code == "simulation_invocation_failed" + + +def test_cancellation_error_does_not_mask_the_child_failure(stepped_clock) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + wait_errors={"baseline": LookupError("baseline worker failed")}, + cancel_errors={"reform": ConnectionError("cancel RPC failed")}, + ) + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(store, artifacts, invoker) + + assert isinstance(error.value.__context__, LookupError) + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + # A sibling that may still be running is not recorded as cancelled. + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.RUNNING + + +def test_sibling_state_error_does_not_mask_the_child_failure(stepped_clock) -> None: + class SiblingUpdateFailingStore(FakeStore): + def replace_simulation(self, record): + if record.error_code == "simulation_cancelled": + raise ConnectionError("database unavailable") + return super().replace_simulation(record) + + store = SiblingUpdateFailingStore() + artifacts = FakeArtifacts() + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + wait_errors={"baseline": LookupError("baseline worker failed")}, + ) + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(store, artifacts, invoker) + + assert isinstance(error.value.__context__, LookupError) + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" From c5e7b793bc63a4cc7f5ff120c7095d222854dc0e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:38:50 +0400 Subject: [PATCH 2/2] Open the Stage 12 child wait window at coordinator entry Anchor the shared child-wait deadline when coordinate_report starts, so claiming, output planning, retained-artifact reads and dispatch shorten the children's waits rather than the 1,500 s left for the production-result wait, aggregation, persistence and failure recording. Add evaluation_id and the child's simulation_execution_id to the cancellation warnings, and cover finished siblings, restored calls and the ImmediateCall cancel double. The coordinator timeout test now runs the coordinator with simulated setup time against the declared Modal timeouts, so it lives beside the other coordinator tests. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 29 ++- .../tests/test_stage12_modal_app.py | 27 -- .../tests/test_stage12_runtime.py | 232 ++++++++++++++++-- 3 files changed, 236 insertions(+), 52 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 493c0153e..fb6137377 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -62,10 +62,13 @@ logger = logging.getLogger(__name__) -# One wait window shared by every child simulation, equal to the workers' -# own 3,000 s Modal timeout. The coordinator's 4,500 s Modal timeout then -# leaves 1,500 s: the 900 s production-result wait plus 600 s to aggregate, -# persist, and record any failure before Modal stops the coordinator. +# One window, opened when `coordinate_report` starts, covers coordinator setup +# (claiming, output planning, retained-artifact reads, and dispatch) and every +# child simulation wait. It equals the workers' own 3,000 s Modal timeout, so +# setup time shortens the children's waits instead of the time after them: the +# coordinator's 4,500 s Modal timeout always leaves 1,500 s after the window, +# the 900 s production-result wait plus 600 s to aggregate, persist, and record +# any failure before Modal stops the coordinator. SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 @@ -211,13 +214,19 @@ def _cancel_unfinished_children( for role, call in calls.items(): if role in finished_roles: continue + simulation = next(item for item in simulations if item.role.value == role) + child_identity = { + "evaluation_id": str(simulation.evaluation_id), + "simulation_execution_id": str(simulation.simulation_execution_id), + "simulation_role": role, + } try: call.cancel() except Exception as cancel_error: # noqa: BLE001 logger.warning( "stage12_child_cancel_failed", extra={ - "simulation_role": role, + **child_identity, "error_type": type(cancel_error).__name__, }, ) @@ -226,7 +235,6 @@ def _cancel_unfinished_children( # The caller records the outcome of the child it was waiting on. continue # A cancelled worker cannot record its own outcome, so record it here. - simulation = next(item for item in simulations if item.role.value == role) try: sibling = persistence.get_simulation(simulation.simulation_execution_id) if sibling.status is ComparisonRunLifecycleStatus.RUNNING: @@ -248,7 +256,7 @@ def _cancel_unfinished_children( logger.warning( "stage12_cancelled_child_state_update_failed", extra={ - "simulation_role": role, + **child_identity, "error_type": type(persistence_error).__name__, }, ) @@ -269,6 +277,8 @@ def coordinate_report( | None = None, runtime: ObservabilityRuntime | None = None, ) -> dict[str, Any]: + # Open the shared child-wait window before any setup work. + wait_deadline = monotonic() + SIMULATION_WAIT_TIMEOUT_SECONDS report = ReportExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) submitted_parent = ComparisonReportRecord.model_validate(parent_payload) @@ -459,9 +469,8 @@ def stage_scope( updated_at=datetime.now(UTC), ) # Every required call has been started before the coordinator waits. - # The calls share one deadline, so a slow baseline never gives reform - # a second full window beyond the coordinator's own timeout. - wait_deadline = monotonic() + SIMULATION_WAIT_TIMEOUT_SECONDS + # Each wait receives only what is left of the shared window, so a slow + # baseline never gives reform a second full window. for role, call in calls.items(): try: wait_span = ( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index d9242a545..ade681c12 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -207,33 +207,6 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( assert "asgi_app" not in vars(module) -def test_coordinator_timeout_leaves_time_to_record_after_the_child_window( - monkeypatch, -) -> None: - from policyengine_simulation_executor.stage12_runtime.comparison import ( - PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, - ) - from policyengine_simulation_executor.stage12_runtime.coordination import ( - SIMULATION_WAIT_TIMEOUT_SECONDS, - ) - - module = _load(monkeypatch) - functions = {name: options for name, options in module.app.function_calls} - minimum_aggregation_and_persistence_seconds = 600 - - assert ( - SIMULATION_WAIT_TIMEOUT_SECONDS - == functions["run_single_simulation_us"]["timeout"] - == functions["run_single_simulation_uk"]["timeout"] - ) - assert ( - functions["coordinate_report"]["timeout"] - - SIMULATION_WAIT_TIMEOUT_SECONDS - - PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS - >= minimum_aggregation_and_persistence_seconds - ) - - @pytest.mark.parametrize( "function_name", [ diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 60c5c55ab..91e710f4c 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -2,8 +2,11 @@ from __future__ import annotations +import importlib import io import json +import logging +import sys import time from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager, nullcontext @@ -14,6 +17,7 @@ import pandas as pd import pytest +from fixtures.fake_modal import install_fake_modal from policyengine_observability import ( DeploymentIdentity, LoggingConfig, @@ -81,6 +85,9 @@ from policyengine_simulation_executor.stage12_runtime.aggregation import ( validate_uk_local_authority_metadata, ) +from policyengine_simulation_executor.stage12_runtime.comparison import ( + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, +) from policyengine_simulation_executor.uk_local_authority_metadata import ( detect_uk_local_authority_metadata, ) @@ -928,11 +935,15 @@ def __init__(self, object_id, result, observed_timeouts): self.object_id = object_id self.result = result self.observed_timeouts = observed_timeouts + self.cancel_count = 0 def get(self, *, timeout=None): self.observed_timeouts.append(timeout) return self.result + def cancel(self): + self.cancel_count += 1 + class ConcurrentInvoker: def __init__( @@ -1580,6 +1591,7 @@ def __init__( wait_error=None, elapsed_seconds=0.0, cancel_error=None, + during_wait=None, ): self.object_id = object_id self.result = result @@ -1587,12 +1599,15 @@ def __init__( self.wait_error = wait_error self.elapsed_seconds = elapsed_seconds self.cancel_error = cancel_error + self.during_wait = during_wait self.wait_timeouts = [] self.cancel_count = 0 def get(self, *, timeout=None): self.wait_timeouts.append(timeout) self.clock.advance(self.elapsed_seconds) + if self.during_wait is not None: + self.during_wait() if self.wait_error is not None: raise self.wait_error return self.result @@ -1612,19 +1627,31 @@ def __init__( wait_errors=None, elapsed_seconds=None, cancel_errors=None, + during_wait=None, ): self.artifacts = artifacts self.clock = clock self.wait_errors = wait_errors or {} self.elapsed_seconds = elapsed_seconds or {} self.cancel_errors = cancel_errors or {} + self.during_wait = during_wait or {} self.calls = {} + self.restored_roles = [] def spawn(self, *, simulation, **_): - parsed = PlannedSimulationExecutionInput.model_validate(simulation) - role = parsed.role.value + return self._scripted_call( + PlannedSimulationExecutionInput.model_validate(simulation) + ) + + def restore(self, invocation_id): + role = invocation_id.removeprefix("call-") + self.restored_roles.append(role) + return self._scripted_call(_planned_simulation(SimulationRole(role))) + + def _scripted_call(self, simulation): + role = simulation.role.value descriptor = self.artifacts.add_simulation( - parsed, _frames(100.0 if role == "baseline" else 120.0) + simulation, _frames(100.0 if role == "baseline" else 120.0) ) call = ScriptedChildCall( f"call-{role}", @@ -1633,13 +1660,11 @@ def spawn(self, *, simulation, **_): wait_error=self.wait_errors.get(role), elapsed_seconds=self.elapsed_seconds.get(role, 0.0), cancel_error=self.cancel_errors.get(role), + during_wait=self.during_wait.get(role), ) self.calls[role] = call return call - def restore(self, invocation_id): - raise AssertionError(f"unexpected restored invocation {invocation_id}") - @pytest.fixture def stepped_clock(monkeypatch): @@ -1648,7 +1673,7 @@ def stepped_clock(monkeypatch): return clock -def _coordinate_with(store, artifacts, invoker): +def _coordinate_with(store, artifacts, invoker, **coordinator_options): return coordinate_report( _report().model_dump(mode="json"), _context().model_dump(mode="json"), @@ -1659,9 +1684,27 @@ def _coordinate_with(store, artifacts, invoker): artifacts=artifacts, invoker=invoker, aggregator=lambda **_: {"result": "complete"}, + **coordinator_options, ) +def _logged_warning(caplog, event_name): + (warning,) = [ + record + for record in caplog.records + if record.levelno == logging.WARNING and record.getMessage() == event_name + ] + return warning + + +def _assert_warning_identifies_the_reform_child(warning, error_type): + assert warning.evaluation_id == str(EVALUATION_ID) + assert warning.simulation_execution_id == str(REFORM_ID) + assert warning.simulation_role == "reform" + assert warning.error_type == error_type + assert warning.exc_info is None + + def test_children_share_one_wait_window(stepped_clock) -> None: store = FakeStore() artifacts = FakeArtifacts() @@ -1702,6 +1745,64 @@ def test_exhausted_wait_window_gives_the_next_child_no_extra_time( assert invoker.calls["reform"].wait_timeouts == [0.0] +def _declared_modal_timeouts(monkeypatch): + install_fake_modal(monkeypatch) + sys.modules.pop("src.modal.v2_app", None) + modal_app = importlib.import_module("src.modal.v2_app") + return { + name: options.get("timeout") for name, options in modal_app.app.function_calls + } + + +def test_coordinator_timeout_leaves_time_to_record_after_the_child_window( + monkeypatch, + stepped_clock, +) -> None: + declared_timeouts = _declared_modal_timeouts(monkeypatch) + child_window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + setup_seconds = 400.0 + minimum_aggregation_and_persistence_seconds = 600 + coordinator_started_at = stepped_clock.monotonic() + + def slow_output_planning(_report): + # Loading the country model is coordinator setup, not child time. + stepped_clock.advance(setup_seconds) + return _output_plan() + + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + # The baseline uses all the time it is given; reform is already done. + elapsed_seconds={"baseline": child_window_seconds - setup_seconds}, + ) + + _coordinate_with( + store, artifacts, invoker, output_plan_resolver=slow_output_planning + ) + + # Setup shortens the children's waits rather than the time after them. + assert invoker.calls["baseline"].wait_timeouts == [ + child_window_seconds - setup_seconds + ] + assert invoker.calls["reform"].wait_timeouts == [0.0] + # Nothing after the waits advances the clock, so this is when they ended. + child_waits_ended_after = stepped_clock.monotonic() - coordinator_started_at + assert child_waits_ended_after == child_window_seconds + assert ( + declared_timeouts["coordinate_report"] + - child_waits_ended_after + - PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + >= minimum_aggregation_and_persistence_seconds + ) + assert ( + child_window_seconds + == declared_timeouts["run_single_simulation_us"] + == declared_timeouts["run_single_simulation_uk"] + ) + + def test_baseline_failure_cancels_and_records_the_reform_child( stepped_clock, ) -> None: @@ -1768,7 +1869,10 @@ def test_reform_failure_leaves_the_finished_baseline_alone(stepped_clock) -> Non assert store.children[REFORM_ID].error_code == "simulation_invocation_failed" -def test_cancellation_error_does_not_mask_the_child_failure(stepped_clock) -> None: +def test_cancellation_error_does_not_mask_the_child_failure( + stepped_clock, + caplog, +) -> None: store = FakeStore() artifacts = FakeArtifacts() invoker = ScriptedInvoker( @@ -1778,18 +1882,29 @@ def test_cancellation_error_does_not_mask_the_child_failure(stepped_clock) -> No cancel_errors={"reform": ConnectionError("cancel RPC failed")}, ) - with pytest.raises( - RuntimeError, match="Stage 12 report coordination failed" - ) as error: + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error, + ): _coordinate_with(store, artifacts, invoker) assert isinstance(error.value.__context__, LookupError) assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" # A sibling that may still be running is not recorded as cancelled. assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.RUNNING + _assert_warning_identifies_the_reform_child( + _logged_warning(caplog, "stage12_child_cancel_failed"), + "ConnectionError", + ) + assert "cancel RPC failed" not in caplog.text -def test_sibling_state_error_does_not_mask_the_child_failure(stepped_clock) -> None: +def test_sibling_state_error_does_not_mask_the_child_failure( + stepped_clock, + caplog, +) -> None: class SiblingUpdateFailingStore(FakeStore): def replace_simulation(self, record): if record.error_code == "simulation_cancelled": @@ -1804,11 +1919,98 @@ def replace_simulation(self, record): wait_errors={"baseline": LookupError("baseline worker failed")}, ) - with pytest.raises( - RuntimeError, match="Stage 12 report coordination failed" - ) as error: + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error, + ): _coordinate_with(store, artifacts, invoker) assert isinstance(error.value.__context__, LookupError) assert invoker.calls["reform"].cancel_count == 1 assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + _assert_warning_identifies_the_reform_child( + _logged_warning(caplog, "stage12_cancelled_child_state_update_failed"), + "ConnectionError", + ) + assert "database unavailable" not in caplog.text + + +@pytest.mark.parametrize( + ("worker_status", "worker_error_code"), + [ + (ComparisonRunLifecycleStatus.SUCCEEDED, None), + (ComparisonRunLifecycleStatus.FAILED, "simulation_execution_failed"), + ], +) +def test_cancellation_leaves_a_sibling_that_already_finished_unchanged( + stepped_clock, + worker_status, + worker_error_code, +) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + worker_records = [] + + def reform_worker_records_its_outcome(): + # The reform worker finishes while the coordinator waits on baseline. + finished_at = NOW + timedelta(minutes=5) + worker_records.append( + store.replace_simulation( + store.get_simulation(REFORM_ID).model_copy( + update={ + "status": worker_status, + "error_code": worker_error_code, + "updated_at": finished_at, + "completed_at": finished_at, + } + ) + ) + ) + + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + wait_errors={"baseline": LookupError("baseline worker failed")}, + during_wait={"baseline": reform_worker_records_its_outcome}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(store, artifacts, invoker) + + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[REFORM_ID] == worker_records[0] + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + + +def test_baseline_failure_cancels_a_reform_call_restored_on_retry( + stepped_clock, +) -> None: + report = _report() + store = FakeStore() + artifacts = FakeArtifacts() + # An earlier coordinator attempt started the reform child before stopping. + store.children[REFORM_ID] = _child(report.reform).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": "call-reform", + "started_at": NOW, + } + ) + invoker = ScriptedInvoker( + artifacts, + stepped_clock, + wait_errors={"baseline": LookupError("baseline worker failed")}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(store, artifacts, invoker) + + assert invoker.restored_roles == ["reform"] + assert invoker.calls["reform"].cancel_count == 1 + assert invoker.calls["reform"].wait_timeouts == [] + cancelled_reform = store.children[REFORM_ID] + assert cancelled_reform.status is ComparisonRunLifecycleStatus.FAILED + assert cancelled_reform.error_code == "simulation_cancelled" + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed"