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..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 @@ -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,15 @@ ) from .simulation import descriptor_from_record, simulation_input_sha256 +logger = logging.getLogger(__name__) + +# 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 @@ -187,6 +198,70 @@ 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 + 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={ + **child_identity, + "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. + 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={ + **child_identity, + "error_type": type(persistence_error).__name__, + }, + ) + + def coordinate_report( payload: object, context_payload: object, @@ -202,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) @@ -392,6 +469,8 @@ def stage_scope( updated_at=datetime.now(UTC), ) # Every required call has been started before the coordinator waits. + # 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 = ( @@ -404,9 +483,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_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index f989c3a90..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, @@ -75,9 +79,15 @@ 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, ) +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, ) @@ -916,17 +926,24 @@ 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): 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__( @@ -1505,6 +1522,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 +1568,449 @@ 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, + during_wait=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.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 + + 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, + 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, **_): + 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( + simulation, _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), + during_wait=self.during_wait.get(role), + ) + self.calls[role] = call + return call + + +@pytest.fixture +def stepped_clock(monkeypatch): + clock = SteppedClock() + monkeypatch.setattr(coordination_module, "monotonic", clock.monotonic) + return clock + + +def _coordinate_with(store, artifacts, invoker, **coordinator_options): + 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"}, + **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() + 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 _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: + 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, + caplog, +) -> 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 ( + 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, + caplog, +) -> 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 ( + 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"