From 7eec512d36b0a3dfff71e484bc25d0a8b773323c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:48:47 +0400 Subject: [PATCH 1/5] Port the Public API's household input preparation Add the input handling the v2 household functions need, reproducing the Public API's in-process household calculation so v2 results can be compared with it field by field: - country-model metadata read from a tax-benefit system (variables with the Public API's default values, entities and role plurals, gov parameters); - the household year and the calculate-full addition of missing yearly, monthly and eternal variables; - dropping deprecated inputs, with the same warnings; - unrecognized household, axis and policy inputs, with the same types, paths and messages; - the list of values to calculate, with and without axes. Co-Authored-By: Claude Opus 5.5 --- .../household_inputs.py | 436 ++++++++++++++++++ .../tests/test_household_inputs.py | 301 ++++++++++++ 2 files changed, 737 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..e87702cb9 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_inputs.py @@ -0,0 +1,436 @@ +"""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 warnings 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_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"}) + + +@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: Any) -> HouseholdModelMetadata: # noqa: ANN401 # a country model's TaxBenefitSystem, read by attribute only + """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() + } + parameters = frozenset( + parameter.name + for parameter in system.parameters.get_descendants() + 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: Any) -> bool | int | float | str | None: # noqa: ANN401 # a country model's Variable, read by attribute only + if variable.value_type.__name__ == "Enum": + return variable.default_value.name + if isinstance(variable.default_value, (int, float, bool)): + return variable.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. + """ + you = situation.get("people", {}).get("you", {}) + 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": + 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" or not isinstance(entity_group, dict): + continue + for entity_id, variables in entity_group.items(): + if not isinstance(variables, dict): + continue + 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] + _drop_deprecated_axes(household, warnings) + return DeprecatedInputsResult(household=household, warnings=tuple(warnings)) + + +def _drop_deprecated_axes( + household: dict[str, Any], warnings: list[DeprecatedVariableWarning] +) -> None: + axes = household.get("axes") + 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): + location = f"{group_index}][{axis_index}" + if _is_deprecated_axis(axis, location, warnings): + changed = True + continue + retained_axes.append(axis) + if retained_axes: + retained_groups.append(retained_axes) + if not changed: + return + if retained_groups: + household["axes"] = retained_groups + else: + del household["axes"] + + +def _is_deprecated_axis( + axis: object, location: str, warnings: list[DeprecatedVariableWarning] +) -> bool: + if not isinstance(axis, dict): + return False + variable_name = axis.get("name") + hint = DEPRECATED_VARIABLES.get(str(variable_name)) + if hint is None: + return False + warnings.append( + DeprecatedVariableWarning( + variable=str(variable_name), + entity_plural="axes", + entity_id=location, + hint=hint, + ) + ) + return True + + +def _input_error( + input_type: InputErrorType, + name: str, + path: str, + *, + expected_entity_plural: str | None = None, + actual_entity_plural: str | None = None, +) -> HouseholdInputError: + if input_type == "household_entity": + message = f"Unrecognized household entity group `{name}` at `{path}`." + elif input_type == "household_variable_wrong_entity": + message = ( + f"Household variable `{name}` belongs on `{expected_entity_plural}`, " + f"not `{actual_entity_plural}`, at `{path}`." + ) + elif input_type == "household_axis_variable": + message = f"Unrecognized household axis variable `{name}` at `{path}`." + elif input_type == "policy_parameter": + message = f"Unrecognized policy parameter `{name}` at `{path}`." + else: + message = f"Unrecognized household variable `{name}` at `{path}`." + return HouseholdInputError( + type=input_type, + name=name, + path=path, + message=message, + expected_entity_plural=expected_entity_plural, + actual_entity_plural=actual_entity_plural, + ) + + +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": + 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 not isinstance(entity_inputs, dict): + continue + 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, + expected_entity_plural=expected_entity.plural, + actual_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.""" + 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[tuple[str, str, str, str]]: + """List the (entity plural, entity id, 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[tuple[str, str, str, str]] = [] + for entity_plural, entities in situation.items(): + if not isinstance(entities, dict): + continue + for entity_id, variables in entities.items(): + if not isinstance(variables, dict): + continue + for variable_name, periods in variables.items(): + if not isinstance(periods, dict): + continue + if variable_names is not None and variable_name not in variable_names: + continue + computations.extend( + (entity_plural, entity_id, variable_name, str(period)) + for period, value in periods.items() + if include_provided_values or value is None + ) + return computations 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..425d5ccd0 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_household_inputs.py @@ -0,0 +1,301 @@ +"""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", exclude_none=True) 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", exclude_none=True) 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_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 f56b694fb15db087bac0ea69a590324626f87a97 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:55:50 +0400 Subject: [PATCH 2/5] Calculate households under several policies like the Public API Add the calculation the v2 household functions run. calculate_household takes one HouseholdExecutionRequest and calculates the household under each requested policy in order, so a baseline and its reform share one invocation, returning per-policy results or a typed failure. It reproduces the Public API's in-process household calculation: the same simulation construction and reform value conversion, requested values, value serialization (infinities as "Infinity"), values that cannot be calculated (null, with a warning when axes are used), warnings, and SPM selection, error codes and compact receipt. Failure codes and messages match the Public API's error responses: bundle mismatch, SPM selection, unrecognized inputs, unparseable household, and any other calculation error. Also: - register the v2_household stage plan for the household function's spans; - add HouseholdBundleProvenance.from_bundle and household_callable_name to the contract; - expose the executor's certified SPM forecast loader. Co-Authored-By: Claude Opus 5.5 --- .../household_execution.py | 36 ++ .../tests/test_household_execution.py | 55 ++ .../stages.py | 17 + .../tests/test_stages.py | 12 + .../household_calculation.py | 574 ++++++++++++++++++ .../household_spm.py | 360 +++++++++++ .../policyengine_simulation_executor/spm.py | 9 + .../tests/test_household_calculation.py | 496 +++++++++++++++ .../tests/test_household_spm.py | 254 ++++++++ 9 files changed, 1813 insertions(+) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py create mode 100644 projects/policyengine-simulation-executor/tests/test_household_calculation.py create mode 100644 projects/policyengine-simulation-executor/tests/test_household_spm.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 index 63cbb0b2a..fd99613c5 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 @@ -57,7 +57,12 @@ if TYPE_CHECKING: from pydantic import ModelWrapValidatorHandler, SerializerFunctionWrapHandler + from policyengine_simulation_contract.stage12_bundle import Stage12BundleManifest + 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 entity and value limits match the Public API's stored v2 # households; the variable limit leaves room for a complete calculated # household sent back as input, which holds every variable of each entity. @@ -441,6 +446,37 @@ class HouseholdBundleProvenance(StrictHouseholdModel): country_package_version: ContractText bundle_manifest_sha256: Sha256Digest + @classmethod + def from_bundle( + cls, + bundle: Stage12BundleManifest, + *, + bundle_manifest_sha256: str, + country: str, + ) -> HouseholdBundleProvenance: + """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 _HouseholdCalculationInputs(StrictHouseholdModel): """Fields shared by the public submission and the execution request. diff --git a/libs/policyengine-simulation-contract/tests/test_household_execution.py b/libs/policyengine-simulation-contract/tests/test_household_execution.py index dcfd0806d..ba4a0e5e2 100644 --- a/libs/policyengine-simulation-contract/tests/test_household_execution.py +++ b/libs/policyengine-simulation-contract/tests/test_household_execution.py @@ -26,6 +26,7 @@ RunningHouseholdCalculation, SubmitHouseholdCalculationRequest, SucceededHouseholdCalculation, + household_callable_name, parse_household_calculation_poll_response, parse_household_execution_outcome, ) @@ -445,3 +446,57 @@ def test_poll_responses_are_discriminated_by_status() -> None: parse_household_calculation_poll_response( {"status": "succeeded", "job_id": str(JOB_ID)} ) + + +def test_bundle_provenance_is_read_from_a_normalized_bundle() -> None: + from policyengine_simulation_contract.stage12_bundle import ( + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, + ) + + 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" + ) + assert household_callable_name("uk") == "calculate_household_uk" 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" 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..c52ae4fce --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_calculation.py @@ -0,0 +1,574 @@ +"""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 and country module): the same input +preparation, simulation construction, reform value conversion, requested +values, value serialization, 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. +""" + +from __future__ import annotations + +import copy +import importlib +import inspect +import json +import logging +import math +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass +from datetime import UTC, date, datetime +from functools import lru_cache +from typing import TYPE_CHECKING, Any, Protocol + +from policyengine_core.enums import Enum +from policyengine_core.errors import SituationParsingError +from policyengine_core.parameters import get_parameter +from policyengine_core.periods import instant +from policyengine_simulation_contract.household_execution import ( + CalculatedHousehold, + HouseholdBundleProvenance, + HouseholdCalculationError, + HouseholdCalculationOutput, + HouseholdExecutionFailure, + HouseholdExecutionResult, + HouseholdPolicyResult, +) +from policyengine_simulation_contract.spm import SPMInputError +from policyengine_simulation_observability.stages import V2_HOUSEHOLD_STAGES, Stage + +from policyengine_simulation_executor.household_inputs import ( + HouseholdModelMetadata, + 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 types import ModuleType + + from policyengine_simulation_contract.household_execution import ( + HouseholdErrorCode, + HouseholdExecutionRequest, + HouseholdPolicy, + HouseholdWorkerIdentity, + ) + +logger = logging.getLogger(__name__) + + +class SpanFactory(Protocol): + """Opens a named tracing span with attributes, as the observability runtime does.""" + + def __call__( + self, name: str, attributes: Mapping[str, object] + ) -> AbstractContextManager[object]: + """Return a context manager that measures the work inside it.""" + ... + + +def _no_span( + name: str, attributes: Mapping[str, object] +) -> AbstractContextManager[object]: # noqa: ARG001 # matches SpanFactory; records nothing + return nullcontext() + + +@dataclass(frozen=True) +class CountryHouseholdModel: + """A loaded country model package with its tax-benefit system and metadata.""" + + country: str + package: ModuleType | Any + tax_benefit_system: Any + 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. + """ + bundle = load_stage12_bundle().bundle + country_bundle = next(item for item in bundle.countries if item.country == country) + package = importlib.import_module( + country_bundle.country_package_name.replace("-", "_") + ) + system = package.CountryTaxBenefitSystem() + return CountryHouseholdModel( + country=country, + package=package, + 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, + ) + + +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 _error( + code: HouseholdErrorCode, + message: str, + **details: object, +) -> _CalculationFailedError: + return _CalculationFailedError( + HouseholdCalculationError.model_validate( + {"code": code, "message": message, **details} + ) + ) + + +def _spm_failure( + error: BaseException, *, policy_key: str | None = None +) -> _CalculationFailedError | None: + detail = household_spm_error_detail(error) + if detail is None: + return None + return _error( + "invalid_spm_selection", + detail.message, + spm_error=detail.model_dump(), + policy_key=policy_key, + ) + + +def _serialize_float(value: object) -> float | str: + """Serialize a number as the Public API does, naming non-finite values. + + Infinities become "Infinity" and "-Infinity" as in the Public API; an + undefined number becomes "NaN", because the contract carries only finite + numbers. + """ + 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: + if isinstance(value, (date, datetime)): + return value.isoformat() + return value + + +def _serialize_axis_result( + result: Any, # noqa: ANN401 # a NumPy or EnumArray result from the country model + variable: Any, # noqa: ANN401 # a country model's Variable + entity_index: int, + count_entities: int, +) -> list[object]: + if variable.value_type == Enum: + result = result.decode_to_str() + entity_values: list[object] = ( + result.reshape((-1, count_entities)).T[entity_index].tolist() + ) + if variable.value_type is float: + return [_serialize_float(value) for value in entity_values] + if variable.value_type is str: + return [str(value) for value in entity_values] + return [_serialize_scalar(value) for value in entity_values] + + +def _serialize_entity_result( + result: Any, # noqa: ANN401 # a NumPy or EnumArray result from the country model + variable: Any, # noqa: ANN401 # a country model's Variable + entity_index: int, +) -> object: + if variable.value_type == Enum: + return result.decode()[entity_index].name + if variable.value_type is float: + return _serialize_float(result[entity_index]) + if variable.value_type is str: + return str(result[entity_index]) + return _serialize_scalar(result.tolist()[entity_index]) + + +def _normalize_reform_values( + reform: Mapping[str, Mapping[str, object]], + system: Any, # noqa: ANN401 # a country model's 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. + """ + normalized: dict[str, dict[str, object]] = {} + for parameter_name, parameter_updates in reform.items(): + parameter = get_parameter(system.parameters, parameter_name) + normalized[parameter_name] = {} + for time_period, value in parameter_updates.items(): + node_type = type(parameter.values_list[-1].value) + if node_type is int: + node_type = float + converted: Any = value + try: + converted = float(converted) + except (TypeError, ValueError, OverflowError): + pass + normalized[parameter_name][time_period] = node_type(converted) + return normalized + + +def _apply_reform_to_system( + system: Any, # noqa: ANN401 # a country model's TaxBenefitSystem + reform: Mapping[str, Mapping[str, object]], +) -> None: + for parameter_name, parameter_updates in reform.items(): + parameter = get_parameter(system.parameters, parameter_name) + for time_period, value in parameter_updates.items(): + start_instant, end_instant = time_period.split(".") + parameter.update( + start=instant(start_instant), + stop=instant(end_instant), + value=value, + ) + + +def _create_simulation( + model: CountryHouseholdModel, + situation: dict[str, Any], + reform: Mapping[str, Mapping[str, object]], + spm_selection: Mapping[str, object] | None, +) -> tuple[Any, Any]: + """Build the country model's simulation of ``situation`` under ``reform``. + + A country ``Simulation`` that takes a tax-benefit system gets a reformed + copy of the model's system; any other gets the converted reform values. + """ + normalized_reform = None + if reform: + system = model.tax_benefit_system.clone() + normalized_reform = _normalize_reform_values(reform, system) + else: + system = model.tax_benefit_system + simulation_type = model.package.Simulation + spm_options = {"spm": dict(spm_selection)} if spm_selection is not None else {} + if "tax_benefit_system" in inspect.signature(simulation_type).parameters: + if normalized_reform: + _apply_reform_to_system(system, normalized_reform) + simulation = simulation_type( + tax_benefit_system=system, situation=situation, **spm_options + ) + return simulation, simulation.tax_benefit_system + simulation_options: dict[str, object] = {"situation": situation, **spm_options} + if normalized_reform: + simulation_options["reform"] = normalized_reform + simulation = simulation_type(**simulation_options) + return simulation, simulation.tax_benefit_system + + +def _entity_position(household: Mapping[str, Any], plural: str, entity_id: str) -> int: + return list(household[plural]).index(entity_id) + + +@dataclass +class _PolicyCalculation: + """One policy's simulation and the household its values are written into.""" + + simulation: Any + system: Any + household: dict[str, Any] + has_axes: bool + spm_requested: bool + warnings: list[str] + + def calculate_all(self) -> None: + """Calculate every requested value, as the Public API's country module does.""" + computations = requested_computations( + self.household, + include_provided_values=self.has_axes, + variable_names=set(self.system.variables) if self.has_axes else None, + ) + for entity_plural, entity_id, variable_name, period in computations: + population = self.simulation.get_population(entity_plural) + try: + self._calculate_value( + population, entity_plural, entity_id, variable_name, period + ) + except Exception as error: # noqa: BLE001 # any model error leaves this value uncalculated, as in the Public API + 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( + population, entity_plural, entity_id, variable_name, period, error + ) + + def _calculate_value( + self, + population: Any, # noqa: ANN401 # a country model's Population + entity_plural: str, + entity_id: str, + variable_name: str, + period: str, + ) -> None: + variable = self.system.get_variable(variable_name) + result = self.simulation.calculate(variable_name, period) + cells = self.household[entity_plural][entity_id][variable_name] + if self.has_axes: + cells[period] = _serialize_axis_result( + result, + variable, + _entity_position(self.household, entity_plural, entity_id), + len(self.household[entity_plural]), + ) + return + cells[period] = _serialize_entity_result( + result, variable, population.get_index(entity_id) + ) + + def _record_failure( # noqa: PLR0913 # the Public API's failure record takes each coordinate + self, + population: Any, # noqa: ANN401 # a country model's Population + entity_plural: str, + entity_id: str, + variable_name: str, + period: str, + error: Exception, + ) -> None: + cells = self.household[entity_plural][entity_id][variable_name] + if self.has_axes: + 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) + return + cells[period] = None + logger.warning( + "Unable to calculate %s for %s in %s: %s", + variable_name, + entity_id, + period, + error, + ) + + +def _calculate_policy( # noqa: PLR0913 # each input of one policy's calculation is distinct + model: CountryHouseholdModel, + situation: Mapping[str, Any], + policy: HouseholdPolicy, + *, + spm_selection: Mapping[str, object] | None, + spm_requested: bool, + input_warnings: tuple[str, ...], +) -> HouseholdPolicyResult: + """Calculate the household under one policy, raising a typed failure.""" + policy_key = policy.policy_key + try: + simulation, system = _create_simulation( + model, copy.deepcopy(dict(situation)), policy.parameters, spm_selection + ) + calculation = _PolicyCalculation( + simulation=simulation, + system=system, + household=json.loads(json.dumps(situation)), + has_axes="axes" in situation, + spm_requested=spm_requested, + warnings=[], + ) + calculation.calculate_all() + spm_provenance = household_spm_provenance(simulation, spm_selection) + except SituationParsingError as error: + failure = _spm_failure(error, policy_key=policy_key) or _error( + "invalid_household", + f"Invalid household payload: {error}", + policy_key=policy_key, + ) + raise failure from error + except Exception as error: + failure = _spm_failure(error, policy_key=policy_key) + if failure is None: + logger.exception("Error calculating household under policy %s", policy_key) + failure = _error( + "calculation_failed", + f"Error calculating household under policy: {error}", + policy_key=policy_key, + ) + raise failure from error + return HouseholdPolicyResult( + policy_key=policy_key, + household=CalculatedHousehold.model_validate(calculation.household), + warnings=(*input_warnings, *calculation.warnings), + spm_provenance=spm_provenance, + ) + + +def _prepared_situation( + request: HouseholdExecutionRequest, + model: CountryHouseholdModel, + *, + today: date, +) -> tuple[dict[str, Any], tuple[str, ...]]: + """Prepare the household as the Public API does; raise for unknown names.""" + 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: + failure = _error( + "invalid_inputs", + unrecognized_inputs_message(input_errors), + input_errors=[error.model_dump() for error in input_errors], + policy_key=policy.policy_key if policy_errors else None, + ) + raise failure + return deprecated.household, tuple( + warning.message for warning in deprecated.warnings + ) + + +def _calculate_output( + request: HouseholdExecutionRequest, + bundle: HouseholdBundleProvenance, + *, + model_loader: Callable[[str], CountryHouseholdModel], + today: date, + span: SpanFactory, +) -> HouseholdCalculationOutput: + """Calculate every policy of ``request``, raising a typed failure.""" + if request.bundle != bundle: + mismatch = _error( + "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", + ) + raise mismatch + try: + spm_selection = resolve_household_spm_selection( + request.country, + request.spm, + saved_household=request.spm_selection_source == "saved_household", + ) + except SPMInputError as error: + failure = _error( + "invalid_spm_selection", + error.message, + spm_error={"code": error.code, "message": error.message}, + ) + raise failure from error + country_attributes = {"country": request.country} + with span( + V2_HOUSEHOLD_STAGES.name(Stage.V2_HOUSEHOLD_COUNTRY_MODEL_LOAD), + country_attributes, + ): + model = model_loader(request.country) + with span( + V2_HOUSEHOLD_STAGES.name(Stage.V2_HOUSEHOLD_INPUT_NORMALIZATION), + country_attributes, + ): + situation, input_warnings = _prepared_situation(request, model, today=today) + results = [] + for policy in request.policies: + with span( + V2_HOUSEHOLD_STAGES.name(Stage.V2_HOUSEHOLD_CALCULATION), + {**country_attributes, "policy_key": policy.policy_key}, + ): + results.append( + _calculate_policy( + model, + situation, + policy, + spm_selection=spm_selection, + spm_requested=request.spm_requested, + input_warnings=input_warnings, + ) + ) + return HouseholdCalculationOutput( + country=request.country, results=tuple(results), bundle=bundle + ) + + +def calculate_household( # noqa: PLR0913 # the seams tests and the Modal function replace are keyword-only + request: HouseholdExecutionRequest, + *, + worker: HouseholdWorkerIdentity, + installed_bundle: HouseholdBundleProvenance | None = None, + model_loader: Callable[[str], CountryHouseholdModel] = load_country_household_model, + today: date | None = None, + span: SpanFactory = _no_span, +) -> 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. Checks run + in the Public API's order: SPM selection, inputs, then each policy. + ``today`` sets the default household year (the current UTC date when + omitted). + """ + bundle = installed_bundle or installed_household_bundle(request.country) + try: + output = _calculate_output( + request, + bundle, + model_loader=model_loader, + today=today or datetime.now(UTC).date(), + span=span, + ) + 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/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..22c2df058 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/household_spm.py @@ -0,0 +1,360 @@ +"""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 + +from policyengine_simulation_contract.spm import ( + SPMErrorDetail, + SPMInputError, + SPMProvenance, + SPMRuntimeVersions, + SPMSelection, + build_spm_provenance, +) +from pydantic import ValidationError + +if TYPE_CHECKING: + 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" + + +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 _bundle_spm_configuration() -> object: + """Return the installed bundle's SPM measurement configuration, if any.""" + from policyengine.bundle import get_current_bundle + + measurements = get_current_bundle().get("measurements") + return measurements.get("spm") if isinstance(measurements, Mapping) else None + + +def _simulation_supports_spm(simulation_type: Any) -> bool: # noqa: ANN401 # a country model's Simulation class, read by attribute only + """Whether a country ``Simulation`` constructor implements canonical SPM.""" + try: + parameters = inspect.signature(simulation_type).parameters + except (TypeError, ValueError): + parameters = {} + 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 artifact 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 _unsupported_country(selection: SPMSelection | None) -> None: + if selection is not None: + error_code = "SPM_SETTINGS_UNSUPPORTED" + raise SPMInputError(error_code, "SPM settings are only available for the US") + + +def _merged_selection(chosen: SPMSelection, defaults: SPMSelection) -> dict[str, Any]: + 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 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 an + artifact hash the bundle no longer carries is reported. + """ + if country != "us": + _unsupported_country(selection) + return None + chosen = selection if selection is not None else SPMSelection() + configured = _bundle_spm_configuration() + if not isinstance(configured, Mapping): + if _us_model_supports_spm(): + # A model that can run canonical SPM must never run it uncertified. + error_code = "SPM_CONFIGURATION_UNAVAILABLE" + raise SPMInputError( + error_code, + "The installed bundle has no certified SPM measurement configuration", + ) + if selection is None: + return None + error_code = "SPM_SETTINGS_UNSUPPORTED" + raise SPMInputError( + error_code, + "This US model bundle does not support canonical SPM settings", + ) + try: + defaults, forecast = _certified_defaults(configured) + except (AttributeError, ImportError, ValueError, TypeError, OSError) as error: + error_code = "SPM_CONFIGURATION_UNAVAILABLE" + raise SPMInputError(error_code, _readable_error(error)) from error + if chosen.forecast_content_sha256 not in (None, defaults.forecast_content_sha256): + if saved_household: + error_code = "SPM_CONFIGURATION_UNAVAILABLE" + raise SPMInputError( + error_code, + "This deployment does not have the SPM artifact this household's " + "saved selection was measured against", + ) + error_code = "SPM_SETTINGS_INVALID" + raise SPMInputError( + error_code, + "SPM selection does not match this bundle's artifact hash", + ) + try: + resolved = _validated_selection(_merged_selection(chosen, defaults), forecast) + except ValueError as error: + error_code = "SPM_SETTINGS_INVALID" + raise SPMInputError(error_code, _readable_error(error)) from error + if resolved.geography_kind == "metro": + _validate_metro_selection(forecast, resolved) + return resolved.model_dump(mode="json") + + +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``. + """ + 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 _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.get("years") + if not isinstance(source_years, Mapping): + msg = "Country SPM receipt requires a years mapping" + raise ValueError(msg) # noqa: TRY004 # the caller maps only ValueError to SPM_CONFIGURATION_UNAVAILABLE + if any(not isinstance(year, str) for year in source_years): + msg = "Country SPM receipt years must be strings" + raise ValueError(msg) + source_versions = receipt.get("runtime_versions") + if not isinstance(source_versions, Mapping): + msg = "Country SPM receipt requires runtime_versions" + raise ValueError(msg) # noqa: TRY004 # the caller maps only ValueError to SPM_CONFIGURATION_UNAVAILABLE + 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(source_versions), + ) + + +def _country_receipt(simulation: Any) -> SPMProvenance: # noqa: ANN401 # a country model's Simulation, read by attribute only + """Reduce the simulation's receipt; raise ValueError if it cannot be certified.""" + selection = _complete_settings(simulation.spm_config) + if selection is None: + msg = "Country SPM configuration is incomplete" + raise ValueError(msg) + source = simulation.spm_provenance() + if not isinstance(source, Mapping): + msg = "Country SPM receipt must be an object" + raise ValueError(msg) # noqa: TRY004 # the caller maps only ValueError to SPM_CONFIGURATION_UNAVAILABLE + 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 ValueError as error: + error_code = "SPM_CONFIGURATION_UNAVAILABLE" + raise SPMInputError( + error_code, + "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 0f84c3c77..536d01c4c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -33,6 +33,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_calculation.py b/projects/policyengine-simulation-executor/tests/test_household_calculation.py new file mode 100644 index 000000000..c976f10e5 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_household_calculation.py @@ -0,0 +1,496 @@ +"""Tests for v2 household calculations, ported from the Public API's behaviour.""" + +from __future__ import annotations + +from contextlib import contextmanager +from datetime import date +from types import SimpleNamespace +from uuid import UUID + +import numpy as np +from policyengine_simulation_contract.household_execution import ( + HouseholdBundleProvenance, + HouseholdExecutionFailure, + HouseholdExecutionRequest, + HouseholdExecutionResult, + HouseholdWorkerIdentity, +) +from policyengine_simulation_executor import household_calculation +from policyengine_simulation_executor.household_calculation import ( + CountryHouseholdModel, + calculate_household, + installed_household_bundle, +) +from policyengine_simulation_executor.household_inputs import ( + HouseholdEntity, + HouseholdModelMetadata, + HouseholdVariable, +) + +CALCULATION_ID = UUID("00000000-0000-4000-8000-0000000000c1") +TODAY = date(2026, 6, 1) +WORKER = HouseholdWorkerIdentity( + application_name="policyengine-simulation-v2-py6-7-5", + function_name="calculate_household_us", +) + + +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 RecordingSpans: + def __init__(self) -> None: + self.spans: list[tuple[str, dict]] = [] + + @contextmanager + def __call__(self, name: str, attributes: dict): + self.spans.append((name, dict(attributes))) + yield + + +# Fake country model: enough of a country package to exercise failure paths +# without importing a real one. + + +class FakePopulation: + 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 FakeSystem: + 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]) + + +def _fake_model(simulation_type) -> CountryHouseholdModel: + value_types = {"age": int, "net_income": float, "spm_resources": float} + metadata = HouseholdModelMetadata( + variables={ + "age": HouseholdVariable( + "age", "person", "year", is_input=True, default_value=40 + ), + "net_income": HouseholdVariable( + "net_income", "household", "year", is_input=False, default_value=0.0 + ), + "spm_resources": HouseholdVariable( + "spm_resources", + "household", + "year", + is_input=False, + default_value=0.0, + ), + }, + entities={ + "person": HouseholdEntity("person", "people", frozenset()), + "household": HouseholdEntity( + "household", "households", frozenset({"members"}) + ), + }, + parameters=frozenset(), + ) + system = FakeSystem(value_types) + simulation_type.tax_benefit_system = system + return CountryHouseholdModel( + country="uk", + package=SimpleNamespace(Simulation=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 _fake_simulation(outcomes: dict[str, object], *, axis_points: int = 1): + class FakeSimulation: + tax_benefit_system: FakeSystem + + def __init__(self, *, situation: dict, **_: object) -> None: + self.situation = situation + people = list(situation["people"]) + households = list(situation.get("households", {})) + self.populations = { + "people": FakePopulation(people, len(people) * axis_points), + "households": FakePopulation(households, len(households) * axis_points), + } + + def get_population(self, plural: str) -> FakePopulation: + 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 FakeSimulation + + +FAKE_HOUSEHOLD = { + "people": {"you": {"age": {"2026": 30}}}, + "households": { + "home": { + "members": ["you"], + "net_income": {"2026": None}, + "spm_resources": {"2026": None}, + } + }, +} + + +def _calculate_fake(simulation_type, household=None, *, country="uk", **overrides): + model = _fake_model(simulation_type) + return calculate_household( + _request(country, household or FAKE_HOUSEHOLD, **overrides), + worker=WORKER, + model_loader=lambda country: model, + today=TODAY, + ) + + +def test_a_value_that_cannot_be_calculated_is_null_without_a_warning() -> None: + outcome = _calculate_fake( + _fake_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 == () + + +def test_with_axes_a_failed_value_is_one_null_per_point_with_a_warning() -> None: + household = { + **FAKE_HOUSEHOLD, + "axes": [[{"name": "age", "count": 3, "min": 0, "max": 2}]], + } + + outcome = _calculate_fake( + _fake_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_an_spm_input_error_fails_only_a_chosen_measurement() -> None: + error = CountryInputError("SPM_GEOGRAPHY_REQUIRED", "SPM needs a county") + outcomes = {"net_income": np.array([1.0]), "spm_resources": error} + + inherited = _calculate_fake(_fake_simulation(outcomes)) + chosen = _calculate_fake(_fake_simulation(outcomes), country="us", spm={}) + + assert isinstance(inherited, HouseholdExecutionResult) + assert inherited.output.results[0].household.to_situation()["households"]["home"][ + "spm_resources" + ] == {"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 chosen.error.message == "SPM needs a county" + + +def test_simulation_build_failures_use_public_api_messages(monkeypatch) -> None: + from policyengine_core.errors import SituationParsingError + + def parsing_failure(**_: object): + raise SituationParsingError(["people"], "unknown person `ghost`") + + def runtime_failure(**_: object): + msg = "model exploded" + raise RuntimeError(msg) + + invalid = _calculate_fake(parsing_failure) + failed = _calculate_fake(runtime_failure) + + assert isinstance(invalid, HouseholdExecutionFailure) + assert invalid.error.code == "invalid_household" + assert invalid.error.message.startswith("Invalid household payload: ") + assert invalid.error.policy_key == "baseline" + assert isinstance(failed, HouseholdExecutionFailure) + assert failed.error.code == "calculation_failed" + assert failed.error.message == ( + "Error calculating household under policy: model exploded" + ) + assert failed.worker == WORKER + + +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_fake( + _fake_simulation({}), bundle=bundle.model_dump(mode="json") + ) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert outcome.error.code == "bundle_mismatch" + + +def test_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_run_in_one_request() -> None: + spans = RecordingSpans() + request = _request( + "us", + US_FAMILY, + policies=[ + {"policy_key": "baseline", "parameters": {}}, + {"policy_key": "reform", "parameters": CTC_REFORM}, + ], + ) + + outcome = calculate_household(request, worker=WORKER, today=TODAY, span=spans) + + 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 baseline_situation.keys() == US_FAMILY.keys() + 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 reform_ctc > baseline_ctc + # A US household inherits the certified SPM selection and carries its receipt. + assert baseline.spm_provenance is not None + assert baseline.spm_provenance.geography_kind == "county" + assert [name for name, _ in spans.spans] == [ + "v2_household_country_model_load", + "v2_household_input_normalization", + "v2_household_calculation", + "v2_household_calculation", + ] + assert spans.spans[3][1] == {"country": "us", "policy_key": "reform"} + + +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, 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 + + +def test_calculate_full_adds_every_missing_yearly_variable() -> None: + outcome = calculate_household( + _request("uk", UK_SINGLE, add_missing_variables=True), + worker=WORKER, + today=TODAY, + ) + + assert isinstance(outcome, HouseholdExecutionResult), outcome + you = outcome.output.results[0].household.to_situation()["people"]["you"] + assert "income_tax" in you + assert isinstance(you["income_tax"]["2026"], float) + + +def test_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, today=TODAY) + + 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_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, 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] + assert ( + "uc_has_limited_capability_for_work" + not in (result.household.to_situation()["people"]["you"]) + ) + + +def test_spm_settings_for_the_uk_are_refused() -> None: + outcome = calculate_household( + _request("uk", UK_SINGLE, spm={}), worker=WORKER, 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_a_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, 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 = household_calculation.load_country_household_model("uk") + + assert household_calculation.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 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..cde8b01ed --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_household_spm.py @@ -0,0 +1,254 @@ +"""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, *, scenario, kind, geoid, as_of): + if geoid not in self.metro_areas: + msg = f"unknown area {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_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 From 6675b00061ce7ac0e00e0d665cfe07be8e275b06 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:01:44 +0400 Subject: [PATCH 3/5] Add v2 Modal household functions for every bundle country The v2 Modal application now defines calculate_household_{country} for each country in the installed PolicyEngine.py bundle (today calculate_household_us and calculate_household_uk). Each takes one HouseholdExecutionRequest and returns the per-policy result or typed failure from calculate_household. - The household image holds the frozen simulation dependency group and no datasets or cache, and the functions get no Modal secrets: household calculations read no private data, Cloud Storage or database. - Sizing: 2 CPUs, 4 GiB reserved with a 16 GiB limit for large households calculated with every variable at many axis points, a 900-second timeout, no retries and at most 20 containers. - Deployment validation calls each household function with a one-adult current-law household and requires a successful calculation with this deployment's bundle; the release record keeps the invocation IDs. - The manifest publisher records a country's household function only when the release record shows it was validated. Nothing calls the household functions yet. Co-Authored-By: Claude Opus 5.5 --- .../modal/utils/update_v2_version_manifest.py | 13 +- .../modal/utils/validate_v2_worker_release.py | 122 ++++++++++++- .../src/modal/v2_app.py | 164 +++++++++++++++--- .../tests/test_stage12_modal_app.py | 122 +++++++++++++ .../tests/test_update_v2_version_manifest.py | 21 ++- .../tests/test_validate_v2_worker_release.py | 118 ++++++++++++- 6 files changed, 521 insertions(+), 39 deletions(-) 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 8157504cd..f75b59674 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,9 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: """Build the installed bundle's v2 worker entry from a validation record. Every bundle country runs report simulations, so each is recorded as - society-wide capable. Raises ValueError if the record is invalid or the - entry carries a credential-like key. + society-wide capable. A country's household function is recorded when the + record shows it was validated. 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 +42,11 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: country=country.country, society_wide_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 aff32b395..2ea408a29 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 small current-law household. +The release record names every invocation, so the manifest records only +functions that ran successfully in the deployed application. +""" from __future__ import annotations @@ -10,12 +16,21 @@ from uuid import uuid4 import modal +from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionRequest, + HouseholdExecutionResult, + household_callable_name, + parse_household_execution_outcome, +) from policyengine_simulation_contract.stage12_manifest import ( ManifestText, Stage12CacheRelease, V2WorkerValidation, v2_application_name, ) +from policyengine_simulation_executor.household_calculation import ( + installed_household_bundle, +) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle if TYPE_CHECKING: @@ -27,13 +42,14 @@ def _spawn_validation( application_name: str, function_name: str, environment: str, + payload: dict | None = None, ) -> tuple[Any, str]: function = modal.Function.from_name( application_name, function_name, environment_name=environment, ) - call = function.spawn() + call = function.spawn() if payload is None else function.spawn(payload) 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,102 @@ def _spawn_validation( return result, invocation_id +def household_validation_request( + country: str, *, year: int +) -> HouseholdExecutionRequest: + """Return a one-adult, current-law household request for ``country``.""" + group_plural = "spm_units" if country == "us" else "benunits" + household: dict[str, Any] = { + "people": { + "you": { + "age": {str(year): 40}, + "employment_income": {str(year): 30_000}, + } + }, + group_plural: {"your unit": {"members": ["you"]}}, + "households": { + "your household": { + "members": ["you"], + "household_net_income": {str(year): 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, + application_name: str, + function_name: str, +) -> None: + """Raise RuntimeError unless the household function calculated the request. + + It must succeed with this deployment's bundle, name this application and + function, and calculate a numeric household net income. + """ + parsed = parse_household_execution_outcome(outcome) + if not isinstance(parsed, HouseholdExecutionResult): + msg = f"{function_name} did not calculate the validation household" + raise RuntimeError(msg) # noqa: TRY004 # a failed calculation is a failed release, not a type error + if parsed.output.bundle != request.bundle: + msg = f"{function_name} calculated with another bundle" + raise RuntimeError(msg) + if (parsed.worker.application_name, parsed.worker.function_name) != ( + application_name, + function_name, + ): + msg = f"{function_name} reported another function" + raise RuntimeError(msg) + situation = parsed.output.results[0].household.to_situation() + net_income = situation["households"]["your household"]["household_net_income"] + if not all(isinstance(value, float) for value in net_income.values()): + msg = f"{function_name} returned no household net income" + raise RuntimeError(msg) + + +def _validate_household_functions( + *, application_name: str, environment: str, countries: tuple[CountryId, ...] +) -> dict[CountryId, ManifestText]: + invocation_ids: dict[CountryId, ManifestText] = {} + year = datetime.now(timezone.utc).year + for country in countries: + function_name = household_callable_name(country) + request = household_validation_request(country, year=year) + outcome, invocation_id = _spawn_validation( + application_name=application_name, + function_name=function_name, + environment=environment, + payload=request.model_dump(mode="json"), + ) + check_household_validation( + outcome, + request=request, + 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 +197,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, + countries=tuple(country.country for country in resolved.bundle.countries), + ) return V2WorkerValidation( validated=True, application_name=application_name, @@ -100,6 +209,7 @@ def validate_release( validated_at=datetime.now(timezone.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/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 7485cec2a..80134bdbd 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 one +household function for every country in the PolicyEngine.py bundle, +``calculate_household_{country}``, which calculates one household under one or +more policies. 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,26 +147,31 @@ def _country_bundle(country: CountryId) -> Stage12CountryBundle: def build_v2_image( - countries: tuple[CountryId, ...], *, include_us_cache: bool = False + countries: tuple[CountryId, ...], + *, + include_us_cache: bool = False, + include_bundle_data: bool = True, ) -> modal.Image: """Build a Stage 12 image with the bundle's certified data for ``countries``. The bundle manifest digest and per-country bundle values are baked into the image environment, which worker validation checks at run time. With - ``include_us_cache``, the image also holds the published US cache. + ``include_us_cache``, the image also holds the published US cache. Without + ``include_bundle_data``, the image holds no datasets, which household + calculations do not need. """ country_values = { country: _country_bundle(country).model_dump(mode="json") for country in countries } - image = ( - modal.Image.debian_slim(python_version="3.13") - .uv_sync( - uv_project_dir=_UV_PROJECT_DIR, - frozen=True, - extra_options="--only-group modal-simulation-image", - ) - .run_commands( + image = modal.Image.debian_slim(python_version="3.13").uv_sync( + uv_project_dir=_UV_PROJECT_DIR, + frozen=True, + extra_options="--only-group modal-simulation-image", + ) + data_environment: dict[str, str] = {} + if include_bundle_data: + image = image.run_commands( bundle_data_install_command( RESOLVED_BUNDLE.bundle.policyengine_version, countries=countries, @@ -165,23 +179,21 @@ def build_v2_image( ), secrets=[hf_secret], ) - .env( - { - **modal_image_environment(), - "POLICYENGINE_DATA_FOLDER": STAGE12_DATA_DIR, - "STAGE12_BUNDLE_MANIFEST_SHA256": ( - RESOLVED_BUNDLE.bundle_manifest_sha256 - ), - "STAGE12_BUNDLE_VALUES": json.dumps( - country_values, - sort_keys=True, - separators=(",", ":"), - ), - "STAGE12_CACHE_MANIFEST_DIGEST": os.environ.get( - "STAGE12_CACHE_MANIFEST_DIGEST", "" - ), - } - ) + data_environment["POLICYENGINE_DATA_FOLDER"] = STAGE12_DATA_DIR + image = image.env( + { + **modal_image_environment(), + **data_environment, + "STAGE12_BUNDLE_MANIFEST_SHA256": (RESOLVED_BUNDLE.bundle_manifest_sha256), + "STAGE12_BUNDLE_VALUES": json.dumps( + country_values, + sort_keys=True, + separators=(",", ":"), + ), + "STAGE12_CACHE_MANIFEST_DIGEST": os.environ.get( + "STAGE12_CACHE_MANIFEST_DIGEST", "" + ), + } ) if include_us_cache: image = image.run_function( @@ -206,6 +218,10 @@ def build_v2_image( 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_COUNTRIES: tuple[CountryId, ...] = tuple( + country.country for country in RESOLVED_BUNDLE.bundle.countries +) +household_image = build_v2_image(HOUSEHOLD_COUNTRIES, include_bundle_data=False) def _validate(country: CountryId) -> dict: @@ -481,3 +497,95 @@ def coordinate_report( return result finally: runtime.shutdown() + + +# Household function sizing. A household calculation holds one country model +# and, for each policy, one simulation of the household at every axis point. +# Each container reserves 4 GiB and may use up to 16 GiB, the memory of a +# whole Public API instance, 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.""" + 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) + with runtime.operation( + V2_HOUSEHOLD_STAGES.name(Stage.V2_HOUSEHOLD_EXECUTION), + attributes={"country": country}, + remote_context=propagated, + ): + request = HouseholdExecutionRequest.model_validate(payload) + if request.country != country: + msg = f"this function calculates {country} households only" + raise ValueError(msg) + outcome = household_calculation.calculate_household( + request, + worker=HouseholdWorkerIdentity( + application_name=APP_NAME, + function_name=household_callable_name(country), + ), + span=lambda name, attributes: runtime.span(name, attributes=attributes), + ) + return outcome.model_dump(mode="json") + finally: + runtime.shutdown() + + +def _household_function(country: CountryId) -> Callable[..., dict]: + """Return the household function for ``country``, named for Modal. + + Modal imports a deployed function by its qualified name, so the function + is given its global name and stored in this module under that name. + """ + + def calculate_household( + payload: dict, + *, + observability_context: dict | None = None, + ) -> dict: + return _run_household_calculation(country, payload, observability_context) + + function_name = household_callable_name(country) + calculate_household.__name__ = function_name + calculate_household.__qualname__ = function_name + calculate_household.__doc__ = ( + f"Calculate one {country} household under each requested policy." + ) + return calculate_household + + +for _household_country in HOUSEHOLD_COUNTRIES: + globals()[household_callable_name(_household_country)] = 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=[], + name=household_callable_name(_household_country), + )(_household_function(_household_country)) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index 6d0acb14d..d198b1618 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -217,6 +217,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 @@ -289,6 +291,126 @@ def fail(*_, **__): assert runtime.shutdown_calls == 1 +def test_household_functions_are_generated_for_every_bundle_country( + monkeypatch, +) -> None: + from policyengine_simulation_contract.household_execution import ( + household_callable_name, + ) + + module = _load(monkeypatch) + functions = {name: options for name, options in module.app.function_calls} + bundle_countries = [ + country.country for country in module.RESOLVED_BUNDLE.bundle.countries + ] + + for country in bundle_countries: + name = household_callable_name(country) + options = functions[name] + assert getattr(module, name).__qualname__ == name + assert options["name"] == name + assert options["image"] is module.household_image + assert options["cpu"] == 2.0 + # Reserve 4 GiB; allow up to 16 GiB for large households calculated + # with every variable at many axis points. + assert options["memory"] == (4096, 16384) + assert options["timeout"] == 900 + assert options["retries"] == 0 + assert options["max_containers"] == 20 + # Household calculations read no private data, Cloud Storage or + # database, so they get no Modal secrets. + assert options["secrets"] == [] + + +def test_household_image_holds_no_datasets_or_cache(monkeypatch) -> None: + module = _load(monkeypatch) + calls = module.household_image.calls + + assert not any(call[0] == "run_commands" for call in calls) + assert not any(call[0] == "run_function" for call in calls) + local_source = next(call for call in calls if call[0] == "add_local_python_source") + assert "policyengine_simulation_contract" in local_source[1] + environment = next(call for call in calls if call[0] == "env")[1] + assert environment["STAGE12_BUNDLE_MANIFEST_SHA256"] == ( + module.RESOLVED_BUNDLE.bundle_manifest_sha256 + ) + assert "POLICYENGINE_DATA_FOLDER" not in environment + + +def _household_payload(country: str) -> dict: + from policyengine_simulation_executor.household_calculation import ( + installed_household_bundle, + ) + + return { + "calculation_id": "00000000-0000-4000-8000-0000000000c1", + "country": country, + "household": {"people": {"you": {"age": {"2026": 40}}}}, + "policies": [{"policy_key": "baseline", "parameters": {}}], + "bundle": installed_household_bundle(country).model_dump(mode="json"), + } + + +@pytest.mark.parametrize("country", ["us", "uk"]) +def test_household_function_calculates_its_country_and_closes_runtime( + monkeypatch, country +) -> None: + from policyengine_simulation_contract.household_execution import ( + HouseholdExecutionFailure, + HouseholdExecutionRequest, + ) + from policyengine_simulation_executor import household_calculation + + module = _load(monkeypatch) + runtime = RecordingRuntime() + runtime.span = lambda name, *, attributes: runtime.operation( + name, attributes=attributes + ) + monkeypatch.setattr(module, "init_process_observability", lambda **_: runtime) + calls = [] + + def fake_calculate(request, *, worker, span): + calls.append((request, worker)) + with span("v2_household_calculation", {"policy_key": "baseline"}): + pass + return HouseholdExecutionFailure( + calculation_id=request.calculation_id, + error={"code": "calculation_failed", "message": "recorded"}, + worker=worker, + ) + + monkeypatch.setattr(household_calculation, "calculate_household", fake_calculate) + + result = getattr(module, f"calculate_household_{country}")( + _household_payload(country), + observability_context=_observability_context(), + ) + + assert result["status"] == "failed" + request, worker = calls[0] + assert isinstance(request, HouseholdExecutionRequest) + assert request.country == country + assert worker.application_name == module.APP_NAME + assert worker.function_name == f"calculate_household_{country}" + assert runtime.operations[0]["name"] == "v2_household_execution" + assert runtime.operations[0]["remote_context"]["observability_id"] == ( + OBSERVABILITY_ID + ) + assert runtime.operations[1]["name"] == "v2_household_calculation" + assert runtime.shutdown_calls == 1 + + +def test_household_function_refuses_another_country(monkeypatch) -> None: + module = _load(monkeypatch) + runtime = RecordingRuntime() + monkeypatch.setattr(module, "init_process_observability", lambda **_: runtime) + + with pytest.raises(ValueError, match="calculates uk households"): + module.calculate_household_uk(_household_payload("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") 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 15d254a80..a50ff7fd1 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 @@ -9,7 +9,7 @@ from src.modal.utils import update_v2_version_manifest as publisher -def _validation() -> V2WorkerValidation: +def _validation(*, household_functions: bool = False) -> V2WorkerValidation: resolved = publisher.load_stage12_bundle() application_name = publisher.v2_application_name( resolved.bundle.policyengine_version @@ -24,6 +24,11 @@ def _validation() -> V2WorkerValidation: "us": "validation-us-1", "uk": "validation-uk-1", }, + household_validation_invocation_ids=( + {"us": "household-us-1", "uk": "household-uk-1"} + if household_functions + else {} + ), ) @@ -53,6 +58,20 @@ def test_publisher_marks_every_bundle_country_society_wide_capable() -> None: ] +def test_publisher_records_the_validated_household_functions() -> None: + worker = publisher.build_worker( + _validation(household_functions=True).model_dump(mode="json") + ) + + assert [ + (country.country, country.household_callable) for country in worker.countries + ] == [ + ("us", "calculate_household_us"), + ("uk", "calculate_household_uk"), + ] + assert all(country.society_wide_capable for country in worker.countries) + + 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..30971f81c 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,14 +3,78 @@ from __future__ import annotations import pytest +from policyengine_simulation_contract.household_execution import ( + CalculatedHousehold, + HouseholdCalculationOutput, + HouseholdExecutionFailure, + HouseholdExecutionRequest, + HouseholdExecutionResult, + HouseholdPolicyResult, + HouseholdWorkerIdentity, +) +from policyengine_simulation_executor.household_calculation import ( + installed_household_bundle, +) from src.modal.utils import validate_v2_worker_release as validation +def _household_outcome( + payload: dict, + *, + application_name: str, + function_name: str, + succeeded: bool = True, + bundle_overrides: dict | None = None, +) -> dict: + request = HouseholdExecutionRequest.model_validate(payload) + worker = HouseholdWorkerIdentity( + application_name=application_name, function_name=function_name + ) + if not succeeded: + return HouseholdExecutionFailure( + calculation_id=request.calculation_id, + error={"code": "calculation_failed", "message": "broken"}, + worker=worker, + ).model_dump(mode="json") + calculated = request.household.to_situation() + for group in calculated["households"].values(): + group["household_net_income"] = { + period: 25_000.0 for period in group["household_net_income"] + } + return HouseholdExecutionResult( + calculation_id=request.calculation_id, + output=HouseholdCalculationOutput( + country=request.country, + results=( + HouseholdPolicyResult( + policy_key="baseline", + household=CalculatedHousehold.model_validate(calculated), + ), + ), + bundle=request.bundle.model_copy(update=bundle_overrides or {}), + ), + worker=worker, + ).model_dump(mode="json") + + def test_all_country_results_must_match_the_deployment_bundle(monkeypatch) -> None: invocation_ids: list[str] = [] - def fake_spawn_validation(*, application_name, function_name, environment): + def fake_spawn_validation( + *, application_name, function_name, environment, payload=None + ): country = function_name.rsplit("_", maxsplit=1)[1] + if function_name.startswith("calculate_household_"): + invocation_id = f"household-call-{country}" + invocation_ids.append(invocation_id) + return ( + _household_outcome( + payload, + application_name=application_name, + function_name=function_name, + ), + invocation_id, + ) resolved = validation.load_stage12_bundle() bundle = next( item for item in resolved.bundle.countries if item.country == country @@ -56,7 +120,16 @@ def fake_spawn_validation(*, application_name, function_name, environment): "us": "call-us", "uk": "call-uk", } - assert invocation_ids == ["call-us", "call-uk"] + assert result.household_validation_invocation_ids == { + "us": "household-call-us", + "uk": "household-call-uk", + } + assert invocation_ids == [ + "call-us", + "call-uk", + "household-call-us", + "household-call-uk", + ] assert result.cache is not None assert result.cache.manifest_sha256 == "c" * 64 @@ -80,3 +153,44 @@ def test_country_result_mismatch_stops_publication(monkeypatch) -> None: environment="staging", expected_cache_manifest_sha256="c" * 64, ) + + +def test_validation_household_requests_current_law_net_income() -> None: + request = validation.household_validation_request("us", year=2026) + + assert request.country == "us" + assert request.bundle == installed_household_bundle("us") + assert [policy.policy_key for policy in request.policies] == ["baseline"] + situation = request.household.to_situation() + assert situation["households"]["your household"]["household_net_income"] == { + "2026": None + } + + +@pytest.mark.parametrize( + ("outcome_options", "message"), + [ + ({"succeeded": False}, "did not calculate"), + ({"bundle_overrides": {"policyengine_version": "0.0.1"}}, "another bundle"), + ({"function_name": "calculate_household_other"}, "another function"), + ], +) +def test_household_function_failures_stop_publication(outcome_options, message) -> None: + application_name = validation.v2_application_name( + validation.load_stage12_bundle().bundle.policyengine_version + ) + request = validation.household_validation_request("uk", year=2026) + options = {"function_name": "calculate_household_uk", **outcome_options} + outcome = _household_outcome( + request.model_dump(mode="json"), + application_name=application_name, + **options, + ) + + with pytest.raises(RuntimeError, match=message): + validation.check_household_validation( + outcome, + request=request, + application_name=application_name, + function_name="calculate_household_uk", + ) From 08131f9a59343e1611888a14c448e6b66867096b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:20:37 +0400 Subject: [PATCH 4/5] Test that unrecognized inputs keep the Public API's order A household with an unknown person variable and an unknown axis variable, sent with `axes` last, must list the person's error first and join the message in that order, as the Public API does. The test failed before the contract serialized households in the caller's key order (the fix in the pull request below) and passes with it. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_household_calculation.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_household_calculation.py b/projects/policyengine-simulation-executor/tests/test_household_calculation.py index c976f10e5..958ba9bea 100644 --- a/projects/policyengine-simulation-executor/tests/test_household_calculation.py +++ b/projects/policyengine-simulation-executor/tests/test_household_calculation.py @@ -261,6 +261,30 @@ def runtime_failure(**_: object): assert failed.worker == WORKER +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`` last, so the axis error follows the person's. + household = { + "people": {"you": {"age": {"2026": 30}, "not_a_variable": {"2026": 1}}}, + "households": {"home": {"members": ["you"], "net_income": {"2026": None}}}, + "axes": [[{"name": "not_an_axis_variable", "count": 2, "min": 0, "max": 1}]], + } + + outcome = _calculate_fake(_fake_simulation({}), household) + + assert isinstance(outcome, HouseholdExecutionFailure) + assert [error.path for error in outcome.error.input_errors] == [ + "household.people.you.not_a_variable", + "household.axes[0][0].name", + ] + assert outcome.error.message == ( + "Unrecognized calculate input(s): Unrecognized household variable " + "`not_a_variable` at `household.people.you.not_a_variable`.; " + "Unrecognized household axis variable `not_an_axis_variable` at " + "`household.axes[0][0].name`." + ) + + def test_a_request_for_another_bundle_is_refused() -> None: bundle = installed_household_bundle("uk").model_copy( update={"policyengine_version": "0.0.1"} From a4b68728d5803bc83b9c5981ee1440c7cf6881cf Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:24:37 +0400 Subject: [PATCH 5/5] Document the v2 household functions and their release validation Describe the household functions in the v2 worker document: one per bundle country, their image without datasets, no secrets, their limits, and that deployment validation calculates a household with each one. A failed household validation stops manifest publication and the Cloud Run promotion, as any v2 worker validation failure does. Documentation only: no test can fail first. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 846698940..90bf7ee60 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -414,6 +414,35 @@ 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 one household function for each +country in the PolicyEngine.py bundle, named `calculate_household_{country}` +(currently `calculate_household_us` and `calculate_household_uk`). Each call +calculates one household, in the Public API's situation shape, under one or +more policies, and returns a `HouseholdExecutionResult` or a typed +`HouseholdExecutionFailure` from +`policyengine_simulation_contract.household_execution`. A function refuses a +request for another country or for another bundle. Nothing calls these +functions until the Simulation Entrypoint has household routes. + +- Image: the frozen `modal-simulation-image` dependency group and the source + packages, without the bundle's datasets or the Stage 12 cache. +- Secrets: none. Household calculations read no private data, Cloud Storage + or database. +- 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. + +Deployment validation calls each household function with a one-adult, +current-law household and requires a successful calculation with the deployed +bundle and a numeric household net income. The validation record keeps each +invocation ID, and the manifest publisher records a country's +`household_callable` only when the record shows that its function was +validated. A failed household validation fails v2 worker validation like any +other validation failure: it stops manifest publication and, as described +under "Release and deployment constraints", prevents promotion of the Cloud Run +revision. + ## Temporary direct runner endpoint > **Temporary Stage 12 interface:** The authenticated routes in this section