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 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/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/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_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/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..958ba9bea --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_household_calculation.py @@ -0,0 +1,520 @@ +"""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_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"} + ) + + 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_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"}) 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 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", + )