From e533dc764cf524436cc8e61590a21f323e34ab96 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:37:05 +0400 Subject: [PATCH 1/3] Let report simulations return their artifact without database rows When the dispatch context's result_delivery is return_to_caller, run_single_simulation reads and writes no Stage 12 row and builds no database store. It checks the policy, calculates the simulation (a segmented US run is never told to stop through the database), stores the input and Parquet artifact as before, and returns the artifact descriptor. A failure raises ReportExecutionError with the same redacted error code and summary a Stage 12 row would record. The default stage12_database delivery runs the same steps as before; the calculation and failure-detail steps move into helpers both modes share. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/failures.py | 26 +- .../stage12_runtime/simulation.py | 236 +++++++++++++----- .../tests/test_stage12_runtime.py | 170 +++++++++++++ 3 files changed, 366 insertions(+), 66 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py index a42f9c05b..ad97388b5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py @@ -4,14 +4,21 @@ from collections.abc import Mapping from datetime import date -from typing import Any +from typing import TYPE_CHECKING, Any +from policyengine_simulation_contract.report_execution import ( + ReportExecutionError, + ReportExecutionFailure, +) from policyengine_simulation_contract.spm import spm_error_detail from policyengine_simulation_observability.errors import make_correlation_id from pydantic import BaseModel, ConfigDict from .failure_causes import StageProgress, failure_summary +if TYPE_CHECKING: + from uuid import UUID + ERROR_CODE_MAX_LENGTH = 64 ERROR_SUMMARY_MAX_LENGTH = 512 @@ -83,6 +90,23 @@ def failure_detail_from_record( ) +def report_execution_error( + evaluation_id: UUID, + detail: Stage12FailureDetail, +) -> ReportExecutionError: + """Return the typed error that carries ``detail`` to a returned report's caller. + + It holds the error code and summary a Stage 12 row would record. + """ + return ReportExecutionError( + ReportExecutionFailure( + evaluation_id=evaluation_id, + error_code=detail.error_code, + error_summary=detail.error_summary or "Simulation failed", + ) + ) + + def redacted_failure_detail( error: BaseException, *, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index e18e95638..0cd1bbf9d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -19,6 +19,7 @@ ComparisonRunLifecycleStatus, ComparisonSimulationRecord, PlannedSimulationExecutionInput, + ReportResultDelivery, RowIdentity, SimulationArtifactDescriptor, SimulationExecutionInput, @@ -48,6 +49,7 @@ Stage12CancellationError, Stage12ExecutionError, Stage12FailureDetail, + report_execution_error, validate_policy_periods, ) from .output_planning import apply_output_plan, validate_output_frames @@ -394,6 +396,144 @@ def _mark_running_simulation_failed( ) +type SimulationCalculator = Callable[ + [PlannedSimulationExecutionInput], + Mapping[str, pd.DataFrame] | SimulationCalculation, +] +type SegmentedSimulationCalculator = Callable[ + [PlannedSimulationExecutionInput, Callable[[], bool]], SimulationCalculation +] + + +def _calculate_and_store( + simulation: PlannedSimulationExecutionInput, + context: Stage12InvocationContext, + *, + artifact_storage: Stage12ArtifactStore, + calculator: SimulationCalculator, + segmented_calculator: SegmentedSimulationCalculator | None, + cancellation_requested: Callable[[], bool], + runtime: ObservabilityRuntime | None, + progress: StageProgress, +) -> SimulationArtifactDescriptor: + """Check the policy, calculate the simulation and store its input and artifact. + + An eligible US simulation goes to ``segmented_calculator``, which checks + ``cancellation_requested`` while it waits for its segments. + """ + stage_span = partial( + tracked_stage, + stage_plan=STAGE12_SIMULATION_STAGES, + progress=progress, + runtime=runtime, + ) + validate_policy_periods(simulation.policy) + with stage_span(Stage.STAGE12_INPUT_WRITE): + artifact_storage.write_input( + prefix=context.artifact_prefix, + simulation=simulation, + ) + if segmented_calculator is not None and should_segment_simulation(simulation): + calculated = segmented_calculator(simulation, cancellation_requested) + elif calculator is calculate_simulation_frames: + calculated = calculate_simulation_frames( + simulation, + runtime=runtime, + progress=progress, + ) + else: + with stage_span(Stage.STAGE12_CALCULATION): + calculated = calculator(simulation) + if isinstance(calculated, SimulationCalculation): + frames = calculated.frames + calculation_provenance = calculated.calculation_provenance + uk_local_authority_metadata = calculated.uk_local_authority_metadata + else: + frames = calculated + calculation_provenance = None + uk_local_authority_metadata = None + validate_output_frames(frames, simulation.output_plan) + with stage_span(Stage.STAGE12_SIMULATION_ARTIFACT_WRITE): + return artifact_storage.write_simulation( + prefix=context.artifact_prefix, + simulation=simulation, + frames=frames, + calculation_provenance=calculation_provenance, + uk_local_authority_metadata=uk_local_authority_metadata, + ) + + +def _simulation_failure_detail( + error: Exception, + *, + simulation: PlannedSimulationExecutionInput, + runtime: ObservabilityRuntime | None, + progress: StageProgress, +) -> Stage12FailureDetail: + """Report a calculation failure once and return its redacted detail.""" + if isinstance(error, Stage12ExecutionError): + # A segmented calculation reports its segment failure where it + # occurs, before raising this error. + return error.detail + return report_failure( + error, + runtime=runtime, + failure_scope="simulation_execution", + default_code="simulation_execution_failed", + attributes={ + "simulation_role": simulation.role.value, + "evaluation_id": str(simulation.evaluation_id), + "simulation_execution_id": str(simulation.simulation_execution_id), + }, + progress=progress, + ) + + +def _never_cancelled() -> bool: + """Report no cancellation: nothing records one for a returned simulation.""" + return False + + +def _run_simulation_for_caller( + simulation: PlannedSimulationExecutionInput, + context: Stage12InvocationContext, + *, + artifact_storage: Stage12ArtifactStore, + calculator: SimulationCalculator, + segmented_calculator: SegmentedSimulationCalculator | None, + runtime: ObservabilityRuntime | None, +) -> dict[str, Any]: + """Calculate one simulation for a returned report, without database access. + + The simulation is always calculated. No record can carry a cancellation, + so the coordinator stops this call through Modal instead. A failure + raises ``ReportExecutionError`` with the redacted detail a Stage 12 row + would record. + """ + progress = StageProgress(role=simulation.role.value) + try: + descriptor = _calculate_and_store( + simulation, + context, + artifact_storage=artifact_storage, + calculator=calculator, + segmented_calculator=segmented_calculator, + cancellation_requested=_never_cancelled, + runtime=runtime, + progress=progress, + ) + except Exception as error: # noqa: BLE001 # every calculation failure reaches the caller as a typed error with its redacted detail + detail = _simulation_failure_detail( + error, + simulation=simulation, + runtime=runtime, + progress=progress, + ) + record_failure_cause(runtime, error, detail.error_summary) + raise report_execution_error(simulation.evaluation_id, detail) from None + return descriptor.model_dump(mode="json") + + def run_single_simulation( payload: object, context_payload: object, @@ -401,14 +541,8 @@ def run_single_simulation( required_country: CountryId, store: ComparisonStore | None = None, artifacts: Stage12ArtifactStore | None = None, - calculator: Callable[ - [PlannedSimulationExecutionInput], - Mapping[str, pd.DataFrame] | SimulationCalculation, - ] = calculate_simulation_frames, - segmented_calculator: Callable[ - [PlannedSimulationExecutionInput, Callable[[], bool]], SimulationCalculation - ] - | None = None, + calculator: SimulationCalculator = calculate_simulation_frames, + segmented_calculator: SegmentedSimulationCalculator | None = None, runtime: ObservabilityRuntime | None = None, ) -> dict[str, Any]: """Calculate one planned simulation, store its artifact, and return its descriptor. @@ -416,10 +550,23 @@ def run_single_simulation( A child already marked succeeded returns its stored descriptor without recalculating. A failure after the run is marked running marks it FAILED and raises a ``RuntimeError`` that carries no calculation detail. + + When the context's ``result_delivery`` returns the report to its caller, + no Stage 12 row is read or written, ``store`` is not used, and a failure + raises ``ReportExecutionError`` with its redacted detail instead. """ simulation = PlannedSimulationExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) _require_context(simulation, context, required_country=required_country) + if context.result_delivery is ReportResultDelivery.RETURN_TO_CALLER: + return _run_simulation_for_caller( + simulation, + context, + artifact_storage=artifacts or artifact_store(), + calculator=calculator, + segmented_calculator=segmented_calculator, + runtime=runtime, + ) persistence = store or runtime_store(runtime) artifact_storage = artifacts or artifact_store() child = persistence.get_simulation(simulation.simulation_execution_id) @@ -462,47 +609,17 @@ def cancellation_requested() -> bool: return latest.status is not ComparisonRunLifecycleStatus.RUNNING progress = StageProgress(role=simulation.role.value) - stage_span = partial( - tracked_stage, - stage_plan=STAGE12_SIMULATION_STAGES, - progress=progress, - runtime=runtime, - ) try: - validate_policy_periods(simulation.policy) - with stage_span(Stage.STAGE12_INPUT_WRITE): - artifact_storage.write_input( - prefix=context.artifact_prefix, - simulation=simulation, - ) - if segmented_calculator is not None and should_segment_simulation(simulation): - calculated = segmented_calculator(simulation, cancellation_requested) - elif calculator is calculate_simulation_frames: - calculated = calculate_simulation_frames( - simulation, - runtime=runtime, - progress=progress, - ) - else: - with stage_span(Stage.STAGE12_CALCULATION): - calculated = calculator(simulation) - if isinstance(calculated, SimulationCalculation): - frames = calculated.frames - calculation_provenance = calculated.calculation_provenance - uk_local_authority_metadata = calculated.uk_local_authority_metadata - else: - frames = calculated - calculation_provenance = None - uk_local_authority_metadata = None - validate_output_frames(frames, simulation.output_plan) - with stage_span(Stage.STAGE12_SIMULATION_ARTIFACT_WRITE): - descriptor = artifact_storage.write_simulation( - prefix=context.artifact_prefix, - simulation=simulation, - frames=frames, - calculation_provenance=calculation_provenance, - uk_local_authority_metadata=uk_local_authority_metadata, - ) + descriptor = _calculate_and_store( + simulation, + context, + artifact_storage=artifact_storage, + calculator=calculator, + segmented_calculator=segmented_calculator, + cancellation_requested=cancellation_requested, + runtime=runtime, + progress=progress, + ) completed = datetime.now(UTC) _, replaced = persistence.replace_simulation_if_status( running.model_copy( @@ -528,23 +645,12 @@ def cancellation_requested() -> bool: # Persist a redacted cause for any country-package calculation failure # before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 # every calculation failure is persisted before a stable exception returns to Modal - if isinstance(error, Stage12ExecutionError): - # A segmented calculation reports its segment failure where it - # occurs, before raising this error. - detail = error.detail - else: - detail = report_failure( - error, - runtime=runtime, - failure_scope="simulation_execution", - default_code="simulation_execution_failed", - attributes={ - "simulation_role": simulation.role.value, - "evaluation_id": str(simulation.evaluation_id), - "simulation_execution_id": str(simulation.simulation_execution_id), - }, - progress=progress, - ) + detail = _simulation_failure_detail( + error, + simulation=simulation, + runtime=runtime, + progress=progress, + ) record_failure_cause(runtime, error, detail.error_summary) failed_at = datetime.now(UTC) persist_failure_record( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 405d12663..2b71bd7a7 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -32,6 +32,7 @@ configure, ) from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from policyengine_simulation_contract.report_execution import ReportExecutionError from policyengine_simulation_contract.spm import ( SPMComparisonProvenance, SPMInputError, @@ -60,6 +61,7 @@ ReportAggregate, ReportExecutionInput, ReportOutputRequirements, + ReportResultDelivery, ResultComparisonStatus, RowIdentity, SimulationArtifactDescriptor, @@ -80,6 +82,8 @@ ) from policyengine_simulation_executor.stage12_runtime import ( SimulationCalculation, + Stage12ExecutionError, + Stage12FailureDetail, Stage12SegmentRunner, _build_spm_result, build_aggregate_report, @@ -3919,3 +3923,169 @@ def setup_that_uses_the_whole_window(_report): 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" + + +# Return-to-caller delivery: the coordinator and its simulations write no +# database rows, return typed values and raise a typed error on failure. + + +def _returning_context() -> Stage12InvocationContext: + return Stage12InvocationContext.model_validate( + { + **_context().model_dump(mode="json"), + "result_delivery": ReportResultDelivery.RETURN_TO_CALLER.value, + } + ) + + +class NoDatabaseStore: + """A store that fails the test on any use: a returned report writes no rows.""" + + def __getattr__(self, name): + pytest.fail(f"a returned report used the Stage 12 store ({name})") + + +@pytest.fixture +def forbid_database(monkeypatch): + """Fail the test if a returned report builds a Stage 12 database store.""" + + def refuse_store(*_args, **_kwargs): + pytest.fail("a returned report built a Stage 12 database store") + + monkeypatch.setattr(simulation_module, "runtime_store", refuse_store) + monkeypatch.setattr(coordination_module, "runtime_store", refuse_store) + + +def _run_returned_simulation(simulation, **options): + return run_single_simulation( + simulation.model_dump(mode="json"), + _returning_context().model_dump(mode="json"), + required_country="us", + **{"artifacts": FakeArtifacts(), **options}, + ) + + +def test_returned_simulation_stores_its_artifact_without_the_database( + forbid_database, +) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + artifacts = FakeArtifacts() + + result = _run_returned_simulation( + simulation, + artifacts=artifacts, + calculator=lambda _: _frames(), + ) + + descriptor = SimulationArtifactDescriptor.model_validate(result) + assert descriptor.role is SimulationRole.BASELINE + assert descriptor.simulation_execution_id == BASELINE_ID + stored = artifacts.read(descriptor.artifact.uri) + assert sha256(stored).hexdigest() == descriptor.artifact.content_sha256 + assert artifacts.input_writes == [ + (_context().artifact_prefix, SimulationRole.BASELINE) + ] + + +def test_returned_simulation_never_uses_a_stage12_store_it_was_given() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + + result = _run_returned_simulation( + simulation, + store=NoDatabaseStore(), + calculator=lambda _: _frames(), + ) + + assert result["role"] == "baseline" + + +def test_returned_simulation_failure_raises_its_redacted_detail( + forbid_database, +) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + runtime = RecordingRuntime() + + def fail(_simulation): + msg = "could not reach postgresql://runtime:database-secret@db.example/postgres" + raise RuntimeError(msg) + + with pytest.raises(ReportExecutionError) as raised: + _run_returned_simulation(simulation, calculator=fail, runtime=runtime) + + failure = raised.value.failure + assert failure.evaluation_id == EVALUATION_ID + assert failure.error_code == "simulation_execution_failed" + assert _failure_cause(failure.error_summary) == ( + "RuntimeError in stage12_calculation (baseline): " + "could not reach postgresql://[REDACTED]@db.example/postgres" + ) + assert str(raised.value) == "Report execution failed" + assert raised.value.__cause__ is None + assert raised.value.__suppress_context__ is True + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == failure.error_summary + + +def test_returned_simulation_keeps_a_typed_input_failure(forbid_database) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={ + "policy": {"gov.irs.credits.ctc.amount.base[0].amount": {"2026": 3_000}} + } + ) + + with pytest.raises(ReportExecutionError) as raised: + _run_returned_simulation( + simulation, + calculator=lambda _: pytest.fail("calculation must not start"), + ) + + assert raised.value.failure.error_code == "invalid_policy_period" + assert raised.value.failure.error_summary.startswith( + "Policy parameter 'gov.irs.credits.ctc.amount.base[0].amount' has invalid " + "effective period '2026'" + ) + + +def test_returned_segmented_simulation_is_never_told_to_stop_by_the_database( + forbid_database, +) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + received = [] + + def segmented_calculator(value, cancellation_requested): + received.append(value) + assert cancellation_requested() is False + return SimulationCalculation(frames=_frames()) + + _run_returned_simulation( + simulation, + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + ) + + assert received == [simulation] + + +def test_returned_segmented_simulation_keeps_the_segment_failure( + forbid_database, +) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + segment_failure = Stage12FailureDetail( + error_code="segment_execution_failed", + error_summary="Segment 3: MemoryError in stage12_calculation", + ) + + def segmented_calculator(_value, _cancellation_requested): + raise Stage12ExecutionError(segment_failure) + + with pytest.raises(ReportExecutionError) as raised: + _run_returned_simulation( + simulation, + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + ) + + assert raised.value.failure.error_code == "segment_execution_failed" + assert raised.value.failure.error_summary == ( + "Segment 3: MemoryError in stage12_calculation" + ) From cd93bd5c946b4223ef6f2df0d0c22a7c43049a41 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:49:10 +0400 Subject: [PATCH 2/3] Add the report coordinator's return mode When the dispatch context's result_delivery is return_to_caller, the report coordinator writes no Stage 12 rows and builds no database store. It takes no parent record, does not claim the report (the caller deduplicates), plans the output, starts the baseline and reform simulations with the same delivery, waits for both within the shared window, reads and checks their artifacts as the Stage 12 path does, and aggregates them. It writes no aggregate artifact. It returns a ReportExecutionResult whose aggregate is validated as SingleYearMacroOutput and must describe the report. A failure is reported once, recorded as the operation's cause and raised as ReportExecutionError with the redacted detail a Stage 12 row would record; a simulation's own ReportExecutionError takes precedence over its call error. A child that must stop is cancelled through Modal, since no record can tell a segmented child to stop. The Modal coordinate_report function's parent argument becomes optional, and a returned report's coordinator operation carries a result_delivery attribute. The default stage12_database path is unchanged; its planning, child wait and artifact checks move into helpers both modes share. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/v2_app.py | 25 +- .../stage12_runtime/coordination.py | 533 +++++++++++++++--- .../stage12_runtime/failures.py | 10 + .../tests/test_stage12_modal_app.py | 62 ++ .../tests/test_stage12_runtime.py | 321 ++++++++++- 5 files changed, 861 insertions(+), 90 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 7485cec2a..5ad4a2473 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -12,6 +12,7 @@ from pathlib import Path from typing import TYPE_CHECKING +from policyengine_simulation_contract.stage12_execution import ReportResultDelivery from policyengine_simulation_contract.stage12_manifest import v2_application_name from policyengine_simulation_executor import stage12_runtime from policyengine_simulation_executor.hf_access_validation import ( @@ -423,7 +424,7 @@ def run_single_simulation_uk( def coordinate_report( payload: dict, context: dict, - parent: dict, + parent: dict | None = None, *, observability_context: dict | None = None, ) -> dict: @@ -432,6 +433,10 @@ def coordinate_report( A ``production_function_call_id`` in ``context`` marks a shadow run whose result is compared with that production job's. A report another invocation already claimed returns a ``deduplicated`` marker instead. + + A ``result_delivery`` of ``return_to_caller`` in ``context`` takes no + ``parent`` and writes no database rows: the report returns a + ``ReportExecutionResult`` payload or raises ``ReportExecutionError``. """ from policyengine_simulation_executor.stage12_runtime import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time coordinate_report as run_report_coordinator, @@ -452,19 +457,21 @@ def coordinate_report( if context.get("production_function_call_id") else STAGE12_CANONICAL_REPORT_STAGES ) + attributes = { + "runner_name": "stage12", + "execution_mode": ( + "shadow" if context.get("production_function_call_id") else "authoritative" + ), + } + if context.get("result_delivery") == ReportResultDelivery.RETURN_TO_CALLER: + # Telemetry tells returned reports apart from Stage 12 runs. + attributes["result_delivery"] = ReportResultDelivery.RETURN_TO_CALLER.value try: propagated = normalize_observability_context(observability_context) result = None with runtime.operation( stage_plan.name(Stage.STAGE12_COORDINATOR_EXECUTION), - attributes={ - "runner_name": "stage12", - "execution_mode": ( - "shadow" - if context.get("production_function_call_id") - else "authoritative" - ), - }, + attributes=attributes, remote_context=propagated, ): result = run_report_coordinator( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 7f81ded06..7fe4b5311 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -4,14 +4,23 @@ import logging import time +from dataclasses import dataclass from datetime import UTC, datetime from functools import partial from hashlib import sha256 from typing import TYPE_CHECKING, Any from uuid import uuid4 +from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from policyengine_simulation_contract.report_execution import ( + ReportExecutionError, + ReportExecutionProvenance, + ReportExecutionResult, + report_aggregate_sha256, +) from policyengine_simulation_contract.stage12_execution import ( AggregateReportArtifactDescriptor, + AggregateReportArtifactPayload, ComparisonReportRecord, ComparisonRunAggregationStatus, ComparisonRunLifecycleStatus, @@ -33,6 +42,7 @@ from policyengine_simulation_executor.stage12_artifacts import ( Stage12ArtifactStore, + comparison_run_prefix, deserialize_calculation_provenance, deserialize_simulation_frames, deserialize_uk_local_authority_metadata, @@ -65,6 +75,8 @@ Stage12ExecutionError, Stage12FailureDetail, failure_detail_from_record, + failure_detail_from_returned_error, + report_execution_error, ) from .output_planning import ( plan_simulation_input, @@ -77,7 +89,11 @@ if TYPE_CHECKING: from collections.abc import Callable + import pandas as pd from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.uk_geography import UKLocalAuthorityMetadata + + from .output_planning import StageScope SIMULATION_POLL_INITIAL_SECONDS = 0.25 SIMULATION_POLL_MAX_SECONDS = 1.0 @@ -300,12 +316,66 @@ def _child_failure_detail( return detail +def _returned_child_failure_detail( + *, + error: Exception, + simulation: PlannedSimulationExecutionInput, + runtime: ObservabilityRuntime | None, + progress: StageProgress, +) -> Stage12FailureDetail: + """Report a returned report's failed child call; return its report's failure. + + The failure the child raised as ``ReportExecutionError`` takes precedence + over the redacted call error, as a Stage 12 child's recorded failure does. + """ + detail = report_failure( + error, + runtime=runtime, + failure_scope="child_wait", + default_code="simulation_invocation_failed", + attributes=_simulation_failure_attributes(simulation), + progress=progress, + ) + if isinstance(error, ReportExecutionError): + return failure_detail_from_returned_error(error) + return detail + + +def _cancel_returned_simulations( + *, + calls: dict[str, ChildCall], + simulations: dict[str, PlannedSimulationExecutionInput], + roles: set[str], + detail: Stage12FailureDetail | None = None, +) -> None: + """Stop each started, unfinished child of a returned report through Modal. + + No record exists to carry ``detail`` or to tell a segmented child to stop, + so every child is cancelled through Modal. Modal does not cancel the + region-group segment calls a segmented child started; they run until they + finish or time out. A cancellation error is logged rather than raised, so + it never replaces the failure that caused the cancellation. + """ + del detail + for role in sorted(roles): + try: + calls[role].cancel() + except Exception as error: # noqa: BLE001 # a failed Modal cancellation is logged and must not replace the failure that caused it + _log_child_operation_failure( + logging.WARNING, + "Stage 12 could not cancel %s simulation %s in report %s (%s)", + simulation=simulations[role], + error=error, + ) + + def _observe_simulations( *, calls: dict[str, ChildCall], simulations: dict[str, PlannedSimulationExecutionInput], descriptors: dict[str, SimulationArtifactDescriptor], - store: ComparisonStore, + child_failure_detail: Callable[..., Stage12FailureDetail], + cancel_simulations: Callable[..., None], runtime: ObservabilityRuntime | None, progress: StageProgress, wait_stage_name: str, @@ -314,6 +384,12 @@ def _observe_simulations( deadline: float, timeout_seconds: float, ) -> None: + """Wait for every started child within the shared window, or stop them all. + + ``child_failure_detail(error=, simulation=)`` reports a failed child and + returns the report's failure; ``cancel_simulations(roles=, detail=)`` + stops the children still running. + """ pending = set(calls) poll_interval = SIMULATION_POLL_INITIAL_SECONDS @@ -329,18 +405,12 @@ def _observe_simulations( except TimeoutError: continue except Exception as error: # noqa: BLE001 # a child call's failure can be any exception type; it is recorded and stops the other child - detail = _child_failure_detail( + detail = child_failure_detail( error=error, simulation=simulations[role], - store=store, - runtime=runtime, - progress=progress, ) - _cancel_pending_simulations( - calls=calls, - simulations=simulations, + cancel_simulations( roles=pending - {role}, - store=store, detail=_cancellation(f"the {role} simulation failed"), ) raise Stage12ExecutionError(detail) from None @@ -366,11 +436,8 @@ def _observe_simulations( runtime=runtime, failure_scope="child_wait_timeout", ) - _cancel_pending_simulations( - calls=calls, - simulations=simulations, + cancel_simulations( roles=pending, - store=store, detail=_cancellation( f"the {timeout_seconds:g}-second wait window closed" ), @@ -527,6 +594,24 @@ def _require_artifact_digest( raise ValueError(mismatch_message) +def _validate_dispatch( + *, + report: ReportExecutionInput, + context: Stage12InvocationContext, + application_name: str, +) -> None: + """Reject a dispatch whose application, worker or bundle differs from the report.""" + if application_name != context.modal_application: + msg = "report coordinator application differs from dispatch context" + raise ValueError(msg) + if report.baseline.bundle.policyengine_version != context.worker_version: + msg = "report worker version differs from dispatch context" + raise ValueError(msg) + if report.baseline.bundle.bundle_manifest_sha256 != context.bundle_manifest_sha256: + msg = "report bundle digest differs from dispatch context" + raise ValueError(msg) + + def _child_context_payload(context: Stage12InvocationContext) -> dict[str, Any]: """Return the dispatch context a child simulation receives. @@ -540,10 +625,304 @@ def _child_context_payload(context: Stage12InvocationContext) -> dict[str, Any]: return context.model_dump(mode="json") +def _plan_simulations( + report: ReportExecutionInput, + *, + output_plan_resolver: Callable[[ReportExecutionInput], Stage12OutputPlan] | None, + stage_scope: StageScope, +) -> tuple[Stage12OutputPlan, tuple[PlannedSimulationExecutionInput, ...]]: + """Resolve the report's shared output plan and attach it to both inputs.""" + with stage_scope( + Stage.STAGE12_OUTPUT_PLANNING, + {"country": report.baseline.geography.country}, + ): + if output_plan_resolver is None: + output_plan = resolve_report_output_plan( + report, + stage_scope=stage_scope, + ) + else: + output_plan = output_plan_resolver(report) + planned_inputs: list[PlannedSimulationExecutionInput] = [] + for simulation in (report.baseline, report.reform): + with stage_scope( + Stage.STAGE12_CHILD_INPUT_PLANNING, + {"simulation_role": simulation.role.value}, + ): + planned_inputs.append(plan_simulation_input(simulation, output_plan)) + return output_plan, tuple(planned_inputs) + + +@dataclass(frozen=True) +class _SimulationOutputs: + """The checked baseline and reform outputs one report aggregates.""" + + baseline: SimulationArtifactDescriptor + reform: SimulationArtifactDescriptor + baseline_frames: dict[str, pd.DataFrame] + reform_frames: dict[str, pd.DataFrame] + uk_local_authority_metadata: UKLocalAuthorityMetadata | None + + def aggregate( + self, + report: ReportExecutionInput, + aggregator: Callable[..., dict[str, Any]], + ) -> dict[str, Any]: + """Return ``aggregator``'s aggregate report payload for these outputs.""" + return aggregator( + report=report, + baseline_frames=self.baseline_frames, + reform_frames=self.reform_frames, + baseline_descriptor=self.baseline, + reform_descriptor=self.reform, + uk_local_authority_metadata=self.uk_local_authority_metadata, + ) + + +def _read_simulation_outputs( + *, + report: ReportExecutionInput, + output_plan: Stage12OutputPlan, + descriptors: dict[str, SimulationArtifactDescriptor], + artifact_storage: Stage12ArtifactStore, + stage_span: StageScope, +) -> _SimulationOutputs: + """Read both children's artifacts and check they form this report's pair. + + Raises ``ValueError`` when the artifacts do not align, a digest does not + match, the UK authority metadata disagrees, or a frame misses the plan. + """ + baseline = descriptors["baseline"] + reform = descriptors["reform"] + validate_aligned_outputs(report, output_plan, baseline, reform) + with stage_span(Stage.STAGE12_ARTIFACT_READ): + baseline_payload = artifact_storage.read(baseline.artifact.uri) + reform_payload = artifact_storage.read(reform.artifact.uri) + _require_artifact_digest( + baseline_payload, + baseline.artifact.content_sha256, + mismatch_message="baseline artifact digest mismatch", + ) + _require_artifact_digest( + reform_payload, + reform.artifact.content_sha256, + mismatch_message="reform artifact digest mismatch", + ) + baseline_frames = deserialize_simulation_frames(baseline_payload) + reform_frames = deserialize_simulation_frames(reform_payload) + uk_local_authority_metadata = validate_uk_local_authority_metadata( + report.baseline.geography.country, + deserialize_uk_local_authority_metadata(baseline_payload), + deserialize_uk_local_authority_metadata(reform_payload), + region_code=report.baseline.geography.region, + ) + validate_output_frames(baseline_frames, output_plan) + validate_output_frames(reform_frames, output_plan) + return _SimulationOutputs( + baseline=baseline, + reform=reform, + baseline_frames=baseline_frames, + reform_frames=reform_frames, + uk_local_authority_metadata=uk_local_authority_metadata, + ) + + +def _returned_result( + *, + report: ReportExecutionInput, + context: Stage12InvocationContext, + output_plan: Stage12OutputPlan, + outputs: _SimulationOutputs, + aggregate: dict[str, Any], + coordinator_invocation_id: str, +) -> ReportExecutionResult: + """Check the aggregate and build the result a returned report delivers. + + The aggregate passes the checks a Stage 12 aggregate artifact write + applies, must describe this report, and its result must be a + ``SingleYearMacroOutput``. + """ + payload = AggregateReportArtifactPayload.model_validate(aggregate) + if ( + payload.evaluation_id != report.evaluation_id + or payload.bundle != report.baseline.bundle + or payload.requested_aggregates != report.requested_aggregates + ): + msg = "aggregate does not describe this report" + raise ValueError(msg) + macro_output = SingleYearMacroOutput.model_validate(payload.result) + return ReportExecutionResult( + evaluation_id=report.evaluation_id, + requested_aggregates=report.requested_aggregates, + aggregate=macro_output, + aggregate_sha256=report_aggregate_sha256(macro_output), + output_plan=output_plan, + baseline=outputs.baseline, + reform=outputs.reform, + provenance=ReportExecutionProvenance( + bundle=report.baseline.bundle, + worker_version=context.worker_version, + modal_environment=context.modal_environment, + modal_application=context.modal_application, + simulation_callable=context.simulation_callable, + version_manifest_sha256=context.version_manifest_sha256, + coordinator_invocation_id=coordinator_invocation_id, + ), + completed_at=datetime.now(UTC), + ) + + +def _coordinate_report_for_caller( + report: ReportExecutionInput, + context: Stage12InvocationContext, + *, + application_name: str, + coordinator_invocation_id: str, + artifact_storage: Stage12ArtifactStore, + child_invoker: ChildInvoker, + aggregator: Callable[..., dict[str, Any]], + output_plan_resolver: Callable[[ReportExecutionInput], Stage12OutputPlan] | None, + runtime: ObservabilityRuntime | None, + sleep: Callable[[float], None], + monotonic: Callable[[], float], + child_wait_deadline: float, + wait_timeout_seconds: float, +) -> dict[str, Any]: + """Run one report and return its ``ReportExecutionResult``, writing no rows. + + Nothing claims the report, so a repeated dispatch runs it again; the + caller deduplicates. Artifacts go under the Stage 12 layout for the + context's environment, creation time and report. Both children start + with the same delivery, so they write no rows either, and a child that + must stop is cancelled through Modal. No aggregate artifact is written: + the aggregate is returned inline. Every failure is reported once, + recorded as the operation's cause, and raised as ``ReportExecutionError`` + with its redacted detail. + """ + stage_plan = STAGE12_CANONICAL_REPORT_STAGES + progress = StageProgress() + stage_span = partial( + tracked_stage, + stage_plan=stage_plan, + progress=progress, + runtime=runtime, + ) + stage_scope = partial(stage_span, operation=True) + descriptors: dict[str, SimulationArtifactDescriptor] = {} + calls: dict[str, ChildCall] = {} + simulations: dict[str, PlannedSimulationExecutionInput] = {} + cancel_simulations = partial( + _cancel_returned_simulations, + calls=calls, + simulations=simulations, + ) + try: + with stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): + context = context.model_copy( + update={ + "artifact_prefix": comparison_run_prefix( + environment=context.environment, + created_at=context.created_at, + evaluation_id=report.evaluation_id, + ) + } + ) + output_plan, planned_simulations = _plan_simulations( + report, + output_plan_resolver=output_plan_resolver, + stage_scope=stage_scope, + ) + simulations.update( + {simulation.role.value: simulation for simulation in planned_simulations} + ) + for simulation in planned_simulations: + try: + with stage_span( + Stage.STAGE12_CHILD_DISPATCH, + {"simulation_role": simulation.role.value}, + ): + call = child_invoker.spawn( + application_name=application_name, + function_name=context.simulation_callable, + environment=context.modal_environment, + simulation=simulation.model_dump(mode="json"), + context=_child_context_payload(context), + observability_context=( + runtime.capture_context() if runtime is not None else None + ), + ) + except Exception as error: # noqa: BLE001 # any dispatch failure stops the started peer and returns to the caller as a typed error + detail = report_failure( + error, + runtime=runtime, + failure_scope="child_dispatch", + default_code="simulation_dispatch_failed", + attributes=_simulation_failure_attributes(simulation), + progress=progress, + ) + cancel_simulations(roles=set(calls)) + raise Stage12ExecutionError(detail) from None + calls[simulation.role.value] = call + with stage_span(Stage.STAGE12_CHILD_WAIT): + _observe_simulations( + calls=calls, + simulations=simulations, + descriptors=descriptors, + child_failure_detail=partial( + _returned_child_failure_detail, + runtime=runtime, + progress=progress, + ), + cancel_simulations=cancel_simulations, + runtime=runtime, + progress=progress, + wait_stage_name=stage_plan.name(Stage.STAGE12_CHILD_WAIT), + sleep=sleep, + monotonic=monotonic, + deadline=child_wait_deadline, + timeout_seconds=wait_timeout_seconds, + ) + outputs = _read_simulation_outputs( + report=report, + output_plan=output_plan, + descriptors=descriptors, + artifact_storage=artifact_storage, + stage_span=stage_span, + ) + with stage_span(Stage.STAGE12_AGGREGATION): + result = _returned_result( + report=report, + context=context, + output_plan=output_plan, + outputs=outputs, + aggregate=outputs.aggregate(report, aggregator), + coordinator_invocation_id=coordinator_invocation_id, + ) + return result.model_dump(mode="json") + except Exception as error: # noqa: BLE001 # every failure returns to the caller as a typed error with its redacted detail + if isinstance(error, Stage12ExecutionError): + # Child, dispatch and wait-timeout failures raise + # Stage12ExecutionError only after reporting the original failure + # and stopping their unfinished peers. + detail = error.detail + else: + detail = report_failure( + error, + runtime=runtime, + failure_scope="report_coordination", + default_code="report_coordination_failed", + attributes={"evaluation_id": str(report.evaluation_id)}, + progress=progress, + ) + cancel_simulations(roles=set(calls) - set(descriptors)) + record_failure_cause(runtime, error, detail.error_summary) + raise report_execution_error(report.evaluation_id, detail) from None + + def coordinate_report( payload: object, context_payload: object, - parent_payload: object, + parent_payload: object | None = None, *, application_name: str, coordinator_invocation_id: str, @@ -568,21 +947,46 @@ def coordinate_report( raises ``ValueError`` before the claim; any later failure stops the unfinished children, marks the report FAILED and raises a ``RuntimeError`` that carries no calculation detail. + + When the context's ``result_delivery`` returns the report to its caller, + ``parent_payload`` must be absent and ``store`` is not used: the report + returns a ``ReportExecutionResult`` payload or raises + ``ReportExecutionError`` (see ``_coordinate_report_for_caller``). """ # Open the shared child-wait window before any setup work. child_wait_deadline = monotonic() + wait_timeout_seconds report = ReportExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) + if context.result_delivery is ReportResultDelivery.RETURN_TO_CALLER: + if parent_payload is not None: + msg = "a report returned to its caller takes no Stage 12 parent record" + raise ValueError(msg) + _validate_dispatch( + report=report, + context=context, + application_name=application_name, + ) + return _coordinate_report_for_caller( + report, + context, + application_name=application_name, + coordinator_invocation_id=coordinator_invocation_id, + artifact_storage=artifacts or artifact_store(), + child_invoker=invoker or ModalChildInvoker(), + aggregator=aggregator, + output_plan_resolver=output_plan_resolver, + runtime=runtime, + sleep=sleep, + monotonic=monotonic, + child_wait_deadline=child_wait_deadline, + wait_timeout_seconds=wait_timeout_seconds, + ) submitted_parent = ComparisonReportRecord.model_validate(parent_payload) - if application_name != context.modal_application: - msg = "report coordinator application differs from dispatch context" - raise ValueError(msg) - if report.baseline.bundle.policyengine_version != context.worker_version: - msg = "report worker version differs from dispatch context" - raise ValueError(msg) - if report.baseline.bundle.bundle_manifest_sha256 != context.bundle_manifest_sha256: - msg = "report bundle digest differs from dispatch context" - raise ValueError(msg) + _validate_dispatch( + report=report, + context=context, + application_name=application_name, + ) _validate_submitted_parent( report=report, context=context, @@ -633,27 +1037,11 @@ def coordinate_report( } ) function_name = context.simulation_callable - with stage_scope( - Stage.STAGE12_OUTPUT_PLANNING, - {"country": report.baseline.geography.country}, - ): - if output_plan_resolver is None: - output_plan = resolve_report_output_plan( - report, - stage_scope=stage_scope, - ) - else: - output_plan = output_plan_resolver(report) - planned_inputs: list[PlannedSimulationExecutionInput] = [] - for simulation in (report.baseline, report.reform): - with stage_scope( - Stage.STAGE12_CHILD_INPUT_PLANNING, - {"simulation_role": simulation.role.value}, - ): - planned_inputs.append( - plan_simulation_input(simulation, output_plan) - ) - planned_simulations = tuple(planned_inputs) + output_plan, planned_simulations = _plan_simulations( + report, + output_plan_resolver=output_plan_resolver, + stage_scope=stage_scope, + ) simulations = { simulation.role.value: simulation for simulation in planned_simulations } @@ -768,7 +1156,18 @@ def coordinate_report( calls=calls, simulations=simulations, descriptors=descriptors, - store=persistence, + child_failure_detail=partial( + _child_failure_detail, + store=persistence, + runtime=runtime, + progress=progress, + ), + cancel_simulations=partial( + _cancel_pending_simulations, + calls=calls, + simulations=simulations, + store=persistence, + ), runtime=runtime, progress=progress, wait_stage_name=stage_plan.name(Stage.STAGE12_CHILD_WAIT), @@ -777,41 +1176,15 @@ def coordinate_report( deadline=child_wait_deadline, timeout_seconds=wait_timeout_seconds, ) - baseline = descriptors["baseline"] - reform = descriptors["reform"] - validate_aligned_outputs(report, output_plan, baseline, reform) - with stage_span(Stage.STAGE12_ARTIFACT_READ): - baseline_payload = artifact_storage.read(baseline.artifact.uri) - reform_payload = artifact_storage.read(reform.artifact.uri) - _require_artifact_digest( - baseline_payload, - baseline.artifact.content_sha256, - mismatch_message="baseline artifact digest mismatch", - ) - _require_artifact_digest( - reform_payload, - reform.artifact.content_sha256, - mismatch_message="reform artifact digest mismatch", + outputs = _read_simulation_outputs( + report=report, + output_plan=output_plan, + descriptors=descriptors, + artifact_storage=artifact_storage, + stage_span=stage_span, ) - baseline_frames = deserialize_simulation_frames(baseline_payload) - reform_frames = deserialize_simulation_frames(reform_payload) - uk_local_authority_metadata = validate_uk_local_authority_metadata( - report.baseline.geography.country, - deserialize_uk_local_authority_metadata(baseline_payload), - deserialize_uk_local_authority_metadata(reform_payload), - region_code=report.baseline.geography.region, - ) - validate_output_frames(baseline_frames, output_plan) - validate_output_frames(reform_frames, output_plan) with stage_span(Stage.STAGE12_AGGREGATION): - aggregate = aggregator( - report=report, - baseline_frames=baseline_frames, - reform_frames=reform_frames, - baseline_descriptor=baseline, - reform_descriptor=reform, - uk_local_authority_metadata=uk_local_authority_metadata, - ) + aggregate = outputs.aggregate(report, aggregator) with stage_span(Stage.STAGE12_AGGREGATE_ARTIFACT_WRITE): aggregate_artifact = artifact_storage.write_aggregate( prefix=context.artifact_prefix, @@ -820,8 +1193,8 @@ def coordinate_report( descriptor = AggregateReportArtifactDescriptor( evaluation_id=report.evaluation_id, artifact=aggregate_artifact, - baseline_artifact_sha256=baseline.artifact.content_sha256, - reform_artifact_sha256=reform.artifact.content_sha256, + baseline_artifact_sha256=outputs.baseline.artifact.content_sha256, + reform_artifact_sha256=outputs.reform.artifact.content_sha256, bundle=report.baseline.bundle, ) completed = datetime.now(UTC) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py index ad97388b5..75f34ab56 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py @@ -107,6 +107,16 @@ def report_execution_error( ) +def failure_detail_from_returned_error( + error: ReportExecutionError, +) -> Stage12FailureDetail: + """Recreate the failure a returned simulation raised, keeping its values.""" + return failure_detail_from_record( + error_code=error.failure.error_code, + error_summary=error.failure.error_summary, + ) + + def redacted_failure_detail( error: BaseException, *, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index 6d0acb14d..c8db53584 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -289,6 +289,68 @@ def fail(*_, **__): assert runtime.shutdown_calls == 1 +def _invoke_coordinator(monkeypatch, module, context, *parent): + from policyengine_simulation_executor import stage12_runtime + + runtime = RecordingRuntime() + received = {} + monkeypatch.setattr(module, "init_process_observability", lambda **_: runtime) + monkeypatch.setattr( + module.modal, + "current_function_call_id", + lambda: "call-1", + raising=False, + ) + + def implementation(payload, context_payload, parent_payload, **options): + received.update( + parent=parent_payload, + coordinator_invocation_id=options["coordinator_invocation_id"], + ) + return {"ok": True} + + monkeypatch.setattr(stage12_runtime, "coordinate_report", implementation) + result = module.coordinate_report( + {}, + context, + *parent, + observability_context=_observability_context(), + ) + return result, received, runtime + + +def test_coordinator_runs_a_returned_report_without_a_parent(monkeypatch) -> None: + module = _load(monkeypatch) + + result, received, runtime = _invoke_coordinator( + monkeypatch, module, {"result_delivery": "return_to_caller"} + ) + + assert result == {"ok": True} + assert received == {"parent": None, "coordinator_invocation_id": "call-1"} + assert runtime.operations[0]["attributes"] == { + "runner_name": "stage12", + "execution_mode": "authoritative", + "result_delivery": "return_to_caller", + } + + +def test_coordinator_operation_attributes_are_unchanged_for_stage12( + monkeypatch, +) -> None: + module = _load(monkeypatch) + + _, received, runtime = _invoke_coordinator( + monkeypatch, module, {}, {"evaluation_id": "parent"} + ) + + assert received["parent"] == {"evaluation_id": "parent"} + assert runtime.operations[0]["attributes"] == { + "runner_name": "stage12", + "execution_mode": "authoritative", + } + + def test_external_package_values_are_assertions_only(monkeypatch) -> None: install_fake_modal(monkeypatch) monkeypatch.setenv("STAGE12_EXPECT_US_VERSION", "not-the-bundled-version") diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 2b71bd7a7..0b15e840c 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -32,7 +32,13 @@ configure, ) from policyengine_simulation_contract.macro_output import SingleYearMacroOutput -from policyengine_simulation_contract.report_execution import ReportExecutionError +from policyengine_simulation_contract.report_execution import ( + ReportExecutionError, + ReportExecutionFailure, + ReportExecutionProvenance, + ReportExecutionResult, + report_aggregate_sha256, +) from policyengine_simulation_contract.spm import ( SPMComparisonProvenance, SPMInputError, @@ -116,6 +122,7 @@ detect_uk_local_authority_metadata, ) from policyengine_stage12_persistence import store as persistence_store_module +from pydantic import ValidationError from sqlalchemy.exc import OperationalError NOW = datetime(2026, 9, 15, tzinfo=UTC) @@ -4089,3 +4096,315 @@ def segmented_calculator(_value, _cancellation_requested): assert raised.value.failure.error_summary == ( "Segment 3: MemoryError in stage12_calculation" ) + + +def _macro_aggregate(**_values): + return _aggregate_payload(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + +def _coordinate_returned_report(invoker, *, context=None, report=None, **options): + context = context or _returning_context() + options.setdefault("artifacts", FakeArtifacts()) + options.setdefault("aggregator", _macro_aggregate) + return coordinate_report( + (report or _report()).model_dump(mode="json"), + context.model_dump(mode="json"), + None, + application_name=context.modal_application, + coordinator_invocation_id="coordinator-1", + invoker=invoker, + **options, + ) + + +def test_returned_report_returns_its_typed_result_without_the_database( + forbid_database, +) -> None: + artifacts = FakeArtifacts() + invoker = ContextRecordingInvoker(artifacts) + observed = {} + + def aggregate(**values): + observed.update(values) + return _macro_aggregate() + + # The coordinator, not the caller, chooses where artifacts go. + caller_context = _returning_context().model_copy( + update={"artifact_prefix": f"caller-chosen/{EVALUATION_ID}"} + ) + + returned = _coordinate_returned_report( + invoker, + context=caller_context, + artifacts=artifacts, + aggregator=aggregate, + ) + + result = ReportExecutionResult.model_validate(returned) + assert returned == result.model_dump(mode="json") + assert result.evaluation_id == EVALUATION_ID + assert result.requested_aggregates == tuple(ReportAggregate) + assert result.aggregate == SingleYearMacroOutput.model_validate( + CURRENT_SINGLE_YEAR_MACRO_RESULT + ) + assert result.aggregate_sha256 == report_aggregate_sha256(result.aggregate) + assert result.output_plan == _output_plan() + assert result.baseline.simulation_execution_id == BASELINE_ID + assert result.reform.simulation_execution_id == REFORM_ID + for descriptor in (result.baseline, result.reform): + stored = artifacts.read(descriptor.artifact.uri) + assert sha256(stored).hexdigest() == descriptor.artifact.content_sha256 + assert result.provenance == ReportExecutionProvenance( + bundle=_bundle(), + worker_version="5.2.0", + modal_environment="staging", + modal_application="policyengine-simulation-v2-py5-2-0", + simulation_callable="run_single_simulation_us", + version_manifest_sha256="b" * 64, + coordinator_invocation_id="coordinator-1", + ) + assert result.completed_at.tzinfo is not None + assert observed["baseline_frames"]["household"][ + "household_net_income" + ].tolist() == [100.0] + assert observed["reform_frames"]["household"]["household_net_income"].tolist() == [ + 120.0 + ] + assert artifacts.aggregate_writes == [] + assert invoker.events == ["spawn:baseline", "spawn:reform"] + 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 [context["result_delivery"] for context in invoker.child_contexts] == [ + "return_to_caller", + "return_to_caller", + ] + assert [context["artifact_prefix"] for context in invoker.child_contexts] == [ + _context().artifact_prefix, + _context().artifact_prefix, + ] + + +def test_returned_report_takes_no_stage12_parent_record(forbid_database) -> None: + invoker = ConcurrentInvoker(FakeArtifacts()) + + with pytest.raises(ValueError, match="no Stage 12 parent record"): + coordinate_report( + _report().model_dump(mode="json"), + _returning_context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + artifacts=FakeArtifacts(), + invoker=invoker, + aggregator=_macro_aggregate, + ) + + assert invoker.events == [] + + +def test_stage12_report_still_requires_its_parent_record() -> None: + invoker = ConcurrentInvoker(FakeArtifacts()) + + with pytest.raises(ValidationError): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + None, + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=FakeStore(), + artifacts=FakeArtifacts(), + invoker=invoker, + aggregator=lambda **_: _aggregate_payload(), + ) + + assert invoker.events == [] + + +def test_returned_report_checks_the_dispatch_context_before_starting() -> None: + invoker = ConcurrentInvoker(FakeArtifacts()) + + with pytest.raises(ValueError, match="application differs"): + coordinate_report( + _report().model_dump(mode="json"), + _returning_context().model_dump(mode="json"), + None, + application_name="policyengine-simulation-v2-py9-9-9", + coordinator_invocation_id="coordinator-1", + store=NoDatabaseStore(), + artifacts=FakeArtifacts(), + invoker=invoker, + aggregator=_macro_aggregate, + ) + + assert invoker.events == [] + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_returned_report_raises_the_failed_childs_detail_and_stops_its_peer( + forbid_database, + segmented, +) -> None: + child_failure = ReportExecutionFailure( + evaluation_id=EVALUATION_ID, + error_code="SPM_YEAR_UNAVAILABLE", + error_summary="No SPM values for 2026", + ) + invoker = ScriptedChildInvoker( + wait_errors={"baseline": ReportExecutionError(child_failure)} + ) + runtime = RecordingRuntime() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + invoker, + report=_report_with_segmentation(segmented=segmented), + runtime=runtime, + ) + + assert raised.value.failure == child_failure + assert raised.value.__cause__ is None + assert raised.value.__suppress_context__ is True + # No record can tell a segmented child to stop, so Modal cancels it. + assert invoker.calls["reform"].cancel_count == 1 + assert invoker.calls["baseline"].cancel_count == 0 + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == "No SPM values for 2026" + + +def test_returned_report_redacts_a_child_call_failure_without_detail( + forbid_database, +) -> None: + invoker = ScriptedChildInvoker( + wait_errors={"reform": LookupError("worker lost with token=opaque-token")} + ) + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report(invoker) + + failure = raised.value.failure + assert failure.error_code == "simulation_invocation_failed" + assert _failure_cause(failure.error_summary) == ( + "LookupError in stage12_child_wait (reform): worker lost with token=[REDACTED]" + ) + assert invoker.calls["baseline"].cancel_count == 1 + assert invoker.calls["reform"].cancel_count == 0 + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_returned_report_wait_timeout_stops_every_child_through_modal( + forbid_database, + segmented, +) -> None: + invoker = ScriptedChildInvoker() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + invoker, + report=_report_with_segmentation(segmented=segmented), + wait_timeout_seconds=0, + ) + + assert raised.value.failure.error_code == "simulation_wait_timeout" + assert raised.value.failure.error_summary == ( + "Stage 12 simulations did not finish within 0 seconds of the report " + "coordinator starting." + ) + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + + +def test_returned_report_dispatch_failure_stops_the_started_child( + forbid_database, +) -> None: + invoker = ReformDispatchFailureInvoker() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report(invoker) + + failure = raised.value.failure + assert failure.error_code == "simulation_dispatch_failed" + assert _failure_cause(failure.error_summary) == ( + "ConnectionError in stage12_child_dispatch (reform): " + "dispatch contains sensitive details" + ) + assert invoker.events == ["spawn:baseline", "spawn:reform"] + assert invoker.calls["baseline"].cancel_count == 1 + + +def test_returned_report_cancellation_error_does_not_replace_the_failure( + forbid_database, + caplog, +) -> None: + invoker = ScriptedChildInvoker( + wait_errors={"reform": LookupError("reform worker failed")}, + cancel_errors={"baseline": ConnectionError("cancel contains sensitive data")}, + ) + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report(invoker) + + assert raised.value.failure.error_code == "simulation_invocation_failed" + assert invoker.calls["baseline"].cancel_count == 1 + _assert_log_identifies_child( + _coordination_log(caplog, logging.WARNING), + role="baseline", + error_type="ConnectionError", + ) + + +def test_returned_report_rejects_an_aggregate_outside_the_macro_output_schema( + forbid_database, +) -> None: + artifacts = FakeArtifacts() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + ConcurrentInvoker(artifacts), + artifacts=artifacts, + aggregator=lambda **_: _aggregate_payload({"budget": {}}), + ) + + failure = raised.value.failure + assert failure.error_code == "report_coordination_failed" + assert _failure_cause(failure.error_summary).startswith( + "ValidationError in stage12_aggregation: " + ) + assert artifacts.aggregate_writes == [] + + +def test_returned_report_rejects_an_aggregate_for_another_report( + forbid_database, +) -> None: + artifacts = FakeArtifacts() + + def aggregate_for_another_report(**_values): + return {**_macro_aggregate(), "evaluation_id": str(uuid4())} + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + ConcurrentInvoker(artifacts), + artifacts=artifacts, + aggregator=aggregate_for_another_report, + ) + + assert raised.value.failure.error_code == "report_coordination_failed" + assert "does not describe this report" in raised.value.failure.error_summary + + +def test_returned_report_rejects_artifacts_that_do_not_satisfy_the_plan( + forbid_database, +) -> None: + artifacts = FakeArtifacts() + + with pytest.raises(ReportExecutionError) as raised: + _coordinate_returned_report( + ConcurrentInvoker(artifacts, wrong_plan_digest=True), + artifacts=artifacts, + aggregator=lambda **_: pytest.fail("aggregation must not start"), + ) + + assert raised.value.failure.error_code == "report_coordination_failed" + assert "do not satisfy the report output plan" in ( + raised.value.failure.error_summary + ) From cbb35b162373dbfa2f060c11cb38bea3a6d207ab Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:11:31 +0400 Subject: [PATCH 3/3] Type the child wait callbacks of both report coordinator modes _observe_simulations takes the failure and cancellation handlers that the Stage 12 and returned-report paths supply. They were typed Callable[..., Stage12FailureDetail] and Callable[..., None], so a type checker accepted a handler with any arguments, and _cancel_returned_simulations could drop its detail keyword without any check noticing. Two private Protocols now state the keyword arguments each handler receives; pyright reports a handler that does not accept them. _read_simulation_outputs declared its stage_span as StageScope, which requires an attributes argument, but calls it with the stage alone. Its annotation now names the one-argument form it uses. Typing-only change: behavior is unchanged, so no test can fail first. The existing Stage 12 and returned-report tests pass, and pyright reports no errors in the coordination module. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/coordination.py | 43 ++++++++++++++++--- 1 file changed, 36 insertions(+), 7 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index 7fe4b5311..719df1197 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -8,7 +8,7 @@ from datetime import UTC, datetime from functools import partial from hashlib import sha256 -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Protocol from uuid import uuid4 from policyengine_simulation_contract.macro_output import SingleYearMacroOutput @@ -88,6 +88,7 @@ if TYPE_CHECKING: from collections.abc import Callable + from contextlib import AbstractContextManager import pandas as pd from policyengine_observability import ObservabilityRuntime @@ -369,13 +370,42 @@ def _cancel_returned_simulations( ) +class _ChildFailureHandler(Protocol): + """Report a failed child call and return the failure its report records. + + The Stage 12 path reads and records the child's row; a returned report + uses the ``ReportExecutionError`` its child raised. + """ + + def __call__( + self, + *, + error: Exception, + simulation: PlannedSimulationExecutionInput, + ) -> Stage12FailureDetail: + """Return the report's failure for ``simulation``'s failed call.""" + ... + + +class _ChildCanceller(Protocol): + """Stop the started, unfinished children with the given roles. + + ``detail`` is the cancellation a Stage 12 child's row records; a returned + report's children have no row to record it on. + """ + + def __call__(self, *, roles: set[str], detail: Stage12FailureDetail) -> None: + """Stop each child named in ``roles``.""" + ... + + def _observe_simulations( *, calls: dict[str, ChildCall], simulations: dict[str, PlannedSimulationExecutionInput], descriptors: dict[str, SimulationArtifactDescriptor], - child_failure_detail: Callable[..., Stage12FailureDetail], - cancel_simulations: Callable[..., None], + child_failure_detail: _ChildFailureHandler, + cancel_simulations: _ChildCanceller, runtime: ObservabilityRuntime | None, progress: StageProgress, wait_stage_name: str, @@ -386,9 +416,8 @@ def _observe_simulations( ) -> None: """Wait for every started child within the shared window, or stop them all. - ``child_failure_detail(error=, simulation=)`` reports a failed child and - returns the report's failure; ``cancel_simulations(roles=, detail=)`` - stops the children still running. + ``child_failure_detail`` reports a failed child and returns the report's + failure; ``cancel_simulations`` stops the children still running. """ pending = set(calls) poll_interval = SIMULATION_POLL_INITIAL_SECONDS @@ -685,7 +714,7 @@ def _read_simulation_outputs( output_plan: Stage12OutputPlan, descriptors: dict[str, SimulationArtifactDescriptor], artifact_storage: Stage12ArtifactStore, - stage_span: StageScope, + stage_span: Callable[[Stage], AbstractContextManager[object]], ) -> _SimulationOutputs: """Read both children's artifacts and check they form this report's pair.