From a0e4e75ef20b4b2e2d6aecf85057f66904f7fd7c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:17:43 +0400 Subject: [PATCH 1/3] Release Stage 12 database sessions and keep original failures visible Stage 12 coordinators and workers held pooled database sessions open while they calculated or waited, so concurrent runs exhausted the session pooler's 15-client limit (EMAXCONNSESSION). The runtime engine now uses NullPool, closing each physical connection when its operation finishes. Coordinator and worker claims retry, with bounded jittered backoff (60 s budget), only when a database connection could not be opened. No statement reached the database in that case, so the retry cannot duplicate a write. Failed statements and validation errors still fail immediately. Coordinator and worker failure handlers now report the original error to telemetry before writing the failure record, and a failed write no longer replaces the original error. A child failure is reported once. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 1 + .../__init__.py | 2 + .../policyengine_stage12_persistence/store.py | 23 +- .../tests/test_store.py | 73 ++++- .../stage12_runtime/coordination.py | 142 +++++++--- .../stage12_runtime/database_retry.py | 104 +++++++ .../stage12_runtime/failure_recording.py | 76 ++++++ .../stage12_runtime/simulation.py | 77 ++++-- .../tests/test_stage12_database_retry.py | 201 ++++++++++++++ .../tests/test_stage12_runtime.py | 254 ++++++++++++++++++ 10 files changed, 881 insertions(+), 72 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 066e85eb1..d65885efd 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 Stage 12 claims that could not open a database connection, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot replace 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..d85893ca1 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/__init__.py @@ -3,6 +3,7 @@ from policyengine_stage12_persistence.store import ( Stage12PersistenceStore, create_stage12_engine, + is_connection_unavailable_error, normalized_database_url, ) from policyengine_stage12_persistence.validation import ( @@ -14,6 +15,7 @@ "Stage12PersistenceStore", "Stage12RuntimeAccessError", "create_stage12_engine", + "is_connection_unavailable_error", "normalized_database_url", "validate_runtime_database", ] 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..2adcbd1fd 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py @@ -18,6 +18,8 @@ from sqlalchemy import Engine, create_engine, select, update from sqlalchemy.dialects.postgresql import insert from sqlalchemy.engine import URL, make_url +from sqlalchemy.exc import OperationalError +from sqlalchemy.pool import NullPool from policyengine_stage12_persistence.tables import ( REPORT_IDENTITY_CONSTRAINT, @@ -115,14 +117,27 @@ 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}, ) +def is_connection_unavailable_error(error: BaseException) -> bool: + """Return whether a database connection could not be opened. + + SQLAlchemy wraps a DBAPI failure raised while opening a connection without + a statement. The database therefore received no work, so the operation can + be retried without risking a duplicate write. Session-pooler connection + limits (``EMAXCONNSESSION``) and connect timeouts fail this way; a failed + statement does not. + """ + + return isinstance(error, OperationalError) and error.statement is None + + def _database_value(value: Any) -> Any: if isinstance(value, StrEnum): return value.value diff --git a/libs/policyengine-stage12-persistence/tests/test_store.py b/libs/policyengine-stage12-persistence/tests/test_store.py index 9ef406e0b..54fa75746 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 @@ -12,8 +13,15 @@ ComparisonSimulationRecord, SimulationRole, ) -from policyengine_stage12_persistence import Stage12PersistenceStore +from policyengine_stage12_persistence import ( + Stage12PersistenceStore, + is_connection_unavailable_error, +) +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 +30,69 @@ 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 test_connection_failures_are_distinguished_from_failed_statements( + tmp_path, +) -> None: + unreachable_database = tmp_path / "missing-directory" / "stage12.db" + unreachable_engine = create_engine(f"sqlite:///{unreachable_database}") + reachable_engine = create_engine("sqlite://") + try: + with ( + pytest.raises(OperationalError) as connection_failure, + unreachable_engine.begin(), + ): + pass + with ( + pytest.raises(OperationalError) as statement_failure, + reachable_engine.begin() as connection, + ): + connection.exec_driver_sql("SELECT * FROM missing_table") + finally: + unreachable_engine.dispose() + reachable_engine.dispose() + + # No statement reached the database, so the operation is safe to retry. + assert is_connection_unavailable_error(connection_failure.value) + assert not is_connection_unavailable_error(statement_failure.value) + assert not is_connection_unavailable_error(RuntimeError("not a database error")) + + def _report() -> ComparisonReportRecord: return ComparisonReportRecord( evaluation_id=REPORT_ID, 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..5ce418da5 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 ( @@ -43,6 +44,7 @@ validate_uk_local_authority_metadata, ) from .comparison import compare_completed_report +from .database_retry import retry_when_connection_unavailable from .dependencies import ( ChildCall, ChildInvoker, @@ -51,6 +53,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 +190,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, @@ -230,10 +257,15 @@ def coordinate_report( else nullcontext() ) with claim_span: - parent, should_execute = _claim_parent( - submitted=submitted_parent, - coordinator_invocation_id=coordinator_invocation_id, - store=persistence, + parent, should_execute = retry_when_connection_unavailable( + partial( + _claim_parent, + submitted=submitted_parent, + coordinator_invocation_id=coordinator_invocation_id, + store=persistence, + ), + operation_name="coordinator_claim", + runtime=runtime, ) if not should_execute: return { @@ -243,6 +275,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 +404,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 +446,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 +552,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..471b29997 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/database_retry.py @@ -0,0 +1,104 @@ +"""Bounded retry for Stage 12 database work that could not open a connection.""" + +from __future__ import annotations + +import logging +import random +import time +from collections.abc import Callable +from contextlib import suppress +from typing import TypeVar + +from policyengine_observability import ObservabilityRuntime +from policyengine_stage12_persistence import is_connection_unavailable_error + +logger = logging.getLogger(__name__) + +OperationResult = TypeVar("OperationResult") + +# Coordinators and workers start in bursts against a session pooler with a +# fixed client limit. A short, jittered wait lets concurrent starts spread out +# instead of failing a report before any simulation work has begun. +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 retry_when_connection_unavailable( + operation: Callable[[], OperationResult], + *, + 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, +) -> OperationResult: + """Run ``operation``, retrying only failures to open a database connection. + + The database received no statement when a connection could not be opened, + so repeating the operation cannot duplicate a write. Every other error, + including a failed statement, propagates immediately. Retries stop once + the retry budget is spent and the final connection error propagates. + """ + + started = monotonic() + attempt = 1 + delay_ceiling = INITIAL_RETRY_DELAY_SECONDS + while True: + try: + return operation() + except Exception as error: + if not is_connection_unavailable_error(error): + raise + 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/failure_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py new file mode 100644 index 000000000..469c9b88b --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py @@ -0,0 +1,76 @@ +"""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 re-raises the original failure after this returns, whether or + not the record was written. + """ + + 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..8fe0e8b0e 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,8 +6,10 @@ 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 +from uuid import UUID import pandas as pd from policyengine_observability import ObservabilityRuntime @@ -37,7 +39,9 @@ detect_uk_local_authority_metadata, ) +from .database_retry import retry_when_connection_unavailable 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 @@ -288,6 +292,29 @@ def descriptor_from_record( ) +def _claim_simulation( + persistence: ComparisonStore, + simulation_execution_id: UUID, +) -> tuple[ComparisonSimulationRecord, ComparisonSimulationRecord | None]: + """Return the stored child and, unless it already succeeded, mark it running.""" + + child = persistence.get_simulation(simulation_execution_id) + if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: + return child, None + started = datetime.now(UTC) + running = child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "started_at": child.started_at or started, + "updated_at": started, + "error_code": None, + "error_summary": None, + } + ) + persistence.replace_simulation(running) + return child, running + + def run_single_simulation( payload: object, context_payload: object, @@ -306,20 +333,13 @@ def run_single_simulation( _require_context(simulation, context, required_country=required_country) persistence = store or runtime_store() artifact_storage = artifacts or artifact_store() - child = persistence.get_simulation(simulation.simulation_execution_id) - if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: - return descriptor_from_record(child, simulation).model_dump(mode="json") - started = datetime.now(UTC) - running = child.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "started_at": child.started_at or started, - "updated_at": started, - "error_code": None, - "error_summary": None, - } + child, running = retry_when_connection_unavailable( + partial(_claim_simulation, persistence, simulation.simulation_execution_id), + operation_name="simulation_claim", + runtime=runtime, ) - persistence.replace_simulation(running) + if running is None: + return descriptor_from_record(child, simulation).model_dump(mode="json") try: input_span = ( runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_INPUT_WRITE)) @@ -385,16 +405,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..5a29b4592 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py @@ -0,0 +1,201 @@ +"""Tests for retrying Stage 12 database work that could not connect.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_executor.stage12_runtime import database_retry +from policyengine_simulation_executor.stage12_runtime.database_retry import ( + retry_when_connection_unavailable, +) +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"), + ) + + +def _failed_statement_error() -> OperationalError: + return OperationalError( + "UPDATE comparison_reports SET status = %(status)s", + {"status": "failed"}, + Exception("server closed the connection unexpectedly"), + ) + + +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, + } + ) + + +def _upper_bound_delay(lower_bound: float, upper_bound: float) -> float: + return upper_bound + + +def test_connection_failures_are_retried_with_growing_bounded_delays() -> None: + clock = ManualClock() + runtime = RetryLogRuntime() + attempts: list[int] = [] + + def claim_after_two_connection_failures() -> str: + attempts.append(len(attempts) + 1) + if len(attempts) <= 2: + raise _connection_limit_error() + return "claimed" + + result = retry_when_connection_unavailable( + claim_after_two_connection_failures, + operation_name="coordinator_claim", + runtime=runtime, + sleep=clock.sleep, + monotonic=clock.monotonic, + choose_delay=_upper_bound_delay, + ) + + assert result == "claimed" + assert attempts == [1, 2, 3] + assert clock.sleeps == [1.0, 2.0] + assert [record["attributes"] for record in runtime.records] == [ + { + "stage12_database_operation": "coordinator_claim", + "stage12_database_failed_attempt": 1, + "stage12_database_retry_delay_seconds": 1.0, + }, + { + "stage12_database_operation": "coordinator_claim", + "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() -> str: + attempts.append(1) + if len(attempts) < 3: + raise _connection_limit_error() + return "connected" + + retry_when_connection_unavailable( + connect_on_third_attempt, + operation_name="simulation_claim", + 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] + + +@pytest.mark.parametrize( + "error_factory", + [_failed_statement_error, lambda: ValueError("invalid comparison record")], +) +def test_statement_and_application_errors_are_not_retried(error_factory) -> None: + clock = ManualClock() + attempts: list[int] = [] + + def fail_without_connection_problem() -> None: + attempts.append(1) + raise error_factory() + + with pytest.raises(type(error_factory())): + retry_when_connection_unavailable( + fail_without_connection_problem, + operation_name="coordinator_claim", + 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"): + retry_when_connection_unavailable( + never_connect, + operation_name="coordinator_claim", + 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() + attempts: list[int] = [] + + def connect_on_second_attempt() -> str: + attempts.append(1) + if len(attempts) == 1: + raise _connection_limit_error() + return "connected" + + assert ( + retry_when_connection_unavailable( + connect_on_second_attempt, + operation_name="coordinator_claim", + runtime=FailingRuntime(), + sleep=clock.sleep, + monotonic=clock.monotonic, + ) + == "connected" + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index f989c3a90..d7d8e52e9 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -69,6 +69,7 @@ SimulationCalculation, _build_spm_result, build_aggregate_report, + database_retry, run_single_simulation, simulation_input_sha256, ) @@ -81,6 +82,7 @@ from policyengine_simulation_executor.uk_local_authority_metadata import ( detect_uk_local_authority_metadata, ) +from sqlalchemy.exc import OperationalError NOW = datetime(2026, 9, 15, tzinfo=UTC) EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") @@ -702,6 +704,108 @@ 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_retries_a_claim_that_could_not_connect(monkeypatch) -> None: + monkeypatch.setattr(database_retry, "INITIAL_RETRY_DELAY_SECONDS", 0.0) + + class BusyPoolerStore(FakeStore): + def __init__(self): + super().__init__() + self.read_attempts = 0 + + def get_simulation(self, simulation_execution_id): + self.read_attempts += 1 + if self.read_attempts == 1: + raise _connection_limit_error() + return super().get_simulation(simulation_execution_id) + + store = BusyPoolerStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + + result = run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: _frames(), + ) + + assert store.read_attempts == 2 + assert result["role"] == "baseline" + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +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 +1517,156 @@ def create_or_resolve_simulation(self, record): assert store.parent.error_summary == "RuntimeError" +def test_coordinator_retries_a_claim_that_could_not_connect(monkeypatch) -> None: + monkeypatch.setattr(database_retry, "INITIAL_RETRY_DELAY_SECONDS", 0.0) + + class BusyPoolerStore(FakeStore): + def __init__(self): + super().__init__() + self.claim_attempts = 0 + + def create_or_resolve_report(self, record): + self.claim_attempts += 1 + if self.claim_attempts <= 2: + raise _connection_limit_error() + return super().create_or_resolve_report(record) + + store = BusyPoolerStore() + artifacts = FakeArtifacts() + invoker = ConcurrentInvoker(artifacts) + + 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"}, + ) + + assert store.claim_attempts == 3 + assert invoker.events == ["spawn:baseline", "spawn:reform"] + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + + +def test_coordinator_does_not_retry_a_claim_that_failed_after_connecting() -> None: + class RejectingStore(FakeStore): + def __init__(self): + super().__init__() + self.claim_attempts = 0 + + def create_or_resolve_report(self, record): + self.claim_attempts += 1 + raise OperationalError( + "INSERT INTO comparison_reports", + {}, + Exception("server closed the connection unexpectedly"), + ) + + store = RejectingStore() + + with pytest.raises(OperationalError): + 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=ConcurrentInvoker(FakeArtifacts()), + ) + + assert store.claim_attempts == 1 + + +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" + + def test_coordinator_rejects_incompatible_row_identity() -> None: store = FakeStore() artifacts = FakeArtifacts() From 9c3098fd29c42d9e64b1c318343becc2d3037de6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:42:42 +0400 Subject: [PATCH 2/3] Retry only the opening of Stage 12 database connections The claim retry decided whether to retry from the error alone: an OperationalError without a statement. SQLAlchemy raises a failed COMMIT, ROLLBACK or BEGIN the same way, and a COMMIT can fail after the database has applied it. Re-running the coordinator's claim then found its own running row and returned a deduplicated success without dispatching the simulations, leaving the report running. Each persistence store operation now opens its connection through one connection opener, which receives only the connect step. Statements, COMMIT and ROLLBACK run after the opener returns, so no opener can retry them. The coordinator and worker build their store with an opener that retries a refused connection with the same bounded, jittered backoff (1 s doubling to 15 s, 60 s budget) and WARNING telemetry naming the store operation. Every coordinator and worker store operation therefore retries a refused connection, not only the claims. Under NullPool each operation opens its own pooler session, so child-state creation, dispatch bookkeeping and the later updates could each fail a report on one refusal. The Cloud Run entrypoint keeps the default opener and does not retry, as before. Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 2 +- .../__init__.py | 4 +- .../policyengine_stage12_persistence/store.py | 74 ++++-- .../tests/test_store.py | 175 ++++++++++--- .../stage12_runtime/coordination.py | 16 +- .../stage12_runtime/database_retry.py | 37 ++- .../stage12_runtime/dependencies.py | 15 +- .../stage12_runtime/simulation.py | 46 ++-- .../tests/test_stage12_database_retry.py | 108 ++++---- .../tests/test_stage12_runtime.py | 230 ++++++++++++++---- 10 files changed, 489 insertions(+), 218 deletions(-) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index d65885efd..7c708f43c 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -10,4 +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 Stage 12 claims that could not open a database connection, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot replace the original error + - 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 d85893ca1..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,9 +1,9 @@ """Direct DML for temporary Stage 12 comparison state.""" from policyengine_stage12_persistence.store import ( + ConnectionOpener, Stage12PersistenceStore, create_stage12_engine, - is_connection_unavailable_error, normalized_database_url, ) from policyengine_stage12_persistence.validation import ( @@ -12,10 +12,10 @@ ) __all__ = [ + "ConnectionOpener", "Stage12PersistenceStore", "Stage12RuntimeAccessError", "create_stage12_engine", - "is_connection_unavailable_error", "normalized_database_url", "validate_runtime_database", ] 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 2adcbd1fd..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,8 +19,7 @@ ) from sqlalchemy import Engine, create_engine, select, update from sqlalchemy.dialects.postgresql import insert -from sqlalchemy.engine import URL, make_url -from sqlalchemy.exc import OperationalError +from sqlalchemy.engine import URL, Connection, make_url from sqlalchemy.pool import NullPool from policyengine_stage12_persistence.tables import ( @@ -125,17 +126,31 @@ def create_stage12_engine(database_url: str) -> Engine: ) -def is_connection_unavailable_error(error: BaseException) -> bool: - """Return whether a database connection could not be opened. +class ConnectionOpener(Protocol): + """Open the database connection that one store operation uses. - SQLAlchemy wraps a DBAPI failure raised while opening a connection without - a statement. The database therefore received no work, so the operation can - be retried without risking a duplicate write. Session-pooler connection - limits (``EMAXCONNSESSION``) and connect timeouts fail this way; a failed - statement does not. + ``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. """ - return isinstance(error, OperationalError) and error.statement is None + 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: @@ -171,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, @@ -189,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( @@ -218,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( @@ -244,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) @@ -253,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) @@ -295,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) @@ -346,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( @@ -395,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( @@ -417,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) @@ -479,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 54fa75746..9df7938b2 100644 --- a/libs/policyengine-stage12-persistence/tests/test_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_store.py @@ -13,10 +13,7 @@ ComparisonSimulationRecord, SimulationRole, ) -from policyengine_stage12_persistence import ( - Stage12PersistenceStore, - is_connection_unavailable_error, -) +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 @@ -66,33 +63,6 @@ def local_engine(_url, **options): engine.dispose() -def test_connection_failures_are_distinguished_from_failed_statements( - tmp_path, -) -> None: - unreachable_database = tmp_path / "missing-directory" / "stage12.db" - unreachable_engine = create_engine(f"sqlite:///{unreachable_database}") - reachable_engine = create_engine("sqlite://") - try: - with ( - pytest.raises(OperationalError) as connection_failure, - unreachable_engine.begin(), - ): - pass - with ( - pytest.raises(OperationalError) as statement_failure, - reachable_engine.begin() as connection, - ): - connection.exec_driver_sql("SELECT * FROM missing_table") - finally: - unreachable_engine.dispose() - reachable_engine.dispose() - - # No statement reached the database, so the operation is safe to retry. - assert is_connection_unavailable_error(connection_failure.value) - assert not is_connection_unavailable_error(statement_failure.value) - assert not is_connection_unavailable_error(RuntimeError("not a database error")) - - def _report() -> ComparisonReportRecord: return ComparisonReportRecord( evaluation_id=REPORT_ID, @@ -164,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) @@ -176,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 5ce418da5..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 @@ -44,7 +44,6 @@ validate_uk_local_authority_metadata, ) from .comparison import compare_completed_report -from .database_retry import retry_when_connection_unavailable from .dependencies import ( ChildCall, ChildInvoker, @@ -243,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 = ( @@ -257,15 +256,10 @@ def coordinate_report( else nullcontext() ) with claim_span: - parent, should_execute = retry_when_connection_unavailable( - partial( - _claim_parent, - submitted=submitted_parent, - coordinator_invocation_id=coordinator_invocation_id, - store=persistence, - ), - operation_name="coordinator_claim", - runtime=runtime, + parent, should_execute = _claim_parent( + submitted=submitted_parent, + coordinator_invocation_id=coordinator_invocation_id, + store=persistence, ) if not should_execute: return { 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 index 471b29997..a432cf3fe 100644 --- 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 @@ -1,4 +1,4 @@ -"""Bounded retry for Stage 12 database work that could not open a connection.""" +"""Bounded retry for Stage 12 database connections that could not be opened.""" from __future__ import annotations @@ -7,39 +7,40 @@ import time from collections.abc import Callable from contextlib import suppress -from typing import TypeVar from policyengine_observability import ObservabilityRuntime -from policyengine_stage12_persistence import is_connection_unavailable_error +from sqlalchemy.engine import Connection +from sqlalchemy.exc import OperationalError logger = logging.getLogger(__name__) -OperationResult = TypeVar("OperationResult") - # Coordinators and workers start in bursts against a session pooler with a -# fixed client limit. A short, jittered wait lets concurrent starts spread out -# instead of failing a report before any simulation work has begun. +# 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 retry_when_connection_unavailable( - operation: Callable[[], OperationResult], +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, -) -> OperationResult: - """Run ``operation``, retrying only failures to open a database connection. +) -> Connection: + """Open a Stage 12 database connection, retrying while it is refused. - The database received no statement when a connection could not be opened, - so repeating the operation cannot duplicate a write. Every other error, - including a failed statement, propagates immediately. Retries stop once - the retry budget is spent and the final connection error propagates. + 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() @@ -47,10 +48,8 @@ def retry_when_connection_unavailable( delay_ceiling = INITIAL_RETRY_DELAY_SECONDS while True: try: - return operation() - except Exception as error: - if not is_connection_unavailable_error(error): - raise + return connect() + except OperationalError as error: remaining_budget = CONNECTION_RETRY_BUDGET_SECONDS - (monotonic() - started) if remaining_budget <= 0: raise 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/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 8fe0e8b0e..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 @@ -9,7 +9,6 @@ from functools import partial from hashlib import sha256 from typing import Any -from uuid import UUID import pandas as pd from policyengine_observability import ObservabilityRuntime @@ -39,7 +38,6 @@ detect_uk_local_authority_metadata, ) -from .database_retry import retry_when_connection_unavailable 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 @@ -292,29 +290,6 @@ def descriptor_from_record( ) -def _claim_simulation( - persistence: ComparisonStore, - simulation_execution_id: UUID, -) -> tuple[ComparisonSimulationRecord, ComparisonSimulationRecord | None]: - """Return the stored child and, unless it already succeeded, mark it running.""" - - child = persistence.get_simulation(simulation_execution_id) - if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: - return child, None - started = datetime.now(UTC) - running = child.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "started_at": child.started_at or started, - "updated_at": started, - "error_code": None, - "error_summary": None, - } - ) - persistence.replace_simulation(running) - return child, running - - def run_single_simulation( payload: object, context_payload: object, @@ -331,15 +306,22 @@ 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, running = retry_when_connection_unavailable( - partial(_claim_simulation, persistence, simulation.simulation_execution_id), - operation_name="simulation_claim", - runtime=runtime, - ) - if running is None: + child = persistence.get_simulation(simulation.simulation_execution_id) + if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: return descriptor_from_record(child, simulation).model_dump(mode="json") + started = datetime.now(UTC) + running = child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "started_at": child.started_at or started, + "updated_at": started, + "error_code": None, + "error_summary": None, + } + ) + persistence.replace_simulation(running) try: input_span = ( runtime.span(STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_INPUT_WRITE)) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py b/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py index 5a29b4592..ccfe6d40b 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_database_retry.py @@ -1,11 +1,13 @@ -"""Tests for retrying Stage 12 database work that could not connect.""" +"""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 ( - retry_when_connection_unavailable, + open_connection_with_retry, ) from sqlalchemy.exc import OperationalError @@ -19,14 +21,6 @@ def _connection_limit_error() -> OperationalError: ) -def _failed_statement_error() -> OperationalError: - return OperationalError( - "UPDATE comparison_reports SET status = %(status)s", - {"status": "failed"}, - Exception("server closed the connection unexpectedly"), - ) - - class ManualClock: def __init__(self) -> None: self.now = 0.0 @@ -55,41 +49,46 @@ def log(self, message, *, severity="INFO", attributes=None, error=None): ) +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_connection_failures_are_retried_with_growing_bounded_delays() -> None: +def test_refused_connections_are_retried_with_growing_bounded_delays() -> None: clock = ManualClock() runtime = RetryLogRuntime() + opened_connection = OpenedConnection() attempts: list[int] = [] - def claim_after_two_connection_failures() -> str: + def connect_after_two_refusals() -> OpenedConnection: attempts.append(len(attempts) + 1) if len(attempts) <= 2: raise _connection_limit_error() - return "claimed" + return opened_connection - result = retry_when_connection_unavailable( - claim_after_two_connection_failures, - operation_name="coordinator_claim", + 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 == "claimed" + 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": "coordinator_claim", + "stage12_database_operation": "create_or_resolve_simulation", "stage12_database_failed_attempt": 1, "stage12_database_retry_delay_seconds": 1.0, }, { - "stage12_database_operation": "coordinator_claim", + "stage12_database_operation": "create_or_resolve_simulation", "stage12_database_failed_attempt": 2, "stage12_database_retry_delay_seconds": 2.0, }, @@ -110,15 +109,15 @@ def lowest_delay(lower_bound: float, upper_bound: float) -> float: attempts: list[int] = [] - def connect_on_third_attempt() -> str: + def connect_on_third_attempt() -> OpenedConnection: attempts.append(1) if len(attempts) < 3: raise _connection_limit_error() - return "connected" + return OpenedConnection() - retry_when_connection_unavailable( + open_connection_with_retry( connect_on_third_attempt, - operation_name="simulation_claim", + operation_name="get_simulation", sleep=clock.sleep, monotonic=clock.monotonic, choose_delay=lowest_delay, @@ -128,22 +127,18 @@ def connect_on_third_attempt() -> str: assert clock.sleeps == [0.5, 1.0] -@pytest.mark.parametrize( - "error_factory", - [_failed_statement_error, lambda: ValueError("invalid comparison record")], -) -def test_statement_and_application_errors_are_not_retried(error_factory) -> None: +def test_errors_other_than_a_refused_connection_are_not_retried() -> None: clock = ManualClock() attempts: list[int] = [] - def fail_without_connection_problem() -> None: + def fail_with_invalid_connection_options() -> None: attempts.append(1) - raise error_factory() + raise ValueError("invalid connection option") - with pytest.raises(type(error_factory())): - retry_when_connection_unavailable( - fail_without_connection_problem, - operation_name="coordinator_claim", + 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, ) @@ -161,9 +156,9 @@ def never_connect() -> None: raise _connection_limit_error() with pytest.raises(OperationalError, match="EMAXCONNSESSION"): - retry_when_connection_unavailable( + open_connection_with_retry( never_connect, - operation_name="coordinator_claim", + operation_name="create_or_resolve_report", sleep=clock.sleep, monotonic=clock.monotonic, choose_delay=_upper_bound_delay, @@ -181,21 +176,52 @@ def log(self, *_args, **_kwargs): raise RuntimeError("telemetry exporter unavailable") clock = ManualClock() + opened_connection = OpenedConnection() attempts: list[int] = [] - def connect_on_second_attempt() -> str: + def connect_on_second_attempt() -> OpenedConnection: attempts.append(1) if len(attempts) == 1: raise _connection_limit_error() - return "connected" + return opened_connection assert ( - retry_when_connection_unavailable( + open_connection_with_retry( connect_on_second_attempt, - operation_name="coordinator_claim", + operation_name="create_or_resolve_report", runtime=FailingRuntime(), sleep=clock.sleep, monotonic=clock.monotonic, ) - == "connected" + 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_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index d7d8e52e9..0154d92c0 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -76,12 +76,22 @@ 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) @@ -726,37 +736,30 @@ def log(self, message, *, severity="INFO", attributes=None, error=None): ) -def test_single_worker_retries_a_claim_that_could_not_connect(monkeypatch) -> None: - monkeypatch.setattr(database_retry, "INITIAL_RETRY_DELAY_SECONDS", 0.0) - - class BusyPoolerStore(FakeStore): - def __init__(self): - super().__init__() - self.read_attempts = 0 - - def get_simulation(self, simulation_execution_id): - self.read_attempts += 1 - if self.read_attempts == 1: - raise _connection_limit_error() - return super().get_simulation(simulation_execution_id) - - store = BusyPoolerStore() +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 = [] - result = run_single_simulation( + 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", - store=store, artifacts=FakeArtifacts(), calculator=lambda _: _frames(), + runtime=runtime, ) - assert store.read_attempts == 2 - assert result["role"] == "baseline" - child = store.children[simulation.simulation_execution_id] - assert child.status is ComparisonRunLifecycleStatus.SUCCEEDED + # 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: @@ -1517,23 +1520,19 @@ def create_or_resolve_simulation(self, record): assert store.parent.error_summary == "RuntimeError" -def test_coordinator_retries_a_claim_that_could_not_connect(monkeypatch) -> None: - monkeypatch.setattr(database_retry, "INITIAL_RETRY_DELAY_SECONDS", 0.0) - - class BusyPoolerStore(FakeStore): - def __init__(self): - super().__init__() - self.claim_attempts = 0 +def test_coordinator_builds_its_store_with_the_coordinator_runtime( + monkeypatch, +) -> None: + store = FakeStore() + store_runtimes = [] - def create_or_resolve_report(self, record): - self.claim_attempts += 1 - if self.claim_attempts <= 2: - raise _connection_limit_error() - return super().create_or_resolve_report(record) + def record_store_runtime(runtime): + store_runtimes.append(runtime) + return store - store = BusyPoolerStore() + monkeypatch.setattr(coordination_module, "runtime_store", record_store_runtime) artifacts = FakeArtifacts() - invoker = ConcurrentInvoker(artifacts) + runtime = RecordingRuntime() coordinate_report( _report().model_dump(mode="json"), @@ -1541,34 +1540,45 @@ def create_or_resolve_report(self, record): _parent().model_dump(mode="json"), application_name=_context().modal_application, coordinator_invocation_id="coordinator-1", - store=store, artifacts=artifacts, - invoker=invoker, + invoker=ConcurrentInvoker(artifacts), aggregator=lambda **_: {"result": "complete"}, + runtime=runtime, ) - assert store.claim_attempts == 3 - assert invoker.events == ["spawn:baseline", "spawn:reform"] + # 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_retry_a_claim_that_failed_after_connecting() -> None: - class RejectingStore(FakeStore): +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 - raise OperationalError( - "INSERT INTO comparison_reports", - {}, - Exception("server closed the connection unexpectedly"), - ) + claim = super().create_or_resolve_report(record) + if self.claim_attempts == 1: + raise lost_commit_reply + return claim - store = RejectingStore() + store = LostCommitReplyStore() + artifacts = FakeArtifacts() + invoker = ConcurrentInvoker(artifacts) - with pytest.raises(OperationalError): + with pytest.raises(OperationalError) as error: coordinate_report( _report().model_dump(mode="json"), _context().model_dump(mode="json"), @@ -1576,11 +1586,131 @@ def create_or_resolve_report(self, record): application_name=_context().modal_application, coordinator_invocation_id="coordinator-1", store=store, - artifacts=FakeArtifacts(), - invoker=ConcurrentInvoker(FakeArtifacts()), + 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: From 0d57ee4ee24cada28d191f40ec63ac186e1c5b17 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:42:42 +0400 Subject: [PATCH 3/3] Test Stage 12 failure reports without a runtime and on dispatch failure Cover the logger fallbacks used when no observability runtime is passed, and check that a failed child dispatch is reported once. Correct the failure-record docstring: the report and simulation handlers raise their stable RuntimeError, so the original cause is visible only in telemetry. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/failure_recording.py | 6 ++- .../tests/test_stage12_failure_recording.py | 53 ++++++++++++++++++ .../tests/test_stage12_runtime.py | 54 +++++++++++++++++++ 3 files changed, 111 insertions(+), 2 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py 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 index 469c9b88b..342c315f4 100644 --- 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 @@ -41,8 +41,10 @@ def persist_failure_record( ) -> None: """Attempt a failure write without letting its error replace the original. - The caller re-raises the original failure after this returns, whether or - not the record was written. + 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: 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 0154d92c0..c7c9431a2 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -1797,6 +1797,60 @@ def get_simulation(self, simulation_execution_id): 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()