From 264a7f556bda37c79ce16be1e5ab51c5a5e73fc9 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:37:12 +0400 Subject: [PATCH 1/4] Add versioned household calculation contracts Add the wire models for v2 household calculations to the shared contract library: - HouseholdDocument keeps the situation shape the Public API's household routes accept (people, group collections with members, variable values by period, optional axes), so results can be compared field by field. - HouseholdExecutionRequest names one or more policies, so a baseline and its reform run in one invocation, plus the expected PolicyEngine.py bundle, the calculate-full option and the SPM selection. An omitted SPM selection inherits the bundle's defaults and survives serialization; an explicit null is rejected. - HouseholdExecutionResult and HouseholdExecutionFailure, told apart by status, carry per-policy calculated households, warnings and SPM receipts, or a typed error with the Public API's input-error fields. - SubmitHouseholdCalculationRequest is the public submission, and the pending, running, succeeded and failed poll responses form a discriminated union. Nothing uses these models yet. Co-Authored-By: Claude Opus 5.5 --- .../household_execution.py | 684 ++++++++++++++++++ .../tests/test_household_execution.py | 422 +++++++++++ 2 files changed, 1106 insertions(+) create mode 100644 libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py create mode 100644 libs/policyengine-simulation-contract/tests/test_household_execution.py diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py new file mode 100644 index 000000000..2f6bac4b0 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py @@ -0,0 +1,684 @@ +"""Versioned household calculation contracts for the v2 Modal household functions. + +A household travels in the situation shape the Public API's household +calculation routes accept and return: ``people`` maps each person to their +variables, every other top-level key is a group entity collection (such as +``households`` or ``tax_units``) whose groups list their ``members``, each +variable maps periods to values, and an optional ``axes`` list varies inputs +across points. A calculated household keeps that shape, with each requested +value filled in, so a v2 result can be compared field by field with the +Public API's in-process calculation. + +These models check structure, types and size only. Whether an entity +collection, variable or policy parameter exists is decided by the country +model inside the household function, which reports unknown names with the +same typed input errors as the Public API. + +Wire models: + +- ``HouseholdExecutionRequest``: Simulation Entrypoint to a Modal household + function. It names one or more policies, so a baseline and its reform run in + one invocation. +- ``HouseholdExecutionResult`` or ``HouseholdExecutionFailure``: what the + household function returns, told apart by ``status``. +- ``SubmitHouseholdCalculationRequest``: the public submission to the + Simulation Entrypoint. +- The poll responses ``PendingHouseholdCalculation``, + ``RunningHouseholdCalculation``, ``SucceededHouseholdCalculation`` and + ``FailedHouseholdCalculation``, told apart by ``status``. +""" + +from __future__ import annotations + +import math +from typing import TYPE_CHECKING, Annotated, Any, Literal, Self +from uuid import UUID + +from pydantic import ( + AfterValidator, + BaseModel, + BeforeValidator, + ConfigDict, + Field, + TypeAdapter, + field_validator, + model_serializer, + model_validator, +) + +from policyengine_simulation_contract.spm import ( + SPMErrorDetail, + SPMProvenance, + SPMSelection, +) + +if TYPE_CHECKING: + from pydantic import SerializerFunctionWrapHandler + +HOUSEHOLD_CONTRACT_SCHEMA_VERSION = 1 +# Size limits. The entity and value limits match the Public API's stored v2 +# households; the variable limit leaves room for a complete calculated +# household sent back as input, which holds every variable of each entity. +MAXIMUM_PEOPLE = 1_000 +MAXIMUM_GROUPS_PER_COLLECTION = 1_000 +MAXIMUM_GROUP_COLLECTIONS = 32 +MAXIMUM_VARIABLES_PER_ENTITY = 10_000 +MAXIMUM_PERIODS_PER_VARIABLE = 1_000 +MAXIMUM_AXIS_GROUPS = 8 +MAXIMUM_PARALLEL_AXES = 32 +# The number of points a household is calculated at: the product of the point +# counts of its perpendicular axis groups. +MAXIMUM_AXIS_POINTS = 10_000 +MAXIMUM_POLICIES_PER_REQUEST = 10 +MAXIMUM_POLICY_PARAMETERS = 10_000 +MAXIMUM_WARNINGS = 10_000 + +EntityKey = Annotated[str, Field(min_length=1, max_length=255)] +EntityPlural = Annotated[str, Field(min_length=1, max_length=255)] +VariableName = Annotated[str, Field(min_length=1, max_length=255)] +PeriodKey = Annotated[str, Field(min_length=1, max_length=64)] +ParameterPath = Annotated[str, Field(min_length=1, max_length=512)] +ParameterPeriod = Annotated[str, Field(min_length=1, max_length=64)] +ContractText = Annotated[str, Field(min_length=1, max_length=255)] +PolicyKey = Annotated[ + str, + Field( + pattern=r"^[a-z][a-z0-9_-]{0,63}$", + description=( + "Caller-chosen name for one policy in the request, such as " + "'baseline' or 'reform'; each result carries the same name." + ), + ), +] +Sha256Digest = Annotated[str, Field(pattern=r"^[0-9a-f]{64}$")] +CountryId = Annotated[ + Literal["us", "uk"], + Field(description="Supported PolicyEngine country ID"), +] + +# One input value. ``None`` asks the household function to calculate it. +InputScalar = bool | int | float | str | None +PeriodInputs = Annotated[ + dict[PeriodKey, InputScalar], + Field(max_length=MAXIMUM_PERIODS_PER_VARIABLE), +] +# A variable's values by period, or one value for the model's default period. +VariableInput = PeriodInputs | InputScalar +EntityInputs = Annotated[ + dict[VariableName, VariableInput], + Field(max_length=MAXIMUM_VARIABLES_PER_ENTITY), +] +# One calculated value. Households with axes hold one value per point. +# Infinite and undefined numbers are the strings "Infinity", "-Infinity" and +# "NaN", as the Public API serializes infinities. +OutputScalar = bool | int | float | str | None +OutputValue = OutputScalar | list[OutputScalar] +PeriodOutputs = Annotated[ + dict[PeriodKey, OutputValue], + Field(max_length=MAXIMUM_PERIODS_PER_VARIABLE), +] +VariableOutput = PeriodOutputs | OutputScalar +EntityOutputs = Annotated[ + dict[VariableName, VariableOutput], + Field(max_length=MAXIMUM_VARIABLES_PER_ENTITY), +] +ParameterValue = bool | int | float | str | None + +HouseholdErrorCode = Literal[ + # A household entity collection, variable, axis variable or policy + # parameter the country model does not have. ``input_errors`` lists them. + "invalid_inputs", + # The country model could not build a simulation from the household. + "invalid_household", + # The SPM selection, or the SPM configuration it needs, is unusable. + # ``spm_error`` carries the SPM error code. + "invalid_spm_selection", + # The request names a bundle other than the one this function runs. + "bundle_mismatch", + # Any other failure while calculating. + "calculation_failed", +] +InputErrorType = Literal[ + "household_entity", + "household_variable", + "household_variable_wrong_entity", + "household_axis_variable", + "policy_parameter", +] + + +class StrictHouseholdModel(BaseModel): + """Base for household contracts; rejects unknown fields and non-finite numbers.""" + + model_config = ConfigDict(extra="forbid", allow_inf_nan=False, frozen=True) + + +def _require_unique_members(members: tuple[str, ...]) -> tuple[str, ...]: + if len(set(members)) != len(members): + msg = "group members must be unique" + raise ValueError(msg) + return members + + +GroupMembers = Annotated[ + tuple[EntityKey, ...], + Field(min_length=1, max_length=MAXIMUM_PEOPLE), + AfterValidator(_require_unique_members), +] + + +class HouseholdGroup(BaseModel): + """One group entity instance, such as a household or a tax unit. + + ``members`` names the people in the group. Every other key is a variable + mapped to its values by period, exactly as in the Public API's situation + shape. + """ + + model_config = ConfigDict(extra="allow", allow_inf_nan=False, frozen=True) + + members: GroupMembers + __pydantic_extra__: dict[VariableName, VariableInput] = Field(init=False) + + @model_validator(mode="after") + def limit_variables(self) -> Self: + """Bound the number of variables one group may carry.""" + if len(self.model_extra or {}) > MAXIMUM_VARIABLES_PER_ENTITY: + msg = "a group carries too many variables" + raise ValueError(msg) + return self + + +class CalculatedHouseholdGroup(BaseModel): + """One calculated group entity instance: its members and calculated variables.""" + + model_config = ConfigDict(extra="allow", allow_inf_nan=False, frozen=True) + + members: GroupMembers + __pydantic_extra__: dict[VariableName, VariableOutput] = Field(init=False) + + @model_validator(mode="after") + def limit_variables(self) -> Self: + """Bound the number of variables one group may carry.""" + if len(self.model_extra or {}) > MAXIMUM_VARIABLES_PER_ENTITY: + msg = "a group carries too many variables" + raise ValueError(msg) + return self + + +class HouseholdAxis(StrictHouseholdModel): + """One input varied across points, as policyengine-core reads an axis. + + The variable ``name`` takes ``count`` evenly spaced values from ``min`` to + ``max`` in ``period`` (the model's default period when omitted), for the + entity member at position ``index`` (the first when omitted). + """ + + name: VariableName + count: Annotated[int, Field(ge=1, le=MAXIMUM_AXIS_POINTS)] + min: float + max: float + period: PeriodKey | int | None = None + index: Annotated[int, Field(ge=0)] | None = None + + @model_serializer(mode="wrap") + def serialize_axis(self, handler: SerializerFunctionWrapHandler) -> dict[str, Any]: + """Leave out optional keys that were not given, as the Public API does.""" + dumped: dict[str, Any] = handler(self) + for name in ("period", "index"): + if dumped.get(name) is None and name not in self.model_fields_set: + dumped.pop(name, None) + return dumped + + +AxisGroup = Annotated[ + tuple[HouseholdAxis, ...], + Field(min_length=1, max_length=MAXIMUM_PARALLEL_AXES), +] + + +def _axis_point_count(axes: tuple[tuple[HouseholdAxis, ...], ...] | None) -> int: + if not axes: + return 1 + return math.prod(group[0].count for group in axes) + + +def _validated_axes( + axes: tuple[tuple[HouseholdAxis, ...], ...] | None, +) -> tuple[tuple[HouseholdAxis, ...], ...] | None: + if axes is None: + return None + for group in axes: + if len({axis.count for axis in group}) != 1: + msg = "parallel axes in one group must share one point count" + raise ValueError(msg) + if _axis_point_count(axes) > MAXIMUM_AXIS_POINTS: + msg = ( + "a household may be calculated at no more than " + f"{MAXIMUM_AXIS_POINTS} axis points" + ) + raise ValueError(msg) + return axes + + +Axes = Annotated[ + tuple[AxisGroup, ...] | None, + Field(max_length=MAXIMUM_AXIS_GROUPS), + AfterValidator(_validated_axes), +] + + +class _SituationModel(BaseModel): + """Situation-shaped household: ``people``, group collections and ``axes``.""" + + model_config = ConfigDict(extra="allow", allow_inf_nan=False, frozen=True) + + axes: Axes = None + + @model_validator(mode="after") + def limit_collections(self) -> Self: + """Bound the number of group collections a household may carry.""" + if len(self.model_extra or {}) > MAXIMUM_GROUP_COLLECTIONS: + msg = "a household carries too many group collections" + raise ValueError(msg) + return self + + @model_serializer(mode="wrap") + def serialize_situation( + self, handler: SerializerFunctionWrapHandler + ) -> dict[str, Any]: + """Leave out ``axes`` when there are none, so the shape matches the input.""" + dumped: dict[str, Any] = handler(self) + if self.axes is None: + dumped.pop("axes", None) + return dumped + + @property + def has_axes(self) -> bool: + """Whether the household is calculated at more than one set of inputs.""" + return self.axes is not None + + def axis_point_count(self) -> int: + """Return how many points the household is calculated at (1 without axes).""" + return _axis_point_count(self.axes) + + def to_situation(self) -> dict[str, Any]: + """Return the household as the JSON situation a country model reads.""" + return self.model_dump(mode="json") + + def group_collections(self) -> dict[str, Any]: + """Return each group entity collection by its plural, in JSON form.""" + situation = self.to_situation() + return { + plural: situation[plural] + for plural in situation + if plural not in {"people", "axes"} + } + + +class HouseholdDocument(_SituationModel): + """A household to calculate, in the Public API's situation shape. + + ``people`` maps each person's name to their variables. Every other key + except ``axes`` is a group entity collection, by its plural, mapping each + group's name to a ``HouseholdGroup``. A variable maps periods (such as + ``"2026"``) to values; ``None`` asks for the value to be calculated. + """ + + people: Annotated[ + dict[EntityKey, EntityInputs], Field(min_length=1, max_length=MAXIMUM_PEOPLE) + ] + __pydantic_extra__: dict[ + EntityPlural, + Annotated[ + dict[EntityKey, HouseholdGroup], + Field(max_length=MAXIMUM_GROUPS_PER_COLLECTION), + ], + ] = Field(init=False) + + +class CalculatedHousehold(_SituationModel): + """A calculated household, in the shape the Public API returns as ``result``. + + It holds the request's people, groups and axes with each requested value + filled in. With axes, a calculated variable holds one value per point. + """ + + people: Annotated[ + dict[EntityKey, EntityOutputs], Field(min_length=1, max_length=MAXIMUM_PEOPLE) + ] + __pydantic_extra__: dict[ + EntityPlural, + Annotated[ + dict[EntityKey, CalculatedHouseholdGroup], + Field(max_length=MAXIMUM_GROUPS_PER_COLLECTION), + ], + ] = Field(init=False) + + +class HouseholdPolicy(StrictHouseholdModel): + """One policy to calculate the household under. + + ``parameters`` maps each changed parameter path to its values by period + range (``"YYYY-MM-DD.YYYY-MM-DD"``), as in the Public API's policy JSON. + Empty parameters mean current law. + """ + + policy_key: PolicyKey + parameters: Annotated[ + dict[ParameterPath, dict[ParameterPeriod, ParameterValue]], + Field(max_length=MAXIMUM_POLICY_PARAMETERS), + ] + + +def _unique_policy_keys( + policies: tuple[HouseholdPolicy, ...], +) -> tuple[HouseholdPolicy, ...]: + keys = [policy.policy_key for policy in policies] + if len(set(keys)) != len(keys): + msg = "policy keys must be unique within a request" + raise ValueError(msg) + return policies + + +Policies = Annotated[ + tuple[HouseholdPolicy, ...], + Field(min_length=1, max_length=MAXIMUM_POLICIES_PER_REQUEST), + AfterValidator(_unique_policy_keys), +] + + +def _reject_explicit_null_spm(value: object) -> object: + # Validators do not run on the default, so only an explicit null gets here. + if value is None: + msg = "spm must be an object; omit it to inherit the certified defaults." + raise ValueError(msg) + return value + + +SPMChoice = Annotated[ + SPMSelection | None, + BeforeValidator(_reject_explicit_null_spm), +] + + +class HouseholdBundleProvenance(StrictHouseholdModel): + """The PolicyEngine.py bundle a household calculation runs with. + + Names the PolicyEngine.py and policyengine-core versions, the country model + package and version, and the SHA-256 of the normalized bundle manifest. + """ + + policyengine_version: ContractText + core_package_version: ContractText + country_package_name: ContractText + country_package_version: ContractText + bundle_manifest_sha256: Sha256Digest + + +class _HouseholdCalculationInputs(StrictHouseholdModel): + """Fields shared by the public submission and the execution request. + + ``add_missing_variables`` adds every yearly, monthly and eternal variable + the household's entities lack before calculating, as the Public API's + ``/calculate-full`` route does. ``spm`` selects SPM settings for a US + household: omit it to inherit the bundle's certified defaults; supplying + it, even as an empty object, is a choice, and then a missing SPM input is + an error rather than an unavailable value. ``spm_selection_source`` says + whether the selection came from this request or from a saved household, + which changes only how an SPM artifact the bundle no longer carries is + reported. + """ + + country: CountryId + household: HouseholdDocument + policies: Policies + add_missing_variables: bool = False + spm: SPMChoice = Field( + default=None, + json_schema_extra=lambda schema: schema.pop("default", None), + ) + spm_selection_source: Literal["request", "saved_household"] = "request" + + @model_serializer(mode="wrap") + def serialize_inputs( + self, handler: SerializerFunctionWrapHandler + ) -> dict[str, Any]: + """Leave out ``spm`` when no selection was made, so omission is kept.""" + dumped: dict[str, Any] = handler(self) + if self.spm is None: + dumped.pop("spm", None) + return dumped + + @property + def spm_requested(self) -> bool: + """Whether the caller chose SPM settings rather than inheriting the defaults.""" + return self.spm is not None + + +class HouseholdExecutionRequest(_HouseholdCalculationInputs): + """Simulation Entrypoint to Modal household function: calculate one household. + + ``calculation_id`` is the durable job identity. ``bundle`` is the bundle + the entrypoint resolved from the version manifest; the household function + refuses a request for any other bundle. + """ + + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + calculation_id: UUID + bundle: HouseholdBundleProvenance + + +class SubmitHouseholdCalculationRequest(_HouseholdCalculationInputs): + """Public submission of one household calculation to the Simulation Entrypoint. + + ``policyengine_version`` pins a PolicyEngine.py bundle; omit it to use the + version manifest's default. + """ + + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + policyengine_version: ContractText | None = None + + def execution_request( + self, + *, + calculation_id: UUID, + bundle: HouseholdBundleProvenance, + ) -> HouseholdExecutionRequest: + """Return the request for the household function that runs ``bundle``. + + Raises ValueError when a pinned PolicyEngine.py version differs from + the bundle's. + """ + if ( + self.policyengine_version is not None + and self.policyengine_version != bundle.policyengine_version + ): + msg = "the pinned PolicyEngine.py version is not the resolved bundle's" + raise ValueError(msg) + inputs = self.model_dump( + mode="json", + exclude={"schema_version", "policyengine_version"}, + ) + return HouseholdExecutionRequest.model_validate( + { + **inputs, + "calculation_id": calculation_id, + "bundle": bundle.model_dump(mode="json"), + } + ) + + +class HouseholdPolicyResult(StrictHouseholdModel): + """The calculated household under one policy. + + ``warnings`` lists, in order, the deprecated inputs that were ignored and + the values that could not be calculated, worded as the Public API words + them. ``spm_provenance`` is the SPM calculation receipt for a US household + calculated with SPM settings. + """ + + policy_key: PolicyKey + household: CalculatedHousehold + warnings: Annotated[tuple[str, ...], Field(max_length=MAXIMUM_WARNINGS)] = () + spm_provenance: SPMProvenance | None = None + + +class HouseholdCalculationOutput(StrictHouseholdModel): + """A finished household calculation: one result per policy, in request order.""" + + country: CountryId + results: Annotated[ + tuple[HouseholdPolicyResult, ...], + Field(min_length=1, max_length=MAXIMUM_POLICIES_PER_REQUEST), + ] + bundle: HouseholdBundleProvenance + + @field_validator("results") + @classmethod + def require_unique_policy_keys( + cls, results: tuple[HouseholdPolicyResult, ...] + ) -> tuple[HouseholdPolicyResult, ...]: + """Reject results that repeat a policy key.""" + keys = [result.policy_key for result in results] + if len(set(keys)) != len(keys): + msg = "result policy keys must be unique" + raise ValueError(msg) + return results + + +class HouseholdInputError(StrictHouseholdModel): + """One household or policy name the country model does not recognize. + + Field names and messages match the entries of the Public API's ``errors`` + list. The entity plurals are set only for a variable placed on the wrong + entity. + """ + + type: InputErrorType + name: str + path: str + message: str + expected_entity_plural: str | None = None + actual_entity_plural: str | None = None + + +class HouseholdCalculationError(StrictHouseholdModel): + """Why a household calculation failed. + + ``message`` is worded as the Public API words the matching error response. + ``input_errors`` lists unrecognized names for ``invalid_inputs``; + ``spm_error`` holds the SPM error code for ``invalid_spm_selection``; + ``policy_key`` names the policy being prepared or calculated, when the + failure belongs to one. + """ + + code: HouseholdErrorCode + message: Annotated[str, Field(min_length=1)] + input_errors: tuple[HouseholdInputError, ...] = () + spm_error: SPMErrorDetail | None = None + policy_key: PolicyKey | None = None + + +class HouseholdWorkerIdentity(StrictHouseholdModel): + """The deployed Modal application and function that ran the calculation.""" + + application_name: ContractText + function_name: ContractText + + +class HouseholdExecutionResult(StrictHouseholdModel): + """Modal household function to Simulation Entrypoint: the calculation succeeded.""" + + status: Literal["succeeded"] = "succeeded" + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + calculation_id: UUID + output: HouseholdCalculationOutput + worker: HouseholdWorkerIdentity + + +class HouseholdExecutionFailure(StrictHouseholdModel): + """Modal household function to Simulation Entrypoint: a typed failure.""" + + status: Literal["failed"] = "failed" + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + calculation_id: UUID + error: HouseholdCalculationError + worker: HouseholdWorkerIdentity + + +HouseholdExecutionOutcome = Annotated[ + HouseholdExecutionResult | HouseholdExecutionFailure, + Field(discriminator="status"), +] +_EXECUTION_OUTCOME_ADAPTER: TypeAdapter[ + HouseholdExecutionResult | HouseholdExecutionFailure +] = TypeAdapter(HouseholdExecutionOutcome) + + +def parse_household_execution_outcome( + payload: object, +) -> HouseholdExecutionResult | HouseholdExecutionFailure: + """Validate a household function's return value as a success or a failure.""" + return _EXECUTION_OUTCOME_ADAPTER.validate_python(payload) + + +class PendingHouseholdCalculation(StrictHouseholdModel): + """Poll response: the calculation is accepted and not yet running.""" + + status: Literal["pending"] = "pending" + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + job_id: UUID + + +class RunningHouseholdCalculation(StrictHouseholdModel): + """Poll response: the household function is calculating.""" + + status: Literal["running"] = "running" + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + job_id: UUID + + +class SucceededHouseholdCalculation(StrictHouseholdModel): + """Poll response: the calculation finished; ``output`` holds the results.""" + + status: Literal["succeeded"] = "succeeded" + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + job_id: UUID + output: HouseholdCalculationOutput + + +class FailedHouseholdCalculation(StrictHouseholdModel): + """Poll response: the calculation failed; ``error`` says why.""" + + status: Literal["failed"] = "failed" + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + job_id: UUID + error: HouseholdCalculationError + + +HouseholdCalculationPollResponse = Annotated[ + PendingHouseholdCalculation + | RunningHouseholdCalculation + | SucceededHouseholdCalculation + | FailedHouseholdCalculation, + Field(discriminator="status"), +] +_POLL_RESPONSE_ADAPTER: TypeAdapter[ + PendingHouseholdCalculation + | RunningHouseholdCalculation + | SucceededHouseholdCalculation + | FailedHouseholdCalculation +] = TypeAdapter(HouseholdCalculationPollResponse) + + +def parse_household_calculation_poll_response( + payload: object, +) -> ( + PendingHouseholdCalculation + | RunningHouseholdCalculation + | SucceededHouseholdCalculation + | FailedHouseholdCalculation +): + """Validate a household calculation poll response by its ``status``.""" + return _POLL_RESPONSE_ADAPTER.validate_python(payload) diff --git a/libs/policyengine-simulation-contract/tests/test_household_execution.py b/libs/policyengine-simulation-contract/tests/test_household_execution.py new file mode 100644 index 000000000..ed98df376 --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_household_execution.py @@ -0,0 +1,422 @@ +"""Tests for the versioned household calculation contracts.""" + +from __future__ import annotations + +import json +from uuid import UUID + +import pytest +from policyengine_simulation_contract.household_execution import ( + MAXIMUM_AXIS_POINTS, + MAXIMUM_POLICIES_PER_REQUEST, + CalculatedHousehold, + FailedHouseholdCalculation, + HouseholdBundleProvenance, + HouseholdCalculationError, + HouseholdCalculationOutput, + HouseholdDocument, + HouseholdExecutionFailure, + HouseholdExecutionRequest, + HouseholdExecutionResult, + HouseholdInputError, + HouseholdPolicy, + HouseholdPolicyResult, + HouseholdWorkerIdentity, + PendingHouseholdCalculation, + RunningHouseholdCalculation, + SubmitHouseholdCalculationRequest, + SucceededHouseholdCalculation, + parse_household_calculation_poll_response, + parse_household_execution_outcome, +) +from policyengine_simulation_contract.spm import SPMErrorDetail +from pydantic import ValidationError + +CALCULATION_ID = UUID("00000000-0000-4000-8000-0000000000c1") +JOB_ID = UUID("00000000-0000-4000-8000-0000000000d2") + +# A household in the situation shape the Public API's calculate routes accept. +US_SITUATION = { + "people": { + "you": {"age": {"2026": 40}, "employment_income": {"2026": 50_000}}, + "your partner": {"age": {"2026": 38}}, + }, + "households": { + "your household": { + "members": ["you", "your partner"], + "state_name": {"2026": "CA"}, + "household_net_income": {"2026": None}, + } + }, + "tax_units": {"your tax unit": {"members": ["you", "your partner"]}}, + "axes": [ + [ + { + "name": "employment_income", + "count": 11, + "min": 0, + "max": 100_000, + "period": "2026", + } + ] + ], +} + + +def _bundle() -> HouseholdBundleProvenance: + return HouseholdBundleProvenance( + policyengine_version="6.7.5", + core_package_version="3.33.0", + country_package_name="policyengine-us", + country_package_version="2.2.1", + bundle_manifest_sha256="a" * 64, + ) + + +def _policies() -> tuple[HouseholdPolicy, ...]: + return ( + HouseholdPolicy(policy_key="baseline", parameters={}), + HouseholdPolicy( + policy_key="reform", + parameters={ + "gov.irs.credits.ctc.amount.base[0].amount": { + "2026-01-01.2100-12-31": 3_000 + } + }, + ), + ) + + +def _request(**overrides: object) -> HouseholdExecutionRequest: + values: dict[str, object] = { + "calculation_id": CALCULATION_ID, + "country": "us", + "household": US_SITUATION, + "policies": _policies(), + "bundle": _bundle(), + } + values.update(overrides) + return HouseholdExecutionRequest.model_validate(values) + + +def _worker() -> HouseholdWorkerIdentity: + return HouseholdWorkerIdentity( + application_name="policyengine-simulation-v2-py6-7-5", + function_name="calculate_household_us", + ) + + +def test_household_document_keeps_the_public_api_situation_shape() -> None: + document = HouseholdDocument.model_validate(US_SITUATION) + + assert document.to_situation() == US_SITUATION + assert document.group_collections() == { + "households": document.to_situation()["households"], + "tax_units": document.to_situation()["tax_units"], + } + assert document.has_axes + + +def test_household_without_axes_never_gains_an_axes_key() -> None: + situation = {key: value for key, value in US_SITUATION.items() if key != "axes"} + + document = HouseholdDocument.model_validate(situation) + + assert not document.has_axes + assert "axes" not in document.to_situation() + assert "axes" not in json.loads(document.model_dump_json()) + + +def test_household_document_accepts_values_without_periods() -> None: + situation = { + "people": {"you": {"age": 40, "is_blind": False}}, + "households": {"your household": {"members": ["you"]}}, + } + + assert HouseholdDocument.model_validate(situation).to_situation() == situation + + +@pytest.mark.parametrize( + "situation", + [ + {"households": {"your household": {"members": ["you"]}}}, + {"people": {}}, + {"people": {"you": {}}, "households": {"your household": {}}}, + {"people": {"you": {}}, "households": {"your household": {"members": []}}}, + {"people": {"you": {"age": {"2026": float("inf")}}}}, + {"people": {"you": {"age": {"2026": float("nan")}}}}, + {"people": {"you": {"age": {"2026": [40, 41]}}}}, + {"people": {"you": {}}, "axes": [[]]}, + {"people": {"you": {}}, "axes": [{"name": "age", "count": 2}]}, + ], +) +def test_household_document_rejects_malformed_situations(situation: dict) -> None: + with pytest.raises(ValidationError): + HouseholdDocument.model_validate(situation) + + +def test_axes_are_bounded_by_their_total_point_count() -> None: + axis = {"name": "employment_income", "min": 0, "max": 1, "period": "2026"} + within = { + "people": {"you": {}}, + "axes": [[{**axis, "count": 100}], [{**axis, "count": 100}]], + } + beyond = { + "people": {"you": {}}, + "axes": [ + [{**axis, "count": MAXIMUM_AXIS_POINTS}], + [{**axis, "count": 2}], + ], + } + + assert HouseholdDocument.model_validate(within).axis_point_count() == 10_000 + with pytest.raises(ValidationError, match="axis points"): + HouseholdDocument.model_validate(beyond) + + +def test_parallel_axes_must_share_one_point_count() -> None: + axis = {"name": "employment_income", "min": 0, "max": 1} + with pytest.raises(ValidationError, match="parallel axes"): + HouseholdDocument.model_validate( + { + "people": {"you": {}, "partner": {}}, + "axes": [[{**axis, "count": 3}, {**axis, "count": 4, "index": 1}]], + } + ) + + +def test_one_request_runs_baseline_and_reform_together() -> None: + request = _request() + + assert [policy.policy_key for policy in request.policies] == [ + "baseline", + "reform", + ] + assert request.schema_version == 1 + assert request.add_missing_variables is False + assert request.spm_selection_source == "request" + + +def test_request_needs_at_least_one_policy_with_unique_keys() -> None: + with pytest.raises(ValidationError): + _request(policies=()) + with pytest.raises(ValidationError, match="unique"): + _request( + policies=( + HouseholdPolicy(policy_key="reform", parameters={}), + HouseholdPolicy(policy_key="reform", parameters={}), + ) + ) + with pytest.raises(ValidationError): + _request( + policies=tuple( + HouseholdPolicy(policy_key=f"policy-{index}", parameters={}) + for index in range(MAXIMUM_POLICIES_PER_REQUEST + 1) + ) + ) + + +def test_omitted_spm_inherits_defaults_and_is_not_a_choice() -> None: + request = _request() + + assert request.spm is None + assert request.spm_requested is False + dumped = request.model_dump(mode="json") + assert "spm" not in dumped + assert HouseholdExecutionRequest.model_validate(dumped).spm_requested is False + + +def test_an_empty_spm_selection_is_still_a_choice() -> None: + request = _request(spm={}) + + assert request.spm_requested is True + dumped = request.model_dump(mode="json") + assert dumped["spm"] == {} + assert HouseholdExecutionRequest.model_validate(dumped).spm_requested is True + + +def test_explicit_null_spm_is_rejected() -> None: + with pytest.raises(ValidationError, match="omit it"): + _request(spm=None) + + +def test_request_rejects_unknown_fields_and_other_schema_versions() -> None: + with pytest.raises(ValidationError): + _request(unexpected=True) + with pytest.raises(ValidationError): + _request(schema_version=2) + + +def test_submit_request_becomes_an_execution_request_without_losing_choices() -> None: + submitted = SubmitHouseholdCalculationRequest.model_validate( + { + "country": "us", + "household": US_SITUATION, + "policies": [policy.model_dump(mode="json") for policy in _policies()], + "add_missing_variables": True, + "spm": {"geography_kind": "national"}, + "spm_selection_source": "saved_household", + "policyengine_version": "6.7.5", + } + ) + + request = submitted.execution_request( + calculation_id=CALCULATION_ID, bundle=_bundle() + ) + + assert request.calculation_id == CALCULATION_ID + assert request.household == submitted.household + assert request.policies == submitted.policies + assert request.add_missing_variables is True + assert request.spm is not None + assert request.spm.model_dump(mode="json") == {"geography_kind": "national"} + assert request.spm_selection_source == "saved_household" + assert request.bundle == _bundle() + + +def test_submit_request_rejects_a_pinned_version_the_bundle_does_not_carry() -> None: + submitted = SubmitHouseholdCalculationRequest( + country="us", + household=HouseholdDocument.model_validate(US_SITUATION), + policies=_policies(), + policyengine_version="6.0.0", + ) + + with pytest.raises(ValueError, match="pinned"): + submitted.execution_request(calculation_id=CALCULATION_ID, bundle=_bundle()) + + +def test_calculated_household_keeps_axis_results_and_non_finite_strings() -> None: + calculated = { + "people": { + "you": { + "age": {"2026": [40, 40]}, + "marginal_tax_rate": {"2026": ["Infinity", 0.25]}, + } + }, + "households": { + "your household": { + "members": ["you"], + "state_name": {"2026": ["CA", "CA"]}, + } + }, + "axes": [[{"name": "employment_income", "count": 2, "min": 0, "max": 1}]], + } + + household = CalculatedHousehold.model_validate(calculated) + + assert household.to_situation() == calculated + + +def _succeeded() -> HouseholdExecutionResult: + return HouseholdExecutionResult( + calculation_id=CALCULATION_ID, + output=HouseholdCalculationOutput( + country="us", + results=( + HouseholdPolicyResult( + policy_key="baseline", + household=CalculatedHousehold.model_validate( + { + "people": {"you": {"age": {"2026": 40}}}, + "households": {"your household": {"members": ["you"]}}, + } + ), + warnings=("Input `x` was ignored.",), + ), + ), + bundle=_bundle(), + ), + worker=_worker(), + ) + + +def _failed() -> HouseholdExecutionFailure: + return HouseholdExecutionFailure( + calculation_id=CALCULATION_ID, + error=HouseholdCalculationError( + code="invalid_inputs", + message="Unrecognized calculate input(s): Unrecognized household...", + input_errors=( + HouseholdInputError( + type="household_variable", + name="not_a_variable", + path="household.people.you.not_a_variable", + message="Unrecognized household variable ...", + ), + ), + ), + worker=_worker(), + ) + + +def test_execution_outcomes_are_discriminated_by_status() -> None: + succeeded = parse_household_execution_outcome(_succeeded().model_dump(mode="json")) + failed = parse_household_execution_outcome(_failed().model_dump(mode="json")) + + assert isinstance(succeeded, HouseholdExecutionResult) + assert succeeded == _succeeded() + assert isinstance(failed, HouseholdExecutionFailure) + assert failed == _failed() + with pytest.raises(ValidationError): + parse_household_execution_outcome({"status": "running"}) + + +def test_input_errors_serialize_with_the_public_api_field_names() -> None: + error = HouseholdInputError( + type="household_variable_wrong_entity", + name="state_name", + path="household.people.you.state_name", + message="Household variable `state_name` belongs on `households` ...", + expected_entity_plural="households", + actual_entity_plural="people", + ) + + assert error.model_dump(mode="json", exclude_none=True) == { + "type": "household_variable_wrong_entity", + "name": "state_name", + "path": "household.people.you.state_name", + "message": "Household variable `state_name` belongs on `households` ...", + "expected_entity_plural": "households", + "actual_entity_plural": "people", + } + + +def test_result_policy_keys_must_be_unique() -> None: + result = _succeeded().output.results[0] + + with pytest.raises(ValidationError, match="unique"): + HouseholdCalculationOutput( + country="us", results=(result, result), bundle=_bundle() + ) + + +def test_poll_responses_are_discriminated_by_status() -> None: + responses = [ + PendingHouseholdCalculation(job_id=JOB_ID), + RunningHouseholdCalculation(job_id=JOB_ID), + SucceededHouseholdCalculation(job_id=JOB_ID, output=_succeeded().output), + FailedHouseholdCalculation( + job_id=JOB_ID, + error=HouseholdCalculationError( + code="invalid_spm_selection", + message="SPM settings are only available for the US", + spm_error=SPMErrorDetail( + code="SPM_SETTINGS_UNSUPPORTED", + message="SPM settings are only available for the US", + ), + ), + ), + ] + + for response in responses: + parsed = parse_household_calculation_poll_response( + response.model_dump(mode="json") + ) + assert type(parsed) is type(response) + assert parsed == response + assert parsed.schema_version == 1 + with pytest.raises(ValidationError): + parse_household_calculation_poll_response( + {"status": "succeeded", "job_id": str(JOB_ID)} + ) From ce8eec891fa39a3f9906eafcbd1fb7ad42795782 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:42:46 +0400 Subject: [PATCH 2/4] Record country capabilities in v2 manifest schema version 2 The v2 version manifest now records, for each country a worker application serves, whether it is society-wide capable (it has a single-simulation worker) and which function calculates its households. - Schema version 2 is written under a new active-schema-2 key. The same manifest is still written in schema version 1 form under active-with-cache and active, for readers released before this change, whenever schema version 1 can describe it. - Readers prefer the schema version 2 key and upgrade a schema version 1 document, marking every listed country society-wide capable with no household function, so the first publish keeps earlier versions. - The worker validation receipt gains household validation invocation IDs; every household function must be validated, and simulation-worker validation must cover exactly the society-wide countries. - The publisher marks every bundle country society-wide capable and records no household functions yet. - The Simulation Entrypoint's comparison adapter skips any country the manifest does not mark as society-wide capable. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 28 +- .../stage12_manifest.py | 189 +++++++++++-- .../tests/test_stage12_manifest.py | 251 +++++++++++++++--- .../stage12_adapter.py | 20 +- .../tests/stage12_fixtures.py | 20 +- .../tests/test_stage12_adapter.py | 19 ++ .../tests/test_stage12_backend.py | 69 ++++- .../modal/utils/update_v2_version_manifest.py | 8 +- .../tests/test_update_v2_version_manifest.py | 15 ++ 9 files changed, 522 insertions(+), 97 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index fed0d09e4..846698940 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -384,14 +384,28 @@ and coordinator images do not contain these files. Deployment validation runs one real 2026 region-group baseline and requires it to load the precomputed file without recalculation. -Readers released before cache releases reject fields they do not know, and +Manifest schema version 2 records, for each country a worker application +serves, whether it is society-wide capable (it has a single-simulation worker, +so the report coordinator may be given its reports) and which function +calculates its households. The Simulation Entrypoint's comparison adapter skips +any country that the manifest does not mark as society-wide capable, with the +reason `unsupported_country`. + +Readers released before a schema change reject fields they do not know, and the manifest is published before the Cloud Run revision that reads it is -promoted. The publisher therefore writes the manifest under the -`active-with-cache` key and the same manifest without cache releases under the -`active` key, which those readers keep reading. Current readers read -`active-with-cache`, falling back to `active` only before the first publish -that writes both keys. Removing the `active` copy needs a later change, once -no deployed or rollback revision reads it. +promoted. The publisher therefore writes the schema version 2 manifest under +the `active-schema-2` key, the same manifest in schema version 1 form (without +the capability and household fields) under the `active-with-cache` key, and the +schema version 1 form without cache releases under the `active` key. Readers +released before schema version 2 keep reading `active-with-cache`, and readers +released before cache releases keep reading `active`. Current readers read +`active-schema-2`, falling back to `active-with-cache` and then `active` only +before the first schema version 2 publish; a schema version 1 document read +that way marks every listed country society-wide capable, with no household +function. Schema version 1 cannot describe a country that is not society-wide +capable, so a manifest holding one is written only under `active-schema-2`, +and the older keys keep the last manifest written there. Removing the older +copies needs a later change, once no deployed or rollback revision reads them. Only an empty-policy baseline using the bundle's default dataset and an exact region-group scope can read a baseline cache file. Reform calculations always diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index f0af8a3cf..55f8896ab 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -1,4 +1,12 @@ -"""Strict, separately stored Stage 12 v2 worker version manifest.""" +"""Strict, separately stored Stage 12 v2 worker version manifest. + +Schema version 2 records, for each country a worker application serves, +whether it can run society-wide (economy) reports and which function +calculates households. It is stored under its own key. The same manifest is +also written in schema version 1 form, without those fields, under the keys +that readers released before schema version 2 read, whenever schema version 1 +can describe it. +""" from __future__ import annotations @@ -21,9 +29,15 @@ V1_ROUTING_STATE_NAME = "simulation-api-routing-state" V2_VERSION_MANIFEST_NAME = "simulation-api-v2-version-manifest" +# The schema version 2 manifest. Readers released with schema version 2 read +# this key first. +SCHEMA_2_MANIFEST_KEY = "active-schema-2" +# Readers released before schema version 2 forbid unknown fields, so they keep +# reading this key, which holds the same manifest in schema version 1 form. ACTIVE_MANIFEST_KEY = "active-with-cache" # Readers released before Stage 12 cache releases forbid unknown fields, so -# they keep reading this key, which holds the same manifest without them. +# they keep reading this key, which holds the schema version 1 form without +# cache fields. PRE_CACHE_MANIFEST_KEY = "active" V2_APPLICATION_PREFIX = "policyengine-simulation-v2-py" V1_APPLICATION_PREFIX = "policyengine-simulation-py" @@ -40,10 +54,33 @@ class StrictManifestModel(BaseModel): class V2CountryWorker(StrictManifestModel): - """A country served by a v2 worker, and the callable that runs its simulations.""" + """A country served by a v2 worker application, and what it can calculate. + + A society-wide capable country names the callable that runs one of its + report simulations; only such a country may be given to the report + coordinator. ``household_callable`` names the function that calculates + its households, and is None for applications deployed before household + functions existed. + """ country: CountryId - single_simulation_callable: CallableName + society_wide_capable: bool + single_simulation_callable: CallableName | None = None + household_callable: CallableName | None = None + + @model_validator(mode="after") + def validate_capabilities(self) -> V2CountryWorker: + """Require a single-simulation callable exactly for society-wide countries. + + A country must also have at least one capability. + """ + if self.society_wide_capable != (self.single_simulation_callable is not None): + msg = "a society-wide country, and only one, names a simulation callable" + raise ValueError(msg) + if not self.society_wide_capable and self.household_callable is None: + msg = "a country needs at least one capability" + raise ValueError(msg) + return self class Stage12CacheRelease(StrictManifestModel): @@ -65,7 +102,8 @@ class V2WorkerValidation(StrictManifestModel): """Receipt for a successful validation of a deployed v2 worker application. Names the application and bundle digest that were validated, when, and the - validation invocation IDs overall and for each country. + validation invocation IDs overall, for each society-wide country's + simulation worker, and for each country's household function. """ validated: Literal[True] @@ -74,6 +112,9 @@ class V2WorkerValidation(StrictManifestModel): validated_at: datetime validation_invocation_id: ManifestText country_validation_invocation_ids: dict[CountryId, ManifestText] + household_validation_invocation_ids: dict[CountryId, ManifestText] = Field( + default_factory=dict + ) cache: Stage12CacheRelease | None = None @@ -81,7 +122,8 @@ class V2WorkerVersion(StrictManifestModel): """One deployed v2 worker application and the bundle it runs. Lists its report coordinator callable, its per-country workers and its - validation receipt; the countries must match the bundle's, in order. + validation receipt; the countries must match the bundle's, in order, and + every country's functions must have been validated. """ application_name: ManifestText @@ -128,11 +170,44 @@ def validate_worker(self) -> V2WorkerVersion: if countries != bundle_countries: msg = "v2 worker countries do not match the normalized bundle" raise ValueError(msg) - if set(self.validation.country_validation_invocation_ids) != set(countries): - msg = "worker validation does not cover every country" + society_wide = { + worker.country for worker in self.countries if worker.society_wide_capable + } + if set(self.validation.country_validation_invocation_ids) != society_wide: + msg = "worker validation does not cover exactly the society-wide countries" + raise ValueError(msg) + households = { + worker.country + for worker in self.countries + if worker.household_callable is not None + } + if set(self.validation.household_validation_invocation_ids) != households: + msg = "worker validation does not cover exactly the household functions" raise ValueError(msg) return self + def society_wide_worker(self, country: str) -> V2CountryWorker | None: + """Return ``country``'s worker when it can run society-wide reports.""" + return next( + ( + worker + for worker in self.countries + if worker.country == country and worker.society_wide_capable + ), + None, + ) + + def household_worker(self, country: str) -> V2CountryWorker | None: + """Return ``country``'s worker when it has a household function.""" + return next( + ( + worker + for worker in self.countries + if worker.country == country and worker.household_callable is not None + ), + None, + ) + class V2VersionManifest(StrictManifestModel): """Registry of deployed v2 worker versions, keyed by PolicyEngine.py version. @@ -141,7 +216,7 @@ class V2VersionManifest(StrictManifestModel): ``generation`` identifies the publish that produced this manifest. """ - schema_version: Literal[1] = 1 + schema_version: Literal[2] = 2 generation: ManifestText default_version: ManifestText versions: dict[str, V2WorkerVersion] @@ -215,20 +290,80 @@ def _is_newer(candidate: str, current: str | None) -> bool: return False -def pre_cache_manifest_payload(manifest: V2VersionManifest) -> dict[str, object]: - """Return the manifest as readers released before cache releases accept it.""" - return manifest.model_dump( +_SCHEMA_1_COUNTRY_FIELDS = frozenset({"country", "single_simulation_callable"}) + + +def schema_1_manifest_payload( + manifest: V2VersionManifest, + *, + include_cache: bool, +) -> dict[str, object] | None: + """Return the manifest as readers released before schema version 2 accept it. + + Without ``include_cache`` the cache fields are left out too, for readers + released before cache releases. Returns None when a country is not + society-wide capable, which schema version 1 cannot describe. + """ + if any( + not country.society_wide_capable + for worker in manifest.versions.values() + for country in worker.countries + ): + return None + validation_exclude: set[str] = {"household_validation_invocation_ids"} + worker_exclude: dict[str, object] = { + "countries": {"__all__": {"society_wide_capable", "household_callable"}}, + "validation": validation_exclude, + } + if not include_cache: + worker_exclude["cache"] = True + validation_exclude.add("cache") + payload = manifest.model_dump( mode="json", - exclude={"versions": {"__all__": {"cache": True, "validation": {"cache"}}}}, + exclude={"versions": {"__all__": worker_exclude}}, ) + payload["schema_version"] = 1 + return payload + + +def _upgraded_country(country: object) -> object: + if isinstance(country, Mapping) and set(country) == _SCHEMA_1_COUNTRY_FIELDS: + # Schema version 1 recorded only countries that run report + # simulations, and no household functions. + return {**country, "society_wide_capable": True} + return country + + +def _upgraded_schema_1_payload(value: Mapping[str, object]) -> dict[str, object]: + payload = deepcopy(dict(value)) + payload["schema_version"] = 2 + versions = payload.get("versions") + if isinstance(versions, dict): + for worker in versions.values(): + if isinstance(worker, dict) and isinstance(worker.get("countries"), list): + worker["countries"] = [ + _upgraded_country(country) for country in worker["countries"] + ] + return payload + + +def parse_v2_manifest(value: object) -> V2VersionManifest: + """Validate a stored manifest, upgrading a schema version 1 document. + + Raises ValueError (as a pydantic ValidationError) for an invalid document. + """ + if isinstance(value, Mapping) and value.get("schema_version") == 1: + value = _upgraded_schema_1_payload(value) + return V2VersionManifest.model_validate(value) def _current_manifest(store: V2ManifestStore) -> object: - current = store.get(ACTIVE_MANIFEST_KEY) - if current is None: - # Until the first publish that writes both keys. - current = store.get(PRE_CACHE_MANIFEST_KEY) - return current + for key in (SCHEMA_2_MANIFEST_KEY, ACTIVE_MANIFEST_KEY, PRE_CACHE_MANIFEST_KEY): + # Until the first schema version 2 publish, only the older keys exist. + current = store.get(key) + if current is not None: + return current + return None def build_next_v2_manifest( @@ -247,7 +382,7 @@ def build_next_v2_manifest( versions: dict[str, V2WorkerVersion] = {} previous_default = None else: - parsed = V2VersionManifest.model_validate(current) + parsed = parse_v2_manifest(current) versions = deepcopy(parsed.versions) previous_default = parsed.default_version version = worker.bundle.policyengine_version @@ -274,8 +409,10 @@ def publish_v2_manifest( ) -> V2VersionManifest: """Add ``worker`` to the stored manifest, write it back and return it. - The whole manifest is written as one store value, so readers never see a - partial update. + The whole manifest is written as one store value per key, so readers never + see a partial update. The schema version 1 keys are rewritten only when + schema version 1 can describe the manifest; otherwise their readers keep + the last manifest written there. """ next_manifest = build_next_v2_manifest( current=_current_manifest(store), @@ -284,8 +421,12 @@ def publish_v2_manifest( ) # Each key holds one entire validated document, and a reader reads only # one key, so no reader observes a partial update assembled from several. - store[ACTIVE_MANIFEST_KEY] = next_manifest.model_dump(mode="json") - store[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(next_manifest) + store[SCHEMA_2_MANIFEST_KEY] = next_manifest.model_dump(mode="json") + with_cache = schema_1_manifest_payload(next_manifest, include_cache=True) + without_cache = schema_1_manifest_payload(next_manifest, include_cache=False) + if with_cache is not None and without_cache is not None: + store[ACTIVE_MANIFEST_KEY] = with_cache + store[PRE_CACHE_MANIFEST_KEY] = without_cache return next_manifest @@ -305,7 +446,7 @@ def load(self) -> V2VersionManifest: """ value = _current_manifest(self._store) try: - manifest = V2VersionManifest.model_validate(value) + manifest = parse_v2_manifest(value) except (TypeError, ValueError): if self._last_valid is None: msg = "no valid v2 version manifest is available" diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 7ab548a54..c0926a5e0 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -14,6 +14,7 @@ from policyengine_simulation_contract.stage12_manifest import ( ACTIVE_MANIFEST_KEY, PRE_CACHE_MANIFEST_KEY, + SCHEMA_2_MANIFEST_KEY, V1_ROUTING_STATE_NAME, V2_VERSION_MANIFEST_NAME, Stage12CacheRelease, @@ -23,8 +24,9 @@ V2WorkerVersion, assert_separate_manifest_names, manifest_digest, - pre_cache_manifest_payload, + parse_v2_manifest, publish_v2_manifest, + schema_1_manifest_payload, v2_application_name, validate_manifest_has_no_credentials, ) @@ -86,22 +88,35 @@ def _resolved_bundle(version: str = "5.2.0") -> ResolvedStage12Bundle: ) -def _worker(version: str = "5.2.0") -> V2WorkerVersion: +def _worker( + version: str = "5.2.0", + *, + with_household_functions: bool = False, + household_only: tuple[str, ...] = (), +) -> V2WorkerVersion: resolved = _resolved_bundle(version) application_name = v2_application_name(version) + countries = tuple( + V2CountryWorker( + country=country, + society_wide_capable=country not in household_only, + single_simulation_callable=( + None + if country in household_only + else f"run_single_simulation_{country}" + ), + household_callable=( + f"calculate_household_{country}" + if with_household_functions or country in household_only + else None + ), + ) + for country in ("us", "uk") + ) return V2WorkerVersion( application_name=application_name, report_coordinator_callable="coordinate_report", - countries=( - V2CountryWorker( - country="us", - single_simulation_callable="run_single_simulation_us", - ), - V2CountryWorker( - country="uk", - single_simulation_callable="run_single_simulation_uk", - ), - ), + countries=countries, bundle=resolved.bundle, bundle_manifest_sha256=resolved.bundle_manifest_sha256, validation=V2WorkerValidation( @@ -111,8 +126,14 @@ def _worker(version: str = "5.2.0") -> V2WorkerVersion: validated_at=datetime(2026, 9, 14, tzinfo=UTC), validation_invocation_id="validation-123", country_validation_invocation_ids={ - "us": "validation-us-123", - "uk": "validation-uk-123", + worker.country: f"validation-{worker.country}-123" + for worker in countries + if worker.society_wide_capable + }, + household_validation_invocation_ids={ + worker.country: f"household-validation-{worker.country}-123" + for worker in countries + if worker.household_callable is not None }, ), ) @@ -123,8 +144,10 @@ def test_v2_application_name_cannot_collide_with_v1() -> None: assert not v2_application_name("5.2.0").startswith("policyengine-simulation-py") -def _cached_worker(version: str = "5.2.0") -> V2WorkerVersion: - worker = _worker(version) +def _cached_worker( + version: str = "5.2.0", *, with_household_functions: bool = False +) -> V2WorkerVersion: + worker = _worker(version, with_household_functions=with_household_functions) cache = Stage12CacheRelease( manifest_sha256="c" * 64, partition_sha256="d" * 64, @@ -141,30 +164,84 @@ def _cached_worker(version: str = "5.2.0") -> V2WorkerVersion: def test_publish_writes_one_complete_v2_document_per_key() -> None: store = FakeStore() - manifest = publish_v2_manifest(store=store, worker=_cached_worker()) + manifest = publish_v2_manifest( + store=store, worker=_cached_worker(with_household_functions=True) + ) assert store.writes == [ - (ACTIVE_MANIFEST_KEY, manifest.model_dump(mode="json")), - (PRE_CACHE_MANIFEST_KEY, pre_cache_manifest_payload(manifest)), + (SCHEMA_2_MANIFEST_KEY, manifest.model_dump(mode="json")), + ( + ACTIVE_MANIFEST_KEY, + schema_1_manifest_payload(manifest, include_cache=True), + ), + ( + PRE_CACHE_MANIFEST_KEY, + schema_1_manifest_payload(manifest, include_cache=False), + ), ] + assert manifest.schema_version == 2 assert manifest.default_version == "5.2.0" assert manifest.versions["5.2.0"].validation.validated is True - assert V2ManifestLoader(store).load().versions["5.2.0"].cache is not None + loaded = V2ManifestLoader(store).load() + assert loaded == manifest + assert loaded.versions["5.2.0"].cache is not None + assert loaded.versions["5.2.0"].household_worker("us") is not None + + +def _schema_1_fields(payload: dict) -> None: + # The fields of each schema version 1 model. Readers released before + # schema version 2 forbid any other field, so adding one here would stop + # them loading. + assert payload["schema_version"] == 1 + assert set(payload) == { + "default_version", + "generation", + "schema_version", + "versions", + } + for country in payload["versions"]["5.2.0"]["countries"]: + assert set(country) == {"country", "single_simulation_callable"} + + +def test_readers_released_before_schema_2_get_only_fields_they_accept() -> None: + store = FakeStore() + publish_v2_manifest( + store=store, worker=_cached_worker(with_household_functions=True) + ) + + with_cache = store.values[ACTIVE_MANIFEST_KEY] + _schema_1_fields(with_cache) + worker = with_cache["versions"]["5.2.0"] + assert set(worker) == { + "application_name", + "bundle", + "bundle_manifest_sha256", + "cache", + "countries", + "report_coordinator_callable", + "validation", + } + assert set(worker["validation"]) == { + "application_name", + "bundle_manifest_sha256", + "cache", + "country_validation_invocation_ids", + "validated", + "validated_at", + "validation_invocation_id", + } def test_readers_released_before_the_cache_get_only_fields_they_accept() -> None: store = FakeStore() - publish_v2_manifest(store=store, worker=_cached_worker()) + publish_v2_manifest( + store=store, worker=_cached_worker(with_household_functions=True) + ) # The fields of each model before Stage 12 cache releases. Those readers # forbid any other field, so adding one here would stop them loading. pre_cache = store.values[PRE_CACHE_MANIFEST_KEY] - assert set(pre_cache) == { - "default_version", - "generation", - "schema_version", - "versions", - } + _schema_1_fields(pre_cache) worker = pre_cache["versions"]["5.2.0"] assert set(worker) == { "application_name", @@ -184,35 +261,87 @@ def test_readers_released_before_the_cache_get_only_fields_they_accept() -> None } -def test_first_publish_with_cache_keeps_versions_from_the_pre_cache_key() -> None: - earlier = _worker("5.1.0").model_dump( - mode="json", exclude={"cache": True, "validation": {"cache"}} - ) - store = FakeStore( - { - PRE_CACHE_MANIFEST_KEY: { - "schema_version": 1, - "generation": "5.1.0:earlier", - "default_version": "5.1.0", - "versions": {"5.1.0": earlier}, +def _schema_1_store(key: str, version: str = "5.1.0") -> FakeStore: + earlier = schema_1_manifest_payload( + parse_v2_manifest( + { + "schema_version": 2, + "generation": f"{version}:earlier", + "default_version": version, + "versions": {version: _worker(version).model_dump(mode="json")}, } - } + ), + include_cache=key == ACTIVE_MANIFEST_KEY, ) + return FakeStore({key: earlier}) + + +@pytest.mark.parametrize("key", [ACTIVE_MANIFEST_KEY, PRE_CACHE_MANIFEST_KEY]) +def test_loader_upgrades_a_schema_1_document(key: str) -> None: + store = _schema_1_store(key) + + manifest = V2ManifestLoader(store).load() + + assert manifest.schema_version == 2 + worker = manifest.versions["5.1.0"] + assert [country.society_wide_capable for country in worker.countries] == [ + True, + True, + ] + assert worker.household_worker("us") is None + assert worker.society_wide_worker("uk") is not None + assert worker.validation.household_validation_invocation_ids == {} + + +def test_loader_prefers_the_schema_2_document() -> None: + store = _schema_1_store(ACTIVE_MANIFEST_KEY) + newer = publish_v2_manifest( + store=FakeStore(), worker=_worker(with_household_functions=True) + ) + store.values[SCHEMA_2_MANIFEST_KEY] = newer.model_dump(mode="json") + + assert V2ManifestLoader(store).load() == newer + + +@pytest.mark.parametrize("key", [ACTIVE_MANIFEST_KEY, PRE_CACHE_MANIFEST_KEY]) +def test_first_schema_2_publish_keeps_versions_from_a_schema_1_key(key: str) -> None: + store = _schema_1_store(key) assert V2ManifestLoader(store).load().default_version == "5.1.0" - manifest = publish_v2_manifest(store=store, worker=_cached_worker()) + manifest = publish_v2_manifest( + store=store, worker=_cached_worker(with_household_functions=True) + ) assert sorted(manifest.versions) == ["5.1.0", "5.2.0"] + assert manifest.versions["5.1.0"].household_worker("us") is None + assert manifest.versions["5.2.0"].household_worker("us") is not None assert V2ManifestLoader(store).load() == manifest +def test_household_only_country_keeps_schema_1_readers_on_their_last_manifest() -> None: + store = _schema_1_store(ACTIVE_MANIFEST_KEY, version="5.1.0") + legacy_before = store.values[ACTIVE_MANIFEST_KEY] + + manifest = publish_v2_manifest(store=store, worker=_worker(household_only=("uk",))) + + # Schema version 1 cannot describe a country without a single-simulation + # function, so only the schema version 2 document changes. + assert [key for key, _ in store.writes] == [SCHEMA_2_MANIFEST_KEY] + assert store.values[ACTIVE_MANIFEST_KEY] == legacy_before + assert schema_1_manifest_payload(manifest, include_cache=True) is None + worker = manifest.versions["5.2.0"] + assert worker.society_wide_worker("uk") is None + assert worker.household_worker("uk") is not None + assert worker.society_wide_worker("us") is not None + + def test_loader_retains_its_last_valid_v2_document() -> None: store = FakeStore() expected = publish_v2_manifest(store=store, worker=_worker()) loader = V2ManifestLoader(store) assert loader.load() == expected - store.values[ACTIVE_MANIFEST_KEY] = {"schema_version": 999} + store.values[SCHEMA_2_MANIFEST_KEY] = {"schema_version": 999} assert loader.load() == expected @@ -243,7 +372,7 @@ def test_v1_and_v2_publication_and_removal_are_bidirectionally_isolated() -> Non manifest ) - v2_store.values.pop(ACTIVE_MANIFEST_KEY) + v2_store.values.pop(SCHEMA_2_MANIFEST_KEY) assert v1_store.values == {ACTIVE_MANIFEST_KEY: {"v1": "new"}} @@ -274,6 +403,44 @@ def test_cache_metadata_is_optional_for_old_workers_and_must_match_validation() V2WorkerVersion.model_validate(value) +@pytest.mark.parametrize( + ("society_wide_capable", "single_simulation_callable", "household_callable"), + [ + (True, None, None), + (False, "run_single_simulation_us", None), + (False, None, None), + ], +) +def test_country_capabilities_must_name_their_functions( + society_wide_capable: bool, + single_simulation_callable: str | None, + household_callable: str | None, +) -> None: + with pytest.raises(ValueError, match="society-wide|capability"): + V2CountryWorker( + country="us", + society_wide_capable=society_wide_capable, + single_simulation_callable=single_simulation_callable, + household_callable=household_callable, + ) + + +def test_household_functions_need_a_validation_receipt() -> None: + value = _worker(with_household_functions=True).model_dump(mode="json") + value["validation"]["household_validation_invocation_ids"] = {"us": "only-us"} + + with pytest.raises(ValueError, match="household"): + V2WorkerVersion.model_validate(value) + + +def test_society_wide_countries_need_a_validation_receipt() -> None: + value = _worker(household_only=("uk",)).model_dump(mode="json") + value["validation"]["country_validation_invocation_ids"]["uk"] = "not-needed" + + with pytest.raises(ValueError, match="society-wide"): + V2WorkerVersion.model_validate(value) + + def test_manifest_rejects_credential_like_fields() -> None: with pytest.raises(ValueError, match="credential-like"): validate_manifest_has_no_credentials( diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py index e0e8f249e..104a8bc99 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py @@ -22,10 +22,7 @@ if TYPE_CHECKING: from policyengine_simulation_contract.stage12_bundle import CountryId - from policyengine_simulation_contract.stage12_manifest import ( - V2CountryWorker, - V2WorkerVersion, - ) + from policyengine_simulation_contract.stage12_manifest import V2WorkerVersion class ComparisonSkipReason(StrEnum): @@ -87,20 +84,17 @@ def _skip(reason: ComparisonSkipReason) -> AdaptedComparisonRun: return AdaptedComparisonRun(report=None, skip_reason=reason) -def _country_worker(worker: V2WorkerVersion, country: str) -> V2CountryWorker | None: - return next( - (item for item in worker.countries if item.country == country), - None, - ) - - def adapt_annual_comparison( payload: dict[str, Any], *, evaluation_id: UUID, worker: V2WorkerVersion, ) -> AdaptedComparisonRun: - """Convert only the reviewed production request shape without guessing.""" + """Convert only the reviewed production request shape without guessing. + + Only a country the manifest marks as society-wide capable is converted; + any other country is skipped as unsupported. + """ if any( value is not None for key, value in payload.items() @@ -116,7 +110,7 @@ def adapt_annual_comparison( if segmented is not None and not isinstance(segmented, bool): return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE) country = payload.get("country") - if not isinstance(country, str) or _country_worker(worker, country) is None: + if not isinstance(country, str) or worker.society_wide_worker(country) is None: return _skip(ComparisonSkipReason.UNSUPPORTED_COUNTRY) bundle_country = next( item for item in worker.bundle.countries if item.country == country diff --git a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py index 46ebe97e7..23eb49f83 100644 --- a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py +++ b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py @@ -28,7 +28,8 @@ EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") -def worker() -> V2WorkerVersion: +def worker(*, household_only: tuple[str, ...] = ()) -> V2WorkerVersion: + """Return a validated 5.2.0 worker; ``household_only`` countries run no reports.""" countries = [] country_workers = [] for country, package, package_version, dataset in ( @@ -59,10 +60,17 @@ def worker() -> V2WorkerVersion: ), ) ) + society_wide = country not in household_only country_workers.append( V2CountryWorker( country=country, - single_simulation_callable=f"run_single_simulation_{country}", + society_wide_capable=society_wide, + single_simulation_callable=( + f"run_single_simulation_{country}" if society_wide else None + ), + household_callable=( + None if society_wide else f"calculate_household_{country}" + ), ) ) application_name = v2_application_name("5.2.0") @@ -85,8 +93,12 @@ def worker() -> V2WorkerVersion: validated_at=datetime(2026, 9, 14, tzinfo=UTC), validation_invocation_id="validation-1", country_validation_invocation_ids={ - "us": "validation-us", - "uk": "validation-uk", + country_worker.country: f"validation-{country_worker.country}" + for country_worker in country_workers + if country_worker.society_wide_capable + }, + household_validation_invocation_ids={ + country: f"household-validation-{country}" for country in household_only }, ), ) diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index 53a449ec9..290abd514 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -121,6 +121,25 @@ def test_segmentation_selection_is_forwarded_to_both_simulations(segmented) -> N assert result.report.reform.options == {"segmented": segmented} +def test_adapter_accepts_only_countries_the_manifest_marks_society_wide() -> None: + payload = {**eligible_payload(), "country": "uk", "region": "uk", "version": None} + + household_only = adapt_annual_comparison( + payload, + evaluation_id=EVALUATION_ID, + worker=worker(household_only=("uk",)), + ) + society_wide = adapt_annual_comparison( + payload, + evaluation_id=EVALUATION_ID, + worker=worker(), + ) + + assert household_only.report is None + assert household_only.skip_reason is ComparisonSkipReason.UNSUPPORTED_COUNTRY + assert society_wide.eligible + + @pytest.mark.parametrize( ("update", "reason"), [ diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_backend.py b/projects/policyengine-simulation-entry/tests/test_stage12_backend.py index 0145e54dc..773994b1b 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_backend.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_backend.py @@ -26,8 +26,8 @@ class FakeLoader: - def __init__(self): - self.worker = worker() + def __init__(self, *, household_only: tuple[str, ...] = ()): + self.worker = worker(household_only=household_only) self.resolved_versions = [] def resolve(self, version=None): @@ -65,7 +65,14 @@ async def spawn(self, **call): return "modal-call-1" -def _backend(store, invoker, *, environment="staging", modal_environment="staging"): +def _backend( + store, + invoker, + *, + environment="staging", + modal_environment="staging", + household_only=(), +): runtime = MagicMock() runtime.capture_context.return_value = { "observability_id": "00000000-0000-4000-8000-000000000099", @@ -76,7 +83,7 @@ def _backend(store, invoker, *, environment="staging", modal_environment="stagin environment=environment, stage12_v2_manifest_environment=modal_environment, ), - manifest_loader=FakeLoader(), + manifest_loader=FakeLoader(household_only=household_only), store=store, invoker=invoker, runtime=runtime, @@ -344,6 +351,60 @@ def test_temporary_submission_rejects_unsupported_input_without_side_effects() - assert invoker.calls == [] +UK_PAYLOAD = {**eligible_payload(), "country": "uk", "region": "uk", "version": None} + + +def test_temporary_submission_rejects_a_country_without_society_wide_reports() -> None: + store = FakeStore() + invoker = FakeInvoker() + + with pytest.raises(TemporaryStage12UnsupportedRequestError) as error: + asyncio.run( + _backend(store, invoker, household_only=("uk",)).submit_temporary_report( + request_payload=UK_PAYLOAD, + request_id="request-1", + observability_id=OBSERVABILITY_ID, + ) + ) + + assert error.value.reason == "unsupported_country" + assert store.events == [] + assert invoker.calls == [] + + +def test_automatic_submission_skips_a_country_without_society_wide_reports() -> None: + invoker = FakeInvoker() + + asyncio.run( + _backend( + FakeStore(), invoker, household_only=("uk",) + ).dispatch_after_production( + request_payload=UK_PAYLOAD, + production_response=_response(), + request_id="request-1", + observability_id=OBSERVABILITY_ID, + ) + ) + + assert invoker.calls == [] + + +def test_society_wide_country_dispatches_its_own_simulation_callable() -> None: + invoker = FakeInvoker() + + asyncio.run( + _backend(FakeStore(), invoker, household_only=("us",)).submit_temporary_report( + request_payload=UK_PAYLOAD, + request_id="request-1", + observability_id=OBSERVABILITY_ID, + ) + ) + + assert invoker.calls[0]["context_payload"]["simulation_callable"] == ( + "run_single_simulation_uk" + ) + + def test_temporary_status_reads_coordinator_owned_postgres_state() -> None: store = FakeStore() invoker = FakeInvoker() diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py index a8f8801ff..8157504cd 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py @@ -24,8 +24,9 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: """Build the installed bundle's v2 worker entry from a validation record. - Raises ValueError if the record is invalid or the entry carries a - credential-like key. + Every bundle country runs report simulations, so each is recorded as + society-wide capable. Raises ValueError if the record is invalid or the + entry carries a credential-like key. """ resolved = load_stage12_bundle() validation = V2WorkerValidation.model_validate(validation_payload) @@ -35,7 +36,8 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: countries=tuple( V2CountryWorker( country=country.country, - single_simulation_callable=(f"run_single_simulation_{country.country}"), + society_wide_capable=True, + single_simulation_callable=f"run_single_simulation_{country.country}", ) for country in resolved.bundle.countries ), diff --git a/projects/policyengine-simulation-executor/tests/test_update_v2_version_manifest.py b/projects/policyengine-simulation-executor/tests/test_update_v2_version_manifest.py index 7d074a446..15d254a80 100644 --- a/projects/policyengine-simulation-executor/tests/test_update_v2_version_manifest.py +++ b/projects/policyengine-simulation-executor/tests/test_update_v2_version_manifest.py @@ -38,6 +38,21 @@ def test_publisher_uses_only_packaged_bundle_values() -> None: ) +def test_publisher_marks_every_bundle_country_society_wide_capable() -> None: + worker = publisher.build_worker(_validation().model_dump(mode="json")) + resolved = publisher.load_stage12_bundle() + + assert [country.model_dump(mode="json") for country in worker.countries] == [ + { + "country": country.country, + "society_wide_capable": True, + "single_simulation_callable": f"run_single_simulation_{country.country}", + "household_callable": None, + } + for country in resolved.bundle.countries + ] + + def test_publisher_rejects_validation_for_another_application() -> None: payload = _validation().model_dump(mode="json") payload["application_name"] = "policyengine-simulation-v2-py-another" From 347faef0a0cbebdd6c4587dc8d88ab3ae0a3dc51 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:19:07 +0400 Subject: [PATCH 3/4] Serialize households with their keys in the caller's order Pydantic dumps declared fields before extra ones, so a household document always came back with `axes` first and `people` before every group collection, whatever order the caller sent. The Public API walks a household in the order it was sent, so code that walks the serialized situation (the unrecognized-input checks, deprecated-input warnings and requested values) listed its results in a different order. The situation models now record the caller's top-level key order and serialize in it. Co-Authored-By: Claude Opus 5.5 --- .../household_execution.py | 34 ++++++++++++++++--- .../tests/test_household_execution.py | 25 ++++++++++++++ 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py index 2f6bac4b0..17b0af101 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py @@ -31,6 +31,7 @@ from __future__ import annotations import math +from collections.abc import Mapping from typing import TYPE_CHECKING, Annotated, Any, Literal, Self from uuid import UUID @@ -40,6 +41,7 @@ BeforeValidator, ConfigDict, Field, + PrivateAttr, TypeAdapter, field_validator, model_serializer, @@ -53,7 +55,7 @@ ) if TYPE_CHECKING: - from pydantic import SerializerFunctionWrapHandler + from pydantic import ModelWrapValidatorHandler, SerializerFunctionWrapHandler HOUSEHOLD_CONTRACT_SCHEMA_VERSION = 1 # Size limits. The entity and value limits match the Public API's stored v2 @@ -269,11 +271,30 @@ def _validated_axes( class _SituationModel(BaseModel): - """Situation-shaped household: ``people``, group collections and ``axes``.""" + """Situation-shaped household: ``people``, group collections and ``axes``. + + It serializes its top-level keys in the order the caller gave them. The + Public API walks a household in that order, so the order of unrecognized + inputs, warnings and calculated values depends on it. + """ model_config = ConfigDict(extra="allow", allow_inf_nan=False, frozen=True) axes: Axes = None + # The caller's top-level keys, in order. Pydantic itself would dump the + # declared fields (``axes`` and ``people``) before the group collections. + _key_order: tuple[str, ...] = PrivateAttr(default=()) + + @model_validator(mode="wrap") + @classmethod + def remember_key_order( + cls, data: object, handler: ModelWrapValidatorHandler[Self] + ) -> Self: + """Record the order of the caller's top-level keys for serialization.""" + situation = handler(data) + if isinstance(data, Mapping): + situation._key_order = tuple(str(key) for key in data) # noqa: SLF001 # set on the instance this validator just built + return situation @model_validator(mode="after") def limit_collections(self) -> Self: @@ -287,11 +308,16 @@ def limit_collections(self) -> Self: def serialize_situation( self, handler: SerializerFunctionWrapHandler ) -> dict[str, Any]: - """Leave out ``axes`` when there are none, so the shape matches the input.""" + """Return the caller's keys in the caller's order, without ``axes`` if none. + + Keys the caller did not give follow, in Pydantic's order. + """ dumped: dict[str, Any] = handler(self) if self.axes is None: dumped.pop("axes", None) - return dumped + ordered = {key: dumped[key] for key in self._key_order if key in dumped} + ordered.update(dumped) + return ordered @property def has_axes(self) -> bool: diff --git a/libs/policyengine-simulation-contract/tests/test_household_execution.py b/libs/policyengine-simulation-contract/tests/test_household_execution.py index ed98df376..dcfd0806d 100644 --- a/libs/policyengine-simulation-contract/tests/test_household_execution.py +++ b/libs/policyengine-simulation-contract/tests/test_household_execution.py @@ -117,6 +117,31 @@ def test_household_document_keeps_the_public_api_situation_shape() -> None: assert document.has_axes +GROUPS_BEFORE_PEOPLE = { + "households": US_SITUATION["households"], + "people": US_SITUATION["people"], + "tax_units": US_SITUATION["tax_units"], + "axes": US_SITUATION["axes"], +} + + +@pytest.mark.parametrize("situation", [US_SITUATION, GROUPS_BEFORE_PEOPLE]) +def test_households_serialize_their_keys_in_the_callers_order(situation) -> None: + # The Public API walks a household in the order its caller sent it, so + # unrecognized inputs and warnings are listed in that order. + document = HouseholdDocument.model_validate(situation) + request = _request(household=situation) + transported = HouseholdExecutionRequest.model_validate( + json.loads(request.model_dump_json()) + ) + calculated = CalculatedHousehold.model_validate(situation) + + assert list(document.to_situation()) == list(situation) + assert list(json.loads(document.model_dump_json())) == list(situation) + assert list(transported.household.to_situation()) == list(situation) + assert list(calculated.to_situation()) == list(situation) + + def test_household_without_axes_never_gains_an_axes_key() -> None: situation = {key: value for key, value in US_SITUATION.items() if key != "axes"} From 824767bd043d16c2e280ff23ec2ba8f4bad8efb4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:19:29 +0400 Subject: [PATCH 4/4] Document how a saved household is calculated like the Public API The Public API's stored-household route always adds missing variables and treats a saved SPM selection as chosen, while the contract ties `add_missing_variables` only to `/calculate-full`. State which fields a saved-household replay sets, so callers built later match that route. Documentation only: the change is a docstring, so no test can fail first. Co-Authored-By: Claude Opus 5.5 --- .../household_execution.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py index 17b0af101..63cbb0b2a 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py @@ -454,6 +454,13 @@ class _HouseholdCalculationInputs(StrictHouseholdModel): whether the selection came from this request or from a saved household, which changes only how an SPM artifact the bundle no longer carries is reported. + + To calculate a saved household as the Public API's stored-household + route (``GET /{country}/household/{id}/policy/{policy_id}``) does, set + ``add_missing_variables``, because that route always adds missing + variables; pass the household's saved SPM selection, if it has one, as + ``spm``, so it counts as chosen; and set ``spm_selection_source`` to + ``"saved_household"``. """ country: CountryId