Add v2 Modal household functions that calculate one household under several policies

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

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.

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"

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]

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"

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
         ),