Materialize Stage 12 report output variables

docs/migration/stage-12-modal-worker-foundation.md

@@ -165,6 +165,43 @@ identifies an artifact declared by that same bundle. The request model converts
 the Public API's `_telemetry` field to `telemetry`; the adapter accepts that
 correlation metadata but does not include it in either simulation input.
 
+## Report output planning
+
+Before starting either child simulation, the Modal report coordinator resolves
+one immutable, strongly typed output plan for the complete Stage 12 report. The
+plan records the requested aggregate profile, whether cliff analysis was
+requested, whether either policy activates labor-supply responses, and the
+required columns for every country-model entity.
+
+The resolver constructs data-free planning `Simulation` objects and invokes
+the output-configuration functions supplied by the pinned PolicyEngine 5.2.0
+bundle. For US reports this includes the budgetary-impact variables; both
+countries use PolicyEngine's conditional cliff and labor-supply configuration.
+Planning never loads a dataset or executes a simulation. The former
+`variables=("*",)` internal marker is not an output expansion mechanism and is
+no longer used.
+
+The coordinator includes the exact same output plan in the baseline and reform
+child inputs. Each child adds the plan's additional variables before calling
+`Simulation.ensure()`, then verifies that every planned entity and column is
+present before writing its Parquet artifact. The artifact descriptor records a
+digest of the plan. After both calls complete, the coordinator checks both plan
+digests and independently validates both loaded Parquet schemas before it runs
+aggregate calculations. Extra columns are permitted; missing planned columns
+fail the temporary Stage 12 report without affecting the production result.
+
+The plan distinguishes calculated variables from columns copied directly from
+the source dataset. UK geographic reports require the household columns
+`constituency_code_oa` and `la_code_oa`. Workers do not request those columns as
+calculated PolicyEngine variables, but they must be present in each materialized
+artifact; otherwise validation fails before geographic aggregation begins.
+
+Aggregation uses the retained output tables without rerunning either model.
+The precomputed baseline and reform simulations retain their original policy
+metadata so PolicyEngine can correctly recognize conditional labor-supply
+analysis. `include_cliffs=true` is supported by this path and is carried through
+both output planning and aggregation.
+
 ## Temporary direct runner endpoint
 
 > **Temporary Stage 12 interface:** The authenticated routes in this section

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

@@ -2,9 +2,11 @@
 
 from __future__ import annotations
 
+import json
 from dataclasses import dataclass
 from datetime import datetime, timedelta
 from enum import StrEnum
+from hashlib import sha256
 from typing import Annotated, Literal
 from uuid import UUID
 
@@ -155,18 +157,6 @@ def require_complete_filter(self) -> GeographySelection:
         return self
 
 
-class RequestedSimulationOutput(StrictContractModel):
-    schema_version: Literal[1] = 1
-    variables: Annotated[tuple[ContractText, ...], Field(min_length=1)]
-
-    @field_validator("variables")
-    @classmethod
-    def require_unique_variables(cls, value: tuple[str, ...]) -> tuple[str, ...]:
-        if len(value) != len(set(value)):
-            raise ValueError("requested variables must be unique")
-        return value
-
-
 class SimulationExecutionInput(StrictContractModel):
     contract_version: Literal[1] = 1
     evaluation_id: UUID
@@ -177,7 +167,6 @@ class SimulationExecutionInput(StrictContractModel):
     year: Annotated[int, Field(ge=1900, le=2200)]
     geography: GeographySelection
     options: dict[str, JsonValue] = Field(default_factory=dict)
-    requested_output: RequestedSimulationOutput
     bundle: BundleProvenance
 
     @model_validator(mode="before")
@@ -190,6 +179,127 @@ def reject_combined_policy_input(cls, value: object) -> object:
         return value
 
 
