From 587c559105f6db2fa54e200a39ecdc13fd53bd87 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 27 Oct 2026 17:19:56 +0400 Subject: [PATCH] Split the Stage 12 runtime tests into smaller files tests/test_stage12_runtime.py (4,709 lines) was over the file-length limit for test code. Its tests move, unchanged, into thirteen files by behaviour: - test_stage12_runtime.py: loading a simulation's dataset and region - test_stage12_dataset_provenance.py: the dataset a calculation loads - test_stage12_single_worker.py, test_stage12_single_worker_failures.py - test_stage12_aggregate_reports.py - test_stage12_coordinator.py, test_stage12_coordinator_failures.py, test_stage12_coordinator_failure_records.py - test_stage12_runtime_store_retries.py - test_stage12_child_cancellation.py, test_stage12_child_state_errors.py - test_stage12_wait_window.py The helpers several files use move to three support modules: stage12_runtime_support.py (inputs, records, runtimes and failure checks), stage12_fake_stores.py (the in-memory stores) and stage12_child_call_support.py (fake child calls and invokers). Their private names are now public, under names that no test uses for a local variable: for example _report is report_execution_input, _simulation is simulation_execution_input, _context is invocation_context and _parent is parent_record. The seven other Stage 12 test modules that imported them from test_stage12_runtime.py now import them from the support modules, and coordinate_report_impl from the package itself. All 1,895 tests keep their names and bodies, apart from the renamed helpers. Pytest names two parametrized cases of test_installed_bundle_accepts_the_dataset_a_region_loads after the bundle helper, so they are now [us-bundle_provenance] and [state/ut-bundle_provenance]. The file-length baseline drops the file (count 8 to 7). Co-Authored-By: Claude Opus 5.5 --- changelog_entry.yaml | 1 + docs/generated/rules.md | 2 +- .../tests/stage12_child_call_support.py | 409 ++ .../tests/stage12_fake_stores.py | 207 + .../tests/stage12_runtime_support.py | 369 ++ .../tests/test_stage12_aggregate_reports.py | 452 ++ .../tests/test_stage12_cache_runtime.py | 12 +- .../tests/test_stage12_child_cancellation.py | 406 ++ .../tests/test_stage12_child_state_errors.py | 400 ++ .../tests/test_stage12_coordination.py | 83 +- .../tests/test_stage12_coordinator.py | 413 ++ ...est_stage12_coordinator_failure_records.py | 402 ++ .../test_stage12_coordinator_failures.py | 527 ++ .../tests/test_stage12_dataset_provenance.py | 335 ++ .../tests/test_stage12_qualification.py | 22 +- .../tests/test_stage12_runtime.py | 4436 +---------------- .../test_stage12_runtime_store_retries.py | 135 + .../tests/test_stage12_segment_calculation.py | 26 +- .../tests/test_stage12_segment_fanout.py | 4 +- .../test_stage12_segmentation_foundation.py | 4 +- .../tests/test_stage12_simulation_worker.py | 32 +- .../tests/test_stage12_single_worker.py | 231 + .../test_stage12_single_worker_failures.py | 335 ++ .../tests/test_stage12_wait_window.py | 149 + rules/baselines/python-file-length.txt | 1 - rules/registry.toml | 2 +- 26 files changed, 4885 insertions(+), 4510 deletions(-) create mode 100644 projects/policyengine-simulation-executor/tests/stage12_child_call_support.py create mode 100644 projects/policyengine-simulation-executor/tests/stage12_fake_stores.py create mode 100644 projects/policyengine-simulation-executor/tests/stage12_runtime_support.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_aggregate_reports.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_child_cancellation.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_child_state_errors.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_coordinator.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_coordinator_failure_records.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_coordinator_failures.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_dataset_provenance.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_runtime_store_retries.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_single_worker.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_single_worker_failures.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_wait_window.py diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 43e2c2382..5fc93bea0 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -16,6 +16,7 @@ - Ran Stage 12 US national simulations as 20 region-group calls each, and built an independent, immutable Stage 12 cache of three annual US datasets and 60 current-law baselines before each deployment, publishing it only after a recorded determinism check and recording its release in the v2 version manifest under a new key, while the existing key keeps a copy that earlier readers accept - Made the budget refuse a household count outside its country's plausible range (above 200 million in the US or 60 million in the UK, negative or not finite), raising ImplausibleEntityCountError instead of publishing a count like #767's changed: + - Split the 4,709-line Stage 12 runtime test file into thirteen test files by behaviour (dataset loading, provenance, the single worker and its failures, aggregates, the coordinator and its failures, runtime store retries, failure records, cancellation, child state errors and the wait window) and three support modules for the shared builders, fake stores and fake child calls, whose helpers are now public; the other Stage 12 tests import them from those modules, every test keeps its name and body, and two parametrized cases are named after the renamed bundle helper - Split Stage 12 report coordination (stage12_runtime/coordination.py, 1,091 lines) into modules for the report coordinator class, its child simulations' records, failures and cancellation, and the parent report's checks and claim, with coordinate_report and every public name still in coordination.py and child failures still logged under its logger; and turned the contract's stage12_execution.py (929 lines) into a stage12_execution package of modules for the shared field types and simulation input, report inputs and output plans, result artifacts, and comparison-run records, which re-exports every name from the same import path, with unchanged JSON schemas; models subclass the base models through the package path ruff.toml lists, so no ruff setting changes - Split the gateway's endpoints.py (1,454 lines) into modules for the shared router and request helpers, route resolution from the active routing state, Modal reads of job metadata and routing, bundle payloads, and the job submission, job status and version routes, which register their routes in the same order, so the OpenAPI document is unchanged; endpoints.py re-exports every public name, the routes keep logging under the policyengine_simulation_gateway.endpoints logger, and the private helpers the new modules share are now public; its 1,801-line test file is split into five test files, with the shared expected-response helpers in fixtures/gateway_responses.py - Split the routing publisher (src/modal/utils/update_version_registry.py, 1,584 lines) into a src/modal/routing_publisher package of modules for the bundle manifest, the routing state, routing plans and their publication record, the check that only serialized workflows write routing, Modal Dict access, plan publication, the legacy seed and the command-line options, and its 2,728-line test file into a tests/version_registry directory of nine test files with shared fixtures; update_version_registry.py keeps the command line, still runs with python -m src.modal.utils.update_version_registry and re-exports every public name; behaviour and output are unchanged, and the private helpers the new modules share are now public diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 78e27a08c..71b7fbb04 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -66,7 +66,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt` - `sim.microdata.weight-stripping-confined`: 0 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` -- `sim.quality.python-file-length`: 8 remaining in `rules/baselines/python-file-length.txt` +- `sim.quality.python-file-length`: 7 remaining in `rules/baselines/python-file-length.txt` ### Exceptions and review dates diff --git a/projects/policyengine-simulation-executor/tests/stage12_child_call_support.py b/projects/policyengine-simulation-executor/tests/stage12_child_call_support.py new file mode 100644 index 000000000..097888f11 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/stage12_child_call_support.py @@ -0,0 +1,409 @@ +"""Fake child calls and invokers, and checks on them, for Stage 12 coordinator tests.""" + +from __future__ import annotations + +import time +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + PlannedSimulationExecutionInput, + ReportAggregate, + ReportExecutionInput, + RowIdentity, + SimulationExecutionInput, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + coordination as coordination_module, +) +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames +from stage12_runtime_support import ( + BASELINE_ID, + EVALUATION_ID, + REFORM_ID, + aggregate_payload, + coordinate_report, + invocation_context, + parent_record, + planned_simulation, + report_execution_input, + simulation_execution_input, + summary_failure_cause, +) + +if TYPE_CHECKING: + from uuid import UUID + + +class FutureCall: + def __init__(self, object_id, future): + self.object_id = object_id + self.future = future + + def get(self, *, timeout=None): + return self.future.result(timeout=timeout) + + def cancel(self): + self.future.cancel() + + +class ImmediateCall: + def __init__(self, object_id, result, observed_timeouts): + self.object_id = object_id + self.result = result + self.observed_timeouts = observed_timeouts + self.cancel_count = 0 + + def get(self, *, timeout=None): + self.observed_timeouts.append(timeout) + return self.result + + def cancel(self): + self.cancel_count += 1 + + +@dataclass(frozen=True) +class InjectedChildFaults: + """Faults a ConcurrentInvoker injects into the child simulations it runs.""" + + fail_role: str | None = None + incompatible: bool = False + wrong_plan_digest: bool = False + missing_column: bool = False + + +class ConcurrentInvoker: + def __init__( + self, + artifacts, + *, + faults: InjectedChildFaults | None = None, + production_result=None, + ): + faults = faults or InjectedChildFaults() + self.artifacts = artifacts + self.fail_role = faults.fail_role + self.incompatible = faults.incompatible + self.wrong_plan_digest = faults.wrong_plan_digest + self.missing_column = faults.missing_column + self.executor = ThreadPoolExecutor(max_workers=2) + self.intervals = {} + self.events = [] + self.environments = [] + self.production_result = production_result + self.production_wait_timeouts = [] + self.output_plans = [] + + def spawn(self, request): + parsed = PlannedSimulationExecutionInput.model_validate(request.simulation) + role = parsed.role.value + self.events.append(f"spawn:{role}") + self.environments.append(request.environment) + self.output_plans.append(parsed.output_plan) + + return self._submit(parsed, role) + + def restore(self, invocation_id): + if invocation_id == "production-job-1": + self.events.append("restore:production") + return ImmediateCall( + invocation_id, + self.production_result, + self.production_wait_timeouts, + ) + role = invocation_id.removeprefix("call-") + parsed = planned_simulation(SimulationRole(role)) + self.events.append(f"restore:{role}") + return self._submit(parsed, role) + + def _submit(self, parsed, role): + + def run(): + started = time.monotonic() + time.sleep(0.05) + if role == self.fail_role: + msg = "child calculation failed" + raise RuntimeError(msg) + frames = simulation_frames(100.0 if role == "baseline" else 120.0) + if self.missing_column and role == "reform": + frames["household"] = frames["household"].drop( + columns=["household_net_income"] + ) + descriptor = self.artifacts.add_simulation(parsed, frames) + if self.incompatible and role == "reform": + descriptor = descriptor.model_copy( + update={ + "row_identity": RowIdentity( + identifier_columns=("household.household_id",), + row_count=1, + identity_sha256="f" * 64, + ) + } + ) + if self.wrong_plan_digest and role == "reform": + descriptor = descriptor.model_copy( + update={"output_plan_sha256": "f" * 64} + ) + self.intervals[role] = (started, time.monotonic()) + return descriptor.model_dump(mode="json") + + return FutureCall(f"call-{role}", self.executor.submit(run)) + + +class RefusingDispatchInvoker: + def __init__(self, dispatch_error): + self.dispatch_error = dispatch_error + self.events = [] + + def spawn(self, request): + role = PlannedSimulationExecutionInput.model_validate( + request.simulation + ).role.value + self.events.append(f"spawn:{role}") + raise self.dispatch_error + + def restore(self, invocation_id): + msg = f"unexpected restored invocation {invocation_id}" + raise AssertionError(msg) + + +class TimeoutCall: + def __init__(self, object_id): + self.object_id = object_id + self.cancelled = 0 + + def get(self, *, timeout=None): + msg = "contains sensitive timeout details" + raise TimeoutError(msg) + + def cancel(self): + self.cancelled += 1 + + +class TimeoutInvoker: + def __init__(self): + self.events = [] + + def spawn(self, request): + role = PlannedSimulationExecutionInput.model_validate( + request.simulation + ).role.value + self.events.append(f"spawn:{role}") + return TimeoutCall(f"call-{role}") + + def restore(self, invocation_id): + msg = f"unexpected restored invocation {invocation_id}" + raise AssertionError(msg) + + +class StartupWriteError(Exception): + """A coordinator persistence failure after a child may have started.""" + + +@dataclass(frozen=True) +class ChildCallScript: + """What a scripted child call does when probed or cancelled. + + It returns ``result`` or raises ``wait_error`` once probed (``during_wait`` + runs first), and raises ``cancel_error`` when cancelled; with none set, it + keeps running. + """ + + result: object = None + wait_error: Exception | None = None + cancel_error: Exception | None = None + during_wait: object = None + + +class ScriptedChildCall: + """A started child that keeps running unless the test scripts an outcome.""" + + def __init__(self, object_id, *, clock=None, script=None): + script = script or ChildCallScript() + self.object_id = object_id + self.clock = clock + self.result = script.result + self.wait_error = script.wait_error + self.cancel_error = script.cancel_error + self.during_wait = script.during_wait + self.probed_at = [] + self.cancel_count = 0 + + def get(self, *, timeout=None): + self.probed_at.append( + self.clock.monotonic() if self.clock is not None else None + ) + if self.during_wait is not None: + self.during_wait() + if self.wait_error is not None: + raise self.wait_error + if self.result is not None: + return self.result + msg = "still running" + raise TimeoutError(msg) + + def cancel(self): + self.cancel_count += 1 + if self.cancel_error is not None: + raise self.cancel_error + + +class ScriptedChildInvoker: + def __init__( + self, + *, + clock=None, + results=None, + wait_errors=None, + cancel_errors=None, + during_wait=None, + ): + self.clock = clock + self.results = results or {} + self.wait_errors = wait_errors or {} + self.cancel_errors = cancel_errors or {} + self.during_wait = during_wait or {} + self.events = [] + self.calls = {} + + def spawn(self, request): + role = PlannedSimulationExecutionInput.model_validate( + request.simulation + ).role.value + self.events.append(f"spawn:{role}") + return self._call(role) + + def restore(self, invocation_id): + role = invocation_id.removeprefix("call-") + self.events.append(f"restore:{role}") + return self._call(role) + + def _call(self, role): + call = ScriptedChildCall( + f"call-{role}", + clock=self.clock, + script=ChildCallScript( + result=self.results.get(role), + wait_error=self.wait_errors.get(role), + cancel_error=self.cancel_errors.get(role), + during_wait=self.during_wait.get(role), + ), + ) + self.calls[role] = call + return call + + +class PlaceholderWriteFailingStore(FakeStore): + def __init__(self, failing_simulation_id): + super().__init__() + self.failing_simulation_id = failing_simulation_id + + def replace_simulation(self, record): + # The coordinator's only plain replacement is the dispatch placeholder. + if record.simulation_execution_id == self.failing_simulation_id: + msg = "placeholder write contains sensitive details" + raise StartupWriteError(msg) + return super().replace_simulation(record) + + +class AttachFailingStore(FakeStore): + def __init__(self, failing_simulation_id): + super().__init__() + self.failing_simulation_id = failing_simulation_id + + def attach_simulation_invocation(self, simulation_execution_id, **values): + if simulation_execution_id == self.failing_simulation_id: + msg = "attach contains sensitive details" + raise StartupWriteError(msg) + return super().attach_simulation_invocation(simulation_execution_id, **values) + + +class CancellationStateFailingStore(FakeStore): + def replace_simulation_if_status(self, record, *, expected_status): + if record.error_code == "simulation_cancelled": + msg = "database contains sensitive details" + raise ConnectionError(msg) + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + +def role_simulation_id(role: str) -> UUID: + return BASELINE_ID if role == "baseline" else REFORM_ID + + +def report_with_segmentation(*, segmented: bool) -> ReportExecutionInput: + if segmented: + return report_execution_input() + + def unsegmented(role: SimulationRole) -> SimulationExecutionInput: + return simulation_execution_input(role).model_copy( + update={"options": {"segmented": False}} + ) + + return ReportExecutionInput( + evaluation_id=EVALUATION_ID, + baseline=unsegmented(SimulationRole.BASELINE), + reform=unsegmented(SimulationRole.REFORM), + requested_aggregates=tuple(ReportAggregate), + ) + + +def coordinate_with(report, store, invoker, **coordinator_options): + coordinator_options.setdefault("artifacts", FakeArtifacts()) + return coordinate_report( + report.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + invoker=invoker, + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + **coordinator_options, + ) + + +def assert_cancelled(child, *, after: str) -> None: + assert child.status is ComparisonRunLifecycleStatus.FAILED + assert child.error_code == "simulation_cancelled" + assert child.error_summary == f"Cancelled after {after}." + assert child.completed_at is not None + + +def assert_original_coordinator_failure(error, store) -> None: + assert isinstance(error.value.__context__, StartupWriteError) + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_code == "report_coordination_failed" + original_error = error.value.__context__ + assert summary_failure_cause(store.parent.error_summary).startswith( + f"{type(original_error).__name__} " + ) + assert summary_failure_cause(store.parent.error_summary).endswith( + f": {original_error}" + ) + + +def coordination_log_record(caplog, level): + (record,) = [ + record + for record in caplog.records + if record.name == coordination_module.__name__ and record.levelno == level + ] + return record + + +def assert_log_identifies_child(record, *, role: str, error_type: str) -> None: + simulation_execution_id = str(role_simulation_id(role)) + assert record.evaluation_id == str(EVALUATION_ID) + assert record.simulation_execution_id == simulation_execution_id + assert record.simulation_role == role + assert record.error_type == error_type + message = record.getMessage() + assert str(EVALUATION_ID) in message + assert simulation_execution_id in message + assert "sensitive" not in message + assert record.exc_info is None diff --git a/projects/policyengine-simulation-executor/tests/stage12_fake_stores.py b/projects/policyengine-simulation-executor/tests/stage12_fake_stores.py new file mode 100644 index 000000000..dd56cdb8d --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/stage12_fake_stores.py @@ -0,0 +1,207 @@ +"""In-memory comparison and artifact stores for the Stage 12 tests.""" + +from __future__ import annotations + +from hashlib import sha256 +from threading import Lock + +import pandas as pd +from fixtures.fake_object_store import InMemoryObjectStore +from policyengine_simulation_contract.stage12_execution import ( + ArtifactMediaType, + ArtifactReference, + ComparisonReportPersistenceResult, + ComparisonSimulationPersistenceResult, + SimulationArtifactDescriptor, + stage12_output_plan_sha256, +) +from policyengine_simulation_executor.stage12_artifacts import ( + Stage12ArtifactStore, + canonical_json_bytes, + serialize_simulation_frames, +) +from stage12_runtime_support import EVALUATION_ID, shared_output_plan + + +class FakeStore: + def __init__(self): + self.parent = None + self.children = {} + self.lock = Lock() + + def create_or_resolve_report(self, record): + with self.lock: + if self.parent is None: + self.parent = record + return ComparisonReportPersistenceResult( + record=record, + created=True, + ) + return ComparisonReportPersistenceResult( + record=self.parent, + created=False, + ) + + def get_report(self, evaluation_id): + assert evaluation_id == EVALUATION_ID + assert self.parent is not None + return self.parent + + def replace_report(self, record): + self.parent = record + return record + + def replace_report_result_comparison(self, record): + self.parent = record + return record + + def create_or_resolve_simulation(self, record): + with self.lock: + existing = self.children.get(record.simulation_execution_id) + if existing is None: + self.children[record.simulation_execution_id] = record + return ComparisonSimulationPersistenceResult( + record=record, created=True + ) + return ComparisonSimulationPersistenceResult( + record=existing, + created=False, + ) + + def get_simulation(self, simulation_execution_id): + return self.children[simulation_execution_id] + + def replace_simulation(self, record): + with self.lock: + self.children[record.simulation_execution_id] = record + return record + + def replace_simulation_if_status(self, record, *, expected_status): + with self.lock: + current = self.children[record.simulation_execution_id] + if current.status is not expected_status: + return current, False + self.children[record.simulation_execution_id] = record + return record, True + + def attach_simulation_invocation( + self, + simulation_execution_id, + *, + expected_placeholder, + modal_invocation_id, + updated_at, + ): + with self.lock: + child = self.children[simulation_execution_id] + assert child.modal_invocation_id == expected_placeholder + child = child.model_copy( + update={ + "modal_invocation_id": modal_invocation_id, + "updated_at": updated_at, + } + ) + self.children[simulation_execution_id] = child + return child + + +class FakeArtifacts: + def __init__(self): + self.payloads = {} + self.aggregate_writes = [] + self.comparison_writes = [] + self.input_writes = [] + # Aggregates go through the real store, so every coordinator test + # enforces the aggregate artifact contract that production enforces. + self.report_store = Stage12ArtifactStore("private", store=InMemoryObjectStore()) + + def write_input(self, *, prefix, simulation): + self.input_writes.append((prefix, simulation.role)) + + def add_simulation(self, simulation, frames=None): + payload, row_identity = serialize_simulation_frames( + frames or simulation_frames() + ) + digest = sha256(payload).hexdigest() + uri = f"gs://private/{simulation.role.value}.parquet" + self.payloads[uri] = payload + return SimulationArtifactDescriptor( + evaluation_id=simulation.evaluation_id, + simulation_execution_id=simulation.simulation_execution_id, + role=simulation.role, + artifact=ArtifactReference( + uri=uri, + media_type=ArtifactMediaType.PARQUET, + content_sha256=digest, + size_bytes=len(payload), + ), + output_plan_sha256=stage12_output_plan_sha256(shared_output_plan()), + row_identity=row_identity, + bundle=simulation.bundle, + ) + + def write_simulation( + self, + *, + prefix, + simulation, + frames, + calculation_provenance=None, + uk_local_authority_metadata=None, + ): + payload, row_identity = serialize_simulation_frames( + frames, + calculation_provenance=calculation_provenance, + uk_local_authority_metadata=uk_local_authority_metadata, + ) + uri = f"gs://private/{simulation.role.value}.parquet" + self.payloads[uri] = payload + descriptor = SimulationArtifactDescriptor( + evaluation_id=simulation.evaluation_id, + simulation_execution_id=simulation.simulation_execution_id, + role=simulation.role, + artifact=ArtifactReference( + uri=uri, + media_type=ArtifactMediaType.PARQUET, + content_sha256=sha256(payload).hexdigest(), + size_bytes=len(payload), + ), + output_plan_sha256=stage12_output_plan_sha256(shared_output_plan()), + row_identity=row_identity, + bundle=simulation.bundle, + ) + return descriptor.model_copy( + update={"calculation_provenance": calculation_provenance} + ) + + def read(self, uri): + if uri in self.payloads: + return self.payloads[uri] + return self.report_store.read(uri) + + def write_aggregate(self, *, prefix, payload): + reference = self.report_store.write_aggregate(prefix=prefix, payload=payload) + self.aggregate_writes.append(payload) + return reference + + def write_comparison(self, *, prefix, payload): + encoded = canonical_json_bytes(payload) + digest = sha256(encoded).hexdigest() + uri = "gs://private/comparison.json" + self.payloads[uri] = encoded + self.comparison_writes.append(payload) + return ArtifactReference( + uri=uri, + media_type=ArtifactMediaType.JSON, + content_sha256=digest, + size_bytes=len(encoded), + ) + + +def simulation_frames(value=100.0): + return { + "household": pd.DataFrame( + {"household_id": [1], "household_net_income": [value]} + ), + "person": pd.DataFrame({"person_id": [1], "household_id": [1], "age": [40]}), + } diff --git a/projects/policyengine-simulation-executor/tests/stage12_runtime_support.py b/projects/policyengine-simulation-executor/tests/stage12_runtime_support.py new file mode 100644 index 000000000..d51d2a64e --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/stage12_runtime_support.py @@ -0,0 +1,369 @@ +"""Inputs, records, runtimes and failure checks shared by the Stage 12 tests.""" + +from __future__ import annotations + +import re +from contextlib import contextmanager, nullcontext +from datetime import UTC, datetime, timedelta +from types import SimpleNamespace +from uuid import UUID + +import pytest +from policyengine_simulation_contract.spm import ( + SPMRuntimeVersions, + SPMSelection, + build_spm_provenance, + validate_spm_calculation_provenance, +) +from policyengine_simulation_contract.stage12_execution import ( + AggregateReportArtifactPayload, + BundleProvenance, + ComparisonReportRecord, + ComparisonRunAggregationStatus, + ComparisonRunLifecycleStatus, + DatasetArtifactMediaType, + DatasetArtifactReference, + DatasetPopulationInput, + DatasetProvenance, + EntityOutputPlan, + GeographySelection, + PlannedSimulationExecutionInput, + ReportAggregate, + ReportExecutionInput, + ReportOutputRequirements, + ResultComparisonStatus, + SimulationExecutionInput, + SimulationRole, + Stage12InvocationContext, + Stage12OutputPlan, +) +from policyengine_simulation_executor.stage12_runtime import ( + CoordinatorDependencies, + CoordinatorInvocation, + simulation_input_sha256, +) +from policyengine_simulation_executor.stage12_runtime import ( + coordinate_report as coordinate_report_impl, +) +from sqlalchemy.exc import OperationalError + +NOW = datetime(2026, 9, 15, tzinfo=UTC) +EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") +BASELINE_ID = UUID("00000000-0000-0000-0000-000000000002") +REFORM_ID = UUID("00000000-0000-0000-0000-000000000003") + + +def spm_calculation_provenance(): + selection = SPMSelection( + forecast_content_sha256="f" * 64, + scenario="official", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ) + receipt = build_spm_provenance( + forecast_id="forecast-1", + forecast_sha256="f" * 64, + selection=selection, + years=("2026",), + runtime_versions=SPMRuntimeVersions.model_validate( + { + "policyengine": "6.2.1", + "policyengine-core": "3.32.10", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", + } + ), + ) + return validate_spm_calculation_provenance(config=selection, receipt=receipt) + + +def bundle_provenance() -> BundleProvenance: + return BundleProvenance( + policyengine_version="5.2.0", + country_package_name="policyengine-us", + country_package_version="1.764.6", + dataset=DatasetProvenance( + identity="populace_us_2024", + uri="hf://policyengine/data/populace_us_2024.h5@revision", + artifact_revision="revision", + data_package_name="populace-data", + data_package_version="0.1.0", + ), + bundle_manifest_sha256="b" * 64, + ) + + +def simulation_execution_input(role: SimulationRole) -> SimulationExecutionInput: + return SimulationExecutionInput( + evaluation_id=EVALUATION_ID, + simulation_execution_id=( + BASELINE_ID if role is SimulationRole.BASELINE else REFORM_ID + ), + role=role, + policy={} if role is SimulationRole.BASELINE else {"reform": {"2026": 1}}, + population=DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=bundle_provenance().dataset.uri, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256="a" * 64, + ) + ), + year=2026, + geography=GeographySelection(country="us", region="us"), + bundle=bundle_provenance(), + ) + + +def shared_output_plan() -> Stage12OutputPlan: + return Stage12OutputPlan( + country="us", + requirements=ReportOutputRequirements( + aggregates=tuple(ReportAggregate), + include_cliff_impacts=False, + labor_supply_response_active=False, + ), + entities=( + EntityOutputPlan( + entity="household", + materialized_variables=("household_id", "household_net_income"), + ), + EntityOutputPlan( + entity="person", + materialized_variables=("age", "household_id", "person_id"), + ), + ), + ) + + +def planned_simulation(role: SimulationRole) -> PlannedSimulationExecutionInput: + return PlannedSimulationExecutionInput.model_validate( + { + **simulation_execution_input(role).model_dump(mode="json"), + "output_plan": shared_output_plan().model_dump(mode="json"), + } + ) + + +def report_execution_input() -> ReportExecutionInput: + return ReportExecutionInput( + evaluation_id=EVALUATION_ID, + baseline=simulation_execution_input(SimulationRole.BASELINE), + reform=simulation_execution_input(SimulationRole.REFORM), + requested_aggregates=tuple(ReportAggregate), + ) + + +def aggregate_payload(result=None) -> dict: + """Return a stand-in aggregate that satisfies the artifact contract. + + Aggregators that must not run return one too, so an unexpected run shows + up as a written aggregate instead of a rejected one. + """ + report = report_execution_input() + return AggregateReportArtifactPayload( + evaluation_id=report.evaluation_id, + requested_aggregates=report.requested_aggregates, + bundle=report.baseline.bundle, + result={} if result is None else result, + ).model_dump(mode="json") + + +def coordinate_report(*payloads, application_name, coordinator_invocation_id, **seams): + seams.setdefault("output_plan_resolver", lambda _: shared_output_plan()) + return coordinate_report_impl( + *payloads, + invocation=CoordinatorInvocation( + application_name=application_name, + invocation_id=coordinator_invocation_id, + ), + dependencies=CoordinatorDependencies(**seams), + ) + + +class RecordingRuntime: + def __init__(self) -> None: + self.operations: list[tuple[str, dict[str, object]]] = [] + self.recorded_exceptions: list[Exception] = [] + + @contextmanager + def operation(self, name, *, attributes=None, **_kwargs): + self.operations.append((name, dict(attributes or {}))) + yield + + def span(self, *_args, **_kwargs): + return nullcontext() + + def capture_context(self): + return {} + + def record_exception(self, error, *, handled, status_code=None): + assert handled + self.recorded_exceptions.append(error) + + +def invocation_context() -> Stage12InvocationContext: + return Stage12InvocationContext( + request_id="request-1", + environment="staging", + modal_environment="staging", + worker_version="5.2.0", + modal_application="policyengine-simulation-v2-py5-2-0", + simulation_callable="run_single_simulation_us", + version_manifest_sha256="b" * 64, + bundle_manifest_sha256="b" * 64, + artifact_prefix=( + "stage-12-runs/staging/2026/09/00000000-0000-0000-0000-000000000001" + ), + created_at=NOW, + retention_expires_at=NOW + timedelta(days=30), + ) + + +def parent_record() -> ComparisonReportRecord: + return ComparisonReportRecord( + evaluation_id=EVALUATION_ID, + status=ComparisonRunLifecycleStatus.PENDING, + aggregation_status=ComparisonRunAggregationStatus.NOT_STARTED, + environment="staging", + calculation_flow="economy", + originating_request_id="request-1", + production_identity=f"direct:{EVALUATION_ID}", + incumbent_execution_id=None, + worker_version="5.2.0", + modal_application="policyengine-simulation-v2-py5-2-0", + report_coordinator_callable="coordinate_report", + version_manifest_sha256="b" * 64, + policyengine_version="5.2.0", + country_package_name="policyengine-us", + country_package_version="1.764.6", + country="us", + dataset_identity="populace_us_2024", + dataset_uri=bundle_provenance().dataset.uri, + data_package_name="populace-data", + data_package_version="0.1.0", + data_artifact_revision="revision", + created_at=NOW, + updated_at=NOW, + retention_expires_at=NOW + timedelta(days=30), + ) + + +def child_record(simulation: SimulationExecutionInput): + from policyengine_simulation_contract.stage12_execution import ( + ComparisonSimulationRecord, + ) + + return ComparisonSimulationRecord( + simulation_execution_id=simulation.simulation_execution_id, + evaluation_id=simulation.evaluation_id, + role=simulation.role, + input_sha256=simulation_input_sha256(simulation), + worker_version="5.2.0", + modal_application="policyengine-simulation-v2-py5-2-0", + simulation_callable="run_single_simulation_us", + version_manifest_sha256="b" * 64, + status=ComparisonRunLifecycleStatus.PENDING, + created_at=NOW, + updated_at=NOW, + retention_expires_at=NOW + timedelta(days=30), + ) + + +def automatic_parent_record() -> ComparisonReportRecord: + return parent_record().model_copy( + update={ + "production_identity": "production-job-1", + "incumbent_execution_id": "production-job-1", + "comparison_status": ResultComparisonStatus.PENDING, + } + ) + + +class _StoppedAfterLoadingError(Exception): + """Raised by a stubbed simulation build once the dataset has loaded.""" + + +def run_until_dataset_loaded( + monkeypatch, *, region, selection, baked_folder, planned=None +): + """Run a Stage 12 calculation up to its dataset load; return what it did.""" + from policyengine_simulation_executor import simulation_runtime + from policyengine_simulation_executor.stage12_cache import state_datasets + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + loads = [] + lookups = [] + + def load_dataset(_params, *, selection, country_module, data_folder=None): + loads.append({"selection": selection, "data_folder": data_folder}) + return object() + + def lookup(**kwargs): + lookups.append(kwargs) + return baked_folder + + def stop_after_loading(*_args, **_kwargs): + raise _StoppedAfterLoadingError + + monkeypatch.setattr(worker, "_require_installed_bundle", lambda _: None) + monkeypatch.setattr(simulation_runtime, "setup_gcp_credentials", nullcontext) + monkeypatch.setattr( + simulation_runtime, "_country_module", lambda _country: object() + ) + monkeypatch.setattr(simulation_runtime, "_resolve_region", lambda **_: region) + monkeypatch.setattr( + simulation_runtime, + "_resolve_dataset_selection", + lambda _params, *, region_resolution: selection, + ) + monkeypatch.setattr(simulation_runtime, "_load_dataset", load_dataset) + monkeypatch.setattr(state_datasets, "baked_state_data_folder", lookup) + monkeypatch.setattr( + worker, "detect_uk_run_dataset_metadata", lambda *_a, **_k: None + ) + monkeypatch.setattr(worker, "build_stage12_simulation", stop_after_loading) + + planned = planned or planned_simulation(SimulationRole.BASELINE) + with pytest.raises(_StoppedAfterLoadingError): + worker.calculate_simulation_frames(planned) + return SimpleNamespace(loads=loads, lookups=lookups, year=planned.year) + + +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"), + ) + + +CORRELATION_SUFFIX = re.compile(r" \(correlation_id=(?P[0-9a-f]+)\)\Z") + + +def summary_correlation_id(error_summary: str) -> str: + """Return the correlation ID that ends a redacted failure summary.""" + match = CORRELATION_SUFFIX.search(error_summary) + assert match is not None, error_summary + return match["correlation_id"] + + +def summary_failure_cause(error_summary: str) -> str: + """Return the redacted cause before a failure summary's correlation ID.""" + summary_correlation_id(error_summary) + return CORRELATION_SUFFIX.sub("", error_summary) + + +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)) + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_aggregate_reports.py b/projects/policyengine-simulation-executor/tests/test_stage12_aggregate_reports.py new file mode 100644 index 000000000..3304280cf --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_aggregate_reports.py @@ -0,0 +1,452 @@ +"""Tests for the aggregate report a Stage 12 coordinator writes.""" + +from __future__ import annotations + +import json +from hashlib import sha256 +from uuid import uuid4 + +import pandas as pd +import pytest +from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from policyengine_simulation_contract.stage12_execution import ( + ReportExecutionInput, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + RetainedSimulationOutput, + SimulationWorkerDependencies, + _build_spm_result, + build_aggregate_report, + run_single_simulation, + simulation_input_sha256, +) +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames +from stage12_runtime_support import ( + EVALUATION_ID, + bundle_provenance, + invocation_context, + planned_simulation, + report_execution_input, + simulation_execution_input, + spm_calculation_provenance, +) + + +def test_aggregate_combines_detached_spm_receipts() -> None: + provenance = spm_calculation_provenance() + + report = report_execution_input() + report = report.model_copy( + update={ + "baseline": report.baseline.model_copy( + update={"options": {"spm": {"scenario": "official"}}} + ), + "reform": report.reform.model_copy( + update={"options": {"spm": {"scenario": "official"}}} + ), + } + ) + baseline = ( + FakeArtifacts() + .add_simulation(report.baseline) + .model_copy(update={"calculation_provenance": provenance}) + ) + reform = ( + FakeArtifacts() + .add_simulation(report.reform) + .model_copy(update={"calculation_provenance": provenance}) + ) + + result = _build_spm_result( + report=report, + baseline_descriptor=baseline, + reform_descriptor=reform, + ) + + assert "spm_config" not in result + assert result["spm_provenance"]["baseline"]["execution_count"] == 1 + assert result["spm_provenance"]["reform"]["execution_count"] == 1 + assert result["spm_provenance"]["baseline"]["receipt"]["scenario"] == "official" + + +def test_aggregate_stand_ins_preserve_policy_and_cliff_options(monkeypatch) -> None: + from policyengine.outputs import labor_supply_response_is_active + from policyengine_simulation_executor import simulation_output_builder + + report = report_execution_input() + options = {"include_cliffs": True} + report = report.model_copy( + update={ + "baseline": report.baseline.model_copy(update={"options": options}), + "reform": report.reform.model_copy( + update={ + "options": options, + "policy": { + "gov.simulation.labor_supply_responses.elasticities.income": 0.1 + }, + } + ), + } + ) + frames = { + entity: pd.DataFrame({f"{entity}_id": [1]}) + for entity in ( + "person", + "marital_unit", + "family", + "spm_unit", + "tax_unit", + "household", + ) + } + observed = {} + + class CapturingBuilder: + def __init__(self, **values): + observed.update(values) + + def serialize(self): + return {"captured": True} + + monkeypatch.setattr( + simulation_output_builder, + "SimulationOutputBuilder", + CapturingBuilder, + ) + artifacts = FakeArtifacts() + + result = build_aggregate_report( + report=report, + baseline=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.reform) + ), + ) + + assert result["result"] == {"captured": True} + assert observed["simulation_params"]["include_cliffs"] is True + assert observed["uk_local_authority_metadata"] is None + assert labor_supply_response_is_active( + observed["baseline"], + observed["reform"], + country_code="us", + ) + + +def test_aggregate_stand_ins_keep_bounded_policy_end_dates(monkeypatch) -> None: + from datetime import datetime + + from policyengine_simulation_executor import simulation_output_builder + + path = "gov.simulation.labor_supply_responses.elasticities.income" + report = report_execution_input() + report = report.model_copy( + update={ + "baseline": report.baseline.model_copy( + update={"policy": {path: {"2026-01-01.2026-12-31": 0.05}}} + ), + "reform": report.reform.model_copy( + update={"policy": {path: {"2026-01-01.2027-12-31": 0.1}}} + ), + } + ) + frames = { + entity: pd.DataFrame({f"{entity}_id": [1]}) + for entity in ( + "person", + "marital_unit", + "family", + "spm_unit", + "tax_unit", + "household", + ) + } + observed = {} + + class CapturingBuilder: + def __init__(self, **values): + observed.update(values) + + def serialize(self): + return {"captured": True} + + monkeypatch.setattr( + simulation_output_builder, "SimulationOutputBuilder", CapturingBuilder + ) + artifacts = FakeArtifacts() + + build_aggregate_report( + report=report, + baseline=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.reform) + ), + ) + + for role, stop in (("baseline", "2026-12-31"), ("reform", "2027-12-31")): + (value,) = observed[role].policy.parameter_values + assert value.end_date == datetime.strptime(stop, "%Y-%m-%d"), role + + +def _uk_report() -> ReportExecutionInput: + uk_bundle = bundle_provenance().model_copy( + update={ + "country_package_name": "policyengine-uk", + "country_package_version": "2.102.3", + } + ) + report_payload = report_execution_input().model_dump(mode="json") + for role, policy in ( + ("baseline", {}), + ( + "reform", + { + "gov.hmrc.income_tax.charges.CB_HITC.phase_out_start": { + "2026-01-01.2100-12-31": 70_000 + } + }, + ), + ): + report_payload[role].update( + { + "policy": policy, + "geography": {"country": "uk", "region": "uk"}, + "bundle": uk_bundle.model_dump(mode="json"), + } + ) + return ReportExecutionInput.model_validate(report_payload) + + +def test_uk_aggregate_without_spm_provenance_is_written(monkeypatch) -> None: + from policyengine_simulation_executor import simulation_output_builder + + report = _uk_report() + assert report.baseline.options.get("spm") is None + uk_output = { + **{ + key: value + for key, value in CURRENT_SINGLE_YEAR_MACRO_RESULT.items() + if key != "spm_provenance" + }, + "congressional_district_impact": None, + } + + class UKOutputBuilder: + def __init__(self, **_values): + pass + + def serialize(self): + # The real builder dumps SingleYearMacroOutput, which serializes + # the unset optional SPM receipt as an explicit null. + return SingleYearMacroOutput.model_validate(uk_output).model_dump( + mode="json" + ) + + monkeypatch.setattr( + simulation_output_builder, "SimulationOutputBuilder", UKOutputBuilder + ) + uk_frames = { + entity: pd.DataFrame({f"{entity}_id": [1]}) + for entity in ("person", "benunit", "household") + } + artifacts = FakeArtifacts() + + aggregate = build_aggregate_report( + report=report, + baseline=RetainedSimulationOutput( + uk_frames, artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + uk_frames, artifacts.add_simulation(report.reform) + ), + ) + assert aggregate["result"]["spm_provenance"] is None + + aggregate_reference = artifacts.write_aggregate( + prefix=invocation_context().artifact_prefix, payload=aggregate + ) + + stored_aggregate = artifacts.read(aggregate_reference.uri) + assert aggregate_reference.content_sha256 == sha256(stored_aggregate).hexdigest() + assert json.loads(stored_aggregate)["result"]["spm_provenance"] is None + + +@pytest.mark.parametrize( + ("region", "routed", "omitted"), + [ + ("uk", True, True), + ("country/wales", True, True), + ("uk", False, False), + ("constituency/E14001063", True, False), + ], +) +def test_uk_aggregate_omits_area_breakdowns_only_for_routed_national_reports( + monkeypatch, region: str, *, routed: bool, omitted: bool +) -> None: + from policyengine_simulation_executor import simulation_output_builder + + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + observed = {} + + class CapturingBuilder: + def __init__(self, **values): + observed.update(values) + + def serialize(self): + return {} + + monkeypatch.setattr( + simulation_output_builder, "SimulationOutputBuilder", CapturingBuilder + ) + payload = _uk_report().model_dump(mode="json") + for role in ("baseline", "reform"): + payload[role]["geography"] = {"country": "uk", "region": region} + report = ReportExecutionInput.model_validate(payload) + frames = { + entity: pd.DataFrame({f"{entity}_id": [1]}) + for entity in ("person", "benunit", "household") + } + artifacts = FakeArtifacts() + + build_aggregate_report( + report=report, + baseline=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + frames, artifacts.add_simulation(report.reform) + ), + ) + + assert observed["omit_uk_area_breakdowns"] is omitted + # The runtime omits them from exactly the same runs. + from policyengine_simulation_executor import uk_local_authority_metadata + + assert uk_local_authority_metadata.uk_area_breakdowns_omitted("uk", region) is ( + omitted + ) + + +def test_aggregate_report_wraps_builder_output_with_report_provenance( + monkeypatch, +) -> None: + from policyengine_simulation_executor import simulation_output_builder + + report = report_execution_input() + report = report.model_copy( + update={"reform": report.reform.model_copy(update={"policy": {}})} + ) + entities = ("person", "marital_unit", "family", "spm_unit", "tax_unit") + + def frames(income): + return { + **{entity: pd.DataFrame({f"{entity}_id": [1]}) for entity in entities}, + "household": pd.DataFrame( + {"household_id": [1], "household_net_income": [income]} + ), + } + + observed = {} + + class CapturingBuilder: + def __init__(self, **values): + observed.update(values) + + def serialize(self): + return {"budget": {"budgetary_impact": 1.0}} + + monkeypatch.setattr( + simulation_output_builder, "SimulationOutputBuilder", CapturingBuilder + ) + artifacts = FakeArtifacts() + + result = build_aggregate_report( + report=report, + baseline=RetainedSimulationOutput( + frames(100.0), artifacts.add_simulation(report.baseline) + ), + reform=RetainedSimulationOutput( + frames(120.0), artifacts.add_simulation(report.reform) + ), + uk_local_authority_metadata=None, + ) + + assert result == { + "contract_version": 1, + "aggregate_schema_version": 1, + "evaluation_id": str(EVALUATION_ID), + "requested_aggregates": [item.value for item in report.requested_aggregates], + "bundle": report.baseline.bundle.model_dump(mode="json"), + "result": {"budget": {"budgetary_impact": 1.0}}, + } + assert list(observed) == [ + "country", + "simulation_params", + "country_module", + "dataset", + "baseline", + "reform", + "resolved_data_version", + "resolved_region_code", + "uk_local_authority_metadata", + "omit_uk_area_breakdowns", + ] + assert observed["country"] == "us" + assert observed["simulation_params"] == { + "country": "us", + "scope": "macro", + "data": report.baseline.bundle.dataset.uri, + "data_version": report.baseline.bundle.dataset.artifact_revision, + "time_period": str(report.baseline.year), + "region": report.baseline.geography.region, + } + assert observed["dataset"] is observed["baseline"].output_dataset + for role, income in (("baseline", 100.0), ("reform", 120.0)): + household = observed[role].output_dataset.data.household + assert household["household_net_income"].tolist() == [income] + assert observed["resolved_data_version"] == ( + report.baseline.bundle.dataset.artifact_revision + ) + assert observed["resolved_region_code"] == report.baseline.geography.region + assert observed["uk_local_authority_metadata"] is None + assert observed["omit_uk_area_breakdowns"] is False + + +def test_single_worker_rejects_combined_baseline_and_reform_input() -> None: + payload = planned_simulation(SimulationRole.BASELINE).model_dump(mode="json") + payload["baseline"] = {} + payload["reform"] = {} + + with pytest.raises(ValueError, match="cannot contain baseline or reform"): + run_single_simulation( + payload, + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=FakeStore(), + artifacts=FakeArtifacts(), + calculator=lambda _: simulation_frames(), + ), + ) + + +def test_simulation_input_digest_excludes_temporary_record_identifiers() -> None: + original = simulation_execution_input(SimulationRole.BASELINE) + recreated = original.model_copy( + update={ + "evaluation_id": uuid4(), + "simulation_execution_id": uuid4(), + } + ) + + assert simulation_input_sha256(original) == simulation_input_sha256(recreated) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py index afed4c94e..16b1b2468 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py @@ -26,7 +26,7 @@ Stage12CachedSimulation, qualifying_cache_identity, ) -from test_stage12_runtime import _planned_simulation +from stage12_runtime_support import planned_simulation # A resolved selection, as normalize_runtime_spm returns for every US request. SPM_SELECTION = { @@ -84,7 +84,7 @@ def _strategy(): def test_cache_writer_and_reader_share_full_identity() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + simulation = planned_simulation(SimulationRole.BASELINE).model_copy( update={"options": {"region_group": ["state/ca"]}} ) resolved = _resolved() @@ -121,7 +121,7 @@ def test_dataset_and_baseline_inputs_rotate_cache_identity() -> None: later = collect_dataset_identity(resolved, country="us", year=2027) assert first.digest != later.digest - simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + simulation = planned_simulation(SimulationRole.BASELINE).model_copy( update={"options": {"region_group": ["state/ca"]}} ) ca = collect_baseline_identity( @@ -177,7 +177,7 @@ def test_dataset_and_baseline_inputs_rotate_cache_identity() -> None: def test_only_current_law_default_dataset_baseline_qualifies() -> None: - baseline = _planned_simulation(SimulationRole.BASELINE).model_copy( + baseline = planned_simulation(SimulationRole.BASELINE).model_copy( update={"options": {"region_group": ["state/ca"]}} ) kwargs = { @@ -256,7 +256,7 @@ def cache_key(self) -> str: def test_only_a_state_region_group_with_a_scope_cache_key_qualifies( options, strategy_name ) -> None: - baseline = _planned_simulation(SimulationRole.BASELINE).model_copy( + baseline = planned_simulation(SimulationRole.BASELINE).model_copy( update={"options": options} ) @@ -284,7 +284,7 @@ def _cached_model() -> Stage12CachedSimulation: dynamic=None, ) model.configure_stage12_cache( - _planned_simulation(SimulationRole.BASELINE).output_plan + planned_simulation(SimulationRole.BASELINE).output_plan ) return model diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_child_cancellation.py b/projects/policyengine-simulation-executor/tests/test_stage12_child_cancellation.py new file mode 100644 index 000000000..8719e9451 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_child_cancellation.py @@ -0,0 +1,406 @@ +"""Tests for cancelling a failed Stage 12 report's unfinished child simulations.""" + +from __future__ import annotations + +import logging +from datetime import timedelta + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, +) +from policyengine_simulation_executor.stage12_runtime import ( + coordination as coordination_module, +) +from stage12_child_call_support import ( + AttachFailingStore, + CancellationStateFailingStore, + PlaceholderWriteFailingStore, + ScriptedChildInvoker, + assert_cancelled, + assert_log_identifies_child, + assert_original_coordinator_failure, + coordinate_with, + coordination_log_record, + report_with_segmentation, + role_simulation_id, +) +from stage12_fake_stores import FakeArtifacts, FakeStore +from stage12_runtime_support import BASELINE_ID, NOW, REFORM_ID, child_record + + +def test_baseline_failure_cancels_and_records_the_reform_child() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline worker failed")} + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_with( + report_with_segmentation(segmented=False), + store, + invoker, + artifacts=artifacts, + ) + + reform_call = invoker.calls["reform"] + assert reform_call.cancel_count == 1 + assert reform_call.probed_at == [] + assert artifacts.aggregate_writes == [] + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert store.parent.error_code == "simulation_invocation_failed" + + +def test_reform_failure_leaves_the_finished_baseline_alone() -> None: + report = report_with_segmentation(segmented=False) + store = FakeStore() + artifacts = FakeArtifacts() + baseline_descriptor = artifacts.add_simulation(report.baseline) + invoker = ScriptedChildInvoker( + results={"baseline": baseline_descriptor.model_dump(mode="json")}, + wait_errors={"reform": LookupError("reform worker failed")}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_with(report, store, invoker, artifacts=artifacts) + + assert invoker.calls["baseline"].cancel_count == 0 + assert store.children[BASELINE_ID].status is ComparisonRunLifecycleStatus.RUNNING + assert store.children[BASELINE_ID].error_code is None + assert store.children[REFORM_ID].error_code == "simulation_invocation_failed" + + +def test_baseline_failure_cancels_a_reform_call_restored_on_retry() -> None: + report = report_with_segmentation(segmented=False) + store = FakeStore() + # An earlier coordinator attempt started the reform child before stopping. + store.children[REFORM_ID] = child_record(report.reform).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": "call-reform", + "started_at": NOW, + } + ) + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline worker failed")} + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_with(report, store, invoker) + + assert invoker.events == ["spawn:baseline", "restore:reform"] + assert invoker.calls["reform"].cancel_count == 1 + assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + + +@pytest.mark.parametrize( + ("worker_status", "worker_error_code"), + [ + (ComparisonRunLifecycleStatus.SUCCEEDED, None), + (ComparisonRunLifecycleStatus.FAILED, "simulation_execution_failed"), + ], +) +def test_cancellation_leaves_a_sibling_that_already_finished_unchanged( + worker_status, + worker_error_code, +) -> None: + store = FakeStore() + worker_records = [] + + def reform_worker_records_its_outcome(): + # The reform worker finishes while the coordinator probes baseline. + finished_at = NOW + timedelta(minutes=5) + worker_records.append( + store.replace_simulation( + store.get_simulation(REFORM_ID).model_copy( + update={ + "status": worker_status, + "error_code": worker_error_code, + "updated_at": finished_at, + "completed_at": finished_at, + } + ) + ) + ) + + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline worker failed")}, + during_wait={"baseline": reform_worker_records_its_outcome}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_with(report_with_segmentation(segmented=False), store, invoker) + + # Cancelling a finished Modal call is harmless; its record is not touched. + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[REFORM_ID] == worker_records[0] + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + + +def test_peer_cancellation_failure_does_not_replace_the_child_failure( + caplog, +) -> None: + store = FakeStore() + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline failure is sensitive")}, + cancel_errors={ + "reform": ConnectionError("cancel request contains sensitive details") + }, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + coordinate_with(report_with_segmentation(segmented=False), store, invoker) + + assert invoker.calls["reform"].cancel_count == 1 + assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert store.parent.error_code == "simulation_invocation_failed" + assert_log_identifies_child( + coordination_log_record(caplog, logging.WARNING), + role="reform", + error_type="ConnectionError", + ) + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_failed_cancellation_state_write_still_cancels_through_modal( + caplog, + segmented, +) -> None: + store = CancellationStateFailingStore() + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline failure is sensitive")}, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + coordinate_with(report_with_segmentation(segmented=segmented), store, invoker) + + # A segmented worker cannot read a cancellation that was never written, so + # it is cancelled through Modal as well. + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.RUNNING + assert store.parent.error_code == "simulation_invocation_failed" + assert_log_identifies_child( + coordination_log_record(caplog, logging.ERROR), + role="reform", + error_type="ConnectionError", + ) + + +def test_wait_timeout_cancels_and_records_every_unfinished_child() -> None: + store = FakeStore() + invoker = ScriptedChildInvoker() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_with( + report_with_segmentation(segmented=False), + store, + invoker, + wait_timeout_seconds=0, + ) + + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + for role in ("baseline", "reform"): + assert_cancelled( + store.children[role_simulation_id(role)], + after="the 0-second wait window closed", + ) + assert store.parent.error_code == "simulation_wait_timeout" + + +def test_timeout_cancellation_failure_log_identifies_the_report_and_child( + caplog, +) -> None: + store = FakeStore() + invoker = ScriptedChildInvoker( + cancel_errors={ + "reform": ConnectionError("cancel request contains sensitive details") + }, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + coordinate_with( + report_with_segmentation(segmented=False), + store, + invoker, + wait_timeout_seconds=0, + ) + + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + assert store.parent.error_code == "simulation_wait_timeout" + assert_log_identifies_child( + coordination_log_record(caplog, logging.WARNING), + role="reform", + error_type="ConnectionError", + ) + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_failed_placeholder_write_stops_the_started_baseline(segmented) -> None: + store = PlaceholderWriteFailingStore(REFORM_ID) + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + coordinate_with(report_with_segmentation(segmented=segmented), store, invoker) + + assert_original_coordinator_failure(error, store) + assert invoker.events == ["spawn:baseline"] + assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") + # A segmented worker stops itself, and its region-group calls, once it + # reads the cancelled record; Modal does not cancel spawned calls + # recursively. + assert invoker.calls["baseline"].cancel_count == (0 if segmented else 1) + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.PENDING + + +@pytest.mark.parametrize( + ("failing_role", "started_roles"), + [ + ("baseline", ["baseline"]), + ("reform", ["baseline", "reform"]), + ], +) +def test_failed_invocation_attach_stops_every_started_child( + failing_role, + started_roles, +) -> None: + store = AttachFailingStore(role_simulation_id(failing_role)) + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + coordinate_with(report_with_segmentation(segmented=False), store, invoker) + + assert_original_coordinator_failure(error, store) + assert invoker.events == [f"spawn:{role}" for role in started_roles] + # The child whose invocation could not be attached is stopped as well. + for role in started_roles: + assert_cancelled( + store.children[role_simulation_id(role)], after="the coordinator failed" + ) + assert invoker.calls[role].cancel_count == 1 + if failing_role == "baseline": + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.PENDING + + +def test_startup_failure_stops_a_call_restored_from_an_earlier_attempt() -> None: + report = report_with_segmentation(segmented=False) + store = PlaceholderWriteFailingStore(REFORM_ID) + store.children[BASELINE_ID] = child_record(report.baseline).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": "call-baseline", + "started_at": NOW, + } + ) + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + coordinate_with(report, store, invoker) + + assert_original_coordinator_failure(error, store) + assert invoker.events == ["restore:baseline"] + assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") + assert invoker.calls["baseline"].cancel_count == 1 + + +def test_startup_failure_leaves_a_finished_child_unchanged() -> None: + report = report_with_segmentation(segmented=False) + store = AttachFailingStore(REFORM_ID) + artifacts = FakeArtifacts() + baseline_descriptor = artifacts.add_simulation(report.baseline) + finished_baseline = child_record(report.baseline).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": baseline_descriptor.artifact.uri, + "output_sha256": baseline_descriptor.artifact.content_sha256, + "output_schema_version": 1, + "row_identity_columns": baseline_descriptor.row_identity.identifier_columns, + "row_count": baseline_descriptor.row_identity.row_count, + "row_identity_sha256": baseline_descriptor.row_identity.identity_sha256, + "completed_at": NOW, + } + ) + store.children[BASELINE_ID] = finished_baseline + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + coordinate_with(report, store, invoker, artifacts=artifacts) + + assert_original_coordinator_failure(error, store) + assert invoker.events == ["spawn:reform"] + assert store.children[BASELINE_ID] == finished_baseline + assert_cancelled(store.children[REFORM_ID], after="the coordinator failed") + assert invoker.calls["reform"].cancel_count == 1 + + +def test_startup_cancellation_error_does_not_replace_the_original_failure( + caplog, +) -> None: + store = PlaceholderWriteFailingStore(REFORM_ID) + invoker = ScriptedChildInvoker( + cancel_errors={ + "baseline": ConnectionError("cancel request contains sensitive details") + } + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error, + ): + coordinate_with(report_with_segmentation(segmented=False), store, invoker) + + assert_original_coordinator_failure(error, store) + assert invoker.calls["baseline"].cancel_count == 1 + assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") + assert_log_identifies_child( + coordination_log_record(caplog, logging.WARNING), + role="baseline", + error_type="ConnectionError", + ) + + +def test_startup_cancellation_state_error_does_not_replace_the_original_failure( + caplog, +) -> None: + class FailingStore(PlaceholderWriteFailingStore, CancellationStateFailingStore): + pass + + store = FailingStore(REFORM_ID) + invoker = ScriptedChildInvoker() + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error, + ): + coordinate_with(report_with_segmentation(segmented=False), store, invoker) + + assert_original_coordinator_failure(error, store) + assert invoker.calls["baseline"].cancel_count == 1 + assert_log_identifies_child( + coordination_log_record(caplog, logging.ERROR), + role="baseline", + error_type="ConnectionError", + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_child_state_errors.py b/projects/policyengine-simulation-executor/tests/test_stage12_child_state_errors.py new file mode 100644 index 000000000..244a15dac --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_child_state_errors.py @@ -0,0 +1,400 @@ +"""Tests that child state errors never replace the failure a report records.""" + +from __future__ import annotations + +import logging +from types import SimpleNamespace + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + PlannedSimulationExecutionInput, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + SegmentRunnerSeams, + SimulationWorkerDependencies, + Stage12SegmentRunner, + run_single_simulation, +) +from policyengine_simulation_executor.stage12_runtime import ( + coordination as coordination_module, +) +from policyengine_simulation_executor.stage12_runtime.partition import US_REGION_GROUPS +from stage12_child_call_support import ( + ChildCallScript, + ConcurrentInvoker, + RefusingDispatchInvoker, + ScriptedChildCall, + ScriptedChildInvoker, + assert_cancelled, + assert_log_identifies_child, + coordinate_with, + coordination_log_record, + report_with_segmentation, +) +from stage12_fake_stores import FakeArtifacts, FakeStore +from stage12_runtime_support import ( + BASELINE_ID, + EVALUATION_ID, + REFORM_ID, + RecordingRuntime, + child_record, + coordinate_report, + invocation_context, + parent_record, + planned_simulation, + report_execution_input, + summary_correlation_id, +) + + +class ChildStateError(Exception): + """A database failure while reading or writing a child's own record.""" + + +class OwnChildStateFailingStore(FakeStore): + """Fail one read, or one failed-status write, of one child's record once. + + Cancellation writes are left alone, so only the failed child's own state + update can fail. + """ + + def __init__(self, *, failing_simulation_id, failing_operation): + super().__init__() + self.failing_simulation_id = failing_simulation_id + self.failing_operation = failing_operation + self.remaining_failures = 1 + + def _fail_once(self, operation, simulation_execution_id): + if ( + operation == self.failing_operation + and simulation_execution_id == self.failing_simulation_id + and self.remaining_failures > 0 + ): + self.remaining_failures -= 1 + msg = "database contains sensitive details" + raise ChildStateError(msg) + + def get_simulation(self, simulation_execution_id): + self._fail_once("read", simulation_execution_id) + return super().get_simulation(simulation_execution_id) + + def replace_simulation_if_status(self, record, *, expected_status): + if record.error_code != "simulation_cancelled": + self._fail_once("write", record.simulation_execution_id) + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + +@pytest.mark.parametrize( + ("failing_operation", "baseline_status", "baseline_error_code"), + [ + # The failure is still recorded: the read before the conditional + # write succeeds. + ( + "read", + ComparisonRunLifecycleStatus.FAILED, + "simulation_invocation_failed", + ), + ("write", ComparisonRunLifecycleStatus.RUNNING, None), + ], +) +def test_failed_child_state_error_does_not_replace_the_child_failure( + caplog, + failing_operation, + baseline_status, + baseline_error_code, +) -> None: + store = OwnChildStateFailingStore( + failing_simulation_id=BASELINE_ID, + failing_operation=failing_operation, + ) + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline failure is sensitive")}, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + coordinate_with(report_with_segmentation(segmented=False), store, invoker) + + assert store.parent.error_code == "simulation_invocation_failed" + summary_correlation_id(store.parent.error_summary) + baseline = store.children[BASELINE_ID] + assert baseline.status is baseline_status + assert baseline.error_code == baseline_error_code + if baseline_error_code is not None: + assert baseline.error_summary == store.parent.error_summary + # The failed child is neither cancelled nor recorded as cancelled. + assert invoker.calls["baseline"].cancel_count == 0 + assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert invoker.calls["reform"].cancel_count == 1 + assert_log_identifies_child( + coordination_log_record(caplog, logging.ERROR), + role="baseline", + error_type="ChildStateError", + ) + + +class ReformDispatchFailureInvoker(ScriptedChildInvoker): + def spawn(self, request): + role = PlannedSimulationExecutionInput.model_validate(request.simulation).role + if role is SimulationRole.REFORM: + self.events.append("spawn:reform") + msg = "dispatch contains sensitive details" + raise ConnectionError(msg) + return super().spawn(request) + + +@pytest.mark.parametrize("failing_operation", ["read", "write"]) +def test_dispatch_failure_state_error_does_not_replace_the_child_failure( + caplog, + failing_operation, +) -> None: + store = OwnChildStateFailingStore( + failing_simulation_id=REFORM_ID, + failing_operation=failing_operation, + ) + invoker = ReformDispatchFailureInvoker() + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + coordinate_with(report_with_segmentation(segmented=False), store, invoker) + + assert invoker.events == ["spawn:baseline", "spawn:reform"] + assert store.parent.error_code == "simulation_dispatch_failed" + summary_correlation_id(store.parent.error_summary) + # The reform record could not be updated, so it keeps its dispatch + # placeholder rather than a cancellation. + reform = store.children[REFORM_ID] + assert reform.status is ComparisonRunLifecycleStatus.RUNNING + assert reform.error_code is None + assert_cancelled(store.children[BASELINE_ID], after="the reform simulation failed") + assert invoker.calls["baseline"].cancel_count == 1 + assert_log_identifies_child( + coordination_log_record(caplog, logging.ERROR), + role="reform", + error_type="ChildStateError", + ) + + +class ExceptionRecordingRuntime(RecordingRuntime): + """Record every telemetry call that can report a failure.""" + + def __init__(self) -> None: + super().__init__() + self.reports: list[tuple[str, object, dict, object]] = [] + + def log(self, message, *, severity="INFO", attributes=None, error=None): + self.reports.append(("log", message, dict(attributes or {}), error)) + + def record_exception(self, error, *, handled, status_code=None): + self.reports.append(("exception", None, {}, error)) + + def event(self, name, *, attributes=None, **_options): + self.reports.append(("event", name, dict(attributes or {}), None)) + + def set_context(self, **_attributes): + return None + + +def _fail_simulation_calculation(runtime): + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + calculation_error = ValueError("calculation contains sensitive details") + + def fail(_simulation): + raise calculation_error + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + runtime=runtime, + ), + ) + return calculation_error, store.children[BASELINE_ID].error_summary + + +def _fail_child_call(runtime): + store = FakeStore() + child_error = LookupError("baseline failure contains sensitive details") + invoker = ScriptedChildInvoker(wait_errors={"baseline": child_error}) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_with( + report_with_segmentation(segmented=False), + store, + invoker, + runtime=runtime, + ) + return child_error, store.parent.error_summary + + +def _fail_child_dispatch(runtime): + store = FakeStore() + dispatch_error = ConnectionError("dispatch contains sensitive details") + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_with( + report_execution_input(), + store, + RefusingDispatchInvoker(dispatch_error), + runtime=runtime, + ) + return dispatch_error, store.parent.error_summary + + +def _fail_report_aggregation(runtime): + store = FakeStore() + artifacts = FakeArtifacts() + aggregation_error = ValueError("aggregation contains sensitive details") + + def fail_aggregation(**_values): + raise aggregation_error + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=fail_aggregation, + runtime=runtime, + ) + return aggregation_error, store.parent.error_summary + + +@pytest.mark.parametrize( + "fail", + [ + _fail_simulation_calculation, + _fail_child_call, + _fail_child_dispatch, + _fail_report_aggregation, + ], +) +@pytest.mark.parametrize("with_runtime", [True, False], ids=["runtime", "no-runtime"]) +def test_an_unexpected_failure_is_reported_exactly_once( + caplog, + fail, + with_runtime, +) -> None: + runtime = ExceptionRecordingRuntime() if with_runtime else None + + with caplog.at_level(logging.DEBUG): + original_error, error_summary = fail(runtime) + + correlation_id = summary_correlation_id(error_summary) + logged_failures = [record for record in caplog.records if record.exc_info] + if runtime is not None: + [(kind, message, attributes, reported_error)] = [ + report for report in runtime.reports if report[0] == "log" + ] + assert (kind, message) == ("log", "Stage 12 execution failed") + assert reported_error is original_error + assert attributes["correlation_id"] == correlation_id + # The redacted cause, never the original error, is also attached to + # the failing operation, once. + [recorded_cause] = [ + report[3] for report in runtime.reports if report[0] == "exception" + ] + assert str(recorded_cause) == error_summary + assert sorted(report[0] for report in runtime.reports) == ["exception", "log"] + assert logged_failures == [] + else: + [logged_failure] = logged_failures + assert logged_failure.exc_info[1] is original_error + assert f"correlation_id={correlation_id}" in logged_failure.getMessage() + + +def test_a_segment_failure_is_reported_exactly_once(caplog) -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + runtime = ExceptionRecordingRuntime() + segment_error = RuntimeError("segment contains sensitive details") + + class SegmentFunction: + def spawn(self, payload, *, observability_context): + segment_index = payload["segment_index"] + # The first segment fails on every probe; the others keep running. + return ScriptedChildCall( + f"segment-{segment_index}", + script=ChildCallScript( + wait_error=segment_error if segment_index == 0 else None + ), + ) + + class SegmentFunctionLookup: + @staticmethod + def from_name(_app_name, _function_name): + return SegmentFunction() + + def segmented_calculator(value, cancellation_requested): + return Stage12SegmentRunner( + value, + app_name="stage12-app", + runtime=runtime, + cancellation_requested=cancellation_requested, + seams=SegmentRunnerSeams( + modal_module=SimpleNamespace(Function=SegmentFunctionLookup), + groups=US_REGION_GROUPS, + sleep=lambda _seconds: None, + ), + ).run() + + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"), + ): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + runtime=runtime, + ), + ) + + child = store.children[BASELINE_ID] + assert child.error_code == "segment_execution_failed" + [(kind, message, attributes, reported_error)] = [ + report for report in runtime.reports if report[0] == "log" + ] + [recorded_cause] = [ + report[3] for report in runtime.reports if report[0] == "exception" + ] + assert str(recorded_cause) == child.error_summary + assert (kind, message) == ("log", "Stage 12 execution failed") + assert reported_error is segment_error + assert attributes == { + "stage12_failure_scope": "segment_call", + "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "segment_index": 0, + "correlation_id": summary_correlation_id( + child.error_summary.removeprefix("Segment 0: ") + ), + } + assert [record for record in caplog.records if record.exc_info] == [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py index 3f547dcb9..edbcda133 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordination.py @@ -21,21 +21,20 @@ parent_claim, report_coordination, ) -from test_stage12_runtime import ( +from stage12_child_call_support import ConcurrentInvoker +from stage12_fake_stores import FakeArtifacts, FakeStore +from stage12_runtime_support import ( BASELINE_ID, EVALUATION_ID, NOW, REFORM_ID, - ConcurrentInvoker, - FakeArtifacts, - FakeStore, - _aggregate_payload, - _child, - _context, - _output_plan, - _parent, - _report, + aggregate_payload, + child_record, coordinate_report, + invocation_context, + parent_record, + report_execution_input, + shared_output_plan, ) CLAIMED_AT = datetime(2026, 9, 15, 1, 0, tzinfo=UTC) @@ -113,23 +112,23 @@ def restore(self, invocation_id): def _coordinate(context=None, parent=None, **dependencies): - context = context or _context() + context = context or invocation_context() dependencies.setdefault("application_name", context.modal_application) dependencies.setdefault("coordinator_invocation_id", "coordinator-1") - dependencies.setdefault("aggregator", lambda **_: _aggregate_payload()) + dependencies.setdefault("aggregator", lambda **_: aggregate_payload()) dependencies.setdefault("artifacts", FakeArtifacts()) dependencies.setdefault("invoker", ConcurrentInvoker(dependencies["artifacts"])) return coordinate_report( - _report().model_dump(mode="json"), + report_execution_input().model_dump(mode="json"), context.model_dump(mode="json"), - (parent or _parent()).model_dump(mode="json"), + (parent or parent_record()).model_dump(mode="json"), **dependencies, ) def _succeeded_child(simulation, artifacts): descriptor = artifacts.add_simulation(simulation) - return _child(simulation).model_copy( + return child_record(simulation).model_copy( update={ "status": ComparisonRunLifecycleStatus.SUCCEEDED, "output_uri": descriptor.artifact.uri, @@ -167,7 +166,7 @@ def test_dispatch_context_must_match_the_report( context_update, application_name, message ): store = FakeStore() - context = _context().model_copy(update=context_update) + context = invocation_context().model_copy(update=context_update) with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): _coordinate( @@ -225,7 +224,9 @@ def test_submitted_direct_parent_must_describe_the_invocation(field, value): message = messages.get(field, f"submitted parent {field} does not match invocation") with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): - _coordinate(parent=_parent().model_copy(update={field: value}), store=store) + _coordinate( + parent=parent_record().model_copy(update={field: value}), store=store + ) assert store.parent is None @@ -240,10 +241,10 @@ def test_submitted_direct_parent_must_describe_the_invocation(field, value): ) def test_submitted_automatic_parent_must_name_the_production_job(update): store = FakeStore() - context = _context().model_copy( + context = invocation_context().model_copy( update={"production_function_call_id": "production-job-1"} ) - parent = _parent().model_copy( + parent = parent_record().model_copy( update={ "production_identity": "production-job-1", "incumbent_execution_id": "production-job-1", @@ -263,7 +264,7 @@ def test_submitted_automatic_parent_must_name_the_production_job(update): def test_existing_parent_with_another_identifier_is_rejected_unwrapped(): store = FakeStore() - store.parent = _parent().model_copy( + store.parent = parent_record().model_copy( update={"evaluation_id": UUID("00000000-0000-0000-0000-0000000000aa")} ) @@ -279,7 +280,7 @@ def test_existing_parent_with_another_identifier_is_rejected_unwrapped(): def test_failed_parent_is_reclaimed_keeping_its_first_start(previous_start): store = FakeStore() artifacts = FakeArtifacts() - store.parent = _parent().model_copy( + store.parent = parent_record().model_copy( update={ "status": ComparisonRunLifecycleStatus.FAILED, "aggregation_status": ComparisonRunAggregationStatus.FAILED, @@ -323,7 +324,9 @@ def record_replace(record): ) def test_claimed_or_finished_parent_is_deduplicated(status): store = FakeStore() - existing = _parent().model_copy(update={"status": status, "error_code": "skipped"}) + existing = parent_record().model_copy( + update={"status": status, "error_code": "skipped"} + ) if status is ComparisonRunLifecycleStatus.SUCCEEDED: existing = existing.model_copy( update={ @@ -349,7 +352,7 @@ def test_claimed_or_finished_parent_is_deduplicated(status): def test_retained_child_with_a_changed_artifact_fails_the_report(): - report = _report() + report = report_execution_input() store = FakeStore() artifacts = FakeArtifacts() succeeded = _succeeded_child(report.baseline, artifacts) @@ -394,7 +397,7 @@ def test_dispatch_failure_marks_the_child_and_report_failed(): def test_dispatch_failure_keeps_a_child_that_succeeded_meanwhile(): - report = _report() + report = report_execution_input() store = FakeStore() artifacts = FakeArtifacts() succeeded = _succeeded_child(report.baseline, artifacts) @@ -408,7 +411,9 @@ def test_dispatch_failure_keeps_a_child_that_succeeded_meanwhile(): def test_default_output_plan_resolver_measures_its_own_stages(monkeypatch): - from test_stage12_runtime import coordinate_report_impl + from policyengine_simulation_executor.stage12_runtime import ( + coordinate_report as coordinate_report_impl, + ) store = FakeStore() artifacts = FakeArtifacts() @@ -418,23 +423,23 @@ def test_default_output_plan_resolver_measures_its_own_stages(monkeypatch): def resolve(report, *, stage_scope): with stage_scope(coordination.Stage.STAGE12_OUTPUT_PLANNING, {"probe": 1}): resolved.append(report.evaluation_id) - return _output_plan() + return shared_output_plan() monkeypatch.setattr(report_coordination, "resolve_report_output_plan", resolve) coordinate_report_impl( - _report().model_dump(mode="json"), - _context().model_dump(mode="json"), - _parent().model_dump(mode="json"), + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), invocation=CoordinatorInvocation( - application_name=_context().modal_application, + application_name=invocation_context().modal_application, invocation_id="coordinator-1", ), dependencies=CoordinatorDependencies( store=store, artifacts=artifacts, invoker=ConcurrentInvoker(artifacts), - aggregator=lambda **_: _aggregate_payload(), + aggregator=lambda **_: aggregate_payload(), runtime=runtime, ), ) @@ -452,7 +457,7 @@ def test_successful_coordination_records_every_stage_and_state(): def aggregate(**values): aggregated.append(sorted(values)) - return _aggregate_payload({"budget": 1}) + return aggregate_payload({"budget": 1}) result = _coordinate( store=store, @@ -489,12 +494,14 @@ def aggregate(**values): "artifact": aggregate_artifact, "baseline_artifact_sha256": result["baseline_artifact_sha256"], "reform_artifact_sha256": result["reform_artifact_sha256"], - "bundle": _report().baseline.bundle.model_dump(mode="json"), + "bundle": report_execution_input().baseline.bundle.model_dump(mode="json"), } - aggregate_uri = f"gs://private/{_context().artifact_prefix}/reports/aggregate.json" + aggregate_uri = ( + f"gs://private/{invocation_context().artifact_prefix}/reports/aggregate.json" + ) assert aggregate_artifact["uri"] == aggregate_uri parent = store.parent.model_dump(mode="json") - expected_parent = _parent().model_dump(mode="json") + expected_parent = parent_record().model_dump(mode="json") expected_parent.update( { "status": "succeeded", @@ -579,10 +586,10 @@ def import_modal(module_name): def _shadow_run(store, artifacts, invoker, runtime=None): - context = _context().model_copy( + context = invocation_context().model_copy( update={"production_function_call_id": "production-job-1"} ) - parent = _parent().model_copy( + parent = parent_record().model_copy( update={ "production_identity": "production-job-1", "incumbent_execution_id": "production-job-1", @@ -595,7 +602,7 @@ def _shadow_run(store, artifacts, invoker, runtime=None): store=store, artifacts=artifacts, invoker=invoker, - aggregator=lambda **_: _aggregate_payload({"budget": {"total": 100.0}}), + aggregator=lambda **_: aggregate_payload({"budget": {"total": 100.0}}), runtime=runtime, ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordinator.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordinator.py new file mode 100644 index 000000000..fda435620 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordinator.py @@ -0,0 +1,413 @@ +"""Tests for Stage 12 report coordination: starting, aggregating and comparing.""" + +from __future__ import annotations + +import io +import json + +from policyengine_observability import ( + DeploymentIdentity, + LoggingConfig, + ObservabilityConfig, + OTelConfig, + ServiceIdentity, + StdoutLogDestination, + configure, +) +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunAggregationStatus, + ComparisonRunLifecycleStatus, + ResultComparisonStatus, +) +from policyengine_simulation_executor.stage12_runtime import ( + CoordinatorDependencies, + CoordinatorInvocation, +) +from policyengine_simulation_executor.stage12_runtime import ( + coordinate_report as coordinate_report_impl, +) +from stage12_child_call_support import ConcurrentInvoker +from stage12_fake_stores import FakeArtifacts, FakeStore +from stage12_runtime_support import ( + EVALUATION_ID, + NOW, + RecordingRuntime, + aggregate_payload, + automatic_parent_record, + coordinate_report, + invocation_context, + parent_record, + report_execution_input, + shared_output_plan, +) + + +def test_coordinator_starts_both_children_before_waiting_and_aggregates() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ConcurrentInvoker(artifacts) + observed = {} + + def aggregate(**values): + observed.update(values) + return aggregate_payload() + + context = invocation_context().model_copy( + update={ + "environment": "production", + "modal_environment": "main", + "artifact_prefix": ( + "stage-12-runs/production/2026/09/00000000-0000-0000-0000-000000000001" + ), + } + ) + parent = parent_record().model_copy(update={"environment": "production"}) + resolver_calls = [] + + def resolver(report): + resolver_calls.append(report.evaluation_id) + return shared_output_plan() + + result = coordinate_report_impl( + report_execution_input().model_dump(mode="json"), + context.model_dump(mode="json"), + parent.model_dump(mode="json"), + invocation=CoordinatorInvocation( + application_name=context.modal_application, + invocation_id="coordinator-1", + ), + dependencies=CoordinatorDependencies( + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=aggregate, + output_plan_resolver=resolver, + ), + ) + + assert invoker.events == ["spawn:baseline", "spawn:reform"] + assert invoker.environments == ["main", "main"] + assert resolver_calls == [EVALUATION_ID] + assert invoker.output_plans == [shared_output_plan(), shared_output_plan()] + latest_start = max(interval[0] for interval in invoker.intervals.values()) + earliest_end = min(interval[1] for interval in invoker.intervals.values()) + assert latest_start < earliest_end + assert result["artifact"]["uri"] == ( + f"gs://private/{context.artifact_prefix}/reports/aggregate.json" + ) + assert observed["baseline"].frames["household"][ + "household_net_income" + ].tolist() == [100.0] + assert observed["reform"].frames["household"]["household_net_income"].tolist() == [ + 120.0 + ] + assert observed["uk_local_authority_metadata"] is None + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert store.parent.aggregation_status is ComparisonRunAggregationStatus.SUCCEEDED + assert store.parent.coordinator_invocation_id == "coordinator-1" + + +def test_duplicate_coordinator_submission_does_not_start_duplicate_children() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + context = invocation_context() + parent = parent_record() + first_invoker = ConcurrentInvoker(artifacts) + + coordinate_report( + report_execution_input().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=first_invoker, + aggregator=lambda **_: aggregate_payload(), + ) + duplicate_invoker = ConcurrentInvoker(artifacts) + duplicate = coordinate_report( + report_execution_input().model_dump(mode="json"), + context.model_dump(mode="json"), + parent.model_dump(mode="json"), + application_name=context.modal_application, + coordinator_invocation_id="coordinator-2", + store=store, + artifacts=artifacts, + invoker=duplicate_invoker, + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + ) + + assert duplicate == { + "deduplicated": True, + "evaluation_id": str(EVALUATION_ID), + "status": "succeeded", + } + assert duplicate_invoker.events == [] + + +def test_coordinator_measures_output_and_child_input_planning() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + runtime = RecordingRuntime() + + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=lambda **_: aggregate_payload(), + runtime=runtime, + ) + + planning_operations = [ + operation + for operation in runtime.operations + if operation[0] + in { + "stage12_coordinator_preparation", + "stage12_output_planning", + "stage12_child_input_planning", + } + ] + assert planning_operations == [ + ("stage12_coordinator_preparation", {}), + ("stage12_output_planning", {"country": "us"}), + ("stage12_child_input_planning", {"simulation_role": "baseline"}), + ("stage12_child_input_planning", {"simulation_role": "reform"}), + ] + + +def test_coordinator_planning_operations_preserve_observability_id() -> None: + output = io.StringIO() + runtime = configure( + ObservabilityConfig( + service=ServiceIdentity( + name="stage12-test", + namespace="policyengine.api-v1", + version="test", + role="coordinator", + ), + deployment=DeploymentIdentity( + environment="test", + platform="local", + region="us-central1", + instance_id="test-process", + ), + logging=LoggingConfig( + destinations=(StdoutLogDestination(),), + ), + otel=OTelConfig(enabled=False), + dispatch_attribute_keys=frozenset({"observability_id"}), + ) + ) + runtime._delivery._stdout = output + observability_id = "00000000-0000-4000-8000-000000000004" + artifacts = FakeArtifacts() + + try: + with runtime.operation( + "stage12_report", + remote_context={ + "captured_at": NOW.isoformat(), + "observability_id": observability_id, + }, + ): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=FakeStore(), + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=lambda **_: aggregate_payload(), + runtime=runtime, + ) + finally: + runtime.shutdown() + + records = [json.loads(line) for line in output.getvalue().splitlines()] + planning_names = { + "stage12_coordinator_preparation", + "stage12_output_planning", + "stage12_child_input_planning", + } + planning_records = [ + record for record in records if record.get("operation.name") in planning_names + ] + assert len(planning_records) == 4 + assert all( + record["attributes"]["observability_id"] == observability_id + for record in planning_records + ) + + +def test_coordinator_compares_automatic_run_after_successful_aggregation() -> None: + store = FakeStore() + parent = automatic_parent_record() + artifacts = FakeArtifacts() + production_result = {"budget": {"total": 100.0, "count": 2}} + stage12_result = {"budget": {"total": 101.0, "count": 2}} + invoker = ConcurrentInvoker( + artifacts, + production_result=production_result, + ) + context = invocation_context().model_copy( + update={"production_function_call_id": "production-job-1"} + ) + + coordinate_report( + report_execution_input().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 **_: aggregate_payload(stage12_result), + ) + + assert invoker.events[-1] == "restore:production" + assert invoker.production_wait_timeouts == [900] + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert store.parent.comparison_status is ResultComparisonStatus.DIFFERENT + assert store.parent.comparison_output_uri == "gs://private/comparison.json" + receipt = artifacts.comparison_writes[0] + assert receipt["difference_count"] == 1 + assert receipt["differences"][0]["path"] == "/budget/total" + assert receipt["differences"][0]["production_value"] == 100.0 + assert receipt["differences"][0]["stage12_value"] == 101.0 + assert "production_result" not in receipt + assert "stage12_result" not in receipt + + +def test_comparison_failure_does_not_change_successful_stage12_report() -> None: + store = FakeStore() + parent = automatic_parent_record() + artifacts = FakeArtifacts() + invoker = ConcurrentInvoker(artifacts, production_result="not-an-object") + context = invocation_context().model_copy( + update={"production_function_call_id": "production-job-1"} + ) + + coordinate_report( + report_execution_input().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 **_: aggregate_payload({"budget": 100}), + ) + + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert store.parent.aggregation_status is ComparisonRunAggregationStatus.SUCCEEDED + assert store.parent.comparison_status is ResultComparisonStatus.FAILED + assert store.parent.comparison_error_code == "result_comparison_failed" + assert artifacts.comparison_writes == [] + + +class UnavailableProductionCall: + object_id = "production-job-1" + + def get(self, *, timeout=None): + msg = ( + "production result unavailable at " + "https://storage.example.com/result.json?X-Goog-Signature=signed-value" + ) + raise ConnectionError(msg) + + +class UnavailableProductionInvoker(ConcurrentInvoker): + def restore(self, invocation_id): + if invocation_id == "production-job-1": + return UnavailableProductionCall() + return super().restore(invocation_id) + + +def test_comparison_failure_stores_a_redacted_cause() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + context = invocation_context().model_copy( + update={"production_function_call_id": "production-job-1"} + ) + + coordinate_report( + report_execution_input().model_dump(mode="json"), + context.model_dump(mode="json"), + automatic_parent_record().model_dump(mode="json"), + application_name=context.modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=UnavailableProductionInvoker(artifacts), + aggregator=lambda **_: aggregate_payload({"budget": 100}), + ) + + assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert store.parent.comparison_status is ResultComparisonStatus.FAILED + assert store.parent.comparison_error_summary == ( + "ConnectionError: production result unavailable at " + "https://storage.example.com/result.json?[REDACTED]" + ) + + +def test_repeated_submission_does_not_retry_failed_comparison() -> None: + store = FakeStore() + parent = automatic_parent_record() + artifacts = FakeArtifacts() + context = invocation_context().model_copy( + update={"production_function_call_id": "production-job-1"} + ) + + first_invoker = ConcurrentInvoker(artifacts, production_result="not-an-object") + coordinate_report( + report_execution_input().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=first_invoker, + aggregator=lambda **_: aggregate_payload({"budget": {"total": 100.0}}), + ) + failed_at = store.parent.comparison_completed_at + failed_summary = store.parent.comparison_error_summary + + second_invoker = ConcurrentInvoker( + artifacts, + production_result={"budget": {"total": 100.0}}, + ) + result = coordinate_report( + report_execution_input().model_dump(mode="json"), + context.model_dump(mode="json"), + parent.model_dump(mode="json"), + application_name=context.modal_application, + coordinator_invocation_id="coordinator-2", + store=store, + artifacts=artifacts, + invoker=second_invoker, + aggregator=lambda **_: aggregate_payload({"budget": {"total": 100.0}}), + ) + + assert result == { + "deduplicated": True, + "evaluation_id": str(EVALUATION_ID), + "status": "succeeded", + } + assert second_invoker.events == [] + assert store.parent.comparison_status is ResultComparisonStatus.FAILED + assert store.parent.comparison_completed_at == failed_at + assert store.parent.comparison_error_code == "result_comparison_failed" + assert store.parent.comparison_error_summary == failed_summary + assert artifacts.comparison_writes == [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordinator_failure_records.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordinator_failure_records.py new file mode 100644 index 000000000..88821f1eb --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordinator_failure_records.py @@ -0,0 +1,402 @@ +"""Tests for the failure records and reports of a failed Stage 12 coordination.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, +) +from sqlalchemy.exc import OperationalError +from stage12_child_call_support import ( + ConcurrentInvoker, + InjectedChildFaults, + RefusingDispatchInvoker, + TimeoutInvoker, +) +from stage12_fake_stores import FakeArtifacts, FakeStore +from stage12_runtime_support import ( + BASELINE_ID, + EVALUATION_ID, + NOW, + RecordingRuntime, + TelemetryRecordingRuntime, + aggregate_payload, + child_record, + connection_limit_error, + coordinate_report, + invocation_context, + parent_record, + report_execution_input, + summary_correlation_id, + summary_failure_cause, +) + + +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_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_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", + "evaluation_id": str(EVALUATION_ID), + "correlation_id": summary_correlation_id(failure_write[1]), + } + assert original_report[1][2] is aggregation_error + assert failure_write[0] == "failure_write" + assert summary_failure_cause(failure_write[1]) == ( + "ValueError in stage12_aggregation: aggregate failed validation" + ) + 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_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(fail_role="baseline") + ), + aggregator=lambda **_: aggregate_payload({"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", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "correlation_id": summary_correlation_id(store.parent.error_summary), + } + ] + 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_code == "simulation_invocation_failed" + assert summary_failure_cause(store.parent.error_summary) == ( + "RuntimeError in stage12_child_wait (baseline): child calculation failed" + ) + + +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_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=invoker, + aggregator=lambda **_: aggregate_payload({"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" + ] + baseline_child = store.children[BASELINE_ID] + assert [(attributes, error) for _message, attributes, error in failure_reports] == [ + ( + { + "stage12_failure_scope": "child_dispatch", + "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "correlation_id": summary_correlation_id(baseline_child.error_summary), + }, + dispatch_error, + ) + ] + assert invoker.events == ["spawn:baseline"] + assert baseline_child.status is ComparisonRunLifecycleStatus.FAILED + assert baseline_child.error_code == "simulation_dispatch_failed" + assert summary_failure_cause(baseline_child.error_summary) == ( + "ConnectionError in stage12_child_dispatch (baseline): " + "Modal refused the dispatch" + ) + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_summary == baseline_child.error_summary + + +def test_coordinator_records_a_redacted_dispatch_failure_cause() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = RefusingDispatchInvoker( + ConnectionError("Modal rejected the call with token=opaque-token") + ) + runtime = RecordingRuntime() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + runtime=runtime, + ) + + assert invoker.events == ["spawn:baseline"] + assert artifacts.aggregate_writes == [] + failed_child = store.children[BASELINE_ID] + assert failed_child.status is ComparisonRunLifecycleStatus.FAILED + assert failed_child.error_code == "simulation_dispatch_failed" + assert summary_failure_cause(failed_child.error_summary) == ( + "ConnectionError in stage12_child_dispatch (baseline): " + "Modal rejected the call with token=[REDACTED]" + ) + assert store.parent.error_code == "simulation_dispatch_failed" + assert store.parent.error_summary == failed_child.error_summary + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == failed_child.error_summary + + +def test_coordinator_rejects_incompatible_row_identity() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(incompatible=True) + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + ) + assert artifacts.aggregate_writes == [] + + +def test_coordinator_reuses_matching_successful_child() -> None: + report = report_execution_input() + store = FakeStore() + artifacts = FakeArtifacts() + baseline_descriptor = artifacts.add_simulation(report.baseline) + baseline = child_record(report.baseline).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": baseline_descriptor.artifact.uri, + "output_sha256": baseline_descriptor.artifact.content_sha256, + "output_schema_version": 1, + "row_identity_columns": baseline_descriptor.row_identity.identifier_columns, + "row_count": baseline_descriptor.row_identity.row_count, + "row_identity_sha256": baseline_descriptor.row_identity.identity_sha256, + "completed_at": NOW, + } + ) + store.children[report.baseline.simulation_execution_id] = baseline + invoker = ConcurrentInvoker(artifacts) + + coordinate_report( + report.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: aggregate_payload(), + ) + + assert invoker.events == ["spawn:reform"] + + +def test_coordinator_resumes_a_matching_running_child_invocation() -> None: + report = report_execution_input() + store = FakeStore() + artifacts = FakeArtifacts() + store.children[report.baseline.simulation_execution_id] = child_record( + report.baseline + ).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": "call-baseline", + "started_at": NOW, + } + ) + invoker = ConcurrentInvoker(artifacts) + + coordinate_report( + report.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: aggregate_payload(), + ) + + assert invoker.events == ["restore:baseline", "spawn:reform"] + + +def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = TimeoutInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + wait_timeout_seconds=0, + ) + + assert "sensitive" not in str(error.value) + + assert invoker.events == ["spawn:baseline", "spawn:reform"] + assert artifacts.aggregate_writes == [] + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_code == "simulation_wait_timeout" + assert store.parent.error_summary == ( + "Stage 12 simulations did not finish within 0 seconds of the report " + "coordinator starting." + ) + assert all( + child.status is ComparisonRunLifecycleStatus.FAILED + for child in store.children.values() + ) + assert all( + child.error_code == "simulation_cancelled" for child in store.children.values() + ) + + +def test_closed_wait_window_is_reported_before_its_failure_records() -> None: + ordered_events: list[tuple[str, object]] = [] + + class UnwritableFailureStore(FakeStore): + def replace_simulation_if_status(self, record, *, expected_status): + if record.status is ComparisonRunLifecycleStatus.FAILED: + ordered_events.append(("failure_write", record.error_code)) + raise connection_limit_error() + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + def replace_report(self, record): + if record.status is ComparisonRunLifecycleStatus.FAILED: + ordered_events.append(("failure_write", record.error_code)) + raise connection_limit_error() + return super().replace_report(record) + + store = UnwritableFailureStore() + invoker = TimeoutInvoker() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=invoker, + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + runtime=TelemetryRecordingRuntime(ordered_events), + wait_timeout_seconds=0, + ) + + original_report, *failure_writes, persistence_report = ordered_events + event_kind, (message, attributes, reported_error) = original_report + assert event_kind == "telemetry" + assert message == "Stage 12 execution failed" + assert attributes == {"stage12_failure_scope": "child_wait_timeout"} + assert reported_error.detail.error_code == "simulation_wait_timeout" + # Both cancellation records and the report record were attempted after + # the closed window had already been reported. + assert failure_writes == [ + ("failure_write", "simulation_cancelled"), + ("failure_write", "simulation_cancelled"), + ("failure_write", "simulation_wait_timeout"), + ] + assert persistence_report[1][0] == ( + "Stage 12 failure record could not be persisted" + ) + assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_coordinator_failures.py b/projects/policyengine-simulation-executor/tests/test_stage12_coordinator_failures.py new file mode 100644 index 000000000..6430b3434 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_coordinator_failures.py @@ -0,0 +1,527 @@ +"""Tests for how the Stage 12 coordinator handles a failing child or artifact.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from typing import override + +import pytest +from policyengine_simulation_contract.spm import SPMComparisonProvenance +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + PlannedSimulationExecutionInput, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + coordination as coordination_module, +) +from sqlalchemy.exc import OperationalError +from stage12_child_call_support import ( + ConcurrentInvoker, + ImmediateCall, + InjectedChildFaults, + TimeoutCall, +) +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames +from stage12_runtime_support import ( + BASELINE_ID, + NOW, + REFORM_ID, + RecordingRuntime, + aggregate_payload, + coordinate_report, + invocation_context, + parent_record, + report_execution_input, + summary_failure_cause, +) + + +def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(fail_role="reform") + ) + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + ) + + assert "child calculation failed" not in str(error.value) + + assert artifacts.aggregate_writes == [] + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert summary_failure_cause(store.parent.error_summary) == ( + "RuntimeError in stage12_child_wait (reform): child calculation failed" + ) + failed_child = store.children[REFORM_ID] + assert failed_child.status is ComparisonRunLifecycleStatus.FAILED + assert failed_child.error_code == "simulation_invocation_failed" + assert failed_child.error_summary == store.parent.error_summary + + +@pytest.mark.parametrize( + ("failed_role", "failed_id", "pending_role", "pending_id"), + [ + ("reform", REFORM_ID, "baseline", BASELINE_ID), + ("baseline", BASELINE_ID, "reform", REFORM_ID), + ], +) +def test_coordinator_observes_one_role_failure_while_the_peer_is_pending( + failed_role, + failed_id, + pending_role, + pending_id, +) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + + class PendingCall: + object_id = f"call-{pending_role}" + + def __init__(self): + self.timeouts = [] + self.cancel_count = 0 + + def get(self, *, timeout=None): + self.timeouts.append(timeout) + raise TimeoutError + + def cancel(self): + self.cancel_count += 1 + + class FailedCall: + object_id = f"call-{failed_role}" + + def __init__(self): + self.timeouts = [] + + def get(self, *, timeout=None): + self.timeouts.append(timeout) + child = store.children[failed_id] + failed_at = datetime.now(UTC) + store.replace_simulation( + child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + "updated_at": failed_at, + "completed_at": failed_at, + } + ) + ) + msg = "remote wrapper" + raise RuntimeError(msg) + + def cancel(self): + return None + + baseline_call = PendingCall() + failed_call = FailedCall() + + class ConcurrentFailureInvoker: + def spawn(self, request): + role = PlannedSimulationExecutionInput.model_validate( + request.simulation + ).role + return failed_call if role.value == failed_role else baseline_call + + def restore(self, invocation_id): + msg = f"unexpected restored invocation {invocation_id}" + raise AssertionError(msg) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=ConcurrentFailureInvoker(), + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + sleep=lambda _: pytest.fail("failure should be observed before sleeping"), + ) + + assert failed_call.timeouts == [0] + assert baseline_call.timeouts == ([0] if pending_role == "baseline" else []) + # A segmented peer reads its cancelled record and stops its own + # region-group calls, so it is not cancelled through Modal. + assert baseline_call.cancel_count == 0 + pending = store.children[pending_id] + assert pending.status is ComparisonRunLifecycleStatus.FAILED + assert pending.error_code == "simulation_cancelled" + assert pending.error_summary == ( + f"Cancelled after the {failed_role} simulation failed." + ) + assert store.parent.error_code == "invalid_policy_period" + assert store.parent.error_summary == "The reform period is invalid." + assert artifacts.aggregate_writes == [] + + +def test_dispatch_failure_cancels_an_already_started_peer() -> None: + store = FakeStore() + baseline_call = TimeoutCall("call-baseline") + + class DispatchFailureInvoker: + def spawn(self, request): + role = PlannedSimulationExecutionInput.model_validate( + request.simulation + ).role + if role is SimulationRole.REFORM: + msg = "dispatch contains sensitive details" + raise ConnectionError(msg) + return baseline_call + + def restore(self, invocation_id): + msg = f"unexpected restored invocation {invocation_id}" + raise AssertionError(msg) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=DispatchFailureInvoker(), + ) + + baseline = store.children[BASELINE_ID] + reform = store.children[REFORM_ID] + assert baseline_call.cancelled == 0 + assert baseline.status is ComparisonRunLifecycleStatus.FAILED + assert baseline.error_code == "simulation_cancelled" + assert baseline.error_summary == "Cancelled after the reform simulation failed." + assert reform.status is ComparisonRunLifecycleStatus.FAILED + assert reform.error_code == "simulation_dispatch_failed" + assert store.parent.error_code == "simulation_dispatch_failed" + + +def test_coordinator_propagates_an_already_persisted_child_failure() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + + class PersistedFailureCall: + object_id = "call-reform" + + def get(self, *, timeout=None): + child = store.children[REFORM_ID] + failed_at = datetime.now(UTC) + store.replace_simulation( + child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + "updated_at": failed_at, + "completed_at": failed_at, + } + ) + ) + msg = "remote wrapper must not replace the child failure" + raise RuntimeError(msg) + + class PersistedFailureInvoker: + def spawn(self, request): + parsed = PlannedSimulationExecutionInput.model_validate(request.simulation) + if parsed.role is SimulationRole.REFORM: + return PersistedFailureCall() + descriptor = artifacts.add_simulation(parsed, simulation_frames(100.0)) + return ImmediateCall( + "call-baseline", descriptor.model_dump(mode="json"), [] + ) + + def restore(self, invocation_id): + msg = f"unexpected restored invocation {invocation_id}" + raise AssertionError(msg) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=PersistedFailureInvoker(), + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + ) + + assert store.children[REFORM_ID].error_code == "invalid_policy_period" + assert store.children[REFORM_ID].error_summary == "The reform period is invalid." + assert store.parent.error_code == "invalid_policy_period" + assert store.parent.error_summary == "The reform period is invalid." + + +@pytest.mark.parametrize( + "invoker", + [ + lambda artifacts: ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(wrong_plan_digest=True) + ), + lambda artifacts: ConcurrentInvoker( + artifacts, faults=InjectedChildFaults(missing_column=True) + ), + ], +) +def test_coordinator_rejects_artifacts_that_do_not_satisfy_the_shared_plan( + invoker, +) -> None: + store = FakeStore() + artifacts = FakeArtifacts() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker(artifacts), + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + ) + + assert artifacts.aggregate_writes == [] + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + + +def test_coordinator_records_failure_when_child_persistence_fails() -> None: + class FailingChildStore(FakeStore): + @override + def create_or_resolve_simulation(self, record): + msg = "database rejected the child record" + raise RuntimeError(msg) + + store = FailingChildStore() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=ConcurrentInvoker(FakeArtifacts()), + ) + + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_code == "report_coordination_failed" + assert summary_failure_cause(store.parent.error_summary) == ( + "RuntimeError in stage12_child_state_create: database rejected the child record" + ) + + +class WorkerRecordedFailureCall: + """A child whose worker persisted its cause before Modal raised.""" + + def __init__(self, object_id, store, simulation_execution_id, worker_summary): + self.object_id = object_id + self.store = store + self.simulation_execution_id = simulation_execution_id + self.worker_summary = worker_summary + + def get(self, *, timeout=None): + child = self.store.children[self.simulation_execution_id] + self.store.children[self.simulation_execution_id] = child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_execution_failed", + "error_summary": self.worker_summary, + "completed_at": NOW, + } + ) + msg = "Stage 12 simulation execution failed" + raise RuntimeError(msg) + + +class UnreachedCall: + def __init__(self, object_id): + self.object_id = object_id + + def get(self, *, timeout=None): + msg = "the coordinator observes baseline first" + raise AssertionError(msg) + + def cancel(self): + return None + + +class WorkerRecordedFailureInvoker: + def __init__(self, store, worker_summary): + self.store = store + self.worker_summary = worker_summary + + def spawn(self, request): + planned = PlannedSimulationExecutionInput.model_validate(request.simulation) + if planned.role is SimulationRole.BASELINE: + return WorkerRecordedFailureCall( + "call-baseline", + self.store, + planned.simulation_execution_id, + self.worker_summary, + ) + return UnreachedCall("call-reform") + + def restore(self, invocation_id): + msg = f"unexpected restored invocation {invocation_id}" + raise AssertionError(msg) + + +def test_coordinator_keeps_the_cause_recorded_by_a_failed_worker() -> None: + store = FakeStore() + worker_summary = ( + "ValueError in stage12_calculation (baseline): " + "Variable 'country' not found in household data " + "(correlation_id=0123456789abcdef0123456789abcdef)" + ) + runtime = RecordingRuntime() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=WorkerRecordedFailureInvoker(store, worker_summary), + aggregator=lambda **_: aggregate_payload({"must": "not run"}), + runtime=runtime, + ) + + failed_child = store.children[BASELINE_ID] + assert failed_child.error_code == "simulation_execution_failed" + assert failed_child.error_summary == worker_summary + assert store.parent.error_code == "simulation_execution_failed" + assert store.parent.error_summary == worker_summary + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == worker_summary + + +def test_coordinator_names_the_stage_that_rejected_the_aggregate() -> None: + class RejectingAggregateArtifacts(FakeArtifacts): + def write_aggregate(self, *, prefix, payload): + SPMComparisonProvenance.model_validate(payload["result"]["spm_provenance"]) + + store = FakeStore() + artifacts = RejectingAggregateArtifacts() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + report_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=lambda **_: {"result": {"spm_provenance": None}}, + ) + + assert summary_failure_cause(store.parent.error_summary) == ( + "ValidationError in stage12_aggregate_artifact_write: " + "1 validation error for SPMComparisonProvenance: " + "Input should be a valid dictionary or instance " + "of SPMComparisonProvenance [model_type]" + ) + + +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_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_context().modal_application, + coordinator_invocation_id="coordinator-1", + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=lambda **_: aggregate_payload(), + 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() -> ( # noqa: E501 # the test name spells out the lost-commit scenario + 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_execution_input().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + parent_record().model_dump(mode="json"), + application_name=invocation_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 == [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_dataset_provenance.py b/projects/policyengine-simulation-executor/tests/test_stage12_dataset_provenance.py new file mode 100644 index 000000000..aa913b906 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_dataset_provenance.py @@ -0,0 +1,335 @@ +"""Tests for the dataset a Stage 12 calculation loads, and the provenance it checks.""" + +from __future__ import annotations + +import pandas as pd +import pytest +from policyengine_simulation_contract.stage12_execution import ( + BundleProvenance, + DatasetArtifactMediaType, + DatasetArtifactReference, + DatasetPopulationInput, + DatasetProvenance, + GeographySelection, + SimulationRole, +) +from policyengine_simulation_contract.uk_geography import ( + UKLocalAuthorityBoundaryVersion, + UKLocalAuthorityMetadata, +) +from policyengine_simulation_executor.stage12_runtime.aggregation import ( + validate_uk_local_authority_metadata, +) +from policyengine_simulation_executor.uk_local_authority_metadata import ( + detect_uk_local_authority_metadata, +) +from stage12_runtime_support import ( + bundle_provenance, + planned_simulation, + run_until_dataset_loaded, +) + +REGIONAL_URI = "hf://policyengine/data/populace_us_2024_acs_local.h5@regional-revision" + + +def _regional_provenance() -> BundleProvenance: + return bundle_provenance().model_copy( + update={ + "dataset": DatasetProvenance( + identity="populace_us_2024_acs_local", + uri=REGIONAL_URI, + artifact_revision="regional-revision", + data_package_name="populace-data", + data_package_version="0.1.0", + ) + } + ) + + +def _installed_bundle_with_regions(): + from policyengine_simulation_contract.stage12_bundle import ( + ResolvedStage12Bundle, + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, + Stage12RegionDataset, + ) + + national = Stage12Dataset( + identity="populace_us_2024", + uri=bundle_provenance().dataset.uri, + artifact_revision="revision", + sha256="a" * 64, + repo_type="dataset", + ) + regional = Stage12Dataset( + identity="populace_us_2024_acs_local", + uri=REGIONAL_URI, + artifact_revision="regional-revision", + sha256="c" * 64, + repo_type="dataset", + ) + return ResolvedStage12Bundle( + bundle=Stage12BundleManifest( + policyengine_version="5.2.0", + policyengine_requirement="policyengine==5.2.0", + core_package_version="3.0.0", + core_package_requirement="policyengine-core==3.0.0", + countries=( + Stage12CountryBundle( + country="us", + country_package_name="policyengine-us", + country_package_version="1.764.6", + country_package_requirement="policyengine-us==1.764.6", + data_package_name="populace-data", + data_package_version="0.1.0", + data_release_version="revision", + data_artifact_revision="revision", + default_dataset=national.identity, + default_dataset_uri=national.uri, + datasets=(national, regional), + region_datasets=( + Stage12RegionDataset( + region_type="national", dataset=national.identity + ), + Stage12RegionDataset( + region_type="state", dataset=regional.identity + ), + Stage12RegionDataset( + region_type="congressional_district", + dataset=regional.identity, + ), + ), + ), + ), + ), + bundle_manifest_sha256="b" * 64, + ) + + +def _simulation_for(region: str, provenance: BundleProvenance): + simulation = planned_simulation(SimulationRole.BASELINE) + return simulation.model_copy( + update={ + "geography": GeographySelection(country="us", region=region), + "bundle": provenance, + "population": DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=provenance.dataset.uri, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256="a" * 64, + ) + ), + } + ) + + +@pytest.mark.parametrize( + ("region", "provenance"), + [ + ("state/ut", _regional_provenance), + ("State/UT", _regional_provenance), + ("ut", _regional_provenance), + ("congressional_district/UT-01", _regional_provenance), + ("us", bundle_provenance), + # Entry services released before regional provenance send the default + # dataset for state runs until they are updated. + ("state/ut", bundle_provenance), + ], +) +def test_installed_bundle_accepts_the_dataset_a_region_loads( + monkeypatch, region, provenance +) -> None: + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + monkeypatch.setattr(worker, "load_stage12_bundle", _installed_bundle_with_regions) + + worker._require_installed_bundle(_simulation_for(region, provenance())) + + +def test_uk_regions_keep_the_default_dataset_when_the_bundle_names_region_types( + monkeypatch, +) -> None: + from policyengine_simulation_contract.stage12_bundle import Stage12RegionDataset + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + resolved = _installed_bundle_with_regions() + us = resolved.bundle.countries[0] + uk = us.model_copy( + update={ + "country": "uk", + "region_datasets": ( + Stage12RegionDataset( + region_type="national", dataset=us.default_dataset + ), + ), + } + ) + resolved = resolved.model_copy( + update={"bundle": resolved.bundle.model_copy(update={"countries": (us, uk)})} + ) + monkeypatch.setattr(worker, "load_stage12_bundle", lambda: resolved) + simulation = _simulation_for("us", bundle_provenance()).model_copy( + update={"geography": GeographySelection(country="uk", region="country/wales")} + ) + + worker._require_installed_bundle(simulation) + with pytest.raises(RuntimeError, match="differs from the installed bundle"): + worker._require_installed_bundle( + simulation.model_copy(update={"bundle": _regional_provenance()}) + ) + + +@pytest.mark.parametrize( + ("region", "dataset"), + [ + ("us", "populace_us_2024_acs_local"), + ("state/ut", "source_enrichment"), + ], +) +def test_installed_bundle_refuses_a_dataset_the_region_does_not_load( + monkeypatch, region, dataset +) -> None: + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + monkeypatch.setattr(worker, "load_stage12_bundle", _installed_bundle_with_regions) + provenance = _regional_provenance() + provenance = provenance.model_copy( + update={"dataset": provenance.dataset.model_copy(update={"identity": dataset})} + ) + + with pytest.raises(RuntimeError, match="differs from the installed bundle"): + worker._require_installed_bundle(_simulation_for(region, provenance)) + + +def _selection(name: str, uri: str, *, is_default: bool): + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + + return DatasetSelection(name=name, uri=uri, is_default=is_default) + + +def test_a_calculation_loads_the_dataset_its_provenance_names(monkeypatch) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + loaded = run_until_dataset_loaded( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_selection( + "populace_us_2024_acs_local", REGIONAL_URI, is_default=False + ), + baked_folder=None, + planned=_simulation_for("state/ut", _regional_provenance()), + ) + + assert loaded.loads[0]["selection"].name == "populace_us_2024_acs_local" + + +def test_a_calculation_refuses_provenance_naming_another_dataset( + monkeypatch, +) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + with pytest.raises(RuntimeError, match="dataset its region loads"): + run_until_dataset_loaded( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_selection( + "populace_us_2024", bundle_provenance().dataset.uri, is_default=True + ), + baked_folder=None, + planned=_simulation_for("state/ut", _regional_provenance()), + ) + + +def test_default_provenance_from_an_older_entry_service_is_still_loaded( + monkeypatch, +) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + loaded = run_until_dataset_loaded( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_selection( + "populace_us_2024_acs_local", REGIONAL_URI, is_default=False + ), + baked_folder=None, + planned=_simulation_for("state/ut", bundle_provenance()), + ) + + assert loaded.loads[0]["selection"].name == "populace_us_2024_acs_local" + + +def test_uk_boundary_version_is_detected_from_the_unscoped_dataset() -> None: + dataset = type( + "Dataset", + (), + { + "data": type( + "Data", + (), + { + "entity_data": { + "household": pd.DataFrame( + {"la_code_oa": ["E06000001", "E07000026"]} + ) + } + }, + )() + }, + )() + + metadata = detect_uk_local_authority_metadata("uk", dataset) + + assert metadata == UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 + ) + + +def test_country_metadata_alignment_rejects_missing_or_mismatched_uk_values() -> None: + lad22 = UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 + ) + lad23 = UKLocalAuthorityMetadata( + boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 + ) + + with pytest.raises(ValueError, match="missing"): + validate_uk_local_authority_metadata("uk", None, lad22) + with pytest.raises(ValueError, match="do not match"): + validate_uk_local_authority_metadata("uk", lad22, lad23) + assert validate_uk_local_authority_metadata("uk", lad22, lad22) == lad22 + assert validate_uk_local_authority_metadata("us", None, None) is None + + +@pytest.mark.parametrize( + ("region_code", "routed", "accepted"), + [ + ("uk", True, True), + ("country/england", True, True), + ("uk", False, False), + ("constituency/E14001063", True, False), + ("local_authority/E06000063", True, False), + (None, True, False), + ], +) +def test_uk_artifacts_without_metadata_need_a_routed_national_request( + monkeypatch, region_code: str | None, *, routed: bool, accepted: bool +) -> None: + monkeypatch.setattr( + "policyengine_simulation_executor.uk_local_authority_metadata." + "uk_area_regions_routed", + lambda: routed, + ) + + if accepted: + assert ( + validate_uk_local_authority_metadata( + "uk", None, None, region_code=region_code + ) + is None + ) + else: + with pytest.raises(ValueError, match="metadata is missing"): + validate_uk_local_authority_metadata( + "uk", None, None, region_code=region_code + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py index 1b81480d7..5b67d0e5f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py @@ -25,7 +25,7 @@ from policyengine_simulation_executor.stage12_runtime.partition import ( US_REGION_GROUPS, ) -from test_stage12_runtime import _planned_simulation, _report +from stage12_runtime_support import planned_simulation, report_execution_input def _frames(value: float) -> dict[str, pd.DataFrame]: @@ -56,7 +56,7 @@ def __init__(self, frames: dict[str, pd.DataFrame]): def _existing_request() -> dict[str, Any]: - report = _report() + report = report_execution_input() return { "country": "us", "scope": "macro", @@ -70,7 +70,7 @@ def _existing_request() -> dict[str, Any]: def test_qualification_runner_exercises_all_segments_and_exact_merge() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) calls = [] def run_segment(segment): @@ -108,7 +108,7 @@ def run_segment(segment): def test_qualification_runs_matching_inputs_through_both_complete_paths() -> None: - report = _report() + report = report_execution_input() baseline = _frames(100.0) reform = _frames(120.0) simulation_calls = [] @@ -173,7 +173,7 @@ def aggregate_builder(**arguments): def test_qualification_fails_with_bounded_diagnostic_without_output_values() -> None: - report = _report() + report = report_execution_input() baseline = _frames(100.0) reform = _frames(120.0) @@ -221,14 +221,14 @@ def test_qualification_rejects_different_input_provenance(field, value, code) -> request = {**_existing_request(), field: value} with pytest.raises(ParityMismatchError) as error: - qualify_report_parity(request, _report()) + qualify_report_parity(request, report_execution_input()) assert error.value.code == code assert error.value.path == field def test_qualification_rejects_existing_executor_provenance_mismatch() -> None: - report = _report() + report = report_execution_input() baseline = _frames(100.0) reform = _frames(120.0) @@ -282,7 +282,7 @@ def single_runner(simulation): def test_qualification_receipt_records_every_digest_and_tolerance_field() -> None: from policyengine_simulation_executor.stage12_parity import NumericalTolerance - report = _report() + report = report_execution_input() baseline = _frames(100.0) reform = _frames(120.0) @@ -362,7 +362,7 @@ def test_qualification_receipt_records_every_digest_and_tolerance_field() -> Non def test_existing_result_is_checked_before_v2_simulations_run( overrides, code, path ) -> None: - report = _report() + report = report_execution_input() baseline = _frames(100.0) reform = _frames(120.0) calls = [] @@ -385,7 +385,7 @@ def test_existing_result_is_checked_before_v2_simulations_run( @pytest.mark.parametrize("v2_report", [{}, {"result": None}, {"result": []}]) def test_v2_aggregate_without_a_result_object_is_rejected(v2_report) -> None: - report = _report() + report = report_execution_input() baseline = _frames(100.0) reform = _frames(120.0) @@ -415,7 +415,7 @@ def test_default_existing_runner_uses_and_shuts_down_a_local_runtime( stage12_qualification, ) - report = _report() + report = report_execution_input() baseline = _frames(100.0) reform = _frames(120.0) events = [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index eca8612bf..2d5134ab2 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -1,540 +1,23 @@ -"""Tests for independent Stage 12 simulations and report coordination.""" +"""Tests for how a Stage 12 simulation loads its dataset and region.""" from __future__ import annotations -import importlib -import io -import json -import logging -import re -import sys -import time -from concurrent.futures import ThreadPoolExecutor -from contextlib import contextmanager, nullcontext -from dataclasses import dataclass -from datetime import UTC, datetime, timedelta -from hashlib import sha256 -from threading import Lock +from contextlib import nullcontext from types import SimpleNamespace -from typing import ClassVar, override -from uuid import UUID, uuid4 +from typing import ClassVar import pandas as pd import pytest -from fixtures.fake_modal import install_fake_modal -from fixtures.fake_object_store import InMemoryObjectStore -from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT -from policyengine_observability import ( - DeploymentIdentity, - LoggingConfig, - ObservabilityConfig, - OTelConfig, - ServiceIdentity, - StdoutLogDestination, - configure, -) -from policyengine_simulation_contract.macro_output import SingleYearMacroOutput -from policyengine_simulation_contract.spm import ( - SPMComparisonProvenance, - SPMInputError, - SPMRuntimeVersions, - SPMSelection, - build_spm_provenance, - validate_spm_calculation_provenance, -) from policyengine_simulation_contract.stage12_execution import ( - AggregateReportArtifactPayload, - ArtifactMediaType, - ArtifactReference, - BundleProvenance, - ComparisonReportPersistenceResult, - ComparisonReportRecord, - ComparisonRunAggregationStatus, - ComparisonRunLifecycleStatus, - ComparisonSimulationPersistenceResult, - DatasetArtifactMediaType, - DatasetArtifactReference, - DatasetPopulationInput, - DatasetProvenance, - EntityOutputPlan, - GeographySelection, PlannedSimulationExecutionInput, - ReportAggregate, - ReportExecutionInput, - ReportOutputRequirements, - ResultComparisonStatus, - RowIdentity, - SimulationArtifactDescriptor, - SimulationExecutionInput, SimulationRole, - Stage12InvocationContext, - Stage12OutputPlan, - stage12_output_plan_sha256, -) -from policyengine_simulation_contract.uk_geography import ( - UKLocalAuthorityBoundaryVersion, - UKLocalAuthorityMetadata, -) -from policyengine_simulation_executor.stage12_artifacts import ( - Stage12ArtifactStore, - canonical_json_bytes, - serialize_simulation_frames, -) -from policyengine_simulation_executor.stage12_runtime import ( - CoordinatorDependencies, - CoordinatorInvocation, - RetainedSimulationOutput, - SegmentRunnerSeams, - SimulationCalculation, - SimulationWorkerDependencies, - Stage12SegmentRunner, - _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 stage12_fake_stores import simulation_frames +from stage12_runtime_support import ( + bundle_provenance, + planned_simulation, + run_until_dataset_loaded, ) -from policyengine_simulation_executor.stage12_runtime.comparison import ( - PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, -) -from policyengine_simulation_executor.stage12_runtime.dependencies import ( - runtime_store, -) -from policyengine_simulation_executor.stage12_runtime.partition import ( - US_REGION_GROUPS, -) -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") -BASELINE_ID = UUID("00000000-0000-0000-0000-000000000002") -REFORM_ID = UUID("00000000-0000-0000-0000-000000000003") - - -def _spm_calculation_provenance(): - selection = SPMSelection( - forecast_content_sha256="f" * 64, - scenario="official", - geography_kind="national", - geography_id=None, - county_vintage="2020", - as_of=None, - ) - receipt = build_spm_provenance( - forecast_id="forecast-1", - forecast_sha256="f" * 64, - selection=selection, - years=("2026",), - runtime_versions=SPMRuntimeVersions.model_validate( - { - "policyengine": "6.2.1", - "policyengine-core": "3.32.10", - "policyengine-us": "2.2.1", - "spm-calculator": "1.0.0", - } - ), - ) - return validate_spm_calculation_provenance(config=selection, receipt=receipt) - - -def _bundle() -> BundleProvenance: - return BundleProvenance( - policyengine_version="5.2.0", - country_package_name="policyengine-us", - country_package_version="1.764.6", - dataset=DatasetProvenance( - identity="populace_us_2024", - uri="hf://policyengine/data/populace_us_2024.h5@revision", - artifact_revision="revision", - data_package_name="populace-data", - data_package_version="0.1.0", - ), - bundle_manifest_sha256="b" * 64, - ) - - -def _simulation(role: SimulationRole) -> SimulationExecutionInput: - return SimulationExecutionInput( - evaluation_id=EVALUATION_ID, - simulation_execution_id=( - BASELINE_ID if role is SimulationRole.BASELINE else REFORM_ID - ), - role=role, - policy={} if role is SimulationRole.BASELINE else {"reform": {"2026": 1}}, - population=DatasetPopulationInput( - artifact=DatasetArtifactReference( - uri=_bundle().dataset.uri, - media_type=DatasetArtifactMediaType.HDF5, - content_sha256="a" * 64, - ) - ), - year=2026, - geography=GeographySelection(country="us", region="us"), - bundle=_bundle(), - ) - - -def _output_plan() -> Stage12OutputPlan: - return Stage12OutputPlan( - country="us", - requirements=ReportOutputRequirements( - aggregates=tuple(ReportAggregate), - include_cliff_impacts=False, - labor_supply_response_active=False, - ), - entities=( - EntityOutputPlan( - entity="household", - materialized_variables=("household_id", "household_net_income"), - ), - EntityOutputPlan( - entity="person", - materialized_variables=("age", "household_id", "person_id"), - ), - ), - ) - - -def _planned_simulation(role: SimulationRole) -> PlannedSimulationExecutionInput: - return PlannedSimulationExecutionInput.model_validate( - { - **_simulation(role).model_dump(mode="json"), - "output_plan": _output_plan().model_dump(mode="json"), - } - ) - - -def _report() -> ReportExecutionInput: - return ReportExecutionInput( - evaluation_id=EVALUATION_ID, - baseline=_simulation(SimulationRole.BASELINE), - reform=_simulation(SimulationRole.REFORM), - requested_aggregates=tuple(ReportAggregate), - ) - - -def _aggregate_payload(result=None) -> dict: - """Return a stand-in aggregate that satisfies the artifact contract. - - Aggregators that must not run return one too, so an unexpected run shows - up as a written aggregate instead of a rejected one. - """ - report = _report() - return AggregateReportArtifactPayload( - evaluation_id=report.evaluation_id, - requested_aggregates=report.requested_aggregates, - bundle=report.baseline.bundle, - result={} if result is None else result, - ).model_dump(mode="json") - - -def coordinate_report(*payloads, application_name, coordinator_invocation_id, **seams): - seams.setdefault("output_plan_resolver", lambda _: _output_plan()) - return coordinate_report_impl( - *payloads, - invocation=CoordinatorInvocation( - application_name=application_name, - invocation_id=coordinator_invocation_id, - ), - dependencies=CoordinatorDependencies(**seams), - ) - - -class RecordingRuntime: - def __init__(self) -> None: - self.operations: list[tuple[str, dict[str, object]]] = [] - self.recorded_exceptions: list[Exception] = [] - - @contextmanager - def operation(self, name, *, attributes=None, **_kwargs): - self.operations.append((name, dict(attributes or {}))) - yield - - def span(self, *_args, **_kwargs): - return nullcontext() - - def capture_context(self): - return {} - - def record_exception(self, error, *, handled, status_code=None): - assert handled - self.recorded_exceptions.append(error) - - -def _context() -> Stage12InvocationContext: - return Stage12InvocationContext( - request_id="request-1", - environment="staging", - modal_environment="staging", - worker_version="5.2.0", - modal_application="policyengine-simulation-v2-py5-2-0", - simulation_callable="run_single_simulation_us", - version_manifest_sha256="b" * 64, - bundle_manifest_sha256="b" * 64, - artifact_prefix=( - "stage-12-runs/staging/2026/09/00000000-0000-0000-0000-000000000001" - ), - created_at=NOW, - retention_expires_at=NOW + timedelta(days=30), - ) - - -def _parent() -> ComparisonReportRecord: - return ComparisonReportRecord( - evaluation_id=EVALUATION_ID, - status=ComparisonRunLifecycleStatus.PENDING, - aggregation_status=ComparisonRunAggregationStatus.NOT_STARTED, - environment="staging", - calculation_flow="economy", - originating_request_id="request-1", - production_identity=f"direct:{EVALUATION_ID}", - incumbent_execution_id=None, - worker_version="5.2.0", - modal_application="policyengine-simulation-v2-py5-2-0", - report_coordinator_callable="coordinate_report", - version_manifest_sha256="b" * 64, - policyengine_version="5.2.0", - country_package_name="policyengine-us", - country_package_version="1.764.6", - country="us", - dataset_identity="populace_us_2024", - dataset_uri=_bundle().dataset.uri, - data_package_name="populace-data", - data_package_version="0.1.0", - data_artifact_revision="revision", - created_at=NOW, - updated_at=NOW, - retention_expires_at=NOW + timedelta(days=30), - ) - - -def _child(simulation: SimulationExecutionInput): - from policyengine_simulation_contract.stage12_execution import ( - ComparisonSimulationRecord, - ) - - return ComparisonSimulationRecord( - simulation_execution_id=simulation.simulation_execution_id, - evaluation_id=simulation.evaluation_id, - role=simulation.role, - input_sha256=simulation_input_sha256(simulation), - worker_version="5.2.0", - modal_application="policyengine-simulation-v2-py5-2-0", - simulation_callable="run_single_simulation_us", - version_manifest_sha256="b" * 64, - status=ComparisonRunLifecycleStatus.PENDING, - created_at=NOW, - updated_at=NOW, - retention_expires_at=NOW + timedelta(days=30), - ) - - -def _automatic_parent() -> ComparisonReportRecord: - return _parent().model_copy( - update={ - "production_identity": "production-job-1", - "incumbent_execution_id": "production-job-1", - "comparison_status": ResultComparisonStatus.PENDING, - } - ) - - -class FakeStore: - def __init__(self): - self.parent = None - self.children = {} - self.lock = Lock() - - def create_or_resolve_report(self, record): - with self.lock: - if self.parent is None: - self.parent = record - return ComparisonReportPersistenceResult( - record=record, - created=True, - ) - return ComparisonReportPersistenceResult( - record=self.parent, - created=False, - ) - - def get_report(self, evaluation_id): - assert evaluation_id == EVALUATION_ID - assert self.parent is not None - return self.parent - - def replace_report(self, record): - self.parent = record - return record - - def replace_report_result_comparison(self, record): - self.parent = record - return record - - def create_or_resolve_simulation(self, record): - with self.lock: - existing = self.children.get(record.simulation_execution_id) - if existing is None: - self.children[record.simulation_execution_id] = record - return ComparisonSimulationPersistenceResult( - record=record, created=True - ) - return ComparisonSimulationPersistenceResult( - record=existing, - created=False, - ) - - def get_simulation(self, simulation_execution_id): - return self.children[simulation_execution_id] - - def replace_simulation(self, record): - with self.lock: - self.children[record.simulation_execution_id] = record - return record - - def replace_simulation_if_status(self, record, *, expected_status): - with self.lock: - current = self.children[record.simulation_execution_id] - if current.status is not expected_status: - return current, False - self.children[record.simulation_execution_id] = record - return record, True - - def attach_simulation_invocation( - self, - simulation_execution_id, - *, - expected_placeholder, - modal_invocation_id, - updated_at, - ): - with self.lock: - child = self.children[simulation_execution_id] - assert child.modal_invocation_id == expected_placeholder - child = child.model_copy( - update={ - "modal_invocation_id": modal_invocation_id, - "updated_at": updated_at, - } - ) - self.children[simulation_execution_id] = child - return child - - -class FakeArtifacts: - def __init__(self): - self.payloads = {} - self.aggregate_writes = [] - self.comparison_writes = [] - self.input_writes = [] - # Aggregates go through the real store, so every coordinator test - # enforces the aggregate artifact contract that production enforces. - self.report_store = Stage12ArtifactStore("private", store=InMemoryObjectStore()) - - def write_input(self, *, prefix, simulation): - self.input_writes.append((prefix, simulation.role)) - - def add_simulation(self, simulation, frames=None): - payload, row_identity = serialize_simulation_frames(frames or _frames()) - digest = sha256(payload).hexdigest() - uri = f"gs://private/{simulation.role.value}.parquet" - self.payloads[uri] = payload - return SimulationArtifactDescriptor( - evaluation_id=simulation.evaluation_id, - simulation_execution_id=simulation.simulation_execution_id, - role=simulation.role, - artifact=ArtifactReference( - uri=uri, - media_type=ArtifactMediaType.PARQUET, - content_sha256=digest, - size_bytes=len(payload), - ), - output_plan_sha256=stage12_output_plan_sha256(_output_plan()), - row_identity=row_identity, - bundle=simulation.bundle, - ) - - def write_simulation( - self, - *, - prefix, - simulation, - frames, - calculation_provenance=None, - uk_local_authority_metadata=None, - ): - payload, row_identity = serialize_simulation_frames( - frames, - calculation_provenance=calculation_provenance, - uk_local_authority_metadata=uk_local_authority_metadata, - ) - uri = f"gs://private/{simulation.role.value}.parquet" - self.payloads[uri] = payload - descriptor = SimulationArtifactDescriptor( - evaluation_id=simulation.evaluation_id, - simulation_execution_id=simulation.simulation_execution_id, - role=simulation.role, - artifact=ArtifactReference( - uri=uri, - media_type=ArtifactMediaType.PARQUET, - content_sha256=sha256(payload).hexdigest(), - size_bytes=len(payload), - ), - output_plan_sha256=stage12_output_plan_sha256(_output_plan()), - row_identity=row_identity, - bundle=simulation.bundle, - ) - return descriptor.model_copy( - update={"calculation_provenance": calculation_provenance} - ) - - def read(self, uri): - if uri in self.payloads: - return self.payloads[uri] - return self.report_store.read(uri) - - def write_aggregate(self, *, prefix, payload): - reference = self.report_store.write_aggregate(prefix=prefix, payload=payload) - self.aggregate_writes.append(payload) - return reference - - def write_comparison(self, *, prefix, payload): - encoded = canonical_json_bytes(payload) - digest = sha256(encoded).hexdigest() - uri = "gs://private/comparison.json" - self.payloads[uri] = encoded - self.comparison_writes.append(payload) - return ArtifactReference( - uri=uri, - media_type=ArtifactMediaType.JSON, - content_sha256=digest, - size_bytes=len(encoded), - ) - - -def _frames(value=100.0): - return { - "household": pd.DataFrame( - {"household_id": [1], "household_net_income": [value]} - ), - "person": pd.DataFrame({"person_id": [1], "household_id": [1], "age": [40]}), - } def test_calculator_uses_the_current_dataset_selection_contract(monkeypatch) -> None: @@ -548,14 +31,17 @@ def test_calculator_uses_the_current_dataset_selection_contract(monkeypatch) -> country_module = object() region = RegionResolution(code="us") selection = DatasetSelection( - name="populace_us_2024", uri=_bundle().dataset.uri, is_default=True + name="populace_us_2024", uri=bundle_provenance().dataset.uri, is_default=True ) dataset = object() received: dict[str, object] = {} class ModelVersion: def resolve_entity_variables(self, _simulation): - return {entity: list(frame.columns) for entity, frame in _frames().items()} + return { + entity: list(frame.columns) + for entity, frame in simulation_frames().items() + } class Model: spm_config = None @@ -564,7 +50,7 @@ class Model: output_dataset = type( "OutputDataset", (), - {"data": type("OutputData", (), {"entity_data": _frames()})()}, + {"data": type("OutputData", (), {"entity_data": simulation_frames()})()}, )() def ensure(self) -> None: @@ -613,7 +99,7 @@ def build_simulation(_params, *, resolved, **_kwargs): monkeypatch.setattr(worker, "build_stage12_simulation", build_simulation) result = worker.calculate_simulation_frames( - _planned_simulation(SimulationRole.BASELINE) + planned_simulation(SimulationRole.BASELINE) ) assert "data" not in received["params"] @@ -629,61 +115,13 @@ def build_simulation(_params, *, resolved, **_kwargs): assert set(result.frames) == {"household", "person"} -class _StoppedAfterLoadingError(Exception): - """Raised by a stubbed simulation build once the dataset has loaded.""" - - -def _stage12_load(monkeypatch, *, region, selection, baked_folder, planned=None): - """Run a Stage 12 calculation up to its dataset load; return what it did.""" - from policyengine_simulation_executor import simulation_runtime - from policyengine_simulation_executor.stage12_cache import state_datasets - from policyengine_simulation_executor.stage12_runtime import simulation as worker - - loads = [] - lookups = [] - - def load_dataset(_params, *, selection, country_module, data_folder=None): - loads.append({"selection": selection, "data_folder": data_folder}) - return object() - - def lookup(**kwargs): - lookups.append(kwargs) - return baked_folder - - def stop_after_loading(*_args, **_kwargs): - raise _StoppedAfterLoadingError - - monkeypatch.setattr(worker, "_require_installed_bundle", lambda _: None) - monkeypatch.setattr(simulation_runtime, "setup_gcp_credentials", nullcontext) - monkeypatch.setattr( - simulation_runtime, "_country_module", lambda _country: object() - ) - monkeypatch.setattr(simulation_runtime, "_resolve_region", lambda **_: region) - monkeypatch.setattr( - simulation_runtime, - "_resolve_dataset_selection", - lambda _params, *, region_resolution: selection, - ) - monkeypatch.setattr(simulation_runtime, "_load_dataset", load_dataset) - monkeypatch.setattr(state_datasets, "baked_state_data_folder", lookup) - monkeypatch.setattr( - worker, "detect_uk_run_dataset_metadata", lambda *_a, **_k: None - ) - monkeypatch.setattr(worker, "build_stage12_simulation", stop_after_loading) - - planned = planned or _planned_simulation(SimulationRole.BASELINE) - with pytest.raises(_StoppedAfterLoadingError): - worker.calculate_simulation_frames(planned) - return SimpleNamespace(loads=loads, lookups=lookups, year=planned.year) - - _UK_LOCAL_AREA_DATASET = "microcosm_uk_2024_25_dense" _UK_LOCAL_AREA_URI = "hf://policyengine/uk/microcosm_uk_2024_25_dense.h5@revision" def _uk_area_planned_simulation(region_code: str) -> PlannedSimulationExecutionInput: """Return a planned UK area simulation whose provenance names the dense file.""" - payload = _planned_simulation(SimulationRole.BASELINE).model_dump(mode="json") + payload = planned_simulation(SimulationRole.BASELINE).model_dump(mode="json") payload["geography"] = {"country": "uk", "region": region_code} payload["output_plan"]["country"] = "uk" payload["bundle"]["country_package_name"] = "policyengine-uk" @@ -793,7 +231,7 @@ def test_stage12_state_and_district_runs_read_their_states_baked_folder( resolve_data_folder, ) - loaded = _stage12_load( + loaded = run_until_dataset_loaded( monkeypatch, region=RegionResolution(code=region_code, state_code="UT"), selection=_regional_selection(), @@ -820,7 +258,7 @@ def test_stage12_state_and_district_runs_read_their_states_baked_folder( def test_stage12_state_run_without_a_baked_file_loads_on_demand(monkeypatch) -> None: from policyengine_simulation_executor.simulation_runtime import RegionResolution - loaded = _stage12_load( + loaded = run_until_dataset_loaded( monkeypatch, region=RegionResolution(code="state/ut", state_code="UT"), selection=_regional_selection(), @@ -830,3844 +268,6 @@ def test_stage12_state_run_without_a_baked_file_loads_on_demand(monkeypatch) -> assert loaded.loads[0]["data_folder"] is None -REGIONAL_URI = "hf://policyengine/data/populace_us_2024_acs_local.h5@regional-revision" - - -def _regional_provenance() -> BundleProvenance: - return _bundle().model_copy( - update={ - "dataset": DatasetProvenance( - identity="populace_us_2024_acs_local", - uri=REGIONAL_URI, - artifact_revision="regional-revision", - data_package_name="populace-data", - data_package_version="0.1.0", - ) - } - ) - - -def _installed_bundle_with_regions(): - from policyengine_simulation_contract.stage12_bundle import ( - ResolvedStage12Bundle, - Stage12BundleManifest, - Stage12CountryBundle, - Stage12Dataset, - Stage12RegionDataset, - ) - - national = Stage12Dataset( - identity="populace_us_2024", - uri=_bundle().dataset.uri, - artifact_revision="revision", - sha256="a" * 64, - repo_type="dataset", - ) - regional = Stage12Dataset( - identity="populace_us_2024_acs_local", - uri=REGIONAL_URI, - artifact_revision="regional-revision", - sha256="c" * 64, - repo_type="dataset", - ) - return ResolvedStage12Bundle( - bundle=Stage12BundleManifest( - policyengine_version="5.2.0", - policyengine_requirement="policyengine==5.2.0", - core_package_version="3.0.0", - core_package_requirement="policyengine-core==3.0.0", - countries=( - Stage12CountryBundle( - country="us", - country_package_name="policyengine-us", - country_package_version="1.764.6", - country_package_requirement="policyengine-us==1.764.6", - data_package_name="populace-data", - data_package_version="0.1.0", - data_release_version="revision", - data_artifact_revision="revision", - default_dataset=national.identity, - default_dataset_uri=national.uri, - datasets=(national, regional), - region_datasets=( - Stage12RegionDataset( - region_type="national", dataset=national.identity - ), - Stage12RegionDataset( - region_type="state", dataset=regional.identity - ), - Stage12RegionDataset( - region_type="congressional_district", - dataset=regional.identity, - ), - ), - ), - ), - ), - bundle_manifest_sha256="b" * 64, - ) - - -def _simulation_for(region: str, provenance: BundleProvenance): - simulation = _planned_simulation(SimulationRole.BASELINE) - return simulation.model_copy( - update={ - "geography": GeographySelection(country="us", region=region), - "bundle": provenance, - "population": DatasetPopulationInput( - artifact=DatasetArtifactReference( - uri=provenance.dataset.uri, - media_type=DatasetArtifactMediaType.HDF5, - content_sha256="a" * 64, - ) - ), - } - ) - - -@pytest.mark.parametrize( - ("region", "provenance"), - [ - ("state/ut", _regional_provenance), - ("State/UT", _regional_provenance), - ("ut", _regional_provenance), - ("congressional_district/UT-01", _regional_provenance), - ("us", _bundle), - # Entry services released before regional provenance send the default - # dataset for state runs until they are updated. - ("state/ut", _bundle), - ], -) -def test_installed_bundle_accepts_the_dataset_a_region_loads( - monkeypatch, region, provenance -) -> None: - from policyengine_simulation_executor.stage12_runtime import simulation as worker - - monkeypatch.setattr(worker, "load_stage12_bundle", _installed_bundle_with_regions) - - worker._require_installed_bundle(_simulation_for(region, provenance())) - - -def test_uk_regions_keep_the_default_dataset_when_the_bundle_names_region_types( - monkeypatch, -) -> None: - from policyengine_simulation_contract.stage12_bundle import Stage12RegionDataset - from policyengine_simulation_executor.stage12_runtime import simulation as worker - - resolved = _installed_bundle_with_regions() - us = resolved.bundle.countries[0] - uk = us.model_copy( - update={ - "country": "uk", - "region_datasets": ( - Stage12RegionDataset( - region_type="national", dataset=us.default_dataset - ), - ), - } - ) - resolved = resolved.model_copy( - update={"bundle": resolved.bundle.model_copy(update={"countries": (us, uk)})} - ) - monkeypatch.setattr(worker, "load_stage12_bundle", lambda: resolved) - simulation = _simulation_for("us", _bundle()).model_copy( - update={"geography": GeographySelection(country="uk", region="country/wales")} - ) - - worker._require_installed_bundle(simulation) - with pytest.raises(RuntimeError, match="differs from the installed bundle"): - worker._require_installed_bundle( - simulation.model_copy(update={"bundle": _regional_provenance()}) - ) - - -@pytest.mark.parametrize( - ("region", "dataset"), - [ - ("us", "populace_us_2024_acs_local"), - ("state/ut", "source_enrichment"), - ], -) -def test_installed_bundle_refuses_a_dataset_the_region_does_not_load( - monkeypatch, region, dataset -) -> None: - from policyengine_simulation_executor.stage12_runtime import simulation as worker - - monkeypatch.setattr(worker, "load_stage12_bundle", _installed_bundle_with_regions) - provenance = _regional_provenance() - provenance = provenance.model_copy( - update={"dataset": provenance.dataset.model_copy(update={"identity": dataset})} - ) - - with pytest.raises(RuntimeError, match="differs from the installed bundle"): - worker._require_installed_bundle(_simulation_for(region, provenance)) - - -def _selection(name: str, uri: str, *, is_default: bool): - from policyengine_simulation_executor.simulation_runtime import DatasetSelection - - return DatasetSelection(name=name, uri=uri, is_default=is_default) - - -def test_a_calculation_loads_the_dataset_its_provenance_names(monkeypatch) -> None: - from policyengine_simulation_executor.simulation_runtime import RegionResolution - - loaded = _stage12_load( - monkeypatch, - region=RegionResolution(code="state/ut", state_code="UT"), - selection=_selection( - "populace_us_2024_acs_local", REGIONAL_URI, is_default=False - ), - baked_folder=None, - planned=_simulation_for("state/ut", _regional_provenance()), - ) - - assert loaded.loads[0]["selection"].name == "populace_us_2024_acs_local" - - -def test_a_calculation_refuses_provenance_naming_another_dataset( - monkeypatch, -) -> None: - from policyengine_simulation_executor.simulation_runtime import RegionResolution - - with pytest.raises(RuntimeError, match="dataset its region loads"): - _stage12_load( - monkeypatch, - region=RegionResolution(code="state/ut", state_code="UT"), - selection=_selection( - "populace_us_2024", _bundle().dataset.uri, is_default=True - ), - baked_folder=None, - planned=_simulation_for("state/ut", _regional_provenance()), - ) - - -def test_default_provenance_from_an_older_entry_service_is_still_loaded( - monkeypatch, -) -> None: - from policyengine_simulation_executor.simulation_runtime import RegionResolution - - loaded = _stage12_load( - monkeypatch, - region=RegionResolution(code="state/ut", state_code="UT"), - selection=_selection( - "populace_us_2024_acs_local", REGIONAL_URI, is_default=False - ), - baked_folder=None, - planned=_simulation_for("state/ut", _bundle()), - ) - - assert loaded.loads[0]["selection"].name == "populace_us_2024_acs_local" - - -def test_uk_boundary_version_is_detected_from_the_unscoped_dataset() -> None: - dataset = type( - "Dataset", - (), - { - "data": type( - "Data", - (), - { - "entity_data": { - "household": pd.DataFrame( - {"la_code_oa": ["E06000001", "E07000026"]} - ) - } - }, - )() - }, - )() - - metadata = detect_uk_local_authority_metadata("uk", dataset) - - assert metadata == UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 - ) - - -def test_country_metadata_alignment_rejects_missing_or_mismatched_uk_values() -> None: - lad22 = UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD22 - ) - lad23 = UKLocalAuthorityMetadata( - boundary_version=UKLocalAuthorityBoundaryVersion.LAD23 - ) - - with pytest.raises(ValueError, match="missing"): - validate_uk_local_authority_metadata("uk", None, lad22) - with pytest.raises(ValueError, match="do not match"): - validate_uk_local_authority_metadata("uk", lad22, lad23) - assert validate_uk_local_authority_metadata("uk", lad22, lad22) == lad22 - assert validate_uk_local_authority_metadata("us", None, None) is None - - -@pytest.mark.parametrize( - ("region_code", "routed", "accepted"), - [ - ("uk", True, True), - ("country/england", True, True), - ("uk", False, False), - ("constituency/E14001063", True, False), - ("local_authority/E06000063", True, False), - (None, True, False), - ], -) -def test_uk_artifacts_without_metadata_need_a_routed_national_request( - monkeypatch, region_code: str | None, *, routed: bool, accepted: bool -) -> None: - monkeypatch.setattr( - "policyengine_simulation_executor.uk_local_authority_metadata." - "uk_area_regions_routed", - lambda: routed, - ) - - if accepted: - assert ( - validate_uk_local_authority_metadata( - "uk", None, None, region_code=region_code - ) - is None - ) - else: - with pytest.raises(ValueError, match="metadata is missing"): - validate_uk_local_authority_metadata( - "uk", None, None, region_code=region_code - ) - - -def test_single_worker_accepts_one_policy_and_persists_one_artifact() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - artifacts = FakeArtifacts() - - result = run_single_simulation( - simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=artifacts, - calculator=lambda _: _frames(), - ), - ) - - assert result["role"] == "baseline" - child = store.children[simulation.simulation_execution_id] - assert child.status is ComparisonRunLifecycleStatus.SUCCEEDED - assert child.output_uri == result["artifact"]["uri"] - assert artifacts.input_writes == [ - (_context().artifact_prefix, SimulationRole.BASELINE) - ] - - -def test_single_worker_routes_eligible_us_calculation_to_segment_runner() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - received = [] - - def segmented_calculator(value, cancellation_requested): - received.append(value) - assert cancellation_requested() is False - return SimulationCalculation(frames=_frames()) - - run_single_simulation( - simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=lambda _: pytest.fail("monolithic calculator was selected"), - segmented_calculator=segmented_calculator, - ), - ) - - assert received == [simulation] - - -def test_single_worker_honors_explicit_segmentation_opt_out() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( - update={"options": {"segmented": False}} - ) - store.children[simulation.simulation_execution_id] = _child(simulation) - segmented_calls = [] - - def unexpected_segmented(value, cancellation_requested): - segmented_calls.append(value) - return SimulationCalculation(frames=_frames()) - - run_single_simulation( - simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=lambda _: _frames(), - segmented_calculator=unexpected_segmented, - ), - ) - - assert segmented_calls == [] - - -def test_single_worker_refuses_a_simulation_that_is_no_longer_active() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation).model_copy( - update={"status": ComparisonRunLifecycleStatus.FAILED} - ) - artifacts = FakeArtifacts() - - with pytest.raises( - RuntimeError, match=r"^Stage 12 simulation is no longer active$" - ): - run_single_simulation( - simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=artifacts, - calculator=lambda _: pytest.fail("a finished simulation ran again"), - ), - ) - - assert artifacts.input_writes == [] - - -def test_single_worker_stops_when_another_writer_claimed_the_simulation() -> None: - class ClaimLosingStore(FakeStore): - def replace_simulation_if_status(self, record, *, expected_status): - return self.children[record.simulation_execution_id], False - - store = ClaimLosingStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - artifacts = FakeArtifacts() - - with pytest.raises(RuntimeError, match=r"^Stage 12 simulation was cancelled$"): - run_single_simulation( - simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=artifacts, - calculator=lambda _: pytest.fail("an unclaimed simulation ran"), - ), - ) - - assert artifacts.input_writes == [] - - -def test_single_worker_treats_a_failed_cancellation_check_as_not_cancelled( - caplog, -) -> None: - class StatusReadFailingStore(FakeStore): - def __init__(self): - super().__init__() - self.fail_reads = False - - def get_simulation(self, simulation_execution_id): - if self.fail_reads: - msg = "status read contains sensitive details" - raise ConnectionError(msg) - return super().get_simulation(simulation_execution_id) - - store = StatusReadFailingStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - checks = [] - - def segmented_calculator(_value, cancellation_requested): - store.fail_reads = True - checks.append(cancellation_requested()) - store.fail_reads = False - return SimulationCalculation(frames=_frames()) - - with caplog.at_level( - logging.WARNING, - logger="policyengine_simulation_executor.stage12_runtime.simulation", - ): - run_single_simulation( - simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - segmented_calculator=segmented_calculator, - ), - ) - - assert checks == [False] - assert ( - f"Stage 12 cancellation check failed for " - f"{simulation.simulation_execution_id} (ConnectionError)" - ) in caplog.text - assert "sensitive" not in caplog.text - assert ( - store.children[simulation.simulation_execution_id].status - is ComparisonRunLifecycleStatus.SUCCEEDED - ) - - -def test_single_worker_retains_detached_calculation_provenance() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - artifacts = FakeArtifacts() - provenance = _spm_calculation_provenance() - - result = run_single_simulation( - simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=artifacts, - calculator=lambda _: SimulationCalculation( - frames=_frames(), - calculation_provenance=provenance, - ), - ), - ) - - assert result["calculation_provenance"] == provenance.model_dump( - mode="json", by_alias=True - ) - assert "spm_config" not in result["calculation_provenance"] - - -def test_single_worker_exposes_only_a_bounded_failure() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - - def fail(_simulation): - msg = "country package calculation failed" - raise RuntimeError(msg) - - with pytest.raises( - RuntimeError, match="Stage 12 simulation execution failed" - ) as error: - run_single_simulation( - simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=fail, - ), - ) - - assert "country package calculation failed" not in str(error.value) - child = store.children[simulation.simulation_execution_id] - assert child.status is ComparisonRunLifecycleStatus.FAILED - assert child.error_code == "simulation_execution_failed" - assert _failure_cause(child.error_summary) == ( - "RuntimeError in stage12_calculation (baseline): " - "country package calculation failed" - ) - - -def test_single_worker_persists_and_records_a_redacted_cause() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - runtime = RecordingRuntime() - - def fail(_simulation): - msg = "could not reach postgresql://runtime:database-secret@db.example/postgres" - raise RuntimeError(msg) - - 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", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=fail, - runtime=runtime, - ), - ) - - child = store.children[simulation.simulation_execution_id] - assert _failure_cause(child.error_summary) == ( - "RuntimeError in stage12_calculation (baseline): " - "could not reach postgresql://[REDACTED]@db.example/postgres" - ) - assert "database-secret" not in child.error_summary - [recorded_cause] = runtime.recorded_exceptions - assert str(recorded_cause) == child.error_summary - - -def test_single_worker_names_the_failing_artifact_write_stage() -> None: - class FailingSimulationArtifacts(FakeArtifacts): - def write_simulation(self, **_kwargs): - msg = "bucket write rejected" - raise OSError(msg) - - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - - 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", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FailingSimulationArtifacts(), - calculator=lambda _: _frames(), - ), - ) - - assert _failure_cause( - store.children[simulation.simulation_execution_id].error_summary - ) == ( - "OSError in stage12_simulation_artifact_write (baseline): bucket write rejected" - ) - - -def test_single_worker_preserves_typed_failure_values() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - - def fail(_simulation): - msg = "SPM_YEAR_UNAVAILABLE" - raise SPMInputError(msg, "No SPM values for 2026") - - 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", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=fail, - ), - ) - - child = store.children[simulation.simulation_execution_id] - assert child.error_code == "SPM_YEAR_UNAVAILABLE" - assert child.error_summary == "No SPM values for 2026" - - -def test_single_worker_rejects_bare_year_policy_period_before_calculation() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( - update={ - "policy": {"gov.irs.credits.ctc.amount.base[0].amount": {"2026": 3_000}} - } - ) - store.children[simulation.simulation_execution_id] = _child(simulation) - - 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", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=lambda _: pytest.fail("calculation must not start"), - ), - ) - - child = store.children[simulation.simulation_execution_id] - assert child.error_code == "invalid_policy_period" - assert child.error_summary == ( - "Policy parameter 'gov.irs.credits.ctc.amount.base[0].amount' has invalid " - "effective period '2026'; expected YYYY-MM-DD or " - "YYYY-MM-DD.YYYY-MM-DD." - ) - - -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"), - ) - - -CORRELATION_SUFFIX = re.compile(r" \(correlation_id=(?P[0-9a-f]+)\)\Z") - - -def _correlation_id(error_summary: str) -> str: - """Return the correlation ID that ends a redacted failure summary.""" - match = CORRELATION_SUFFIX.search(error_summary) - assert match is not None, error_summary - return match["correlation_id"] - - -def _failure_cause(error_summary: str) -> str: - """Return the redacted cause before a failure summary's correlation ID.""" - _correlation_id(error_summary) - return CORRELATION_SUFFIX.sub("", error_summary) - - -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", - dependencies=SimulationWorkerDependencies( - 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_if_status(self, record, *, expected_status): - if record.status is ComparisonRunLifecycleStatus.FAILED: - ordered_events.append(("failure_write", record.error_summary)) - raise _connection_limit_error() - return super().replace_simulation_if_status( - record, expected_status=expected_status - ) - - 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", - dependencies=SimulationWorkerDependencies( - 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", - "evaluation_id": str(EVALUATION_ID), - "simulation_execution_id": str(BASELINE_ID), - "correlation_id": _correlation_id(failure_write[1]), - } - assert failure_write[0] == "failure_write" - assert _failure_cause(failure_write[1]) == ( - "ValueError in stage12_calculation (baseline): " - "country column missing from household data" - ) - 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_late_simulation_success_does_not_overwrite_parent_cancellation() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - - def segmented_calculator(value, cancellation_requested): - current = store.children[value.simulation_execution_id] - cancelled_at = datetime.now(UTC) - store.replace_simulation( - current.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": "simulation_cancelled", - "error_summary": "Cancelled after the reform simulation failed.", - "updated_at": cancelled_at, - "completed_at": cancelled_at, - } - ) - ) - assert cancellation_requested() is True - return SimulationCalculation(frames=_frames()) - - with pytest.raises(RuntimeError, match="Stage 12 simulation was cancelled"): - run_single_simulation( - simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=lambda _: pytest.fail("monolithic calculator was selected"), - segmented_calculator=segmented_calculator, - ), - ) - - child = store.children[simulation.simulation_execution_id] - assert child.status is ComparisonRunLifecycleStatus.FAILED - assert child.error_code == "simulation_cancelled" - - -def test_single_worker_rejects_frames_that_do_not_satisfy_the_output_plan() -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - frames = _frames() - frames["person"] = frames["person"].drop(columns=["age"]) - - 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", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=lambda _: frames, - ), - ) - - assert store.children[simulation.simulation_execution_id].status is ( - ComparisonRunLifecycleStatus.FAILED - ) - - -def test_aggregate_combines_detached_spm_receipts() -> None: - provenance = _spm_calculation_provenance() - - report = _report() - report = report.model_copy( - update={ - "baseline": report.baseline.model_copy( - update={"options": {"spm": {"scenario": "official"}}} - ), - "reform": report.reform.model_copy( - update={"options": {"spm": {"scenario": "official"}}} - ), - } - ) - baseline = ( - FakeArtifacts() - .add_simulation(report.baseline) - .model_copy(update={"calculation_provenance": provenance}) - ) - reform = ( - FakeArtifacts() - .add_simulation(report.reform) - .model_copy(update={"calculation_provenance": provenance}) - ) - - result = _build_spm_result( - report=report, - baseline_descriptor=baseline, - reform_descriptor=reform, - ) - - assert "spm_config" not in result - assert result["spm_provenance"]["baseline"]["execution_count"] == 1 - assert result["spm_provenance"]["reform"]["execution_count"] == 1 - assert result["spm_provenance"]["baseline"]["receipt"]["scenario"] == "official" - - -def test_aggregate_stand_ins_preserve_policy_and_cliff_options(monkeypatch) -> None: - from policyengine.outputs import labor_supply_response_is_active - from policyengine_simulation_executor import simulation_output_builder - - report = _report() - options = {"include_cliffs": True} - report = report.model_copy( - update={ - "baseline": report.baseline.model_copy(update={"options": options}), - "reform": report.reform.model_copy( - update={ - "options": options, - "policy": { - "gov.simulation.labor_supply_responses.elasticities.income": 0.1 - }, - } - ), - } - ) - frames = { - entity: pd.DataFrame({f"{entity}_id": [1]}) - for entity in ( - "person", - "marital_unit", - "family", - "spm_unit", - "tax_unit", - "household", - ) - } - observed = {} - - class CapturingBuilder: - def __init__(self, **values): - observed.update(values) - - def serialize(self): - return {"captured": True} - - monkeypatch.setattr( - simulation_output_builder, - "SimulationOutputBuilder", - CapturingBuilder, - ) - artifacts = FakeArtifacts() - - result = build_aggregate_report( - report=report, - baseline=RetainedSimulationOutput( - frames, artifacts.add_simulation(report.baseline) - ), - reform=RetainedSimulationOutput( - frames, artifacts.add_simulation(report.reform) - ), - ) - - assert result["result"] == {"captured": True} - assert observed["simulation_params"]["include_cliffs"] is True - assert observed["uk_local_authority_metadata"] is None - assert labor_supply_response_is_active( - observed["baseline"], - observed["reform"], - country_code="us", - ) - - -def test_aggregate_stand_ins_keep_bounded_policy_end_dates(monkeypatch) -> None: - from datetime import datetime - - from policyengine_simulation_executor import simulation_output_builder - - path = "gov.simulation.labor_supply_responses.elasticities.income" - report = _report() - report = report.model_copy( - update={ - "baseline": report.baseline.model_copy( - update={"policy": {path: {"2026-01-01.2026-12-31": 0.05}}} - ), - "reform": report.reform.model_copy( - update={"policy": {path: {"2026-01-01.2027-12-31": 0.1}}} - ), - } - ) - frames = { - entity: pd.DataFrame({f"{entity}_id": [1]}) - for entity in ( - "person", - "marital_unit", - "family", - "spm_unit", - "tax_unit", - "household", - ) - } - observed = {} - - class CapturingBuilder: - def __init__(self, **values): - observed.update(values) - - def serialize(self): - return {"captured": True} - - monkeypatch.setattr( - simulation_output_builder, "SimulationOutputBuilder", CapturingBuilder - ) - artifacts = FakeArtifacts() - - build_aggregate_report( - report=report, - baseline=RetainedSimulationOutput( - frames, artifacts.add_simulation(report.baseline) - ), - reform=RetainedSimulationOutput( - frames, artifacts.add_simulation(report.reform) - ), - ) - - for role, stop in (("baseline", "2026-12-31"), ("reform", "2027-12-31")): - (value,) = observed[role].policy.parameter_values - assert value.end_date == datetime.strptime(stop, "%Y-%m-%d"), role - - -def _uk_report() -> ReportExecutionInput: - uk_bundle = _bundle().model_copy( - update={ - "country_package_name": "policyengine-uk", - "country_package_version": "2.102.3", - } - ) - report_payload = _report().model_dump(mode="json") - for role, policy in ( - ("baseline", {}), - ( - "reform", - { - "gov.hmrc.income_tax.charges.CB_HITC.phase_out_start": { - "2026-01-01.2100-12-31": 70_000 - } - }, - ), - ): - report_payload[role].update( - { - "policy": policy, - "geography": {"country": "uk", "region": "uk"}, - "bundle": uk_bundle.model_dump(mode="json"), - } - ) - return ReportExecutionInput.model_validate(report_payload) - - -def test_uk_aggregate_without_spm_provenance_is_written(monkeypatch) -> None: - from policyengine_simulation_executor import simulation_output_builder - - report = _uk_report() - assert report.baseline.options.get("spm") is None - uk_output = { - **{ - key: value - for key, value in CURRENT_SINGLE_YEAR_MACRO_RESULT.items() - if key != "spm_provenance" - }, - "congressional_district_impact": None, - } - - class UKOutputBuilder: - def __init__(self, **_values): - pass - - def serialize(self): - # The real builder dumps SingleYearMacroOutput, which serializes - # the unset optional SPM receipt as an explicit null. - return SingleYearMacroOutput.model_validate(uk_output).model_dump( - mode="json" - ) - - monkeypatch.setattr( - simulation_output_builder, "SimulationOutputBuilder", UKOutputBuilder - ) - uk_frames = { - entity: pd.DataFrame({f"{entity}_id": [1]}) - for entity in ("person", "benunit", "household") - } - artifacts = FakeArtifacts() - - aggregate = build_aggregate_report( - report=report, - baseline=RetainedSimulationOutput( - uk_frames, artifacts.add_simulation(report.baseline) - ), - reform=RetainedSimulationOutput( - uk_frames, artifacts.add_simulation(report.reform) - ), - ) - assert aggregate["result"]["spm_provenance"] is None - - aggregate_reference = artifacts.write_aggregate( - prefix=_context().artifact_prefix, payload=aggregate - ) - - stored_aggregate = artifacts.read(aggregate_reference.uri) - assert aggregate_reference.content_sha256 == sha256(stored_aggregate).hexdigest() - assert json.loads(stored_aggregate)["result"]["spm_provenance"] is None - - -@pytest.mark.parametrize( - ("region", "routed", "omitted"), - [ - ("uk", True, True), - ("country/wales", True, True), - ("uk", False, False), - ("constituency/E14001063", True, False), - ], -) -def test_uk_aggregate_omits_area_breakdowns_only_for_routed_national_reports( - monkeypatch, region: str, *, routed: bool, omitted: bool -) -> None: - from policyengine_simulation_executor import simulation_output_builder - - monkeypatch.setattr( - "policyengine_simulation_executor.uk_local_authority_metadata." - "uk_area_regions_routed", - lambda: routed, - ) - observed = {} - - class CapturingBuilder: - def __init__(self, **values): - observed.update(values) - - def serialize(self): - return {} - - monkeypatch.setattr( - simulation_output_builder, "SimulationOutputBuilder", CapturingBuilder - ) - payload = _uk_report().model_dump(mode="json") - for role in ("baseline", "reform"): - payload[role]["geography"] = {"country": "uk", "region": region} - report = ReportExecutionInput.model_validate(payload) - frames = { - entity: pd.DataFrame({f"{entity}_id": [1]}) - for entity in ("person", "benunit", "household") - } - artifacts = FakeArtifacts() - - build_aggregate_report( - report=report, - baseline=RetainedSimulationOutput( - frames, artifacts.add_simulation(report.baseline) - ), - reform=RetainedSimulationOutput( - frames, artifacts.add_simulation(report.reform) - ), - ) - - assert observed["omit_uk_area_breakdowns"] is omitted - # The runtime omits them from exactly the same runs. - from policyengine_simulation_executor import uk_local_authority_metadata - - assert uk_local_authority_metadata.uk_area_breakdowns_omitted("uk", region) is ( - omitted - ) - - -def test_aggregate_report_wraps_builder_output_with_report_provenance( - monkeypatch, -) -> None: - from policyengine_simulation_executor import simulation_output_builder - - report = _report() - report = report.model_copy( - update={"reform": report.reform.model_copy(update={"policy": {}})} - ) - entities = ("person", "marital_unit", "family", "spm_unit", "tax_unit") - - def frames(income): - return { - **{entity: pd.DataFrame({f"{entity}_id": [1]}) for entity in entities}, - "household": pd.DataFrame( - {"household_id": [1], "household_net_income": [income]} - ), - } - - observed = {} - - class CapturingBuilder: - def __init__(self, **values): - observed.update(values) - - def serialize(self): - return {"budget": {"budgetary_impact": 1.0}} - - monkeypatch.setattr( - simulation_output_builder, "SimulationOutputBuilder", CapturingBuilder - ) - artifacts = FakeArtifacts() - - result = build_aggregate_report( - report=report, - baseline=RetainedSimulationOutput( - frames(100.0), artifacts.add_simulation(report.baseline) - ), - reform=RetainedSimulationOutput( - frames(120.0), artifacts.add_simulation(report.reform) - ), - uk_local_authority_metadata=None, - ) - - assert result == { - "contract_version": 1, - "aggregate_schema_version": 1, - "evaluation_id": str(EVALUATION_ID), - "requested_aggregates": [item.value for item in report.requested_aggregates], - "bundle": report.baseline.bundle.model_dump(mode="json"), - "result": {"budget": {"budgetary_impact": 1.0}}, - } - assert list(observed) == [ - "country", - "simulation_params", - "country_module", - "dataset", - "baseline", - "reform", - "resolved_data_version", - "resolved_region_code", - "uk_local_authority_metadata", - "omit_uk_area_breakdowns", - ] - assert observed["country"] == "us" - assert observed["simulation_params"] == { - "country": "us", - "scope": "macro", - "data": report.baseline.bundle.dataset.uri, - "data_version": report.baseline.bundle.dataset.artifact_revision, - "time_period": str(report.baseline.year), - "region": report.baseline.geography.region, - } - assert observed["dataset"] is observed["baseline"].output_dataset - for role, income in (("baseline", 100.0), ("reform", 120.0)): - household = observed[role].output_dataset.data.household - assert household["household_net_income"].tolist() == [income] - assert observed["resolved_data_version"] == ( - report.baseline.bundle.dataset.artifact_revision - ) - assert observed["resolved_region_code"] == report.baseline.geography.region - assert observed["uk_local_authority_metadata"] is None - assert observed["omit_uk_area_breakdowns"] is False - - -def test_single_worker_rejects_combined_baseline_and_reform_input() -> None: - payload = _planned_simulation(SimulationRole.BASELINE).model_dump(mode="json") - payload["baseline"] = {} - payload["reform"] = {} - - with pytest.raises(ValueError, match="cannot contain baseline or reform"): - run_single_simulation( - payload, - _context().model_dump(mode="json"), - required_country="us", - dependencies=SimulationWorkerDependencies( - store=FakeStore(), - artifacts=FakeArtifacts(), - calculator=lambda _: _frames(), - ), - ) - - -def test_simulation_input_digest_excludes_temporary_record_identifiers() -> None: - original = _simulation(SimulationRole.BASELINE) - recreated = original.model_copy( - update={ - "evaluation_id": uuid4(), - "simulation_execution_id": uuid4(), - } - ) - - assert simulation_input_sha256(original) == simulation_input_sha256(recreated) - - -class FutureCall: - def __init__(self, object_id, future): - self.object_id = object_id - self.future = future - - def get(self, *, timeout=None): - return self.future.result(timeout=timeout) - - def cancel(self): - self.future.cancel() - - -class ImmediateCall: - def __init__(self, object_id, result, observed_timeouts): - self.object_id = object_id - self.result = result - self.observed_timeouts = observed_timeouts - self.cancel_count = 0 - - def get(self, *, timeout=None): - self.observed_timeouts.append(timeout) - return self.result - - def cancel(self): - self.cancel_count += 1 - - -@dataclass(frozen=True) -class InjectedChildFaults: - """Faults a ConcurrentInvoker injects into the child simulations it runs.""" - - fail_role: str | None = None - incompatible: bool = False - wrong_plan_digest: bool = False - missing_column: bool = False - - -class ConcurrentInvoker: - def __init__( - self, - artifacts, - *, - faults: InjectedChildFaults | None = None, - production_result=None, - ): - faults = faults or InjectedChildFaults() - self.artifacts = artifacts - self.fail_role = faults.fail_role - self.incompatible = faults.incompatible - self.wrong_plan_digest = faults.wrong_plan_digest - self.missing_column = faults.missing_column - self.executor = ThreadPoolExecutor(max_workers=2) - self.intervals = {} - self.events = [] - self.environments = [] - self.production_result = production_result - self.production_wait_timeouts = [] - self.output_plans = [] - - def spawn(self, request): - parsed = PlannedSimulationExecutionInput.model_validate(request.simulation) - role = parsed.role.value - self.events.append(f"spawn:{role}") - self.environments.append(request.environment) - self.output_plans.append(parsed.output_plan) - - return self._submit(parsed, role) - - def restore(self, invocation_id): - if invocation_id == "production-job-1": - self.events.append("restore:production") - return ImmediateCall( - invocation_id, - self.production_result, - self.production_wait_timeouts, - ) - role = invocation_id.removeprefix("call-") - parsed = _planned_simulation(SimulationRole(role)) - self.events.append(f"restore:{role}") - return self._submit(parsed, role) - - def _submit(self, parsed, role): - - def run(): - started = time.monotonic() - time.sleep(0.05) - if role == self.fail_role: - msg = "child calculation failed" - raise RuntimeError(msg) - frames = _frames(100.0 if role == "baseline" else 120.0) - if self.missing_column and role == "reform": - frames["household"] = frames["household"].drop( - columns=["household_net_income"] - ) - descriptor = self.artifacts.add_simulation(parsed, frames) - if self.incompatible and role == "reform": - descriptor = descriptor.model_copy( - update={ - "row_identity": RowIdentity( - identifier_columns=("household.household_id",), - row_count=1, - identity_sha256="f" * 64, - ) - } - ) - if self.wrong_plan_digest and role == "reform": - descriptor = descriptor.model_copy( - update={"output_plan_sha256": "f" * 64} - ) - self.intervals[role] = (started, time.monotonic()) - return descriptor.model_dump(mode="json") - - return FutureCall(f"call-{role}", self.executor.submit(run)) - - -def test_coordinator_starts_both_children_before_waiting_and_aggregates() -> None: - store = FakeStore() - artifacts = FakeArtifacts() - invoker = ConcurrentInvoker(artifacts) - observed = {} - - def aggregate(**values): - observed.update(values) - return _aggregate_payload() - - context = _context().model_copy( - update={ - "environment": "production", - "modal_environment": "main", - "artifact_prefix": ( - "stage-12-runs/production/2026/09/00000000-0000-0000-0000-000000000001" - ), - } - ) - parent = _parent().model_copy(update={"environment": "production"}) - resolver_calls = [] - - def resolver(report): - resolver_calls.append(report.evaluation_id) - return _output_plan() - - result = coordinate_report_impl( - _report().model_dump(mode="json"), - context.model_dump(mode="json"), - parent.model_dump(mode="json"), - invocation=CoordinatorInvocation( - application_name=context.modal_application, - invocation_id="coordinator-1", - ), - dependencies=CoordinatorDependencies( - store=store, - artifacts=artifacts, - invoker=invoker, - aggregator=aggregate, - output_plan_resolver=resolver, - ), - ) - - assert invoker.events == ["spawn:baseline", "spawn:reform"] - assert invoker.environments == ["main", "main"] - assert resolver_calls == [EVALUATION_ID] - assert invoker.output_plans == [_output_plan(), _output_plan()] - latest_start = max(interval[0] for interval in invoker.intervals.values()) - earliest_end = min(interval[1] for interval in invoker.intervals.values()) - assert latest_start < earliest_end - assert result["artifact"]["uri"] == ( - f"gs://private/{context.artifact_prefix}/reports/aggregate.json" - ) - assert observed["baseline"].frames["household"][ - "household_net_income" - ].tolist() == [100.0] - assert observed["reform"].frames["household"]["household_net_income"].tolist() == [ - 120.0 - ] - assert observed["uk_local_authority_metadata"] is None - assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED - assert store.parent.aggregation_status is ComparisonRunAggregationStatus.SUCCEEDED - assert store.parent.coordinator_invocation_id == "coordinator-1" - - -def test_duplicate_coordinator_submission_does_not_start_duplicate_children() -> None: - store = FakeStore() - artifacts = FakeArtifacts() - context = _context() - parent = _parent() - first_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=first_invoker, - aggregator=lambda **_: _aggregate_payload(), - ) - duplicate_invoker = ConcurrentInvoker(artifacts) - duplicate = 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-2", - store=store, - artifacts=artifacts, - invoker=duplicate_invoker, - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - ) - - assert duplicate == { - "deduplicated": True, - "evaluation_id": str(EVALUATION_ID), - "status": "succeeded", - } - assert duplicate_invoker.events == [] - - -def test_coordinator_measures_output_and_child_input_planning() -> None: - store = FakeStore() - 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", - store=store, - artifacts=artifacts, - invoker=ConcurrentInvoker(artifacts), - aggregator=lambda **_: _aggregate_payload(), - runtime=runtime, - ) - - planning_operations = [ - operation - for operation in runtime.operations - if operation[0] - in { - "stage12_coordinator_preparation", - "stage12_output_planning", - "stage12_child_input_planning", - } - ] - assert planning_operations == [ - ("stage12_coordinator_preparation", {}), - ("stage12_output_planning", {"country": "us"}), - ("stage12_child_input_planning", {"simulation_role": "baseline"}), - ("stage12_child_input_planning", {"simulation_role": "reform"}), - ] - - -def test_coordinator_planning_operations_preserve_observability_id() -> None: - output = io.StringIO() - runtime = configure( - ObservabilityConfig( - service=ServiceIdentity( - name="stage12-test", - namespace="policyengine.api-v1", - version="test", - role="coordinator", - ), - deployment=DeploymentIdentity( - environment="test", - platform="local", - region="us-central1", - instance_id="test-process", - ), - logging=LoggingConfig( - destinations=(StdoutLogDestination(),), - ), - otel=OTelConfig(enabled=False), - dispatch_attribute_keys=frozenset({"observability_id"}), - ) - ) - runtime._delivery._stdout = output - observability_id = "00000000-0000-4000-8000-000000000004" - artifacts = FakeArtifacts() - - try: - with runtime.operation( - "stage12_report", - remote_context={ - "captured_at": NOW.isoformat(), - "observability_id": observability_id, - }, - ): - 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=FakeStore(), - artifacts=artifacts, - invoker=ConcurrentInvoker(artifacts), - aggregator=lambda **_: _aggregate_payload(), - runtime=runtime, - ) - finally: - runtime.shutdown() - - records = [json.loads(line) for line in output.getvalue().splitlines()] - planning_names = { - "stage12_coordinator_preparation", - "stage12_output_planning", - "stage12_child_input_planning", - } - planning_records = [ - record for record in records if record.get("operation.name") in planning_names - ] - assert len(planning_records) == 4 - assert all( - record["attributes"]["observability_id"] == observability_id - for record in planning_records - ) - - -def test_coordinator_compares_automatic_run_after_successful_aggregation() -> None: - store = FakeStore() - parent = _automatic_parent() - artifacts = FakeArtifacts() - production_result = {"budget": {"total": 100.0, "count": 2}} - stage12_result = {"budget": {"total": 101.0, "count": 2}} - invoker = ConcurrentInvoker( - artifacts, - production_result=production_result, - ) - context = _context().model_copy( - update={"production_function_call_id": "production-job-1"} - ) - - 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 **_: _aggregate_payload(stage12_result), - ) - - assert invoker.events[-1] == "restore:production" - assert invoker.production_wait_timeouts == [900] - assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED - assert store.parent.comparison_status is ResultComparisonStatus.DIFFERENT - assert store.parent.comparison_output_uri == "gs://private/comparison.json" - receipt = artifacts.comparison_writes[0] - assert receipt["difference_count"] == 1 - assert receipt["differences"][0]["path"] == "/budget/total" - assert receipt["differences"][0]["production_value"] == 100.0 - assert receipt["differences"][0]["stage12_value"] == 101.0 - assert "production_result" not in receipt - assert "stage12_result" not in receipt - - -def test_comparison_failure_does_not_change_successful_stage12_report() -> None: - store = FakeStore() - parent = _automatic_parent() - artifacts = FakeArtifacts() - invoker = ConcurrentInvoker(artifacts, production_result="not-an-object") - context = _context().model_copy( - update={"production_function_call_id": "production-job-1"} - ) - - 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 **_: _aggregate_payload({"budget": 100}), - ) - - assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED - assert store.parent.aggregation_status is ComparisonRunAggregationStatus.SUCCEEDED - assert store.parent.comparison_status is ResultComparisonStatus.FAILED - assert store.parent.comparison_error_code == "result_comparison_failed" - assert artifacts.comparison_writes == [] - - -class UnavailableProductionCall: - object_id = "production-job-1" - - def get(self, *, timeout=None): - msg = ( - "production result unavailable at " - "https://storage.example.com/result.json?X-Goog-Signature=signed-value" - ) - raise ConnectionError(msg) - - -class UnavailableProductionInvoker(ConcurrentInvoker): - def restore(self, invocation_id): - if invocation_id == "production-job-1": - return UnavailableProductionCall() - return super().restore(invocation_id) - - -def test_comparison_failure_stores_a_redacted_cause() -> None: - store = FakeStore() - artifacts = FakeArtifacts() - context = _context().model_copy( - update={"production_function_call_id": "production-job-1"} - ) - - coordinate_report( - _report().model_dump(mode="json"), - context.model_dump(mode="json"), - _automatic_parent().model_dump(mode="json"), - application_name=context.modal_application, - coordinator_invocation_id="coordinator-1", - store=store, - artifacts=artifacts, - invoker=UnavailableProductionInvoker(artifacts), - aggregator=lambda **_: _aggregate_payload({"budget": 100}), - ) - - assert store.parent.status is ComparisonRunLifecycleStatus.SUCCEEDED - assert store.parent.comparison_status is ResultComparisonStatus.FAILED - assert store.parent.comparison_error_summary == ( - "ConnectionError: production result unavailable at " - "https://storage.example.com/result.json?[REDACTED]" - ) - - -def test_repeated_submission_does_not_retry_failed_comparison() -> None: - store = FakeStore() - parent = _automatic_parent() - artifacts = FakeArtifacts() - context = _context().model_copy( - update={"production_function_call_id": "production-job-1"} - ) - - first_invoker = ConcurrentInvoker(artifacts, production_result="not-an-object") - 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=first_invoker, - aggregator=lambda **_: _aggregate_payload({"budget": {"total": 100.0}}), - ) - failed_at = store.parent.comparison_completed_at - failed_summary = store.parent.comparison_error_summary - - second_invoker = ConcurrentInvoker( - artifacts, - production_result={"budget": {"total": 100.0}}, - ) - result = 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-2", - store=store, - artifacts=artifacts, - invoker=second_invoker, - aggregator=lambda **_: _aggregate_payload({"budget": {"total": 100.0}}), - ) - - assert result == { - "deduplicated": True, - "evaluation_id": str(EVALUATION_ID), - "status": "succeeded", - } - assert second_invoker.events == [] - assert store.parent.comparison_status is ResultComparisonStatus.FAILED - assert store.parent.comparison_completed_at == failed_at - assert store.parent.comparison_error_code == "result_comparison_failed" - assert store.parent.comparison_error_summary == failed_summary - assert artifacts.comparison_writes == [] - - -def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None: - store = FakeStore() - artifacts = FakeArtifacts() - invoker = ConcurrentInvoker( - artifacts, faults=InjectedChildFaults(fail_role="reform") - ) - - 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=invoker, - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - ) - - assert "child calculation failed" not in str(error.value) - - assert artifacts.aggregate_writes == [] - assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert _failure_cause(store.parent.error_summary) == ( - "RuntimeError in stage12_child_wait (reform): child calculation failed" - ) - failed_child = store.children[REFORM_ID] - assert failed_child.status is ComparisonRunLifecycleStatus.FAILED - assert failed_child.error_code == "simulation_invocation_failed" - assert failed_child.error_summary == store.parent.error_summary - - -@pytest.mark.parametrize( - ("failed_role", "failed_id", "pending_role", "pending_id"), - [ - ("reform", REFORM_ID, "baseline", BASELINE_ID), - ("baseline", BASELINE_ID, "reform", REFORM_ID), - ], -) -def test_coordinator_observes_one_role_failure_while_the_peer_is_pending( - failed_role, - failed_id, - pending_role, - pending_id, -) -> None: - store = FakeStore() - artifacts = FakeArtifacts() - - class PendingCall: - object_id = f"call-{pending_role}" - - def __init__(self): - self.timeouts = [] - self.cancel_count = 0 - - def get(self, *, timeout=None): - self.timeouts.append(timeout) - raise TimeoutError - - def cancel(self): - self.cancel_count += 1 - - class FailedCall: - object_id = f"call-{failed_role}" - - def __init__(self): - self.timeouts = [] - - def get(self, *, timeout=None): - self.timeouts.append(timeout) - child = store.children[failed_id] - failed_at = datetime.now(UTC) - store.replace_simulation( - child.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": "invalid_policy_period", - "error_summary": "The reform period is invalid.", - "updated_at": failed_at, - "completed_at": failed_at, - } - ) - ) - msg = "remote wrapper" - raise RuntimeError(msg) - - def cancel(self): - return None - - baseline_call = PendingCall() - failed_call = FailedCall() - - class ConcurrentFailureInvoker: - def spawn(self, request): - role = PlannedSimulationExecutionInput.model_validate( - request.simulation - ).role - return failed_call if role.value == failed_role else baseline_call - - def restore(self, invocation_id): - msg = f"unexpected restored invocation {invocation_id}" - raise AssertionError(msg) - - 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=ConcurrentFailureInvoker(), - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - sleep=lambda _: pytest.fail("failure should be observed before sleeping"), - ) - - assert failed_call.timeouts == [0] - assert baseline_call.timeouts == ([0] if pending_role == "baseline" else []) - # A segmented peer reads its cancelled record and stops its own - # region-group calls, so it is not cancelled through Modal. - assert baseline_call.cancel_count == 0 - pending = store.children[pending_id] - assert pending.status is ComparisonRunLifecycleStatus.FAILED - assert pending.error_code == "simulation_cancelled" - assert pending.error_summary == ( - f"Cancelled after the {failed_role} simulation failed." - ) - assert store.parent.error_code == "invalid_policy_period" - assert store.parent.error_summary == "The reform period is invalid." - assert artifacts.aggregate_writes == [] - - -def test_dispatch_failure_cancels_an_already_started_peer() -> None: - store = FakeStore() - baseline_call = TimeoutCall("call-baseline") - - class DispatchFailureInvoker: - def spawn(self, request): - role = PlannedSimulationExecutionInput.model_validate( - request.simulation - ).role - if role is SimulationRole.REFORM: - msg = "dispatch contains sensitive details" - raise ConnectionError(msg) - return baseline_call - - def restore(self, invocation_id): - msg = f"unexpected restored invocation {invocation_id}" - raise AssertionError(msg) - - 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=DispatchFailureInvoker(), - ) - - baseline = store.children[BASELINE_ID] - reform = store.children[REFORM_ID] - assert baseline_call.cancelled == 0 - assert baseline.status is ComparisonRunLifecycleStatus.FAILED - assert baseline.error_code == "simulation_cancelled" - assert baseline.error_summary == "Cancelled after the reform simulation failed." - assert reform.status is ComparisonRunLifecycleStatus.FAILED - assert reform.error_code == "simulation_dispatch_failed" - assert store.parent.error_code == "simulation_dispatch_failed" - - -def test_coordinator_propagates_an_already_persisted_child_failure() -> None: - store = FakeStore() - artifacts = FakeArtifacts() - - class PersistedFailureCall: - object_id = "call-reform" - - def get(self, *, timeout=None): - child = store.children[REFORM_ID] - failed_at = datetime.now(UTC) - store.replace_simulation( - child.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": "invalid_policy_period", - "error_summary": "The reform period is invalid.", - "updated_at": failed_at, - "completed_at": failed_at, - } - ) - ) - msg = "remote wrapper must not replace the child failure" - raise RuntimeError(msg) - - class PersistedFailureInvoker: - def spawn(self, request): - parsed = PlannedSimulationExecutionInput.model_validate(request.simulation) - if parsed.role is SimulationRole.REFORM: - return PersistedFailureCall() - descriptor = artifacts.add_simulation(parsed, _frames(100.0)) - return ImmediateCall( - "call-baseline", descriptor.model_dump(mode="json"), [] - ) - - def restore(self, invocation_id): - msg = f"unexpected restored invocation {invocation_id}" - raise AssertionError(msg) - - 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=PersistedFailureInvoker(), - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - ) - - assert store.children[REFORM_ID].error_code == "invalid_policy_period" - assert store.children[REFORM_ID].error_summary == "The reform period is invalid." - assert store.parent.error_code == "invalid_policy_period" - assert store.parent.error_summary == "The reform period is invalid." - - -@pytest.mark.parametrize( - "invoker", - [ - lambda artifacts: ConcurrentInvoker( - artifacts, faults=InjectedChildFaults(wrong_plan_digest=True) - ), - lambda artifacts: ConcurrentInvoker( - artifacts, faults=InjectedChildFaults(missing_column=True) - ), - ], -) -def test_coordinator_rejects_artifacts_that_do_not_satisfy_the_shared_plan( - invoker, -) -> None: - store = FakeStore() - 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=invoker(artifacts), - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - ) - - assert artifacts.aggregate_writes == [] - assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - - -def test_coordinator_records_failure_when_child_persistence_fails() -> None: - class FailingChildStore(FakeStore): - @override - def create_or_resolve_simulation(self, record): - msg = "database rejected the child record" - raise RuntimeError(msg) - - store = FailingChildStore() - - 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=ConcurrentInvoker(FakeArtifacts()), - ) - - assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_code == "report_coordination_failed" - assert _failure_cause(store.parent.error_summary) == ( - "RuntimeError in stage12_child_state_create: database rejected the child record" - ) - - -class WorkerRecordedFailureCall: - """A child whose worker persisted its cause before Modal raised.""" - - def __init__(self, object_id, store, simulation_execution_id, worker_summary): - self.object_id = object_id - self.store = store - self.simulation_execution_id = simulation_execution_id - self.worker_summary = worker_summary - - def get(self, *, timeout=None): - child = self.store.children[self.simulation_execution_id] - self.store.children[self.simulation_execution_id] = child.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": "simulation_execution_failed", - "error_summary": self.worker_summary, - "completed_at": NOW, - } - ) - msg = "Stage 12 simulation execution failed" - raise RuntimeError(msg) - - -class UnreachedCall: - def __init__(self, object_id): - self.object_id = object_id - - def get(self, *, timeout=None): - msg = "the coordinator observes baseline first" - raise AssertionError(msg) - - def cancel(self): - return None - - -class WorkerRecordedFailureInvoker: - def __init__(self, store, worker_summary): - self.store = store - self.worker_summary = worker_summary - - def spawn(self, request): - planned = PlannedSimulationExecutionInput.model_validate(request.simulation) - if planned.role is SimulationRole.BASELINE: - return WorkerRecordedFailureCall( - "call-baseline", - self.store, - planned.simulation_execution_id, - self.worker_summary, - ) - return UnreachedCall("call-reform") - - def restore(self, invocation_id): - msg = f"unexpected restored invocation {invocation_id}" - raise AssertionError(msg) - - -def test_coordinator_keeps_the_cause_recorded_by_a_failed_worker() -> None: - store = FakeStore() - worker_summary = ( - "ValueError in stage12_calculation (baseline): " - "Variable 'country' not found in household data " - "(correlation_id=0123456789abcdef0123456789abcdef)" - ) - runtime = RecordingRuntime() - - 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=WorkerRecordedFailureInvoker(store, worker_summary), - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - runtime=runtime, - ) - - failed_child = store.children[BASELINE_ID] - assert failed_child.error_code == "simulation_execution_failed" - assert failed_child.error_summary == worker_summary - assert store.parent.error_code == "simulation_execution_failed" - assert store.parent.error_summary == worker_summary - [recorded_cause] = runtime.recorded_exceptions - assert str(recorded_cause) == worker_summary - - -def test_coordinator_names_the_stage_that_rejected_the_aggregate() -> None: - class RejectingAggregateArtifacts(FakeArtifacts): - def write_aggregate(self, *, prefix, payload): - SPMComparisonProvenance.model_validate(payload["result"]["spm_provenance"]) - - store = FakeStore() - artifacts = RejectingAggregateArtifacts() - - 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), - aggregator=lambda **_: {"result": {"spm_provenance": None}}, - ) - - assert _failure_cause(store.parent.error_summary) == ( - "ValidationError in stage12_aggregate_artifact_write: " - "1 validation error for SPMComparisonProvenance: " - "Input should be a valid dictionary or instance " - "of SPMComparisonProvenance [model_type]" - ) - - -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 **_: _aggregate_payload(), - 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() -> ( # noqa: E501 # the test name spells out the lost-commit scenario - 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", - "evaluation_id": str(EVALUATION_ID), - "correlation_id": _correlation_id(failure_write[1]), - } - assert original_report[1][2] is aggregation_error - assert failure_write[0] == "failure_write" - assert _failure_cause(failure_write[1]) == ( - "ValueError in stage12_aggregation: aggregate failed validation" - ) - 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, faults=InjectedChildFaults(fail_role="baseline") - ), - aggregator=lambda **_: _aggregate_payload({"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", - "evaluation_id": str(EVALUATION_ID), - "simulation_execution_id": str(BASELINE_ID), - "correlation_id": _correlation_id(store.parent.error_summary), - } - ] - 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_code == "simulation_invocation_failed" - assert _failure_cause(store.parent.error_summary) == ( - "RuntimeError in stage12_child_wait (baseline): child calculation failed" - ) - - -class RefusingDispatchInvoker: - def __init__(self, dispatch_error): - self.dispatch_error = dispatch_error - self.events = [] - - def spawn(self, request): - role = PlannedSimulationExecutionInput.model_validate( - request.simulation - ).role.value - self.events.append(f"spawn:{role}") - raise self.dispatch_error - - def restore(self, invocation_id): - msg = f"unexpected restored invocation {invocation_id}" - raise AssertionError(msg) - - -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 **_: _aggregate_payload({"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" - ] - baseline_child = store.children[BASELINE_ID] - assert [(attributes, error) for _message, attributes, error in failure_reports] == [ - ( - { - "stage12_failure_scope": "child_dispatch", - "simulation_role": "baseline", - "evaluation_id": str(EVALUATION_ID), - "simulation_execution_id": str(BASELINE_ID), - "correlation_id": _correlation_id(baseline_child.error_summary), - }, - dispatch_error, - ) - ] - assert invoker.events == ["spawn:baseline"] - assert baseline_child.status is ComparisonRunLifecycleStatus.FAILED - assert baseline_child.error_code == "simulation_dispatch_failed" - assert _failure_cause(baseline_child.error_summary) == ( - "ConnectionError in stage12_child_dispatch (baseline): " - "Modal refused the dispatch" - ) - assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == baseline_child.error_summary - - -def test_coordinator_records_a_redacted_dispatch_failure_cause() -> None: - store = FakeStore() - artifacts = FakeArtifacts() - invoker = RefusingDispatchInvoker( - ConnectionError("Modal rejected the call with token=opaque-token") - ) - runtime = RecordingRuntime() - - 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=invoker, - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - runtime=runtime, - ) - - assert invoker.events == ["spawn:baseline"] - assert artifacts.aggregate_writes == [] - failed_child = store.children[BASELINE_ID] - assert failed_child.status is ComparisonRunLifecycleStatus.FAILED - assert failed_child.error_code == "simulation_dispatch_failed" - assert _failure_cause(failed_child.error_summary) == ( - "ConnectionError in stage12_child_dispatch (baseline): " - "Modal rejected the call with token=[REDACTED]" - ) - assert store.parent.error_code == "simulation_dispatch_failed" - assert store.parent.error_summary == failed_child.error_summary - [recorded_cause] = runtime.recorded_exceptions - assert str(recorded_cause) == failed_child.error_summary - - -def test_coordinator_rejects_incompatible_row_identity() -> None: - store = FakeStore() - artifacts = FakeArtifacts() - invoker = ConcurrentInvoker( - artifacts, faults=InjectedChildFaults(incompatible=True) - ) - - 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=invoker, - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - ) - assert artifacts.aggregate_writes == [] - - -def test_coordinator_reuses_matching_successful_child() -> None: - report = _report() - store = FakeStore() - artifacts = FakeArtifacts() - baseline_descriptor = artifacts.add_simulation(report.baseline) - baseline = _child(report.baseline).model_copy( - update={ - "status": ComparisonRunLifecycleStatus.SUCCEEDED, - "output_uri": baseline_descriptor.artifact.uri, - "output_sha256": baseline_descriptor.artifact.content_sha256, - "output_schema_version": 1, - "row_identity_columns": baseline_descriptor.row_identity.identifier_columns, - "row_count": baseline_descriptor.row_identity.row_count, - "row_identity_sha256": baseline_descriptor.row_identity.identity_sha256, - "completed_at": NOW, - } - ) - store.children[report.baseline.simulation_execution_id] = baseline - 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 **_: _aggregate_payload(), - ) - - assert invoker.events == ["spawn:reform"] - - -def test_coordinator_resumes_a_matching_running_child_invocation() -> None: - report = _report() - store = FakeStore() - artifacts = FakeArtifacts() - store.children[report.baseline.simulation_execution_id] = _child( - report.baseline - ).model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "modal_invocation_id": "call-baseline", - "started_at": NOW, - } - ) - 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 **_: _aggregate_payload(), - ) - - assert invoker.events == ["restore:baseline", "spawn:reform"] - - -class TimeoutCall: - def __init__(self, object_id): - self.object_id = object_id - self.cancelled = 0 - - def get(self, *, timeout=None): - msg = "contains sensitive timeout details" - raise TimeoutError(msg) - - def cancel(self): - self.cancelled += 1 - - -class TimeoutInvoker: - def __init__(self): - self.events = [] - - def spawn(self, request): - role = PlannedSimulationExecutionInput.model_validate( - request.simulation - ).role.value - self.events.append(f"spawn:{role}") - return TimeoutCall(f"call-{role}") - - def restore(self, invocation_id): - msg = f"unexpected restored invocation {invocation_id}" - raise AssertionError(msg) - - -def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: - store = FakeStore() - artifacts = FakeArtifacts() - invoker = TimeoutInvoker() - - 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=invoker, - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - wait_timeout_seconds=0, - ) - - assert "sensitive" not in str(error.value) - - assert invoker.events == ["spawn:baseline", "spawn:reform"] - assert artifacts.aggregate_writes == [] - assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_code == "simulation_wait_timeout" - assert store.parent.error_summary == ( - "Stage 12 simulations did not finish within 0 seconds of the report " - "coordinator starting." - ) - assert all( - child.status is ComparisonRunLifecycleStatus.FAILED - for child in store.children.values() - ) - assert all( - child.error_code == "simulation_cancelled" for child in store.children.values() - ) - - -def test_closed_wait_window_is_reported_before_its_failure_records() -> None: - ordered_events: list[tuple[str, object]] = [] - - class UnwritableFailureStore(FakeStore): - def replace_simulation_if_status(self, record, *, expected_status): - if record.status is ComparisonRunLifecycleStatus.FAILED: - ordered_events.append(("failure_write", record.error_code)) - raise _connection_limit_error() - return super().replace_simulation_if_status( - record, expected_status=expected_status - ) - - def replace_report(self, record): - if record.status is ComparisonRunLifecycleStatus.FAILED: - ordered_events.append(("failure_write", record.error_code)) - raise _connection_limit_error() - return super().replace_report(record) - - store = UnwritableFailureStore() - invoker = TimeoutInvoker() - - 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 **_: _aggregate_payload({"must": "not run"}), - runtime=TelemetryRecordingRuntime(ordered_events), - wait_timeout_seconds=0, - ) - - original_report, *failure_writes, persistence_report = ordered_events - event_kind, (message, attributes, reported_error) = original_report - assert event_kind == "telemetry" - assert message == "Stage 12 execution failed" - assert attributes == {"stage12_failure_scope": "child_wait_timeout"} - assert reported_error.detail.error_code == "simulation_wait_timeout" - # Both cancellation records and the report record were attempted after - # the closed window had already been reported. - assert failure_writes == [ - ("failure_write", "simulation_cancelled"), - ("failure_write", "simulation_cancelled"), - ("failure_write", "simulation_wait_timeout"), - ] - assert persistence_report[1][0] == ( - "Stage 12 failure record could not be persisted" - ) - assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING - - -class StartupWriteError(Exception): - """A coordinator persistence failure after a child may have started.""" - - -@dataclass(frozen=True) -class ChildCallScript: - """What a scripted child call does when probed or cancelled. - - It returns ``result`` or raises ``wait_error`` once probed (``during_wait`` - runs first), and raises ``cancel_error`` when cancelled; with none set, it - keeps running. - """ - - result: object = None - wait_error: Exception | None = None - cancel_error: Exception | None = None - during_wait: object = None - - -class ScriptedChildCall: - """A started child that keeps running unless the test scripts an outcome.""" - - def __init__(self, object_id, *, clock=None, script=None): - script = script or ChildCallScript() - self.object_id = object_id - self.clock = clock - self.result = script.result - self.wait_error = script.wait_error - self.cancel_error = script.cancel_error - self.during_wait = script.during_wait - self.probed_at = [] - self.cancel_count = 0 - - def get(self, *, timeout=None): - self.probed_at.append( - self.clock.monotonic() if self.clock is not None else None - ) - if self.during_wait is not None: - self.during_wait() - if self.wait_error is not None: - raise self.wait_error - if self.result is not None: - return self.result - msg = "still running" - raise TimeoutError(msg) - - def cancel(self): - self.cancel_count += 1 - if self.cancel_error is not None: - raise self.cancel_error - - -class ScriptedChildInvoker: - def __init__( - self, - *, - clock=None, - results=None, - wait_errors=None, - cancel_errors=None, - during_wait=None, - ): - self.clock = clock - self.results = results or {} - self.wait_errors = wait_errors or {} - self.cancel_errors = cancel_errors or {} - self.during_wait = during_wait or {} - self.events = [] - self.calls = {} - - def spawn(self, request): - role = PlannedSimulationExecutionInput.model_validate( - request.simulation - ).role.value - self.events.append(f"spawn:{role}") - return self._call(role) - - def restore(self, invocation_id): - role = invocation_id.removeprefix("call-") - self.events.append(f"restore:{role}") - return self._call(role) - - def _call(self, role): - call = ScriptedChildCall( - f"call-{role}", - clock=self.clock, - script=ChildCallScript( - result=self.results.get(role), - wait_error=self.wait_errors.get(role), - cancel_error=self.cancel_errors.get(role), - during_wait=self.during_wait.get(role), - ), - ) - self.calls[role] = call - return call - - -class PlaceholderWriteFailingStore(FakeStore): - def __init__(self, failing_simulation_id): - super().__init__() - self.failing_simulation_id = failing_simulation_id - - def replace_simulation(self, record): - # The coordinator's only plain replacement is the dispatch placeholder. - if record.simulation_execution_id == self.failing_simulation_id: - msg = "placeholder write contains sensitive details" - raise StartupWriteError(msg) - return super().replace_simulation(record) - - -class AttachFailingStore(FakeStore): - def __init__(self, failing_simulation_id): - super().__init__() - self.failing_simulation_id = failing_simulation_id - - def attach_simulation_invocation(self, simulation_execution_id, **values): - if simulation_execution_id == self.failing_simulation_id: - msg = "attach contains sensitive details" - raise StartupWriteError(msg) - return super().attach_simulation_invocation(simulation_execution_id, **values) - - -class CancellationStateFailingStore(FakeStore): - def replace_simulation_if_status(self, record, *, expected_status): - if record.error_code == "simulation_cancelled": - msg = "database contains sensitive details" - raise ConnectionError(msg) - return super().replace_simulation_if_status( - record, expected_status=expected_status - ) - - -def _simulation_id(role: str) -> UUID: - return BASELINE_ID if role == "baseline" else REFORM_ID - - -def _report_with_segmentation(*, segmented: bool) -> ReportExecutionInput: - if segmented: - return _report() - - def unsegmented(role: SimulationRole) -> SimulationExecutionInput: - return _simulation(role).model_copy(update={"options": {"segmented": False}}) - - return ReportExecutionInput( - evaluation_id=EVALUATION_ID, - baseline=unsegmented(SimulationRole.BASELINE), - reform=unsegmented(SimulationRole.REFORM), - requested_aggregates=tuple(ReportAggregate), - ) - - -def _coordinate_with(report, store, invoker, **coordinator_options): - coordinator_options.setdefault("artifacts", FakeArtifacts()) - return coordinate_report( - report.model_dump(mode="json"), - _context().model_dump(mode="json"), - _parent().model_dump(mode="json"), - application_name=_context().modal_application, - coordinator_invocation_id="coordinator-1", - store=store, - invoker=invoker, - aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - **coordinator_options, - ) - - -def _assert_cancelled(child, *, after: str) -> None: - assert child.status is ComparisonRunLifecycleStatus.FAILED - assert child.error_code == "simulation_cancelled" - assert child.error_summary == f"Cancelled after {after}." - assert child.completed_at is not None - - -def _assert_original_coordinator_failure(error, store) -> None: - assert isinstance(error.value.__context__, StartupWriteError) - assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_code == "report_coordination_failed" - original_error = error.value.__context__ - assert _failure_cause(store.parent.error_summary).startswith( - f"{type(original_error).__name__} " - ) - assert _failure_cause(store.parent.error_summary).endswith(f": {original_error}") - - -def _coordination_log(caplog, level): - (record,) = [ - record - for record in caplog.records - if record.name == coordination_module.__name__ and record.levelno == level - ] - return record - - -def _assert_log_identifies_child(record, *, role: str, error_type: str) -> None: - simulation_execution_id = str(_simulation_id(role)) - assert record.evaluation_id == str(EVALUATION_ID) - assert record.simulation_execution_id == simulation_execution_id - assert record.simulation_role == role - assert record.error_type == error_type - message = record.getMessage() - assert str(EVALUATION_ID) in message - assert simulation_execution_id in message - assert "sensitive" not in message - assert record.exc_info is None - - -def test_baseline_failure_cancels_and_records_the_reform_child() -> None: - store = FakeStore() - artifacts = FakeArtifacts() - invoker = ScriptedChildInvoker( - wait_errors={"baseline": LookupError("baseline worker failed")} - ) - - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - _coordinate_with( - _report_with_segmentation(segmented=False), - store, - invoker, - artifacts=artifacts, - ) - - reform_call = invoker.calls["reform"] - assert reform_call.cancel_count == 1 - assert reform_call.probed_at == [] - assert artifacts.aggregate_writes == [] - assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" - _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") - assert store.parent.error_code == "simulation_invocation_failed" - - -def test_reform_failure_leaves_the_finished_baseline_alone() -> None: - report = _report_with_segmentation(segmented=False) - store = FakeStore() - artifacts = FakeArtifacts() - baseline_descriptor = artifacts.add_simulation(report.baseline) - invoker = ScriptedChildInvoker( - results={"baseline": baseline_descriptor.model_dump(mode="json")}, - wait_errors={"reform": LookupError("reform worker failed")}, - ) - - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - _coordinate_with(report, store, invoker, artifacts=artifacts) - - assert invoker.calls["baseline"].cancel_count == 0 - assert store.children[BASELINE_ID].status is ComparisonRunLifecycleStatus.RUNNING - assert store.children[BASELINE_ID].error_code is None - assert store.children[REFORM_ID].error_code == "simulation_invocation_failed" - - -def test_baseline_failure_cancels_a_reform_call_restored_on_retry() -> None: - report = _report_with_segmentation(segmented=False) - store = FakeStore() - # An earlier coordinator attempt started the reform child before stopping. - store.children[REFORM_ID] = _child(report.reform).model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "modal_invocation_id": "call-reform", - "started_at": NOW, - } - ) - invoker = ScriptedChildInvoker( - wait_errors={"baseline": LookupError("baseline worker failed")} - ) - - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - _coordinate_with(report, store, invoker) - - assert invoker.events == ["spawn:baseline", "restore:reform"] - assert invoker.calls["reform"].cancel_count == 1 - _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") - assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" - - -@pytest.mark.parametrize( - ("worker_status", "worker_error_code"), - [ - (ComparisonRunLifecycleStatus.SUCCEEDED, None), - (ComparisonRunLifecycleStatus.FAILED, "simulation_execution_failed"), - ], -) -def test_cancellation_leaves_a_sibling_that_already_finished_unchanged( - worker_status, - worker_error_code, -) -> None: - store = FakeStore() - worker_records = [] - - def reform_worker_records_its_outcome(): - # The reform worker finishes while the coordinator probes baseline. - finished_at = NOW + timedelta(minutes=5) - worker_records.append( - store.replace_simulation( - store.get_simulation(REFORM_ID).model_copy( - update={ - "status": worker_status, - "error_code": worker_error_code, - "updated_at": finished_at, - "completed_at": finished_at, - } - ) - ) - ) - - invoker = ScriptedChildInvoker( - wait_errors={"baseline": LookupError("baseline worker failed")}, - during_wait={"baseline": reform_worker_records_its_outcome}, - ) - - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) - - # Cancelling a finished Modal call is harmless; its record is not touched. - assert invoker.calls["reform"].cancel_count == 1 - assert store.children[REFORM_ID] == worker_records[0] - assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" - - -def test_peer_cancellation_failure_does_not_replace_the_child_failure( - caplog, -) -> None: - store = FakeStore() - invoker = ScriptedChildInvoker( - wait_errors={"baseline": LookupError("baseline failure is sensitive")}, - cancel_errors={ - "reform": ConnectionError("cancel request contains sensitive details") - }, - ) - - with ( - caplog.at_level(logging.WARNING, logger=coordination_module.__name__), - pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), - ): - _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) - - assert invoker.calls["reform"].cancel_count == 1 - _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") - assert store.parent.error_code == "simulation_invocation_failed" - _assert_log_identifies_child( - _coordination_log(caplog, logging.WARNING), - role="reform", - error_type="ConnectionError", - ) - - -@pytest.mark.parametrize("segmented", [True, False]) -def test_failed_cancellation_state_write_still_cancels_through_modal( - caplog, - segmented, -) -> None: - store = CancellationStateFailingStore() - invoker = ScriptedChildInvoker( - wait_errors={"baseline": LookupError("baseline failure is sensitive")}, - ) - - with ( - caplog.at_level(logging.WARNING, logger=coordination_module.__name__), - pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), - ): - _coordinate_with(_report_with_segmentation(segmented=segmented), store, invoker) - - # A segmented worker cannot read a cancellation that was never written, so - # it is cancelled through Modal as well. - assert invoker.calls["reform"].cancel_count == 1 - assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.RUNNING - assert store.parent.error_code == "simulation_invocation_failed" - _assert_log_identifies_child( - _coordination_log(caplog, logging.ERROR), - role="reform", - error_type="ConnectionError", - ) - - -def test_wait_timeout_cancels_and_records_every_unfinished_child() -> None: - store = FakeStore() - invoker = ScriptedChildInvoker() - - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - _coordinate_with( - _report_with_segmentation(segmented=False), - store, - invoker, - wait_timeout_seconds=0, - ) - - assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] - for role in ("baseline", "reform"): - _assert_cancelled( - store.children[_simulation_id(role)], - after="the 0-second wait window closed", - ) - assert store.parent.error_code == "simulation_wait_timeout" - - -def test_timeout_cancellation_failure_log_identifies_the_report_and_child( - caplog, -) -> None: - store = FakeStore() - invoker = ScriptedChildInvoker( - cancel_errors={ - "reform": ConnectionError("cancel request contains sensitive details") - }, - ) - - with ( - caplog.at_level(logging.WARNING, logger=coordination_module.__name__), - pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), - ): - _coordinate_with( - _report_with_segmentation(segmented=False), - store, - invoker, - wait_timeout_seconds=0, - ) - - assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] - assert store.parent.error_code == "simulation_wait_timeout" - _assert_log_identifies_child( - _coordination_log(caplog, logging.WARNING), - role="reform", - error_type="ConnectionError", - ) - - -@pytest.mark.parametrize("segmented", [True, False]) -def test_failed_placeholder_write_stops_the_started_baseline(segmented) -> None: - store = PlaceholderWriteFailingStore(REFORM_ID) - invoker = ScriptedChildInvoker() - - with pytest.raises( - RuntimeError, match="Stage 12 report coordination failed" - ) as error: - _coordinate_with(_report_with_segmentation(segmented=segmented), store, invoker) - - _assert_original_coordinator_failure(error, store) - assert invoker.events == ["spawn:baseline"] - _assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") - # A segmented worker stops itself, and its region-group calls, once it - # reads the cancelled record; Modal does not cancel spawned calls - # recursively. - assert invoker.calls["baseline"].cancel_count == (0 if segmented else 1) - assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.PENDING - - -@pytest.mark.parametrize( - ("failing_role", "started_roles"), - [ - ("baseline", ["baseline"]), - ("reform", ["baseline", "reform"]), - ], -) -def test_failed_invocation_attach_stops_every_started_child( - failing_role, - started_roles, -) -> None: - store = AttachFailingStore(_simulation_id(failing_role)) - invoker = ScriptedChildInvoker() - - with pytest.raises( - RuntimeError, match="Stage 12 report coordination failed" - ) as error: - _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) - - _assert_original_coordinator_failure(error, store) - assert invoker.events == [f"spawn:{role}" for role in started_roles] - # The child whose invocation could not be attached is stopped as well. - for role in started_roles: - _assert_cancelled( - store.children[_simulation_id(role)], after="the coordinator failed" - ) - assert invoker.calls[role].cancel_count == 1 - if failing_role == "baseline": - assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.PENDING - - -def test_startup_failure_stops_a_call_restored_from_an_earlier_attempt() -> None: - report = _report_with_segmentation(segmented=False) - store = PlaceholderWriteFailingStore(REFORM_ID) - store.children[BASELINE_ID] = _child(report.baseline).model_copy( - update={ - "status": ComparisonRunLifecycleStatus.RUNNING, - "modal_invocation_id": "call-baseline", - "started_at": NOW, - } - ) - invoker = ScriptedChildInvoker() - - with pytest.raises( - RuntimeError, match="Stage 12 report coordination failed" - ) as error: - _coordinate_with(report, store, invoker) - - _assert_original_coordinator_failure(error, store) - assert invoker.events == ["restore:baseline"] - _assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") - assert invoker.calls["baseline"].cancel_count == 1 - - -def test_startup_failure_leaves_a_finished_child_unchanged() -> None: - report = _report_with_segmentation(segmented=False) - store = AttachFailingStore(REFORM_ID) - artifacts = FakeArtifacts() - baseline_descriptor = artifacts.add_simulation(report.baseline) - finished_baseline = _child(report.baseline).model_copy( - update={ - "status": ComparisonRunLifecycleStatus.SUCCEEDED, - "output_uri": baseline_descriptor.artifact.uri, - "output_sha256": baseline_descriptor.artifact.content_sha256, - "output_schema_version": 1, - "row_identity_columns": baseline_descriptor.row_identity.identifier_columns, - "row_count": baseline_descriptor.row_identity.row_count, - "row_identity_sha256": baseline_descriptor.row_identity.identity_sha256, - "completed_at": NOW, - } - ) - store.children[BASELINE_ID] = finished_baseline - invoker = ScriptedChildInvoker() - - with pytest.raises( - RuntimeError, match="Stage 12 report coordination failed" - ) as error: - _coordinate_with(report, store, invoker, artifacts=artifacts) - - _assert_original_coordinator_failure(error, store) - assert invoker.events == ["spawn:reform"] - assert store.children[BASELINE_ID] == finished_baseline - _assert_cancelled(store.children[REFORM_ID], after="the coordinator failed") - assert invoker.calls["reform"].cancel_count == 1 - - -def test_startup_cancellation_error_does_not_replace_the_original_failure( - caplog, -) -> None: - store = PlaceholderWriteFailingStore(REFORM_ID) - invoker = ScriptedChildInvoker( - cancel_errors={ - "baseline": ConnectionError("cancel request contains sensitive details") - } - ) - - with ( - caplog.at_level(logging.WARNING, logger=coordination_module.__name__), - pytest.raises( - RuntimeError, match="Stage 12 report coordination failed" - ) as error, - ): - _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) - - _assert_original_coordinator_failure(error, store) - assert invoker.calls["baseline"].cancel_count == 1 - _assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") - _assert_log_identifies_child( - _coordination_log(caplog, logging.WARNING), - role="baseline", - error_type="ConnectionError", - ) - - -def test_startup_cancellation_state_error_does_not_replace_the_original_failure( - caplog, -) -> None: - class FailingStore(PlaceholderWriteFailingStore, CancellationStateFailingStore): - pass - - store = FailingStore(REFORM_ID) - invoker = ScriptedChildInvoker() - - with ( - caplog.at_level(logging.WARNING, logger=coordination_module.__name__), - pytest.raises( - RuntimeError, match="Stage 12 report coordination failed" - ) as error, - ): - _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) - - _assert_original_coordinator_failure(error, store) - assert invoker.calls["baseline"].cancel_count == 1 - _assert_log_identifies_child( - _coordination_log(caplog, logging.ERROR), - role="baseline", - error_type="ConnectionError", - ) - - -class ChildStateError(Exception): - """A database failure while reading or writing a child's own record.""" - - -class OwnChildStateFailingStore(FakeStore): - """Fail one read, or one failed-status write, of one child's record once. - - Cancellation writes are left alone, so only the failed child's own state - update can fail. - """ - - def __init__(self, *, failing_simulation_id, failing_operation): - super().__init__() - self.failing_simulation_id = failing_simulation_id - self.failing_operation = failing_operation - self.remaining_failures = 1 - - def _fail_once(self, operation, simulation_execution_id): - if ( - operation == self.failing_operation - and simulation_execution_id == self.failing_simulation_id - and self.remaining_failures > 0 - ): - self.remaining_failures -= 1 - msg = "database contains sensitive details" - raise ChildStateError(msg) - - def get_simulation(self, simulation_execution_id): - self._fail_once("read", simulation_execution_id) - return super().get_simulation(simulation_execution_id) - - def replace_simulation_if_status(self, record, *, expected_status): - if record.error_code != "simulation_cancelled": - self._fail_once("write", record.simulation_execution_id) - return super().replace_simulation_if_status( - record, expected_status=expected_status - ) - - -@pytest.mark.parametrize( - ("failing_operation", "baseline_status", "baseline_error_code"), - [ - # The failure is still recorded: the read before the conditional - # write succeeds. - ( - "read", - ComparisonRunLifecycleStatus.FAILED, - "simulation_invocation_failed", - ), - ("write", ComparisonRunLifecycleStatus.RUNNING, None), - ], -) -def test_failed_child_state_error_does_not_replace_the_child_failure( - caplog, - failing_operation, - baseline_status, - baseline_error_code, -) -> None: - store = OwnChildStateFailingStore( - failing_simulation_id=BASELINE_ID, - failing_operation=failing_operation, - ) - invoker = ScriptedChildInvoker( - wait_errors={"baseline": LookupError("baseline failure is sensitive")}, - ) - - with ( - caplog.at_level(logging.WARNING, logger=coordination_module.__name__), - pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), - ): - _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) - - assert store.parent.error_code == "simulation_invocation_failed" - _correlation_id(store.parent.error_summary) - baseline = store.children[BASELINE_ID] - assert baseline.status is baseline_status - assert baseline.error_code == baseline_error_code - if baseline_error_code is not None: - assert baseline.error_summary == store.parent.error_summary - # The failed child is neither cancelled nor recorded as cancelled. - assert invoker.calls["baseline"].cancel_count == 0 - _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") - assert invoker.calls["reform"].cancel_count == 1 - _assert_log_identifies_child( - _coordination_log(caplog, logging.ERROR), - role="baseline", - error_type="ChildStateError", - ) - - -class ReformDispatchFailureInvoker(ScriptedChildInvoker): - def spawn(self, request): - role = PlannedSimulationExecutionInput.model_validate(request.simulation).role - if role is SimulationRole.REFORM: - self.events.append("spawn:reform") - msg = "dispatch contains sensitive details" - raise ConnectionError(msg) - return super().spawn(request) - - -@pytest.mark.parametrize("failing_operation", ["read", "write"]) -def test_dispatch_failure_state_error_does_not_replace_the_child_failure( - caplog, - failing_operation, -) -> None: - store = OwnChildStateFailingStore( - failing_simulation_id=REFORM_ID, - failing_operation=failing_operation, - ) - invoker = ReformDispatchFailureInvoker() - - with ( - caplog.at_level(logging.WARNING, logger=coordination_module.__name__), - pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), - ): - _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) - - assert invoker.events == ["spawn:baseline", "spawn:reform"] - assert store.parent.error_code == "simulation_dispatch_failed" - _correlation_id(store.parent.error_summary) - # The reform record could not be updated, so it keeps its dispatch - # placeholder rather than a cancellation. - reform = store.children[REFORM_ID] - assert reform.status is ComparisonRunLifecycleStatus.RUNNING - assert reform.error_code is None - _assert_cancelled(store.children[BASELINE_ID], after="the reform simulation failed") - assert invoker.calls["baseline"].cancel_count == 1 - _assert_log_identifies_child( - _coordination_log(caplog, logging.ERROR), - role="reform", - error_type="ChildStateError", - ) - - -class ExceptionRecordingRuntime(RecordingRuntime): - """Record every telemetry call that can report a failure.""" - - def __init__(self) -> None: - super().__init__() - self.reports: list[tuple[str, object, dict, object]] = [] - - def log(self, message, *, severity="INFO", attributes=None, error=None): - self.reports.append(("log", message, dict(attributes or {}), error)) - - def record_exception(self, error, *, handled, status_code=None): - self.reports.append(("exception", None, {}, error)) - - def event(self, name, *, attributes=None, **_options): - self.reports.append(("event", name, dict(attributes or {}), None)) - - def set_context(self, **_attributes): - return None - - -def _fail_simulation_calculation(runtime): - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - calculation_error = ValueError("calculation contains sensitive details") - - 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", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=fail, - runtime=runtime, - ), - ) - return calculation_error, store.children[BASELINE_ID].error_summary - - -def _fail_child_call(runtime): - store = FakeStore() - child_error = LookupError("baseline failure contains sensitive details") - invoker = ScriptedChildInvoker(wait_errors={"baseline": child_error}) - - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - _coordinate_with( - _report_with_segmentation(segmented=False), - store, - invoker, - runtime=runtime, - ) - return child_error, store.parent.error_summary - - -def _fail_child_dispatch(runtime): - store = FakeStore() - dispatch_error = ConnectionError("dispatch contains sensitive details") - - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - _coordinate_with( - _report(), - store, - RefusingDispatchInvoker(dispatch_error), - runtime=runtime, - ) - return dispatch_error, store.parent.error_summary - - -def _fail_report_aggregation(runtime): - store = FakeStore() - artifacts = FakeArtifacts() - aggregation_error = ValueError("aggregation contains sensitive details") - - def fail_aggregation(**_values): - raise aggregation_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=artifacts, - invoker=ConcurrentInvoker(artifacts), - aggregator=fail_aggregation, - runtime=runtime, - ) - return aggregation_error, store.parent.error_summary - - -@pytest.mark.parametrize( - "fail", - [ - _fail_simulation_calculation, - _fail_child_call, - _fail_child_dispatch, - _fail_report_aggregation, - ], -) -@pytest.mark.parametrize("with_runtime", [True, False], ids=["runtime", "no-runtime"]) -def test_an_unexpected_failure_is_reported_exactly_once( - caplog, - fail, - with_runtime, -) -> None: - runtime = ExceptionRecordingRuntime() if with_runtime else None - - with caplog.at_level(logging.DEBUG): - original_error, error_summary = fail(runtime) - - correlation_id = _correlation_id(error_summary) - logged_failures = [record for record in caplog.records if record.exc_info] - if runtime is not None: - [(kind, message, attributes, reported_error)] = [ - report for report in runtime.reports if report[0] == "log" - ] - assert (kind, message) == ("log", "Stage 12 execution failed") - assert reported_error is original_error - assert attributes["correlation_id"] == correlation_id - # The redacted cause, never the original error, is also attached to - # the failing operation, once. - [recorded_cause] = [ - report[3] for report in runtime.reports if report[0] == "exception" - ] - assert str(recorded_cause) == error_summary - assert sorted(report[0] for report in runtime.reports) == ["exception", "log"] - assert logged_failures == [] - else: - [logged_failure] = logged_failures - assert logged_failure.exc_info[1] is original_error - assert f"correlation_id={correlation_id}" in logged_failure.getMessage() - - -def test_a_segment_failure_is_reported_exactly_once(caplog) -> None: - store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) - runtime = ExceptionRecordingRuntime() - segment_error = RuntimeError("segment contains sensitive details") - - class SegmentFunction: - def spawn(self, payload, *, observability_context): - segment_index = payload["segment_index"] - # The first segment fails on every probe; the others keep running. - return ScriptedChildCall( - f"segment-{segment_index}", - script=ChildCallScript( - wait_error=segment_error if segment_index == 0 else None - ), - ) - - class SegmentFunctionLookup: - @staticmethod - def from_name(_app_name, _function_name): - return SegmentFunction() - - def segmented_calculator(value, cancellation_requested): - return Stage12SegmentRunner( - value, - app_name="stage12-app", - runtime=runtime, - cancellation_requested=cancellation_requested, - seams=SegmentRunnerSeams( - modal_module=SimpleNamespace(Function=SegmentFunctionLookup), - groups=US_REGION_GROUPS, - sleep=lambda _seconds: None, - ), - ).run() - - with ( - caplog.at_level(logging.DEBUG), - 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", - dependencies=SimulationWorkerDependencies( - store=store, - artifacts=FakeArtifacts(), - calculator=lambda _: pytest.fail("monolithic calculator was selected"), - segmented_calculator=segmented_calculator, - runtime=runtime, - ), - ) - - child = store.children[BASELINE_ID] - assert child.error_code == "segment_execution_failed" - [(kind, message, attributes, reported_error)] = [ - report for report in runtime.reports if report[0] == "log" - ] - [recorded_cause] = [ - report[3] for report in runtime.reports if report[0] == "exception" - ] - assert str(recorded_cause) == child.error_summary - assert (kind, message) == ("log", "Stage 12 execution failed") - assert reported_error is segment_error - assert attributes == { - "stage12_failure_scope": "segment_call", - "simulation_role": "baseline", - "evaluation_id": str(EVALUATION_ID), - "simulation_execution_id": str(BASELINE_ID), - "segment_index": 0, - "correlation_id": _correlation_id( - child.error_summary.removeprefix("Segment 0: ") - ), - } - assert [record for record in caplog.records if record.exc_info] == [] - - -class SteppedClock: - def __init__(self, start_seconds=10_000.0): - self.current_seconds = start_seconds - - def monotonic(self): - return self.current_seconds - - def advance(self, elapsed_seconds): - self.current_seconds += elapsed_seconds - - -def _declared_modal_timeouts(monkeypatch): - install_fake_modal(monkeypatch) - sys.modules.pop("src.modal.v2_app", None) - modal_app = importlib.import_module("src.modal.v2_app") - return { - name: options.get("timeout") for name, options in modal_app.app.function_calls - } - - -# After the child window closes and the production-result wait ends, the -# coordinator still has to read both simulation artifacts, aggregate, write its -# artifacts, and record the outcome before Modal stops it. -MINIMUM_FINISHING_SECONDS = 600 - - -def test_declared_modal_timeouts_fit_the_shared_wait_window(monkeypatch) -> None: - declared_timeouts = _declared_modal_timeouts(monkeypatch) - window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS - child_worker_timeouts = { - name: timeout - for name, timeout in declared_timeouts.items() - if name.startswith("run_single_simulation") - } - - assert set(child_worker_timeouts) == { - "run_single_simulation_us", - "run_single_simulation_uk", - "run_single_simulation_segment_us", - } - # The window covers every worker's own Modal timeout. - assert all(window_seconds >= timeout for timeout in child_worker_timeouts.values()) - assert ( - declared_timeouts["coordinate_report"] - >= window_seconds - + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS - + MINIMUM_FINISHING_SECONDS - ) - - -@pytest.mark.parametrize("setup_seconds", [0.0, 400.0, 2_000.0]) -def test_child_wait_window_opens_when_the_coordinator_starts( - monkeypatch, - setup_seconds, -) -> None: - declared_timeouts = _declared_modal_timeouts(monkeypatch) - window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS - clock = SteppedClock() - coordinator_started_at = clock.monotonic() - - def slow_output_planning(_report): - # Loading the country model is coordinator setup, not child time. - clock.advance(setup_seconds) - return _output_plan() - - store = FakeStore() - invoker = ScriptedChildInvoker(clock=clock) - - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - _coordinate_with( - _report(), - store, - invoker, - output_plan_resolver=slow_output_planning, - monotonic=clock.monotonic, - sleep=clock.advance, - ) - - probed_at = invoker.calls["baseline"].probed_at - child_waits_began_after = probed_at[0] - coordinator_started_at - window_closed_after = probed_at[-1] - coordinator_started_at - assert child_waits_began_after == setup_seconds - # Setup shortens the children's window rather than the time after it. - assert window_closed_after == pytest.approx(window_seconds) - assert probed_at[-1] - probed_at[0] == pytest.approx(window_seconds - setup_seconds) - assert ( - declared_timeouts["coordinate_report"] - - window_closed_after - - PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS - >= MINIMUM_FINISHING_SECONDS - ) - assert store.parent.error_code == "simulation_wait_timeout" - for role in ("baseline", "reform"): - _assert_cancelled( - store.children[_simulation_id(role)], - after=f"the {window_seconds:g}-second wait window closed", - ) - - -def test_exhausted_wait_window_gives_the_children_no_extra_time() -> None: - window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS - clock = SteppedClock() - - def setup_that_uses_the_whole_window(_report): - clock.advance(window_seconds + 60.0) - return _output_plan() - - store = FakeStore() - invoker = ScriptedChildInvoker(clock=clock) - - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - _coordinate_with( - _report_with_segmentation(segmented=False), - store, - invoker, - output_plan_resolver=setup_that_uses_the_whole_window, - monotonic=clock.monotonic, - sleep=lambda _: pytest.fail("an exhausted window must not wait"), - ) - - assert [len(call.probed_at) for call in invoker.calls.values()] == [1, 1] - assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] - assert store.parent.error_code == "simulation_wait_timeout" - - def _runtime_region_type(country: str, region: str | None) -> str | None: """Return the region type the runtime's own region normalisation implies.""" from policyengine_simulation_executor.simulation_runtime import ( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime_store_retries.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime_store_retries.py new file mode 100644 index 000000000..4ae09c3e5 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime_store_retries.py @@ -0,0 +1,135 @@ +"""Tests for the Stage 12 runtime store's retry of refused database connections.""" + +from __future__ import annotations + +from contextlib import contextmanager + +import pytest +from policyengine_simulation_contract.stage12_execution import SimulationRole +from policyengine_simulation_executor.stage12_runtime import database_retry +from policyengine_simulation_executor.stage12_runtime.dependencies import runtime_store +from policyengine_stage12_persistence import store as persistence_store_module +from sqlalchemy.exc import OperationalError +from stage12_runtime_support import ( + TelemetryRecordingRuntime, + child_record, + connection_limit_error, + parent_record, + planned_simulation, +) + + +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_record(planned_simulation(SimulationRole.BASELINE)) + engine = BusyPoolerEngine( + [parent_record().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_record() + 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 == [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py b/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py index 792302e56..128b7447f 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py @@ -27,7 +27,7 @@ merge_segment_results, should_segment_simulation, ) -from test_stage12_runtime import _planned_simulation +from stage12_runtime_support import planned_simulation def _frames(index: int) -> dict[str, pd.DataFrame]: @@ -93,7 +93,7 @@ def _spm_provenance() -> SPMProvenance: def test_merge_preserves_one_identical_compact_spm_receipt() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) receipt = _spm_provenance() calculated = merge_segment_results(simulation, _results(simulation, receipt)) @@ -106,7 +106,7 @@ def test_merge_preserves_one_identical_compact_spm_receipt() -> None: def test_merge_rejects_missing_or_incompatible_segment_spm_receipts( difference: str, ) -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) receipt = _spm_provenance() results = _results(simulation, receipt) alternative = { @@ -136,7 +136,7 @@ def test_merge_rejects_missing_or_incompatible_segment_spm_receipts( def test_eligibility_matches_stage12_supported_national_shape() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) requirements = simulation.output_plan.requirements assert should_segment_simulation(simulation) is True assert ( @@ -176,7 +176,7 @@ def test_eligibility_matches_stage12_supported_national_shape() -> None: def test_builds_twenty_single_policy_segment_inputs() -> None: - simulation = _planned_simulation(SimulationRole.REFORM) + simulation = planned_simulation(SimulationRole.REFORM) segments = build_segment_inputs(simulation) @@ -187,7 +187,7 @@ def test_builds_twenty_single_policy_segment_inputs() -> None: def test_segment_calculation_scopes_one_group_and_preserves_dtypes() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) segment = build_segment_inputs(simulation)[0] received = [] @@ -218,7 +218,7 @@ def calculator(scoped): def test_merge_rejects_missing_duplicate_and_cross_role_results() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) results = _results(simulation) with pytest.raises(ValueError, match="incomplete"): @@ -236,7 +236,7 @@ def test_merge_rejects_missing_duplicate_and_cross_role_results() -> None: def test_merge_rejects_duplicate_entity_ids_across_groups() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) results = _results(simulation) duplicate_payload, _ = serialize_simulation_frames(_frames(0)) results[1] = results[1].model_copy( @@ -263,7 +263,7 @@ def _with_frames( def test_merge_rejects_a_result_set_that_omits_a_segment() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) results = _results(simulation) results[-1] = results[-1].model_copy(update={"segment_index": 99}) @@ -274,7 +274,7 @@ def test_merge_rejects_a_result_set_that_omits_a_segment() -> None: def test_merge_rejects_a_result_for_another_region_group() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) results = _results(simulation) results[0] = results[0].model_copy(update={"region_codes": results[1].region_codes}) @@ -285,7 +285,7 @@ def test_merge_rejects_a_result_for_another_region_group() -> None: def test_merge_rejects_segments_with_different_entities() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) results = _results(simulation) results[1] = _with_frames(results[1], {"household": _frames(1)["household"]}) @@ -296,7 +296,7 @@ def test_merge_rejects_segments_with_different_entities() -> None: def test_merge_rejects_segments_with_different_column_dtypes() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) results = _results(simulation) frames = _frames(1) frames["household"] = frames["household"].astype( @@ -324,7 +324,7 @@ def test_merge_rejects_an_entity_without_its_identifier(monkeypatch) -> None: "deserialize_simulation_frames", lambda _payload: without_identifier, ) - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) with pytest.raises( ValueError, match=r"^Stage 12 segment entity 'household' has no identifier$" diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py index 15284a869..a21794b09 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -21,7 +21,7 @@ from policyengine_simulation_executor.stage12_runtime.partition import ( US_REGION_GROUPS, ) -from test_stage12_runtime import _planned_simulation +from stage12_runtime_support import planned_simulation if TYPE_CHECKING: from policyengine_observability import ObservabilityRuntime @@ -107,7 +107,7 @@ def _results(simulation) -> list[dict]: def _runner(*, errors: dict[int, list[Exception]] | None = None): - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) function = FakeFunction(_results(simulation), errors or {}) class FunctionLookup: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py b/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py index cda82ee58..9b2e6ec5b 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py @@ -17,7 +17,7 @@ Stage12SegmentInput, ) from pydantic import ValidationError -from test_stage12_runtime import _planned_simulation +from stage12_runtime_support import planned_simulation PROJECT = Path(__file__).resolve().parents[1] PACKAGE = PROJECT / "src" / "policyengine_simulation_executor" @@ -74,7 +74,7 @@ def test_new_model_regions_are_added_to_the_final_group_in_canonical_order() -> def test_segment_input_rejects_duplicate_or_non_state_regions() -> None: - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) with pytest.raises(ValidationError, match="must be unique"): Stage12SegmentInput( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py b/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py index a326488c6..904275334 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_simulation_worker.py @@ -18,14 +18,12 @@ from policyengine_simulation_executor.stage12_runtime.failure_causes import ( StageProgress, ) -from test_stage12_runtime import ( +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames +from stage12_runtime_support import ( NOW, - FakeArtifacts, - FakeStore, - _child, - _context, - _frames, - _planned_simulation, + child_record, + invocation_context, + planned_simulation, ) @@ -39,13 +37,13 @@ def span(self, name, *_args, **_kwargs): def _run(store, artifacts, *, calculator=None, runtime=None): - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) seams = {"store": store, "artifacts": artifacts, "runtime": runtime} if calculator is not None: seams["calculator"] = calculator return run_single_simulation( simulation.model_dump(mode="json"), - _context().model_dump(mode="json"), + invocation_context().model_dump(mode="json"), required_country="us", dependencies=SimulationWorkerDependencies(**seams), ) @@ -54,9 +52,9 @@ def _run(store, artifacts, *, calculator=None, runtime=None): def test_succeeded_child_returns_its_stored_descriptor_without_recalculating(): store = FakeStore() artifacts = FakeArtifacts() - simulation = _planned_simulation(SimulationRole.BASELINE) + simulation = planned_simulation(SimulationRole.BASELINE) descriptor = artifacts.add_simulation(simulation) - succeeded = _child(simulation).model_copy( + succeeded = child_record(simulation).model_copy( update={ "status": ComparisonRunLifecycleStatus.SUCCEEDED, "output_uri": descriptor.artifact.uri, @@ -83,14 +81,14 @@ def test_succeeded_child_returns_its_stored_descriptor_without_recalculating(): def test_default_calculator_receives_the_runtime_and_times_itself(monkeypatch): store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) runtime = SpanRecordingRuntime() calls = [] def calculate(planned, *, runtime=None, progress=None): calls.append((planned.role, runtime, progress)) - return _frames() + return simulation_frames() monkeypatch.setattr(simulation_module, "calculate_simulation_frames", calculate) @@ -110,14 +108,14 @@ def calculate(planned, *, runtime=None, progress=None): def test_injected_calculator_is_timed_by_the_worker(): store = FakeStore() - simulation = _planned_simulation(SimulationRole.BASELINE) - store.children[simulation.simulation_execution_id] = _child(simulation) + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) runtime = SpanRecordingRuntime() calls = [] def calculate(planned): calls.append(planned.role) - return _frames() + return simulation_frames() _run(store, FakeArtifacts(), calculator=calculate, runtime=runtime) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_single_worker.py b/projects/policyengine-simulation-executor/tests/test_stage12_single_worker.py new file mode 100644 index 000000000..6bab645c1 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_single_worker.py @@ -0,0 +1,231 @@ +"""Tests for a Stage 12 single-simulation worker's successful and refused runs.""" + +from __future__ import annotations + +import logging + +import pytest +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + SimulationCalculation, + SimulationWorkerDependencies, + run_single_simulation, +) +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames +from stage12_runtime_support import ( + child_record, + invocation_context, + planned_simulation, + spm_calculation_provenance, +) + + +def test_single_worker_accepts_one_policy_and_persists_one_artifact() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + artifacts = FakeArtifacts() + + result = run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=artifacts, + calculator=lambda _: simulation_frames(), + ), + ) + + assert result["role"] == "baseline" + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.SUCCEEDED + assert child.output_uri == result["artifact"]["uri"] + assert artifacts.input_writes == [ + (invocation_context().artifact_prefix, SimulationRole.BASELINE) + ] + + +def test_single_worker_routes_eligible_us_calculation_to_segment_runner() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + received = [] + + def segmented_calculator(value, cancellation_requested): + received.append(value) + assert cancellation_requested() is False + return SimulationCalculation(frames=simulation_frames()) + + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + ), + ) + + assert received == [simulation] + + +def test_single_worker_honors_explicit_segmentation_opt_out() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE).model_copy( + update={"options": {"segmented": False}} + ) + store.children[simulation.simulation_execution_id] = child_record(simulation) + segmented_calls = [] + + def unexpected_segmented(value, cancellation_requested): + segmented_calls.append(value) + return SimulationCalculation(frames=simulation_frames()) + + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: simulation_frames(), + segmented_calculator=unexpected_segmented, + ), + ) + + assert segmented_calls == [] + + +def test_single_worker_refuses_a_simulation_that_is_no_longer_active() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record( + simulation + ).model_copy(update={"status": ComparisonRunLifecycleStatus.FAILED}) + artifacts = FakeArtifacts() + + with pytest.raises( + RuntimeError, match=r"^Stage 12 simulation is no longer active$" + ): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=artifacts, + calculator=lambda _: pytest.fail("a finished simulation ran again"), + ), + ) + + assert artifacts.input_writes == [] + + +def test_single_worker_stops_when_another_writer_claimed_the_simulation() -> None: + class ClaimLosingStore(FakeStore): + def replace_simulation_if_status(self, record, *, expected_status): + return self.children[record.simulation_execution_id], False + + store = ClaimLosingStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + artifacts = FakeArtifacts() + + with pytest.raises(RuntimeError, match=r"^Stage 12 simulation was cancelled$"): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=artifacts, + calculator=lambda _: pytest.fail("an unclaimed simulation ran"), + ), + ) + + assert artifacts.input_writes == [] + + +def test_single_worker_treats_a_failed_cancellation_check_as_not_cancelled( + caplog, +) -> None: + class StatusReadFailingStore(FakeStore): + def __init__(self): + super().__init__() + self.fail_reads = False + + def get_simulation(self, simulation_execution_id): + if self.fail_reads: + msg = "status read contains sensitive details" + raise ConnectionError(msg) + return super().get_simulation(simulation_execution_id) + + store = StatusReadFailingStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + checks = [] + + def segmented_calculator(_value, cancellation_requested): + store.fail_reads = True + checks.append(cancellation_requested()) + store.fail_reads = False + return SimulationCalculation(frames=simulation_frames()) + + with caplog.at_level( + logging.WARNING, + logger="policyengine_simulation_executor.stage12_runtime.simulation", + ): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + segmented_calculator=segmented_calculator, + ), + ) + + assert checks == [False] + assert ( + f"Stage 12 cancellation check failed for " + f"{simulation.simulation_execution_id} (ConnectionError)" + ) in caplog.text + assert "sensitive" not in caplog.text + assert ( + store.children[simulation.simulation_execution_id].status + is ComparisonRunLifecycleStatus.SUCCEEDED + ) + + +def test_single_worker_retains_detached_calculation_provenance() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + artifacts = FakeArtifacts() + provenance = spm_calculation_provenance() + + result = run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=artifacts, + calculator=lambda _: SimulationCalculation( + frames=simulation_frames(), + calculation_provenance=provenance, + ), + ), + ) + + assert result["calculation_provenance"] == provenance.model_dump( + mode="json", by_alias=True + ) + assert "spm_config" not in result["calculation_provenance"] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_single_worker_failures.py b/projects/policyengine-simulation-executor/tests/test_stage12_single_worker_failures.py new file mode 100644 index 000000000..dca921684 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_single_worker_failures.py @@ -0,0 +1,335 @@ +"""Tests for how a Stage 12 single-simulation worker records and reports failures.""" + +from __future__ import annotations + +from datetime import UTC, datetime + +import pytest +from policyengine_simulation_contract.spm import SPMInputError +from policyengine_simulation_contract.stage12_execution import ( + ComparisonRunLifecycleStatus, + SimulationRole, +) +from policyengine_simulation_executor.stage12_runtime import ( + SimulationCalculation, + SimulationWorkerDependencies, + run_single_simulation, +) +from policyengine_simulation_executor.stage12_runtime import ( + simulation as simulation_module, +) +from sqlalchemy.exc import OperationalError +from stage12_fake_stores import FakeArtifacts, FakeStore, simulation_frames +from stage12_runtime_support import ( + BASELINE_ID, + EVALUATION_ID, + RecordingRuntime, + TelemetryRecordingRuntime, + child_record, + connection_limit_error, + invocation_context, + planned_simulation, + summary_correlation_id, + summary_failure_cause, +) + + +def test_single_worker_exposes_only_a_bounded_failure() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + + def fail(_simulation): + msg = "country package calculation failed" + raise RuntimeError(msg) + + with pytest.raises( + RuntimeError, match="Stage 12 simulation execution failed" + ) as error: + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + ), + ) + + assert "country package calculation failed" not in str(error.value) + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.FAILED + assert child.error_code == "simulation_execution_failed" + assert summary_failure_cause(child.error_summary) == ( + "RuntimeError in stage12_calculation (baseline): " + "country package calculation failed" + ) + + +def test_single_worker_persists_and_records_a_redacted_cause() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + runtime = RecordingRuntime() + + def fail(_simulation): + msg = "could not reach postgresql://runtime:database-secret@db.example/postgres" + raise RuntimeError(msg) + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + runtime=runtime, + ), + ) + + child = store.children[simulation.simulation_execution_id] + assert summary_failure_cause(child.error_summary) == ( + "RuntimeError in stage12_calculation (baseline): " + "could not reach postgresql://[REDACTED]@db.example/postgres" + ) + assert "database-secret" not in child.error_summary + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == child.error_summary + + +def test_single_worker_names_the_failing_artifact_write_stage() -> None: + class FailingSimulationArtifacts(FakeArtifacts): + def write_simulation(self, **_kwargs): + msg = "bucket write rejected" + raise OSError(msg) + + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FailingSimulationArtifacts(), + calculator=lambda _: simulation_frames(), + ), + ) + + assert summary_failure_cause( + store.children[simulation.simulation_execution_id].error_summary + ) == ( + "OSError in stage12_simulation_artifact_write (baseline): bucket write rejected" + ) + + +def test_single_worker_preserves_typed_failure_values() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + + def fail(_simulation): + msg = "SPM_YEAR_UNAVAILABLE" + raise SPMInputError(msg, "No SPM values for 2026") + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + ), + ) + + child = store.children[simulation.simulation_execution_id] + assert child.error_code == "SPM_YEAR_UNAVAILABLE" + assert child.error_summary == "No SPM values for 2026" + + +def test_single_worker_rejects_bare_year_policy_period_before_calculation() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE).model_copy( + update={ + "policy": {"gov.irs.credits.ctc.amount.base[0].amount": {"2026": 3_000}} + } + ) + store.children[simulation.simulation_execution_id] = child_record(simulation) + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("calculation must not start"), + ), + ) + + child = store.children[simulation.simulation_execution_id] + assert child.error_code == "invalid_policy_period" + assert child.error_summary == ( + "Policy parameter 'gov.irs.credits.ctc.amount.base[0].amount' has invalid " + "effective period '2026'; expected YYYY-MM-DD or " + "YYYY-MM-DD.YYYY-MM-DD." + ) + + +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_record(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"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + artifacts=FakeArtifacts(), + calculator=lambda _: simulation_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_if_status(self, record, *, expected_status): + if record.status is ComparisonRunLifecycleStatus.FAILED: + ordered_events.append(("failure_write", record.error_summary)) + raise connection_limit_error() + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + store = UnwritableFailureStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(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"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + 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", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "correlation_id": summary_correlation_id(failure_write[1]), + } + assert failure_write[0] == "failure_write" + assert summary_failure_cause(failure_write[1]) == ( + "ValueError in stage12_calculation (baseline): " + "country column missing from household data" + ) + 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_late_simulation_success_does_not_overwrite_parent_cancellation() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + + def segmented_calculator(value, cancellation_requested): + current = store.children[value.simulation_execution_id] + cancelled_at = datetime.now(UTC) + store.replace_simulation( + current.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_cancelled", + "error_summary": "Cancelled after the reform simulation failed.", + "updated_at": cancelled_at, + "completed_at": cancelled_at, + } + ) + ) + assert cancellation_requested() is True + return SimulationCalculation(frames=simulation_frames()) + + with pytest.raises(RuntimeError, match="Stage 12 simulation was cancelled"): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + ), + ) + + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.FAILED + assert child.error_code == "simulation_cancelled" + + +def test_single_worker_rejects_frames_that_do_not_satisfy_the_output_plan() -> None: + store = FakeStore() + simulation = planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = child_record(simulation) + frames = simulation_frames() + frames["person"] = frames["person"].drop(columns=["age"]) + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + invocation_context().model_dump(mode="json"), + required_country="us", + dependencies=SimulationWorkerDependencies( + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: frames, + ), + ) + + assert store.children[simulation.simulation_execution_id].status is ( + ComparisonRunLifecycleStatus.FAILED + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_wait_window.py b/projects/policyengine-simulation-executor/tests/test_stage12_wait_window.py new file mode 100644 index 000000000..4e97b6bfc --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_wait_window.py @@ -0,0 +1,149 @@ +"""Tests for the shared wait window of a Stage 12 report's child simulations.""" + +from __future__ import annotations + +import importlib +import sys + +import pytest +from fixtures.fake_modal import install_fake_modal +from policyengine_simulation_executor.stage12_runtime import ( + coordination as coordination_module, +) +from policyengine_simulation_executor.stage12_runtime.comparison import ( + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, +) +from stage12_child_call_support import ( + ScriptedChildInvoker, + assert_cancelled, + coordinate_with, + report_with_segmentation, + role_simulation_id, +) +from stage12_fake_stores import FakeStore +from stage12_runtime_support import report_execution_input, shared_output_plan + + +class SteppedClock: + def __init__(self, start_seconds=10_000.0): + self.current_seconds = start_seconds + + def monotonic(self): + return self.current_seconds + + def advance(self, elapsed_seconds): + self.current_seconds += elapsed_seconds + + +def _declared_modal_timeouts(monkeypatch): + install_fake_modal(monkeypatch) + sys.modules.pop("src.modal.v2_app", None) + modal_app = importlib.import_module("src.modal.v2_app") + return { + name: options.get("timeout") for name, options in modal_app.app.function_calls + } + + +# After the child window closes and the production-result wait ends, the +# coordinator still has to read both simulation artifacts, aggregate, write its +# artifacts, and record the outcome before Modal stops it. +MINIMUM_FINISHING_SECONDS = 600 + + +def test_declared_modal_timeouts_fit_the_shared_wait_window(monkeypatch) -> None: + declared_timeouts = _declared_modal_timeouts(monkeypatch) + window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + child_worker_timeouts = { + name: timeout + for name, timeout in declared_timeouts.items() + if name.startswith("run_single_simulation") + } + + assert set(child_worker_timeouts) == { + "run_single_simulation_us", + "run_single_simulation_uk", + "run_single_simulation_segment_us", + } + # The window covers every worker's own Modal timeout. + assert all(window_seconds >= timeout for timeout in child_worker_timeouts.values()) + assert ( + declared_timeouts["coordinate_report"] + >= window_seconds + + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + + MINIMUM_FINISHING_SECONDS + ) + + +@pytest.mark.parametrize("setup_seconds", [0.0, 400.0, 2_000.0]) +def test_child_wait_window_opens_when_the_coordinator_starts( + monkeypatch, + setup_seconds, +) -> None: + declared_timeouts = _declared_modal_timeouts(monkeypatch) + window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + clock = SteppedClock() + coordinator_started_at = clock.monotonic() + + def slow_output_planning(_report): + # Loading the country model is coordinator setup, not child time. + clock.advance(setup_seconds) + return shared_output_plan() + + store = FakeStore() + invoker = ScriptedChildInvoker(clock=clock) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_with( + report_execution_input(), + store, + invoker, + output_plan_resolver=slow_output_planning, + monotonic=clock.monotonic, + sleep=clock.advance, + ) + + probed_at = invoker.calls["baseline"].probed_at + child_waits_began_after = probed_at[0] - coordinator_started_at + window_closed_after = probed_at[-1] - coordinator_started_at + assert child_waits_began_after == setup_seconds + # Setup shortens the children's window rather than the time after it. + assert window_closed_after == pytest.approx(window_seconds) + assert probed_at[-1] - probed_at[0] == pytest.approx(window_seconds - setup_seconds) + assert ( + declared_timeouts["coordinate_report"] + - window_closed_after + - PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + >= MINIMUM_FINISHING_SECONDS + ) + assert store.parent.error_code == "simulation_wait_timeout" + for role in ("baseline", "reform"): + assert_cancelled( + store.children[role_simulation_id(role)], + after=f"the {window_seconds:g}-second wait window closed", + ) + + +def test_exhausted_wait_window_gives_the_children_no_extra_time() -> None: + window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + clock = SteppedClock() + + def setup_that_uses_the_whole_window(_report): + clock.advance(window_seconds + 60.0) + return shared_output_plan() + + store = FakeStore() + invoker = ScriptedChildInvoker(clock=clock) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_with( + report_with_segmentation(segmented=False), + store, + invoker, + output_plan_resolver=setup_that_uses_the_whole_window, + monotonic=clock.monotonic, + sleep=lambda _: pytest.fail("an exhausted window must not wait"), + ) + + assert [len(call.probed_at) for call in invoker.calls.values()] == [1, 1] + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + assert store.parent.error_code == "simulation_wait_timeout" diff --git a/rules/baselines/python-file-length.txt b/rules/baselines/python-file-length.txt index a93a18dba..dd28cdefd 100644 --- a/rules/baselines/python-file-length.txt +++ b/rules/baselines/python-file-length.txt @@ -11,4 +11,3 @@ projects/policyengine-simulation-executor/src/policyengine_simulation_executor/s projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py projects/policyengine-simulation-executor/tests/test_precompute.py projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py -projects/policyengine-simulation-executor/tests/test_stage12_runtime.py diff --git a/rules/registry.toml b/rules/registry.toml index 8add82684..858f1ef21 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1190,7 +1190,7 @@ job = "quality-guards" invocation = "run_quality_guards.py" [rule.baseline] file = "rules/baselines/python-file-length.txt" -count = 8 +count = 7 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.python-file-length/bad_files.toml" good_fixture = "rules/fixtures/sim.quality.python-file-length/good_files.toml"