diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 066e85eb1..7c708f43c 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 + - Released each Stage 12 database session after its operation instead of holding pooled sessions through calculations, retried opening each Stage 12 coordinator and worker database connection that the pooler refused without ever retrying a statement or commit, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot hide the original error diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py index 2944b84b8..e1344e788 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -1,6 +1,7 @@ """Direct DML for temporary Stage 12 comparison state.""" from policyengine_stage12_persistence.store import ( + ConnectionOpener, Stage12PersistenceStore, create_stage12_engine, normalized_database_url, @@ -11,6 +12,7 @@ ) __all__ = [ + "ConnectionOpener", "Stage12PersistenceStore", "Stage12RuntimeAccessError", "create_stage12_engine", diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py index b8f36421e..f463d0a32 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py @@ -2,9 +2,11 @@ from __future__ import annotations +from collections.abc import Callable, Iterator +from contextlib import contextmanager from datetime import datetime from enum import StrEnum -from typing import Any +from typing import Any, Protocol from uuid import UUID from policyengine_simulation_contract.stage12_execution import ( @@ -17,7 +19,8 @@ ) from sqlalchemy import Engine, create_engine, select, update from sqlalchemy.dialects.postgresql import insert -from sqlalchemy.engine import URL, make_url +from sqlalchemy.engine import URL, Connection, make_url +from sqlalchemy.pool import NullPool from policyengine_stage12_persistence.tables import ( REPORT_IDENTITY_CONSTRAINT, @@ -115,14 +118,41 @@ def create_stage12_engine(database_url: str) -> Engine: return create_engine( normalized_database_url(database_url), - pool_pre_ping=True, - pool_recycle=300, - pool_size=5, - max_overflow=5, + # Coordinators and workers spend most of their lifetime calculating + # or waiting. Release each database session after its short operation + # instead of reserving a pooler slot throughout that idle interval. + poolclass=NullPool, connect_args={"connect_timeout": 10}, ) +class ConnectionOpener(Protocol): + """Open the database connection that one store operation uses. + + ``connect`` only opens a connection; it sends no application statement. + The operation's statements, and the COMMIT or ROLLBACK that ends it, run + after the opener returns. An opener may therefore retry ``connect`` + without risking a duplicate write, and never sees a later failure. + """ + + def __call__( + self, + connect: Callable[[], Connection], + *, + operation_name: str, + ) -> Connection: ... + + +def open_connection_once( + connect: Callable[[], Connection], + *, + operation_name: str, +) -> Connection: + """Open a connection once, so a refused connection fails the operation.""" + + return connect() + + def _database_value(value: Any) -> Any: if isinstance(value, StrEnum): return value.value @@ -156,12 +186,31 @@ def __init__( database_url: str | None = None, *, engine: Engine | None = None, + open_connection: ConnectionOpener = open_connection_once, ) -> None: if engine is None: if database_url is None: raise ValueError("Stage 12 database URL is required") engine = create_stage12_engine(database_url) self.engine = engine + self._open_connection = open_connection + + @contextmanager + def _connection(self, operation_name: str) -> Iterator[Connection]: + # Opening the connection is the only step the opener sees. Every + # statement, COMMIT and ROLLBACK runs after the opener returns, so + # an opener that retries can never repeat work the database received. + connection = self._open_connection( + self.engine.connect, + operation_name=operation_name, + ) + with connection: + yield connection + + @contextmanager + def _transaction(self, operation_name: str) -> Iterator[Connection]: + with self._connection(operation_name) as connection, connection.begin(): + yield connection def create_or_resolve_report( self, @@ -174,7 +223,7 @@ def create_or_resolve_report( .on_conflict_do_nothing(constraint=REPORT_IDENTITY_CONSTRAINT) .returning(*comparison_reports.c) ) - with self.engine.begin() as connection: + with self._transaction("create_or_resolve_report") as connection: created = connection.execute(statement).mappings().one_or_none() if created is not None: return ComparisonReportPersistenceResult( @@ -203,7 +252,7 @@ def create_or_resolve_report( return ComparisonReportPersistenceResult(record=resolved, created=False) def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: - with self.engine.connect() as connection: + with self._connection("get_report") as connection: row = ( connection.execute( select(comparison_reports).where( @@ -229,7 +278,7 @@ def list_simulations( comparison_simulations.c.simulation_execution_id, ) ) - with self.engine.connect() as connection: + with self._connection("list_simulations") as connection: rows = connection.execute(statement).mappings().all() return tuple(_simulation(row) for row in rows) @@ -238,7 +287,7 @@ def replace_report( record: ComparisonReportRecord, ) -> ComparisonReportRecord: values = _record_values(record) - with self.engine.begin() as connection: + with self._transaction("replace_report") as connection: row = ( connection.execute( select(comparison_reports) @@ -280,7 +329,7 @@ def replace_report_result_comparison( record: ComparisonReportRecord, ) -> ComparisonReportRecord: values = _record_values(record) - with self.engine.begin() as connection: + with self._transaction("replace_report_result_comparison") as connection: row = ( connection.execute( select(comparison_reports) @@ -331,7 +380,7 @@ def create_or_resolve_simulation( .on_conflict_do_nothing(constraint=SIMULATION_IDENTITY_CONSTRAINT) .returning(*comparison_simulations.c) ) - with self.engine.begin() as connection: + with self._transaction("create_or_resolve_simulation") as connection: parent_row = ( connection.execute( select(comparison_reports).where( @@ -380,7 +429,7 @@ def get_simulation( self, simulation_execution_id: UUID, ) -> ComparisonSimulationRecord: - with self.engine.connect() as connection: + with self._connection("get_simulation") as connection: row = ( connection.execute( select(comparison_simulations).where( @@ -402,7 +451,7 @@ def replace_simulation( record: ComparisonSimulationRecord, ) -> ComparisonSimulationRecord: values = _record_values(record) - with self.engine.begin() as connection: + with self._transaction("replace_simulation") as connection: row = ( connection.execute( select(comparison_simulations) @@ -464,7 +513,7 @@ def attach_simulation_invocation( ) .returning(*comparison_simulations.c) ) - with self.engine.begin() as connection: + with self._transaction("attach_simulation_invocation") as connection: row = connection.execute(statement).mappings().one_or_none() if row is None: row = ( diff --git a/libs/policyengine-stage12-persistence/tests/test_store.py b/libs/policyengine-stage12-persistence/tests/test_store.py index 9ef406e0b..9df7938b2 100644 --- a/libs/policyengine-stage12-persistence/tests/test_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_store.py @@ -1,5 +1,6 @@ from __future__ import annotations +import sqlite3 from contextlib import contextmanager from datetime import UTC, datetime, timedelta from uuid import UUID @@ -13,7 +14,11 @@ SimulationRole, ) from policyengine_stage12_persistence import Stage12PersistenceStore +from policyengine_stage12_persistence import store as store_module from policyengine_stage12_persistence.store import normalized_database_url +from sqlalchemy import create_engine +from sqlalchemy.exc import OperationalError +from sqlalchemy.pool import QueuePool from sqlalchemy.sql import ClauseElement NOW = datetime(2026, 9, 22, tzinfo=UTC) @@ -22,6 +27,42 @@ OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000012" +@pytest.mark.parametrize("fail_transaction", [False, True]) +def test_runtime_engine_closes_connections_between_transactions( + monkeypatch, fail_transaction +) -> None: + def local_engine(_url, **options): + # Exercise the runtime pooling configuration with a real local DBAPI. + # PostgreSQL's default is QueuePool; SQLite's memory default differs. + options.pop("connect_args") + options.setdefault("poolclass", QueuePool) + return create_engine("sqlite://", **options) + + monkeypatch.setattr(store_module, "create_engine", local_engine) + engine = store_module.create_stage12_engine( + "postgresql://runtime:secret@db.example/postgres" + ) + try: + try: + with engine.begin() as connection: + physical_connection = connection.connection.driver_connection + assert connection.exec_driver_sql("SELECT 1").scalar_one() == 1 + if fail_transaction: + raise RuntimeError("calculation failed") + except RuntimeError: + assert fail_transaction + + # Long calculations must not hold an idle session after DB work ends. + with pytest.raises(sqlite3.ProgrammingError, match="closed database"): + physical_connection.execute("SELECT 1") + + with engine.begin() as connection: + assert connection.connection.driver_connection is not physical_connection + assert connection.exec_driver_sql("SELECT 1").scalar_one() == 1 + finally: + engine.dispose() + + def _report() -> ComparisonReportRecord: return ComparisonReportRecord( evaluation_id=REPORT_ID, @@ -93,10 +134,33 @@ def all(self): return self.rows or [] +def _refused_connection_error() -> OperationalError: + # SQLAlchemy wraps a DBAPI failure raised while connecting without a + # statement, exactly as it wraps a failed COMMIT. + return OperationalError( + None, + None, + Exception("(EMAXCONNSESSION) max clients reached in session mode"), + ) + + class FakeConnection: - def __init__(self, results): + def __init__(self, results, *, commit_error=None): self.results = list(results) self.statements = [] + self.commit_error = commit_error + + def __enter__(self): + return self + + def __exit__(self, *_exception_details): + return None + + @contextmanager + def begin(self): + yield + if self.commit_error is not None: + raise self.commit_error def execute(self, statement): assert isinstance(statement, ClauseElement) @@ -105,16 +169,118 @@ def execute(self, statement): class FakeEngine: - def __init__(self, results): - self.connection = FakeConnection(results) - - @contextmanager - def begin(self): - yield self.connection + def __init__(self, results, *, refused_connections=0, commit_error=None): + self.connection = FakeConnection(results, commit_error=commit_error) + self.refused_connections = refused_connections + self.connection_attempts = 0 - @contextmanager def connect(self): - yield self.connection + self.connection_attempts += 1 + if self.connection_attempts <= self.refused_connections: + raise _refused_connection_error() + return self.connection + + +class RetryingConnectionOpener: + """Record each operation's connection and retry one refused connection.""" + + def __init__(self): + self.operation_names = [] + self.refusals = [] + + def __call__(self, connect, *, operation_name): + self.operation_names.append(operation_name) + try: + return connect() + except OperationalError as refusal: + self.refusals.append(refusal) + return connect() + + +def _sqlalchemy_commit_failure(monkeypatch) -> OperationalError: + """Return the error SQLAlchemy raises when a COMMIT's reply is lost.""" + + engine = create_engine("sqlite://") + + def commit_then_lose_the_reply(dbapi_connection): + dbapi_connection.commit() + raise sqlite3.OperationalError("server closed the connection unexpectedly") + + monkeypatch.setattr(engine.dialect, "do_commit", commit_then_lose_the_reply) + try: + with ( + pytest.raises(OperationalError) as commit_failure, + engine.begin() as connection, + ): + connection.exec_driver_sql("SELECT 1") + finally: + engine.dispose() + return commit_failure.value + + +def test_operations_open_connections_through_the_opener_before_any_statement() -> None: + report = _report() + simulation = _simulation() + engine = FakeEngine( + [ + FakeResult(row=_row(report)), + FakeResult(row=_row(simulation)), + FakeResult(row=_row(simulation)), + ], + refused_connections=1, + ) + opener = RetryingConnectionOpener() + store = Stage12PersistenceStore( + engine=engine, # type: ignore[arg-type] + open_connection=opener, + ) + + created = store.create_or_resolve_simulation(simulation) + read = store.get_simulation(SIMULATION_ID) + + assert created.record == simulation + assert read == simulation + assert opener.operation_names == ["create_or_resolve_simulation", "get_simulation"] + # The refusal reached the opener before the store sent any statement. + assert len(opener.refusals) == 1 + assert engine.connection_attempts == 3 + assert len(engine.connection.statements) == 3 + + +def test_commit_failure_is_raised_after_the_opener_returns(monkeypatch) -> None: + # A failed COMMIT carries no statement, like a refused connection, but the + # write may already be durable. Only where it occurs tells them apart. + lost_commit_reply = _sqlalchemy_commit_failure(monkeypatch) + assert lost_commit_reply.statement is None + report = _report() + engine = FakeEngine( + [FakeResult(row=_row(report))], + commit_error=lost_commit_reply, + ) + opener = RetryingConnectionOpener() + store = Stage12PersistenceStore( + engine=engine, # type: ignore[arg-type] + open_connection=opener, + ) + + with pytest.raises(OperationalError) as failure: + store.create_or_resolve_report(report) + + assert failure.value is lost_commit_reply + assert opener.operation_names == ["create_or_resolve_report"] + assert opener.refusals == [] + assert engine.connection_attempts == 1 + assert len(engine.connection.statements) == 1 + + +def test_default_store_does_not_retry_a_refused_connection() -> None: + engine = FakeEngine([], refused_connections=1) + store = Stage12PersistenceStore(engine=engine) # type: ignore[arg-type] + + with pytest.raises(OperationalError, match="EMAXCONNSESSION"): + store.get_report(REPORT_ID) + + assert engine.connection_attempts == 1 def test_create_and_read_report_use_sqlalchemy_statements_only() -> None: 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..5ccf33c70 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 @@ -5,9 +5,10 @@ from collections.abc import Callable, Mapping from contextlib import nullcontext from datetime import UTC, datetime +from functools import partial from hashlib import sha256 from typing import Any -from uuid import uuid4 +from uuid import UUID, uuid4 from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( @@ -51,6 +52,7 @@ artifact_store, runtime_store, ) +from .failure_recording import persist_failure_record, report_original_failure from .output_planning import ( plan_simulation_input, resolve_report_output_plan, @@ -187,6 +189,30 @@ def _claim_parent( return claimed, True +def _mark_child_failed( + persistence: ComparisonStore, + simulation_execution_id: UUID, + *, + error_code: str, + error_summary: str, +) -> None: + latest = persistence.get_simulation(simulation_execution_id) + if latest.status is ComparisonRunLifecycleStatus.SUCCEEDED: + return + failed_at = datetime.now(UTC) + persistence.replace_simulation( + latest.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": error_code, + "error_summary": error_summary, + "updated_at": failed_at, + "completed_at": failed_at, + } + ) + ) + + def coordinate_report( payload: object, context_payload: object, @@ -216,7 +242,7 @@ def coordinate_report( context=context, parent=submitted_parent, ) - persistence = store or runtime_store() + persistence = store or runtime_store(runtime) artifact_storage = artifacts or artifact_store() child_invoker = invoker or ModalChildInvoker() stage_plan: StagePlan = ( @@ -243,6 +269,9 @@ def coordinate_report( } descriptors: dict[str, SimulationArtifactDescriptor] = {} calls: dict[str, ChildCall] = {} + # A child failure is reported where it occurs and then propagates to the + # report handler below, which must not report the same error twice. + reported_child_failure: Exception | None = None try: def stage_scope( @@ -369,20 +398,24 @@ def stage_scope( observability_context=child_observability_context, ) except Exception as error: - failed_at = datetime.now(UTC) - latest = persistence.get_simulation(simulation.simulation_execution_id) - if latest.status is not ComparisonRunLifecycleStatus.SUCCEEDED: - persistence.replace_simulation( - latest.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": "simulation_dispatch_failed", - "error_summary": type(error).__name__, - "updated_at": failed_at, - "completed_at": failed_at, - } - ) - ) + report_original_failure( + error, + runtime=runtime, + failure_scope="child_dispatch", + attributes={"simulation_role": simulation.role.value}, + ) + reported_child_failure = error + persist_failure_record( + partial( + _mark_child_failed, + persistence, + simulation.simulation_execution_id, + error_code="simulation_dispatch_failed", + error_summary=type(error).__name__, + ), + runtime=runtime, + failure_scope="child_dispatch", + ) raise calls[simulation.role.value] = call persistence.attach_simulation_invocation( @@ -407,23 +440,27 @@ def stage_scope( call.get(timeout=SIMULATION_WAIT_TIMEOUT_SECONDS) ) except Exception as error: + report_original_failure( + error, + runtime=runtime, + failure_scope="child_wait", + attributes={"simulation_role": role}, + ) + reported_child_failure = error simulation = next( item for item in simulations if item.role.value == role ) - child = persistence.get_simulation(simulation.simulation_execution_id) - if child.status is not ComparisonRunLifecycleStatus.SUCCEEDED: - failed_at = datetime.now(UTC) - persistence.replace_simulation( - child.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": "simulation_invocation_failed", - "error_summary": type(error).__name__, - "updated_at": failed_at, - "completed_at": failed_at, - } - ) - ) + persist_failure_record( + partial( + _mark_child_failed, + persistence, + simulation.simulation_execution_id, + error_code="simulation_invocation_failed", + error_summary=type(error).__name__, + ), + runtime=runtime, + failure_scope="child_wait", + ) raise baseline = descriptors["baseline"] reform = descriptors["reform"] @@ -509,17 +546,28 @@ def stage_scope( # Persist any coordinator, child-call, aggregation, or artifact failure # before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 - failed_at = datetime.now(UTC) - persistence.replace_report( - parent.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "aggregation_status": ComparisonRunAggregationStatus.FAILED, - "error_code": "report_coordination_failed", - "error_summary": type(error).__name__, - "updated_at": failed_at, - "completed_at": failed_at, - } + if error is not reported_child_failure: + report_original_failure( + error, + runtime=runtime, + failure_scope="report_coordination", ) + failed_at = datetime.now(UTC) + persist_failure_record( + partial( + persistence.replace_report, + parent.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "aggregation_status": ComparisonRunAggregationStatus.FAILED, + "error_code": "report_coordination_failed", + "error_summary": type(error).__name__, + "updated_at": failed_at, + "completed_at": failed_at, + } + ), + ), + runtime=runtime, + failure_scope="report_coordination", ) raise RuntimeError("Stage 12 report coordination failed") from None diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py new file mode 100644 index 000000000..a432cf3fe --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py @@ -0,0 +1,103 @@ +"""Bounded retry for Stage 12 database connections that could not be opened.""" + +from __future__ import annotations + +import logging +import random +import time +from collections.abc import Callable +from contextlib import suppress + +from policyengine_observability import ObservabilityRuntime +from sqlalchemy.engine import Connection +from sqlalchemy.exc import OperationalError + +logger = logging.getLogger(__name__) + +# Coordinators and workers start in bursts against a session pooler with a +# fixed client limit. A short, jittered wait lets concurrent connections +# spread out instead of failing a report on one refused connection. +CONNECTION_RETRY_BUDGET_SECONDS = 60.0 +INITIAL_RETRY_DELAY_SECONDS = 1.0 +MAXIMUM_RETRY_DELAY_SECONDS = 15.0 +RETRY_DELAY_GROWTH_FACTOR = 2.0 + + +def open_connection_with_retry( + connect: Callable[[], Connection], + *, + operation_name: str, + runtime: ObservabilityRuntime | None = None, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, + choose_delay: Callable[[float, float], float] = random.uniform, +) -> Connection: + """Open a Stage 12 database connection, retrying while it is refused. + + The persistence store passes only its connect step here. Failing that step + means no connection opened and no statement reached the database, so + opening it again cannot duplicate a write. The store runs the operation's + statements and its COMMIT or ROLLBACK after this returns; their failures, + including a COMMIT whose reply was lost, propagate without a retry. + Errors other than ``OperationalError`` propagate immediately, and retries + stop once the retry budget is spent. + """ + + started = monotonic() + attempt = 1 + delay_ceiling = INITIAL_RETRY_DELAY_SECONDS + while True: + try: + return connect() + except OperationalError as error: + remaining_budget = CONNECTION_RETRY_BUDGET_SECONDS - (monotonic() - started) + if remaining_budget <= 0: + raise + delay = min( + choose_delay(delay_ceiling / 2, delay_ceiling), remaining_budget + ) + _report_connection_retry( + runtime, + operation_name=operation_name, + failed_attempt=attempt, + delay_seconds=delay, + error=error, + ) + sleep(delay) + attempt += 1 + delay_ceiling = min( + delay_ceiling * RETRY_DELAY_GROWTH_FACTOR, + MAXIMUM_RETRY_DELAY_SECONDS, + ) + + +def _report_connection_retry( + runtime: ObservabilityRuntime | None, + *, + operation_name: str, + failed_attempt: int, + delay_seconds: float, + error: Exception, +) -> None: + if runtime is None: + logger.warning( + "Stage 12 database connection unavailable for %s; retrying " + "attempt %d after %.3f seconds", + operation_name, + failed_attempt + 1, + delay_seconds, + exc_info=error, + ) + return + # Telemetry must not change whether the operation is retried. + with suppress(Exception): + runtime.log( + "Stage 12 database connection unavailable; retrying", + severity="WARNING", + attributes={ + "stage12_database_operation": operation_name, + "stage12_database_failed_attempt": failed_attempt, + "stage12_database_retry_delay_seconds": round(delay_seconds, 3), + }, + error=error, + ) 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..e5848ebe0 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 @@ -4,9 +4,11 @@ import os from datetime import datetime +from functools import partial from typing import Any, Protocol from uuid import UUID +from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( ComparisonReportPersistenceResult, ComparisonReportRecord, @@ -17,6 +19,8 @@ from policyengine_simulation_executor.stage12_artifacts import Stage12ArtifactStore +from .database_retry import open_connection_with_retry + class ComparisonStore(Protocol): def create_or_resolve_report( @@ -112,8 +116,15 @@ def restore(self, invocation_id: str) -> ChildCall: return modal.FunctionCall.from_id(invocation_id) -def runtime_store() -> Stage12PersistenceStore: - return Stage12PersistenceStore(os.environ.get("STAGE12_DATABASE_URL", "")) +def runtime_store( + runtime: ObservabilityRuntime | None = None, +) -> Stage12PersistenceStore: + # Every operation opens its own pooler session, so each one retries a + # refused connection rather than failing the report during a burst. + return Stage12PersistenceStore( + os.environ.get("STAGE12_DATABASE_URL", ""), + open_connection=partial(open_connection_with_retry, runtime=runtime), + ) def artifact_store() -> Stage12ArtifactStore: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py new file mode 100644 index 000000000..342c315f4 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py @@ -0,0 +1,78 @@ +"""Failure reporting that does not depend on the Stage 12 database.""" + +from __future__ import annotations + +import logging +from collections.abc import Callable, Mapping +from contextlib import suppress + +from policyengine_observability import ObservabilityRuntime + +logger = logging.getLogger(__name__) + + +def report_original_failure( + error: Exception, + *, + runtime: ObservabilityRuntime | None, + failure_scope: str, + attributes: Mapping[str, object] | None = None, +) -> None: + """Emit the original failure before any database write is attempted. + + Recording a failure in the database can itself fail, for example when the + session pooler has no free connections. Telemetry therefore receives the + original cause first, so it survives a failed or hung database write. + """ + + _emit_error( + "Stage 12 execution failed", + error=error, + runtime=runtime, + attributes={"stage12_failure_scope": failure_scope, **(attributes or {})}, + ) + + +def persist_failure_record( + write_failure_record: Callable[[], object], + *, + runtime: ObservabilityRuntime | None, + failure_scope: str, +) -> None: + """Attempt a failure write without letting its error replace the original. + + The caller raises the same failure after this returns, whether or not the + record was written: a child handler re-raises the original error, and the + report and simulation handlers raise their stable ``RuntimeError``. The + original cause then remains visible only in telemetry. + """ + + try: + write_failure_record() + except Exception as persistence_error: # noqa: BLE001 + _emit_error( + "Stage 12 failure record could not be persisted", + error=persistence_error, + runtime=runtime, + attributes={"stage12_failure_scope": failure_scope}, + ) + + +def _emit_error( + message: str, + *, + error: Exception, + runtime: ObservabilityRuntime | None, + attributes: Mapping[str, object], +) -> None: + if runtime is None: + logger.error( + "%s (%s)", + message, + attributes.get("stage12_failure_scope"), + exc_info=error, + ) + return + # Telemetry must not change the failure the caller propagates. + with suppress(Exception): + runtime.log(message, severity="ERROR", attributes=attributes, error=error) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index cc483437f..139ce06cd 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -6,6 +6,7 @@ from contextlib import nullcontext from dataclasses import dataclass from datetime import UTC, datetime +from functools import partial from hashlib import sha256 from typing import Any @@ -38,6 +39,7 @@ ) from .dependencies import ComparisonStore, artifact_store, runtime_store +from .failure_recording import persist_failure_record, report_original_failure from .output_planning import apply_output_plan, validate_output_frames @@ -304,7 +306,7 @@ def run_single_simulation( simulation = PlannedSimulationExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) _require_context(simulation, context, required_country=required_country) - persistence = store or runtime_store() + persistence = store or runtime_store(runtime) artifact_storage = artifacts or artifact_store() child = persistence.get_simulation(simulation.simulation_execution_id) if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: @@ -385,16 +387,27 @@ def run_single_simulation( # Persist any country-package calculation failure before returning a # stable exception to Modal. except Exception as error: # noqa: BLE001 + report_original_failure( + error, + runtime=runtime, + failure_scope="simulation_execution", + attributes={"simulation_role": simulation.role.value}, + ) failed_at = datetime.now(UTC) - persistence.replace_simulation( - running.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": "simulation_execution_failed", - "error_summary": type(error).__name__, - "updated_at": failed_at, - "completed_at": failed_at, - } - ) + persist_failure_record( + partial( + persistence.replace_simulation, + running.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_execution_failed", + "error_summary": type(error).__name__, + "updated_at": failed_at, + "completed_at": failed_at, + } + ), + ), + runtime=runtime, + failure_scope="simulation_execution", ) raise RuntimeError("Stage 12 simulation execution failed") from None diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py b/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py new file mode 100644 index 000000000..ccfe6d40b --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py @@ -0,0 +1,227 @@ +"""Tests for retrying Stage 12 database connections that could not be opened.""" + +from __future__ import annotations + +import logging + +import pytest +from policyengine_simulation_executor.stage12_runtime import database_retry +from policyengine_simulation_executor.stage12_runtime.database_retry import ( + open_connection_with_retry, +) +from sqlalchemy.exc import OperationalError + + +def _connection_limit_error() -> OperationalError: + # SQLAlchemy wraps a failure to open a connection without a statement. + return OperationalError( + None, + None, + Exception("(EMAXCONNSESSION) max clients reached in session mode"), + ) + + +class ManualClock: + def __init__(self) -> None: + self.now = 0.0 + self.sleeps: list[float] = [] + + def monotonic(self) -> float: + return self.now + + def sleep(self, seconds: float) -> None: + self.sleeps.append(seconds) + self.now += seconds + + +class RetryLogRuntime: + def __init__(self) -> None: + self.records: list[dict[str, object]] = [] + + def log(self, message, *, severity="INFO", attributes=None, error=None): + self.records.append( + { + "message": message, + "severity": severity, + "attributes": dict(attributes or {}), + "error": error, + } + ) + + +class OpenedConnection: + """Stand in for the connection a successful connect step returns.""" + + +def _upper_bound_delay(lower_bound: float, upper_bound: float) -> float: + return upper_bound + + +def test_refused_connections_are_retried_with_growing_bounded_delays() -> None: + clock = ManualClock() + runtime = RetryLogRuntime() + opened_connection = OpenedConnection() + attempts: list[int] = [] + + def connect_after_two_refusals() -> OpenedConnection: + attempts.append(len(attempts) + 1) + if len(attempts) <= 2: + raise _connection_limit_error() + return opened_connection + + result = open_connection_with_retry( + connect_after_two_refusals, + operation_name="create_or_resolve_simulation", + runtime=runtime, + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ) + + assert result is opened_connection + assert attempts == [1, 2, 3] + assert clock.sleeps == [1.0, 2.0] + assert [record["attributes"] for record in runtime.records] == [ + { + "stage12_database_operation": "create_or_resolve_simulation", + "stage12_database_failed_attempt": 1, + "stage12_database_retry_delay_seconds": 1.0, + }, + { + "stage12_database_operation": "create_or_resolve_simulation", + "stage12_database_failed_attempt": 2, + "stage12_database_retry_delay_seconds": 2.0, + }, + ] + assert all(record["severity"] == "WARNING" for record in runtime.records) + assert all( + isinstance(record["error"], OperationalError) for record in runtime.records + ) + + +def test_jittered_delay_stays_between_half_and_all_of_the_ceiling() -> None: + clock = ManualClock() + requested_ranges: list[tuple[float, float]] = [] + + def lowest_delay(lower_bound: float, upper_bound: float) -> float: + requested_ranges.append((lower_bound, upper_bound)) + return lower_bound + + attempts: list[int] = [] + + def connect_on_third_attempt() -> OpenedConnection: + attempts.append(1) + if len(attempts) < 3: + raise _connection_limit_error() + return OpenedConnection() + + open_connection_with_retry( + connect_on_third_attempt, + operation_name="get_simulation", + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=lowest_delay, + ) + + assert requested_ranges == [(0.5, 1.0), (1.0, 2.0)] + assert clock.sleeps == [0.5, 1.0] + + +def test_errors_other_than_a_refused_connection_are_not_retried() -> None: + clock = ManualClock() + attempts: list[int] = [] + + def fail_with_invalid_connection_options() -> None: + attempts.append(1) + raise ValueError("invalid connection option") + + with pytest.raises(ValueError, match="invalid connection option"): + open_connection_with_retry( + fail_with_invalid_connection_options, + operation_name="get_report", + sleep=clock.sleep, + monotonic=clock.monotonic, + ) + + assert attempts == [1] + assert clock.sleeps == [] + + +def test_retries_stop_when_the_budget_is_spent() -> None: + clock = ManualClock() + attempts: list[int] = [] + + def never_connect() -> None: + attempts.append(1) + raise _connection_limit_error() + + with pytest.raises(OperationalError, match="EMAXCONNSESSION"): + open_connection_with_retry( + never_connect, + operation_name="create_or_resolve_report", + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ) + + budget = database_retry.CONNECTION_RETRY_BUDGET_SECONDS + assert sum(clock.sleeps) == pytest.approx(budget) + assert max(clock.sleeps) <= database_retry.MAXIMUM_RETRY_DELAY_SECONDS + assert len(attempts) == len(clock.sleeps) + 1 + + +def test_retry_telemetry_failure_does_not_stop_the_retry() -> None: + class FailingRuntime: + def log(self, *_args, **_kwargs): + raise RuntimeError("telemetry exporter unavailable") + + clock = ManualClock() + opened_connection = OpenedConnection() + attempts: list[int] = [] + + def connect_on_second_attempt() -> OpenedConnection: + attempts.append(1) + if len(attempts) == 1: + raise _connection_limit_error() + return opened_connection + + assert ( + open_connection_with_retry( + connect_on_second_attempt, + operation_name="create_or_resolve_report", + runtime=FailingRuntime(), + sleep=clock.sleep, + monotonic=clock.monotonic, + ) + is opened_connection + ) + + +def test_retry_is_logged_when_no_observability_runtime_is_available( + caplog, +) -> None: + clock = ManualClock() + attempts: list[int] = [] + + def connect_on_second_attempt() -> OpenedConnection: + attempts.append(1) + if len(attempts) == 1: + raise _connection_limit_error() + return OpenedConnection() + + with caplog.at_level(logging.WARNING, logger=database_retry.__name__): + open_connection_with_retry( + connect_on_second_attempt, + operation_name="replace_simulation", + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ) + + [record] = caplog.records + assert record.levelno == logging.WARNING + assert record.getMessage() == ( + "Stage 12 database connection unavailable for replace_simulation; " + "retrying attempt 2 after 1.000 seconds" + ) + assert isinstance(record.exc_info[1], OperationalError) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py b/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py new file mode 100644 index 000000000..7ca71d018 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py @@ -0,0 +1,53 @@ +"""Tests for reporting Stage 12 failures without an observability runtime.""" + +from __future__ import annotations + +import logging + +from policyengine_simulation_executor.stage12_runtime import failure_recording +from policyengine_simulation_executor.stage12_runtime.failure_recording import ( + persist_failure_record, + report_original_failure, +) + + +def test_original_failure_is_logged_when_no_observability_runtime_is_available( + caplog, +) -> None: + calculation_error = ValueError("country column missing from household data") + + with caplog.at_level(logging.ERROR, logger=failure_recording.__name__): + report_original_failure( + calculation_error, + runtime=None, + failure_scope="simulation_execution", + attributes={"simulation_role": "baseline"}, + ) + + [record] = caplog.records + assert record.levelno == logging.ERROR + assert record.getMessage() == "Stage 12 execution failed (simulation_execution)" + assert record.exc_info[1] is calculation_error + + +def test_unwritable_failure_record_is_logged_when_no_runtime_is_available( + caplog, +) -> None: + persistence_error = ConnectionError("database unavailable") + + def fail_to_write_record() -> None: + raise persistence_error + + with caplog.at_level(logging.ERROR, logger=failure_recording.__name__): + persist_failure_record( + fail_to_write_record, + runtime=None, + failure_scope="report_coordination", + ) + + [record] = caplog.records + assert record.levelno == logging.ERROR + assert record.getMessage() == ( + "Stage 12 failure record could not be persisted (report_coordination)" + ) + assert record.exc_info[1] is persistence_error diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index f989c3a90..c7c9431a2 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -69,18 +69,30 @@ SimulationCalculation, _build_spm_result, build_aggregate_report, + database_retry, run_single_simulation, simulation_input_sha256, ) 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 import ( + simulation as simulation_module, +) from policyengine_simulation_executor.stage12_runtime.aggregation import ( validate_uk_local_authority_metadata, ) +from policyengine_simulation_executor.stage12_runtime.dependencies import ( + runtime_store, +) from policyengine_simulation_executor.uk_local_authority_metadata import ( detect_uk_local_authority_metadata, ) +from policyengine_stage12_persistence import store as persistence_store_module +from sqlalchemy.exc import OperationalError NOW = datetime(2026, 9, 15, tzinfo=UTC) EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") @@ -702,6 +714,101 @@ def fail(_simulation): assert child.error_summary == "RuntimeError" +def _connection_limit_error() -> OperationalError: + # SQLAlchemy wraps a failure to open a connection without a statement. + return OperationalError( + None, + None, + Exception("(EMAXCONNSESSION) max clients reached in session mode"), + ) + + +class TelemetryRecordingRuntime(RecordingRuntime): + """Record telemetry in the same ordered list as database writes.""" + + def __init__(self, ordered_events: list[tuple[str, object]]) -> None: + super().__init__() + self.ordered_events = ordered_events + + def log(self, message, *, severity="INFO", attributes=None, error=None): + self.ordered_events.append( + ("telemetry", (message, dict(attributes or {}), error)) + ) + + +def test_single_worker_builds_its_store_with_the_worker_runtime(monkeypatch) -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + store_runtimes = [] + + def record_store_runtime(runtime): + store_runtimes.append(runtime) + return store + + monkeypatch.setattr(simulation_module, "runtime_store", record_store_runtime) + runtime = RecordingRuntime() + + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + artifacts=FakeArtifacts(), + calculator=lambda _: _frames(), + runtime=runtime, + ) + + # The runtime store reports each connection retry through this runtime. + assert store_runtimes == [runtime] + + +def test_single_worker_reports_failure_before_an_unwritable_failure_record() -> None: + ordered_events: list[tuple[str, object]] = [] + + class UnwritableFailureStore(FakeStore): + def replace_simulation(self, record): + if record.status is ComparisonRunLifecycleStatus.FAILED: + ordered_events.append(("failure_write", record.error_summary)) + raise _connection_limit_error() + return super().replace_simulation(record) + + store = UnwritableFailureStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + calculation_error = ValueError("country column missing from household data") + + def fail(_simulation): + raise calculation_error + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + runtime=TelemetryRecordingRuntime(ordered_events), + ) + + original_report, failure_write, persistence_report = ordered_events + assert original_report[0] == "telemetry" + _message, attributes, reported_error = original_report[1] + assert reported_error is calculation_error + assert attributes == { + "stage12_failure_scope": "simulation_execution", + "simulation_role": "baseline", + } + assert failure_write == ("failure_write", "ValueError") + assert persistence_report[0] == "telemetry" + assert persistence_report[1][0] == ( + "Stage 12 failure record could not be persisted" + ) + assert isinstance(persistence_report[1][2], OperationalError) + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.RUNNING + + def test_single_worker_rejects_frames_that_do_not_satisfy_the_output_plan() -> None: store = FakeStore() simulation = _planned_simulation(SimulationRole.BASELINE) @@ -1413,6 +1520,337 @@ def create_or_resolve_simulation(self, record): assert store.parent.error_summary == "RuntimeError" +def test_coordinator_builds_its_store_with_the_coordinator_runtime( + monkeypatch, +) -> None: + store = FakeStore() + store_runtimes = [] + + def record_store_runtime(runtime): + store_runtimes.append(runtime) + return store + + monkeypatch.setattr(coordination_module, "runtime_store", record_store_runtime) + artifacts = FakeArtifacts() + runtime = RecordingRuntime() + + 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", + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=lambda **_: {"result": "complete"}, + runtime=runtime, + ) + + # The runtime store reports each connection retry through this runtime. + assert store_runtimes == [runtime] + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +def test_coordinator_does_not_deduplicate_against_its_own_claim_after_a_lost_commit() -> ( + None +): + # SQLAlchemy raises a failed COMMIT without a statement, like a refused + # connection, although the claim may already be durable. + lost_commit_reply = OperationalError( + None, + None, + Exception("server closed the connection unexpectedly"), + ) + + class LostCommitReplyStore(FakeStore): + def __init__(self): + super().__init__() + self.claim_attempts = 0 + + def create_or_resolve_report(self, record): + self.claim_attempts += 1 + claim = super().create_or_resolve_report(record) + if self.claim_attempts == 1: + raise lost_commit_reply + return claim + + store = LostCommitReplyStore() + artifacts = FakeArtifacts() + invoker = ConcurrentInvoker(artifacts) + + with pytest.raises(OperationalError) as error: + 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, + ) + + # A retry would find this coordinator's own running claim and return a + # deduplicated success without dispatching either simulation. + assert error.value is lost_commit_reply + assert store.claim_attempts == 1 + assert invoker.events == [] + + +class ScriptedResult: + def __init__(self, row): + self.row = row + + def mappings(self): + return self + + def one_or_none(self): + return self.row + + def one(self): + return self.row + + +class ScriptedConnection: + """Answer each statement with the next scripted row.""" + + def __init__(self, rows, *, commit_error=None): + self.rows = list(rows) + self.statement_count = 0 + self.commit_error = commit_error + + def __enter__(self): + return self + + def __exit__(self, *_exception_details): + return None + + @contextmanager + def begin(self): + yield + if self.commit_error is not None: + raise self.commit_error + + def execute(self, _statement): + self.statement_count += 1 + return ScriptedResult(self.rows.pop(0)) + + +class BusyPoolerEngine: + """Refuse the first connections, as a session pooler at its limit does.""" + + def __init__(self, rows, *, refused_connections=0, commit_error=None): + self.connection = ScriptedConnection(rows, commit_error=commit_error) + self.refused_connections = refused_connections + self.connection_attempts = 0 + + def connect(self): + self.connection_attempts += 1 + if self.connection_attempts <= self.refused_connections: + raise _connection_limit_error() + return self.connection + + +def _runtime_store_over(engine, monkeypatch, runtime): + monkeypatch.setattr(database_retry, "INITIAL_RETRY_DELAY_SECONDS", 0.0) + monkeypatch.setattr( + persistence_store_module, + "create_stage12_engine", + lambda _database_url: engine, + ) + return runtime_store(runtime) + + +def test_runtime_store_retries_a_refused_connection_for_child_state_creation( + monkeypatch, +) -> None: + ordered_events: list[tuple[str, object]] = [] + child = _child(_planned_simulation(SimulationRole.BASELINE)) + engine = BusyPoolerEngine( + [_parent().model_dump(mode="python"), child.model_dump(mode="python")], + refused_connections=1, + ) + store = _runtime_store_over( + engine, monkeypatch, TelemetryRecordingRuntime(ordered_events) + ) + + result = store.create_or_resolve_simulation(child) + + assert result.created is True + assert result.record == child + assert engine.connection_attempts == 2 + assert engine.connection.statement_count == 2 + [(event_kind, (message, attributes, error))] = ordered_events + assert event_kind == "telemetry" + assert message == "Stage 12 database connection unavailable; retrying" + assert attributes["stage12_database_operation"] == "create_or_resolve_simulation" + assert attributes["stage12_database_failed_attempt"] == 1 + assert isinstance(error, OperationalError) + + +def test_runtime_store_does_not_retry_a_failed_commit(monkeypatch) -> None: + ordered_events: list[tuple[str, object]] = [] + lost_commit_reply = OperationalError( + None, + None, + Exception("server closed the connection unexpectedly"), + ) + report = _parent() + engine = BusyPoolerEngine( + [report.model_dump(mode="python")], + commit_error=lost_commit_reply, + ) + store = _runtime_store_over( + engine, monkeypatch, TelemetryRecordingRuntime(ordered_events) + ) + + with pytest.raises(OperationalError) as failure: + store.create_or_resolve_report(report) + + assert failure.value is lost_commit_reply + assert engine.connection_attempts == 1 + assert engine.connection.statement_count == 1 + assert ordered_events == [] + + +def test_coordinator_reports_failure_before_an_unwritable_failure_record() -> None: + ordered_events: list[tuple[str, object]] = [] + aggregation_error = ValueError("aggregate failed validation") + + class UnwritableFailureStore(FakeStore): + def replace_report(self, record): + if record.status is ComparisonRunLifecycleStatus.FAILED: + ordered_events.append(("failure_write", record.error_summary)) + raise _connection_limit_error() + return super().replace_report(record) + + def fail_aggregation(**_): + raise aggregation_error + + store = UnwritableFailureStore() + artifacts = FakeArtifacts() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + 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=ConcurrentInvoker(artifacts), + aggregator=fail_aggregation, + runtime=TelemetryRecordingRuntime(ordered_events), + ) + + assert error.value.__cause__ is None + original_report, failure_write, persistence_report = ordered_events + assert original_report[1][1] == {"stage12_failure_scope": "report_coordination"} + assert original_report[1][2] is aggregation_error + assert failure_write == ("failure_write", "ValueError") + assert isinstance(persistence_report[1][2], OperationalError) + assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING + + +def test_coordinator_reports_a_child_failure_once_when_its_record_is_unwritable() -> ( + None +): + ordered_events: list[tuple[str, object]] = [] + + class UnreadableChildStore(FakeStore): + def get_simulation(self, simulation_execution_id): + ordered_events.append(("child_read", str(simulation_execution_id))) + raise _connection_limit_error() + + store = UnreadableChildStore() + artifacts = FakeArtifacts() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + 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=ConcurrentInvoker(artifacts, fail_role="baseline"), + aggregator=lambda **_: {"must": "not run"}, + runtime=TelemetryRecordingRuntime(ordered_events), + ) + + failure_reports = [ + event[1] + for event in ordered_events + if event[0] == "telemetry" and event[1][0] == "Stage 12 execution failed" + ] + assert [report[1] for report in failure_reports] == [ + {"stage12_failure_scope": "child_wait", "simulation_role": "baseline"} + ] + assert ordered_events[0][0] == "telemetry" + assert ordered_events[1] == ("child_read", str(BASELINE_ID)) + # The report record is still written after the child record could not be. + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_summary == "RuntimeError" + + +class RefusingDispatchInvoker: + def __init__(self, dispatch_error): + self.dispatch_error = dispatch_error + self.events = [] + + def spawn(self, *, simulation, **_): + role = PlannedSimulationExecutionInput.model_validate(simulation).role.value + self.events.append(f"spawn:{role}") + raise self.dispatch_error + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") + + +def test_coordinator_reports_a_child_dispatch_failure_once() -> None: + ordered_events: list[tuple[str, object]] = [] + dispatch_error = ConnectionError("Modal refused the dispatch") + store = FakeStore() + invoker = RefusingDispatchInvoker(dispatch_error) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + 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=FakeArtifacts(), + invoker=invoker, + aggregator=lambda **_: {"must": "not run"}, + runtime=TelemetryRecordingRuntime(ordered_events), + ) + + failure_reports = [ + event[1] + for event in ordered_events + if event[0] == "telemetry" and event[1][0] == "Stage 12 execution failed" + ] + assert [(attributes, error) for _message, attributes, error in failure_reports] == [ + ( + {"stage12_failure_scope": "child_dispatch", "simulation_role": "baseline"}, + dispatch_error, + ) + ] + assert invoker.events == ["spawn:baseline"] + baseline_child = store.children[BASELINE_ID] + assert baseline_child.status is ComparisonRunLifecycleStatus.FAILED + assert baseline_child.error_code == "simulation_dispatch_failed" + assert baseline_child.error_summary == "ConnectionError" + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_summary == "ConnectionError" + + def test_coordinator_rejects_incompatible_row_identity() -> None: store = FakeStore() artifacts = FakeArtifacts()