Let the v2 report coordinator return a typed result to its caller

libs/policyengine-simulation-contract/src/policyengine_simulation_contract/report_execution.py

@@ -0,0 +1,226 @@
+"""The typed outcome a report coordinator returns to its caller.
+
+A report coordinator dispatched with ``result_delivery`` set to
+``ReportResultDelivery.RETURN_TO_CALLER`` reads and writes no Stage 12 rows.
+When the report succeeds it returns a ``ReportExecutionResult``, which carries
+the report's aggregate inline, references to the baseline and reform simulation
+artifacts, and the software, data and Modal deployment that produced them.
+When the report fails it raises ``ReportExecutionError``, which carries the
+redacted error code and summary that a Stage 12 report row would record for
+the same failure.
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime
+from hashlib import sha256
+from typing import TYPE_CHECKING, Annotated, Any, Literal
+from uuid import UUID
+
+from pydantic import Field, field_validator, model_validator
+
+from policyengine_simulation_contract.macro_output import SingleYearMacroOutput
+from policyengine_simulation_contract.stage12_execution import (
+    BundleProvenance,
+    ContractText,
+    ErrorCode,
+    ErrorSummary,
+    ReportAggregate,
+    Sha256Digest,
+    SimulationArtifactDescriptor,
+    SimulationRole,
+    Stage12OutputPlan,
+    StrictContractModel,
+    stage12_output_plan_sha256,
+)
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+REPORT_EXECUTION_ERROR_MESSAGE = "Report execution failed"
+
+
+def report_aggregate_sha256(aggregate: SingleYearMacroOutput) -> str:
+    """Return the SHA-256 of the aggregate's canonical JSON encoding.
+
+    The encoding is compact, key-sorted UTF-8 JSON of the aggregate's JSON-mode
+    dump, so an aggregate has the same digest before and after a JSON round
+    trip. Raises ``ValueError`` when the aggregate contains NaN or infinity.
+    """
+    try:
+        payload = json.dumps(
+            aggregate.model_dump(mode="json"),
+            sort_keys=True,
+            separators=(",", ":"),
+            ensure_ascii=False,
+            allow_nan=False,
+        ).encode("utf-8")
+    except ValueError:
+        msg = "aggregate contains a non-finite number"
+        raise ValueError(msg) from None
+    return sha256(payload).hexdigest()
+
+
+class ReportExecutionProvenance(StrictContractModel):
+    """The software, data and Modal deployment that produced one report.
+
+    ``bundle`` names the PolicyEngine.py version, the country model package,
+    the dataset and the bundle manifest they came from. The other fields name
+    the Modal environment, application and simulation function that ran, the
+    digest of the version manifest they were resolved from, and the Modal
+    function call ID of the report coordinator.
+    """
+
+    bundle: BundleProvenance
+    worker_version: ContractText
+    modal_environment: ContractText
+    modal_application: ContractText
+    simulation_callable: ContractText
+    version_manifest_sha256: Sha256Digest
+    coordinator_invocation_id: ContractText
+
+
+class ReportExecutionResult(StrictContractModel):
+    """A completed report, as its coordinator returns it to the caller.
+
+    ``aggregate`` is the report's single-year macro result: its totals, inline.
+    ``aggregate_sha256`` is the digest ``report_aggregate_sha256`` gives for it.
+    ``baseline`` and ``reform`` reference the two simulation artifacts the
+    aggregate was computed from, and ``output_plan`` declares the columns they
+    hold. Both artifacts must belong to this report, carry the baseline and
+    reform roles, follow the output plan, cover the same rows and use the
+    bundle named in ``provenance``. ``requested_aggregates`` must be the
+    aggregates the output plan was built for, in the same order.
+    """
+
+    contract_version: Literal[1] = 1
+    evaluation_id: UUID
+    requested_aggregates: Annotated[tuple[ReportAggregate, ...], Field(min_length=1)]
+    aggregate_schema_version: Literal[1] = 1
+    aggregate: SingleYearMacroOutput
+    aggregate_sha256: Sha256Digest
+    output_plan: Stage12OutputPlan
+    baseline: SimulationArtifactDescriptor
+    reform: SimulationArtifactDescriptor
+    provenance: ReportExecutionProvenance
+    completed_at: datetime
+
+    @field_validator("requested_aggregates")
+    @classmethod
+    def require_unique_aggregates(
+        cls,
+        value: tuple[ReportAggregate, ...],
+    ) -> tuple[ReportAggregate, ...]:
+        """Reject repeated aggregate names."""
+        if len(value) != len(set(value)):
+            msg = "requested aggregates must be unique"
+            raise ValueError(msg)
+        return value
+
+    @field_validator("completed_at")
+    @classmethod
+    def require_aware_completion_time(cls, value: datetime) -> datetime:
+        """Reject a completion time without a timezone."""
+        if value.tzinfo is None:
+            msg = "completed_at must include a timezone"
+            raise ValueError(msg)
+        return value
+
+    @model_validator(mode="after")
+    def require_aggregates_of_the_output_plan(self) -> ReportExecutionResult:
+        """Require the requested aggregates the output plan was built for.
+
+        The coordinator builds the output plan from the report's requested
+        aggregates, so a result naming others contradicts its own plan.
+        """
+        if self.requested_aggregates != self.output_plan.requirements.aggregates:
+            msg = "requested aggregates must match the output plan's aggregates"
+            raise ValueError(msg)
+        return self
+
+    @model_validator(mode="after")
+    def require_artifacts_of_this_report(self) -> ReportExecutionResult:
+        """Require a baseline and a reform artifact that form this report."""
+        if self.baseline.role is not SimulationRole.BASELINE:
+            msg = "the baseline artifact must use the baseline role"
+            raise ValueError(msg)
+        if self.reform.role is not SimulationRole.REFORM:
+            msg = "the reform artifact must use the reform role"
+            raise ValueError(msg)
+        for descriptor in (self.baseline, self.reform):
+            self._require_artifact_of_this_report(descriptor)
+        if self.baseline.simulation_execution_id == self.reform.simulation_execution_id:
+            msg = (
+                "baseline and reform artifacts must have different simulation "
+                "identifiers"
+            )
+            raise ValueError(msg)
+        if self.baseline.row_identity != self.reform.row_identity:
+            msg = "baseline and reform artifacts must cover the same rows"
+            raise ValueError(msg)
+        return self
+
+    def _require_artifact_of_this_report(
+        self,
+        descriptor: SimulationArtifactDescriptor,
+    ) -> None:
+        """Require one artifact to belong to this report, bundle and output plan."""
+        if descriptor.evaluation_id != self.evaluation_id:
+            msg = "simulation artifacts must belong to this report"
+            raise ValueError(msg)
+        if descriptor.bundle != self.provenance.bundle:
+            msg = "simulation artifacts must use the report's bundle"
+            raise ValueError(msg)
+        if descriptor.output_plan_sha256 != stage12_output_plan_sha256(
+            self.output_plan
+        ):
+            msg = "simulation artifacts must follow the report's output plan"
+            raise ValueError(msg)
+
+    @model_validator(mode="after")
+    def require_the_aggregate_digest(self) -> ReportExecutionResult:
+        """Require ``aggregate_sha256`` to be the digest of ``aggregate``."""
+        if report_aggregate_sha256(self.aggregate) != self.aggregate_sha256:
+            msg = "aggregate_sha256 does not match the aggregate"
+            raise ValueError(msg)
+        return self
+
+
+class ReportExecutionFailure(StrictContractModel):
+    """A failed report or simulation: its redacted error code and summary.
+
+    The values are the ones a Stage 12 row records for the same failure, which
+    are safe to store and to show to the caller.
+    """
+
+    contract_version: Literal[1] = 1
+    evaluation_id: UUID
+    error_code: ErrorCode
+    error_summary: ErrorSummary
+
+
+class ReportExecutionError(RuntimeError):
+    """Raised by a coordinator or simulation whose result goes to its caller.
+
+    The message is fixed; the failure is in ``failure``. The error pickles
+    with its failure as plain JSON values, so Modal's ``FunctionCall.get``
+    raises it in the caller as this type with the same failure, provided the
+    caller can import this module, whatever pydantic version each side runs.
+    """
+
+    def __init__(self, failure: ReportExecutionFailure) -> None:
+        """Carry ``failure`` under the fixed message."""
+        self.failure = failure
+        super().__init__(REPORT_EXECUTION_ERROR_MESSAGE)
+
+    def __reduce__(
+        self,
+    ) -> tuple[Callable[[dict[str, Any]], ReportExecutionError], tuple[dict[str, Any]]]:
+        """Rebuild the error from its failure's JSON values when it is unpickled."""
+        return (_report_execution_error, (self.failure.model_dump(mode="json"),))
+
+
+def _report_execution_error(failure: dict[str, Any]) -> ReportExecutionError:
+    """Rebuild a pickled ``ReportExecutionError`` from its failure's JSON values."""
+    return ReportExecutionError(ReportExecutionFailure.model_validate(failure))

libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_execution.py

@@ -587,13 +587,29 @@ def validate_differences(self) -> ResultComparisonArtifactPayload:
         return self
 
 
+class ReportResultDelivery(StrEnum):
+    """How a report coordinator and its simulations deliver their outcome.
+
+    ``STAGE12_DATABASE``, the default, records the report's and each
+    simulation's state in the Stage 12 tables, and the coordinator returns an
+    aggregate artifact descriptor. ``RETURN_TO_CALLER`` reads and writes no
+    Stage 12 rows: the coordinator returns a ``ReportExecutionResult`` (see
+    ``report_execution``) or raises ``ReportExecutionError``.
+    """
+
+    STAGE12_DATABASE = "stage12_database"
+    RETURN_TO_CALLER = "return_to_caller"
+
+
 class Stage12InvocationContext(StrictContractModel):
     """Deployment and retention context for one comparison-run invocation.
 
     Records the originating request, the environments, the worker version and
     callable, the manifest digests and the artifact prefix. Automatic runs also
     name the production function call they are compared with. Artifacts are kept
-    for at most 30 days after ``created_at``.
+    for at most 30 days after ``created_at``. ``result_delivery`` selects how the
+    coordinator and its simulations deliver their outcome; a report returned to
+    its caller cannot be compared with a production job.
     """
 
     contract_version: Literal[1] = 1
