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..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 @@ -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 typing import TYPE_CHECKING, Any, Protocol 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, @@ -76,8 +88,13 @@ if TYPE_CHECKING: from collections.abc import Callable + from contextlib import AbstractContextManager + 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 +317,95 @@ 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, + ) + + +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], - store: ComparisonStore, + child_failure_detail: _ChildFailureHandler, + cancel_simulations: _ChildCanceller, runtime: ObservabilityRuntime | None, progress: StageProgress, wait_stage_name: str, @@ -314,6 +414,11 @@ 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`` 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 @@ -329,18 +434,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 +465,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 +623,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 +654,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: Callable[[Stage], AbstractContextManager[object]], +) -> _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 +976,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 +1066,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 +1185,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 +1205,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 +1222,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 a42f9c05b..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 @@ -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,33 @@ 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 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/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_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 405d12663..0b15e840c 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -32,6 +32,13 @@ configure, ) from policyengine_simulation_contract.macro_output import SingleYearMacroOutput +from policyengine_simulation_contract.report_execution import ( + ReportExecutionError, + ReportExecutionFailure, + ReportExecutionProvenance, + ReportExecutionResult, + report_aggregate_sha256, +) from policyengine_simulation_contract.spm import ( SPMComparisonProvenance, SPMInputError, @@ -60,6 +67,7 @@ ReportAggregate, ReportExecutionInput, ReportOutputRequirements, + ReportResultDelivery, ResultComparisonStatus, RowIdentity, SimulationArtifactDescriptor, @@ -80,6 +88,8 @@ ) from policyengine_simulation_executor.stage12_runtime import ( SimulationCalculation, + Stage12ExecutionError, + Stage12FailureDetail, Stage12SegmentRunner, _build_spm_result, build_aggregate_report, @@ -112,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) @@ -3919,3 +3930,481 @@ 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" + ) + + +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 + )