From 6b1fcf41ccbfe40cb11abf027ccb8bbbc93b56b8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:11:38 +0400 Subject: [PATCH 01/19] Add a versioned contract for household calculations on Modal The v2 household functions take one household in the Public API's situation shape, with optional axes, and one to ten policies, so a baseline and its reform are calculated in one call. The request names the PolicyEngine.py bundle it expects, whether to add missing variables (as the Public API's /calculate-full route does), and an optional SPM selection whose omission is kept distinct from an empty choice. A household keeps its caller's key order through validation and transport, at the top level and in each group, because the Public API walks a household in that order and the order of its input errors, warnings and values follows it. The result carries each policy's calculated household, warnings and SPM receipt, the applied SPM selection and the bundle provenance; a failure carries a typed error whose messages and input errors match the Public API's error responses. Co-Authored-By: Claude Opus 5.5 --- .../household_execution.py | 657 ++++++++++++++++++ .../tests/test_household_execution.py | 539 ++++++++++++++ 2 files changed, 1196 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..df1f17be9 --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/household_execution.py @@ -0,0 +1,657 @@ +"""Versioned household calculation contract 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 except ``axes`` 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 input errors as the Public API. + +Wire models: + +- ``HouseholdExecutionRequest``: the caller to a Modal household function. It + names one to ten policies, so a baseline and its reform are calculated in + one call. +- ``HouseholdExecutionResult`` or ``HouseholdExecutionFailure``: what the + household function returns, told apart by ``status``. +""" + +from __future__ import annotations + +import math +from collections.abc import Mapping +from typing import TYPE_CHECKING, Annotated, Any, Literal, Self +from uuid import UUID + +from pydantic import ( + AfterValidator, + BaseModel, + BeforeValidator, + ConfigDict, + Field, + PrivateAttr, + TypeAdapter, + field_validator, + model_serializer, + model_validator, +) + +from policyengine_simulation_contract.spm import ( + SPMErrorDetail, + SPMProvenance, + SPMResolvedConfiguration, + SPMSelection, +) +from policyengine_simulation_contract.stage12_bundle import ( + CountryId, + Sha256Digest, + Stage12BundleManifest, +) + +if TYPE_CHECKING: + from pydantic import ModelWrapValidatorHandler, SerializerFunctionWrapHandler + +HOUSEHOLD_CONTRACT_SCHEMA_VERSION = 1 +# A v2 worker application's household function for a country is this prefix +# followed by the country ID, such as ``calculate_household_us``. +HOUSEHOLD_CALLABLE_PREFIX = "calculate_household_" +# Size limits. 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'; its result carries the same name." + ), + ), +] + +# 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. A household with axes holds one value per point. +# Infinite and undefined numbers are the strings "Infinity", "-Infinity" and +# "NaN"; the Public API writes infinities the same way. +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), +] +ParameterScalar = bool | int | float | str | None +# A list value sets a parameter whose value is a list, such as a list of +# variable names. +ParameterValue = ParameterScalar | list[ParameterScalar] + +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) + + +class CallerOrderedModel(BaseModel): + """A model that serializes its top-level keys in the order its caller gave them. + + The Public API walks a household in the order the caller sent it, so the + order of unrecognized inputs, warnings and calculated values depends on + it. Pydantic itself would write the declared fields before the extra ones. + Keys the caller did not give follow, in Pydantic's order. + """ + + model_config = ConfigDict(extra="allow", allow_inf_nan=False, frozen=True) + + _caller_key_order: tuple[str, ...] = PrivateAttr(default=()) + + @model_validator(mode="wrap") + @classmethod + def remember_caller_key_order( + cls, data: object, handler: ModelWrapValidatorHandler[Self] + ) -> Self: + """Record the order of the caller's top-level keys for serialization.""" + validated: Self = handler(data) + if isinstance(data, Mapping): + validated._caller_key_order = tuple(str(key) for key in data) + elif isinstance(data, CallerOrderedModel): + validated._caller_key_order = data.caller_key_order + return validated + + @model_serializer(mode="wrap") + def serialize_in_caller_order( + self, handler: SerializerFunctionWrapHandler + ) -> dict[str, Any]: + """Return the caller's keys in the caller's order, then any others.""" + dumped: dict[str, Any] = handler(self) + ordered = {key: dumped[key] for key in self._caller_key_order if key in dumped} + ordered.update(dumped) + left_out = self.keys_left_out() + return {key: value for key, value in ordered.items() if key not in left_out} + + @property + def caller_key_order(self) -> tuple[str, ...]: + """The caller's top-level keys, in the order the caller gave them.""" + return self._caller_key_order + + def keys_left_out(self) -> frozenset[str]: + """Return the keys serialization leaves out; none unless a subclass says.""" + return frozenset() + + +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(CallerOrderedModel): + """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. + """ + + members: GroupMembers + __pydantic_extra__: dict[VariableName, VariableInput] = Field(init=False) # pyright: ignore[reportIncompatibleVariableOverride] # Pydantic's documented way to type extra fields + + @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(CallerOrderedModel): + """One calculated group entity instance: its members and calculated variables.""" + + members: GroupMembers + __pydantic_extra__: dict[VariableName, VariableOutput] = Field(init=False) # pyright: ignore[reportIncompatibleVariableOverride] # Pydantic's documented way to type extra fields + + @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: int | float + max: int | 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 the optional keys the caller did not give.""" + dumped: dict[str, Any] = handler(self) + return { + name: value + for name, value in dumped.items() + if name not in {"period", "index"} or name in self.model_fields_set + } + + +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), +] + + +# One group entity collection: each group of one entity, by the group's name. +InputGroups = Annotated[ + dict[EntityKey, HouseholdGroup], Field(max_length=MAXIMUM_GROUPS_PER_COLLECTION) +] +CalculatedGroups = Annotated[ + dict[EntityKey, CalculatedHouseholdGroup], + Field(max_length=MAXIMUM_GROUPS_PER_COLLECTION), +] + + +class _Situation(CallerOrderedModel): + """A situation-shaped household: ``people``, group collections and ``axes``.""" + + 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 + + def keys_left_out(self) -> frozenset[str]: + """Leave out ``axes`` when the household has none.""" + return frozenset({"axes"}) if self.axes is None else frozenset() + + @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") + + +class HouseholdDocument(_Situation): + """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, InputGroups] = Field(init=False) # pyright: ignore[reportIncompatibleVariableOverride] # Pydantic's documented way to type extra fields + + +class CalculatedHousehold(_Situation): + """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, CalculatedGroups] = Field(init=False) # pyright: ignore[reportIncompatibleVariableOverride] # Pydantic's documented way to type extra fields + + +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 + + @classmethod + def from_bundle( + cls, + bundle: Stage12BundleManifest, + *, + bundle_manifest_sha256: str, + country: str, + ) -> Self: + """Return ``country``'s provenance within a normalized bundle manifest. + + Raises ValueError when the bundle does not carry ``country``. + """ + country_bundle = next( + (item for item in bundle.countries if item.country == country), None + ) + if country_bundle is None: + msg = f"the bundle does not carry country {country!r}" + raise ValueError(msg) + return cls( + policyengine_version=bundle.policyengine_version, + core_package_version=bundle.core_package_version, + country_package_name=country_bundle.country_package_name, + country_package_version=country_bundle.country_package_version, + bundle_manifest_sha256=bundle_manifest_sha256, + ) + + +def household_callable_name(country: str) -> str: + """Return the name of a v2 worker application's household function.""" + return f"{HOUSEHOLD_CALLABLE_PREFIX}{country}" + + +class HouseholdExecutionRequest(StrictHouseholdModel): + """A caller to a Modal household function: calculate one household. + + ``calculation_id`` is the caller's durable identity for the calculation. + ``bundle`` is the bundle the caller resolved from the version manifest; the + household function refuses a request for any other bundle. + + ``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 uncalculated value. ``spm_selection_source`` says + whether the selection came from this request or from a saved household, + which changes only how an SPM forecast 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"``. + """ + + schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION + calculation_id: UUID + country: CountryId + household: HouseholdDocument + policies: Policies + add_missing_variables: bool = False + spm: SPMChoice = None + spm_selection_source: Literal["request", "saved_household"] = "request" + bundle: HouseholdBundleProvenance + + @model_serializer(mode="wrap") + def serialize_request( + 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 HouseholdPolicyResult(StrictHouseholdModel): + """The calculated household under one policy. + + ``warnings`` lists, in order, the deprecated inputs that were ignored and, + for a household with axes, the values that could not be calculated, worded + as the Public API words them. ``spm_provenance`` is the SPM calculation + receipt of 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. + + ``spm_config`` is the complete SPM selection the calculation applied: the + caller's choice merged over the bundle's certified defaults, or None when + no SPM applies (any country but the US). + """ + + country: CountryId + bundle: HouseholdBundleProvenance + spm_config: SPMResolvedConfiguration | None = None + results: Annotated[ + tuple[HouseholdPolicyResult, ...], + Field(min_length=1, max_length=MAXIMUM_POLICIES_PER_REQUEST), + ] + + @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, which name entity plurals 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 + + @model_serializer(mode="wrap") + def serialize_input_error( + self, handler: SerializerFunctionWrapHandler + ) -> dict[str, Any]: + """Leave out unset entity plurals, as the Public API's error entries do.""" + dumped: dict[str, Any] = handler(self) + for name in ("expected_entity_plural", "actual_entity_plural"): + if dumped.get(name) is None: + dumped.pop(name, None) + return dumped + + +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): + """A Modal household function to its caller: 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): + """A Modal household function to its caller: the calculation failed.""" + + 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) 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..3d1ef4f2c --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_household_execution.py @@ -0,0 +1,539 @@ +"""Tests for the versioned household calculation contract.""" + +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, + HouseholdBundleProvenance, + HouseholdCalculationError, + HouseholdCalculationOutput, + HouseholdDocument, + HouseholdExecutionFailure, + HouseholdExecutionRequest, + HouseholdExecutionResult, + HouseholdInputError, + HouseholdPolicy, + HouseholdPolicyResult, + HouseholdWorkerIdentity, + household_callable_name, + parse_household_execution_outcome, +) +from policyengine_simulation_contract.spm import ( + SPMErrorDetail, + SPMResolvedConfiguration, + SPMRuntimeVersions, + SPMSelection, + build_spm_provenance, +) +from policyengine_simulation_contract.stage12_bundle import ( + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, +) +from pydantic import ValidationError + +CALCULATION_ID = UUID("00000000-0000-4000-8000-0000000000c1") + +# 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", + } + ] + ], +} +# The same household with its top-level keys, and one group's keys, in +# another order. +REORDERED_SITUATION = { + "households": { + "your household": { + "state_name": {"2026": "CA"}, + "members": ["you", "your partner"], + "household_net_income": {"2026": None}, + } + }, + "axes": US_SITUATION["axes"], + "people": US_SITUATION["people"], + "tax_units": US_SITUATION["tax_units"], +} +CTC_REFORM = { + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01.2100-12-31": 3_000} +} +SPM_RUNTIME_VERSIONS = SPMRuntimeVersions.model_validate( + { + "policyengine": "0.0.0-test-policyengine", + "policyengine-core": "0.0.0-test-policyengine-core", + "policyengine-us": "0.0.0-test-policyengine-us", + "spm-calculator": "0.0.0-test-spm-calculator", + } +) +RESOLVED_SPM_SELECTION = SPMSelection( + forecast_content_sha256="c" * 64, + scenario="ce_trend", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, +) + + +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() -> list[dict[str, object]]: + return [ + {"policy_key": "baseline", "parameters": {}}, + {"policy_key": "reform", "parameters": CTC_REFORM}, + ] + + +def _request(**overrides: object) -> HouseholdExecutionRequest: + values: dict[str, object] = { + "calculation_id": str(CALCULATION_ID), + "country": "us", + "household": US_SITUATION, + "policies": _policies(), + "bundle": _bundle().model_dump(mode="json"), + } + 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 _calculated(situation: dict) -> CalculatedHousehold: + return CalculatedHousehold.model_validate(situation) + + +def _output(**overrides: object) -> HouseholdCalculationOutput: + values: dict[str, object] = { + "country": "us", + "bundle": _bundle(), + "results": ( + HouseholdPolicyResult( + policy_key="baseline", + household=_calculated( + { + "people": {"you": {"age": {"2026": 40}}}, + "households": {"your household": {"members": ["you"]}}, + } + ), + warnings=("Input `x` was ignored.",), + ), + ), + } + values.update(overrides) + return HouseholdCalculationOutput.model_validate(values) + + +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.has_axes + assert document.axis_point_count() == 11 + + +def test_households_keep_the_callers_key_order_through_transport() -> None: + # The Public API walks a household in the order its caller sent it, so the + # order of unrecognized inputs, warnings and calculated values follows it. + request = _request(household=REORDERED_SITUATION) + transported = HouseholdExecutionRequest.model_validate( + json.loads(request.model_dump_json()) + ) + + for situation in ( + request.household.to_situation(), + transported.household.to_situation(), + json.loads(transported.household.model_dump_json()), + _calculated(REORDERED_SITUATION).to_situation(), + ): + assert list(situation) == list(REORDERED_SITUATION) + assert list(situation["households"]["your household"]) == [ + "state_name", + "members", + "household_net_income", + ] + assert situation == REORDERED_SITUATION + + +def test_a_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 document.axis_point_count() == 1 + assert document.to_situation() == situation + assert "axes" not in json.loads(document.model_dump_json()) + + +def test_axes_keep_the_values_and_keys_the_caller_sent() -> None: + situation = { + "people": {"you": {}}, + "axes": [[{"name": "employment_income", "count": 3, "min": 0, "max": 2.5}]], + } + + assert HouseholdDocument.model_validate(situation).to_situation() == situation + + +def test_a_household_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": {}}, + "households": {"your household": {"members": ["you", "you"]}}, + }, + {"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}]}, + {"people": {"you": {}}, "axes": [[{"name": "age", "count": 0}]]}, + ], +) +def test_a_household_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_carries_the_baseline_and_its_reform() -> None: + request = _request() + + assert [policy.policy_key for policy in request.policies] == ["baseline", "reform"] + assert request.policies[1].parameters == CTC_REFORM + assert request.schema_version == 1 + assert request.add_missing_variables is False + assert request.spm_selection_source == "request" + assert request.bundle == _bundle() + + +def test_add_missing_variables_asks_for_a_full_calculation() -> None: + assert _request(add_missing_variables=True).add_missing_variables is True + + +def test_a_request_needs_one_to_ten_policies_with_unique_keys() -> None: + eleven_policies = [ + {"policy_key": f"policy-{index}", "parameters": {}} + for index in range(MAXIMUM_POLICIES_PER_REQUEST + 1) + ] + repeated_key = [ + {"policy_key": "reform", "parameters": {}}, + {"policy_key": "reform", "parameters": CTC_REFORM}, + ] + + assert len(_request(policies=eleven_policies[:10]).policies) == 10 + with pytest.raises(ValidationError): + _request(policies=[]) + with pytest.raises(ValidationError): + _request(policies=eleven_policies) + with pytest.raises(ValidationError, match="unique"): + _request(policies=repeated_key) + with pytest.raises(ValidationError): + _request(policies=[{"policy_key": "Not A Key", "parameters": {}}]) + + +def test_omitted_spm_inherits_the_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={}, spm_selection_source="saved_household") + + assert request.spm_requested is True + dumped = request.model_dump(mode="json") + assert dumped["spm"] == {} + assert dumped["spm_selection_source"] == "saved_household" + assert HouseholdExecutionRequest.model_validate(dumped).spm_requested is True + + +def test_an_explicit_null_spm_is_rejected() -> None: + with pytest.raises(ValidationError, match="omit it"): + _request(spm=None) + + +def test_a_request_rejects_unknown_fields_and_other_schema_versions() -> None: + with pytest.raises(ValidationError): + _request(unexpected=True) + with pytest.raises(ValidationError): + _request(schema_version=2) + with pytest.raises(ValidationError): + _request(country="ca") + + +def test_a_calculated_household_keeps_axis_results_and_non_finite_strings() -> None: + calculated = { + "people": { + "you": { + "age": {"2026": [40, 40]}, + "marginal_tax_rate": {"2026": ["Infinity", 0.25]}, + "is_blind": {"2026": [False, False]}, + } + }, + "households": { + "your household": { + "members": ["you"], + "state_name": {"2026": ["CA", "CA"]}, + "spm_unit_net_income": {"2026": [None, None]}, + } + }, + "axes": [[{"name": "employment_income", "count": 2, "min": 0, "max": 1}]], + } + + assert _calculated(calculated).to_situation() == calculated + + +def test_execution_outcomes_are_told_apart_by_status() -> None: + succeeded = HouseholdExecutionResult( + calculation_id=CALCULATION_ID, output=_output(), worker=_worker() + ) + failed = HouseholdExecutionFailure( + calculation_id=CALCULATION_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", + ), + ), + worker=_worker(), + ) + + parsed_success = parse_household_execution_outcome( + succeeded.model_dump(mode="json") + ) + parsed_failure = parse_household_execution_outcome(failed.model_dump(mode="json")) + + assert isinstance(parsed_success, HouseholdExecutionResult) + assert parsed_success == succeeded + assert parsed_success.schema_version == 1 + assert isinstance(parsed_failure, HouseholdExecutionFailure) + assert parsed_failure == failed + with pytest.raises(ValidationError): + parse_household_execution_outcome({"status": "running"}) + + +def test_input_errors_serialize_as_the_public_api_lists_them() -> None: + wrong_entity = 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", + ) + unknown = HouseholdInputError( + type="household_variable", + name="not_a_variable", + path="household.people.you.not_a_variable", + message="Unrecognized household variable ...", + ) + + assert wrong_entity.model_dump(mode="json") == { + "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", + } + # Like the Public API's error entries, an error names entity plurals only + # for a variable on the wrong entity. + assert unknown.model_dump(mode="json") == { + "type": "household_variable", + "name": "not_a_variable", + "path": "household.people.you.not_a_variable", + "message": "Unrecognized household variable ...", + } + assert HouseholdInputError.model_validate(unknown.model_dump()) == unknown + + +def test_output_results_must_have_unique_policy_keys() -> None: + result = _output().results[0] + + with pytest.raises(ValidationError, match="unique"): + _output(results=(result, result)) + + +def test_output_carries_the_applied_spm_selection_and_each_receipt() -> None: + receipt = build_spm_provenance( + forecast_id="forecast-2026", + forecast_sha256="c" * 64, + selection=RESOLVED_SPM_SELECTION, + years=("2026",), + runtime_versions=SPM_RUNTIME_VERSIONS, + ) + applied = SPMResolvedConfiguration.model_validate( + RESOLVED_SPM_SELECTION.model_dump(mode="json") + ) + result = _output().results[0].model_copy(update={"spm_provenance": receipt}) + + output = _output(results=(result,), spm_config=applied) + transported = HouseholdCalculationOutput.model_validate( + json.loads(output.model_dump_json()) + ) + + assert transported.spm_config == applied + assert transported.results[0].spm_provenance == receipt + assert transported.model_dump(mode="json")["spm_config"] == { + "forecast_content_sha256": "c" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + } + assert _output().spm_config is None + + +def test_bundle_provenance_is_read_from_a_normalized_bundle() -> None: + dataset = Stage12Dataset( + identity="populace_uk_2023", + uri="hf://policyengine/data/populace_uk_2023.h5@r1", + artifact_revision="r1", + sha256="c" * 64, + repo_type="dataset", + ) + bundle = Stage12BundleManifest( + policyengine_version="6.7.5", + policyengine_requirement="policyengine==6.7.5", + core_package_version="3.33.0", + core_package_requirement="policyengine-core==3.33.0", + countries=( + Stage12CountryBundle( + country="uk", + country_package_name="policyengine-uk", + country_package_version="2.130.1", + country_package_requirement="policyengine-uk==2.130.1", + data_package_name="populace-data", + data_package_version="0.1.0", + data_release_version="r1", + data_artifact_revision="r1", + default_dataset="populace_uk_2023", + default_dataset_uri=dataset.uri, + datasets=(dataset,), + ), + ), + ) + + provenance = HouseholdBundleProvenance.from_bundle( + bundle, bundle_manifest_sha256="d" * 64, country="uk" + ) + + assert provenance.model_dump() == { + "policyengine_version": "6.7.5", + "core_package_version": "3.33.0", + "country_package_name": "policyengine-uk", + "country_package_version": "2.130.1", + "bundle_manifest_sha256": "d" * 64, + } + with pytest.raises(ValueError, match="does not carry"): + HouseholdBundleProvenance.from_bundle( + bundle, bundle_manifest_sha256="d" * 64, country="us" + ) + + +def test_household_functions_are_named_by_country() -> None: + assert household_callable_name("us") == "calculate_household_us" + assert household_callable_name("uk") == "calculate_household_uk" + + +def test_policy_values_may_be_lists_for_list_parameters() -> None: + policy = HouseholdPolicy.model_validate( + { + "policy_key": "reform", + "parameters": { + "gov.usda.snap.income.sources.earned": { + "2026-01-01.2100-12-31": ["employment_income"] + } + }, + } + ) + + assert policy.parameters["gov.usda.snap.income.sources.earned"] == { + "2026-01-01.2100-12-31": ["employment_income"] + } From 03cda9bc00139d5bce679aba30b39efadab22b02 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 21:21:27 +0400 Subject: [PATCH 02/19] Record each country's capabilities in v2 manifest schema version 2 Schema version 2 of the v2 version manifest records, for each country a worker application serves, whether it is economy-capable (it names a single-simulation worker) and which function calculates its households, and the validation receipt names an invocation for each of those functions. The Simulation Entrypoint's comparison adapter now accepts only a country the manifest marks as economy-capable, so the report coordinator never receives a report for a country that only calculates households. Deployed and rollback revisions read schema version 1 and reject unknown fields, so the publisher writes schema version 2 under a new key and keeps writing the schema version 1 forms (with region datasets, without them, and without cache releases) under the existing keys whenever schema version 1 can describe the manifest. Current readers use the schema version 2 document while the older keys hold its schema version 1 forms, or, for a manifest schema version 1 cannot describe, unless the older keys hold a later validation; otherwise they read the older keys as before, as economy-capable countries without household functions. Code released before schema version 2 that publishes after a rollback therefore leaves no stale document behind, and the next publish starts from what it wrote. The publisher records every bundle country as economy-capable. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 49 ++- .../stage12_manifest.py | 293 ++++++++++++--- .../tests/test_stage12_manifest.py | 336 +++++++++++++++--- .../stage12_adapter.py | 20 +- .../tests/stage12_fixtures.py | 23 +- .../tests/test_stage12_adapter.py | 11 + .../tests/test_stage12_backend.py | 57 ++- .../modal/utils/update_v2_version_manifest.py | 8 +- .../tests/test_update_v2_version_manifest.py | 17 + 9 files changed, 686 insertions(+), 128 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 953ca0d70..0a2eddcbd 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -396,20 +396,38 @@ 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 economy-capable (it has a single-simulation worker, so +its reports may be given to the report coordinator) and which function +calculates its households. The Simulation Entrypoint's comparison adapter +skips any country the manifest does not mark as economy-capable, with the +reason `unsupported_country`, so the coordinator receives only those +countries' reports. + +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 whole manifest under the -`active-with-region-datasets` key, the same manifest without each country -bundle's `region_datasets` under the `active-with-cache` key, and without cache -releases either under the `active` key; readers released before each field keep -reading the key that lacks it. Current readers use the -`active-with-region-datasets` copy only while it, with its region datasets -removed, equals the manifest under `active-with-cache`; otherwise they read -`active-with-cache`, falling back to `active` before the first publish that -writes it. Code released before region datasets republishes only the two older -keys, so after a rollback the out-of-date region dataset copy stops matching and -is no longer read. Removing the older copies needs a later change, once no -deployed or rollback revision reads them. +promoted. The publisher therefore writes the whole schema version 2 manifest +under the `active-schema-2` key and, in schema version 1 form (without the +capability and household fields), the same manifest under the +`active-with-region-datasets` key, without each country bundle's +`region_datasets` under the `active-with-cache` key, and without cache releases +either under the `active` key; readers released before each field keep reading +the key that lacks it. A schema version 1 document marks every listed country +economy-capable, with no household function. + +Current readers use the `active-schema-2` document while each older key holds +its schema version 1 form. Schema version 1 cannot describe a country that is +not economy-capable, so a manifest holding one is written only under +`active-schema-2` and the older keys keep the last manifest written there; such +a document is used unless the document under the older keys was validated +later. Otherwise readers use the `active-with-region-datasets` copy only while +it, with its region datasets removed, equals the manifest under +`active-with-cache`; otherwise they read `active-with-cache`, falling back to +`active` before the first publish that writes it. Code released before schema +version 2 or before region datasets republishes only the older keys it knows, so +after a rollback the out-of-date copies stop matching and are no longer read, +and the next publish starts from what the older code wrote. Removing the older +copies needs a later change, once no deployed or rollback revision reads them. A rollback to code released before region datasets has one window: the older worker is deployed before its publish job runs, and until that publish the @@ -417,8 +435,9 @@ region dataset copy still matches, so updated entry services keep sending regional provenance for state and district runs, which the older worker refuses. Only Stage 12 comparison runs fail in that window, never production responses; if the older deploy's validation fails, the window lasts until a later publish. -To avoid it, delete the `active-with-region-datasets` key from the -`simulation-api-v2-version-manifest` Modal Dict before deploying the older code. +To avoid it, delete the `active-schema-2` and `active-with-region-datasets` keys +from the `simulation-api-v2-version-manifest` Modal Dict before deploying the +older code. `region_datasets` names the dataset each region type loads. The entry service records that dataset as a state or district run's dataset provenance, and the 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 a10d414cc..4cb9562b6 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 economy reports and which function calculates its +households. It is stored under its own key. Whenever schema version 1 can +describe the same manifest, it is also written in that form, without those +fields, under the keys that readers released before schema version 2 read: +with region datasets, without them, and without them or cache releases. +""" from __future__ import annotations @@ -8,7 +16,7 @@ from copy import deepcopy from datetime import datetime from hashlib import sha256 -from typing import Annotated, Literal, Protocol +from typing import TYPE_CHECKING, Annotated, Literal, Protocol from packaging.version import InvalidVersion, Version from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -19,16 +27,24 @@ Stage12BundleManifest, ) +if TYPE_CHECKING: + from pydantic.main import IncEx + V1_ROUTING_STATE_NAME = "simulation-api-routing-state" V2_VERSION_MANIFEST_NAME = "simulation-api-v2-version-manifest" -# The complete manifest, including each country bundle's region datasets. +# 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 read +# the keys below, which hold the same manifest in schema version 1 form. This +# one keeps each country bundle's region datasets. REGION_DATASETS_MANIFEST_KEY = "active-with-region-datasets" # Readers released before region datasets forbid unknown fields, so they keep -# reading this key, which holds the same manifest without them. +# reading this key, which holds the schema version 1 form without them. 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 or -# region datasets. +# they keep reading this key, which holds the schema version 1 form without +# the cache fields or region datasets. PRE_CACHE_MANIFEST_KEY = "active" V2_APPLICATION_PREFIX = "policyengine-simulation-v2-py" V1_APPLICATION_PREFIX = "policyengine-simulation-py" @@ -46,10 +62,35 @@ 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. + + An economy-capable country names the callable that runs one of its report + simulations; only such a country's reports are given to the report + coordinator. ``household_callable`` names the function that calculates + its households, and is None when the application has none or it was not + validated. + """ country: CountryId - single_simulation_callable: CallableName + economy_capable: bool + single_simulation_callable: CallableName | None = None + household_callable: CallableName | None = None + + @model_validator(mode="after") + def validate_capabilities(self) -> V2CountryWorker: + """Require a simulation callable exactly when the country is economy-capable. + + A country must also have at least one capability. + """ + if self.economy_capable != (self.single_simulation_callable is not None): + msg = ( + "an economy-capable country, and only one, names a simulation callable" + ) + raise ValueError(msg) + if not self.economy_capable and self.household_callable is None: + msg = "a country needs at least one capability" + raise ValueError(msg) + return self class Stage12CacheRelease(StrictManifestModel): @@ -71,7 +112,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 economy-capable country's + simulation worker, and for each country's household function. """ validated: Literal[True] @@ -80,14 +122,36 @@ 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 +def _require_function_receipts( + countries: tuple[V2CountryWorker, ...], validation: V2WorkerValidation +) -> None: + """Require a validation invocation for exactly the functions countries name.""" + economy_countries = { + worker.country for worker in countries if worker.economy_capable + } + if set(validation.country_validation_invocation_ids) != economy_countries: + msg = "worker validation does not cover exactly the economy-capable countries" + raise ValueError(msg) + household_countries = { + worker.country for worker in countries if worker.household_callable is not None + } + if set(validation.household_validation_invocation_ids) != household_countries: + msg = "worker validation does not cover exactly the household functions" + raise ValueError(msg) + + 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 function a country names must have been validated. """ application_name: ManifestText @@ -105,7 +169,7 @@ def validate_worker(self) -> V2WorkerVersion: The application name must be the v2 name derived from the bundle's PolicyEngine.py version, the validation must name this application and bundle digest, and the countries must equal the bundle's, in order, - each with a validation invocation. + each named function with a validation invocation. """ if not self.application_name.startswith(V2_APPLICATION_PREFIX): msg = "v2 application uses an invalid name prefix" @@ -134,11 +198,31 @@ 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" - raise ValueError(msg) + _require_function_receipts(self.countries, self.validation) return self + def economy_worker(self, country: str) -> V2CountryWorker | None: + """Return ``country``'s worker when it can run economy reports.""" + return next( + ( + worker + for worker in self.countries + if worker.country == country and worker.economy_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. @@ -147,7 +231,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] @@ -224,37 +308,106 @@ def _is_newer(candidate: str, current: str | None) -> bool: return False -_REGION_DATASETS_FIELD = { - "bundle": {"countries": {"__all__": {"region_datasets"}}}, -} +def schema_1_manifest_payload( + manifest: V2VersionManifest, + *, + include_cache: bool = True, + include_region_datasets: bool = True, +) -> dict[str, object] | None: + """Return the manifest as readers released before schema version 2 accept it. + + Without ``include_region_datasets`` each country bundle's region datasets + are left out, for readers released before them; without ``include_cache`` + the cache fields are left out too, for readers released before cache + releases. Returns None when a country is not economy-capable, which + schema version 1 cannot describe. + """ + if any( + not country.economy_capable + for worker in manifest.versions.values() + for country in worker.countries + ): + return None + validation_exclude = {"household_validation_invocation_ids"} + worker_exclude: dict[str, IncEx | bool] = { + "countries": {"__all__": {"economy_capable", "household_callable"}}, + "validation": validation_exclude, + } + if not include_region_datasets: + worker_exclude["bundle"] = {"countries": {"__all__": {"region_datasets"}}} + if not include_cache: + worker_exclude["cache"] = True + validation_exclude.add("cache") + payload = manifest.model_dump( + mode="json", + exclude={"versions": {"__all__": worker_exclude}}, + ) + payload["schema_version"] = 1 + return payload def pre_region_datasets_manifest_payload( manifest: V2VersionManifest, -) -> dict[str, object]: +) -> dict[str, object] | None: """Return the manifest as readers released before region datasets accept it.""" - return manifest.model_dump( - mode="json", - exclude={"versions": {"__all__": _REGION_DATASETS_FIELD}}, - ) + return schema_1_manifest_payload(manifest, include_region_datasets=False) -def pre_cache_manifest_payload(manifest: V2VersionManifest) -> dict[str, object]: +def pre_cache_manifest_payload(manifest: V2VersionManifest) -> dict[str, object] | None: """Return the manifest as readers released before cache releases accept it.""" - return manifest.model_dump( - mode="json", - exclude={ - "versions": { - "__all__": { - "cache": True, - "validation": {"cache"}, - **_REGION_DATASETS_FIELD, - } - } - }, + return schema_1_manifest_payload( + manifest, include_cache=False, include_region_datasets=False ) +def _schema_1_copies( + manifest: V2VersionManifest, +) -> dict[str, dict[str, object]] | None: + """Return the schema version 1 documents a publish writes, by key. + + Returns None when schema version 1 cannot describe the manifest. + """ + complete = schema_1_manifest_payload(manifest) + without_region_datasets = pre_region_datasets_manifest_payload(manifest) + without_cache = pre_cache_manifest_payload(manifest) + if complete is None or without_region_datasets is None or without_cache is None: + return None + return { + REGION_DATASETS_MANIFEST_KEY: complete, + ACTIVE_MANIFEST_KEY: without_region_datasets, + PRE_CACHE_MANIFEST_KEY: without_cache, + } + + +def _upgraded_schema_1_document(document: Mapping[str, object]) -> dict[str, object]: + """Return a schema version 1 document in schema version 2 form. + + Schema version 1 listed only countries whose report simulations a worker + runs, and no household functions, so every country is economy-capable. + """ + upgraded = deepcopy(dict(document)) + upgraded["schema_version"] = 2 + versions = upgraded.get("versions") + if isinstance(versions, dict): + for worker in versions.values(): + countries = worker.get("countries") if isinstance(worker, dict) else None + if isinstance(countries, list): + for country in countries: + if isinstance(country, dict): + country.setdefault("economy_capable", True) + return upgraded + + +def parse_v2_manifest(value: object) -> V2VersionManifest: + """Validate a stored manifest, reading a schema version 1 document as version 2. + + Raises ValueError (as a pydantic ValidationError) for an invalid document. + """ + if isinstance(value, Mapping) and value.get("schema_version") == 1: + value = _upgraded_schema_1_document(value) + return V2VersionManifest.model_validate(value) + + def _is_current_region_dataset_copy(value: object, cache_copy: object) -> bool: """Return whether ``value`` is the copy the last publish wrote with ``cache_copy``. @@ -268,13 +421,14 @@ def _is_current_region_dataset_copy(value: object, cache_copy: object) -> bool: if cache_copy is None: return True try: - manifest = V2VersionManifest.model_validate(value) + manifest = parse_v2_manifest(value) except (TypeError, ValueError): return False return pre_region_datasets_manifest_payload(manifest) == cache_copy -def _current_manifest(store: V2ManifestStore) -> object: +def _current_schema_1_manifest(store: V2ManifestStore) -> object: + """Return the schema version 1 document the last publish wrote.""" # The newest key the last publish wrote: each older key holds the same # manifest without the fields its readers do not accept. cache_copy = store.get(ACTIVE_MANIFEST_KEY) @@ -287,6 +441,53 @@ def _current_manifest(store: V2ManifestStore) -> object: return store.get(PRE_CACHE_MANIFEST_KEY) +def _newest_validation(manifest: V2VersionManifest) -> datetime: + return max(worker.validation.validated_at for worker in manifest.versions.values()) + + +def _is_current_schema_2_document(store: V2ManifestStore, document: object) -> bool: + """Return whether ``document``, under the schema version 2 key, is the latest. + + Code released before schema version 2 publishes only the older keys, for + example after a rollback, which leaves an out-of-date schema version 2 + document behind. Whenever schema version 1 can describe a manifest, its + publish writes the schema version 1 forms to the older keys, so the + document is current while each older key holds its form. Otherwise (a + country that is not economy-capable) the publish left the older keys + alone, and the document is current unless the document under them was + validated later. An unreadable document counts as current, so readers and + publishers report it. + """ + try: + manifest = parse_v2_manifest(document) + except (TypeError, ValueError): + return True + schema_1_copies = _schema_1_copies(manifest) + if schema_1_copies is not None: + return all( + store.get(key) in (None, payload) + for key, payload in schema_1_copies.items() + ) + older = _current_schema_1_manifest(store) + if older is None: + return True + try: + older_manifest = parse_v2_manifest(older) + except (TypeError, ValueError): + return True + return _newest_validation(older_manifest) <= _newest_validation(manifest) + + +def _current_manifest(store: V2ManifestStore) -> object: + # Until the first schema version 2 publish, only the older keys exist. + schema_2_document = store.get(SCHEMA_2_MANIFEST_KEY) + if schema_2_document is not None and _is_current_schema_2_document( + store, schema_2_document + ): + return schema_2_document + return _current_schema_1_manifest(store) + + def build_next_v2_manifest( *, current: object, @@ -303,7 +504,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 @@ -330,8 +531,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), @@ -341,9 +544,11 @@ def publish_v2_manifest( # Each key holds one entire validated document, and a reader uses one # whole document (the region dataset copy only while it matches the cache # copy), so no reader observes a partial update assembled from several. - store[REGION_DATASETS_MANIFEST_KEY] = next_manifest.model_dump(mode="json") - store[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(next_manifest) - store[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(next_manifest) + store[SCHEMA_2_MANIFEST_KEY] = next_manifest.model_dump(mode="json") + schema_1_copies = _schema_1_copies(next_manifest) + if schema_1_copies is not None: + for key, payload in schema_1_copies.items(): + store[key] = payload return next_manifest @@ -363,7 +568,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 78c1f908e..cd7bccfe7 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -17,6 +17,7 @@ ACTIVE_MANIFEST_KEY, PRE_CACHE_MANIFEST_KEY, REGION_DATASETS_MANIFEST_KEY, + SCHEMA_2_MANIFEST_KEY, V1_ROUTING_STATE_NAME, V2_VERSION_MANIFEST_NAME, Stage12CacheRelease, @@ -26,9 +27,11 @@ V2WorkerVersion, assert_separate_manifest_names, manifest_digest, + parse_v2_manifest, pre_cache_manifest_payload, pre_region_datasets_manifest_payload, publish_v2_manifest, + schema_1_manifest_payload, v2_application_name, validate_manifest_has_no_credentials, ) @@ -90,22 +93,39 @@ def _resolved_bundle(version: str = "5.2.0") -> ResolvedStage12Bundle: ) -def _worker(version: str = "5.2.0") -> V2WorkerVersion: +def _country_worker( + country: str, *, economy: bool = True, household: bool = False +) -> V2CountryWorker: + return V2CountryWorker( + country=country, + economy_capable=economy, + single_simulation_callable=( + f"run_single_simulation_{country}" if economy else None + ), + household_callable=f"calculate_household_{country}" if household else None, + ) + + +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( + _country_worker( + country, + economy=country not in household_only, + household=with_household_functions or country in household_only, + ) + 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( @@ -115,8 +135,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.economy_capable + }, + household_validation_invocation_ids={ + worker.country: f"household-validation-{worker.country}-123" + for worker in countries + if worker.household_callable is not None }, ), ) @@ -127,8 +153,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, @@ -163,9 +191,11 @@ def test_v2_application_name_rejects_versions_a_modal_name_cannot_hold( v2_application_name(version) -def _regional_worker(version: str = "5.2.0") -> V2WorkerVersion: +def _regional_worker( + version: str = "5.2.0", *, with_household_functions: bool = False +) -> V2WorkerVersion: """Return a cached worker whose US bundle names each region type's dataset.""" - worker = _cached_worker(version) + worker = _cached_worker(version, with_household_functions=with_household_functions) countries = tuple( country.model_copy( update={ @@ -191,31 +221,82 @@ def _regional_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=_regional_worker()) + manifest = publish_v2_manifest( + store=store, worker=_regional_worker(with_household_functions=True) + ) assert store.writes == [ - (REGION_DATASETS_MANIFEST_KEY, manifest.model_dump(mode="json")), + (SCHEMA_2_MANIFEST_KEY, manifest.model_dump(mode="json")), + (REGION_DATASETS_MANIFEST_KEY, schema_1_manifest_payload(manifest)), (ACTIVE_MANIFEST_KEY, pre_region_datasets_manifest_payload(manifest)), (PRE_CACHE_MANIFEST_KEY, pre_cache_manifest_payload(manifest)), ] + 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 _assert_schema_1_top_level(payload: object) -> dict: + # The fields of the schema version 1 manifest and country models. Readers + # released before schema version 2 forbid any other field, so adding one + # here would stop them loading. + assert isinstance(payload, dict) + 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"} + return payload + + +@pytest.mark.parametrize("key", [REGION_DATASETS_MANIFEST_KEY, ACTIVE_MANIFEST_KEY]) +def test_readers_released_before_schema_2_get_only_fields_they_accept( + key: str, +) -> None: + store = FakeStore() + publish_v2_manifest( + store=store, worker=_regional_worker(with_household_functions=True) + ) + + with_cache = _assert_schema_1_top_level(store.values[key]) + 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", - } + pre_cache = _assert_schema_1_top_level(store.values[PRE_CACHE_MANIFEST_KEY]) worker = pre_cache["versions"]["5.2.0"] assert set(worker) == { "application_name", @@ -264,16 +345,18 @@ def test_readers_released_before_region_datasets_get_only_fields_they_accept( assert set(country) == _COUNTRY_BUNDLE_FIELDS_BEFORE_REGION_DATASETS -def test_the_region_dataset_key_keeps_region_datasets_and_cache() -> None: +@pytest.mark.parametrize("key", [SCHEMA_2_MANIFEST_KEY, REGION_DATASETS_MANIFEST_KEY]) +def test_the_region_dataset_keys_keep_region_datasets_and_cache(key: str) -> None: store = FakeStore() publish_v2_manifest(store=store, worker=_regional_worker()) - loaded = V2ManifestLoader(store).load().versions["5.2.0"] + loaded = parse_v2_manifest(store.values[key]).versions["5.2.0"] us = next(country for country in loaded.bundle.countries if country.country == "us") assert us.dataset_for_region_type("state") == "populace_us_2024_local" assert us.dataset_for_region_type("national") == "populace_us_2024" assert us.dataset_for_region_type("congressional_district") is None assert loaded.cache is not None + assert V2ManifestLoader(store).load().versions["5.2.0"] == loaded def test_a_bundle_published_before_region_datasets_has_none() -> None: @@ -304,35 +387,111 @@ def test_first_publish_with_region_datasets_keeps_versions_from_the_cache_key() assert V2ManifestLoader(store).load() == manifest -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"}} +def _schema_1_document(*, include_cache: bool, version: str = "5.1.0") -> dict: + """Return a manifest exactly as a publisher released before schema 2 wrote it.""" + worker = _worker(version).model_dump( + mode="json", + exclude={ + "countries": {"__all__": {"economy_capable", "household_callable"}}, + "validation": ( + {"household_validation_invocation_ids"} + | (set() if include_cache else {"cache"}) + ), + **({} if include_cache else {"cache": True}), + }, ) - store = FakeStore( - { - PRE_CACHE_MANIFEST_KEY: { - "schema_version": 1, - "generation": "5.1.0:earlier", - "default_version": "5.1.0", - "versions": {"5.1.0": earlier}, - } - } + return { + "schema_version": 1, + "generation": f"{version}:earlier", + "default_version": version, + "versions": {version: worker}, + } + + +@pytest.mark.parametrize( + ("key", "include_cache"), + [(ACTIVE_MANIFEST_KEY, True), (PRE_CACHE_MANIFEST_KEY, False)], +) +def test_loader_reads_a_schema_1_document_as_economy_capable_countries( + *, key: str, include_cache: bool +) -> None: + store = FakeStore({key: _schema_1_document(include_cache=include_cache)}) + + manifest = V2ManifestLoader(store).load() + + assert manifest.schema_version == 2 + worker = manifest.versions["5.1.0"] + assert [country.economy_capable for country in worker.countries] == [True, True] + assert worker.economy_worker("uk") == _country_worker("uk") + assert worker.household_worker("us") is None + assert worker.validation.household_validation_invocation_ids == {} + + +def test_loader_reads_the_schema_2_document_its_publish_wrote() -> None: + store = FakeStore({ACTIVE_MANIFEST_KEY: _schema_1_document(include_cache=True)}) + newer = publish_v2_manifest( + store=store, worker=_worker(with_household_functions=True) ) + + loaded = V2ManifestLoader(store).load() + + # Only the schema version 2 document names the household functions. + assert loaded == newer + assert loaded.versions["5.2.0"].household_worker("us") is not None + + +@pytest.mark.parametrize( + ("key", "include_cache"), + [(ACTIVE_MANIFEST_KEY, True), (PRE_CACHE_MANIFEST_KEY, False)], +) +def test_first_schema_2_publish_keeps_versions_from_a_schema_1_key( + *, key: str, include_cache: bool +) -> None: + store = FakeStore({key: _schema_1_document(include_cache=include_cache)}) 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_a_household_only_country_leaves_schema_1_readers_on_their_manifest() -> None: + earlier = _schema_1_document(include_cache=True) + store = FakeStore({ACTIVE_MANIFEST_KEY: earlier}) + + 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] == earlier + assert schema_1_manifest_payload(manifest, include_cache=True) is None + worker = manifest.versions["5.2.0"] + assert worker.economy_worker("uk") is None + assert worker.household_worker("uk") is not None + assert worker.economy_worker("us") is not None + + +def test_parse_v2_manifest_rejects_unknown_schema_versions() -> None: + document = _schema_1_document(include_cache=True) + document["schema_version"] = 3 + + with pytest.raises(ValueError, match="schema_version"): + parse_v2_manifest(document) + + 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 @@ -363,7 +522,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"}} @@ -398,6 +557,46 @@ def test_cache_metadata_is_optional_for_old_workers_and_must_match_validation() V2WorkerVersion.model_validate(value) +@pytest.mark.parametrize( + ("economy_capable", "single_simulation_callable", "household_callable"), + [ + (True, None, None), + (False, "run_single_simulation_us", None), + (False, "run_single_simulation_us", "calculate_household_us"), + (False, None, None), + ], +) +def test_country_capabilities_must_name_their_functions( + *, + economy_capable: bool, + single_simulation_callable: str | None, + household_callable: str | None, +) -> None: + with pytest.raises(ValueError, match=r"economy-capable|capability"): + V2CountryWorker( + country="us", + economy_capable=economy_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_only_economy_capable_countries_have_simulation_receipts() -> None: + value = _worker(household_only=("uk",)).model_dump(mode="json") + value["validation"]["country_validation_invocation_ids"]["uk"] = "not-needed" + + with pytest.raises(ValueError, match="economy-capable"): + V2WorkerVersion.model_validate(value) + + def test_manifest_rejects_credential_like_fields() -> None: with pytest.raises(ValueError, match="credential-like"): validate_manifest_has_no_credentials( @@ -479,3 +678,50 @@ def test_a_publish_after_a_rollback_keeps_the_versions_the_older_code_added() -> manifest = publish_v2_manifest(store=store, worker=_regional_worker("5.4.0")) assert sorted(manifest.versions) == ["5.3.0", "5.4.0"] + + +def test_a_schema_2_document_left_behind_by_an_older_publish_is_not_read() -> None: + """After a rollback, code released before schema 2 writes only older keys.""" + store = FakeStore() + publish_v2_manifest( + store=store, worker=_regional_worker(with_household_functions=True) + ) + stale = store.values[SCHEMA_2_MANIFEST_KEY] + older = publish_v2_manifest(store=FakeStore(), worker=_regional_worker("5.3.0")) + store.values[REGION_DATASETS_MANIFEST_KEY] = schema_1_manifest_payload(older) + store.values[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(older) + store.values[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(older) + + loaded = V2ManifestLoader(store).load() + manifest = publish_v2_manifest(store=store, worker=_regional_worker("5.4.0")) + + assert loaded.default_version == "5.3.0" + assert loaded.versions["5.3.0"].household_worker("us") is None + assert sorted(manifest.versions) == ["5.3.0", "5.4.0"] + assert store.values[SCHEMA_2_MANIFEST_KEY] != stale + + +def _validated_on(worker: V2WorkerVersion, day: int) -> V2WorkerVersion: + """Return ``worker`` as validated on ``day`` of September 2026.""" + validation = worker.validation.model_copy( + update={"validated_at": datetime(2026, 9, day, tzinfo=UTC)} + ) + return worker.model_copy(update={"validation": validation}) + + +def test_a_household_only_manifest_is_read_until_an_older_publish_follows() -> None: + # Schema version 1 cannot describe a household-only country, so its + # publish leaves the older keys alone; a later validation under them + # means code released before schema 2 published after it. + store = FakeStore({ACTIVE_MANIFEST_KEY: _schema_1_document(include_cache=True)}) + household_only = publish_v2_manifest( + store=store, worker=_validated_on(_worker(household_only=("uk",)), 15) + ) + assert V2ManifestLoader(store).load() == household_only + + older = publish_v2_manifest( + store=FakeStore(), worker=_validated_on(_cached_worker("5.3.0"), 20) + ) + store.values[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(older) + + assert V2ManifestLoader(store).load().default_version == "5.3.0" 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 c66929968..24abe5ff8 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 @@ -26,10 +26,7 @@ Stage12CountryBundle, Stage12Dataset, ) - from policyengine_simulation_contract.stage12_manifest import ( - V2CountryWorker, - V2WorkerVersion, - ) + from policyengine_simulation_contract.stage12_manifest import V2WorkerVersion class ComparisonSkipReason(StrEnum): @@ -92,13 +89,6 @@ 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 _options_skip_reason(payload: dict[str, Any]) -> ComparisonSkipReason | None: """Return why the request's fields or options are unsupported, if they are.""" if any( @@ -189,12 +179,16 @@ def adapt_annual_comparison( 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 economy-capable is converted; any + other country is skipped as unsupported. + """ skip_reason = _options_skip_reason(payload) if skip_reason is not None: return _skip(skip_reason) country = payload.get("country") - if not isinstance(country, str) or _country_worker(worker, country) is None: + if not isinstance(country, str) or worker.economy_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 fa4e83f6a..557e842a6 100644 --- a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py +++ b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py @@ -29,7 +29,12 @@ EVALUATION_ID = UUID("00000000-0000-0000-0000-000000000001") -def worker() -> V2WorkerVersion: +def worker(*, household_only: tuple[str, ...] = ()) -> V2WorkerVersion: + """Return a validated v2 worker for the US and the UK. + + A country in ``household_only`` has only a household function, so its + economy reports cannot run. + """ countries = [] country_workers = [] for country, package, package_version, dataset in ( @@ -60,10 +65,15 @@ def worker() -> V2WorkerVersion: ), ) ) + economy_capable = country not in household_only country_workers.append( V2CountryWorker( country=country, - single_simulation_callable=f"run_single_simulation_{country}", + economy_capable=economy_capable, + single_simulation_callable=( + f"run_single_simulation_{country}" if economy_capable else None + ), + household_callable=f"calculate_household_{country}", ) ) application_name = v2_application_name("5.2.0") @@ -86,8 +96,13 @@ 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", + item.country: f"validation-{item.country}" + for item in country_workers + if item.economy_capable + }, + household_validation_invocation_ids={ + item.country: f"household-validation-{item.country}" + for item in country_workers }, ), ) diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index 5751a9b24..21ff7a4d0 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -155,6 +155,17 @@ def test_adapter_returns_bounded_skip_reason(update, reason) -> None: assert result.skip_reason is reason +def test_adapter_skips_a_country_the_manifest_marks_household_only() -> None: + result = adapt_annual_comparison( + eligible_payload(), + evaluation_id=EVALUATION_ID, + worker=worker(household_only=("us",)), + ) + + assert result.report is None + assert result.skip_reason is ComparisonSkipReason.UNSUPPORTED_COUNTRY + + def _expected_simulation(role: str, execution_id: str, policy: dict) -> dict: dataset_uri = "hf://policyengine/data/populace_us_2024.h5@populace_us_2024-revision" return { diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_backend.py b/projects/policyengine-simulation-entry/tests/test_stage12_backend.py index ec75a064b..5bb9a6b2d 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_backend.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_backend.py @@ -29,8 +29,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): @@ -68,7 +68,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", @@ -79,7 +86,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, @@ -347,6 +354,48 @@ def test_temporary_submission_rejects_unsupported_input_without_side_effects() - assert invoker.calls == [] +def test_temporary_submission_rejects_a_country_without_economy_reports() -> None: + store = FakeStore() + invoker = FakeInvoker() + + with pytest.raises(TemporaryStage12UnsupportedRequestError) as error: + asyncio.run( + _backend(store, invoker, household_only=("us",)).submit_temporary_report( + request_payload=eligible_payload(), + request_id="request-1", + observability_id=OBSERVABILITY_ID, + ) + ) + + assert error.value.reason == "unsupported_country" + assert store.events == [] + assert invoker.calls == [] + + +def test_automatic_dispatch_skips_a_country_without_economy_reports() -> None: + invoker = FakeInvoker() + backend = _backend(FakeStore(), invoker, household_only=("us",)) + + asyncio.run( + backend.dispatch_after_production( + request_payload=eligible_payload(), + production_response=_response(), + request_id="request-1", + observability_id=OBSERVABILITY_ID, + ) + ) + + assert invoker.calls == [] + backend._runtime.event.assert_called_once_with( + "stage12_comparison_run_skipped", + attributes={ + "request_id": "request-1", + "production_identity": "job-1", + "reason": "unsupported_country", + }, + ) + + 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..00dbab4ce 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 has a single-simulation worker, so each is recorded + as economy-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}"), + economy_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 a808efcbe..cee40de8d 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,23 @@ def test_publisher_uses_only_packaged_bundle_values() -> None: ) +def test_publisher_records_every_bundle_country_as_economy_capable() -> None: + worker = publisher.build_worker(_validation().model_dump(mode="json")) + + assert [ + ( + country.country, + country.economy_capable, + country.single_simulation_callable, + country.household_callable, + ) + for country in worker.countries + ] == [ + (country.country, True, f"run_single_simulation_{country.country}", None) + for country in publisher.load_stage12_bundle().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 0485255fe04397de03dbea9250bd7cdc5b22f201 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:19:44 +0400 Subject: [PATCH 03/19] Port the Public API's household input preparation The v2 household functions prepare a household exactly as the Public API's household calculation does before calculating it: /calculate-full adds every missing yearly, monthly and eternal variable for the year of the first age entered for "you"; deprecated inputs are dropped, including axes that vary them, with the Public API's warning text; unknown entity groups, variables, variables on the wrong entity, axis variables and policy parameters are reported in the order the household was sent, with the Public API's error types and messages; and the requested values are the null values, or with axes every value of a known variable, in document order. Co-Authored-By: Claude Opus 5.5 --- .../household_inputs.py | 477 ++++++++++++++++++ .../tests/test_household_inputs.py | 322 ++++++++++++ 2 files changed, 799 insertions(+) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_inputs.py create mode 100644 projects/policyengine-simulation-executor/tests/test_household_inputs.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_inputs.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_inputs.py new file mode 100644 index 000000000..47a909cd6 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_inputs.py @@ -0,0 +1,477 @@ +"""Household input preparation, matching the Public API's in-process calculation. + +These functions reproduce, for the v2 household functions, how the Public API +prepares a household before calculating it: which period missing variables +are added for, which deprecated inputs are dropped (and the warning that says +so), which names are reported as unrecognized, and which values are +requested. Keeping the wording and order identical lets a v2 result be +compared with the Public API's result field by field. +""" + +from __future__ import annotations + +import copy +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from policyengine_simulation_contract.household_execution import HouseholdInputError + +if TYPE_CHECKING: + from collections.abc import Iterable, Mapping + from datetime import date + + from policyengine_core.taxbenefitsystems import TaxBenefitSystem + from policyengine_core.variables import Variable + from policyengine_simulation_contract.household_execution import InputErrorType + +# Removed or renamed model variables that callers may still send. The value is +# the migration hint shown in the warning. Matches the Public API's registry. +DEPRECATED_VARIABLES: dict[str, str] = { + "medical_out_of_pocket_expenses": ( + "Removed in policyengine-us 1.673.0. Migrate non-premium spending " + "to `other_medical_expenses` and premium spending to " + "`health_insurance_premiums`." + ), + "uc_has_limited_capability_for_work": ( + "Removed in policyengine-uk 2.129.0. Set " + "`uc_limited_capability_for_work` instead." + ), +} +# Variables with these definition periods are added when missing. +_ADDED_DEFINITION_PERIODS = frozenset({"year", "month", "eternity"}) +_AXES_KEY = "axes" + +# One value to calculate: entity plural, entity name, variable and period. +RequestedValue = tuple[str, str, str, str] + + +@dataclass(frozen=True) +class HouseholdVariable: + """One country-model variable, as household preparation needs it. + + ``default_value`` is the value the Public API's metadata reports: the + enum member name for an enum variable, the default for a boolean or + numeric variable, and None otherwise. + """ + + name: str + entity: str + definition_period: str + is_input: bool + default_value: bool | int | float | str | None + + +@dataclass(frozen=True) +class HouseholdEntity: + """One country-model entity: its key, plural and the plurals of its roles.""" + + key: str + plural: str + role_plurals: frozenset[str] + + +@dataclass(frozen=True) +class HouseholdModelMetadata: + """The variables, entities and policy parameters of one country model.""" + + variables: Mapping[str, HouseholdVariable] + entities: Mapping[str, HouseholdEntity] + parameters: frozenset[str] + + @classmethod + def from_tax_benefit_system( + cls, system: TaxBenefitSystem + ) -> HouseholdModelMetadata: + """Read the metadata the Public API builds from a tax-benefit system. + + Variables keep the system's order. Parameters are every descendant + whose name starts with ``gov``, as the Public API's metadata lists. + """ + entities = { + entity.key: HouseholdEntity( + key=entity.key, + plural=entity.plural, + role_plurals=frozenset( + role.plural + for role in getattr(entity, "roles", ()) + if role.plural is not None + ), + ) + for entity in system.entities + } + variables = { + name: HouseholdVariable( + name=name, + entity=variable.entity.key, + definition_period=variable.definition_period, + is_input=variable.is_input_variable(), + default_value=_metadata_default_value(variable), + ) + for name, variable in system.variables.items() + } + root = system.parameters + parameters = frozenset( + parameter.name + for parameter in (root.get_descendants() if root is not None else ()) + if parameter.name[:3] == "gov" + ) + return cls(variables=variables, entities=entities, parameters=parameters) + + def entity_by_plural(self) -> dict[str, HouseholdEntity]: + """Return the entities keyed by their plural.""" + return {entity.plural: entity for entity in self.entities.values()} + + +def _metadata_default_value(variable: Variable) -> bool | int | float | str | None: + default_value: object = variable.default_value + if variable.value_type.__name__ == "Enum": + return str(getattr(default_value, "name", default_value)) + if isinstance(default_value, (int, float, bool)): + return default_value + return None + + +def household_year(situation: Mapping[str, Any], *, today: date) -> str: + """Return the period missing yearly variables are added for. + + As in the Public API, this is the first period of the age entered for the + person named ``you``, or else the current calendar year. An age entered + without periods counts as no age. + """ + people = situation.get("people") + you = people.get("you") if isinstance(people, dict) else None + ages = you.get("age") if isinstance(you, dict) else None + if isinstance(ages, dict) and ages: + return str(next(iter(ages))) + return str(today.year) + + +def add_missing_variables( + situation: dict[str, Any], + metadata: HouseholdModelMetadata, + *, + today: date, +) -> None: + """Add every yearly, monthly and eternal variable each entity lacks, in place. + + Input variables get their default value and calculated variables get + None, which requests them, for the household's year. This is what the + Public API's ``/calculate-full`` route adds. + """ + year = household_year(situation, today=today) + for variable in metadata.variables.values(): + if variable.definition_period not in _ADDED_DEFINITION_PERIODS: + continue + entity_plural = metadata.entities[variable.entity].plural + for entity in situation.get(entity_plural, {}).values(): + if variable.name not in entity: + entity[variable.name] = { + year: variable.default_value if variable.is_input else None + } + + +@dataclass(frozen=True) +class DeprecatedVariableWarning: + """A removed or renamed variable was supplied and dropped before calculation.""" + + variable: str + entity_plural: str + entity_id: str + hint: str + + @property + def message(self) -> str: + """Return the warning naming the ignored input, where it was and the hint.""" + location = f"`{self.entity_plural}/{self.entity_id}`" + if self.entity_plural == _AXES_KEY: + location = f"`axes[{self.entity_id}].name`" + return ( + f"Input `{self.variable}` on {location} is deprecated and was " + f"ignored for this calculation. {self.hint}" + ) + + +@dataclass(frozen=True) +class DeprecatedInputsResult: + """A household copy without deprecated inputs, and one warning per input.""" + + household: dict[str, Any] + warnings: tuple[DeprecatedVariableWarning, ...] + + +def drop_deprecated_inputs(situation: Mapping[str, Any]) -> DeprecatedInputsResult: + """Return a copy of ``situation`` without deprecated inputs, with warnings. + + Each deprecated variable on an entity, and each axis that varies one, is + removed with one warning; an axis group left empty is removed, and + ``axes`` is removed when no group is left. ``situation`` is not changed. + """ + warnings: list[DeprecatedVariableWarning] = [] + household = copy.deepcopy(dict(situation)) + for entity_plural, entity_group in household.items(): + if entity_plural == _AXES_KEY or not isinstance(entity_group, dict): + continue + for entity_id, variables in entity_group.items(): + if isinstance(variables, dict): + warnings.extend( + _drop_deprecated_variables(entity_plural, entity_id, variables) + ) + _drop_deprecated_axes(household, warnings) + return DeprecatedInputsResult(household=household, warnings=tuple(warnings)) + + +def _drop_deprecated_variables( + entity_plural: str, entity_id: str, variables: dict[str, Any] +) -> list[DeprecatedVariableWarning]: + warnings: list[DeprecatedVariableWarning] = [] + for variable_name in list(variables): + hint = DEPRECATED_VARIABLES.get(variable_name) + if hint is None: + continue + warnings.append( + DeprecatedVariableWarning( + variable=variable_name, + entity_plural=entity_plural, + entity_id=entity_id, + hint=hint, + ) + ) + del variables[variable_name] + return warnings + + +def _drop_deprecated_axes( + household: dict[str, Any], warnings: list[DeprecatedVariableWarning] +) -> None: + axes = household.get(_AXES_KEY) + if not isinstance(axes, list): + return + changed = False + retained_groups = [] + for group_index, group in enumerate(axes): + retained_axes = [] + for axis_index, axis in enumerate(group): + warning = _deprecated_axis_warning(axis, f"{group_index}][{axis_index}") + if warning is None: + retained_axes.append(axis) + continue + warnings.append(warning) + changed = True + if retained_axes: + retained_groups.append(retained_axes) + if not changed: + return + if retained_groups: + household[_AXES_KEY] = retained_groups + else: + del household[_AXES_KEY] + + +def _deprecated_axis_warning( + axis: object, location: str +) -> DeprecatedVariableWarning | None: + if not isinstance(axis, dict): + return None + variable_name = str(axis.get("name")) + hint = DEPRECATED_VARIABLES.get(variable_name) + if hint is None: + return None + return DeprecatedVariableWarning( + variable=variable_name, entity_plural=_AXES_KEY, entity_id=location, hint=hint + ) + + +def _input_error_message( + input_type: InputErrorType, + name: str, + path: str, + entity_plurals: tuple[str, str] | None, +) -> str: + if input_type == "household_entity": + return f"Unrecognized household entity group `{name}` at `{path}`." + if input_type == "household_variable_wrong_entity" and entity_plurals: + expected, actual = entity_plurals + return ( + f"Household variable `{name}` belongs on `{expected}`, " + f"not `{actual}`, at `{path}`." + ) + if input_type == "household_axis_variable": + return f"Unrecognized household axis variable `{name}` at `{path}`." + if input_type == "policy_parameter": + return f"Unrecognized policy parameter `{name}` at `{path}`." + return f"Unrecognized household variable `{name}` at `{path}`." + + +def _input_error( + input_type: InputErrorType, + name: str, + path: str, + *, + entity_plurals: tuple[str, str] | None = None, +) -> HouseholdInputError: + """Return one unrecognized input, worded as the Public API words it. + + ``entity_plurals`` is the (expected, actual) pair for a variable placed on + the wrong entity. + """ + expected, actual = entity_plurals or (None, None) + return HouseholdInputError( + type=input_type, + name=name, + path=path, + message=_input_error_message(input_type, name, path, entity_plurals), + expected_entity_plural=expected, + actual_entity_plural=actual, + ) + + +def find_unrecognized_household_inputs( + situation: Mapping[str, Any], + metadata: HouseholdModelMetadata, +) -> list[HouseholdInputError]: + """Find unknown entity groups, variables and axis variables, in document order. + + Also reports a known variable placed on the wrong entity. Role membership + keys (such as ``members``) are not variables and are skipped. + """ + errors: list[HouseholdInputError] = [] + entity_by_plural = metadata.entity_by_plural() + for entity_plural, entity_group in situation.items(): + if entity_plural == _AXES_KEY: + errors.extend(_find_unrecognized_axes(entity_group, metadata)) + continue + entity = entity_by_plural.get(entity_plural) + if entity is None: + errors.append( + _input_error( + "household_entity", entity_plural, f"household.{entity_plural}" + ) + ) + continue + if not isinstance(entity_group, dict): + continue + for entity_id, entity_inputs in entity_group.items(): + if isinstance(entity_inputs, dict): + errors.extend( + _find_unrecognized_entity_inputs( + entity, entity_id, entity_inputs, metadata + ) + ) + return errors + + +def _find_unrecognized_entity_inputs( + entity: HouseholdEntity, + entity_id: str, + entity_inputs: Mapping[str, Any], + metadata: HouseholdModelMetadata, +) -> list[HouseholdInputError]: + errors: list[HouseholdInputError] = [] + for variable_name in entity_inputs: + if variable_name in entity.role_plurals: + continue + variable = metadata.variables.get(variable_name) + path = f"household.{entity.plural}.{entity_id}.{variable_name}" + if variable is None: + errors.append(_input_error("household_variable", variable_name, path)) + continue + expected_entity = metadata.entities.get(variable.entity) + if expected_entity is not None and expected_entity.plural != entity.plural: + errors.append( + _input_error( + "household_variable_wrong_entity", + variable_name, + path, + entity_plurals=(expected_entity.plural, entity.plural), + ) + ) + return errors + + +def _find_unrecognized_axes( + axes: object, metadata: HouseholdModelMetadata +) -> list[HouseholdInputError]: + if not isinstance(axes, list): + return [] + errors: list[HouseholdInputError] = [] + for group_index, group in enumerate(axes): + for axis_index, axis in enumerate(group): + if not isinstance(axis, dict) or "name" not in axis: + continue + variable_name = str(axis["name"]) + if variable_name in metadata.variables: + continue + errors.append( + _input_error( + "household_axis_variable", + variable_name, + f"household.axes[{group_index}][{axis_index}].name", + ) + ) + return errors + + +def find_unrecognized_policy_inputs( + parameters: Iterable[str], + metadata: HouseholdModelMetadata, +) -> list[HouseholdInputError]: + """Find policy parameter paths the country model does not have, in order.""" + return [ + _input_error("policy_parameter", name, f"policy.{name}") + for name in parameters + if name not in metadata.parameters + ] + + +def unrecognized_inputs_message(errors: Iterable[HouseholdInputError]) -> str: + """Join the errors' messages as the Public API's error response does.""" + return "Unrecognized calculate input(s): " + "; ".join( + error.message for error in errors + ) + + +def requested_computations( + situation: Mapping[str, Any], + *, + include_provided_values: bool, + variable_names: set[str] | frozenset[str] | None = None, +) -> list[RequestedValue]: + """List the (entity plural, entity name, variable, period) values to calculate. + + A value is requested when it is None or, with ``include_provided_values``, + whatever it is. ``variable_names`` keeps only those variables. Values are + listed in document order, as the Public API lists them. + """ + computations: list[RequestedValue] = [] + for entity_plural, entities in situation.items(): + if not isinstance(entities, dict): + continue + for entity_id, variables in entities.items(): + if isinstance(variables, dict): + computations.extend( + _requested_entity_values( + (entity_plural, entity_id), + variables, + include_provided_values=include_provided_values, + variable_names=variable_names, + ) + ) + return computations + + +def _requested_entity_values( + entity: tuple[str, str], + variables: Mapping[str, Any], + *, + include_provided_values: bool, + variable_names: set[str] | frozenset[str] | None, +) -> list[RequestedValue]: + entity_plural, entity_id = entity + return [ + (entity_plural, entity_id, variable_name, str(period)) + for variable_name, periods in variables.items() + if isinstance(periods, dict) + and (variable_names is None or variable_name in variable_names) + for period, value in periods.items() + if include_provided_values or value is None + ] diff --git a/projects/policyengine-simulation-executor/tests/test_household_inputs.py b/projects/policyengine-simulation-executor/tests/test_household_inputs.py new file mode 100644 index 000000000..9f29025cf --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_household_inputs.py @@ -0,0 +1,322 @@ +"""Tests for household input preparation, ported from the Public API.""" + +from __future__ import annotations + +from datetime import date +from types import SimpleNamespace + +from policyengine_simulation_executor.household_inputs import ( + HouseholdEntity, + HouseholdModelMetadata, + HouseholdVariable, + add_missing_variables, + drop_deprecated_inputs, + find_unrecognized_household_inputs, + find_unrecognized_policy_inputs, + household_year, + requested_computations, + unrecognized_inputs_message, +) + +TODAY = date(2027, 3, 1) + + +def _metadata() -> HouseholdModelMetadata: + variables = [ + HouseholdVariable("age", "person", "year", is_input=True, default_value=40), + HouseholdVariable( + "employment_income", "person", "year", is_input=True, default_value=0 + ), + HouseholdVariable( + "is_blind", "person", "eternity", is_input=True, default_value=False + ), + HouseholdVariable( + "monthly_rent", "household", "month", is_input=True, default_value=0 + ), + HouseholdVariable( + "state_name", "household", "year", is_input=True, default_value="CA" + ), + HouseholdVariable( + "household_net_income", + "household", + "year", + is_input=False, + default_value=0, + ), + HouseholdVariable( + "daily_amount", "household", "day", is_input=True, default_value=0 + ), + ] + return HouseholdModelMetadata( + variables={variable.name: variable for variable in variables}, + entities={ + "person": HouseholdEntity("person", "people", frozenset()), + "household": HouseholdEntity( + "household", "households", frozenset({"members"}) + ), + }, + parameters=frozenset({"gov.tax.rate", "gov.tax.brackets[0].rate"}), + ) + + +def _situation() -> dict: + return { + "people": {"you": {"age": {"2026": 30}}, "partner": {}}, + "households": {"home": {"members": ["you", "partner"]}}, + } + + +def test_household_year_is_the_first_age_period_of_you() -> None: + assert household_year(_situation(), today=TODAY) == "2026" + assert household_year({"people": {"partner": {}}}, today=TODAY) == "2027" + assert household_year({"people": {"you": {"age": 30}}}, today=TODAY) == "2027" + + +def test_add_missing_variables_matches_the_public_api() -> None: + situation = _situation() + + add_missing_variables(situation, _metadata(), today=TODAY) + + assert situation == { + "people": { + "you": { + "age": {"2026": 30}, + "employment_income": {"2026": 0}, + "is_blind": {"2026": False}, + }, + "partner": { + "age": {"2026": 40}, + "employment_income": {"2026": 0}, + "is_blind": {"2026": False}, + }, + }, + "households": { + "home": { + "members": ["you", "partner"], + "monthly_rent": {"2026": 0}, + "state_name": {"2026": "CA"}, + # Calculated variables are requested, never defaulted. + "household_net_income": {"2026": None}, + } + }, + } + + +def test_deprecated_inputs_are_dropped_with_public_api_warnings() -> None: + situation = { + "people": {"you": {"age": {"2026": 30}, "medical_out_of_pocket_expenses": 5}}, + "axes": [ + [{"name": "uc_has_limited_capability_for_work", "count": 2}], + [{"name": "employment_income", "count": 2}], + ], + } + + result = drop_deprecated_inputs(situation) + + assert result.household == { + "people": {"you": {"age": {"2026": 30}}}, + "axes": [[{"name": "employment_income", "count": 2}]], + } + assert "medical_out_of_pocket_expenses" in situation["people"]["you"] + assert [warning.message for warning in result.warnings] == [ + "Input `medical_out_of_pocket_expenses` on `people/you` is deprecated and " + "was ignored for this calculation. Removed in policyengine-us 1.673.0. " + "Migrate non-premium spending to `other_medical_expenses` and premium " + "spending to `health_insurance_premiums`.", + "Input `uc_has_limited_capability_for_work` on `axes[0][0].name` is " + "deprecated and was ignored for this calculation. Removed in " + "policyengine-uk 2.129.0. Set `uc_limited_capability_for_work` instead.", + ] + + +def test_axes_left_empty_by_deprecated_inputs_are_removed() -> None: + situation = { + "people": {"you": {}}, + "axes": [[{"name": "medical_out_of_pocket_expenses", "count": 2}]], + } + + assert "axes" not in drop_deprecated_inputs(situation).household + + +def test_unrecognized_household_inputs_use_public_api_types_and_messages() -> None: + situation = { + "people": { + "you": {"age": {"2026": 30}, "not_a_variable": 1, "state_name": "CA"} + }, + "households": {"home": {"members": ["you"]}}, + "pets": {"rex": {}}, + "axes": [[{"name": "unknown_axis", "count": 2}]], + } + + errors = find_unrecognized_household_inputs(situation, _metadata()) + + assert [error.model_dump(mode="json") for error in errors] == [ + { + "type": "household_variable", + "name": "not_a_variable", + "path": "household.people.you.not_a_variable", + "message": ( + "Unrecognized household variable `not_a_variable` at " + "`household.people.you.not_a_variable`." + ), + }, + { + "type": "household_variable_wrong_entity", + "name": "state_name", + "path": "household.people.you.state_name", + "message": ( + "Household variable `state_name` belongs on `households`, not " + "`people`, at `household.people.you.state_name`." + ), + "expected_entity_plural": "households", + "actual_entity_plural": "people", + }, + { + "type": "household_entity", + "name": "pets", + "path": "household.pets", + "message": ( + "Unrecognized household entity group `pets` at `household.pets`." + ), + }, + { + "type": "household_axis_variable", + "name": "unknown_axis", + "path": "household.axes[0][0].name", + "message": ( + "Unrecognized household axis variable `unknown_axis` at " + "`household.axes[0][0].name`." + ), + }, + ] + assert unrecognized_inputs_message(errors[:1]) == ( + "Unrecognized calculate input(s): Unrecognized household variable " + "`not_a_variable` at `household.people.you.not_a_variable`." + ) + + +def test_unrecognized_policy_parameters_are_reported_by_path() -> None: + errors = find_unrecognized_policy_inputs( + {"gov.tax.rate": {}, "gov.tax.brackets[0].rate": {}, "gov.missing": {}}, + _metadata(), + ) + + assert [error.model_dump(mode="json") for error in errors] == [ + { + "type": "policy_parameter", + "name": "gov.missing", + "path": "policy.gov.missing", + "message": "Unrecognized policy parameter `gov.missing` at " + "`policy.gov.missing`.", + } + ] + + +def test_deprecated_inputs_leave_the_callers_household_unchanged() -> None: + situation = {"people": {"you": {"medical_out_of_pocket_expenses": 5}}} + + result = drop_deprecated_inputs(situation) + + assert result.household == {"people": {"you": {}}} + assert situation == {"people": {"you": {"medical_out_of_pocket_expenses": 5}}} + + +def test_a_household_without_deprecated_inputs_keeps_its_axes() -> None: + situation = { + "people": {"you": {}}, + "axes": [[{"name": "employment_income", "count": 2}]], + } + + result = drop_deprecated_inputs(situation) + + assert result.household == situation + assert result.warnings == () + + +def test_requested_computations_are_the_null_values_without_axes() -> None: + situation = { + "people": {"you": {"age": {"2026": 30}, "employment_income": {"2026": None}}}, + "households": { + "home": {"members": ["you"], "household_net_income": {"2026": None}} + }, + } + + assert requested_computations(situation, include_provided_values=False) == [ + ("people", "you", "employment_income", "2026"), + ("households", "home", "household_net_income", "2026"), + ] + + +def test_requested_computations_with_axes_cover_every_known_variable() -> None: + situation = { + "people": {"you": {"age": {"2026": 30}, "unknown": {"2026": 1}}}, + "households": {"home": {"members": ["you"]}}, + "axes": [[{"name": "age", "count": 2, "min": 0, "max": 1}]], + } + + assert requested_computations( + situation, + include_provided_values=True, + variable_names={"age", "employment_income"}, + ) == [("people", "you", "age", "2026")] + + +def test_metadata_is_read_from_a_tax_benefit_system() -> None: + class Enum: + pass + + class Level: + name = "LOW" + + person = SimpleNamespace(key="person", plural="people", is_person=True) + household = SimpleNamespace( + key="household", + plural="households", + is_person=False, + roles=[SimpleNamespace(plural="members"), SimpleNamespace(plural=None)], + ) + + def variable(name, entity, value_type, default_value, *, is_input=True): + return SimpleNamespace( + entity=entity, + definition_period="year", + value_type=value_type, + default_value=default_value, + is_input_variable=lambda: is_input, + ) + + parameters = SimpleNamespace( + get_descendants=lambda: [ + SimpleNamespace(name="gov.tax.rate"), + SimpleNamespace(name="calibration.factor"), + ] + ) + system = SimpleNamespace( + entities=[person, household], + variables={ + "age": variable("age", person, int, 40), + "level": variable("level", household, Enum, Level()), + "label": variable("label", household, str, "x"), + "start": variable("start", person, date, date(2020, 1, 1)), + "net_income": variable("net_income", household, float, 0.0, is_input=False), + }, + parameters=parameters, + ) + + metadata = HouseholdModelMetadata.from_tax_benefit_system(system) + + assert metadata.entities["household"] == HouseholdEntity( + "household", "households", frozenset({"members"}) + ) + assert metadata.entities["person"].role_plurals == frozenset() + assert [ + (item.name, item.entity, item.default_value, item.is_input) + for item in metadata.variables.values() + ] == [ + ("age", "person", 40, True), + ("level", "household", "LOW", True), + ("label", "household", None, True), + ("start", "person", None, True), + ("net_income", "household", 0.0, False), + ] + assert metadata.parameters == frozenset({"gov.tax.rate"}) From 7aeb91407447b430b59431218370dbde3cedcb3c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 02:22:51 +0400 Subject: [PATCH 04/19] Resolve household SPM selections and receipts like the Public API A US household's SPM selection is merged over the installed bundle's certified defaults with the Public API's rules and error codes: settings for any other country, or for a US model bundle without canonical SPM, are unsupported; a model that can run canonical SPM without a certified configuration fails closed; a forecast hash the bundle lacks is an invalid setting from a request but an unavailable configuration from a saved household; and an unknown scenario, county vintage or area is reported with the Public API's code and message. After each calculation, the country model's SPM receipt is reduced to the compact public receipt and checked against the resolved selection. SPM errors that end a calculation are recognized through wrapped causes, as the Public API recognizes them. The executor's SPM module now exposes the certified forecast it already loads once per process. Co-Authored-By: Claude Opus 5.5 --- .../household_spm.py | 418 ++++++++++++++++++ .../policyengine_simulation_executor/spm.py | 9 + .../tests/test_household_spm.py | 294 ++++++++++++ 3 files changed, 721 insertions(+) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py create mode 100644 projects/policyengine-simulation-executor/tests/test_household_spm.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py new file mode 100644 index 000000000..2c5b61d29 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py @@ -0,0 +1,418 @@ +"""SPM selection and receipts for v2 household calculations. + +This reproduces the Public API's household SPM handling, so that a v2 +household result and error carry the same SPM codes, messages and receipt: + +- ``resolve_household_spm_selection`` merges a caller's selection over the + installed bundle's certified defaults, with the Public API's error codes; +- ``household_spm_error_detail`` recognizes the SPM errors that end a + calculation; +- ``household_spm_provenance`` reduces the country model's receipt to the + compact public receipt and checks it matches the resolved selection. +""" + +from __future__ import annotations + +import inspect +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, NoReturn + +from policyengine_simulation_contract.spm import ( + SPMErrorDetail, + SPMInputError, + SPMProvenance, + SPMRuntimeVersions, + SPMSelection, + build_spm_provenance, +) +from pydantic import ValidationError + +if TYPE_CHECKING: + from collections.abc import Callable + + from spm_calculator.rolling_forecast import SPMForecast + +# Country SPM errors about the household's own inputs. When the caller did not +# choose SPM settings, a value that needs one of these is left uncalculated +# instead of failing the calculation. +HOUSEHOLD_SPM_INPUT_ERROR_CODES = frozenset( + { + "SPM_GEOGRAPHY_REQUIRED", + "SPM_GEOGRAPHY_UNAVAILABLE", + "SPM_COMPOSITION_REQUIRED", + "SPM_YEAR_UNAVAILABLE", + } +) +_SUPPORTED_COUNTY_VINTAGE = "2020" +_CONFIGURATION_UNAVAILABLE = "SPM_CONFIGURATION_UNAVAILABLE" +_SETTINGS_INVALID = "SPM_SETTINGS_INVALID" +_SETTINGS_UNSUPPORTED = "SPM_SETTINGS_UNSUPPORTED" + + +def _readable_error(error: BaseException) -> str: + """Render a validation error as a client-facing message, without its URL.""" + if not isinstance(error, ValidationError): + return str(error) + return "; ".join( + ": ".join( + part + for part in ( + ".".join(str(item) for item in detail["loc"]), + detail["msg"].removeprefix("Value error, "), + ) + if part + ) + for detail in error.errors() + ) + + +def _current_bundle() -> object: + """Return the installed PolicyEngine.py bundle manifest.""" + from policyengine.bundle import get_current_bundle + + return get_current_bundle() + + +def _bundle_spm_configuration() -> object: + """Return the installed bundle's SPM measurement configuration, if any. + + Raises SPMInputError when the bundle manifest cannot be read. + """ + try: + bundle = _current_bundle() + except (AttributeError, ImportError, OSError, TypeError, ValueError): + bundle = None + if not isinstance(bundle, Mapping): + raise SPMInputError( + _CONFIGURATION_UNAVAILABLE, "The installed bundle manifest is invalid" + ) + measurements = bundle.get("measurements") + return measurements.get("spm") if isinstance(measurements, Mapping) else None + + +def _simulation_supports_spm(simulation_type: Callable[..., object]) -> bool: + """Whether a country ``Simulation`` constructor implements canonical SPM.""" + try: + parameters = inspect.signature(simulation_type).parameters + except (TypeError, ValueError): + return False + return "spm" in parameters or ( + hasattr(simulation_type, "spm_config") + and callable(getattr(simulation_type, "spm_provenance", None)) + ) + + +def _us_model_supports_spm() -> bool: + """Whether the installed US model implements canonical SPM. + + A model this build cannot import has no canonical SPM. + """ + try: + from policyengine_us import Simulation + except (AttributeError, ImportError, OSError, TypeError, ValueError): + return False + return _simulation_supports_spm(Simulation) + + +def _forecast(expected_sha256: str) -> SPMForecast: + from policyengine_simulation_executor.spm import certified_spm_forecast + + return certified_spm_forecast(expected_sha256) + + +def _certified_defaults(configured: object) -> tuple[SPMSelection, SPMForecast]: + """Return the bundle's pinned SPM defaults and the forecast they name. + + Raises ValueError (or an import or loading error) when the configuration + pins no forecast hash or scenario, or the US model lacks canonical SPM. + """ + defaults = SPMSelection.model_validate(configured) + if defaults.forecast_content_sha256 is None or defaults.scenario is None: + msg = "The bundle must pin the SPM artifact hash and scenario" + raise ValueError(msg) + if not _us_model_supports_spm(): + msg = "The installed US model does not support canonical SPM" + raise ValueError(msg) + return defaults, _forecast(defaults.forecast_content_sha256) + + +def _validated_selection(values: dict[str, Any], forecast: SPMForecast) -> SPMSelection: + """Validate merged settings against the forecast; raise ValueError if invalid.""" + resolved = SPMSelection.model_validate(values) + forecast.entry(forecast.years[0], scenario=resolved.scenario, as_of=resolved.as_of) + if resolved.county_vintage != _SUPPORTED_COUNTY_VINTAGE: + msg = "Unsupported county vintage: use 2020" + raise ValueError(msg) + return resolved + + +def _validate_metro_selection(forecast: SPMForecast, selection: SPMSelection) -> None: + """Raise SPM_GEOGRAPHY_UNAVAILABLE unless the area is in some forecast year.""" + for year in forecast.years: + try: + forecast.geography_factor( + year, + "renter", + scenario=selection.scenario, + kind="metro", + geoid=selection.geography_id, + as_of=selection.as_of, + ) + except ValueError: + continue + return + error_code = "SPM_GEOGRAPHY_UNAVAILABLE" + raise SPMInputError( + error_code, + f"SPM area is unavailable in the selected artifact: {selection.geography_id}", + ) + + +def _merged_selection(chosen: SPMSelection, defaults: SPMSelection) -> dict[str, Any]: + """Overlay the caller's choices on the defaults, as the Public API does. + + Changing the geography kind also replaces an inherited geography ID, and + the forecast hash always comes from the bundle. + """ + values = {name: getattr(defaults, name) for name in SPMSelection.model_fields} + values.update(chosen.model_dump(exclude_unset=True)) + if ( + "geography_kind" in chosen.model_fields_set + and chosen.geography_kind != defaults.geography_kind + ): + values["geography_id"] = chosen.geography_id + values["forecast_content_sha256"] = defaults.forecast_content_sha256 + values["scenario"] = chosen.scenario or defaults.scenario + return values + + +def _certified_configuration(*, selection_sent: bool) -> Mapping[str, object] | None: + """Return the bundle's certified SPM settings, or None for a legacy bundle. + + Raises SPM_CONFIGURATION_UNAVAILABLE for an unreadable bundle or for a US + model that can run canonical SPM without a certified configuration, and + SPM_SETTINGS_UNSUPPORTED when settings were sent to a legacy bundle. + """ + configured = _bundle_spm_configuration() + if isinstance(configured, Mapping): + return configured + if _us_model_supports_spm(): + # A model that can run canonical SPM must never run it uncertified. + raise SPMInputError( + _CONFIGURATION_UNAVAILABLE, + "The installed bundle has no certified SPM measurement configuration", + ) + if not selection_sent: + return None + raise SPMInputError( + _SETTINGS_UNSUPPORTED, + "This US model bundle does not support canonical SPM settings", + ) + + +def _raise_unavailable_forecast(*, saved_household: bool) -> NoReturn: + """Reject a selection naming a forecast this bundle does not have.""" + if saved_household: + # A saved selection is the household's identity, not a request: this + # deployment lacks the forecast it was measured against. + raise SPMInputError( + _CONFIGURATION_UNAVAILABLE, + "This deployment does not have the SPM artifact this household's " + "saved selection was measured against", + ) + raise SPMInputError( + _SETTINGS_INVALID, + "SPM selection does not match this bundle's artifact hash", + ) + + +def _resolve_certified_selection( + chosen: SPMSelection, + configured: Mapping[str, object], + *, + saved_household: bool, +) -> dict[str, Any]: + """Merge the caller's settings over the certified defaults and validate them.""" + try: + defaults, forecast = _certified_defaults(configured) + except (AttributeError, ImportError, ValueError, TypeError, OSError) as error: + raise SPMInputError( + _CONFIGURATION_UNAVAILABLE, _readable_error(error) + ) from error + if chosen.forecast_content_sha256 not in (None, defaults.forecast_content_sha256): + _raise_unavailable_forecast(saved_household=saved_household) + try: + resolved = _validated_selection(_merged_selection(chosen, defaults), forecast) + except ValueError as error: + raise SPMInputError(_SETTINGS_INVALID, _readable_error(error)) from error + if resolved.geography_kind == "metro": + _validate_metro_selection(forecast, resolved) + return resolved.model_dump(mode="json") + + +def resolve_household_spm_selection( + country: str, + selection: SPMSelection | None, + *, + saved_household: bool, +) -> dict[str, Any] | None: + """Resolve a household's SPM selection against the installed bundle. + + ``selection`` is None when the caller made no choice. Returns the complete + resolved selection, or None when no SPM applies (any country but the US, + or a bundle whose US model predates canonical SPM). Raises SPMInputError + with the Public API's code and message otherwise. ``saved_household`` says + the selection was read from a saved household, which changes only how a + forecast hash the bundle no longer carries is reported. + """ + if country != "us": + if selection is not None: + raise SPMInputError( + _SETTINGS_UNSUPPORTED, "SPM settings are only available for the US" + ) + return None + configured = _certified_configuration(selection_sent=selection is not None) + if configured is None: + return None + return _resolve_certified_selection( + selection if selection is not None else SPMSelection(), + configured, + saved_household=saved_household, + ) + + +def household_spm_error_detail(error: BaseException) -> SPMErrorDetail | None: + """Return the SPM error that ends a calculation, following wrapped causes. + + Recognizes this module's own ``SPMInputError`` and a country model's error + carrying one of ``HOUSEHOLD_SPM_INPUT_ERROR_CODES``, as the Public API does. + """ + visited: set[int] = set() + current: BaseException | None = error + while current is not None and id(current) not in visited: + visited.add(id(current)) + code = getattr(current, "code", None) + if isinstance(current, SPMInputError) or ( + isinstance(current, ValueError) and code in HOUSEHOLD_SPM_INPUT_ERROR_CODES + ): + return SPMErrorDetail(code=str(code), message=str(current)) + current = current.__cause__ or current.__context__ + return None + + +def _required_text(source: Mapping[str, object], field_name: str) -> str: + value = source.get(field_name) + if not isinstance(value, str) or not value: + msg = f"Country SPM receipt requires {field_name}" + raise ValueError(msg) + return value + + +def _complete_settings(settings: object) -> SPMSelection | None: + """Return the selection only when every field is present.""" + try: + selection = SPMSelection.model_validate(settings) + except ValidationError: + return None + if set(SPMSelection.model_fields) - selection.model_fields_set: + return None + return selection + + +def _receipt_mapping(receipt: Mapping[str, object], field_name: str) -> Mapping: + value = receipt.get(field_name) + if not isinstance(value, Mapping): + msg = f"Country SPM receipt requires {field_name}" + raise TypeError(msg) + return value + + +def _compact_receipt( + selection: SPMSelection, receipt: Mapping[str, object] +) -> SPMProvenance: + """Reduce the country model's rich receipt to the compact public receipt.""" + forecast_sha256 = _required_text(receipt, "forecast_sha256") + if ( + forecast_sha256 != selection.forecast_content_sha256 + or _required_text(receipt, "scenario") != selection.scenario + or _required_text(receipt, "geography_kind") != selection.geography_kind + ): + msg = "Country SPM receipt differs from the SPM configuration" + raise ValueError(msg) + source_years = _receipt_mapping(receipt, "years") + if any(not isinstance(year, str) for year in source_years): + msg = "Country SPM receipt years must be strings" + raise ValueError(msg) + return build_spm_provenance( + forecast_id=_required_text(receipt, "forecast_id"), + forecast_sha256=forecast_sha256, + selection=selection, + years=tuple(sorted(source_years)), + runtime_versions=SPMRuntimeVersions.model_validate( + _receipt_mapping(receipt, "runtime_versions") + ), + ) + + +def _country_receipt(simulation: object) -> SPMProvenance: + """Reduce the simulation's receipt. + + Raises ValueError or TypeError if the receipt cannot be certified. + """ + selection = _complete_settings(getattr(simulation, "spm_config", None)) + if selection is None: + msg = "Country SPM configuration is incomplete" + raise ValueError(msg) + receipt_method = getattr(simulation, "spm_provenance", None) + source = receipt_method() if callable(receipt_method) else None + if not isinstance(source, Mapping): + msg = "Country SPM receipt must be an object" + raise TypeError(msg) + return _compact_receipt(selection, source) + + +def _calculation_receipt(simulation: object) -> SPMProvenance | None: + """Return the simulation's compact SPM receipt; None without SPM support.""" + if not hasattr(simulation, "spm_config"): + return None + try: + return _country_receipt(simulation) + except (TypeError, ValueError) as error: + raise SPMInputError( + _CONFIGURATION_UNAVAILABLE, + "The installed country model reported an SPM receipt this API cannot " + f"certify: {_readable_error(error)}", + ) from error + + +def household_spm_provenance( + simulation: object, + expected_selection: Mapping[str, object] | None, +) -> SPMProvenance | None: + """Return the simulation's SPM receipt, checked against the resolved selection. + + Raises SPMInputError for a receipt that cannot be certified, and + ValueError for a receipt where none was expected, a missing receipt, or a + receipt for another selection. + """ + receipt = _calculation_receipt(simulation) + if expected_selection is None: + if receipt is None: + return None + msg = "Unexpected SPM provenance for a calculation without SPM" + raise ValueError(msg) + if receipt is None: + msg = "The calculation produced no SPM receipt" + raise ValueError(msg) + receipt_selection = { + "forecast_content_sha256": receipt.forecast_sha256, + "scenario": receipt.scenario, + "geography_kind": receipt.geography_kind, + "geography_id": receipt.geography_id, + "county_vintage": receipt.county_vintage, + "as_of": receipt.as_of.isoformat() if receipt.as_of is not None else None, + } + if receipt_selection != dict(expected_selection): + msg = "SPM calculation configuration differs from the request" + raise ValueError(msg) + return receipt diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py index b5c2490a7..cbc29f387 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -34,6 +34,15 @@ def _forecast(expected_sha256: str | None) -> SPMForecast: return load_forecast(expected_sha256=expected_sha256) +def certified_spm_forecast(expected_sha256: str) -> SPMForecast: + """Return the packaged SPM forecast whose content has ``expected_sha256``. + + Loaded once per process and digest; raises when the packaged forecast has + another digest. + """ + return _forecast(expected_sha256) + + @lru_cache(maxsize=1) def runtime_spm_capability() -> SPMCapability | None: """Return the installed image's certified SPM capability. diff --git a/projects/policyengine-simulation-executor/tests/test_household_spm.py b/projects/policyengine-simulation-executor/tests/test_household_spm.py new file mode 100644 index 000000000..1a3cd975a --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_household_spm.py @@ -0,0 +1,294 @@ +"""Tests for household SPM selection and receipts, ported from the Public API.""" + +from __future__ import annotations + +import pytest +from policyengine_simulation_contract.spm import SPMInputError, SPMSelection +from policyengine_simulation_executor import household_spm + +FORECAST_SHA = "3d86d5c4c0423480e6b69b75d222ffa4a7a2639e4094df5ba2504af01be17173" +DEFAULTS = { + "as_of": None, + "county_vintage": "2020", + "forecast_content_sha256": FORECAST_SHA, + "geography_id": None, + "geography_kind": "county", + "scenario": "ce_trend", +} +RUNTIME_VERSIONS = { + "policyengine": "6.7.5", + "policyengine-core": "3.33.0", + "policyengine-us": "2.2.1", + "spm-calculator": "1.0.0", +} + + +class FakeForecast: + years = (2025, 2026) + + def __init__(self, *, metro_areas=("31080",)): + self.metro_areas = metro_areas + + def entry(self, year, *, scenario, as_of): + if scenario not in {"ce_trend", "flat"}: + msg = f"Unknown forecast scenario: {scenario}" + raise ValueError(msg) + + def geography_factor(self, year, tenure, **options): + if options["geoid"] not in self.metro_areas: + msg = f"unknown area {options['geoid']}" + raise ValueError(msg) + return 1.0 + + +@pytest.fixture +def certified_bundle(monkeypatch): + monkeypatch.setattr( + household_spm, "_bundle_spm_configuration", lambda: dict(DEFAULTS) + ) + monkeypatch.setattr(household_spm, "_us_model_supports_spm", lambda: True) + monkeypatch.setattr(household_spm, "_forecast", lambda sha: FakeForecast()) + + +def _resolve(selection=None, *, country="us", saved_household=False): + return household_spm.resolve_household_spm_selection( + country, + None if selection is None else SPMSelection.model_validate(selection), + saved_household=saved_household, + ) + + +def _error(selection=None, **options) -> SPMInputError: + with pytest.raises(SPMInputError) as raised: + _resolve(selection, **options) + return raised.value + + +def test_omitted_selection_resolves_the_certified_defaults(certified_bundle) -> None: + assert _resolve() == DEFAULTS + + +def test_explicit_choices_override_defaults(certified_bundle) -> None: + assert _resolve({"geography_kind": "metro", "geography_id": "31080"}) == { + **DEFAULTS, + "geography_kind": "metro", + "geography_id": "31080", + } + assert _resolve({"scenario": "flat"})["scenario"] == "flat" + + +def test_other_countries_accept_no_selection(certified_bundle) -> None: + assert _resolve(country="uk") is None + error = _error({}, country="uk") + assert (error.code, error.message) == ( + "SPM_SETTINGS_UNSUPPORTED", + "SPM settings are only available for the US", + ) + + +def test_a_foreign_artifact_hash_is_reported_by_where_it_came_from( + certified_bundle, +) -> None: + selection = {"forecast_content_sha256": "f" * 64} + + requested = _error(selection) + saved = _error(selection, saved_household=True) + + assert (requested.code, requested.message) == ( + "SPM_SETTINGS_INVALID", + "SPM selection does not match this bundle's artifact hash", + ) + assert saved.code == "SPM_CONFIGURATION_UNAVAILABLE" + assert "saved selection" in saved.message + + +@pytest.mark.parametrize( + ("selection", "code", "message"), + [ + ( + {"scenario": "unknown"}, + "SPM_SETTINGS_INVALID", + "Unknown forecast scenario: unknown", + ), + ( + {"county_vintage": "2010"}, + "SPM_SETTINGS_INVALID", + "Unsupported county vintage: use 2020", + ), + ( + {"geography_kind": "metro", "geography_id": "99999"}, + "SPM_GEOGRAPHY_UNAVAILABLE", + "SPM area is unavailable in the selected artifact: 99999", + ), + ], +) +def test_invalid_selections_use_public_api_codes( + certified_bundle, selection, code, message +) -> None: + error = _error(selection) + + assert (error.code, error.message) == (code, message) + + +def test_a_bundle_without_spm_keeps_legacy_behaviour_only_for_old_models( + monkeypatch, +) -> None: + monkeypatch.setattr(household_spm, "_bundle_spm_configuration", lambda: None) + monkeypatch.setattr(household_spm, "_us_model_supports_spm", lambda: False) + + assert _resolve() is None + assert _error({}).code == "SPM_SETTINGS_UNSUPPORTED" + + monkeypatch.setattr(household_spm, "_us_model_supports_spm", lambda: True) + assert _error().code == "SPM_CONFIGURATION_UNAVAILABLE" + + +def test_an_unreadable_bundle_manifest_fails_closed(monkeypatch) -> None: + def unreadable_bundle(): + msg = "bundle manifest is missing" + raise ValueError(msg) + + monkeypatch.setattr(household_spm, "_current_bundle", unreadable_bundle) + + error = _error() + + assert (error.code, error.message) == ( + "SPM_CONFIGURATION_UNAVAILABLE", + "The installed bundle manifest is invalid", + ) + + +def test_changing_the_geography_kind_drops_an_inherited_area(monkeypatch) -> None: + monkeypatch.setattr( + household_spm, + "_bundle_spm_configuration", + lambda: {**DEFAULTS, "geography_kind": "metro", "geography_id": "31080"}, + ) + monkeypatch.setattr(household_spm, "_us_model_supports_spm", lambda: True) + monkeypatch.setattr(household_spm, "_forecast", lambda sha: FakeForecast()) + + assert _resolve({"geography_kind": "national"}) == { + **DEFAULTS, + "geography_kind": "national", + } + assert _resolve()["geography_id"] == "31080" + + +def test_an_uncertifiable_bundle_configuration_fails_closed(monkeypatch) -> None: + monkeypatch.setattr( + household_spm, + "_bundle_spm_configuration", + lambda: {**DEFAULTS, "scenario": None}, + ) + monkeypatch.setattr(household_spm, "_us_model_supports_spm", lambda: True) + + assert _error().code == "SPM_CONFIGURATION_UNAVAILABLE" + + +class FakeSimulation: + def __init__(self, *, config=None, receipt=None): + self.spm_config = dict(DEFAULTS) if config is None else config + self.receipt = receipt + + def spm_provenance(self): + return self.receipt + + +def _country_receipt(**overrides) -> dict: + return { + "forecast_id": "spm-forecast-2026", + "forecast_sha256": FORECAST_SHA, + "scenario": "ce_trend", + "geography_kind": "county", + "years": {"2027": {}, "2026": {}}, + "runtime_versions": RUNTIME_VERSIONS, + "composition_method": "ignored by the compact receipt", + **overrides, + } + + +def test_the_country_receipt_becomes_the_compact_public_receipt() -> None: + provenance = household_spm.household_spm_provenance( + FakeSimulation(receipt=_country_receipt()), DEFAULTS + ) + + assert provenance is not None + assert provenance.model_dump(mode="json", by_alias=True) == { + "schema_version": "canonical-spm-provenance-v2", + "forecast_id": "spm-forecast-2026", + "forecast_sha256": FORECAST_SHA, + "scenario": "ce_trend", + "geography_kind": "county", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + "years": ["2026", "2027"], + "runtime_versions": RUNTIME_VERSIONS, + } + + +def test_a_country_without_spm_has_no_receipt() -> None: + assert household_spm.household_spm_provenance(object(), None) is None + + +def test_an_uncertifiable_country_receipt_is_an_spm_error() -> None: + with pytest.raises(SPMInputError) as raised: + household_spm.household_spm_provenance( + FakeSimulation(receipt=_country_receipt(scenario="flat")), DEFAULTS + ) + + assert raised.value.code == "SPM_CONFIGURATION_UNAVAILABLE" + assert raised.value.message.startswith( + "The installed country model reported an SPM receipt this API cannot certify" + ) + + +def test_a_receipt_for_another_selection_is_rejected() -> None: + with pytest.raises(ValueError, match="differs from the request"): + household_spm.household_spm_provenance( + FakeSimulation(receipt=_country_receipt()), + {**DEFAULTS, "county_vintage": "2010"}, + ) + with pytest.raises(ValueError, match="without SPM"): + household_spm.household_spm_provenance( + FakeSimulation(receipt=_country_receipt()), None + ) + + +class CountryInputError(ValueError): + def __init__(self, code: str, message: str) -> None: + super().__init__(message) + self.code = code + + +def test_spm_errors_are_recognized_through_wrapping_exceptions() -> None: + cause = CountryInputError("SPM_GEOGRAPHY_REQUIRED", "needs a county") + wrapped = RuntimeError("while calculating spm_unit_spm_threshold") + wrapped.__cause__ = cause + + detail = household_spm.household_spm_error_detail(wrapped) + + assert detail is not None + assert (detail.code, detail.message) == ("SPM_GEOGRAPHY_REQUIRED", "needs a county") + assert ( + household_spm.household_spm_error_detail( + SPMInputError("SPM_SETTINGS_INVALID", "bad") + ) + is not None + ) + assert ( + household_spm.household_spm_error_detail( + CountryInputError("SPM_SCENARIO_UNAVAILABLE", "no scenario") + ) + is None + ) + assert household_spm.household_spm_error_detail(ValueError("plain")) is None + + +def test_the_certified_forecast_is_loaded_once_by_its_digest() -> None: + from policyengine_simulation_executor.spm import certified_spm_forecast + + forecast = certified_spm_forecast(FORECAST_SHA) + + assert certified_spm_forecast(FORECAST_SHA) is forecast + assert forecast.years From bf9d7793c749f065aae1550a8ce03498addfc9d0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:01:34 +0400 Subject: [PATCH 05/19] Register observability stages for the v2 household functions A v2 household function call is measured as one operation that covers the whole call, with nested spans for loading the country model, preparing the household, and calculating it under each policy. The stage plan lists those names so the household functions open spans only from registered stages. Co-Authored-By: Claude Opus 5.5 --- .../stages.py | 17 +++++++++++++++++ .../tests/test_stages.py | 12 ++++++++++++ 2 files changed, 29 insertions(+) diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py index aaa35eadd..14036f212 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py @@ -26,6 +26,7 @@ class RunConfiguration(StrEnum): STAGE12_CANONICAL_REPORT = "stage12_canonical_report" STAGE12_SIMULATION = "stage12_simulation" STAGE12_SEGMENT = "stage12_segment" + V2_HOUSEHOLD = "v2_household" class Stage(StrEnum): @@ -111,6 +112,10 @@ class Stage(StrEnum): STAGE12_AGGREGATION = "stage12_aggregation" STAGE12_AGGREGATE_ARTIFACT_WRITE = "stage12_aggregate_artifact_write" STAGE12_RESULT_COMPARISON = "stage12_result_comparison" + V2_HOUSEHOLD_EXECUTION = "v2_household_execution" + V2_HOUSEHOLD_COUNTRY_MODEL_LOAD = "v2_household_country_model_load" + V2_HOUSEHOLD_INPUT_NORMALIZATION = "v2_household_input_normalization" + V2_HOUSEHOLD_CALCULATION = "v2_household_calculation" @dataclass(frozen=True) @@ -271,6 +276,17 @@ def name(self, stage: Stage) -> str: Stage.STAGE12_CALCULATION, ), ), + # One v2 household function call: loading the country model, + # preparing the household, then one calculation per policy. + RunConfiguration.V2_HOUSEHOLD: StagePlan( + RunConfiguration.V2_HOUSEHOLD, + ( + Stage.V2_HOUSEHOLD_EXECUTION, + Stage.V2_HOUSEHOLD_COUNTRY_MODEL_LOAD, + Stage.V2_HOUSEHOLD_INPUT_NORMALIZATION, + Stage.V2_HOUSEHOLD_CALCULATION, + ), + ), } ) @@ -285,3 +301,4 @@ def name(self, stage: Stage) -> str: ] STAGE12_SIMULATION_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.STAGE12_SIMULATION] STAGE12_SEGMENT_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.STAGE12_SEGMENT] +V2_HOUSEHOLD_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.V2_HOUSEHOLD] diff --git a/libs/policyengine-simulation-observability/tests/test_stages.py b/libs/policyengine-simulation-observability/tests/test_stages.py index fbba9f721..0861363b9 100644 --- a/libs/policyengine-simulation-observability/tests/test_stages.py +++ b/libs/policyengine-simulation-observability/tests/test_stages.py @@ -54,3 +54,15 @@ def test_annual_status_polls_measure_the_job_progress_read() -> None: annual = RUN_STAGE_REGISTRY[RunConfiguration.ANNUAL_IMPACT] assert annual.name(Stage.MODAL_JOB_PROGRESS_READ) == "modal_job_progress_read" + + +def test_v2_household_plan_measures_model_load_inputs_and_each_policy() -> None: + household = RUN_STAGE_REGISTRY[RunConfiguration.V2_HOUSEHOLD] + + assert household.stages == ( + Stage.V2_HOUSEHOLD_EXECUTION, + Stage.V2_HOUSEHOLD_COUNTRY_MODEL_LOAD, + Stage.V2_HOUSEHOLD_INPUT_NORMALIZATION, + Stage.V2_HOUSEHOLD_CALCULATION, + ) + assert household.name(Stage.V2_HOUSEHOLD_CALCULATION) == "v2_household_calculation" From d52d956f0eb1b9259aec1622e7c95a7dd365d29a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:27:01 +0400 Subject: [PATCH 06/19] Calculate a household under several policies as the Public API does calculate_household runs one household request: it checks the request names the installed bundle, resolves the SPM selection, prepares the household once, then builds a simulation and calculates the requested values for each policy in order. It returns each policy's calculated household, warnings and SPM receipt with the applied SPM selection and bundle provenance, or a typed failure. Each step follows the Public API's household calculation: reform values are converted to each parameter's value type, values are written as its country module writes them, a value that cannot be calculated is null (one null per point, with a warning, when the household has axes), and an SPM error ends the calculation only for a chosen measurement or a configuration the build cannot certify. Failures carry the Public API's messages: an invalid household payload, unrecognized inputs, an SPM error, or an error calculating the household under a policy. Country models are loaded once per process. The tests calculate real US and UK households, including US households with axes and with every missing variable added. Co-Authored-By: Claude Opus 5.5 --- .../household_calculation.py | 664 ++++++++++++++ .../tests/test_household_calculation.py | 858 ++++++++++++++++++ 2 files changed, 1522 insertions(+) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py create mode 100644 projects/policyengine-simulation-executor/tests/test_household_calculation.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py new file mode 100644 index 000000000..9a9b29c26 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py @@ -0,0 +1,664 @@ +"""Household calculations for the v2 Modal household functions. + +``calculate_household`` runs one ``HouseholdExecutionRequest``: it calculates +the household under each requested policy, in order, and returns a +``HouseholdExecutionResult`` or a typed ``HouseholdExecutionFailure``. + +The calculation reproduces the Public API's in-process household calculation +(its household calculation service, household routes and country module): the +same input preparation, simulation construction, reform value conversion, +requested values, value formatting, handling of values that cannot be +calculated, warnings and SPM receipt. Failures carry the codes and messages of +the Public API's matching error responses. This keeps v2 results comparable +with the Public API's results while both exist. + +Two values are written differently, because the contract carries only finite +numbers and JSON scalars: an undefined number is the string "NaN" (the Public +API writes a bare ``NaN``, which is not JSON), and a date is the HTTP date +text the Public API's Flask encoder writes. +""" + +from __future__ import annotations + +import importlib +import json +import logging +import math +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass, field +from datetime import UTC, date, datetime, time +from email.utils import format_datetime +from functools import lru_cache +from typing import TYPE_CHECKING, Any + +from policyengine_core.enums import Enum +from policyengine_core.errors import SituationParsingError +from policyengine_core.parameters import get_parameter +from policyengine_core.periods import period as parse_period +from policyengine_simulation_contract.household_execution import ( + CalculatedHousehold, + HouseholdBundleProvenance, + HouseholdCalculationError, + HouseholdCalculationOutput, + HouseholdExecutionFailure, + HouseholdExecutionResult, + HouseholdPolicyResult, +) +from policyengine_simulation_contract.spm import ( + SPMErrorDetail, + SPMInputError, + SPMResolvedConfiguration, +) +from policyengine_simulation_observability import scopes +from policyengine_simulation_observability.stages import V2_HOUSEHOLD_STAGES, Stage + +from policyengine_simulation_executor.household_inputs import ( + HouseholdModelMetadata, + RequestedValue, + add_missing_variables, + drop_deprecated_inputs, + find_unrecognized_household_inputs, + find_unrecognized_policy_inputs, + requested_computations, + unrecognized_inputs_message, +) +from policyengine_simulation_executor.household_spm import ( + HOUSEHOLD_SPM_INPUT_ERROR_CODES, + household_spm_error_detail, + household_spm_provenance, + resolve_household_spm_selection, +) +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping + + from policyengine_core.periods import Period + from policyengine_core.simulations import Simulation + from policyengine_core.taxbenefitsystems import TaxBenefitSystem + from policyengine_observability import ObservabilityRuntime + from policyengine_simulation_contract.household_execution import ( + HouseholdErrorCode, + HouseholdExecutionRequest, + HouseholdPolicy, + HouseholdWorkerIdentity, + ) + from policyengine_simulation_contract.spm import SPMProvenance + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class CountryHouseholdModel: + """A loaded country model: its ``Simulation`` type, tax-benefit system and metadata. + + The tax-benefit system supplies the metadata and the value types of + reform parameters; simulations build their own, as in the Public API. + """ + + country: str + simulation_type: Callable[..., Simulation] + tax_benefit_system: TaxBenefitSystem + metadata: HouseholdModelMetadata + + +@lru_cache(maxsize=4) +def load_country_household_model(country: str) -> CountryHouseholdModel: + """Import ``country``'s model package and build its tax-benefit system once. + + The package is the one the installed PolicyEngine.py bundle names for the + country. The result is kept for the life of the process, so warm Modal + containers reuse it. Raises ValueError for a country the bundle lacks. + """ + bundle = load_stage12_bundle().bundle + country_bundle = next( + (item for item in bundle.countries if item.country == country), None + ) + if country_bundle is None: + msg = f"the installed bundle does not carry country {country!r}" + raise ValueError(msg) + package = importlib.import_module( + country_bundle.country_package_name.replace("-", "_") + ) + system = package.CountryTaxBenefitSystem() + return CountryHouseholdModel( + country=country, + simulation_type=package.Simulation, + tax_benefit_system=system, + metadata=HouseholdModelMetadata.from_tax_benefit_system(system), + ) + + +def installed_household_bundle(country: str) -> HouseholdBundleProvenance: + """Return the provenance of the installed bundle for ``country``.""" + resolved = load_stage12_bundle() + return HouseholdBundleProvenance.from_bundle( + resolved.bundle, + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + country=country, + ) + + +@dataclass(frozen=True) +class HouseholdCalculationDependencies: + """Seams a household calculation uses; unset seams use the defaults. + + The defaults are the installed bundle, the process's cached country + models, today's UTC date (which sets the default household year) and no + telemetry runtime. + """ + + installed_bundle: HouseholdBundleProvenance | None = None + model_loader: Callable[[str], CountryHouseholdModel] = load_country_household_model + today: date | None = None + runtime: ObservabilityRuntime | None = None + + +class _CalculationFailedError(Exception): + """A household calculation ended with a typed error.""" + + def __init__(self, error: HouseholdCalculationError) -> None: + super().__init__(error.message) + self.error = error + + +def _failure( + code: HouseholdErrorCode, + message: str, + *, + policy_key: str | None = None, +) -> _CalculationFailedError: + return _CalculationFailedError( + HouseholdCalculationError(code=code, message=message, policy_key=policy_key) + ) + + +def _spm_failure( + error: BaseException, *, policy_key: str | None +) -> _CalculationFailedError | None: + detail = household_spm_error_detail(error) + if detail is None: + return None + return _CalculationFailedError( + HouseholdCalculationError( + code="invalid_spm_selection", + message=detail.message, + spm_error=detail, + policy_key=policy_key, + ) + ) + + +def _invalid_household_failure( + error: SituationParsingError, *, policy_key: str +) -> _CalculationFailedError: + """Report a household the country model cannot read, as the Public API does.""" + return _failure( + "invalid_household", + f"Invalid household payload: {error}", + policy_key=policy_key, + ) + + +def _bundle_mismatch_failure( + bundle: HouseholdBundleProvenance, +) -> _CalculationFailedError: + """Report a request for a bundle other than ``bundle``, the one installed.""" + return _failure( + "bundle_mismatch", + "This household function runs PolicyEngine.py " + f"{bundle.policyengine_version} with {bundle.country_package_name} " + f"{bundle.country_package_version} " + f"(bundle {bundle.bundle_manifest_sha256}), not the requested bundle", + ) + + +def _unexpected_failure( + error: BaseException, *, policy_key: str +) -> _CalculationFailedError: + """Report an SPM error as such, and anything else as a failed calculation.""" + failure = _spm_failure(error, policy_key=policy_key) + if failure is not None: + return failure + logger.error( + "Error calculating household under policy %s", policy_key, exc_info=error + ) + return _failure( + "calculation_failed", + f"Error calculating household under policy: {error}", + policy_key=policy_key, + ) + + +def _stage_scope( + runtime: ObservabilityRuntime | None, + stage: Stage, + attributes: Mapping[str, object], +) -> AbstractContextManager[object]: + if runtime is None: + return nullcontext() + return scopes.span( + runtime, V2_HOUSEHOLD_STAGES.name(stage), attributes=dict(attributes) + ) + + +# Value formatting, as the Public API's country module and JSON encoder +# write values. + + +def _serialize_float(value: object) -> float | str: + """Write a number as the Public API does, naming non-finite values.""" + serialized = float(str(value)) + if math.isinf(serialized): + return "Infinity" if serialized > 0 else "-Infinity" + if math.isnan(serialized): + return "NaN" + return serialized + + +def _serialize_scalar(value: object) -> object: + """Write a date as the HTTP date text the Public API's Flask encoder writes.""" + if isinstance(value, datetime): + moment = value if value.tzinfo is not None else value.replace(tzinfo=UTC) + return format_datetime(moment.astimezone(UTC), usegmt=True) + if isinstance(value, date): + return format_datetime(datetime.combine(value, time(), tzinfo=UTC), usegmt=True) + return value + + +def _serialize_axis_values( + result: Any, # noqa: ANN401 # a NumPy array or EnumArray from the country model + value_type: type, + entity_index: int, + entity_count: int, +) -> list[object]: + if value_type == Enum: + result = result.decode_to_str() + entity_values: list[object] = ( + result.reshape((-1, entity_count)).T[entity_index].tolist() + ) + if value_type is float: + return [_serialize_float(value) for value in entity_values] + if value_type is str: + return [str(value) for value in entity_values] + return [_serialize_scalar(value) for value in entity_values] + + +def _serialize_entity_value( + result: Any, # noqa: ANN401 # a NumPy array or EnumArray from the country model + value_type: type, + entity_index: int, +) -> object: + if value_type == Enum: + return result.decode()[entity_index].name + if value_type is float: + return _serialize_float(result[entity_index]) + if value_type is str: + return str(result[entity_index]) + return _serialize_scalar(result.tolist()[entity_index]) + + +@dataclass(frozen=True) +class _PreparedHousehold: + """A household ready to calculate under each policy. + + ``situation`` has its missing variables added and deprecated inputs + removed; ``input_warnings`` say which inputs were removed. + ``spm_selection`` is the resolved SPM selection (None when no SPM + applies) and ``spm_requested`` whether the caller chose it. + """ + + situation: dict[str, Any] + input_warnings: tuple[str, ...] + spm_selection: dict[str, Any] | None + spm_requested: bool + + +def _normalized_reform( + reform: Mapping[str, Mapping[str, object]], + system: TaxBenefitSystem, +) -> dict[str, dict[str, object]]: + """Convert reform values to each parameter's current value type. + + Integers become floats, and numeric text becomes a number first, as in the + Public API. + """ + root = system.parameters + if root is None: + msg = "the country model has no parameters" + raise LookupError(msg) + normalized: dict[str, dict[str, object]] = {} + for parameter_name, parameter_updates in reform.items(): + parameter = get_parameter(root, parameter_name) + convert = _reform_value_type(parameter.values_list[-1].value) + normalized[parameter_name] = { + time_period: convert(_numeric_if_possible(value)) + for time_period, value in parameter_updates.items() + } + return normalized + + +def _reform_value_type(current_value: object) -> Callable[[Any], object]: + """Return the type a reform value becomes: the current value's, int as float.""" + value_type = type(current_value) + return float if value_type is int else value_type + + +def _numeric_if_possible(value: object) -> object: + """Return ``value`` as a float when it reads as one, as the Public API does.""" + if not isinstance(value, (bool, int, float, str)): + return value + try: + return float(value) + except (ValueError, OverflowError): + return value + + +def _create_simulation( + model: CountryHouseholdModel, + prepared: _PreparedHousehold, + policy: HouseholdPolicy, +) -> Simulation: + """Build the country model's simulation of the household under ``policy``. + + The country ``Simulation`` takes the reform as converted parameter values + and the resolved SPM selection, as the Public API passes them. + """ + options: dict[str, object] = { + "situation": json.loads(json.dumps(prepared.situation)) + } + if prepared.spm_selection is not None: + options["spm"] = dict(prepared.spm_selection) + if policy.parameters: + options["reform"] = _normalized_reform( + policy.parameters, model.tax_benefit_system + ) + return model.simulation_type(**options) + + +def _model_period(text: str) -> Period: + """Read a household period key as policyengine-core's ``calculate`` reads it. + + Raises ValueError for text that names no period. + """ + parsed = parse_period(text) + if parsed is None: + msg = f"{text!r} names no period" + raise ValueError(msg) + return parsed + + +@dataclass +class _PolicyCalculation: + """One policy's simulation and the household its values are written into.""" + + simulation: Simulation + household: dict[str, Any] + spm_requested: bool + warnings: list[str] = field(default_factory=list) + + @property + def has_axes(self) -> bool: + """Whether the household is calculated at several axis points.""" + return "axes" in self.household + + def calculate_all(self) -> None: + """Calculate every requested value, as the Public API's country module does.""" + system = self.simulation.tax_benefit_system + computations = requested_computations( + self.household, + include_provided_values=self.has_axes, + variable_names=set(system.variables) if self.has_axes else None, + ) + for requested in computations: + try: + self._calculate_value(requested) + except Exception as error: + detail = household_spm_error_detail(error) + # A chosen measurement reports its missing inputs. An + # inherited one leaves its dependants uncalculated, but only + # for a missing input: a configuration failure always ends + # the calculation. + if detail is not None and ( + self.spm_requested + or detail.code not in HOUSEHOLD_SPM_INPUT_ERROR_CODES + ): + raise + self._record_failure(requested, error) + + def _calculate_value(self, requested: RequestedValue) -> None: + entity_plural, entity_id, variable_name, period = requested + population = self.simulation.get_population(entity_plural) + variable = self.simulation.tax_benefit_system.get_variable(variable_name) + if variable is None: + msg = f"the country model has no variable {variable_name}" + raise LookupError(msg) + value_type = variable.value_type + result = self.simulation.calculate(variable_name, _model_period(period)) + entities = self.household[entity_plural] + if self.has_axes: + value: object = _serialize_axis_values( + result, value_type, list(entities).index(entity_id), len(entities) + ) + else: + value = _serialize_entity_value( + result, value_type, population.get_index(entity_id) + ) + entities[entity_id][variable_name][period] = value + + def _record_failure(self, requested: RequestedValue, error: Exception) -> None: + entity_plural, entity_id, variable_name, period = requested + cells = self.household[entity_plural][entity_id][variable_name] + if not self.has_axes: + cells[period] = None + logger.warning( + "Unable to calculate %s for %s in %s: %s", + variable_name, + entity_id, + period, + error, + ) + return + population = self.simulation.get_population(entity_plural) + axis_point_count = population.count // len(self.household[entity_plural]) + cells[period] = [None] * axis_point_count + warning = ( + f"Unable to calculate {variable_name} for {entity_id} in " + f"{period}; returned {axis_point_count} null axis values: {error}" + ) + self.warnings.append(warning) + logger.warning("%s", warning) + + +def _parsed_policy_calculation( + model: CountryHouseholdModel, + prepared: _PreparedHousehold, + policy: HouseholdPolicy, +) -> _PolicyCalculation: + """Build the policy's simulation, failing as the Public API's parse stage does.""" + try: + simulation = _create_simulation(model, prepared, policy) + except SituationParsingError as error: + raise _invalid_household_failure(error, policy_key=policy.policy_key) from error + except Exception as error: + raise _unexpected_failure(error, policy_key=policy.policy_key) from error + return _PolicyCalculation( + simulation=simulation, + household=json.loads(json.dumps(prepared.situation)), + spm_requested=prepared.spm_requested, + ) + + +def _calculated_policy_result( + calculation: _PolicyCalculation, + prepared: _PreparedHousehold, + policy: HouseholdPolicy, +) -> HouseholdPolicyResult: + """Calculate the requested values and return the policy's result.""" + calculation.calculate_all() + spm_provenance: SPMProvenance | None = household_spm_provenance( + calculation.simulation, prepared.spm_selection + ) + return HouseholdPolicyResult( + policy_key=policy.policy_key, + household=CalculatedHousehold.model_validate(calculation.household), + warnings=(*prepared.input_warnings, *calculation.warnings), + spm_provenance=spm_provenance, + ) + + +def _calculate_policy( + model: CountryHouseholdModel, + prepared: _PreparedHousehold, + policy: HouseholdPolicy, +) -> HouseholdPolicyResult: + """Calculate the household under one policy, raising a typed failure.""" + calculation = _parsed_policy_calculation(model, prepared, policy) + try: + return _calculated_policy_result(calculation, prepared, policy) + except SituationParsingError as error: + failure = _spm_failure( + error, policy_key=policy.policy_key + ) or _invalid_household_failure(error, policy_key=policy.policy_key) + raise failure from error + except Exception as error: + raise _unexpected_failure(error, policy_key=policy.policy_key) from error + + +def _resolved_spm_selection( + request: HouseholdExecutionRequest, +) -> dict[str, Any] | None: + try: + return resolve_household_spm_selection( + request.country, + request.spm, + saved_household=request.spm_selection_source == "saved_household", + ) + except SPMInputError as error: + raise _CalculationFailedError( + HouseholdCalculationError( + code="invalid_spm_selection", + message=error.message, + spm_error=SPMErrorDetail(code=error.code, message=error.message), + ) + ) from error + + +def _prepared_household( + request: HouseholdExecutionRequest, + model: CountryHouseholdModel, + *, + today: date, + spm_selection: dict[str, Any] | None, +) -> _PreparedHousehold: + """Prepare the household as the Public API does; raise for unknown names. + + Unknown household names are reported with the first policy, as the Public + API's baseline request would report them; unknown policy parameters are + reported with the policy that names them. + """ + situation = request.household.to_situation() + if request.add_missing_variables: + add_missing_variables(situation, model.metadata, today=today) + deprecated = drop_deprecated_inputs(situation) + household_errors = find_unrecognized_household_inputs( + deprecated.household, model.metadata + ) + for policy in request.policies: + policy_errors = find_unrecognized_policy_inputs( + policy.parameters, model.metadata + ) + input_errors = (*household_errors, *policy_errors) + if input_errors: + raise _CalculationFailedError( + HouseholdCalculationError( + code="invalid_inputs", + message=unrecognized_inputs_message(input_errors), + input_errors=input_errors, + policy_key=policy.policy_key if policy_errors else None, + ) + ) + return _PreparedHousehold( + situation=deprecated.household, + input_warnings=tuple(warning.message for warning in deprecated.warnings), + spm_selection=spm_selection, + spm_requested=request.spm_requested, + ) + + +def _calculate_output( + request: HouseholdExecutionRequest, + bundle: HouseholdBundleProvenance, + dependencies: HouseholdCalculationDependencies, +) -> HouseholdCalculationOutput: + """Calculate every policy of ``request``, raising a typed failure. + + Checks run in the Public API's order: the bundle, the SPM selection, the + household's and policies' names, then each policy. + """ + if request.bundle != bundle: + raise _bundle_mismatch_failure(bundle) + spm_selection = _resolved_spm_selection(request) + country = {"country": request.country} + runtime = dependencies.runtime + with _stage_scope(runtime, Stage.V2_HOUSEHOLD_COUNTRY_MODEL_LOAD, country): + model = dependencies.model_loader(request.country) + with _stage_scope(runtime, Stage.V2_HOUSEHOLD_INPUT_NORMALIZATION, country): + prepared = _prepared_household( + request, + model, + today=dependencies.today or datetime.now(UTC).date(), + spm_selection=spm_selection, + ) + results = [] + for policy in request.policies: + attributes = {**country, "policy_key": policy.policy_key} + with _stage_scope(runtime, Stage.V2_HOUSEHOLD_CALCULATION, attributes): + results.append(_calculate_policy(model, prepared, policy)) + return HouseholdCalculationOutput( + country=request.country, + bundle=bundle, + spm_config=( + SPMResolvedConfiguration.model_validate(spm_selection) + if spm_selection is not None + else None + ), + results=tuple(results), + ) + + +def calculate_household( + request: HouseholdExecutionRequest, + *, + worker: HouseholdWorkerIdentity, + required_country: str | None = None, + dependencies: HouseholdCalculationDependencies | None = None, +) -> HouseholdExecutionResult | HouseholdExecutionFailure: + """Calculate one household under each requested policy. + + Returns the per-policy results, or a typed failure for a request for + another bundle, an unusable SPM selection, unrecognized inputs, a household + the country model cannot build, or any other calculation error. Raises + ValueError when ``required_country`` is set and the request names another + country; a household function calculates its own country only. + """ + if required_country is not None and request.country != required_country: + msg = f"this function calculates {required_country} households only" + raise ValueError(msg) + dependencies = dependencies or HouseholdCalculationDependencies() + bundle = dependencies.installed_bundle or installed_household_bundle( + request.country + ) + try: + output = _calculate_output(request, bundle, dependencies) + except _CalculationFailedError as failure: + return HouseholdExecutionFailure( + calculation_id=request.calculation_id, + error=failure.error, + worker=worker, + ) + return HouseholdExecutionResult( + calculation_id=request.calculation_id, + output=output, + worker=worker, + ) diff --git a/projects/policyengine-simulation-executor/tests/test_household_calculation.py b/projects/policyengine-simulation-executor/tests/test_household_calculation.py new file mode 100644 index 000000000..c6874b31e --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_household_calculation.py @@ -0,0 +1,858 @@ +"""Tests for v2 household calculations, ported from the Public API's behaviour. + +The first tests use a stand-in country model to cover the failure paths. The +later ones calculate with the installed US and UK models, as the Public API's +in-process calculation does; the US model takes about a minute to load once. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from datetime import date +from types import SimpleNamespace +from uuid import UUID + +import numpy as np +import pytest +from policyengine_core.errors import SituationParsingError +from policyengine_simulation_contract.household_execution import ( + HouseholdBundleProvenance, + HouseholdExecutionFailure, + HouseholdExecutionRequest, + HouseholdExecutionResult, + HouseholdWorkerIdentity, +) +from policyengine_simulation_contract.spm import SPMInputError +from policyengine_simulation_executor import household_calculation +from policyengine_simulation_executor.household_calculation import ( + CountryHouseholdModel, + HouseholdCalculationDependencies, + calculate_household, + installed_household_bundle, + load_country_household_model, +) +from policyengine_simulation_executor.household_inputs import ( + HouseholdEntity, + HouseholdModelMetadata, + HouseholdVariable, +) + +CALCULATION_ID = UUID("00000000-0000-4000-8000-0000000000c1") +TODAY = date(2026, 6, 1) + + +def _worker(country: str) -> HouseholdWorkerIdentity: + return HouseholdWorkerIdentity( + application_name="policyengine-simulation-v2-py6-7-5", + function_name=f"calculate_household_{country}", + ) + + +def _request(country: str, household: dict, **overrides) -> HouseholdExecutionRequest: + values = { + "calculation_id": CALCULATION_ID, + "country": country, + "household": household, + "policies": [{"policy_key": "baseline", "parameters": {}}], + "bundle": installed_household_bundle(country).model_dump(mode="json"), + } + values.update(overrides) + return HouseholdExecutionRequest.model_validate(values) + + +class RecordingRuntime: + """Records the spans a calculation opens, in order.""" + + def __init__(self) -> None: + self.spans: list[tuple[str, dict]] = [] + + @contextmanager + def span(self, name: str, *, attributes: dict | None = None): + self.spans.append((name, dict(attributes or {}))) + yield + + +# A stand-in country model: enough of a country package to exercise the +# failure paths without importing a real one. + + +class StandInPopulation: + def __init__(self, ids: list[str], count: int) -> None: + self.ids = ids + self.count = count + + def get_index(self, entity_id: str) -> int: + return self.ids.index(entity_id) + + +class StandInSystem: + def __init__(self, value_types: dict[str, type]) -> None: + self.value_types = value_types + self.variables = dict.fromkeys(value_types) + + def get_variable(self, name: str): + return SimpleNamespace(value_type=self.value_types[name]) + + +STAND_IN_VALUE_TYPES = { + "age": int, + "net_income": float, + "spm_resources": float, + "registration_date": date, +} + + +def _stand_in_model(simulation_type) -> CountryHouseholdModel: + metadata = HouseholdModelMetadata( + variables={ + name: HouseholdVariable( + name, + "person" if name == "age" else "household", + "year", + is_input=name == "age", + default_value=40 if name == "age" else None, + ) + for name in STAND_IN_VALUE_TYPES + }, + entities={ + "person": HouseholdEntity("person", "people", frozenset()), + "household": HouseholdEntity( + "household", "households", frozenset({"members"}) + ), + }, + parameters=frozenset(), + ) + system = StandInSystem(STAND_IN_VALUE_TYPES) + simulation_type.tax_benefit_system = system + return CountryHouseholdModel( + country="uk", + simulation_type=simulation_type, + tax_benefit_system=system, + metadata=metadata, + ) + + +class CountryInputError(ValueError): + def __init__(self, code: str, message: str) -> None: + super().__init__(message) + self.code = code + + +def _stand_in_simulation(outcomes: dict[str, object], *, axis_points: int = 1): + class StandInSimulation: + tax_benefit_system: StandInSystem + + def __init__(self, *, situation: dict, **options: object) -> None: + self.situation = situation + self.options = options + people = list(situation["people"]) + households = list(situation.get("households", {})) + self.populations = { + "people": StandInPopulation(people, len(people) * axis_points), + "households": StandInPopulation( + households, len(households) * axis_points + ), + } + + def get_population(self, plural: str) -> StandInPopulation: + return self.populations[plural] + + def calculate(self, variable_name: str, period: str): + outcome = outcomes[variable_name] + if isinstance(outcome, Exception): + raise outcome + return outcome + + return StandInSimulation + + +STAND_IN_HOUSEHOLD = { + "people": {"you": {"age": {"2026": 30}}}, + "households": { + "home": { + "members": ["you"], + "net_income": {"2026": None}, + "spm_resources": {"2026": None}, + } + }, +} + + +def _calculate_with_stand_in( + simulation_type, household=None, *, country="uk", **overrides +): + model = _stand_in_model(simulation_type) + return calculate_household( + _request(country, household or STAND_IN_HOUSEHOLD, **overrides), + worker=_worker(country), + dependencies=HouseholdCalculationDependencies( + model_loader=lambda _country: model, today=TODAY + ), + ) + + +def test_a_value_that_cannot_be_calculated_is_null_without_a_warning() -> None: + outcome = _calculate_with_stand_in( + _stand_in_simulation( + {"net_income": np.array([1.5]), "spm_resources": KeyError("missing")} + ) + ) + + assert isinstance(outcome, HouseholdExecutionResult) + result = outcome.output.results[0] + home = result.household.to_situation()["households"]["home"] + assert home["net_income"] == {"2026": 1.5} + assert home["spm_resources"] == {"2026": None} + assert result.warnings == () + assert outcome.output.spm_config is None + assert result.spm_provenance is None + + +def test_with_axes_a_failed_value_is_one_null_per_point_with_a_warning() -> None: + household = { + **STAND_IN_HOUSEHOLD, + "axes": [[{"name": "age", "count": 3, "min": 0, "max": 2}]], + } + + outcome = _calculate_with_stand_in( + _stand_in_simulation( + { + "age": np.array([0, 1, 2]), + "net_income": np.array([1.0, float("inf"), float("nan")]), + "spm_resources": KeyError("missing"), + }, + axis_points=3, + ), + household, + ) + + assert isinstance(outcome, HouseholdExecutionResult) + result = outcome.output.results[0] + situation = result.household.to_situation() + assert situation["people"]["you"]["age"] == {"2026": [0, 1, 2]} + assert situation["households"]["home"]["net_income"] == { + "2026": [1.0, "Infinity", "NaN"] + } + assert situation["households"]["home"]["spm_resources"] == { + "2026": [None, None, None] + } + assert situation["axes"] == household["axes"] + assert result.warnings == ( + "Unable to calculate spm_resources for home in 2026; returned 3 null " + "axis values: 'missing'", + ) + + +def test_dates_are_written_as_the_public_api_writes_them() -> None: + # The Public API's Flask JSON encoder writes a date as an HTTP date. + household = { + "people": {"you": {"age": {"2026": 30}}}, + "households": { + "home": {"members": ["you"], "registration_date": {"2026": None}} + }, + } + + outcome = _calculate_with_stand_in( + _stand_in_simulation( + {"registration_date": np.array(["1970-01-01"], dtype="datetime64[D]")} + ), + household, + ) + + assert isinstance(outcome, HouseholdExecutionResult) + home = outcome.output.results[0].household.to_situation()["households"]["home"] + assert home["registration_date"] == {"2026": "Thu, 01 Jan 1970 00:00:00 GMT"} + + +def test_an_spm_input_error_fails_only_a_chosen_measurement(monkeypatch) -> None: + error = CountryInputError("SPM_GEOGRAPHY_REQUIRED", "SPM needs a county") + outcomes = {"net_income": np.array([1.0]), "spm_resources": error} + resolved = { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "county", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + } + monkeypatch.setattr( + household_calculation, + "resolve_household_spm_selection", + lambda country, selection, *, saved_household: dict(resolved), + ) + monkeypatch.setattr( + household_calculation, + "household_spm_provenance", + lambda simulation, expected_selection: None, + ) + + inherited = _calculate_with_stand_in(_stand_in_simulation(outcomes), country="us") + chosen = _calculate_with_stand_in( + _stand_in_simulation(outcomes), country="us", spm={} + ) + + assert isinstance(inherited, HouseholdExecutionResult) + home = inherited.output.results[0].household.to_situation()["households"]["home"] + assert home["spm_resources"] == {"2026": None} + assert inherited.output.spm_config is not None + assert inherited.output.spm_config.model_dump(mode="json") == resolved + assert isinstance(chosen, HouseholdExecutionFailure) + assert chosen.error.code == "invalid_spm_selection" + assert chosen.error.spm_error is not None + assert chosen.error.spm_error.code == "SPM_GEOGRAPHY_REQUIRED" + assert chosen.error.message == "SPM needs a county" + assert chosen.error.policy_key == "baseline" + + +def test_the_selection_and_its_receipt_reach_each_simulation(monkeypatch) -> None: + resolved = { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + } + created: list[dict] = [] + receipts: list[tuple[object, object]] = [] + + class RecordingSimulation(_stand_in_simulation({"net_income": np.array([1.0])})): + def __init__(self, **options: object) -> None: + super().__init__(**options) + created.append(self.options) + + monkeypatch.setattr( + household_calculation, + "resolve_household_spm_selection", + lambda country, selection, *, saved_household: resolved, + ) + + def record_receipt(simulation, expected_selection): + receipts.append((simulation, expected_selection)) + + monkeypatch.setattr( + household_calculation, "household_spm_provenance", record_receipt + ) + household = { + "people": {"you": {"age": {"2026": 30}}}, + "households": {"home": {"members": ["you"], "net_income": {"2026": None}}}, + } + + _calculate_with_stand_in(RecordingSimulation, household, country="us") + + assert [options["spm"] for options in created] == [resolved] + assert [expected for _, expected in receipts] == [resolved] + + +def test_an_spm_configuration_error_fails_even_an_inherited_measurement( + monkeypatch, +) -> None: + # Only a missing SPM input leaves an inherited measurement's value null; + # a configuration the build cannot certify always ends the calculation. + error = SPMInputError("SPM_CONFIGURATION_UNAVAILABLE", "no certified artifact") + monkeypatch.setattr( + household_calculation, + "resolve_household_spm_selection", + lambda country, selection, *, saved_household: { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + }, + ) + + outcome = _calculate_with_stand_in( + _stand_in_simulation({"net_income": np.array([1.0]), "spm_resources": error}), + country="us", + ) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "invalid_spm_selection" + assert outcome.error.spm_error is not None + assert outcome.error.spm_error.code == "SPM_CONFIGURATION_UNAVAILABLE" + assert outcome.error.message == "no certified artifact" + + +@pytest.mark.parametrize( + ("simulation_error", "code", "message"), + [ + ( + SituationParsingError(["people"], "unknown person `ghost`"), + "invalid_household", + "Invalid household payload: ", + ), + ( + RuntimeError("model exploded"), + "calculation_failed", + "Error calculating household under policy: model exploded", + ), + ( + SPMInputError("SPM_SETTINGS_INVALID", "bad SPM settings"), + "invalid_spm_selection", + "bad SPM settings", + ), + ], +) +def test_simulation_build_failures_use_public_api_messages( + simulation_error: Exception, code: str, message: str +) -> None: + def failing_simulation(**_: object): + raise simulation_error + + outcome = _calculate_with_stand_in(failing_simulation) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == code + assert outcome.error.message.startswith(message) + assert outcome.error.policy_key == "baseline" + assert outcome.worker == _worker("uk") + assert outcome.calculation_id == CALCULATION_ID + + +def test_a_receipt_that_cannot_be_certified_is_an_spm_failure(monkeypatch) -> None: + def uncertifiable(simulation, expected_selection): + code = "SPM_CONFIGURATION_UNAVAILABLE" + raise SPMInputError(code, "cannot certify") + + monkeypatch.setattr( + household_calculation, "household_spm_provenance", uncertifiable + ) + + outcome = _calculate_with_stand_in( + _stand_in_simulation( + {"net_income": np.array([1.0]), "spm_resources": np.array([1.0])} + ) + ) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "invalid_spm_selection" + assert outcome.error.message == "cannot certify" + + +def test_unrecognized_inputs_are_listed_in_the_order_the_household_was_sent() -> None: + # As the Public API lists them: it walks the household as sent, here with + # ``axes`` first, so the axis error comes before the person's. + household = { + "axes": [[{"name": "not_an_axis_variable", "count": 2, "min": 0, "max": 1}]], + "people": {"you": {"age": {"2026": 30}, "not_a_variable": {"2026": 1}}}, + "households": {"home": {"members": ["you"], "net_income": {"2026": None}}}, + } + + outcome = _calculate_with_stand_in(_stand_in_simulation({}), household) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "invalid_inputs" + assert [error.path for error in outcome.error.input_errors] == [ + "household.axes[0][0].name", + "household.people.you.not_a_variable", + ] + assert outcome.error.message == ( + "Unrecognized calculate input(s): Unrecognized household axis variable " + "`not_an_axis_variable` at `household.axes[0][0].name`.; Unrecognized " + "household variable `not_a_variable` at " + "`household.people.you.not_a_variable`." + ) + + +def test_a_request_for_another_bundle_is_refused() -> None: + bundle = installed_household_bundle("uk").model_copy( + update={"policyengine_version": "0.0.1"} + ) + + outcome = _calculate_with_stand_in( + _stand_in_simulation({}), bundle=bundle.model_dump(mode="json") + ) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "bundle_mismatch" + assert outcome.error.policy_key is None + + +def test_a_function_refuses_another_countrys_household() -> None: + model = _stand_in_model(_stand_in_simulation({})) + + with pytest.raises(ValueError, match="calculates us households only"): + calculate_household( + _request("uk", STAND_IN_HOUSEHOLD), + worker=_worker("us"), + required_country="us", + dependencies=HouseholdCalculationDependencies( + model_loader=lambda _country: model + ), + ) + + +def test_the_installed_bundle_names_the_country_package() -> None: + bundle = installed_household_bundle("us") + + assert isinstance(bundle, HouseholdBundleProvenance) + assert bundle.country_package_name == "policyengine-us" + + +# Installed country models: real calculations, compared with the Public API's +# result shape. + +US_FAMILY = { + "people": { + "you": {"age": {"2026": 35}, "employment_income": {"2026": 30_000}}, + "your first dependent": {"age": {"2026": 5}}, + }, + "tax_units": { + "your tax unit": { + "members": ["you", "your first dependent"], + "ctc": {"2026": None}, + } + }, + "spm_units": {"your spm unit": {"members": ["you", "your first dependent"]}}, + "families": {"your family": {"members": ["you", "your first dependent"]}}, + "marital_units": { + "your marital unit": {"members": ["you"]}, + "your first dependent's marital unit": {"members": ["your first dependent"]}, + }, + "households": { + "your household": { + "members": ["you", "your first dependent"], + "state_name": {"2026": "TX"}, + "county_fips": {"2026": "48201"}, + "household_net_income": {"2026": None}, + } + }, +} +CTC_REFORM = { + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01.2100-12-31": 5_000} +} + + +def test_us_baseline_and_reform_are_calculated_in_one_request() -> None: + runtime = RecordingRuntime() + request = _request( + "us", + US_FAMILY, + policies=[ + {"policy_key": "baseline", "parameters": {}}, + {"policy_key": "reform", "parameters": CTC_REFORM}, + ], + ) + + outcome = calculate_household( + request, + worker=_worker("us"), + dependencies=HouseholdCalculationDependencies(today=TODAY, runtime=runtime), + ) + + assert isinstance(outcome, HouseholdExecutionResult), outcome + assert outcome.calculation_id == CALCULATION_ID + assert outcome.output.bundle == installed_household_bundle("us") + baseline, reform = outcome.output.results + assert (baseline.policy_key, reform.policy_key) == ("baseline", "reform") + baseline_situation = baseline.household.to_situation() + reform_situation = reform.household.to_situation() + # Same structure as the input, with each requested value filled in. + assert list(baseline_situation) == list(US_FAMILY) + assert baseline_situation["people"] == US_FAMILY["people"] + baseline_ctc = baseline_situation["tax_units"]["your tax unit"]["ctc"]["2026"] + reform_ctc = reform_situation["tax_units"]["your tax unit"]["ctc"]["2026"] + assert isinstance(baseline_ctc, float) + assert isinstance(reform_ctc, float) + assert reform_ctc > baseline_ctc + # A US household inherits the certified SPM selection and carries its receipt. + assert outcome.output.spm_config is not None + assert outcome.output.spm_config.geography_kind == "county" + assert baseline.spm_provenance is not None + assert reform.spm_provenance is not None + assert baseline.spm_provenance.scenario == outcome.output.spm_config.scenario + assert [name for name, _ in runtime.spans] == [ + "v2_household_country_model_load", + "v2_household_input_normalization", + "v2_household_calculation", + "v2_household_calculation", + ] + assert runtime.spans[3][1] == {"country": "us", "policy_key": "reform"} + + +def test_us_axes_calculate_one_value_per_point() -> None: + household = { + **US_FAMILY, + "axes": [ + [ + { + "name": "employment_income", + "count": 5, + "min": 0, + "max": 80_000, + "period": "2026", + } + ] + ], + } + + outcome = calculate_household( + _request("us", household), + worker=_worker("us"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + assert isinstance(outcome, HouseholdExecutionResult), outcome + situation = outcome.output.results[0].household.to_situation() + incomes = situation["households"]["your household"]["household_net_income"] + assert len(incomes["2026"]) == 5 + assert all(isinstance(income, float) for income in incomes["2026"]) + assert incomes["2026"][0] < incomes["2026"][4] + # With axes, every value of a known variable is calculated at each point, + # and the people's provided values are the axis values. + assert situation["people"]["you"]["employment_income"] == { + "2026": [0.0, 20_000.0, 40_000.0, 60_000.0, 80_000.0] + } + assert situation["people"]["you"]["age"] == {"2026": [35.0] * 5} + assert situation["households"]["your household"]["state_name"] == { + "2026": ["TX"] * 5 + } + assert situation["axes"] == household["axes"] + + +def test_us_calculate_full_adds_and_calculates_every_missing_variable() -> None: + outcome = calculate_household( + _request("us", US_FAMILY, add_missing_variables=True), + worker=_worker("us"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + assert isinstance(outcome, HouseholdExecutionResult), outcome + situation = outcome.output.results[0].household.to_situation() + you = situation["people"]["you"] + tax_unit = situation["tax_units"]["your tax unit"] + model = load_country_household_model("us") + person_variables = { + variable.name + for variable in model.metadata.variables.values() + if variable.entity == "person" + and variable.definition_period in {"year", "month", "eternity"} + } + assert person_variables <= set(you) + assert isinstance(you["is_blind"]["2026"], bool) + assert isinstance(tax_unit["income_tax"]["2026"], float) + assert isinstance(tax_unit["eitc"]["2026"], float) + assert isinstance( + situation["households"]["your household"]["household_net_income"]["2026"], + float, + ) + + +UK_SINGLE = { + "people": {"you": {"age": {"2026": 40}, "employment_income": {"2026": 30_000}}}, + "benunits": {"your benefit unit": {"members": ["you"]}}, + "households": { + "your household": {"members": ["you"], "household_net_income": {"2026": None}} + }, +} + + +def test_uk_axes_calculate_one_value_per_point() -> None: + household = { + **UK_SINGLE, + "axes": [ + [ + { + "name": "employment_income", + "count": 3, + "min": 0, + "max": 60_000, + "period": "2026", + } + ] + ], + } + + outcome = calculate_household( + _request("uk", household), + worker=_worker("uk"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + assert isinstance(outcome, HouseholdExecutionResult), outcome + situation = outcome.output.results[0].household.to_situation() + incomes = situation["households"]["your household"]["household_net_income"] + assert len(incomes["2026"]) == 3 + assert incomes["2026"][0] < incomes["2026"][2] + assert situation["people"]["you"]["employment_income"] == { + "2026": [0.0, 30_000.0, 60_000.0] + } + assert outcome.output.results[0].spm_provenance is None + assert outcome.output.spm_config is None + + +def test_uk_calculate_full_adds_every_missing_yearly_variable() -> None: + outcome = calculate_household( + _request("uk", UK_SINGLE, add_missing_variables=True), + worker=_worker("uk"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + assert isinstance(outcome, HouseholdExecutionResult), outcome + you = outcome.output.results[0].household.to_situation()["people"]["you"] + assert isinstance(you["income_tax"]["2026"], float) + + +def test_uk_unrecognized_inputs_fail_with_public_api_errors() -> None: + household = { + **UK_SINGLE, + "people": {"you": {**UK_SINGLE["people"]["you"], "not_a_variable": 1}}, + } + request = _request( + "uk", + household, + policies=[ + {"policy_key": "baseline", "parameters": {}}, + {"policy_key": "reform", "parameters": {"gov.not.a.parameter": {}}}, + ], + ) + + outcome = calculate_household( + request, + worker=_worker("uk"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + # The first policy's request already fails on the household's own input, + # as the Public API's baseline request would. + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "invalid_inputs" + assert outcome.error.message == ( + "Unrecognized calculate input(s): Unrecognized household variable " + "`not_a_variable` at `household.people.you.not_a_variable`." + ) + assert outcome.error.policy_key is None + assert [error.path for error in outcome.error.input_errors] == [ + "household.people.you.not_a_variable" + ] + + +def test_uk_unrecognized_policy_parameters_name_their_policy() -> None: + request = _request( + "uk", + UK_SINGLE, + policies=[ + {"policy_key": "baseline", "parameters": {}}, + {"policy_key": "reform", "parameters": {"gov.not.a.parameter": {}}}, + ], + ) + + outcome = calculate_household( + request, + worker=_worker("uk"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "invalid_inputs" + assert outcome.error.policy_key == "reform" + assert [error.path for error in outcome.error.input_errors] == [ + "policy.gov.not.a.parameter" + ] + + +def test_uk_reform_values_given_as_text_are_converted_to_numbers() -> None: + # The Public API converts a reform value to the parameter's type, reading + # numeric text as a number first. + basic_rate = "gov.hmrc.income_tax.rates.uk[0].rate" + request = _request( + "uk", + UK_SINGLE, + policies=[ + {"policy_key": "baseline", "parameters": {}}, + { + "policy_key": "reform", + "parameters": {basic_rate: {"2026-01-01.2100-12-31": "0.3"}}, + }, + ], + ) + + outcome = calculate_household( + request, + worker=_worker("uk"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + assert isinstance(outcome, HouseholdExecutionResult), outcome + baseline, reform = ( + result.household.to_situation()["households"]["your household"][ + "household_net_income" + ]["2026"] + for result in outcome.output.results + ) + assert isinstance(baseline, float) + assert isinstance(reform, float) + assert reform < baseline + + +def test_uk_deprecated_inputs_warn_in_every_policy_result() -> None: + household = { + **UK_SINGLE, + "people": { + "you": { + **UK_SINGLE["people"]["you"], + "uc_has_limited_capability_for_work": {"2026": True}, + } + }, + } + request = _request( + "uk", + household, + policies=[ + {"policy_key": "baseline", "parameters": {}}, + {"policy_key": "other-baseline", "parameters": {}}, + ], + ) + + outcome = calculate_household( + request, + worker=_worker("uk"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + assert isinstance(outcome, HouseholdExecutionResult), outcome + for result in outcome.output.results: + assert len(result.warnings) == 1 + assert "uc_has_limited_capability_for_work" in result.warnings[0] + you = result.household.to_situation()["people"]["you"] + assert "uc_has_limited_capability_for_work" not in you + + +def test_uk_spm_settings_are_refused() -> None: + outcome = calculate_household( + _request("uk", UK_SINGLE, spm={}), + worker=_worker("uk"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "invalid_spm_selection" + assert outcome.error.spm_error is not None + assert outcome.error.spm_error.code == "SPM_SETTINGS_UNSUPPORTED" + + +def test_uk_member_who_is_not_a_person_is_an_invalid_household() -> None: + household = { + **UK_SINGLE, + "benunits": {"your benefit unit": {"members": ["you", "ghost"]}}, + } + + outcome = calculate_household( + _request("uk", household), + worker=_worker("uk"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "invalid_household" + assert outcome.error.message.startswith("Invalid household payload: ") + + +def test_country_models_are_loaded_once_per_process() -> None: + first = load_country_household_model("uk") + + assert load_country_household_model("uk") is first + assert first.metadata.entities["benunit"].plural == "benunits" + assert "gov.hmrc.income_tax.rates.uk[0].rate" in first.metadata.parameters From 389c6ccbe0eb24fc9f63315a89e53e4eb3d43690 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:28:23 +0400 Subject: [PATCH 07/19] Add calculate_household_us and calculate_household_uk to the v2 Modal app Each function takes a household request, calculates it with the executor's household calculation and returns the result or typed failure as JSON. It continues the caller's trace through the keyword-only observability context and measures the call as one v2 household operation. The functions run in their own image: the CA store, the locked packages and the source, without the bundle's datasets or the Stage 12 cache, so neither the image build nor the functions use a secret. Each reserves 2 CPUs and 4 GiB, may use up to 16 GiB for large households with axes, times out after 900 seconds, is not retried, and runs at most 20 containers. The other v2 images keep their layers and order. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 27 +++ .../src/modal/v2_app.py | 171 ++++++++++++++++-- .../tests/test_stage12_modal_app.py | 114 ++++++++++++ 3 files changed, 292 insertions(+), 20 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 0a2eddcbd..cdbf9fc8f 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -453,6 +453,33 @@ not satisfy the requested output plan is recalculated. Each region-group log records whether it loaded a complete file, recalculated an incomplete file, or computed because no file existed. +## v2 household functions + +The versioned v2 application also defines a household function for each +country, `calculate_household_us` and `calculate_household_uk`. Each call +calculates one household, in the Public API's situation shape and with +optional axes, under one to ten policies, so a baseline and its reform are +calculated in one call. It takes a `HouseholdExecutionRequest` and returns a +`HouseholdExecutionResult` or a typed `HouseholdExecutionFailure` from +`policyengine_simulation_contract.household_execution`, as JSON. The +calculation follows the Public API's household calculation, including its +input errors, warnings and SPM receipt. A function refuses a request for +another country, and returns a `bundle_mismatch` failure for a request that +names another bundle. Nothing calls these functions until the Simulation +Entrypoint has household routes. + +- Image: the CA store, the frozen `modal-simulation-image` dependency group + and the source packages, without the bundle's datasets or the Stage 12 + cache. Building it needs no secret. +- Secrets: none. A household calculation reads no private data, Cloud + Storage or database. Telemetry authenticates with the identity token Modal + gives every container (`MODAL_IDENTITY_TOKEN`) and the workload identity + settings baked into the image at deployment. +- Limits: 2 CPUs, 4,096 MiB of reserved and 16,384 MiB of maximum memory, a + 900-second timeout, no retries, and a Modal `max_containers` value of 20. +- Each container loads a country model the first time it calculates that + country and keeps it for later calls. + ## Temporary direct runner endpoint > **Temporary Stage 12 interface:** The authenticated routes in this section diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 92642d673..396181322 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -1,5 +1,10 @@ # pyright: reportAttributeAccessIssue=false -"""Separately named Stage 12 Modal coordinator and country workers. +"""Separately named v2 Modal application: report coordinator and country workers. + +Besides the Stage 12 report coordinator and simulation workers, it defines a +household function for each country, ``calculate_household_us`` and +``calculate_household_uk``, which calculates one household under one or more +policies in an image without datasets or secrets. This module defines additional resources. It does not rename or modify the existing ``policyengine-simulation-py{version}`` applications. @@ -12,6 +17,9 @@ from pathlib import Path from typing import TYPE_CHECKING +from policyengine_simulation_contract.household_execution import ( + household_callable_name, +) from policyengine_simulation_contract.stage12_manifest import v2_application_name from policyengine_simulation_executor import stage12_runtime from policyengine_simulation_executor.hf_access_validation import ( @@ -39,6 +47,7 @@ STAGE12_SEGMENT_STAGES, STAGE12_SHADOW_REPORT_STAGES, STAGE12_SIMULATION_STAGES, + V2_HOUSEHOLD_STAGES, Stage, ) from policyengine_simulation_observability.telemetry import ( @@ -138,6 +147,35 @@ def _country_bundle(country: CountryId) -> Stage12CountryBundle: ) +def _v2_package_image() -> modal.Image: + """Return the base of every v2 image: the CA store and the locked packages.""" + return ( + modal.Image.debian_slim(python_version="3.13") + # Upgrade Debian's CA store before anything is downloaded, and make + # grpcio verify servers against it (src/modal/ca_certificates.py). + .apt_install(OPERATING_SYSTEM_CA_PACKAGE) + .env(GRPC_CA_BUNDLE_ENV) + .uv_sync( + uv_project_dir=_UV_PROJECT_DIR, + frozen=True, + extra_options="--only-group modal-simulation-image", + ) + ) + + +def _add_v2_sources(image: modal.Image) -> modal.Image: + """Copy the source packages and static runtime files every v2 image runs.""" + image = image.add_local_python_source( + "src.modal", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_simulation_contract", + "policyengine_stage12_persistence", + copy=True, + ) + return add_static_runtime_files(image, uv_project_dir=_UV_PROJECT_DIR) + + def build_v2_image( countries: tuple[CountryId, ...], *, include_us_cache: bool = False ) -> modal.Image: @@ -152,16 +190,7 @@ def build_v2_image( for country in countries } image = ( - modal.Image.debian_slim(python_version="3.13") - # Upgrade Debian's CA store before anything is downloaded, and make - # grpcio verify servers against it (src/modal/ca_certificates.py). - .apt_install(OPERATING_SYSTEM_CA_PACKAGE) - .env(GRPC_CA_BUNDLE_ENV) - .uv_sync( - uv_project_dir=_UV_PROJECT_DIR, - frozen=True, - extra_options="--only-group modal-simulation-image", - ) + _v2_package_image() .run_commands( bundle_data_install_command( RESOLVED_BUNDLE.bundle.policyengine_version, @@ -199,20 +228,23 @@ def build_v2_image( # for each precomputed year), downloaded and hashed in turn. timeout=3600, ) - image = image.add_local_python_source( - "src.modal", - "policyengine_simulation_executor", - "policyengine_simulation_observability", - "policyengine_simulation_contract", - "policyengine_stage12_persistence", - copy=True, - ) - return add_static_runtime_files(image, uv_project_dir=_UV_PROJECT_DIR) + return _add_v2_sources(image) + + +def build_household_image() -> modal.Image: + """Build the household functions' image: packages and source only. + + A household calculation reads no dataset, cache, bucket or database, so + the image downloads none and its build needs no secret. The installed + PolicyEngine.py package names the bundle the functions calculate with. + """ + return _add_v2_sources(_v2_package_image().env(modal_image_environment())) us_worker_image = build_v2_image(("us",), include_us_cache=True) uk_worker_image = build_v2_image(("uk",)) coordinator_image = build_v2_image(("us", "uk")) +household_image = build_household_image() def _validate(country: CountryId) -> dict: @@ -481,3 +513,102 @@ def coordinate_report( return result finally: runtime.shutdown() + + +# Household function sizing. A call holds one country model and, for each +# policy, a simulation of the household at every axis point. Each container +# reserves 4 GiB and may use up to 16 GiB before Modal stops it, so a large +# household calculated with every variable at many axis points still fits. +# The timeout is three times the Public API's 300-second request timeout. +HOUSEHOLD_FUNCTION_CPU = 2.0 +HOUSEHOLD_FUNCTION_MEMORY_MIB = (4096, 16384) +HOUSEHOLD_FUNCTION_TIMEOUT_SECONDS = 900 +HOUSEHOLD_FUNCTION_MAX_CONTAINERS = 20 + + +def _run_household_calculation( + country: CountryId, + payload: dict, + observability_context: dict | None, +) -> dict: + """Calculate one household request for ``country`` and return its outcome. + + The outcome is a ``HouseholdExecutionResult`` or ``HouseholdExecutionFailure`` + as JSON. Raises for a malformed request or a request for another country. + """ + from policyengine_simulation_contract.household_execution import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time + HouseholdExecutionRequest, + HouseholdWorkerIdentity, + ) + from policyengine_simulation_executor import ( # noqa: PLC0415 # Modal function body: loaded in the worker container, not at deploy time + household_calculation, + ) + + runtime = init_process_observability( + service_name=f"policyengine-v2-household-{country}-worker", + service_role="v2_household_worker", + platform="modal", + environment=os.getenv("MODAL_ENVIRONMENT", "local"), + ) + try: + propagated = normalize_observability_context(observability_context) + result = None + with runtime.operation( + V2_HOUSEHOLD_STAGES.name(Stage.V2_HOUSEHOLD_EXECUTION), + attributes={"country": country}, + remote_context=propagated, + ): + outcome = household_calculation.calculate_household( + HouseholdExecutionRequest.model_validate(payload), + worker=HouseholdWorkerIdentity( + application_name=APP_NAME, + function_name=household_callable_name(country), + ), + required_country=country, + dependencies=household_calculation.HouseholdCalculationDependencies( + runtime=runtime + ), + ) + result = outcome.model_dump(mode="json") + if result is None: + msg = f"the {country} household calculation returned no result" + raise RuntimeError(msg) + return result + finally: + runtime.shutdown() + + +@app.function( + image=household_image, + cpu=HOUSEHOLD_FUNCTION_CPU, + memory=HOUSEHOLD_FUNCTION_MEMORY_MIB, + timeout=HOUSEHOLD_FUNCTION_TIMEOUT_SECONDS, + retries=0, + max_containers=HOUSEHOLD_FUNCTION_MAX_CONTAINERS, + secrets=[], +) +def calculate_household_us( + payload: dict, + *, + observability_context: dict | None = None, +) -> dict: + """Calculate one US household under each requested policy.""" + return _run_household_calculation("us", payload, observability_context) + + +@app.function( + image=household_image, + cpu=HOUSEHOLD_FUNCTION_CPU, + memory=HOUSEHOLD_FUNCTION_MEMORY_MIB, + timeout=HOUSEHOLD_FUNCTION_TIMEOUT_SECONDS, + retries=0, + max_containers=HOUSEHOLD_FUNCTION_MAX_CONTAINERS, + secrets=[], +) +def calculate_household_uk( + payload: dict, + *, + observability_context: dict | None = None, +) -> dict: + """Calculate one UK household under each requested policy.""" + return _run_household_calculation("uk", payload, observability_context) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index 00ae66afc..3c757f0ae 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -247,6 +247,8 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( "run_single_simulation_us", "run_single_simulation_uk", "coordinate_report", + "calculate_household_us", + "calculate_household_uk", } assert functions["run_single_simulation_us"]["image"] is module.us_worker_image assert functions["run_single_simulation_uk"]["image"] is module.uk_worker_image @@ -319,6 +321,118 @@ def fail(*_, **__): assert runtime.shutdown_calls == 1 +def test_household_functions_run_without_datasets_or_secrets(monkeypatch) -> None: + module = _load(monkeypatch) + functions = dict(module.app.function_calls) + image_calls = module.household_image.calls + + for country in ("us", "uk"): + declaration = functions[f"calculate_household_{country}"] + assert declaration["image"] is module.household_image + assert declaration["secrets"] == [] + assert declaration["cpu"] == 2.0 + # 4 GiB reserved, up to 16 GiB, for large households with axes. + assert declaration["memory"] == (4096, 16384) + assert declaration["timeout"] == 900 + assert declaration["retries"] == 0 + assert declaration["max_containers"] == 20 + # The image installs the locked packages and the source, but downloads no + # dataset or cache, so building it needs no secret either. + assert [call[0] for call in image_calls[:4]] == [ + "debian_slim", + "apt_install", + "env", + "uv_sync", + ] + assert not any(call[0] in {"run_commands", "run_function"} for call in image_calls) + assert not any( + isinstance(argument, dict) and "secrets" in argument + for call in image_calls + for argument in call[1:] + ) + environment = { + name: value + for call in image_calls + if call[0] == "env" + for name, value in call[1].items() + } + assert "POLICYENGINE_DATA_FOLDER" not in environment + assert "STAGE12_CACHE_MANIFEST_DIGEST" not in environment + assert "OTEL_TRACES_EXPORTER" in environment + local_source = next( + call for call in image_calls if call[0] == "add_local_python_source" + ) + assert "policyengine_simulation_executor" in local_source[1] + + +def _household_payload(country: str) -> dict: + return { + "calculation_id": "00000000-0000-4000-8000-0000000000c1", + "country": country, + "household": {"people": {"you": {"age": {"2026": 40}}}}, + "policies": [{"policy_key": "baseline", "parameters": {}}], + "bundle": { + "policyengine_version": "6.7.5", + "core_package_version": "3.0.0", + "country_package_name": f"policyengine-{country}", + "country_package_version": "1.0.0", + "bundle_manifest_sha256": "a" * 64, + }, + } + + +@pytest.mark.parametrize("country", ["us", "uk"]) +def test_household_functions_calculate_their_own_country_with_telemetry( + monkeypatch, country +) -> None: + from policyengine_simulation_executor import household_calculation + + module = _load(monkeypatch) + runtime = RecordingRuntime() + monkeypatch.setattr(module, "init_process_observability", lambda **_: runtime) + calls: list[dict] = [] + + class Outcome: + def model_dump(self, *, mode): + assert mode == "json" + return {"status": "succeeded"} + + def calculate(request, **keywords): + calls.append({"request": request, **keywords}) + return Outcome() + + monkeypatch.setattr(household_calculation, "calculate_household", calculate) + + result = getattr(module, f"calculate_household_{country}")( + _household_payload(country), + observability_context=_observability_context(), + ) + + assert result == {"status": "succeeded"} + (call,) = calls + assert call["request"].country == country + assert call["required_country"] == country + assert call["worker"].application_name == module.APP_NAME + assert call["worker"].function_name == f"calculate_household_{country}" + assert call["dependencies"].runtime is runtime + (operation,) = runtime.operations + assert operation["name"] == "v2_household_execution" + assert operation["attributes"] == {"country": country} + assert operation["remote_context"]["observability_id"] == OBSERVABILITY_ID + assert runtime.shutdown_calls == 1 + + +def test_household_functions_close_runtime_after_an_error(monkeypatch) -> None: + module = _load(monkeypatch) + runtime = RecordingRuntime() + monkeypatch.setattr(module, "init_process_observability", lambda **_: runtime) + + with pytest.raises(ValueError, match="validation error"): + module.calculate_household_us({"country": "us"}) + + assert runtime.shutdown_calls == 1 + + def test_external_package_values_are_assertions_only(monkeypatch) -> None: install_fake_modal(monkeypatch) monkeypatch.setenv("STAGE12_EXPECT_US_VERSION", "not-the-bundled-version") From 60d778edf0a3b99f0340714d321470cddaaebc2c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:29:45 +0400 Subject: [PATCH 08/19] Validate the household functions at release and record them in the manifest After the country validation functions, release validation calls each deployed household function with a one-adult, current-law household and requires a successful calculation with the deployed bundle, reported by that function, with a number for the household's net income. The release record keeps each household invocation ID, and the manifest publisher records a country's household_callable only when the record shows its function was validated. A household function that fails validation stops the release, as a failed country validation already does. The validation household and its check live in the executor package, so the pre-merge image check can calculate the same household inside the household image. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 10 ++ .../modal/utils/update_v2_version_manifest.py | 14 ++- .../modal/utils/validate_v2_worker_release.py | 95 ++++++++++++-- .../household_validation.py | 93 ++++++++++++++ .../tests/test_household_validation.py | 119 ++++++++++++++++++ .../tests/test_update_v2_version_manifest.py | 13 ++ .../tests/test_validate_v2_worker_release.py | 86 +++++++++++++ 7 files changed, 421 insertions(+), 9 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_validation.py create mode 100644 projects/policyengine-simulation-executor/tests/test_household_validation.py diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index cdbf9fc8f..23b84522b 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -480,6 +480,16 @@ Entrypoint has household routes. - Each container loads a country model the first time it calculates that country and keeps it for later calls. +Deployment validation calls each household function with a one-adult, +current-law household and requires a successful calculation with the deployed +bundle, reported by that function, with a number for the household's net +income. The validation record keeps each invocation ID, and the manifest +publisher records a country's `household_callable` only when the record shows +its function was validated. A failed household validation fails v2 worker +validation like any other validation failure: the manifest is not published, +and, as described under "Release and deployment constraints", the deployment +does not promote the new Simulation Entrypoint revision. + ## Temporary direct runner endpoint > **Temporary Stage 12 interface:** The authenticated routes in this section 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 00dbab4ce..70e6dfd49 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 @@ -7,6 +7,9 @@ from pathlib import Path import modal +from policyengine_simulation_contract.household_execution import ( + household_callable_name, +) from policyengine_simulation_contract.stage12_manifest import ( V1_ROUTING_STATE_NAME, V2_VERSION_MANIFEST_NAME, @@ -25,8 +28,10 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: """Build the installed bundle's v2 worker entry from a validation record. Every bundle country has a single-simulation worker, so each is recorded - as economy-capable. Raises ValueError if the record is invalid or the - entry carries a credential-like key. + as economy-capable. A country's household function is recorded when the + record shows that it calculated the validation household. 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) @@ -38,6 +43,11 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: country=country.country, economy_capable=True, single_simulation_callable=f"run_single_simulation_{country.country}", + household_callable=( + household_callable_name(country.country) + if country.country in validation.household_validation_invocation_ids + else None + ), ) for country in resolved.bundle.countries ), diff --git a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py index 198b019ef..d5fee48f6 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py @@ -1,4 +1,10 @@ -"""Invoke and verify every deployed Stage 12 country validation function.""" +"""Invoke and verify every deployed v2 validation and household function. + +Each country's validation function checks its simulation worker image, and +each country's household function calculates a one-adult, current-law +household. The release record names every invocation, so the manifest names +only functions that succeeded in the deployed application. +""" from __future__ import annotations @@ -10,30 +16,40 @@ from uuid import uuid4 import modal +from policyengine_simulation_contract.household_execution import ( + HouseholdWorkerIdentity, + household_callable_name, +) from policyengine_simulation_contract.stage12_manifest import ( ManifestText, Stage12CacheRelease, V2WorkerValidation, v2_application_name, ) +from policyengine_simulation_executor.household_validation import ( + check_household_validation, + household_validation_request, +) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle if TYPE_CHECKING: from policyengine_simulation_contract.stage12_bundle import CountryId -def _spawn_validation( +def _invoke_deployed_function( *, application_name: str, function_name: str, environment: str, + arguments: tuple[object, ...] = (), ) -> tuple[Any, str]: + """Call a deployed function, wait for its result, and return it with the call ID.""" function = modal.Function.from_name( application_name, function_name, environment_name=environment, ) - call = function.spawn() + call = function.spawn(*arguments) result = call.get(timeout=1_000) invocation_id = getattr(call, "object_id", None) if not isinstance(invocation_id, str) or not invocation_id: @@ -42,14 +58,73 @@ def _spawn_validation( return result, invocation_id +def _spawn_validation( + *, + application_name: str, + function_name: str, + environment: str, +) -> tuple[Any, str]: + return _invoke_deployed_function( + application_name=application_name, + function_name=function_name, + environment=environment, + ) + + +def _call_household_function( + *, + application_name: str, + function_name: str, + environment: str, + payload: dict[str, Any], +) -> tuple[Any, str]: + return _invoke_deployed_function( + application_name=application_name, + function_name=function_name, + environment=environment, + arguments=(payload,), + ) + + +def validate_household_functions( + *, application_name: str, environment: str, year: int +) -> dict[CountryId, ManifestText]: + """Calculate a validation household with each country's household function. + + Returns each country's invocation ID. Raises RuntimeError when a function + does not calculate the household as ``check_household_validation`` requires. + """ + invocation_ids: dict[CountryId, ManifestText] = {} + for country_bundle in load_stage12_bundle().bundle.countries: + country = country_bundle.country + function_name = household_callable_name(country) + request = household_validation_request(country, year=year) + outcome, invocation_id = _call_household_function( + application_name=application_name, + function_name=function_name, + environment=environment, + payload=request.model_dump(mode="json"), + ) + check_household_validation( + outcome, + request=request, + worker=HouseholdWorkerIdentity( + application_name=application_name, function_name=function_name + ), + ) + invocation_ids[country] = invocation_id + return invocation_ids + + def validate_release( *, environment: str, expected_cache_manifest_sha256: str ) -> V2WorkerValidation: - """Invoke each deployed country validation function and return the release record. + """Invoke each deployed validation and household function; return the record. - Raises RuntimeError when a function does not report success, returns no - invocation id, reports bundle values that differ from the local bundle, or - reports a US cache other than ``expected_cache_manifest_sha256``. + Raises RuntimeError when a validation function does not report success, + returns no invocation id, reports bundle values that differ from the local + bundle, or reports a US cache other than ``expected_cache_manifest_sha256``, + or when a household function does not calculate its validation household. """ resolved = load_stage12_bundle() application_name = v2_application_name(resolved.bundle.policyengine_version) @@ -93,6 +168,11 @@ def validate_release( msg = f"{function_name} unexpectedly reported a cache" raise RuntimeError(msg) invocation_ids[country_bundle.country] = invocation_id + household_invocation_ids = validate_household_functions( + application_name=application_name, + environment=environment, + year=datetime.now(UTC).year, + ) return V2WorkerValidation( validated=True, application_name=application_name, @@ -100,6 +180,7 @@ def validate_release( validated_at=datetime.now(UTC), validation_invocation_id=f"stage12-validation-{uuid4()}", country_validation_invocation_ids=invocation_ids, + household_validation_invocation_ids=household_invocation_ids, cache=cache_release, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_validation.py new file mode 100644 index 000000000..76d842744 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_validation.py @@ -0,0 +1,93 @@ +"""The household that checks a v2 household function calculates. + +Release validation calls each deployed household function with this +household, and the pre-merge image check calculates it inside the household +image. Both use ``check_household_validation`` to decide whether it worked. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from uuid import uuid4 + +from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionRequest, + HouseholdExecutionResult, + parse_household_execution_outcome, +) + +from policyengine_simulation_executor.household_calculation import ( + installed_household_bundle, +) + +if TYPE_CHECKING: + from policyengine_simulation_contract.household_execution import ( + HouseholdWorkerIdentity, + ) + from policyengine_simulation_contract.stage12_bundle import CountryId + +# The group entity, besides the household, that holds the adult in each +# country's validation household. +_ADULT_GROUP_PLURALS: dict[str, str] = {"us": "spm_units", "uk": "benunits"} + + +def household_validation_request( + country: CountryId, *, year: int +) -> HouseholdExecutionRequest: + """Return a one-adult, current-law household request for ``country``. + + The request names the installed bundle and asks for the household's net + income in ``year``. + """ + period = str(year) + household = { + "people": {"you": {"age": {period: 40}, "employment_income": {period: 30_000}}}, + _ADULT_GROUP_PLURALS[country]: {"your unit": {"members": ["you"]}}, + "households": { + "your household": { + "members": ["you"], + "household_net_income": {period: None}, + } + }, + } + return HouseholdExecutionRequest.model_validate( + { + "calculation_id": uuid4(), + "country": country, + "household": household, + "policies": [{"policy_key": "baseline", "parameters": {}}], + "bundle": installed_household_bundle(country), + } + ) + + +def check_household_validation( + outcome: object, + *, + request: HouseholdExecutionRequest, + worker: HouseholdWorkerIdentity, +) -> float: + """Return the household net income ``outcome`` calculated for ``request``. + + Raises RuntimeError unless the outcome is a success with the request's + bundle, reported by ``worker``, with a number for the household's net + income. + """ + name = worker.function_name + parsed = parse_household_execution_outcome(outcome) + if not isinstance(parsed, HouseholdExecutionResult): + msg = f"{name} did not calculate the validation household" + raise RuntimeError(msg) # noqa: TRY004 # a failed calculation fails the check; it is not a type error + if parsed.output.bundle != request.bundle: + msg = f"{name} calculated with another bundle" + raise RuntimeError(msg) + if parsed.worker != worker: + msg = f"{name} reported another function" + raise RuntimeError(msg) + situation = parsed.output.results[0].household.to_situation() + net_income = situation["households"]["your household"]["household_net_income"] + values = list(net_income.values()) if isinstance(net_income, dict) else [] + if len(values) != 1 or not isinstance(values[0], float): + msg = f"{name} returned no household net income" + raise RuntimeError(msg) + return values[0] diff --git a/projects/policyengine-simulation-executor/tests/test_household_validation.py b/projects/policyengine-simulation-executor/tests/test_household_validation.py new file mode 100644 index 000000000..a82abb2df --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_household_validation.py @@ -0,0 +1,119 @@ +"""Tests for the household that checks a v2 household function calculates.""" + +from __future__ import annotations + +from copy import deepcopy + +import pytest +from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionRequest, + HouseholdWorkerIdentity, +) +from policyengine_simulation_executor.household_calculation import ( + installed_household_bundle, +) +from policyengine_simulation_executor.household_validation import ( + check_household_validation, + household_validation_request, +) + +WORKER = HouseholdWorkerIdentity( + application_name="policyengine-simulation-v2-py-test", + function_name="calculate_household_uk", +) + + +def _succeeded(request: HouseholdExecutionRequest) -> dict: + situation = request.household.to_situation() + situation["households"]["your household"]["household_net_income"] = { + "2026": 25_000.0 + } + return { + "status": "succeeded", + "schema_version": 1, + "calculation_id": str(request.calculation_id), + "output": { + "country": request.country, + "bundle": request.bundle.model_dump(mode="json"), + "results": [{"policy_key": "baseline", "household": situation}], + }, + "worker": WORKER.model_dump(mode="json"), + } + + +@pytest.mark.parametrize( + ("country", "group_plural"), [("us", "spm_units"), ("uk", "benunits")] +) +def test_the_validation_household_is_one_adult_under_current_law( + country, group_plural +) -> None: + request = household_validation_request(country, year=2026) + + situation = request.household.to_situation() + assert request.country == country + assert request.bundle == installed_household_bundle(country) + assert [policy.parameters for policy in request.policies] == [{}] + assert list(situation["people"]) == ["you"] + assert situation[group_plural] == {"your unit": {"members": ["you"]}} + assert situation["households"]["your household"]["household_net_income"] == { + "2026": None + } + + +def test_a_calculated_net_income_passes_the_check() -> None: + request = household_validation_request("uk", year=2026) + + net_income = check_household_validation( + _succeeded(request), request=request, worker=WORKER + ) + + assert net_income == 25_000.0 + + +def _failed(outcome: dict) -> dict: + return { + "status": "failed", + "schema_version": 1, + "calculation_id": outcome["calculation_id"], + "error": {"code": "calculation_failed", "message": "model exploded"}, + "worker": outcome["worker"], + } + + +def _with_another_bundle(outcome: dict) -> dict: + changed = deepcopy(outcome) + changed["output"]["bundle"]["policyengine_version"] = "0.0.1" + return changed + + +def _from_another_function(outcome: dict) -> dict: + changed = deepcopy(outcome) + changed["worker"]["function_name"] = "calculate_household_us" + return changed + + +def _without_net_income(outcome: dict) -> dict: + changed = deepcopy(outcome) + household = changed["output"]["results"][0]["household"] + household["households"]["your household"]["household_net_income"] = {"2026": None} + return changed + + +@pytest.mark.parametrize( + ("change_outcome", "message"), + [ + (_failed, "did not calculate the validation household"), + (_with_another_bundle, "calculated with another bundle"), + (_from_another_function, "reported another function"), + (_without_net_income, "returned no household net income"), + ], +) +def test_an_outcome_that_is_not_the_calculated_household_fails_the_check( + change_outcome, message +) -> None: + request = household_validation_request("uk", year=2026) + + with pytest.raises(RuntimeError, match=f"calculate_household_uk {message}"): + check_household_validation( + change_outcome(_succeeded(request)), request=request, worker=WORKER + ) 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 cee40de8d..f37e3697b 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 @@ -55,6 +55,19 @@ def test_publisher_records_every_bundle_country_as_economy_capable() -> None: ] +def test_publisher_records_each_validated_household_function() -> None: + payload = _validation().model_dump(mode="json") + payload["household_validation_invocation_ids"] = {"uk": "household-uk-1"} + + worker = publisher.build_worker(payload) + + assert { + country.country: country.household_callable for country in worker.countries + } == {"us": None, "uk": "calculate_household_uk"} + assert worker.household_worker("uk") is not None + assert worker.household_worker("us") is None + + def test_publisher_rejects_validation_for_another_application() -> None: payload = _validation().model_dump(mode="json") payload["application_name"] = "policyengine-simulation-v2-py-another" diff --git a/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py b/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py index c5b555c9d..09e156f0a 100644 --- a/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py +++ b/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py @@ -3,9 +3,54 @@ from __future__ import annotations import pytest +from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionRequest, +) from src.modal.utils import validate_v2_worker_release as validation +def _household_outcome( + request: HouseholdExecutionRequest, *, application_name: str, function_name: str +) -> dict: + """Return the success a household function reports for ``request``.""" + situation = request.household.to_situation() + for period in situation["households"]["your household"]["household_net_income"]: + situation["households"]["your household"]["household_net_income"][period] = ( + 25_000.0 + ) + return { + "status": "succeeded", + "schema_version": 1, + "calculation_id": str(request.calculation_id), + "output": { + "country": request.country, + "bundle": request.bundle.model_dump(mode="json"), + "results": [{"policy_key": "baseline", "household": situation}], + }, + "worker": { + "application_name": application_name, + "function_name": function_name, + }, + } + + +class HouseholdFunctions: + """Stands in for the deployed household functions; records each call.""" + + def __init__(self, change_outcome=None) -> None: + self.calls: list[tuple[str, HouseholdExecutionRequest]] = [] + self.change_outcome = change_outcome or (lambda outcome: outcome) + + def __call__(self, *, application_name, function_name, environment, payload): + assert environment == "staging" + request = HouseholdExecutionRequest.model_validate(payload) + self.calls.append((function_name, request)) + outcome = _household_outcome( + request, application_name=application_name, function_name=function_name + ) + return self.change_outcome(outcome), f"household-{request.country}" + + def test_all_country_results_must_match_the_deployment_bundle(monkeypatch) -> None: invocation_ids: list[str] = [] @@ -45,6 +90,8 @@ def fake_spawn_validation(*, application_name, function_name, environment): ) monkeypatch.setattr(validation, "_spawn_validation", fake_spawn_validation) + household_functions = HouseholdFunctions() + monkeypatch.setattr(validation, "_call_household_function", household_functions) result = validation.validate_release( environment="staging", @@ -59,6 +106,16 @@ def fake_spawn_validation(*, application_name, function_name, environment): assert invocation_ids == ["call-us", "call-uk"] assert result.cache is not None assert result.cache.manifest_sha256 == "c" * 64 + # Each household function calculated a current-law household with the + # deployment's bundle, and the record names each invocation. + assert [name for name, _ in household_functions.calls] == [ + "calculate_household_us", + "calculate_household_uk", + ] + assert result.household_validation_invocation_ids == { + "us": "household-us", + "uk": "household-uk", + } def test_country_result_mismatch_stops_publication(monkeypatch) -> None: @@ -80,3 +137,32 @@ def test_country_result_mismatch_stops_publication(monkeypatch) -> None: environment="staging", expected_cache_manifest_sha256="c" * 64, ) + + +def _fail(outcome: dict) -> dict: + return { + "status": "failed", + "schema_version": 1, + "calculation_id": outcome["calculation_id"], + "error": {"code": "calculation_failed", "message": "model exploded"}, + "worker": outcome["worker"], + } + + +def test_a_household_function_that_does_not_calculate_stops_publication( + monkeypatch, +) -> None: + household_functions = HouseholdFunctions(_fail) + monkeypatch.setattr(validation, "_call_household_function", household_functions) + + with pytest.raises( + RuntimeError, + match="calculate_household_us did not calculate the validation household", + ): + validation.validate_household_functions( + application_name="policyengine-simulation-v2-py-test", + environment="staging", + year=2026, + ) + + assert [name for name, _ in household_functions.calls] == ["calculate_household_us"] From 00b3409d294691a6ca268ad221570f3a548e2e22 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 17:33:36 +0400 Subject: [PATCH 09/19] Build every v2 image in the pre-merge image check The image check now also runs a v2 image smoke. It builds the report coordinator, UK worker and household images the v2 application deploys, and the US worker image without its last layer, the Stage 12 cache, which only a deployment can fetch: the deployment's precompute writes the cache, and the image build reads it with the deployment's Cloud Storage credential. That layer adds data files, not Python packages. In each report image the smoke imports the v2 application module, as Modal does before it runs a function, and the modules the coordinator, workers and validation functions import while handling a request. In the household image it imports the household entry points and calculates the release validation household for each country, which shows the image calculates without datasets. The smoke functions use no secret. Modal builds the smoke's images together. The job's timeout rises from 45 to 90 minutes, because after a relock both the executor smoke and the v2 smoke download certified datasets. Co-Authored-By: Claude Opus 5.5 --- .github/scripts/modal-image-smoke.sh | 12 +- .github/workflows/pr-image-smoke.yml | 7 +- .../README.md | 8 +- .../src/modal/v2_image_smoke.py | 152 ++++++++++++++++++ .../integration/test_image_smoke_modal.py | 7 + .../tests/test_modal_scripts.py | 10 +- .../tests/test_v2_image_smoke.py | 129 +++++++++++++++ 7 files changed, 317 insertions(+), 8 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/modal/v2_image_smoke.py create mode 100644 projects/policyengine-simulation-executor/tests/test_v2_image_smoke.py diff --git a/.github/scripts/modal-image-smoke.sh b/.github/scripts/modal-image-smoke.sh index dc9c0062f..b8aede54e 100755 --- a/.github/scripts/modal-image-smoke.sh +++ b/.github/scripts/modal-image-smoke.sh @@ -3,8 +3,10 @@ # Usage: ./modal-image-smoke.sh # # Each smoke imports the true app entrypoints inside the real (or, for -# the executor, prefix-identical) Modal image — catching in-image -# dependency breakage (issue #602's class) before merge. +# the executor and the v2 US worker, prefix-identical) Modal image — +# catching in-image dependency breakage (issue #602's class) before merge. +# The v2 image smoke builds every image of the v2 application: the report +# coordinator, the country simulation workers and the household functions. set -euo pipefail @@ -47,4 +49,10 @@ echo "=== Stage 12 static runtime file smoke (env: $MODAL_ENV) ===" src/modal/v2_static_runtime_files_smoke.py ) +echo "=== v2 image smoke: every v2 image and its entry points (env: $MODAL_ENV) ===" +( + cd "$REPO_ROOT/projects/policyengine-simulation-executor" + uv run --frozen modal run --env="$MODAL_ENV" src/modal/v2_image_smoke.py +) + echo "=== Image smokes passed ===" diff --git a/.github/workflows/pr-image-smoke.yml b/.github/workflows/pr-image-smoke.yml index 730b517d5..fd00793b4 100644 --- a/.github/workflows/pr-image-smoke.yml +++ b/.github/workflows/pr-image-smoke.yml @@ -33,9 +33,10 @@ jobs: # Same-repository pull requests only, and in a stack only the lowest # unmerged one (see docs/engineering/skills/github-prs.md). if: ${{ github.event.pull_request.head.repo.full_name == github.repository && (!github.event.pull_request.stack || github.event.pull_request.stack.base.ref == github.event.pull_request.base.ref) }} - # Warm cache: ~2 min. After a relock the executor smoke pays the - # certified dataset download layer (~15-20 min). - timeout-minutes: 45 + # Warm cache: a few minutes. After a relock the executor smoke and the + # v2 image smoke each pay their certified dataset download layers + # (~15-20 min each; Modal builds the v2 smoke's images together). + timeout-minutes: 90 steps: - uses: actions/checkout@v6 diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 37f5e91cb..9f6a47558 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -28,8 +28,12 @@ in the project dependencies and `modal-simulation-image`, then run `uv lock`. The automated updater changes only those wrapper requirements; it reads the wrapper's release manifest to verify and report the selected component and data versions. Other image dependencies belong in `modal-simulation-image`. PRs -touching image inputs run an in-image import smoke (`src/modal/smoke_app.py` via -`.github/workflows/pr-image-smoke.yml`). Note that any change to the +touching image inputs run in-image import smokes through +`.github/workflows/pr-image-smoke.yml`: `src/modal/smoke_app.py` for the v1 +worker image, and `src/modal/v2_image_smoke.py` for every v2 image (the report +coordinator, the country simulation workers, the US one without its Stage 12 +cache layer, and the household functions, where it also calculates a +household). Note that any change to the group or lock invalidates the image layer cache, including the artifact fetch layer below. diff --git a/projects/policyengine-simulation-executor/src/modal/v2_image_smoke.py b/projects/policyengine-simulation-executor/src/modal/v2_image_smoke.py new file mode 100644 index 000000000..89ef78504 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/v2_image_smoke.py @@ -0,0 +1,152 @@ +# pyright: reportAttributeAccessIssue=false, reportFunctionMemberAccess=false +"""Pre-merge image smoke: build every v2 Modal image and import its entry points. + +The coordinator, UK worker and household images are the ones the v2 +application deploys. The US worker image is built without its last layer, the +Stage 12 cache: only a deployment can add that layer, because the deployment's +precompute writes the cache and the image build reads it with the deployment's +Cloud Storage credential. The layer adds data files, not Python packages. + +In each image the smoke imports the v2 application module, as Modal does +before it runs any of the application's functions, and the modules those +functions import while handling a request. In the household image it also +calculates the release validation household for each country, which shows the +image calculates households without datasets. + +Usage: + uv run modal run --env=staging src/modal/v2_image_smoke.py +""" + +from src.modal.v2_app import ( + build_v2_image, + coordinator_image, + household_image, + uk_worker_image, +) + +import modal + +SMOKE_APP_NAME = "policyengine-simulation-v2-image-smoke" +app = modal.App(SMOKE_APP_NAME) + +us_worker_image_without_cache = build_v2_image(("us",)) + +# Modules the report coordinator, simulation workers and validation functions +# import when they start or while they handle a request. +REPORT_ENTRY_POINT_MODULES = ( + "src.modal.v2_app", + "policyengine_simulation_executor.stage12_runtime", + "policyengine_simulation_executor.stage12_worker_validation", + "policyengine_simulation_executor.simulation_runtime", + "policyengine_simulation_executor.spm", + "policyengine_simulation_executor.uk_local_authority_metadata", + "microdf", + "policyengine.tax_benefit_models.us.datasets", + "policyengine.tax_benefit_models.uk.datasets", +) +# Modules the household functions import when they start or calculate. +HOUSEHOLD_ENTRY_POINT_MODULES = ( + "src.modal.v2_app", + "policyengine_simulation_executor.household_calculation", + "policyengine_simulation_executor.household_validation", +) + + +def _import_modules(names: tuple[str, ...]) -> list[str]: + """Import each named module, raising on the first that fails.""" + import importlib # noqa: PLC0415 # smoke app: standard-library import kept with the in-image check that uses it + + for name in names: + importlib.import_module(name) + return list(names) + + +def _report_image_smoke() -> dict[str, object]: + """Import the report entry points; return them and the installed bundle version.""" + from policyengine_simulation_executor.stage12_bundle import ( # noqa: PLC0415 # smoke check: imported inside the Modal image, which is what this verifies + load_stage12_bundle, + ) + + return { + "modules": _import_modules(REPORT_ENTRY_POINT_MODULES), + "policyengine_version": load_stage12_bundle().bundle.policyengine_version, + } + + +@app.function(image=coordinator_image, timeout=600, memory=8192) +def smoke_v2_coordinator_image() -> dict[str, object]: + """Import the v2 entry points inside the report coordinator's image.""" + return _report_image_smoke() + + +@app.function(image=uk_worker_image, timeout=600, memory=8192) +def smoke_v2_uk_worker_image() -> dict[str, object]: + """Import the v2 entry points inside the UK simulation worker's image.""" + return _report_image_smoke() + + +@app.function(image=us_worker_image_without_cache, timeout=600, memory=8192) +def smoke_v2_us_worker_image() -> dict[str, object]: + """Import the v2 entry points inside the US simulation worker's image. + + The image has every layer of the deployed image except the Stage 12 cache. + """ + return _report_image_smoke() + + +@app.function(image=household_image, timeout=900, memory=(4096, 16384)) +def smoke_v2_household_image() -> dict[str, object]: + """Import the household entry points and calculate each validation household. + + Raises when a calculation fails or does not return a household net income. + """ + from datetime import ( # noqa: PLC0415 # smoke app: standard-library import kept with the in-image check that uses it + UTC, + datetime, + ) + + from policyengine_simulation_contract.household_execution import ( # noqa: PLC0415 # smoke check: imported inside the Modal image, which is what this verifies + HouseholdWorkerIdentity, + household_callable_name, + ) + from policyengine_simulation_executor import ( # noqa: PLC0415 # smoke check: imported inside the Modal image, which is what this verifies + household_calculation, + ) + from policyengine_simulation_executor.household_validation import ( # noqa: PLC0415 # smoke check: imported inside the Modal image, which is what this verifies + check_household_validation, + household_validation_request, + ) + from policyengine_simulation_executor.stage12_bundle import ( # noqa: PLC0415 # smoke check: imported inside the Modal image, which is what this verifies + load_stage12_bundle, + ) + + modules = _import_modules(HOUSEHOLD_ENTRY_POINT_MODULES) + year = datetime.now(UTC).year + net_incomes: dict[str, float] = {} + for country_bundle in load_stage12_bundle().bundle.countries: + country = country_bundle.country + request = household_validation_request(country, year=year) + worker = HouseholdWorkerIdentity( + application_name=SMOKE_APP_NAME, + function_name=household_callable_name(country), + ) + outcome = household_calculation.calculate_household( + request, worker=worker, required_country=country + ) + net_incomes[country] = check_household_validation( + outcome.model_dump(mode="json"), request=request, worker=worker + ) + return {"modules": modules, "household_net_income": net_incomes} + + +@app.local_entrypoint() +def main() -> None: + """Run every v2 image smoke remotely and print their reports.""" + reports = { + "coordinator": smoke_v2_coordinator_image.remote(), + "uk_worker": smoke_v2_uk_worker_image.remote(), + "us_worker": smoke_v2_us_worker_image.remote(), + "household": smoke_v2_household_image.remote(), + } + print(reports) # noqa: T201 # Modal local entrypoint: runs locally, and stdout is the smoke's report + print("v2 image smoke OK") # noqa: T201 # Modal local entrypoint: runs locally, and stdout is the smoke's report diff --git a/projects/policyengine-simulation-executor/tests/integration/test_image_smoke_modal.py b/projects/policyengine-simulation-executor/tests/integration/test_image_smoke_modal.py index 20e240bc1..560e5334e 100644 --- a/projects/policyengine-simulation-executor/tests/integration/test_image_smoke_modal.py +++ b/projects/policyengine-simulation-executor/tests/integration/test_image_smoke_modal.py @@ -38,3 +38,10 @@ def test_stage12_static_runtime_file_image_smoke(): result = _run_modal_smoke("src/modal/v2_static_runtime_files_smoke.py") assert result.returncode == 0, result.stdout + result.stderr assert "Stage 12 static runtime file smoke OK" in result.stdout + + +@pytest.mark.integration +def test_v2_image_smoke(): + result = _run_modal_smoke("src/modal/v2_image_smoke.py") + assert result.returncode == 0, result.stdout + result.stderr + assert "v2 image smoke OK" in result.stdout diff --git a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py index d427800d9..b29699da8 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py @@ -182,7 +182,15 @@ def test_syncs_expected_hf_key_before_image_smokes(self, tmp_path): "run --frozen modal secret create pe-uk-private-hf-read-token " f"HUGGING_FACE_TOKEN={FAKE_HF_TOKEN} --env=staging --force" ) - assert len(calls) == 4 + assert calls[1:] == [ + "run --frozen modal run --env=staging " + "src/policyengine_simulation_gateway/smoke_app.py", + "run --frozen modal run --env=staging src/modal/smoke_app.py", + "run --frozen modal run --env=staging " + "src/modal/v2_static_runtime_files_smoke.py", + # Builds every v2 image and imports its entry points. + "run --frozen modal run --env=staging src/modal/v2_image_smoke.py", + ] def test_workflow_supplies_hf_token(self): """CI must validate the GitHub value before synchronizing it to Modal.""" diff --git a/projects/policyengine-simulation-executor/tests/test_v2_image_smoke.py b/projects/policyengine-simulation-executor/tests/test_v2_image_smoke.py new file mode 100644 index 000000000..56341d053 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_v2_image_smoke.py @@ -0,0 +1,129 @@ +"""Tests for the pre-merge smoke of every v2 Modal image.""" + +from __future__ import annotations + +import importlib +import sys + +import pytest +from fixtures.fake_modal import install_fake_modal +from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionResult, +) + + +@pytest.fixture +def v2_modules(monkeypatch): + """Import the v2 application and its image smoke under the fake Modal.""" + install_fake_modal(monkeypatch) + sys.modules.pop("src.modal.v2_app", None) + sys.modules.pop("src.modal.v2_image_smoke", None) + v2_app = importlib.import_module("src.modal.v2_app") + return v2_app, importlib.import_module("src.modal.v2_image_smoke") + + +def test_the_smoke_builds_every_v2_image(v2_modules) -> None: + v2_app, smoke = v2_modules + functions = dict(smoke.app.function_calls) + + assert functions["smoke_v2_coordinator_image"]["image"] is v2_app.coordinator_image + assert functions["smoke_v2_uk_worker_image"]["image"] is v2_app.uk_worker_image + assert functions["smoke_v2_household_image"]["image"] is v2_app.household_image + # The US worker image without its last data layer, the Stage 12 cache, + # which only a deployment can fetch: the deployment's precompute writes it + # with the deployment's Cloud Storage credential. + us_image = functions["smoke_v2_us_worker_image"]["image"] + data_install = next(call for call in us_image.calls if call[0] == "run_commands") + assert "--country us" in data_install[1][0] + assert "--country uk" not in data_install[1][0] + assert not any(call[0] == "run_function" for call in us_image.calls) + assert [call[0] for call in us_image.calls] == [ + call[0] for call in v2_app.us_worker_image.calls if call[0] != "run_function" + ] + assert all( + declaration.get("secrets") in (None, []) for declaration in functions.values() + ) + + +@pytest.mark.parametrize( + "function_name", + [ + "smoke_v2_coordinator_image", + "smoke_v2_uk_worker_image", + "smoke_v2_us_worker_image", + ], +) +def test_report_image_smokes_import_the_v2_entry_points( + v2_modules, function_name +) -> None: + _, smoke = v2_modules + + report = getattr(smoke, function_name)() + + assert report["modules"] == list(smoke.REPORT_ENTRY_POINT_MODULES) + assert { + "src.modal.v2_app", + "policyengine_simulation_executor.stage12_runtime", + "policyengine_simulation_executor.stage12_worker_validation", + # Imported only while handling a request. + "policyengine_simulation_executor.simulation_runtime", + "microdf", + } <= set(report["modules"]) + assert report["policyengine_version"] + + +def test_household_image_smoke_calculates_each_validation_household( + monkeypatch, v2_modules +) -> None: + from policyengine_simulation_executor import household_calculation + + _, smoke = v2_modules + calculated: list[str] = [] + + def calculate(request, *, worker, required_country): + assert required_country == request.country + calculated.append(request.country) + situation = request.household.to_situation() + net_income = situation["households"]["your household"]["household_net_income"] + for period in net_income: + net_income[period] = 20_000.0 + return HouseholdExecutionResult.model_validate( + { + "calculation_id": request.calculation_id, + "output": { + "country": request.country, + "bundle": request.bundle, + "results": [{"policy_key": "baseline", "household": situation}], + }, + "worker": worker, + } + ) + + monkeypatch.setattr(household_calculation, "calculate_household", calculate) + + report = smoke.smoke_v2_household_image() + + assert calculated == ["us", "uk"] + assert report["household_net_income"] == {"us": 20_000.0, "uk": 20_000.0} + assert report["modules"] == [ + "src.modal.v2_app", + "policyengine_simulation_executor.household_calculation", + "policyengine_simulation_executor.household_validation", + ] + + +def test_the_household_image_smoke_fails_when_a_calculation_fails( + monkeypatch, v2_modules +) -> None: + from policyengine_simulation_executor import household_calculation + + _, smoke = v2_modules + + def fail(request, *, worker, required_country): + msg = "model exploded" + raise RuntimeError(msg) + + monkeypatch.setattr(household_calculation, "calculate_household", fail) + + with pytest.raises(RuntimeError, match="model exploded"): + smoke.smoke_v2_household_image() From 9063f1045b1755e905b80b5f8cde75672d518356 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:01:37 +0400 Subject: [PATCH 10/19] Word uncertifiable SPM receipts exactly as the Public API does A household's SPM receipt that the installed country model reports in a form that cannot be certified now fails with the Public API's reason, checked in its order: the forecast hash, scenario and geography kind are all required before they are compared with the selection, then the years mapping, the runtime versions and the forecast ID. The contract's receipt model words its year and metro-area errors as the Public API's model does, and a receipt that is not an object, or whose years are not a mapping, is reported the same way. Co-Authored-By: Claude Opus 5.5 --- .../policyengine_simulation_contract/spm.py | 8 +-- .../household_spm.py | 52 +++++++++----- .../tests/test_household_spm.py | 70 +++++++++++++++++++ 3 files changed, 109 insertions(+), 21 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index c2bddcae1..e666da24f 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -193,10 +193,10 @@ def serialize_provenance( # noqa: ANN201 # a return annotation would become th def validate_years(cls, value: tuple[str, ...]) -> tuple[str, ...]: """Require four-digit years, sorted ascending with no duplicates.""" if any(re.fullmatch(r"[0-9]{4}", year) is None for year in value): - msg = "SPM provenance years must use four-digit years" + msg = "years must contain four-digit calendar years" raise ValueError(msg) if value != tuple(sorted(set(value))): - msg = "SPM provenance years must be sorted and unique" + msg = "years must be sorted and unique" raise ValueError(msg) return value @@ -205,10 +205,10 @@ def validate_provenance_location(self) -> "SPMProvenance": """Require ``geography_id`` for metro receipts and forbid it otherwise.""" if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): - msg = "An SPM area receipt requires geography_id" + msg = "A metro SPM receipt requires geography_id" raise ValueError(msg) elif self.geography_id is not None: - msg = "Only an SPM area receipt accepts geography_id" + msg = "Only a metro SPM receipt accepts geography_id" raise ValueError(msg) return self diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py index 2c5b61d29..56860cee3 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py @@ -323,51 +323,69 @@ def _receipt_mapping(receipt: Mapping[str, object], field_name: str) -> Mapping: value = receipt.get(field_name) if not isinstance(value, Mapping): msg = f"Country SPM receipt requires {field_name}" - raise TypeError(msg) + raise ValueError(msg) # noqa: TRY004 # only ValueError means an uncertifiable receipt, as in the Public API return value +def _receipt_years(receipt: Mapping[str, object]) -> tuple[str, ...]: + """Return the receipt's years, sorted; raise ValueError if they are unusable.""" + source_years = receipt.get("years") + if not isinstance(source_years, Mapping): + msg = "Country SPM receipt requires a years mapping" + raise ValueError(msg) # noqa: TRY004 # only ValueError means an uncertifiable receipt, as in the Public API + if any(not isinstance(year, str) for year in source_years): + msg = "Country SPM receipt years must be strings" + raise ValueError(msg) + return tuple(sorted(source_years)) + + def _compact_receipt( selection: SPMSelection, receipt: Mapping[str, object] ) -> SPMProvenance: - """Reduce the country model's rich receipt to the compact public receipt.""" + """Reduce the country model's rich receipt to the compact public receipt. + + Checks run in the Public API's order, so the first problem reported is the + one it reports: the required texts, their match with the selection, the + years, the runtime versions, then the forecast ID. + """ forecast_sha256 = _required_text(receipt, "forecast_sha256") + scenario = _required_text(receipt, "scenario") + geography_kind = _required_text(receipt, "geography_kind") if ( forecast_sha256 != selection.forecast_content_sha256 - or _required_text(receipt, "scenario") != selection.scenario - or _required_text(receipt, "geography_kind") != selection.geography_kind + or scenario != selection.scenario + or geography_kind != selection.geography_kind ): msg = "Country SPM receipt differs from the SPM configuration" raise ValueError(msg) - source_years = _receipt_mapping(receipt, "years") - if any(not isinstance(year, str) for year in source_years): - msg = "Country SPM receipt years must be strings" - raise ValueError(msg) + years = _receipt_years(receipt) + runtime_versions = SPMRuntimeVersions.model_validate( + _receipt_mapping(receipt, "runtime_versions") + ) return build_spm_provenance( forecast_id=_required_text(receipt, "forecast_id"), forecast_sha256=forecast_sha256, selection=selection, - years=tuple(sorted(source_years)), - runtime_versions=SPMRuntimeVersions.model_validate( - _receipt_mapping(receipt, "runtime_versions") - ), + years=years, + runtime_versions=runtime_versions, ) def _country_receipt(simulation: object) -> SPMProvenance: """Reduce the simulation's receipt. - Raises ValueError or TypeError if the receipt cannot be certified. + Raises ValueError if the receipt cannot be certified. A simulation with + ``spm_config`` but no ``spm_provenance`` method raises AttributeError, as + in the Public API. """ selection = _complete_settings(getattr(simulation, "spm_config", None)) if selection is None: msg = "Country SPM configuration is incomplete" raise ValueError(msg) - receipt_method = getattr(simulation, "spm_provenance", None) - source = receipt_method() if callable(receipt_method) else None + source = simulation.spm_provenance() # pyright: ignore[reportAttributeAccessIssue] # a country Simulation with spm_config has this method if not isinstance(source, Mapping): msg = "Country SPM receipt must be an object" - raise TypeError(msg) + raise ValueError(msg) # noqa: TRY004 # only ValueError means an uncertifiable receipt, as in the Public API return _compact_receipt(selection, source) @@ -377,7 +395,7 @@ def _calculation_receipt(simulation: object) -> SPMProvenance | None: return None try: return _country_receipt(simulation) - except (TypeError, ValueError) as error: + except ValueError as error: raise SPMInputError( _CONFIGURATION_UNAVAILABLE, "The installed country model reported an SPM receipt this API cannot " diff --git a/projects/policyengine-simulation-executor/tests/test_household_spm.py b/projects/policyengine-simulation-executor/tests/test_household_spm.py index 1a3cd975a..b20263705 100644 --- a/projects/policyengine-simulation-executor/tests/test_household_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_household_spm.py @@ -243,6 +243,76 @@ def test_an_uncertifiable_country_receipt_is_an_spm_error() -> None: ) +_UNCERTIFIABLE = ( + "The installed country model reported an SPM receipt this API cannot certify: " +) + + +def _receipt_without(*names: str, **overrides: object) -> dict: + receipt = _country_receipt(**overrides) + for name in names: + del receipt[name] + return receipt + + +@pytest.mark.parametrize( + ("receipt", "reason"), + [ + # Each required text is checked before the receipt is compared. + ( + _receipt_without("scenario", forecast_sha256="b" * 64), + "Country SPM receipt requires scenario", + ), + ( + _receipt_without("geography_kind", scenario="flat"), + "Country SPM receipt requires geography_kind", + ), + ( + _country_receipt(years=["2026"]), + "Country SPM receipt requires a years mapping", + ), + ( + _country_receipt(years={2026: {}}), + "Country SPM receipt years must be strings", + ), + # Runtime versions are checked before the forecast ID. + ( + _receipt_without("runtime_versions", "forecast_id"), + "Country SPM receipt requires runtime_versions", + ), + ( + _receipt_without("forecast_id"), + "Country SPM receipt requires forecast_id", + ), + ( + _country_receipt(years={"26": {}}), + "years: years must contain four-digit calendar years", + ), + ], +) +def test_uncertifiable_receipts_are_worded_as_the_public_api_words_them( + receipt, reason +) -> None: + with pytest.raises(SPMInputError) as raised: + household_spm.household_spm_provenance( + FakeSimulation(receipt=receipt), DEFAULTS + ) + + assert raised.value.code == "SPM_CONFIGURATION_UNAVAILABLE" + assert raised.value.message == _UNCERTIFIABLE + reason + + +def test_a_receipt_that_is_not_an_object_is_uncertifiable() -> None: + with pytest.raises(SPMInputError) as raised: + household_spm.household_spm_provenance( + FakeSimulation(receipt=["not", "an", "object"]), DEFAULTS + ) + + assert ( + raised.value.message == _UNCERTIFIABLE + "Country SPM receipt must be an object" + ) + + def test_a_receipt_for_another_selection_is_rejected() -> None: with pytest.raises(ValueError, match="differs from the request"): household_spm.household_spm_provenance( From d9462a3b53149fe1e3414f81597d5115463f07ae Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:06:15 +0400 Subject: [PATCH 11/19] Check a household's SPM settings exactly as sent, with the Public API's codes The request's spm field now carries the caller's SPM settings as any JSON value, and the household function checks them as the Public API's routes do: settings for a country other than the US are unsupported, an explicit null for the US is invalid, and malformed settings are invalid with the Public API's message, before the bundle's certification is read. A metro selection needs an area and any other kind accepts none, judged on the settings as sent, as the Public API's model judges them. Before, a malformed selection or an explicit null failed the request's validation instead of returning a typed SPM failure. The request's documentation now says which parts of a saved household's calculation match the Public API's stored-household route (the calculated household, warnings, SPM receipt and SPM codes) and which do not (two error messages, which that route words with the household and policy IDs). Co-Authored-By: Claude Opus 5.5 --- .../household_execution.py | 68 +++++++++---------- .../tests/test_household_execution.py | 16 ++++- .../household_calculation.py | 1 + .../household_spm.py | 66 +++++++++++++----- .../tests/test_household_calculation.py | 29 +++++++- .../tests/test_household_spm.py | 50 ++++++++++++-- 6 files changed, 169 insertions(+), 61 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 df1f17be9..254df8be1 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 @@ -33,9 +33,9 @@ from pydantic import ( AfterValidator, BaseModel, - BeforeValidator, ConfigDict, Field, + JsonValue, PrivateAttr, TypeAdapter, field_validator, @@ -47,7 +47,6 @@ SPMErrorDetail, SPMProvenance, SPMResolvedConfiguration, - SPMSelection, ) from policyengine_simulation_contract.stage12_bundle import ( CountryId, @@ -416,20 +415,6 @@ def _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. @@ -484,20 +469,33 @@ class HouseholdExecutionRequest(StrictHouseholdModel): ``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 uncalculated value. ``spm_selection_source`` says - whether the selection came from this request or from a saved household, - which changes only how an SPM forecast 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"``. + ``/calculate-full`` route does. ``spm`` carries the caller's SPM settings + exactly as sent, as any JSON value, and the household function checks + them with the Public API's SPM error codes: settings for a country other + than the US are unsupported, and for the US an explicit null or malformed + settings are invalid. Omit ``spm`` to inherit the bundle's certified + defaults; sending it, even as an empty object, is a choice, and then a + missing SPM input is an error rather than an uncalculated value. + ``spm_selection_source`` says whether the settings came from this request + or from a saved household, which changes only how an SPM forecast the + bundle no longer carries is reported. + + To calculate a saved household with the inputs the Public API's + stored-household route (``GET /{country}/household/{id}/policy/{policy_id}``) + uses, set ``add_missing_variables``, because that route always adds + missing variables; pass the household's saved SPM selection as ``spm`` + only if it has one (a saved null is no choice); and set + ``spm_selection_source`` to ``"saved_household"``. The calculated + household, warnings, SPM receipt and SPM error codes then match that + route's. Its error messages do not: the route reports a household it + cannot read, and any other failure except an SPM error, as an error + calculating that household ID under that policy ID, while a household + function reports the request routes' wording (``invalid_household`` and + ``calculation_failed``), so a caller that serves the stored-household + route rewrites those two messages. + + These models check the request's structure; a request they reject (for + example a group without ``members``) is the caller's to report. """ schema_version: Literal[1] = HOUSEHOLD_CONTRACT_SCHEMA_VERSION @@ -506,7 +504,7 @@ class HouseholdExecutionRequest(StrictHouseholdModel): household: HouseholdDocument policies: Policies add_missing_variables: bool = False - spm: SPMChoice = None + spm: JsonValue = None spm_selection_source: Literal["request", "saved_household"] = "request" bundle: HouseholdBundleProvenance @@ -514,16 +512,16 @@ class HouseholdExecutionRequest(StrictHouseholdModel): def serialize_request( self, handler: SerializerFunctionWrapHandler ) -> dict[str, Any]: - """Leave out ``spm`` when no selection was made, so omission is kept.""" + """Leave out ``spm`` when the caller sent none, so omission is kept.""" dumped: dict[str, Any] = handler(self) - if self.spm is None: + if not self.spm_requested: 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 + """Whether the caller sent SPM settings rather than inheriting the defaults.""" + return "spm" in self.model_fields_set class HouseholdPolicyResult(StrictHouseholdModel): diff --git a/libs/policyengine-simulation-contract/tests/test_household_execution.py b/libs/policyengine-simulation-contract/tests/test_household_execution.py index 3d1ef4f2c..ce7aac034 100644 --- a/libs/policyengine-simulation-contract/tests/test_household_execution.py +++ b/libs/policyengine-simulation-contract/tests/test_household_execution.py @@ -330,9 +330,19 @@ def test_an_empty_spm_selection_is_still_a_choice() -> None: assert HouseholdExecutionRequest.model_validate(dumped).spm_requested is True -def test_an_explicit_null_spm_is_rejected() -> None: - with pytest.raises(ValidationError, match="omit it"): - _request(spm=None) +def test_spm_settings_travel_exactly_as_the_caller_sent_them() -> None: + # The household function checks them, so a malformed selection or an + # explicit null fails with the Public API's SPM error codes. + for sent in (None, {"geography_kind": "planet"}, "national"): + request = _request(spm=sent) + + assert request.spm == sent + assert request.spm_requested is True + dumped = request.model_dump(mode="json") + assert dumped["spm"] == sent + transported = HouseholdExecutionRequest.model_validate(dumped) + assert transported.spm == sent + assert transported.spm_requested is True def test_a_request_rejects_unknown_fields_and_other_schema_versions() -> None: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py index 9a9b29c26..647fed5c6 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py @@ -532,6 +532,7 @@ def _resolved_spm_selection( return resolve_household_spm_selection( request.country, request.spm, + sent=request.spm_requested, saved_household=request.spm_selection_source == "saved_household", ) except SPMInputError as error: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py index 56860cee3..94437a276 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py @@ -3,8 +3,9 @@ This reproduces the Public API's household SPM handling, so that a v2 household result and error carry the same SPM codes, messages and receipt: -- ``resolve_household_spm_selection`` merges a caller's selection over the - installed bundle's certified defaults, with the Public API's error codes; +- ``resolve_household_spm_selection`` checks a caller's settings, exactly as + sent, and merges them over the installed bundle's certified defaults, with + the Public API's error codes; - ``household_spm_error_detail`` recognizes the SPM errors that end a calculation; - ``household_spm_provenance`` reduces the country model's receipt to the @@ -250,34 +251,67 @@ def _resolve_certified_selection( return resolved.model_dump(mode="json") +def _requested_selection(selection: object) -> SPMSelection: + """Read the caller's settings; raise SPM_SETTINGS_INVALID if they are malformed. + + As in the Public API, the area rule applies to the settings as sent, with + an omitted geography kind read as county: a metro selection needs an area + and any other kind accepts none. The contract's model checks that rule + only once both fields are set, because Stage 12 requests inherit omitted + fields first. + """ + try: + chosen = SPMSelection.model_validate({} if selection is None else selection) + except ValidationError as error: + raise SPMInputError(_SETTINGS_INVALID, _readable_error(error)) from error + if chosen.geography_kind == "metro": + if not chosen.geography_id or not chosen.geography_id.strip(): + raise SPMInputError( + _SETTINGS_INVALID, "An SPM area selection requires geography_id" + ) + elif chosen.geography_id is not None: + raise SPMInputError( + _SETTINGS_INVALID, "Only an SPM area selection accepts geography_id" + ) + return chosen + + def resolve_household_spm_selection( country: str, - selection: SPMSelection | None, + selection: object, *, + sent: bool, saved_household: bool, ) -> dict[str, Any] | None: - """Resolve a household's SPM selection against the installed bundle. - - ``selection`` is None when the caller made no choice. Returns the complete - resolved selection, or None when no SPM applies (any country but the US, - or a bundle whose US model predates canonical SPM). Raises SPMInputError - with the Public API's code and message otherwise. ``saved_household`` says - the selection was read from a saved household, which changes only how a - forecast hash the bundle no longer carries is reported. + """Resolve a household's SPM settings, exactly as sent, against the bundle. + + ``sent`` says whether the caller sent settings at all; ``selection`` is + what it sent. Returns the complete resolved selection, or None when no SPM + applies (any country but the US, or a bundle whose US model predates + canonical SPM). Raises SPMInputError with the Public API's code and + message otherwise, checking in its order: settings for another country, + an explicit null, malformed settings, then the bundle's certification. + ``saved_household`` says the settings were read from a saved household, + which changes only how a forecast hash the bundle no longer carries is + reported. """ if country != "us": - if selection is not None: + if sent: raise SPMInputError( _SETTINGS_UNSUPPORTED, "SPM settings are only available for the US" ) return None - configured = _certified_configuration(selection_sent=selection is not None) + if sent and selection is None: + raise SPMInputError( + _SETTINGS_INVALID, + "spm must be an object; omit it to inherit the certified defaults.", + ) + chosen = _requested_selection(selection) + configured = _certified_configuration(selection_sent=sent) if configured is None: return None return _resolve_certified_selection( - selection if selection is not None else SPMSelection(), - configured, - saved_household=saved_household, + chosen, configured, saved_household=saved_household ) diff --git a/projects/policyengine-simulation-executor/tests/test_household_calculation.py b/projects/policyengine-simulation-executor/tests/test_household_calculation.py index c6874b31e..1ae011b7f 100644 --- a/projects/policyengine-simulation-executor/tests/test_household_calculation.py +++ b/projects/policyengine-simulation-executor/tests/test_household_calculation.py @@ -278,7 +278,7 @@ def test_an_spm_input_error_fails_only_a_chosen_measurement(monkeypatch) -> None monkeypatch.setattr( household_calculation, "resolve_household_spm_selection", - lambda country, selection, *, saved_household: dict(resolved), + lambda country, selection, *, sent, saved_household: dict(resolved), ) monkeypatch.setattr( household_calculation, @@ -324,7 +324,7 @@ def __init__(self, **options: object) -> None: monkeypatch.setattr( household_calculation, "resolve_household_spm_selection", - lambda country, selection, *, saved_household: resolved, + lambda country, selection, *, sent, saved_household: resolved, ) def record_receipt(simulation, expected_selection): @@ -353,7 +353,7 @@ def test_an_spm_configuration_error_fails_even_an_inherited_measurement( monkeypatch.setattr( household_calculation, "resolve_household_spm_selection", - lambda country, selection, *, saved_household: { + lambda country, selection, *, sent, saved_household: { "forecast_content_sha256": "a" * 64, "scenario": "ce_trend", "geography_kind": "national", @@ -411,6 +411,29 @@ def failing_simulation(**_: object): assert outcome.calculation_id == CALCULATION_ID +@pytest.mark.parametrize( + ("country", "settings", "spm_code"), + [ + ("uk", None, "SPM_SETTINGS_UNSUPPORTED"), + ("us", None, "SPM_SETTINGS_INVALID"), + ("us", {"geography_kind": "planet"}, "SPM_SETTINGS_INVALID"), + ], +) +def test_spm_settings_as_sent_fail_with_public_api_codes( + country, settings, spm_code +) -> None: + outcome = _calculate_with_stand_in( + _stand_in_simulation({}), country=country, spm=settings + ) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "invalid_spm_selection" + assert outcome.error.spm_error is not None + assert outcome.error.spm_error.code == spm_code + assert outcome.error.message == outcome.error.spm_error.message + assert outcome.error.policy_key is None + + def test_a_receipt_that_cannot_be_certified_is_an_spm_failure(monkeypatch) -> None: def uncertifiable(simulation, expected_selection): code = "SPM_CONFIGURATION_UNAVAILABLE" diff --git a/projects/policyengine-simulation-executor/tests/test_household_spm.py b/projects/policyengine-simulation-executor/tests/test_household_spm.py index b20263705..919f09ba6 100644 --- a/projects/policyengine-simulation-executor/tests/test_household_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_household_spm.py @@ -3,7 +3,7 @@ from __future__ import annotations import pytest -from policyengine_simulation_contract.spm import SPMInputError, SPMSelection +from policyengine_simulation_contract.spm import SPMInputError from policyengine_simulation_executor import household_spm FORECAST_SHA = "3d86d5c4c0423480e6b69b75d222ffa4a7a2639e4094df5ba2504af01be17173" @@ -50,15 +50,20 @@ def certified_bundle(monkeypatch): monkeypatch.setattr(household_spm, "_forecast", lambda sha: FakeForecast()) -def _resolve(selection=None, *, country="us", saved_household=False): +# Stands for a request that sends no SPM settings at all. +OMITTED = object() + + +def _resolve(selection=OMITTED, *, country="us", saved_household=False): return household_spm.resolve_household_spm_selection( country, - None if selection is None else SPMSelection.model_validate(selection), + None if selection is OMITTED else selection, + sent=selection is not OMITTED, saved_household=saved_household, ) -def _error(selection=None, **options) -> SPMInputError: +def _error(selection=OMITTED, **options) -> SPMInputError: with pytest.raises(SPMInputError) as raised: _resolve(selection, **options) return raised.value @@ -86,6 +91,43 @@ def test_other_countries_accept_no_selection(certified_bundle) -> None: ) +def test_an_explicit_null_is_invalid_for_the_us_and_unsupported_elsewhere( + certified_bundle, +) -> None: + us_error = _error(None) + uk_error = _error(None, country="uk") + + assert (us_error.code, us_error.message) == ( + "SPM_SETTINGS_INVALID", + "spm must be an object; omit it to inherit the certified defaults.", + ) + assert (uk_error.code, uk_error.message) == ( + "SPM_SETTINGS_UNSUPPORTED", + "SPM settings are only available for the US", + ) + + +@pytest.mark.parametrize( + ("selection", "message"), + [ + ({"geography_kind": "metro"}, "An SPM area selection requires geography_id"), + ({"geography_id": "31080"}, "Only an SPM area selection accepts geography_id"), + ("national", "Input should be a valid dictionary or instance of SPMSelection"), + ], +) +def test_a_malformed_selection_is_invalid_before_the_bundle_is_read( + monkeypatch, selection, message +) -> None: + # As in the Public API, even a bundle without SPM support reports the + # settings' own problem first. + monkeypatch.setattr(household_spm, "_bundle_spm_configuration", lambda: None) + monkeypatch.setattr(household_spm, "_us_model_supports_spm", lambda: False) + + error = _error(selection) + + assert (error.code, error.message) == ("SPM_SETTINGS_INVALID", message) + + def test_a_foreign_artifact_hash_is_reported_by_where_it_came_from( certified_bundle, ) -> None: From db610916aeaa98b244a63ea8f4b169300a8cb766 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:06:50 +0400 Subject: [PATCH 12/19] Load the household SPM forecast without the main path's SPM module The household SPM code now loads the certified forecast with its own per-process loader instead of a wrapper added to the main path's SPM module, which returns to its state on main. The household path and the main path share no SPM code. Co-Authored-By: Claude Opus 5.5 --- .../household_spm.py | 16 ++++++++++++---- .../src/policyengine_simulation_executor/spm.py | 9 --------- .../tests/test_household_spm.py | 12 ++++++------ 3 files changed, 18 insertions(+), 19 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py index 94437a276..cf5b05817 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py @@ -16,6 +16,7 @@ import inspect from collections.abc import Mapping +from functools import lru_cache from typing import TYPE_CHECKING, Any, NoReturn from policyengine_simulation_contract.spm import ( @@ -115,10 +116,17 @@ def _us_model_supports_spm() -> bool: return _simulation_supports_spm(Simulation) -def _forecast(expected_sha256: str) -> SPMForecast: - from policyengine_simulation_executor.spm import certified_spm_forecast +@lru_cache(maxsize=4) +def certified_forecast(expected_sha256: str) -> SPMForecast: + """Return the packaged SPM forecast whose content has ``expected_sha256``. - return certified_spm_forecast(expected_sha256) + Loaded once per process and digest; raises when the packaged forecast has + another digest. The household path keeps this loader apart from the main + path's SPM module. + """ + from spm_calculator import load_forecast + + return load_forecast(expected_sha256=expected_sha256) def _certified_defaults(configured: object) -> tuple[SPMSelection, SPMForecast]: @@ -134,7 +142,7 @@ def _certified_defaults(configured: object) -> tuple[SPMSelection, SPMForecast]: if not _us_model_supports_spm(): msg = "The installed US model does not support canonical SPM" raise ValueError(msg) - return defaults, _forecast(defaults.forecast_content_sha256) + return defaults, certified_forecast(defaults.forecast_content_sha256) def _validated_selection(values: dict[str, Any], forecast: SPMForecast) -> SPMSelection: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py index cbc29f387..b5c2490a7 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -34,15 +34,6 @@ def _forecast(expected_sha256: str | None) -> SPMForecast: return load_forecast(expected_sha256=expected_sha256) -def certified_spm_forecast(expected_sha256: str) -> SPMForecast: - """Return the packaged SPM forecast whose content has ``expected_sha256``. - - Loaded once per process and digest; raises when the packaged forecast has - another digest. - """ - return _forecast(expected_sha256) - - @lru_cache(maxsize=1) def runtime_spm_capability() -> SPMCapability | None: """Return the installed image's certified SPM capability. diff --git a/projects/policyengine-simulation-executor/tests/test_household_spm.py b/projects/policyengine-simulation-executor/tests/test_household_spm.py index 919f09ba6..d1053aa52 100644 --- a/projects/policyengine-simulation-executor/tests/test_household_spm.py +++ b/projects/policyengine-simulation-executor/tests/test_household_spm.py @@ -47,7 +47,7 @@ def certified_bundle(monkeypatch): household_spm, "_bundle_spm_configuration", lambda: dict(DEFAULTS) ) monkeypatch.setattr(household_spm, "_us_model_supports_spm", lambda: True) - monkeypatch.setattr(household_spm, "_forecast", lambda sha: FakeForecast()) + monkeypatch.setattr(household_spm, "certified_forecast", lambda sha: FakeForecast()) # Stands for a request that sends no SPM settings at all. @@ -207,7 +207,7 @@ def test_changing_the_geography_kind_drops_an_inherited_area(monkeypatch) -> Non lambda: {**DEFAULTS, "geography_kind": "metro", "geography_id": "31080"}, ) monkeypatch.setattr(household_spm, "_us_model_supports_spm", lambda: True) - monkeypatch.setattr(household_spm, "_forecast", lambda sha: FakeForecast()) + monkeypatch.setattr(household_spm, "certified_forecast", lambda sha: FakeForecast()) assert _resolve({"geography_kind": "national"}) == { **DEFAULTS, @@ -398,9 +398,9 @@ def test_spm_errors_are_recognized_through_wrapping_exceptions() -> None: def test_the_certified_forecast_is_loaded_once_by_its_digest() -> None: - from policyengine_simulation_executor.spm import certified_spm_forecast + # The household path keeps its own loader, separate from the main path's + # SPM module. + forecast = household_spm.certified_forecast(FORECAST_SHA) - forecast = certified_spm_forecast(FORECAST_SHA) - - assert certified_spm_forecast(FORECAST_SHA) is forecast + assert household_spm.certified_forecast(FORECAST_SHA) is forecast assert forecast.years From 85afc4e58d9cf86dfcc5a62460a6e2493a2dc943 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:07:12 +0400 Subject: [PATCH 13/19] List every known difference from the Public API in one place The household calculation module's documentation now lists each known difference from the Public API's household calculation, which were described in several places or not at all: how undefined numbers and dates are written, the UTC default year, an age given as one value under calculate-full, an explicit null axes, checking every policy's parameters first, the request routes' error wording, and requests the contract rejects. Co-Authored-By: Claude Opus 5.5 --- .../household_calculation.py | 24 +++++++++++++++---- 1 file changed, 20 insertions(+), 4 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py index 647fed5c6..0184783f4 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py @@ -12,10 +12,26 @@ the Public API's matching error responses. This keeps v2 results comparable with the Public API's results while both exist. -Two values are written differently, because the contract carries only finite -numbers and JSON scalars: an undefined number is the string "NaN" (the Public -API writes a bare ``NaN``, which is not JSON), and a date is the HTTP date -text the Public API's Flask encoder writes. +Every known difference from the Public API is listed here: + +- An undefined number is the string "NaN": the contract carries only finite + numbers and JSON scalars, and the Public API writes a bare ``NaN``, which + is not JSON. A date is the HTTP date text the Public API's Flask encoder + writes. +- The default household year, used when ``you`` has no age by period, is the + current UTC year; the Public API uses its host's local date. +- With ``add_missing_variables``, an age given for ``you`` as one value, not + by period, counts as no age; the Public API fails with an internal error. +- An explicit null ``axes`` means no axes, because the contract drops it; the + Public API treats the key's presence as a household with axes. +- Every policy's parameters are checked before any policy is calculated, so + an unknown parameter in a reform fails the whole request. +- Its messages are the request routes' (``/calculate`` and + ``/calculate-full``). The stored-household route words a household it + cannot read, and any other failure except an SPM error, with the household + and policy IDs, which a household function does not have. +- A request the contract's models reject, such as a group without + ``members``, never reaches this module; the caller reports it. """ from __future__ import annotations From 5c700a1ce98ce1cef54b08df330bd91bbb8bb389 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:09:50 +0400 Subject: [PATCH 14/19] Test publishing after a publisher released before schema version 2 A review asked whether a publish after code released before schema version 2 (as a deployed revert would run) could start from a stale schema version 2 document and drop what that code wrote. The schema version 2 reader already uses the document only while it is current; these tests pin that down for a publish: it keeps an application the older code redeployed without household functions, it takes the default from what the older code wrote, and older keys that a household-only schema version 2 publish left alone change nothing. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_manifest.py | 93 +++++++++++++++++-- 1 file changed, 85 insertions(+), 8 deletions(-) diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index cd7bccfe7..5d02bf4dc 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -23,6 +23,7 @@ Stage12CacheRelease, V2CountryWorker, V2ManifestLoader, + V2VersionManifest, V2WorkerValidation, V2WorkerVersion, assert_separate_manifest_names, @@ -477,6 +478,90 @@ def test_a_household_only_country_leaves_schema_1_readers_on_their_manifest() -> assert worker.economy_worker("us") is not None +def _validated_on(worker: V2WorkerVersion, day: int) -> V2WorkerVersion: + """Return ``worker`` as validated on ``day`` of September 2026.""" + validation = worker.validation.model_copy( + update={"validated_at": datetime(2026, 9, day, tzinfo=UTC)} + ) + return worker.model_copy(update={"validation": validation}) + + +def _write_as_a_schema_1_publisher( + store: FakeStore, *workers: V2WorkerVersion, default_version: str +) -> None: + """Write the older keys as a publisher released before schema 2 writes them.""" + manifest = V2VersionManifest( + generation="older-publisher", + default_version=default_version, + versions={worker.bundle.policyengine_version: worker for worker in workers}, + ) + store[ACTIVE_MANIFEST_KEY] = schema_1_manifest_payload(manifest, include_cache=True) + store[PRE_CACHE_MANIFEST_KEY] = schema_1_manifest_payload( + manifest, include_cache=False + ) + + +def test_a_publish_keeps_what_an_older_publisher_wrote_after_it() -> None: + # A deployment of code released before schema 2 (for example after a + # revert) redeploys 5.2.0 without household functions and writes only the + # older keys. The next schema 2 publish must start from that, not from + # its own stale document. + store = FakeStore() + publish_v2_manifest( + store=store, + worker=_validated_on(_worker(with_household_functions=True), 14), + ) + _write_as_a_schema_1_publisher( + store, _validated_on(_worker(), 20), default_version="5.2.0" + ) + + manifest = publish_v2_manifest( + store=store, + worker=_validated_on(_worker("5.3.0", with_household_functions=True), 25), + ) + + assert sorted(manifest.versions) == ["5.2.0", "5.3.0"] + assert manifest.versions["5.2.0"].household_worker("us") is None + assert manifest.versions["5.2.0"].validation.validated_at.day == 20 + assert manifest.default_version == "5.3.0" + + +def test_a_publish_keeps_the_default_of_the_latest_document() -> None: + store = FakeStore() + publish_v2_manifest(store=store, worker=_validated_on(_worker("5.3.0"), 14)) + # An older publisher later made 5.2.0 the default. + _write_as_a_schema_1_publisher( + store, + _validated_on(_worker("5.3.0"), 14), + _validated_on(_worker("5.2.0"), 20), + default_version="5.2.0", + ) + + manifest = publish_v2_manifest( + store=store, worker=_validated_on(_worker("5.2.5"), 25) + ) + + # 5.2.5 is newer than the latest default, 5.2.0, so it becomes the default. + assert manifest.default_version == "5.2.5" + assert sorted(manifest.versions) == ["5.2.0", "5.2.5", "5.3.0"] + + +def test_older_keys_left_behind_by_a_schema_2_publish_change_nothing() -> None: + store = FakeStore({ACTIVE_MANIFEST_KEY: _schema_1_document(include_cache=True)}) + publish_v2_manifest( + store=store, worker=_validated_on(_worker(household_only=("uk",)), 14) + ) + + manifest = publish_v2_manifest( + store=store, + worker=_validated_on(_worker("5.3.0", household_only=("uk",)), 25), + ) + + assert sorted(manifest.versions) == ["5.1.0", "5.2.0", "5.3.0"] + assert manifest.versions["5.2.0"].economy_worker("uk") is None + assert manifest.default_version == "5.3.0" + + def test_parse_v2_manifest_rejects_unknown_schema_versions() -> None: document = _schema_1_document(include_cache=True) document["schema_version"] = 3 @@ -701,14 +786,6 @@ def test_a_schema_2_document_left_behind_by_an_older_publish_is_not_read() -> No assert store.values[SCHEMA_2_MANIFEST_KEY] != stale -def _validated_on(worker: V2WorkerVersion, day: int) -> V2WorkerVersion: - """Return ``worker`` as validated on ``day`` of September 2026.""" - validation = worker.validation.model_copy( - update={"validated_at": datetime(2026, 9, day, tzinfo=UTC)} - ) - return worker.model_copy(update={"validation": validation}) - - def test_a_household_only_manifest_is_read_until_an_older_publish_follows() -> None: # Schema version 1 cannot describe a household-only country, so its # publish leaves the older keys alone; a later validation under them From fd45e2613b01a5b0dd7fa894f931eaf0675a22df Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:11:14 +0400 Subject: [PATCH 15/19] Resolve a comparison report's simulation callable as a typed text A v2 worker version now returns the simulation callable of an economy-capable country, and raises LookupError for any other country. The Simulation Entrypoint's comparison backend stores that callable on the prepared report and sends it to the coordinator, instead of looking up the country's worker by matching and sending a value that schema version 2 allows to be null. The adapter already skips a country that is not economy-capable, so a prepared report always has one. Co-Authored-By: Claude Opus 5.5 --- .../stage12_manifest.py | 14 ++++++ .../tests/test_stage12_manifest.py | 10 ++++ .../stage12_backend.py | 47 ++++++------------- 3 files changed, 39 insertions(+), 32 deletions(-) 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 4cb9562b6..1f935a9f1 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 @@ -212,6 +212,20 @@ def economy_worker(self, country: str) -> V2CountryWorker | None: None, ) + def economy_simulation_callable(self, country: str) -> str: + """Return the callable that runs one of ``country``'s report simulations. + + Raises LookupError when the manifest does not mark ``country`` as + economy-capable. + """ + worker = self.economy_worker(country) + if worker is None or worker.single_simulation_callable is None: + msg = ( + f"country {country!r} is not economy-capable in {self.application_name}" + ) + raise LookupError(msg) + return worker.single_simulation_callable + def household_worker(self, country: str) -> V2CountryWorker | None: """Return ``country``'s worker when it has a household function.""" return next( diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 5d02bf4dc..47c28b6fd 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -562,6 +562,16 @@ def test_older_keys_left_behind_by_a_schema_2_publish_change_nothing() -> None: assert manifest.default_version == "5.3.0" +def test_only_an_economy_capable_country_names_a_simulation_callable() -> None: + worker = _worker(household_only=("uk",)) + + assert worker.economy_simulation_callable("us") == "run_single_simulation_us" + with pytest.raises(LookupError, match="'uk' is not economy-capable"): + worker.economy_simulation_callable("uk") + with pytest.raises(LookupError, match="'fr' is not economy-capable"): + worker.economy_simulation_callable("fr") + + def test_parse_v2_manifest_rejects_unknown_schema_versions() -> None: document = _schema_1_document(include_cache=True) document["schema_version"] = 3 diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py index a1d232941..ed399fb56 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_backend.py @@ -20,7 +20,6 @@ ResultComparisonStatus, ) from policyengine_simulation_contract.stage12_manifest import ( - V2CountryWorker, V2ManifestLoader, V2WorkerVersion, ) @@ -100,12 +99,16 @@ class ReportOrigin: @dataclass(frozen=True) class PreparedReport: - """A comparison run ready to dispatch: its parent row, inputs and worker.""" + """A comparison run ready to dispatch: its parent row, inputs and worker. + + ``simulation_callable`` runs one of the report country's simulations; the + country is economy-capable, or the adapter would have skipped it. + """ parent: ComparisonReportRecord report: ReportExecutionInput worker_version: ResolvedWorkerVersion - country_worker: V2CountryWorker + simulation_callable: str class TemporaryStage12UnsupportedRequestError(ValueError): @@ -291,7 +294,6 @@ def _resolve_report( ) -> tuple[ ReportExecutionInput | None, ResolvedWorkerVersion, - V2CountryWorker, str | None, ]: worker_version = ResolvedWorkerVersion( @@ -303,21 +305,8 @@ def _resolve_report( evaluation_id=evaluation_id, worker=worker, ) - country = request_payload.get("country") - bundle_country = next( - (item for item in worker.bundle.countries if item.country == country), - worker.bundle.countries[0], - ) - country_worker = next( - item for item in worker.countries if item.country == bundle_country.country - ) skip_reason = adapted.skip_reason.value if adapted.skip_reason else None - return ( - adapted.report, - worker_version, - country_worker, - skip_reason, - ) + return adapted.report, worker_version, skip_reason def _parent_record( self, @@ -365,7 +354,7 @@ def _prepare_temporary_report( request_id: str, ) -> PreparedReport: evaluation_id = uuid4() - report, worker_version, country_worker, skip_reason = self._resolve_report( + report, worker_version, skip_reason = self._resolve_report( request_payload=request_payload, evaluation_id=evaluation_id, ) @@ -388,7 +377,9 @@ def _prepare_temporary_report( parent=parent, report=report, worker_version=worker_version, - country_worker=country_worker, + simulation_callable=worker_version.worker.economy_simulation_callable( + str(request_payload.get("country")) + ), ) def _prepare_automatic_report( @@ -433,14 +424,6 @@ def _prepare_automatic_report( }, ) return None - country = request_payload.get("country") - bundle_country = next( - (item for item in worker.bundle.countries if item.country == country), - worker.bundle.countries[0], - ) - country_worker = next( - item for item in worker.countries if item.country == bundle_country.country - ) parent = self._parent_record( evaluation_id=evaluation_id, origin=ReportOrigin( @@ -456,7 +439,9 @@ def _prepare_automatic_report( parent=parent, report=adapted.report, worker_version=worker_version, - country_worker=country_worker, + simulation_callable=worker.economy_simulation_callable( + str(request_payload.get("country")) + ), ) async def _spawn_report( @@ -519,9 +504,7 @@ async def _spawn_report( ), "worker_version": version, "modal_application": worker.application_name, - "simulation_callable": ( - prepared.country_worker.single_simulation_callable - ), + "simulation_callable": prepared.simulation_callable, "version_manifest_sha256": manifest_sha256, "bundle_manifest_sha256": worker.bundle_manifest_sha256, "artifact_prefix": artifact_prefix, From 615955f16ba7954eb28d2db6cf1fd6dc87de6ff0 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:13:03 +0400 Subject: [PATCH 16/19] Test a chosen US SPM selection and the validation household with real models Two real-model tests join the PR checks: a US household without a county leaves the SPM threshold uncalculated under the inherited county selection, fails with SPM_GEOGRAPHY_REQUIRED when the selection is chosen, and is calculated with a national receipt when the national kind is chosen; and the installed US and UK models calculate the household that release validation and the image check send. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_household_calculation.py | 45 +++++++++++++++++++ .../tests/test_household_validation.py | 18 ++++++++ 2 files changed, 63 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_household_calculation.py b/projects/policyengine-simulation-executor/tests/test_household_calculation.py index 1ae011b7f..1abc55163 100644 --- a/projects/policyengine-simulation-executor/tests/test_household_calculation.py +++ b/projects/policyengine-simulation-executor/tests/test_household_calculation.py @@ -663,6 +663,51 @@ def test_us_calculate_full_adds_and_calculates_every_missing_variable() -> None: ) +US_SINGLE_WITHOUT_COUNTY = { + "people": {"you": {"age": {"2026": 35}, "employment_income": {"2026": 30_000}}}, + "spm_units": { + "your spm unit": { + "members": ["you"], + "spm_unit_spm_threshold": {"2026": None}, + } + }, + "households": { + "your household": {"members": ["you"], "state_name": {"2026": "TX"}} + }, +} + + +def test_us_a_chosen_spm_selection_must_have_its_geography() -> None: + # The certified defaults measure by county. Inherited, a household without + # a county leaves the SPM threshold uncalculated; chosen, it is an error. + def calculate(**spm_settings): + return calculate_household( + _request("us", US_SINGLE_WITHOUT_COUNTY, **spm_settings), + worker=_worker("us"), + dependencies=HouseholdCalculationDependencies(today=TODAY), + ) + + inherited = calculate() + chosen = calculate(spm={}) + national = calculate(spm={"geography_kind": "national"}) + + assert isinstance(inherited, HouseholdExecutionResult), inherited + unit = inherited.output.results[0].household.to_situation()["spm_units"] + assert unit["your spm unit"]["spm_unit_spm_threshold"] == {"2026": None} + assert isinstance(chosen, HouseholdExecutionFailure) + assert chosen.error.code == "invalid_spm_selection" + assert chosen.error.spm_error is not None + assert chosen.error.spm_error.code == "SPM_GEOGRAPHY_REQUIRED" + assert isinstance(national, HouseholdExecutionResult), national + unit = national.output.results[0].household.to_situation()["spm_units"] + assert isinstance(unit["your spm unit"]["spm_unit_spm_threshold"]["2026"], float) + assert national.output.spm_config is not None + assert national.output.spm_config.geography_kind == "national" + receipt = national.output.results[0].spm_provenance + assert receipt is not None + assert receipt.geography_kind == "national" + + UK_SINGLE = { "people": {"you": {"age": {"2026": 40}, "employment_income": {"2026": 30_000}}}, "benunits": {"your benefit unit": {"members": ["you"]}}, diff --git a/projects/policyengine-simulation-executor/tests/test_household_validation.py b/projects/policyengine-simulation-executor/tests/test_household_validation.py index a82abb2df..d49d53ed0 100644 --- a/projects/policyengine-simulation-executor/tests/test_household_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_household_validation.py @@ -10,6 +10,7 @@ HouseholdWorkerIdentity, ) from policyengine_simulation_executor.household_calculation import ( + calculate_household, installed_household_bundle, ) from policyengine_simulation_executor.household_validation import ( @@ -117,3 +118,20 @@ def test_an_outcome_that_is_not_the_calculated_household_fails_the_check( check_household_validation( change_outcome(_succeeded(request)), request=request, worker=WORKER ) + + +@pytest.mark.parametrize("country", ["us", "uk"]) +def test_the_installed_models_calculate_the_validation_household(country) -> None: + # Release validation and the image check depend on this household, so the + # PR checks calculate it with the installed country models too. + request = household_validation_request(country, year=2026) + worker = WORKER.model_copy( + update={"function_name": f"calculate_household_{country}"} + ) + + outcome = calculate_household(request, worker=worker, required_country=country) + + net_income = check_household_validation( + outcome.model_dump(mode="json"), request=request, worker=worker + ) + assert net_income > 0 From 7455b2406b733fca136553477bf3feac859abae2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:49:36 +0400 Subject: [PATCH 17/19] Measure the whole country model load in its own stage A US household calculation resolved its SPM selection before the model-load span opened, and that resolution imports the US model, so a cold container's model load was measured in none of the household stages. The country model now loads first, in its own span, and the SPM selection is resolved while the household is prepared, as the Public API resolves it while preparing a calculation. A test checks both, and that a saved household's settings are resolved as saved. The module's list of differences from the Public API now says that a failure under any one policy returns one failure for the whole request. Co-Authored-By: Claude Opus 5.5 --- .../household_calculation.py | 15 ++++-- .../tests/test_household_calculation.py | 53 ++++++++++++++++++- 2 files changed, 62 insertions(+), 6 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py index 0184783f4..576944a2d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py @@ -25,7 +25,11 @@ - An explicit null ``axes`` means no axes, because the contract drops it; the Public API treats the key's presence as a household with axes. - Every policy's parameters are checked before any policy is calculated, so - an unknown parameter in a reform fails the whole request. + an unknown parameter in a reform fails the whole request. Likewise any + failure under one policy, such as a reform value that cannot be converted + or an SPM error only the reform causes, returns one failure for the whole + request, without the policies already calculated; the Public API's + separate request for each policy would still return the others. - Its messages are the request routes' (``/calculate`` and ``/calculate-full``). The stored-household route words a household it cannot read, and any other failure except an SPM error, with the household @@ -610,17 +614,20 @@ def _calculate_output( ) -> HouseholdCalculationOutput: """Calculate every policy of ``request``, raising a typed failure. - Checks run in the Public API's order: the bundle, the SPM selection, the - household's and policies' names, then each policy. + After the bundle check and loading the country model, checks run in the + Public API's order: the SPM selection, the household's and policies' + names, then each policy. The model loads first so its span measures the + whole load, which for the US otherwise happens while the SPM selection is + resolved. """ if request.bundle != bundle: raise _bundle_mismatch_failure(bundle) - spm_selection = _resolved_spm_selection(request) country = {"country": request.country} runtime = dependencies.runtime with _stage_scope(runtime, Stage.V2_HOUSEHOLD_COUNTRY_MODEL_LOAD, country): model = dependencies.model_loader(request.country) with _stage_scope(runtime, Stage.V2_HOUSEHOLD_INPUT_NORMALIZATION, country): + spm_selection = _resolved_spm_selection(request) prepared = _prepared_household( request, model, diff --git a/projects/policyengine-simulation-executor/tests/test_household_calculation.py b/projects/policyengine-simulation-executor/tests/test_household_calculation.py index 1abc55163..bc9a667c3 100644 --- a/projects/policyengine-simulation-executor/tests/test_household_calculation.py +++ b/projects/policyengine-simulation-executor/tests/test_household_calculation.py @@ -61,15 +61,20 @@ def _request(country: str, household: dict, **overrides) -> HouseholdExecutionRe class RecordingRuntime: - """Records the spans a calculation opens, in order.""" + """Records the spans a calculation opens, in order, and which are open.""" def __init__(self) -> None: self.spans: list[tuple[str, dict]] = [] + self.open_spans: list[str] = [] @contextmanager def span(self, name: str, *, attributes: dict | None = None): self.spans.append((name, dict(attributes or {}))) - yield + self.open_spans.append(name) + try: + yield + finally: + self.open_spans.pop() # A stand-in country model: enough of a country package to exercise the @@ -434,6 +439,50 @@ def test_spm_settings_as_sent_fail_with_public_api_codes( assert outcome.error.policy_key is None +@pytest.mark.parametrize( + ("source", "saved_household"), [("request", False), ("saved_household", True)] +) +def test_spm_settings_are_resolved_while_preparing_the_household( + monkeypatch, source, saved_household +) -> None: + # The country model loads first, under its own span, so the time the US + # model takes to import is measured there and not in SPM resolution. + runtime = RecordingRuntime() + events: list[tuple[str, object]] = [] + model = _stand_in_model(_stand_in_simulation({"net_income": np.array([1.0])})) + + def load_model(country): + events.append(("model", list(runtime.open_spans))) + return model + + def resolve(country, selection, *, sent, saved_household): + events.append(("spm", list(runtime.open_spans))) + events.append(("spm options", (selection, sent, saved_household))) + + monkeypatch.setattr( + household_calculation, "resolve_household_spm_selection", resolve + ) + household = { + "people": {"you": {"age": {"2026": 30}}}, + "households": {"home": {"members": ["you"], "net_income": {"2026": None}}}, + } + + outcome = calculate_household( + _request("us", household, spm={}, spm_selection_source=source), + worker=_worker("us"), + dependencies=HouseholdCalculationDependencies( + model_loader=load_model, today=TODAY, runtime=runtime + ), + ) + + assert isinstance(outcome, HouseholdExecutionResult), outcome + assert events == [ + ("model", ["v2_household_country_model_load"]), + ("spm", ["v2_household_input_normalization"]), + ("spm options", ({}, True, saved_household)), + ] + + def test_a_receipt_that_cannot_be_certified_is_an_spm_failure(monkeypatch) -> None: def uncertifiable(simulation, expected_selection): code = "SPM_CONFIGURATION_UNAVAILABLE" From b7601405df4a76651312855be183aa3a43329bb5 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:04:43 +0400 Subject: [PATCH 18/19] Keep the shared SPM receipt model's wording and translate it for households The contract's SPM receipt model is shared with the main path, so it returns to its wording on main. The household SPM code now translates that model's year and metro-area errors into the Public API's wording when it reports a receipt it cannot certify, so household messages still match the Public API and nothing the main path runs changes. Co-Authored-By: Claude Opus 5.5 --- .../policyengine_simulation_contract/spm.py | 8 +++--- .../household_spm.py | 25 +++++++++++++++++-- 2 files changed, 27 insertions(+), 6 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py index e666da24f..c2bddcae1 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -193,10 +193,10 @@ def serialize_provenance( # noqa: ANN201 # a return annotation would become th def validate_years(cls, value: tuple[str, ...]) -> tuple[str, ...]: """Require four-digit years, sorted ascending with no duplicates.""" if any(re.fullmatch(r"[0-9]{4}", year) is None for year in value): - msg = "years must contain four-digit calendar years" + msg = "SPM provenance years must use four-digit years" raise ValueError(msg) if value != tuple(sorted(set(value))): - msg = "years must be sorted and unique" + msg = "SPM provenance years must be sorted and unique" raise ValueError(msg) return value @@ -205,10 +205,10 @@ def validate_provenance_location(self) -> "SPMProvenance": """Require ``geography_id`` for metro receipts and forbid it otherwise.""" if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): - msg = "A metro SPM receipt requires geography_id" + msg = "An SPM area receipt requires geography_id" raise ValueError(msg) elif self.geography_id is not None: - msg = "Only a metro SPM receipt accepts geography_id" + msg = "Only an SPM area receipt accepts geography_id" raise ValueError(msg) return self diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py index cf5b05817..7f7ddad6d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py @@ -51,8 +51,25 @@ _SETTINGS_UNSUPPORTED = "SPM_SETTINGS_UNSUPPORTED" +# The contract's SPM receipt model, which the main path shares, words these +# errors differently from the Public API's model. The household path reports +# the Public API's wording without changing the shared model. +_PUBLIC_API_RECEIPT_MESSAGES = { + "SPM provenance years must use four-digit years": ( + "years must contain four-digit calendar years" + ), + "SPM provenance years must be sorted and unique": "years must be sorted and unique", + "An SPM area receipt requires geography_id": ( + "A metro SPM receipt requires geography_id" + ), + "Only an SPM area receipt accepts geography_id": ( + "Only a metro SPM receipt accepts geography_id" + ), +} + + def _readable_error(error: BaseException) -> str: - """Render a validation error as a client-facing message, without its URL.""" + """Render a validation error as the Public API renders it, without its URL.""" if not isinstance(error, ValidationError): return str(error) return "; ".join( @@ -60,7 +77,7 @@ def _readable_error(error: BaseException) -> str: part for part in ( ".".join(str(item) for item in detail["loc"]), - detail["msg"].removeprefix("Value error, "), + _public_api_message(detail["msg"].removeprefix("Value error, ")), ) if part ) @@ -68,6 +85,10 @@ def _readable_error(error: BaseException) -> str: ) +def _public_api_message(message: str) -> str: + return _PUBLIC_API_RECEIPT_MESSAGES.get(message, message) + + def _current_bundle() -> object: """Return the installed PolicyEngine.py bundle manifest.""" from policyengine.bundle import get_current_bundle From 08f1301629205d8eefa86f25e0b3d726bc81e059 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 22:05:19 +0400 Subject: [PATCH 19/19] Test the manifest with released document shapes and document rollback limits The tests' older-publisher documents now carry only the fields those publishers wrote: no region datasets under the keys that predate them. Two tests join them: the fixtures hold only released fields, and a publish refuses an unreadable schema version 2 document. The migration document now says what a rollback to code released before schema version 2 costs once current code publishes again: the manifest names no household functions for versions published before the rollback until each is published again, and it loses what only the schema version 2 document held. It also says that current readers do not read a publish by code released before cache releases, so a rollback that far needs the newer keys deleted first; the region dataset readers on main already behave this way. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 10 ++++-- .../tests/test_stage12_manifest.py | 33 ++++++++++++++++--- 2 files changed, 36 insertions(+), 7 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 23b84522b..ae9be88a0 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -426,8 +426,14 @@ it, with its region datasets removed, equals the manifest under `active` before the first publish that writes it. Code released before schema version 2 or before region datasets republishes only the older keys it knows, so after a rollback the out-of-date copies stop matching and are no longer read, -and the next publish starts from what the older code wrote. Removing the older -copies needs a later change, once no deployed or rollback revision reads them. +and the next publish starts from what the older code wrote. That document has +no household functions, so the manifest names none for versions published +before the rollback until each is published again, and it lacks whatever only +`active-schema-2` held, such as a version with a household-only country. Code +released before cache releases writes only `active`, which current readers do +not read while `active-with-cache` exists; a rollback that far needs the newer +keys deleted from the Modal Dict first. Removing the older copies needs a later +change, once no deployed or rollback revision reads them. A rollback to code released before region datasets has one window: the older worker is deployed before its publish job runs, and until that publish the diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 47c28b6fd..eb3b464af 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -398,6 +398,7 @@ def _schema_1_document(*, include_cache: bool, version: str = "5.1.0") -> dict: {"household_validation_invocation_ids"} | (set() if include_cache else {"cache"}) ), + "bundle": {"countries": {"__all__": {"region_datasets"}}}, **({} if include_cache else {"cache": True}), }, ) @@ -489,16 +490,14 @@ def _validated_on(worker: V2WorkerVersion, day: int) -> V2WorkerVersion: def _write_as_a_schema_1_publisher( store: FakeStore, *workers: V2WorkerVersion, default_version: str ) -> None: - """Write the older keys as a publisher released before schema 2 writes them.""" + """Write the keys a publisher released before region datasets writes.""" manifest = V2VersionManifest( generation="older-publisher", default_version=default_version, versions={worker.bundle.policyengine_version: worker for worker in workers}, ) - store[ACTIVE_MANIFEST_KEY] = schema_1_manifest_payload(manifest, include_cache=True) - store[PRE_CACHE_MANIFEST_KEY] = schema_1_manifest_payload( - manifest, include_cache=False - ) + store[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(manifest) + store[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(manifest) def test_a_publish_keeps_what_an_older_publisher_wrote_after_it() -> None: @@ -591,6 +590,30 @@ def test_loader_retains_its_last_valid_v2_document() -> None: assert loader.load() == expected +def test_a_publish_refuses_an_unreadable_schema_2_document() -> None: + store = FakeStore() + publish_v2_manifest(store=store, worker=_worker()) + store.values[SCHEMA_2_MANIFEST_KEY] = {"schema_version": 999} + + with pytest.raises(ValueError, match="schema_version"): + publish_v2_manifest(store=store, worker=_worker("5.3.0")) + + +@pytest.mark.parametrize("include_cache", [True, False]) +def test_older_publisher_fixtures_hold_only_released_fields( + include_cache: bool, # noqa: FBT001 # pytest passes parameters by name +) -> None: + # The documents these tests store as older publishers' output carry only + # the fields those publishers wrote. + document = _schema_1_document(include_cache=include_cache) + worker = _assert_schema_1_top_level( + {**document, "versions": {"5.2.0": document["versions"]["5.1.0"]}} + )["versions"]["5.2.0"] + for country in worker["bundle"]["countries"]: + assert set(country) == _COUNTRY_BUNDLE_FIELDS_BEFORE_REGION_DATASETS + assert ("cache" in worker) is include_cache + + def test_loader_never_falls_back_to_v1_state() -> None: v1_store = FakeStore({ACTIVE_MANIFEST_KEY: {"schema_version": 1}}) v2_store = FakeStore()