@@ -615,13 +631,33 @@ class Stage12InvocationContext(StrictContractModel):
     production_function_call_id: ContractText | None = None
     created_at: datetime
     retention_expires_at: datetime
+    # Callers that leave it out, such as the Simulation Entrypoint today, keep
+    # the Stage 12 database behavior.
+    result_delivery: ReportResultDelivery = ReportResultDelivery.STAGE12_DATABASE
 
     @model_validator(mode="after")
     def validate_retention(self) -> Stage12InvocationContext:
         """Require aware timestamps and a retention window of at most 30 days."""
         _validate_retention(self.created_at, self.retention_expires_at)
         return self
 
+    @model_validator(mode="after")
+    def validate_result_delivery(self) -> Stage12InvocationContext:
+        """Reject a production comparison for a report returned to its caller.
+
+        The comparison is recorded on the Stage 12 report row, which a
+        returned report never writes.
+        """
+        if (
+            self.result_delivery is ReportResultDelivery.RETURN_TO_CALLER
+            and self.production_function_call_id is not None
+        ):
+            msg = (
+                "a report returned to its caller cannot request a production comparison"
+            )
+            raise ValueError(msg)
+        return self
+
 
 ErrorCode = Annotated[str, Field(min_length=1, max_length=64)]
 ErrorSummary = Annotated[str, Field(min_length=1, max_length=512)]