Share one Stage 12 child wait deadline and cancel unfinished children

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

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
                 )

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(