+class ReportAggregate(StrEnum):
+    BUDGET = "budget"
+    POVERTY = "poverty"
+    INEQUALITY = "inequality"
+    DISTRIBUTIONAL = "distributional"
+    WINNERS_AND_LOSERS = "winners_and_losers"
+    GEOGRAPHIC = "geographic"
+    PROGRAM_STATISTICS = "program_statistics"
+
+
+class ReportOutputRequirements(StrictContractModel):
+    """Report features that determine which simulation columns must exist."""
+
+    schema_version: Literal[1] = 1
+    aggregates: Annotated[tuple[ReportAggregate, ...], Field(min_length=1)]
+    include_cliff_impacts: bool
+    labor_supply_response_active: bool
+
+    @field_validator("aggregates")
+    @classmethod
+    def require_complete_aggregate_profile(
+        cls,
+        value: tuple[ReportAggregate, ...],
+    ) -> tuple[ReportAggregate, ...]:
+        if len(value) != len(set(value)):
+            raise ValueError("report output aggregates must be unique")
+        if value != tuple(ReportAggregate):
+            raise ValueError(
+                "Stage 12 currently requires the complete aggregate profile"
+            )
+        return value
+
+
+class EntityOutputPlan(StrictContractModel):
+    """Required materialized columns for one country-model entity."""
+
+    entity: ContractText
+    materialized_variables: Annotated[
+        tuple[ContractText, ...],
+        Field(min_length=1),
+    ]
+    additional_variables: tuple[ContractText, ...] = ()
+    dataset_variables: tuple[ContractText, ...] = ()
+
+    @field_validator(
+        "materialized_variables",
+        "additional_variables",
+        "dataset_variables",
+    )
+    @classmethod
+    def require_canonical_variables(cls, value: tuple[str, ...]) -> tuple[str, ...]:
+        if len(value) != len(set(value)):
+            raise ValueError("output-plan variables must be unique")
+        if value != tuple(sorted(value)):
+            raise ValueError("output-plan variables must use canonical order")
+        return value
+
+    @model_validator(mode="after")
+    def require_variable_subsets(self) -> EntityOutputPlan:
+        if not set(self.additional_variables).issubset(self.materialized_variables):
+            raise ValueError(
+                "additional output-plan variables must be a materialized subset"
+            )
+        if not set(self.dataset_variables).issubset(self.materialized_variables):
+            raise ValueError(
+                "dataset output-plan variables must be a materialized subset"
+            )
+        if set(self.additional_variables).intersection(self.dataset_variables):
+            raise ValueError(
+                "calculated additional variables and dataset variables must be disjoint"
+            )
+        return self
+
+
+class Stage12OutputPlan(StrictContractModel):
+    """One immutable output schema shared by both report simulations."""
+
+    schema_version: Literal[1] = 1
+    country: CountryId
+    requirements: ReportOutputRequirements
+    entities: Annotated[tuple[EntityOutputPlan, ...], Field(min_length=1)]
+
+    @field_validator("entities")
+    @classmethod
+    def require_canonical_entities(
+        cls,
+        value: tuple[EntityOutputPlan, ...],
+    ) -> tuple[EntityOutputPlan, ...]:
+        names = tuple(entity.entity for entity in value)
+        if len(names) != len(set(names)):
+            raise ValueError("output-plan entities must be unique")
+        if names != tuple(sorted(names)):
+            raise ValueError("output-plan entities must use canonical order")
+        return value
+
+
+def stage12_output_plan_sha256(plan: Stage12OutputPlan) -> str:
+    """Return the stable digest recorded by each Stage 12 child artifact."""
+
+    payload = json.dumps(
+        plan.model_dump(mode="json"),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+        allow_nan=False,
+    ).encode("utf-8")
+    return sha256(payload).hexdigest()
+
+
+class PlannedSimulationExecutionInput(SimulationExecutionInput):
+    """Internal coordinator-to-worker input with a resolved output schema."""
+
+    output_plan: Stage12OutputPlan
+
+    @model_validator(mode="after")
+    def require_output_plan_country(self) -> PlannedSimulationExecutionInput:
+        if self.output_plan.country != self.geography.country:
+            raise ValueError("output plan country must match simulation geography")
+        return self
+
+
 class RowIdentity(StrictContractModel):
     schema_version: Literal[1] = 1
     identifier_columns: Annotated[tuple[ContractText, ...], Field(min_length=1)]
@@ -211,21 +321,12 @@ class SimulationArtifactDescriptor(StrictContractModel):
     role: SimulationRole
     artifact: ArtifactReference
     output_schema_version: Literal[1] = 1
+    output_plan_sha256: Sha256Digest
     row_identity: RowIdentity
     bundle: BundleProvenance
     calculation_provenance: dict[str, JsonValue] | None = None
 
 
-class ReportAggregate(StrEnum):
-    BUDGET = "budget"
-    POVERTY = "poverty"
-    INEQUALITY = "inequality"
-    DISTRIBUTIONAL = "distributional"
-    WINNERS_AND_LOSERS = "winners_and_losers"
-    GEOGRAPHIC = "geographic"
-    PROGRAM_STATISTICS = "program_statistics"
-
-
 class ReportExecutionInput(StrictContractModel):
     contract_version: Literal[1] = 1
     evaluation_id: UUID
@@ -250,7 +351,6 @@ def require_aligned_simulations(self) -> ReportExecutionInput:
             "year",
             "geography",
             "options",
-            "requested_output",
             "bundle",
         ):
             if getattr(self.baseline, field_name) != getattr(self.reform, field_name):

libs/policyengine-simulation-contract/tests/test_stage12_execution.py

@@ -0,0 +1,117 @@
+"""Tests for typed Stage 12 execution contracts."""
+
+from __future__ import annotations
+
+import pytest
+
+from policyengine_simulation_contract.stage12_execution import (
+    EntityOutputPlan,
+    ReportAggregate,
+    ReportOutputRequirements,
+    Stage12OutputPlan,
+    stage12_output_plan_sha256,
+)
+
+
+def _plan() -> Stage12OutputPlan:
+    return Stage12OutputPlan(
+        country="us",
+        requirements=ReportOutputRequirements(
+            aggregates=tuple(ReportAggregate),
+            include_cliff_impacts=False,
+            labor_supply_response_active=False,
+        ),
+        entities=(
+            EntityOutputPlan(
+                entity="household",
+                materialized_variables=("household_id", "household_net_income"),
+                additional_variables=(),
+            ),
+            EntityOutputPlan(
+                entity="person",
+                materialized_variables=("federal_benefit_cost", "person_id"),
+                additional_variables=("federal_benefit_cost",),
+            ),
+        ),
+    )
+
+
+def test_output_plan_has_a_deterministic_digest() -> None:
+    plan = _plan()
+
+    assert stage12_output_plan_sha256(plan) == stage12_output_plan_sha256(
+        Stage12OutputPlan.model_validate(plan.model_dump(mode="json"))
+    )
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "message"),
+    [
+        (
+            "materialized_variables",
+            ("person_id", "federal_benefit_cost"),
+            "canonical order",
+        ),
+        (
+            "materialized_variables",
+            ("person_id", "person_id"),
+            "unique",
+        ),
+        (
+            "additional_variables",
+            ("missing",),
+            "subset",
+        ),
+        (
+            "dataset_variables",
+            ("missing",),
+            "subset",
+        ),
+    ],
+)
+def test_entity_output_plan_rejects_noncanonical_variables(
+    field: str,
+    value: tuple[str, ...],
+    message: str,
+) -> None:
+    values = {
+        "entity": "person",
+        "materialized_variables": ("federal_benefit_cost", "person_id"),
+        "additional_variables": ("federal_benefit_cost",),
+        "dataset_variables": (),
+    }
+    values[field] = value
+
+    with pytest.raises(ValueError, match=message):
+        EntityOutputPlan.model_validate(values)
+
+
+def test_output_plan_rejects_duplicate_or_unsorted_entities() -> None:
+    entity = _plan().entities[0]
+
+    with pytest.raises(ValueError, match="canonical order"):
+        Stage12OutputPlan(
+            country="us",
+            requirements=_plan().requirements,
+            entities=(
+                _plan().entities[1],
+                entity,
+            ),
+        )
+
+    with pytest.raises(ValueError, match="unique"):
+        Stage12OutputPlan(
+            country="us",
+            requirements=_plan().requirements,
+            entities=(entity, entity),
+        )
+
+
+def test_entity_output_plan_separates_calculated_and_dataset_variables() -> None:
+    with pytest.raises(ValueError, match="must be disjoint"):
+        EntityOutputPlan(
+            entity="household",
+            materialized_variables=("constituency_code_oa", "household_id"),
+            additional_variables=("constituency_code_oa",),
+            dataset_variables=("constituency_code_oa",),
+        )