diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py index ae3b5515d..c451bc3d1 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/budget_window_state.py @@ -254,6 +254,7 @@ def build_batch_status_response( child_jobs=state.child_jobs, result=state.result, error=state.error, + errors=state.errors, resolved_app_name=state.resolved_app_name, policyengine_bundle=state.policyengine_bundle, run_id=state.run_id, diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py index 9ba41c34a..83babe3bd 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/gateway_models.py @@ -14,11 +14,16 @@ model_validator, ) +from policyengine_simulation_contract.spm import ( + SPMSelection, + SPMCapability, + SPMComparisonProvenance, + SPMErrorDetail, +) from policyengine_simulation_contract.json_types import JsonObject from policyengine_simulation_contract.macro_output import SingleYearMacroOutput from policyengine_simulation_observability.telemetry import TelemetryEnvelope - # Hard cap on request body size (bytes). SimulationOptions + telemetry + any # reform/baseline parameter tree should fit comfortably in ~256 KB. A hostile # client that tries to stream a multi-MB reform dict is rejected with 422 @@ -95,6 +100,7 @@ class GatewayRequestBase(BaseModel): """ country: str + spm: Optional[SPMSelection] = None version: Optional[str] = None policyengine_version: Optional[str] = None telemetry: TelemetryEnvelope | None = None @@ -156,6 +162,17 @@ class PolicyEngineBundle(BaseModel): policyengine_version: Optional[str] = None data_version: Optional[str] = None dataset: Optional[str] = None + spm: Optional[SPMCapability] = None + + +class SimulationErrorResponse(BaseModel): + """Existing route errors and public typed SPM input failures.""" + + status: Literal["failed"] | None = None + result: None = None + error: str | None = None + errors: list[SPMErrorDetail] | None = None + detail: str | None = None class JobSubmitResponse(BaseModel): @@ -176,6 +193,7 @@ class JobStatusResponse(BaseModel): status: Literal["running", "complete", "failed"] result: Optional[SingleYearMacroOutput] = None + errors: Optional[list[SPMErrorDetail]] = None error: Optional[str] = None resolved_app_name: Optional[str] = None policyengine_bundle: Optional[PolicyEngineBundle] = None @@ -223,6 +241,8 @@ def validate_end_year(self) -> "BudgetWindowBatchRequest": class BudgetWindowAnnualImpact(BaseModel): """Annual budget-window impact row.""" + spm_config: Optional[SPMSelection] = None + spm_provenance: Optional[SPMComparisonProvenance] = None year: str taxRevenueImpact: float federalTaxRevenueImpact: float @@ -265,6 +285,7 @@ class BatchChildJobStatus(BaseModel): "failed", "cancelled", ] + errors: Optional[list[SPMErrorDetail]] = None error: Optional[str] = None @@ -292,6 +313,7 @@ class BudgetWindowBatchStatusResponse(BaseModel): failed_years: list[str] = Field(default_factory=list) child_jobs: dict[str, BatchChildJobStatus] = Field(default_factory=dict) result: Optional[BudgetWindowResult] = None + errors: Optional[list[SPMErrorDetail]] = None error: Optional[str] = None resolved_app_name: Optional[str] = None policyengine_bundle: Optional[PolicyEngineBundle] = None @@ -323,6 +345,7 @@ class BudgetWindowBatchState(BaseModel): default_factory=dict ) result: Optional[BudgetWindowResult] = None + errors: Optional[list[SPMErrorDetail]] = None error: Optional[str] = None created_at: str updated_at: str @@ -348,6 +371,15 @@ class VersionMap(RootModel[dict[str, str]]): class VersionsResponse(BaseModel): """All supported simulation routing version maps.""" + spm_capabilities: dict[str, SPMCapability] = Field( + default_factory=dict, + description=( + "Canonical SPM capabilities keyed by exact PolicyEngine wrapper version. " + "Resolve a country-model route to its app in the routing maps, then " + "find that app's exact version in the policyengine map. App names, " + "country-model versions and latest aliases are not capability keys." + ), + ) policyengine: VersionMap us: VersionMap uk: VersionMap diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py index 229131bb8..7b6eae156 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/macro_output.py @@ -6,6 +6,7 @@ from pydantic import BaseModel, ConfigDict, RootModel +from policyengine_simulation_contract.spm import SPMSelection, SPMComparisonProvenance T = TypeVar("T") @@ -188,6 +189,9 @@ class CongressionalDistrictImpactOutput(MacroOutputModel): class SingleYearMacroOutput(MacroOutputModel): """Completed response returned by a single-year macro simulation.""" + spm_config: SPMSelection | None = None + spm_provenance: SPMComparisonProvenance | None = None + model_version: str data_version: str budget: BudgetaryImpact diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py new file mode 100644 index 000000000..bbde945da --- /dev/null +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/spm.py @@ -0,0 +1,337 @@ +"""Dependency-light public models for US SPM selection and calculation receipts.""" + +import re +from datetime import date +from typing import Any, Literal, Optional + +from pydantic import ( + BaseModel, + ConfigDict, + Field, + SerializationInfo, + SerializerFunctionWrapHandler, + field_validator, + model_serializer, + model_validator, +) + + +def _selection_schema(schema): + # An omitted option inherits the selected bundle's default. Advertising + # Python attribute defaults here makes generated clients send an explicit + # county selection even when their caller omitted geography entirely. + for field in schema.get("properties", {}).values(): + field.pop("default", None) + + +class SPMSelection(BaseModel): + """Select from the bundle's pinned artifact; national geography is explicit. + + County mode reads the household's observed ``county_fips``. A state alone + does not identify an SPM area. These settings contain no provider or path. + """ + + model_config = ConfigDict( + frozen=True, extra="forbid", json_schema_extra=_selection_schema + ) + + forecast_content_sha256: Optional[str] = Field( + default=None, pattern=r"^[0-9a-f]{64}$" + ) + scenario: Optional[str] = Field(default=None, min_length=1, pattern=r"^\S+$") + geography_kind: Literal["county", "national", "metro"] = "county" + geography_id: Optional[str] = Field(default=None, min_length=1) + county_vintage: str = Field(default="2020", pattern=r"^[0-9]{4}$") + as_of: Optional[str] = None + + @model_serializer(mode="wrap") + def serialize_selection( + self, handler: SerializerFunctionWrapHandler, info: SerializationInfo + ): + """Preserve inherited options through ordinary and nested JSON. + + Presence is the contract: an omitted option inherits the bundle + default, so an option explicitly selected as null has to stay on the + wire. ``exclude_none`` would otherwise turn a completed result's + resolved selection back into a partial request, and the poll routes + apply it to every body via ``response_model_exclude_none``. + """ + dumped = handler(self) + if info.exclude_none: + excluded = info.exclude or frozenset() + dumped = { + name: dumped.get(name) + for name in type(self).model_fields + if name in dumped + or ( + name in self.model_fields_set + and name not in excluded + and getattr(self, name) is None + ) + } + return { + name: value + for name, value in dumped.items() + if name in self.model_fields_set + } + + @field_validator("as_of") + @classmethod + def validate_as_of(cls, value): + if value is not None: + if date.fromisoformat(value).isoformat() != value: + raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)") + return value + + @model_validator(mode="after") + def validate_location(self): + # Request options inherit bundle defaults. Validate coupled options + # once both are explicit (including after defaults are resolved). + if not {"geography_kind", "geography_id"} <= self.model_fields_set: + return self + if self.geography_kind == "metro": + if not self.geography_id or not self.geography_id.strip(): + raise ValueError("An SPM area selection requires geography_id") + elif self.geography_id is not None: + raise ValueError("Only an SPM area selection accepts geography_id") + return self + + +class SPMProvenance(BaseModel): + """Detached calculation receipt; data certification is a separate claim.""" + + model_config = ConfigDict(extra="forbid") + + forecast_id: str + forecast_sha256: str + scenario: str + geography_kind: str + runtime_versions: dict[str, Optional[str]] + years: dict[str, dict[str, Any]] + geographies: list[dict[str, Any]] + composition_method: str + storage_method: str + + +SPM_CONTRACT_VERSION = "canonical-spm-v1" +SPM_ERROR_CODES = frozenset( + { + "SPM_GEOGRAPHY_REQUIRED", + "SPM_GEOGRAPHY_UNAVAILABLE", + "SPM_COMPOSITION_REQUIRED", + "SPM_YEAR_UNAVAILABLE", + "SPM_SCENARIO_UNAVAILABLE", + "SPM_CONFIGURATION_UNAVAILABLE", + "SPM_SETTINGS_INVALID", + } +) + + +class SPMErrorDetail(BaseModel): + code: str + message: str + + +class SPMInputError(ValueError): + """Transportable error shared by the control plane and worker.""" + + def __init__(self, code: str, message: str): + self.code = code + self.message = message + super().__init__(code, message) + + def __str__(self): + return self.message + + def to_dict(self): + return {"code": self.code, "message": self.message} + + +def spm_error_detail(exc: BaseException) -> SPMErrorDetail | None: + """Recognize only public typed errors, including wrapped country errors.""" + seen = set() + current: BaseException | None = exc + while current is not None and id(current) not in seen: + seen.add(id(current)) + code = getattr(current, "code", None) + if isinstance(current, ValueError) and code in SPM_ERROR_CODES: + return SPMErrorDetail(code=code, message=str(current)) + current = current.__cause__ or current.__context__ + return None + + +class SPMCapability(BaseModel): + model_config = ConfigDict(extra="forbid") + contract_version: Literal["canonical-spm-v1"] = SPM_CONTRACT_VERSION + defaults: SPMSelection + + @model_validator(mode="after") + def pinned_defaults(self): + if not self.defaults.forecast_content_sha256 or not self.defaults.scenario: + raise ValueError( + "Certified SPM defaults must pin artifact hash and scenario" + ) + self.defaults = SPMSelection( + **{name: getattr(self.defaults, name) for name in SPMSelection.model_fields} + ) + return self + + +class SPMComparisonProvenance(BaseModel): + """One receipt per executed regional segment, separately for each policy.""" + + model_config = ConfigDict(extra="forbid") + baseline: list[SPMProvenance] = Field(min_length=1) + reform: list[SPMProvenance] = Field(min_length=1) + + +def resolve_spm_selection( + country, + selection, + *, + capability, + policyengine_version, + model_version, + route_provenance=None, +): + """Resolve only certified metadata; unrecognized future bundles fail closed.""" + if country.lower() != "us": + if selection is not None: + raise SPMInputError( + "SPM_SETTINGS_INVALID", "SPM settings are supported only for the US" + ) + return None + if capability is None: + # Existing versioned worker bundles predate canonical SPM. This does + # not certify a future bundle or change the routing registry default. + version_parts = str(policyengine_version or "").split(".") + historical = ( + len(version_parts) == 3 + and all(re.fullmatch(r"[0-9]+", p) for p in version_parts) + and tuple(map(int, version_parts)) < (5, 2, 0) + ) + # The last two pre-canonical wrappers pin US 1.764.6. A route the + # registry carries no manifest for states no model version at all, + # and states nothing that contradicts the wrapper pin; a route that + # states a different model version does. + historical = historical or ( + policyengine_version in {"5.2.0", "5.3.0"} + and model_version in (None, "1.764.6") + ) + # Country-only routes predating the bundle manifests have no wrapper + # version. Require both their route provenance -- derived from route + # shape, not from the registry's rewritable generation marker -- and + # a pre-canonical US model; absence of capability alone proves + # nothing. + model_parts = str(model_version or "").split(".") + historical = historical or ( + policyengine_version is None + and route_provenance in {"legacy-country-dict", "legacy-country-route"} + and len(model_parts) == 3 + and all(re.fullmatch(r"[0-9]+", p) for p in model_parts) + and tuple(map(int, model_parts)) <= (1, 764, 6) + ) + if selection is None and historical: + return None + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "This worker bundle has no certified canonical SPM capability", + ) + try: + cap = SPMCapability.model_validate(capability) + chosen = SPMSelection.model_validate({} if selection is None else selection) + defaults = cap.defaults + if chosen.forecast_content_sha256 not in ( + None, + defaults.forecast_content_sha256, + ): + raise ValueError( + "SPM selection does not match the certified bundle artifact hash" + ) + 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 SPMSelection.model_validate(values).model_dump(mode="json") + except ValueError as exc: + raise SPMInputError("SPM_SETTINGS_INVALID", str(exc)) from exc + + +def validate_spm_result( + result: dict, selection: Any, *, expected_year: int | str | None = None +): + """Do not accept incomplete or mixed-method child/cached output.""" + if selection is None: + if ( + result.get("spm_config") is not None + or result.get("spm_provenance") is not None + ): + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "Unexpected SPM receipt for a historical result", + ) + return None + try: + chosen = SPMSelection.model_validate(selection) + result_selection = SPMSelection.model_validate(result.get("spm_config")) + required = { + "forecast_content_sha256", + "scenario", + "geography_kind", + "county_vintage", + } + if not required <= result_selection.model_fields_set: + raise ValueError("Result has no complete resolved SPM selection") + selection = {name: getattr(chosen, name) for name in SPMSelection.model_fields} + if { + name: getattr(result_selection, name) for name in SPMSelection.model_fields + } != selection: + raise ValueError("Result SPM selection differs from the request") + provenance = SPMComparisonProvenance.model_validate( + result.get("spm_provenance") + ) + for receipt in provenance.baseline + provenance.reform: + if ( + receipt.forecast_sha256 != selection["forecast_content_sha256"] + or receipt.scenario != selection["scenario"] + or receipt.geography_kind != selection["geography_kind"] + ): + raise ValueError("Result SPM provenance differs from the request") + if expected_year is not None and str(expected_year) not in receipt.years: + raise ValueError( + "Result SPM provenance does not cover the requested year" + ) + return provenance + except ValueError as exc: + raise SPMInputError("SPM_CONFIGURATION_UNAVAILABLE", str(exc)) from exc + + +def combine_spm_results( + results: list[dict], + selection: dict | None, + *, + expected_year: int | str | None = None, +) -> dict: + receipts = [ + validate_spm_result(result, selection, expected_year=expected_year) + for result in results + ] + if selection is None: + return {} + validated: list[SPMComparisonProvenance] = [] + for receipt in receipts: + if receipt is None: + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", "Missing canonical SPM receipt" + ) + validated.append(receipt) + combined = SPMComparisonProvenance( + baseline=[r for item in validated for r in item.baseline], + reform=[r for item in validated for r in item.reform], + ) + return {"spm_config": selection, "spm_provenance": combined.model_dump(mode="json")} diff --git a/libs/policyengine-simulation-contract/tests/test_gateway_models.py b/libs/policyengine-simulation-contract/tests/test_gateway_models.py index a6bd60dd3..75d90e918 100644 --- a/libs/policyengine-simulation-contract/tests/test_gateway_models.py +++ b/libs/policyengine-simulation-contract/tests/test_gateway_models.py @@ -416,6 +416,7 @@ def test_versions_response_has_named_maps(self): ) assert response.model_dump(mode="json") == { + "spm_capabilities": {}, "policyengine": {"latest": "4.10.0"}, "us": {"latest": "1.500.0"}, "uk": {"latest": "2.66.0"}, @@ -575,6 +576,7 @@ def test_budget_window_batch_submit_response_serializes_correctly(self): "model_version": "1.500.0", "policyengine_version": None, "data_version": None, + "spm": None, "dataset": "default", }, "run_id": "batch-run-123", @@ -628,5 +630,6 @@ def test_budget_window_batch_status_response_accepts_child_jobs_and_result(self) "job_id": "fc-2026", "status": "complete", "error": None, + "errors": None, } assert dumped["result"]["kind"] == "budgetWindow" diff --git a/libs/policyengine-simulation-contract/tests/test_spm_selection.py b/libs/policyengine-simulation-contract/tests/test_spm_selection.py new file mode 100644 index 000000000..0993b617f --- /dev/null +++ b/libs/policyengine-simulation-contract/tests/test_spm_selection.py @@ -0,0 +1,202 @@ +"""Partial public selections inherit pinned bundle defaults before validation.""" + +import pytest + +from policyengine_simulation_contract.gateway_models import SimulationRequest +from policyengine_simulation_contract.spm import ( + SPMInputError, + SPMSelection, + resolve_spm_selection, +) + + +def resolve(selection, *, kind="metro", area="35620"): + return resolve_spm_selection( + "us", + selection, + capability={ + "defaults": { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": kind, + "geography_id": area, + }, + }, + policyengine_version="5.3.0", + model_version="1.824.7", + ) + + +def test_partial_area_override_inherits_metro_default_through_request(): + request = SimulationRequest(country="us", spm={"geography_id": "31080"}) + payload = request.model_dump(mode="json")["spm"] + assert payload == {"geography_id": "31080"} + resolved = resolve(payload) + assert resolved["geography_kind"] == "metro" + assert resolved["geography_id"] == "31080" + + +@pytest.mark.parametrize("kind", ["county", "national"]) +def test_partial_area_override_rejects_incompatible_defaults(kind): + with pytest.raises(SPMInputError) as error: + resolve({"geography_id": "31080"}, kind=kind, area=None) + assert error.value.code == "SPM_SETTINGS_INVALID" + + +def test_explicit_metro_inherits_area_only_from_metro_defaults(): + assert resolve({"geography_kind": "metro"})["geography_id"] == "35620" + with pytest.raises(SPMInputError) as error: + resolve({"geography_kind": "metro"}, kind="county", area=None) + assert error.value.code == "SPM_SETTINGS_INVALID" + + +def test_explicit_national_discards_default_metro_area(): + resolved = resolve({"geography_kind": "national"}) + assert resolved["geography_kind"] == "national" + assert resolved["geography_id"] is None + + +@pytest.mark.parametrize( + "provenance,wrapper,model,accepted", + [ + ("legacy-country-dict", None, "1.500.0", True), + ("legacy-country-route", None, "1.715.2", True), + ("legacy-country-route", None, "1.764.6", True), + (None, None, "1.715.2", False), + ("unknown", None, "1.715.2", False), + ("legacy-seed", None, "1.715.2", False), + ("legacy-country-route", None, "1.764.7", False), + ("legacy-country-route", None, "unknown", False), + ("legacy-country-dict", "5.3.1", "1.715.2", False), + ], +) +def test_legacy_allowance_requires_route_provenance_and_historical_model( + provenance, wrapper, model, accepted +): + kwargs = dict( + capability=None, + policyengine_version=wrapper, + model_version=model, + route_provenance=provenance, + ) + if accepted: + assert resolve_spm_selection("us", None, **kwargs) is None + else: + with pytest.raises(SPMInputError) as error: + resolve_spm_selection("us", None, **kwargs) + assert error.value.code == "SPM_CONFIGURATION_UNAVAILABLE" + for selection in ({}, {"geography_kind": "national"}): + with pytest.raises(SPMInputError) as error: + resolve_spm_selection("us", selection, **kwargs) + assert error.value.code == "SPM_CONFIGURATION_UNAVAILABLE" + + +@pytest.mark.parametrize("wrapper", ["5.2.0", "5.3.0"]) +@pytest.mark.parametrize("model", [None, "1.764.6"]) +def test_pinned_pre_canonical_wrapper_is_historical_without_a_stated_model( + wrapper, model +): + """A routing entry with no bundle manifest states no model version. + + Requiring the pin's model version there rejected live 5.2.0/5.3.0 + routes, because the response's model version falls back to the wrapper + version when the registry carries no manifest. + """ + assert ( + resolve_spm_selection( + "us", + None, + capability=None, + policyengine_version=wrapper, + model_version=model, + ) + is None + ) + + +@pytest.mark.parametrize( + "wrapper,model", + [ + ("5.2.0", "1.824.7"), + ("5.3.0", "1.824.7"), + ("5.4.0", None), + ("5.4.0", "1.764.6"), + ], +) +def test_unpinned_or_contradicted_wrapper_is_not_historical(wrapper, model): + """A stated model version that contradicts the pin, and any wrapper + outside the pinned pre-canonical set, still fail closed.""" + with pytest.raises(SPMInputError) as error: + resolve_spm_selection( + "us", + None, + capability=None, + policyengine_version=wrapper, + model_version=model, + ) + assert error.value.code == "SPM_CONFIGURATION_UNAVAILABLE" + + +class TestSerializerPreservesExplicitNulls: + """``exclude_none`` must not turn a resolved selection back into a partial. + + Presence is the contract: an omitted option inherits the bundle default, + so an option that was explicitly selected as null is a different request + from one that was omitted. The poll routes serialise every body with + ``response_model_exclude_none``, so the serializer is the only thing + standing between a completed result and a body a client cannot replay. + """ + + resolved = SPMSelection( + forecast_content_sha256="a" * 64, + scenario="ce_trend", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ) + + def test_exclude_none_keeps_explicitly_selected_nulls(self): + assert self.resolved.model_dump(mode="json", exclude_none=True) == { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + } + + def test_exclude_none_still_drops_options_the_caller_omitted(self): + partial = SPMSelection(geography_kind="national") + assert partial.model_dump(mode="json", exclude_none=True) == { + "geography_kind": "national" + } + + def test_exclude_none_survives_nesting_in_a_response_model(self): + nested = SimulationRequest(country="us", spm=self.resolved) + assert nested.model_dump(mode="json", exclude_none=True)["spm"] == { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + } + + def test_explicit_exclude_still_removes_a_null_option(self): + assert "as_of" not in self.resolved.model_dump( + mode="json", exclude_none=True, exclude={"as_of"} + ) + + def test_ordinary_serialization_is_unchanged(self): + assert self.resolved.model_dump(mode="json") == { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + } + assert SPMSelection(geography_kind="national").model_dump(mode="json") == { + "geography_kind": "national" + } diff --git a/projects/policyengine-apis-integ/tests/test_spm_generated_client.py b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py new file mode 100644 index 000000000..31c9e9d49 --- /dev/null +++ b/projects/policyengine-apis-integ/tests/test_spm_generated_client.py @@ -0,0 +1,58 @@ +"""Hermetic wire checks against the actual generated public Python client.""" + +import json + +import pytest + +from policyengine_api_simulation_client.models import ( + SPMCapability, + SPMComparisonProvenance, + SPMSelection, + SimulationRequest, +) + + +def test_omitted_selection_options_stay_omitted_in_generated_requests(): + selection = SPMSelection(scenario="zero_real") + assert selection.to_dict() == {"scenario": "zero_real"} + request = SimulationRequest.from_dict({"country": "us", "spm": selection.to_dict()}) + assert json.loads(json.dumps(request.to_dict()))["spm"] == {"scenario": "zero_real"} + + +def test_explicit_null_date_survives_generated_request_roundtrip(): + selected = {"geography_kind": "national", "as_of": None} + request = SimulationRequest.from_dict({"country": "us", "spm": selected}) + assert request.to_dict()["spm"] == selected + + +def test_generated_capability_keeps_and_enforces_contract_version(): + capability = { + "contract_version": "canonical-spm-v1", + "defaults": { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + }, + } + assert SPMCapability.from_dict(capability).to_dict() == capability + with pytest.raises(ValueError, match="contract_version"): + SPMCapability.from_dict({**capability, "contract_version": "unknown-v2"}) + + +def test_generated_receipts_preserve_arbitrary_dated_metadata(): + receipt = { + "forecast_id": "test-only", + "forecast_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "runtime_versions": {"policyengine-us": "test-only", "optional": None}, + "years": {"2026": {"status": "forecast", "window": [2021, 2025]}}, + "geographies": [{"year": 2026, "county_assignment": None}], + "composition_method": "test-only", + "storage_method": "test-only", + } + provenance = {"baseline": [receipt], "reform": [receipt]} + assert SPMComparisonProvenance.from_dict(provenance).to_dict() == provenance diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index f3bbe511f..fefa907ba 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -24,6 +24,7 @@ PingResponse, ReadinessResponse, SimulationRequest, + SimulationErrorResponse, VersionMap, VersionsResponse, ) @@ -50,7 +51,6 @@ RequestIdentifiers, ) - logger = logging.getLogger(__name__) BACKEND_RESPONSE_HEADER = { "X-PolicyEngine-Simulation-Backend": "old_gateway", @@ -61,9 +61,13 @@ def _model_json(model: BaseModel) -> JsonObject: - return _json_object_adapter.validate_python( - model.model_dump(mode="json", by_alias=True, exclude_none=True) - ) + payload = model.model_dump(mode="json", by_alias=True, exclude_none=True) + selection = getattr(model, "spm", None) + if selection is not None: + # An explicit null information date clears a bundle cutoff. Preserve + # that distinction while leaving omitted request fields omitted. + payload["spm"] = selection.model_dump(mode="json", exclude_unset=True) + return _json_object_adapter.validate_python(payload) def _response(result: BackendResponse) -> Response: @@ -271,7 +275,10 @@ async def forward( response_model_exclude_none=True, responses={ 200: {"description": "Job submitted successfully"}, - 400: {"description": "Invalid request (unknown country/version)"}, + 400: { + "description": "Invalid request or SPM selection", + "model": SimulationErrorResponse, + }, }, dependencies=protected, ) @@ -299,7 +306,10 @@ async def submit_comparison( response_model_exclude_none=True, responses={ 200: {"description": "Budget-window batch submitted successfully"}, - 400: {"description": "Invalid request (unknown country/version/year)"}, + 400: { + "description": "Invalid request or SPM selection", + "model": SimulationErrorResponse, + }, }, dependencies=protected, ) @@ -332,6 +342,10 @@ async def submit_budget_window( responses={ 200: {"description": "Job complete", "model": JobStatusResponse}, 202: {"description": "Job still running"}, + 400: { + "description": "SPM input or configuration error", + "model": JobStatusResponse, + }, 404: {"description": "Job not found"}, 500: {"description": "Job failed"}, }, @@ -360,6 +374,10 @@ async def get_job(job_id: str, request: Request) -> Response: "model": BudgetWindowBatchStatusResponse, }, 202: {"description": "Batch submitted or running"}, + 400: { + "description": "SPM input or configuration error", + "model": BudgetWindowBatchStatusResponse, + }, 404: {"description": "Batch job not found"}, 500: {"description": "Batch failed"}, }, diff --git a/projects/policyengine-simulation-entry/tests/test_app.py b/projects/policyengine-simulation-entry/tests/test_app.py index bbfa7bdea..c43bd6fcb 100644 --- a/projects/policyengine-simulation-entry/tests/test_app.py +++ b/projects/policyengine-simulation-entry/tests/test_app.py @@ -358,3 +358,30 @@ def test_upstream_error_status_and_body_are_preserved(client, backend, status_co assert result.json() == {"detail": "upstream response"} assert result.headers[REQUEST_ID_HEADER] == "request-upstream-error" assert "x-request-id" not in result.headers + + +@pytest.mark.parametrize( + "path,extra", + [ + ("/simulate/economy/comparison", {}), + ( + "/simulate/economy/budget-window", + {"region": "us", "start_year": "2026", "window_size": 2}, + ), + ], +) +@pytest.mark.parametrize("date_fields", [{"as_of": None}, {}]) +def test_submission_preserves_spm_field_presence( + client, backend, path, extra, date_fields +): + selection = {"geography_kind": "national", **date_fields} + + result = client.post(path, json={"country": "us", "spm": selection, **extra}) + + assert result.status_code == 200 + forwarded = backend.requests[-1] + assert forwarded.method == "POST" + assert forwarded.path == path + # Explicit null clears the downstream bundle cutoff; an omitted date inherits it. + # Other omitted fields must remain absent so bundle defaults can fill them. + assert forwarded.json_body["spm"] == selection diff --git a/projects/policyengine-simulation-entry/tests/test_openapi.py b/projects/policyengine-simulation-entry/tests/test_openapi.py index 9b1c16afb..c7cc77054 100644 --- a/projects/policyengine-simulation-entry/tests/test_openapi.py +++ b/projects/policyengine-simulation-entry/tests/test_openapi.py @@ -7,7 +7,6 @@ from policyengine_simulation_entry.app import create_app - type HttpMethod = Literal["get", "post", "put", "patch", "delete"] @@ -60,16 +59,39 @@ def operation_ids(spec: OpenAPIDocument) -> dict[tuple[str, str], str]: def normalized_compatibility_paths(spec: OpenAPIDocument) -> OpenAPIPathMap: - """Remove the two intentional Cloud Run-only OpenAPI additions.""" + """Remove intentional Cloud Run and additive SPM error documentation.""" paths = deepcopy(spec["paths"]) paths.pop("/ready", None) for operations in paths.values(): for method, operation in operations.items(): if method.lower() in {"get", "post", "put", "patch", "delete"}: operation.pop("security", None) + for path, description in { + "/simulate/economy/comparison": "Invalid request (unknown country/version)", + "/simulate/economy/budget-window": "Invalid request (unknown country/version/year)", + }.items(): + paths[path]["post"]["responses"]["400"] = {"description": description} + for path in ("/jobs/{job_id}", "/budget-window-jobs/{batch_job_id}"): + paths[path]["get"]["responses"].pop("400", None) return paths +def test_canonical_spm_extensions_are_public(): + spec = create_app().openapi() + schemas = spec["components"]["schemas"] + assert schemas["SPMSelection"]["additionalProperties"] is False + assert all( + "default" not in field + for field in schemas["SPMSelection"]["properties"].values() + ) + assert "spm" in schemas["SimulationRequest"]["properties"] + assert "spm" in schemas["BudgetWindowBatchRequest"]["properties"] + assert "spm_provenance" in schemas["SingleYearMacroOutput"]["properties"] + assert "spm_provenance" in schemas["BudgetWindowAnnualImpact"]["properties"] + for path in ("/jobs/{job_id}", "/budget-window-jobs/{batch_job_id}"): + assert "400" in spec["paths"][path]["get"]["responses"] + + def test_route_table_is_frozen(): assert methods(cast(OpenAPIDocument, create_app().openapi())) == EXPECTED_ROUTES @@ -110,6 +132,7 @@ def test_normalized_contract_matches_old_gateway(): ) == normalized_compatibility_paths(gateway_spec) cloud_run_schemas = deepcopy(cloud_run_spec["components"]["schemas"]) cloud_run_schemas.pop("ReadinessResponse") + cloud_run_schemas.pop("SimulationErrorResponse") assert cloud_run_schemas == gateway_spec["components"]["schemas"] gateway_operation_ids = operation_ids(gateway_spec) cloud_run_operation_ids = operation_ids(cloud_run_spec) diff --git a/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py b/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py index 23804049d..f296f91dc 100644 --- a/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py +++ b/projects/policyengine-simulation-executor/fixtures/test_simulation_api_contracts.py @@ -1,6 +1,8 @@ """Fixtures for simulation API contract tests.""" CURRENT_SINGLE_YEAR_MACRO_KEYS = { + "spm_config", + "spm_provenance", "model_version", "data_version", "budget", @@ -30,6 +32,8 @@ } CURRENT_SINGLE_YEAR_MACRO_RESULT = { + "spm_config": None, + "spm_provenance": None, "model_version": "1.715.2", "data_version": "1.115.5", "budget": { diff --git a/projects/policyengine-simulation-executor/fixtures/wrapper_spm.py b/projects/policyengine-simulation-executor/fixtures/wrapper_spm.py new file mode 100644 index 000000000..46442315e --- /dev/null +++ b/projects/policyengine-simulation-executor/fixtures/wrapper_spm.py @@ -0,0 +1,56 @@ +"""The canonical wrapper's ``storage_id``, transcribed for hermetic tests. + +The executor pins a pre-canonical ``policyengine``: its ``Simulation`` has +no ``spm`` field and no ``storage_id``, so hermetic CI cannot import the +property that names every canonical baseline artifact. Precompute plans a +store path from ``BaselineArtifactIdentity.storage_id`` and the +in-container worker aborts when the wrapper's value disagrees, so the two +derivations have to be one identifier reached down two paths. + +This is that second path, copied out of the SPM-capable wrapper so the +key-discipline tests and the precompute writer==reader tests state it once. +It proves our side has not drifted from the contract we read; it is not the +wrapper, and only the SPM_NATIVE_SMOKE_SOURCE-gated suites run against one. + +``policyengine/core/simulation.py``:: + + @property + def storage_id(self) -> str: + \"\"\"Include resolved SPM settings in cache and saved-result identity.\"\"\" + config = self.spm_config + if config is None: + return self.id + encoded = json.dumps(config, sort_keys=True, separators=(",", ":")).encode() + return f"{self.id}-spm-{hashlib.sha256(encoded).hexdigest()}" + +Read from ``policyengine-5.3.0-py3-none-any.whl`` sha256 +``8c640d967575dddad70840bcbe938cea251c56958eded647c83eb9c5902735f1``, the +unpublished development wheel the native qualification lane installs. The +hash is the identifier, not the version: another local build carries the +same filename and version and has no ``storage_id`` at all. + +``spm_config`` above is not a stored value. On the canonical wrapper it +refuses a non-US ``tax_benefit_model_version`` and otherwise re-resolves +``spm`` through the installed bundle, so an unset selection becomes the +bundle's defaults rather than None -- which is why this function takes a +config that has already been resolved, and why the agreement claim is +about the digest, not about the resolution. Both are checked against that +wheel in ``TestWrapperStorageIdAgreement``'s recorded out-of-band run. +""" + +import hashlib +import json + + +def wrapper_storage_id(simulation_id: str, spm_config: dict | None) -> str: + if spm_config is None: + return simulation_id + encoded = json.dumps(spm_config, sort_keys=True, separators=(",", ":")).encode() + return f"{simulation_id}-spm-{hashlib.sha256(encoded).hexdigest()}" + + +def installed_wrapper_has_storage_id() -> bool: + """True once an SPM-capable wrapper is pinned and this file can retire.""" + from policyengine.core import Simulation + + return hasattr(Simulation, "storage_id") diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py index a8f93c603..b1c04010c 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_batch.py @@ -13,6 +13,11 @@ def run_budget_window_batch_impl(params: dict[str, Any]) -> dict[str, Any]: + from policyengine_simulation_executor.spm import normalize_runtime_spm + + selection = normalize_runtime_spm(params) + if selection is not None: + params = {**params, "spm": selection} batch_job_id = modal.current_function_call_id() set_attribute("batch_job_id", batch_job_id) with segment(SegmentName.BUDGET_WINDOW_CONTEXT): diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py index 2b9d7ca47..cedff79b8 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_results.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_results.py @@ -3,6 +3,7 @@ from __future__ import annotations from decimal import Decimal +from policyengine_simulation_contract.spm import SPMSelection, validate_spm_result from typing import Any from policyengine_simulation_contract.gateway_models import ( @@ -36,7 +37,9 @@ def extract_annual_impact( *, simulation_year: str, child_result: dict[str, Any], + spm: Any = None, ) -> BudgetWindowAnnualImpact: + receipt = validate_spm_result(child_result, spm, expected_year=simulation_year) budget = child_result.get("budget", {}) if not isinstance(budget, dict): raise ValueError("Malformed budget-window child result: missing budget object") @@ -61,6 +64,8 @@ def extract_annual_impact( state_tax_revenue_impact = 0.0 return BudgetWindowAnnualImpact( + spm_config=SPMSelection.model_validate(spm) if spm is not None else None, + spm_provenance=receipt, year=simulation_year, taxRevenueImpact=tax_revenue_impact, federalTaxRevenueImpact=tax_revenue_impact - state_tax_revenue_impact, diff --git a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py index 262c5a9b5..c54bee76c 100644 --- a/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py @@ -30,6 +30,7 @@ put_batch_job_seed, put_batch_job_state, ) +from policyengine_simulation_contract.spm import spm_error_detail from policyengine_simulation_observability.errors import log_and_redact_exception from policyengine_simulation_observability.observability import SegmentName @@ -186,6 +187,15 @@ def poll_running_children_once(self) -> bool: except TimeoutError: continue except Exception as exc: + detail = spm_error_detail(exc) + if detail: + self.state.errors = [detail] + self.fail_batch_for_child_error( + simulation_year=simulation_year, error=detail.message + ) + self.state.child_jobs[simulation_year].errors = [detail] + put_batch_job_state(self.state) + return False redacted = log_and_redact_exception( exc, scope="budget_window_child_call", @@ -212,8 +222,18 @@ def poll_running_children_once(self) -> bool: annual_impact = extract_annual_impact( simulation_year=simulation_year, child_result=child_result, + spm=self.state.request_payload.get("spm"), ) except Exception as exc: + detail = spm_error_detail(exc) + if detail: + self.state.errors = [detail] + self.fail_batch_for_child_error( + simulation_year=simulation_year, error=detail.message + ) + self.state.child_jobs[simulation_year].errors = [detail] + put_batch_job_state(self.state) + return False redacted = log_and_redact_exception( exc, scope="budget_window_child_result_parsing", diff --git a/projects/policyengine-simulation-executor/src/modal/segmented_national.py b/projects/policyengine-simulation-executor/src/modal/segmented_national.py index 310c1e0da..e2c4cd583 100644 --- a/projects/policyengine-simulation-executor/src/modal/segmented_national.py +++ b/projects/policyengine-simulation-executor/src/modal/segmented_national.py @@ -33,6 +33,11 @@ from policyengine_simulation_executor.segmented_national_reduce import ( build_national_output, ) +from policyengine_simulation_contract.spm import ( + spm_error_detail, + SPMInputError, + combine_spm_results, +) from policyengine_simulation_observability.errors import log_and_redact_exception from policyengine_simulation_observability.observability import SegmentName from policyengine_simulation_observability.telemetry import split_internal_payload @@ -235,6 +240,10 @@ def _collect(self, handles: list[tuple[list[str], Any]]) -> list[dict]: poll_errors.pop(index, None) continue except Exception as exc: + detail = spm_error_detail(exc) + if detail: + self._cancel_all(handles) + raise SPMInputError(detail.code, detail.message) from exc # One transient poll-RPC blip must not kill the job; a # real child failure re-raises on the next probe too. attempts = poll_errors.get(index, 0) + 1 @@ -300,6 +309,13 @@ def _reduce(self, child_results: list[dict], *, country_module) -> dict[str, Any year=_parse_year(simulation_params), resolved_data_version=_requested_data_version(simulation_params), ) + output.update( + combine_spm_results( + child_results, + simulation_params.get("spm"), + expected_year=_parse_year(simulation_params), + ) + ) for key in ("model_version", "data_version"): if output.get(key): set_attribute(key, str(output[key])) @@ -315,6 +331,11 @@ def run_segmented_national_impl( def dispatch_run_simulation(params: dict[str, Any], *, app_name: str) -> dict[str, Any]: """The run_simulation entrypoint's routing: segmented national fan-out for eligible requests, the monolithic path for everything else.""" + from policyengine_simulation_executor.spm import normalize_runtime_spm + + selection = normalize_runtime_spm(params) + if selection is not None: + params = {**params, "spm": selection} if should_run_segmented_national(params): return run_segmented_national_impl(params, app_name=app_name) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py index 78027db9e..649a652a6 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_version_registry.py @@ -5,7 +5,7 @@ import modal from packaging.version import InvalidVersion, Version -from typing import Iterable, TypedDict +from typing import Iterable, TypedDict, NotRequired POLICYENGINE_VERSION_DICT_NAME = "simulation-api-policyengine-versions" US_VERSION_DICT_NAME = "simulation-api-us-versions" @@ -38,6 +38,7 @@ class BundleManifestMetadata(TypedDict): policyengine_version: str us: CountryBundleMetadata uk: CountryBundleMetadata + spm: NotRequired[dict] class LatestVersions(TypedDict): @@ -104,12 +105,18 @@ def build_bundle_manifest_metadata( app_name: str, policyengine_version: str, ) -> BundleManifestMetadata: - return { + from policyengine_simulation_executor.spm import runtime_spm_capability + + capability = runtime_spm_capability() + metadata: BundleManifestMetadata = { "app_name": app_name, "policyengine_version": policyengine_version, "us": _country_bundle_metadata("us"), "uk": _country_bundle_metadata("uk"), } + if capability is not None: + metadata["spm"] = capability.model_dump(mode="json") + return metadata def _empty_routing_state() -> RoutingState: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py index cb83156e5..5cdcffba9 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py @@ -74,6 +74,7 @@ def dataset_key( data_build_fingerprint: Optional[str], model_version: str, policyengine_version: str, + spm: Optional[dict] = None, ) -> str: """Digest identifying one single-year dataset artifact. @@ -86,7 +87,8 @@ def dataset_key( """ return canonical_digest( { - "schema": DATASET_KEY_SCHEMA, + "schema": DATASET_KEY_SCHEMA if spm is None else "ds2-spm", + **({"spm": spm} if spm is not None else {}), "country": country.lower(), "dataset": dataset, "year": int(year), @@ -108,6 +110,7 @@ def baseline_key( dataset_digest: str, model_version: str, policyengine_version: str, + spm: Optional[dict] = None, policy: str = CURRENT_LAW_POLICY, ) -> str: """Digest identifying one precomputed baseline simulation artifact. @@ -119,7 +122,8 @@ def baseline_key( """ return canonical_digest( { - "schema": BASELINE_KEY_SCHEMA, + "schema": BASELINE_KEY_SCHEMA if spm is None else "bl2-spm", + **({"spm": spm} if spm is not None else {}), "country": country.lower(), "region": region, "scope_key": scope_key, @@ -206,11 +210,13 @@ class DatasetArtifactIdentity: data_build_fingerprint: Optional[str] model_version: str policyengine_version: str + spm: Optional[dict] = None @property def digest(self) -> str: return dataset_key( country=self.country, + spm=self.spm, dataset=self.dataset, year=self.year, data_version=self.data_version, @@ -238,12 +244,14 @@ class BaselineArtifactIdentity: region: str scope_key: Optional[str] dataset: DatasetArtifactIdentity + spm: Optional[dict] = None @property def digest(self) -> str: return baseline_key( country=self.country, region=self.region, + spm=self.spm, scope_key=self.scope_key, dataset_digest=self.dataset.digest, model_version=self.dataset.model_version, @@ -254,9 +262,15 @@ def digest(self) -> str: def simulation_id(self) -> str: return baseline_simulation_id(self.digest) + @property + def storage_id(self) -> str: + if self.spm is None: + return self.simulation_id + return f"{self.simulation_id}-spm-{canonical_digest(self.spm)}" + @property def store_path(self) -> str: - return baseline_artifact_path(self.country, self.digest, self.simulation_id) + return baseline_artifact_path(self.country, self.digest, self.storage_id) def _receipt_source_sha256(country: str, data_version: str) -> Optional[str]: @@ -310,7 +324,11 @@ def collect_dataset_identity(country: str, year: int) -> DatasetArtifactIdentity stem = dataset_logical_name( resolve_dataset_reference(bundle.country, bundle.default_dataset) ) + from policyengine_simulation_executor.spm import normalize_runtime_spm + + selection = normalize_runtime_spm({"country": country, "time_period": year}) return DatasetArtifactIdentity( + spm=selection, country=bundle.country, dataset=bundle.default_dataset, stem=stem, @@ -330,8 +348,15 @@ def collect_baseline_identity( *, region: str, scope_key: Optional[str], + spm: Optional[dict] = None, ) -> BaselineArtifactIdentity: + from policyengine_simulation_executor.spm import normalize_runtime_spm + + selection = normalize_runtime_spm( + {"country": country, "time_period": year, "spm": spm} + ) return BaselineArtifactIdentity( + spm=selection, country=country.lower(), region=region, scope_key=scope_key, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py index dadb9672e..f8dfa290c 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/baseline_artifacts.py @@ -33,7 +33,8 @@ # Artifact outcomes reported via the `baseline_artifact` observability # attribute. "hit": loaded (disk or in-process cache) and complete; -# "incomplete": loaded but missing requested columns -> recomputed; +# "incomplete": loaded but missing requested columns or a valid SPM receipt +# -> recomputed; # "miss": no artifact -> computed (today's behavior). OUTCOME_HIT = "hit" OUTCOME_INCOMPLETE = "incomplete" @@ -102,7 +103,11 @@ def qualifying_baseline_identity( return None return artifact_keys.collect_baseline_identity( - country, year, region=region, scope_key=scope_key + country, + year, + region=region, + scope_key=scope_key, + **({"spm": params["spm"]} if params.get("spm") is not None else {}), ) @@ -130,7 +135,11 @@ def deterministic_baseline_id( scoping_strategy=scoping_strategy, year=year, ) - except Exception: + except Exception as exc: + from policyengine_simulation_contract.spm import spm_error_detail + + if spm_error_detail(exc): + raise logger.warning( "Could not collect baseline artifact identity for %s; " "using a random simulation id", @@ -171,23 +180,45 @@ def _record_outcome(self, outcome: str) -> None: self._artifact_outcome = outcome def ensure(self) -> None: + # The wrapper restores selection metadata on load/cache hits. Keep the + # request's selection so an invalid cached entry cannot replace it. + selection = getattr(self, "spm_config", None) + requested_spm = getattr(self, "spm", None) self._computed_this_process = False super().ensure() if self._computed_this_process: self._record_outcome(OUTCOME_MISS) return + from policyengine_simulation_contract.spm import SPMInputError + + from policyengine_simulation_executor.spm import simulation_spm_result + + # A tax-only artifact can legitimately have no receipt years. Check its + # columns first, then treat unusable cached receipts as another gap. missing = self._missing_output_columns() - if not missing: + receipt_error = None + if not missing and selection is not None: + try: + simulation_spm_result( + self, self, selection, expected_year=self.dataset.year + ) + except SPMInputError as exc: + receipt_error = str(exc) + if not missing and receipt_error is None: self._record_outcome(OUTCOME_HIT) return logger.warning( - "Baseline artifact %s is missing output columns %s; recomputing", + "Baseline artifact %s is incomplete (missing columns=%s, " + "SPM receipt=%s); recomputing", self.id, missing, + receipt_error, ) self._record_outcome(OUTCOME_INCOMPLETE) + if selection is not None: + setattr(self, "spm", requested_spm) self.run() self.save() # Replace the in-process cache entry so this request's second @@ -197,8 +228,11 @@ def ensure(self) -> None: # existing key and exposes no remove, so evict directly first. from policyengine.core.simulation import _cache - _cache._cache.pop(self.id, None) - _cache.add(self.id, self) + _cache._cache.pop(getattr(self, "storage_id", self.id), None) + _cache.add( + getattr(self, "storage_id", self.id), + self.model_copy(deep=False) if selection is not None else self, + ) def _missing_output_columns(self) -> list[tuple[str, str]]: data = getattr(self.output_dataset, "data", None) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/compat_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/compat_models.py index 64e880015..12d545d98 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/compat_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/compat_models.py @@ -5,6 +5,7 @@ from typing import Any, Optional from pydantic import BaseModel, ConfigDict +from policyengine_simulation_contract.spm import SPMSelection from policyengine_simulation_executor.simulation_macro_output import ( SingleYearMacroOutput, @@ -15,6 +16,7 @@ class SimulationOptions(BaseModel): """Legacy request schema name kept for generated clients.""" country: str + spm: Optional[SPMSelection] = None scope: Optional[str] = None data: Optional[str] = None time_period: Optional[str | int] = None diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 1232ff54d..f26bc6640 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -152,7 +152,7 @@ def build_manifest(plan: PrecomputePlan) -> ArtifactManifest: ManifestArtifact( type="baseline", path=entry.path, - filename=baseline_artifact_filename(entry.simulation_id), + filename=entry.path.rsplit("/", maxsplit=1)[-1], year=entry.year, digest=entry.digest, ) @@ -310,6 +310,14 @@ def _prepare_cohort_baseline(bucket: str, expected: BaselinePlanEntry): f"({expected.simulation_id} != {baseline.id}); refusing to act " "under a mismatched key." ) + planned_storage_id = Path(expected.path).stem + storage_id = getattr(baseline, "storage_id", baseline.id) + if storage_id != planned_storage_id: + raise RuntimeError( + "Planned and in-container baseline storage ids disagree " + f"({planned_storage_id} != {storage_id}); refusing to act " + "under a mismatched key." + ) # The extras economic_impact_analysis applies unconditionally before # ensure(); the artifact must carry them or every request would fail @@ -341,7 +349,9 @@ def compute_baseline_impl( baseline.ensure() compute_seconds = time.monotonic() - started - artifact_file = data_folder / baseline_artifact_filename(baseline.id) + artifact_file = data_folder / baseline_artifact_filename( + getattr(baseline, "storage_id", baseline.id) + ) if not artifact_file.exists(): raise RuntimeError(f"ensure() left no artifact at {artifact_file}") uploaded = store.upload_file(expected.path, artifact_file) @@ -373,7 +383,9 @@ def verify_determinism_impl( from policyengine_simulation_executor.baseline_artifacts import OUTCOME_HIT store, data_folder, baseline = _prepare_cohort_baseline(bucket, expected) - artifact_file = data_folder / baseline_artifact_filename(baseline.id) + artifact_file = data_folder / baseline_artifact_filename( + getattr(baseline, "storage_id", baseline.id) + ) if not artifact_file.exists(): store.download_file(expected.path, artifact_file) baseline.ensure() @@ -383,7 +395,9 @@ def verify_determinism_impl( f"(outcome={baseline.artifact_outcome})" ) + selection = getattr(baseline, "spm_config", None) fresh = Simulation( + **({"spm": selection} if selection is not None else {}), dataset=baseline.dataset, tax_benefit_model_version=baseline.tax_benefit_model_version, policy=None, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py index a8b80bd3c..309f44315 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation.py @@ -1,6 +1,8 @@ import logging from fastapi import APIRouter +from fastapi.responses import JSONResponse +from policyengine_simulation_contract.spm import spm_error_detail from policyengine_simulation_executor.simulation_runtime import run_simulation_impl from policyengine_simulation_executor.compat_models import ( @@ -15,11 +17,28 @@ def create_router(): router = APIRouter() @router.post("/simulate/economy/comparison", response_model=EconomyComparison) - async def simulate(parameters: SimulationOptions) -> EconomyComparison: + async def simulate( + parameters: SimulationOptions, + ) -> EconomyComparison | JSONResponse: logger.info("Calculating comparison") - result = run_simulation_impl( - parameters.model_dump(mode="json", exclude_none=True) - ) + try: + payload = parameters.model_dump(mode="json", exclude_none=True) + if parameters.spm is not None: + payload["spm"] = parameters.spm.model_dump(mode="json") + result = run_simulation_impl(payload) + except ValueError as exc: + detail = spm_error_detail(exc) + if detail: + return JSONResponse( + status_code=400, + content={ + "status": "failed", + "result": None, + "error": detail.message, + "errors": [detail.model_dump()], + }, + ) + raise logger.info("Comparison complete") return EconomyComparison.model_validate(result) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py index fb804d9f2..c031b2ee7 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_common.py @@ -98,6 +98,10 @@ def _poverty_module_function(name: str): def _try_compute_output(label: str, fn, *args, **kwargs): try: return fn(*args, **kwargs) - except Exception: + except Exception as exc: + from policyengine_simulation_contract.spm import spm_error_detail + + if spm_error_detail(exc) is not None: + raise logger.warning("Unable to calculate %s", label, exc_info=True) return None diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index d75582f4a..5a90b5322 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -136,7 +136,20 @@ def run_simulation_impl(params: dict) -> dict: # Set up GCP credentials if needed. The credentials temp file is # cleaned up on exit so we never leave signed JSON material on disk. with setup_gcp_credentials(): - return _run_simulation_impl_core(params) + try: + return _run_simulation_impl_core(params) + except ValueError as exc: + # Gateway images intentionally do not install country packages. + # Transport the public error through the shared contract instead. + from policyengine_simulation_contract.spm import ( + SPMInputError, + spm_error_detail, + ) + + detail = spm_error_detail(exc) + if detail is not None: + raise SPMInputError(detail.code, detail.message) from None + raise def _parse_year(params: dict[str, Any]) -> int: @@ -490,6 +503,11 @@ def _build_simulation( deterministic_baseline_id, ) + from policyengine_simulation_executor.spm import normalize_runtime_spm + + selection = normalize_runtime_spm(params) + params = {**params, **({"spm": selection} if selection is not None else {})} + spm_kwargs = {"spm": selection} if selection is not None else {} country = params.get("country", "us") country_module = _country_module(country) simulation_id = deterministic_baseline_id( @@ -505,6 +523,7 @@ def _build_simulation( # when one is baked beside the dataset, and the subclass validates # the load (falling back to run()) — see baseline_artifacts. return ArtifactBaselineSimulation( + **spm_kwargs, id=simulation_id, dataset=dataset, tax_benefit_model_version=country_module.model, @@ -512,6 +531,7 @@ def _build_simulation( scoping_strategy=scoping_strategy, ) return Simulation( + **spm_kwargs, dataset=dataset, tax_benefit_model_version=country_module.model, policy=policy, @@ -523,6 +543,14 @@ def _run_simulation_impl_core(params: dict) -> dict: with segment(SegmentName.REQUEST_PARSE): simulation_params, telemetry, metadata = split_internal_payload(params) metadata = metadata or {} + from policyengine_simulation_executor.spm import ( + normalize_runtime_spm, + simulation_spm_result, + ) + + selection = normalize_runtime_spm(simulation_params) + if selection is not None: + simulation_params["spm"] = selection logger.info( "Starting simulation for country=%s run_id=%s process_id=%s", @@ -590,6 +618,11 @@ def _run_simulation_impl_core(params: dict) -> dict: resolved_region_code=region_resolution.code, ) output = builder.serialize() + output.update( + simulation_spm_result( + baseline, reform, selection, expected_year=_parse_year(simulation_params) + ) + ) # ensure() has run inside the builder by now, so the artifact outcome # (hit / incomplete / miss) is known for deterministic-id baselines. artifact_outcome = getattr(baseline, "artifact_outcome", None) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py new file mode 100644 index 000000000..bfb1c311a --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py @@ -0,0 +1,226 @@ +"""Certified SPM runtime selection, independent of caller-supplied metadata.""" + +import json +from functools import lru_cache + +from policyengine_simulation_contract.spm import ( + SPMCapability, + SPMInputError, + SPMProvenance, + combine_spm_results, + resolve_spm_selection, + spm_error_detail, +) + + +@lru_cache(maxsize=4) +def _forecast(expected_sha256): + from spm_calculator import load_forecast + + return load_forecast(expected_sha256=expected_sha256) + + +@lru_cache(maxsize=1) +def runtime_spm_capability(): + """The installed image's certified SPM capability, or None. + + Memoized for the life of the process because every input is a property + of the installed image -- the bundle configuration, the wrapper's model + fields, the country model's methods and the pinned forecast -- none of + which a running container can change. One national request resolves the + selection about six times (the worker entrypoint, ``_build_simulation``, + and the dataset and baseline identity collectors), and each resolution + called this. + + ``lru_cache`` does not memoize exceptions, so every fail-closed path + still re-derives and re-raises: the cache can only ever skip work that + already succeeded. Tests that stub the installed environment must call + :func:`reset_spm_runtime_caches` (the executor suite does, automatically). + """ + from policyengine.bundle import get_current_bundle + from policyengine.core import Simulation + + try: + bundle = get_current_bundle() + configured = bundle.get("measurements", {}).get("spm") + except (ImportError, ValueError, TypeError, AttributeError) as exc: + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "The installed SPM bundle configuration is unavailable", + ) from exc + if configured is None: + return None + if "spm" not in Simulation.model_fields or not hasattr( + Simulation, "spm_provenance" + ): + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "The installed wrapper does not support canonical SPM", + ) + try: + from policyengine_us import Microsimulation + + if not ( + hasattr(Microsimulation, "spm_config") + and callable(getattr(Microsimulation, "spm_provenance", None)) + ): + raise ValueError("The installed US model does not support canonical SPM") + capability = SPMCapability(defaults=configured) + forecast = _forecast(capability.defaults.forecast_content_sha256) + forecast.entry( + forecast.years[0], + scenario=capability.defaults.scenario, + as_of=capability.defaults.as_of, + ) + return capability + except (ImportError, ValueError, TypeError, OSError) as exc: + raise SPMInputError("SPM_CONFIGURATION_UNAVAILABLE", str(exc)) from exc + + +@lru_cache(maxsize=64) +def _prevalidate_selection(selection_json: str, start_year: int, window_size: int): + """Prove a resolved selection can be measured, before anything expensive. + + Keyed on the resolved selection and the year range, which is everything + it reads: the provider is constructed from the selection alone and asked + only for per-year metadata. The same request resolves the same selection + several times over (see :func:`runtime_spm_capability`), and the review + of this change counted about six provider constructions per national + request; memoizing collapses them to one per distinct selection and + range. + + Only successful validations are memoized -- ``lru_cache`` re-runs after + an exception -- so a rejected selection is rejected again, with the same + typed error, every time it is asked for. + """ + from spm_calculator.policyengine_adapter import PolicyEngineSPMProvider + + selection = json.loads(selection_json) + forecast = _forecast(selection["forecast_content_sha256"]) + if selection["county_vintage"] != "2020": + raise ValueError("Unsupported county vintage: use 2020") + provider = PolicyEngineSPMProvider( + forecast, + **{ + key: value + for key, value in selection.items() + if key != "forecast_content_sha256" + }, + ) + for year in range(start_year, start_year + window_size): + # Use the country adapter's typed year contract. This temporary + # provider validates metadata without measuring any SPM amount + # or modifying the actual simulation's calculation receipts. + provider.year_metadata(year) + if selection["geography_kind"] == "metro": + try: + forecast.geography_factor( + year, + "renter", + kind="metro", + geoid=selection["geography_id"], + scenario=selection["scenario"], + as_of=selection["as_of"], + ) + except ValueError as exc: + raise SPMInputError("SPM_GEOGRAPHY_UNAVAILABLE", str(exc)) from None + + +# Bound at import, so clearing still works while a test has replaced one of +# these module attributes with a stub. +_MEMO_CLEARERS = ( + runtime_spm_capability.cache_clear, + _prevalidate_selection.cache_clear, +) + + +def reset_spm_runtime_caches(): + """Forget what was memoized from the installed environment. + + Production never needs this: a container's image is fixed for the life + of the process. Tests that stub the installed bundle, wrapper or + provider do, and the executor suite calls it around every test. + + ``_forecast`` is deliberately not cleared. It is keyed by the content + hash of what it loads, so it cannot go stale, and a test that stubs it + replaces the module attribute rather than filling the cache. + """ + for clear in _MEMO_CLEARERS: + clear() + + +def normalize_runtime_spm(params): + """Resolve before dataset loading, artifact lookup, or child submission.""" + from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + ) + + country = params.get("country", "us").lower() + bundle = get_country_release_bundle(country) + selection = resolve_spm_selection( + country, + params.get("spm"), + capability=runtime_spm_capability() if country == "us" else None, + policyengine_version=bundle.policyengine_version, + model_version=bundle.model_version, + ) + if selection is not None: + canonical = json.dumps(selection, sort_keys=True, separators=(",", ":")) + try: + from policyengine_simulation_executor.simulation_runtime import _parse_year + + try: + start = int(params.get("start_year") or _parse_year(params)) + window = int(params.get("window_size", 1)) + except ValueError: + # Keep the order the range's old parse position gave it: the + # selection was checked before the range, so an unusable year + # never masked an unknown scenario. An empty range asks for + # no year and reports the selection's own defect first. + _prevalidate_selection(canonical, 0, 0) + raise + _prevalidate_selection(canonical, start, window) + except ValueError as exc: + detail = spm_error_detail(exc) + if detail: + raise SPMInputError(detail.code, detail.message) from exc + # The installed provider constructor still exposes the forecast's + # plain scenario error; use the same narrow translation as the + # calculator's Frame and Axiom adapters, leaving other errors alone. + if str(exc).startswith("Unknown forecast scenario:"): + raise SPMInputError("SPM_SCENARIO_UNAVAILABLE", str(exc)) from exc + raise SPMInputError("SPM_SETTINGS_INVALID", str(exc)) from exc + return selection + + +def simulation_spm_result(baseline, reform, selection, *, expected_year=None): + if selection is None: + return {} + receipts = [] + for simulation in (baseline, reform): + if simulation.spm_config != selection: + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "Simulation ignored the requested SPM selection", + ) + try: + receipts.append( + SPMProvenance.model_validate(simulation.spm_provenance()).model_dump( + mode="json" + ) + ) + except ValueError as exc: + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "Simulation has no valid SPM calculation receipt", + ) from exc + return combine_spm_results( + [ + { + "spm_config": selection, + "spm_provenance": {"baseline": [receipts[0]], "reform": [receipts[1]]}, + } + ], + selection, + expected_year=expected_year, + ) diff --git a/projects/policyengine-simulation-executor/tests/conftest.py b/projects/policyengine-simulation-executor/tests/conftest.py index c94efcc3b..466293e88 100644 --- a/projects/policyengine-simulation-executor/tests/conftest.py +++ b/projects/policyengine-simulation-executor/tests/conftest.py @@ -4,6 +4,8 @@ import sys from pathlib import Path +import pytest + pytest_plugins = () project_root = Path(__file__).parent.parent @@ -15,3 +17,19 @@ ).ensure_project_root_on_path ensure_project_root_on_path() + + +@pytest.fixture(autouse=True) +def _forget_spm_runtime_caches(): + """Keep the SPM runtime memoization out of every other test's way. + + ``runtime_spm_capability`` and the selection prevalidation are memoized + on the installed environment, which a container cannot change but a test + can. Clearing on both sides of each test means a stub is never served a + previous test's answer, and never leaves one behind. + """ + from policyengine_simulation_executor.spm import reset_spm_runtime_caches + + reset_spm_runtime_caches() + yield + reset_spm_runtime_caches() diff --git a/projects/policyengine-simulation-executor/tests/native_spm_support.py b/projects/policyengine-simulation-executor/tests/native_spm_support.py new file mode 100644 index 000000000..e88abf6af --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/native_spm_support.py @@ -0,0 +1,36 @@ +"""One-household native dataset shared by installed SPM integration tests.""" + + +def materialize_native_household(source_path, tmp_path): + import pandas as pd + from policyengine.tax_benefit_models.us import ensure_datasets + + tiny = tmp_path / "native.h5" + with pd.HDFStore(source_path, "r") as source, pd.HDFStore(tiny, "w") as dest: + household = source.select("household", start=0, stop=1) + household_id = int(household.household_id.iloc[0]) + people = source.select("person", where=f"person_household_id == {household_id}") + assert 0 < len(people) <= 20 + dest.put("household", household, format="table", data_columns=True) + dest.put("person", people, format="table", data_columns=True) + for entity in ("spm_unit", "tax_unit", "family", "marital_unit"): + ids = [int(value) for value in people[f"person_{entity}_id"].unique()] + dest.put( + entity, + source.select(entity, where=f"{entity}_id in {ids}"), + format="table", + data_columns=True, + ) + dest.put("_time_period", source.select("_time_period"), format="table") + # Unmanaged is an explicit test-only path; production keeps certified bundle + # loading. This exercises the real native-to-year materializer. + datasets = ensure_datasets( + datasets=[str(tiny)], + years=[2024], + data_folder=str(tmp_path / "year-data"), + allow_unmanaged=True, + ) + dataset = next(iter(datasets.values())) + dataset.load() + assert len(dataset.data.person) == len(people) + return dataset diff --git a/projects/policyengine-simulation-executor/tests/test_artifact_keys.py b/projects/policyengine-simulation-executor/tests/test_artifact_keys.py index 955b91413..9ece5ae4e 100644 --- a/projects/policyengine-simulation-executor/tests/test_artifact_keys.py +++ b/projects/policyengine-simulation-executor/tests/test_artifact_keys.py @@ -9,9 +9,15 @@ churn (or, worse, a writer/reader mismatch). """ +from types import SimpleNamespace + import pytest from fixtures.identity_stubs import install_identity_stubs +from fixtures.wrapper_spm import ( + installed_wrapper_has_storage_id, + wrapper_storage_id as _wrapper_storage_id, +) from policyengine_simulation_executor import artifact_keys as ak @@ -175,3 +181,179 @@ def test_baseline_identity_composes(self, stub_identity_sources): assert identity.store_path == ( f"baselines/us/{_BASELINE_GOLDEN}/bl1-f9cac05d94509895.h5" ) + + +_SPM_SELECTION = { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, +} +_INSTALLED_WRAPPER_HAS_STORAGE_ID = installed_wrapper_has_storage_id() +_SPM_STORAGE_GOLDEN = ( + "bl1-21f52b30719e20bb-spm-" + "7396bf5f4876c42bb6cba0f9533478098edc0861cc3657bd0d60f88dbb26ac39" +) + + +class TestWrapperStorageIdAgreement: + """The planner's storage id against the wrapper that names the file. + + Precompute plans a store path from ``BaselineArtifactIdentity.storage_id`` + and the in-container worker refuses to publish when the wrapper's own + ``Simulation.storage_id`` disagrees. The wrapper's value also names the + saved ``.h5``, so the two derivations are not merely compared — they are + the same identifier reached down two independent code paths, and a + disagreement blocks every canonical publish. + + The ``policyengine`` this project pins is pre-canonical: it has neither + an ``spm`` field nor a ``storage_id``, so an SPM-capable wrapper cannot + be imported in hermetic CI and these tests cannot prove agreement with + one. What they do prove: + + * the no-selection arm agrees through the exact accessor ``precompute`` + uses, against the real installed object — the same answer on either + wrapper, which is the point: that arm must not move; + * the SPM arm agrees with the canonical wrapper's expression as read from + the wheel above, so our side cannot drift from the contract without a + reviewable diff, and the digest cannot be quietly reformatted; + * ``test_spm_arm_matches_the_installed_wrapper`` stops skipping and + asserts real equality the moment an SPM-capable wrapper is pinned. It + supplies the two things that wrapper's ``spm_config`` needs and this + project cannot assume — a US model version, and a bundle pinning this + selection — because ``spm_config`` refuses a non-US model and + re-resolves the selection through ``get_current_bundle`` rather than + reading it off the object. + + Agreement with the real wrapper was checked out of band on 2026-09-11 by + running the executor environment with that wheel's ``policyengine`` + shadowing the pinned one, then resolving seven selection shapes (unset, + empty, national, county, metro, a moved ``as_of``, and a non-ASCII + scenario) through both the wheel's ``resolve_spm_selection`` and this + repo's, and comparing the resolved configs and the storage ids. All + seven agreed on both. That is a recorded observation, not coverage — + only the native lane re-runs anything like it. + + The wheel's sha256 is load-bearing, not decoration: two builds both + named ``policyengine-5.3.0-py3-none-any.whl`` are available locally, and + the other one (``962882ea…``) has no ``storage_id`` at all. + """ + + @pytest.fixture + def identity(self, stub_identity_sources): + """Build the identity directly from an already-resolved selection. + + ``collect_baseline_identity`` resolves the selection through the + installed bundle first; that resolution is the contract suite's + subject. What is under test here is only what the identity then + does with the resolved dict. + """ + + def _identity(spm=None): + return ak.BaselineArtifactIdentity( + spm=spm, + country="us", + region="national", + scope_key=_BASELINE_KWARGS["scope_key"], + dataset=ak.collect_dataset_identity("us", 2026), + ) + + return _identity + + @pytest.mark.parametrize( + "selection", + [ + None, + _SPM_SELECTION, + {**_SPM_SELECTION, "geography_kind": "county"}, + {**_SPM_SELECTION, "geography_kind": "metro", "geography_id": "35620"}, + {**_SPM_SELECTION, "as_of": "2025-01-01"}, + # Non-ASCII is the one input where our explicit ensure_ascii=True + # could diverge from the wrapper's json.dumps defaults. + {**_SPM_SELECTION, "scenario": "ce_trend_ü"}, + ], + ) + def test_storage_id_matches_the_wrapper_expression(self, identity, selection): + built = identity(selection) + assert built.storage_id == _wrapper_storage_id(built.simulation_id, selection) + + def test_spm_storage_id_golden(self, identity): + """Freeze the string. Changing it rotates every canonical artifact.""" + built = identity(_SPM_SELECTION) + assert built.storage_id == _SPM_STORAGE_GOLDEN + assert built.store_path.endswith(f"/{_SPM_STORAGE_GOLDEN}.h5") + + @pytest.mark.parametrize("field", sorted(_SPM_SELECTION)) + def test_every_selection_field_rotates_the_storage_id(self, identity, field): + perturbed = {**_SPM_SELECTION, field: "other"} + assert identity(perturbed).storage_id != _SPM_STORAGE_GOLDEN + + def test_storage_id_ignores_selection_key_order(self, identity): + reversed_selection = dict(reversed(list(_SPM_SELECTION.items()))) + assert identity(reversed_selection).storage_id == _SPM_STORAGE_GOLDEN + + def test_no_selection_keeps_the_plain_simulation_id(self, identity): + built = identity(None) + assert built.storage_id == built.simulation_id + assert built.store_path.endswith(f"/{built.simulation_id}.h5") + + def test_legacy_arm_matches_the_wrapper_accessor(self, identity): + """The no-selection arm, through the accessor precompute uses. + + ``precompute`` reads ``getattr(baseline, "storage_id", baseline.id)``. + With no resolved selection that has to be the planned id on *any* + wrapper: pre-canonical, because the attribute is absent and the + fallback is the id; canonical, because the property short-circuits + to the id when ``spm_config`` is None. The assertion is the same + either way, so this pins the accessor's contract rather than telling + the two wrappers apart. + + "No resolved selection" is the condition, not "the caller sent no + selection". On a canonical bundle the wrapper resolves an unset + ``spm`` into the bundle's defaults and returns a suffixed id — and + so does this project, through the same rule in + ``resolve_spm_selection``, which is why they still agree. A plain id + is what a bundle with no SPM measurement produces. + """ + from policyengine.core import Simulation + + built = identity(None) + wrapper = Simulation.model_construct(id=built.simulation_id) + assert getattr(wrapper, "storage_id", wrapper.id) == built.storage_id + + @pytest.mark.skipif( + not _INSTALLED_WRAPPER_HAS_STORAGE_ID, + reason=( + "The pinned policyengine is pre-canonical and has no storage_id; " + "this asserts real equality as soon as an SPM-capable wrapper is " + "pinned, replacing the transcribed expression above." + ), + ) + def test_spm_arm_matches_the_installed_wrapper(self, identity, monkeypatch): + """Real equality against the wrapper, once one can be imported. + + The canonical ``spm_config`` reads the model version's country and + re-resolves the selection through the installed bundle, so a bare + ``model_construct(id=..., spm=...)`` raises "SPM selection is only + supported by the US model" instead of comparing anything. Both are + supplied here so this activates on a pin rather than erroring. + """ + import policyengine.bundle + from policyengine.core import Simulation + + built = identity(_SPM_SELECTION) + # Only after the planner has read the real installed bundle: this + # stub is for the wrapper's re-resolution, not for ours. + monkeypatch.setattr( + policyengine.bundle, + "get_current_bundle", + lambda: {"measurements": {"spm": _SPM_SELECTION}}, + ) + wrapper = Simulation.model_construct( + id=built.simulation_id, + spm=_SPM_SELECTION, + tax_benefit_model_version=SimpleNamespace(country_code="us"), + ) + assert wrapper.storage_id == built.storage_id diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py index 4520ea18f..d4562bdcc 100644 --- a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts.py @@ -7,11 +7,15 @@ simulation cache is swapped per test. """ +from copy import deepcopy from types import SimpleNamespace import pandas as pd import pytest +from policyengine.core import Simulation +from pydantic import PrivateAttr + from policyengine_simulation_executor import artifact_keys as ak from policyengine_simulation_executor import baseline_artifacts as ba @@ -584,3 +588,381 @@ def test_incomplete_disk_artifact_recomputes_and_overwrites( year=2026, ) assert "employment_income" in reloaded.data.person.columns + + +SPM_SELECTION = { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, +} + + +def _spm_receipt(selection=None, *, year="2026"): + """One calculation receipt, the shape ``spm_provenance()`` returns. + + ``simulation_spm_result`` reads one receipt per simulation and pairs the + baseline's with the reform's itself, so the wrapper hands back a flat + receipt rather than a baseline/reform comparison. + """ + selection = selection or SPM_SELECTION + return { + "forecast_id": "test-only", + "forecast_sha256": selection["forecast_content_sha256"], + "scenario": selection["scenario"], + "geography_kind": selection["geography_kind"], + "runtime_versions": {"policyengine-us": "test-only"}, + "years": {year: {"status": "forecast"}}, + "geographies": [], + "composition_method": "classified-inputs", + "storage_method": "formula", + } + + +def _installed_wrapper_supports_spm() -> bool: + return "spm" in Simulation.model_fields and hasattr(Simulation, "spm_provenance") + + +_INSTALLED_WRAPPER_SUPPORTS_SPM = _installed_wrapper_supports_spm() + + +class SPMWrapperSimulation(Simulation): + """Test double for the canonical wrapper's SPM surface. + + The wrapper the canonical bundle installs adds an ``spm`` field, an + ``spm_config`` holding the selection an artifact was built under, and an + ``spm_provenance()`` calculation receipt; it restores that metadata on a + load or a cache hit. The ``policyengine`` this project pins has none of + it, which is why ``ensure()``'s receipt validation was reachable only + from the SPM_NATIVE_SMOKE_SOURCE-gated tests. + + This double reproduces that surface's *shape and restore timing*, which + is all the guard depends on. It is deliberately more permissive than the + real wrapper, whose own load rejects an artifact built under a different + selection: the guard is written not to trust the wrapper, so it is + exercised here as the independent check it is meant to be. + + It is evidence about ``ensure()``'s decisions, never about the wrapper. + It cannot show that a real ``spm_provenance()`` returns a mapping + ``SPMProvenance`` accepts, that a real load or cache hit restores the + artifact's selection at all, or that a real ``storage_id`` separates two + selections. Those three are the native suite's claims; hermetic green + here is not wrapper conformance. + + ``storage_id`` is a plain field rather than a property derived from the + selection, so it stays put while a load rewrites ``spm`` — the double + makes no claim about how the wrapper computes it, only that the artifact + class keys the process cache on it. + """ + + spm: dict | None = None + spm_config: dict | None = None + spm_receipt: dict | None = None + storage_id: str = "" + _provenance_reads: list = PrivateAttr(default_factory=list) + + def spm_provenance(self): + self._provenance_reads.append(deepcopy(self.spm_config)) + return self.spm_receipt + + def ensure(self): + from policyengine.core.simulation import _cache + + cache_key = self.storage_id or self.id + cached = _cache.get(cache_key) + if cached is not None: + self.output_dataset = cached.output_dataset + # The restore the guard exists to survive: a cached entry's + # selection lands on this request's simulation, replacing the + # one the request asked for. + self.spm = deepcopy(getattr(cached, "spm", None)) + self.spm_config = deepcopy(getattr(cached, "spm_config", None)) + self.spm_receipt = deepcopy(getattr(cached, "spm_receipt", None)) + return + try: + self.tax_benefit_model_version.load(self) + except Exception: + self.run() + self.save() + # The wrapper files a snapshot, not the live object, so re-pointing + # this simulation at another selection cannot rewrite the entry an + # earlier key names. + _cache.add(cache_key, self.model_copy(deep=False)) + + +class CanonicalSPMSimulation(ba.ArtifactBaselineSimulation, SPMWrapperSimulation): + """``ArtifactBaselineSimulation`` over a wrapper that supports SPM.""" + + +class SPMModelVersion(FakeModelVersion): + """Restores an artifact's stored SPM metadata on load, as the wrapper does.""" + + def __init__( + self, *, load_result="complete", stored_config=None, stored_receipt=None + ): + super().__init__(load_result=load_result) + self.stored_config = stored_config + self.stored_receipt = stored_receipt + self.spm_at_run = [] + + def load(self, simulation): + super().load(simulation) + # An artifact carries the selection it was built under, and the + # wrapper restores it over the requested one. + simulation.spm = deepcopy(self.stored_config) + simulation.spm_config = deepcopy(self.stored_config) + simulation.spm_receipt = deepcopy(self.stored_receipt) + + def run(self, simulation): + super().run(simulation) + # A real recompute measures SPM under whatever ``spm`` now holds. + self.spm_at_run.append(deepcopy(simulation.spm)) + simulation.spm_config = deepcopy(simulation.spm) + simulation.spm_receipt = _spm_receipt(simulation.spm) + + +def _make_spm_sim( + model_version, *, spm=None, sim_id="bl1-spm", storage_id=None, year=2026 +): + selection = SPM_SELECTION if spm is None else spm + return CanonicalSPMSimulation.model_construct( + id=sim_id, + storage_id=storage_id or f"{sim_id}-{selection['scenario']}", + dataset=SimpleNamespace(year=year), + tax_benefit_model_version=model_version, + policy=None, + dynamic=None, + scoping_strategy=None, + extra_variables={}, + output_dataset=None, + spm=deepcopy(SPM_SELECTION) if spm is None else deepcopy(spm), + # The wrapper carries the selection it was configured with from + # construction; a load or cache hit then overwrites it with whatever + # the artifact was built under, which is exactly what the guard in + # ``ensure()`` compares against the pre-load value. + spm_config=deepcopy(SPM_SELECTION) if spm is None else deepcopy(spm), + spm_receipt=None, + ) + + +class TestEnsureValidatesSPMReceipts: + """``ensure()``'s receipt validation, driven by fixtures rather than a model. + + Only the opt-in native tests reached this block before, so hermetic CI + could not tell a working guard from a dead one. These cases fix the + guard's *decisions*; the native tests remain the only evidence that real + wrapper receipts have the shape the decisions are made on. + """ + + def test_matching_receipt_loads_as_a_hit(self, fresh_cache): + model = SPMModelVersion( + stored_config=deepcopy(SPM_SELECTION), stored_receipt=_spm_receipt() + ) + sim = _make_spm_sim(model) + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_HIT + assert model.calls == ["load"] + + def test_artifact_built_under_another_selection_recomputes(self, fresh_cache): + other = {**SPM_SELECTION, "scenario": "zero_real"} + model = SPMModelVersion(stored_config=other, stored_receipt=_spm_receipt(other)) + sim = _make_spm_sim(model) + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_INCOMPLETE + assert model.calls == ["load", "run", "save"] + # The requested selection is restored before the recompute, so the + # new artifact is measured under the request, not the artifact. + assert model.spm_at_run == [SPM_SELECTION] + assert sim.spm_config == SPM_SELECTION + + def test_artifact_without_a_receipt_recomputes(self, fresh_cache): + model = SPMModelVersion( + stored_config=deepcopy(SPM_SELECTION), stored_receipt=None + ) + sim = _make_spm_sim(model) + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_INCOMPLETE + assert model.calls == ["load", "run", "save"] + + def test_malformed_receipt_recomputes(self, fresh_cache): + malformed = _spm_receipt() + del malformed["composition_method"] + model = SPMModelVersion( + stored_config=deepcopy(SPM_SELECTION), stored_receipt=malformed + ) + sim = _make_spm_sim(model) + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_INCOMPLETE + assert model.calls == ["load", "run", "save"] + + def test_receipt_from_another_forecast_recomputes(self, fresh_cache): + stale = _spm_receipt() + stale["forecast_sha256"] = "b" * 64 + model = SPMModelVersion( + stored_config=deepcopy(SPM_SELECTION), stored_receipt=stale + ) + sim = _make_spm_sim(model) + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_INCOMPLETE + assert model.calls == ["load", "run", "save"] + + def test_receipt_for_another_year_recomputes(self, fresh_cache): + model = SPMModelVersion( + stored_config=deepcopy(SPM_SELECTION), + stored_receipt=_spm_receipt(year="2024"), + ) + sim = _make_spm_sim(model) + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_INCOMPLETE + assert model.calls == ["load", "run", "save"] + + def test_cached_entry_with_an_unusable_receipt_is_revalidated(self, fresh_cache): + """A cache hit skips the load entirely, so the entry's receipt is + the only thing standing between the request and someone else's + output. A tax-only run leaves a receipt with no years under the same + selection; the guard has to notice and recompute.""" + stale = _make_spm_sim(SPMModelVersion()) + stale.output_dataset = _output(_complete_frames()) + stale.spm_receipt = _spm_receipt(year="2024") + fresh_cache.add(stale.storage_id, stale) + + model = SPMModelVersion() + sim = _make_spm_sim(model) + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_INCOMPLETE + # Cache hit short-circuits the load; the guard still forces the run. + assert model.calls == ["run", "save"] + + def test_cached_selection_cannot_replace_the_request(self, fresh_cache): + """Defense in depth: an entry whose selection disagrees with the key + it is filed under. A correct wrapper never writes one — the key is + derived from the selection — but a cache hit restores the entry's + selection onto this request without revalidating, so the guard + keeps the request's own selection and recomputes under it.""" + other = {**SPM_SELECTION, "scenario": "zero_real"} + stale = _make_spm_sim(SPMModelVersion(), spm=other) + stale.output_dataset = _output(_complete_frames()) + stale.spm_receipt = _spm_receipt(other) + request = _make_spm_sim(SPMModelVersion()) + fresh_cache.add(request.storage_id, stale) + + model = SPMModelVersion() + sim = _make_spm_sim(model) + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_INCOMPLETE + assert model.calls == ["run", "save"] + assert model.spm_at_run == [SPM_SELECTION] + assert sim.spm_config == SPM_SELECTION + + def test_recompute_replaces_the_cache_with_this_request_s_selection( + self, fresh_cache + ): + other = {**SPM_SELECTION, "scenario": "zero_real"} + model = SPMModelVersion(stored_config=other, stored_receipt=_spm_receipt(other)) + sim = _make_spm_sim(model) + sim.ensure() + assert fresh_cache.get(sim.storage_id).spm_config == SPM_SELECTION + + # The next request in this container hits the replaced entry and + # validates clean: one recompute for the container, not one each. + later_model = SPMModelVersion(stored_config=other) + later = _make_spm_sim(later_model) + later.ensure() + assert later.artifact_outcome == ba.OUTCOME_HIT + assert later_model.calls == [] + + def test_missing_columns_are_decided_before_the_receipt(self, fresh_cache): + """A tax-only artifact can legitimately carry no usable receipt. + Columns are checked first so the outcome is attributed to the gap + that actually exists, and the receipt is never consulted.""" + model = SPMModelVersion( + load_result="incomplete", + stored_config=deepcopy(SPM_SELECTION), + stored_receipt=None, + ) + sim = _make_spm_sim(model) + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_INCOMPLETE + assert sim._provenance_reads == [] + + def test_a_request_without_a_selection_never_validates_receipts(self, fresh_cache): + """The legacy path: no selection, so a complete artifact is a hit + even though the wrapper surface exists and holds no receipt.""" + model = SPMModelVersion(stored_config=None, stored_receipt=None) + sim = _make_spm_sim(model, sim_id="bl1-legacy") + sim.spm = None + sim.spm_config = None + sim.ensure() + assert sim.artifact_outcome == ba.OUTCOME_HIT + assert model.calls == ["load"] + assert sim._provenance_reads == [] + + def test_the_cache_is_keyed_by_storage_id_not_simulation_id(self, fresh_cache): + """One caller id must not serve two selections. + + The deterministic simulation id is shared by every selection; only + the storage id separates them, and the artifact class evicts and + replaces under it. Keyed on the id instead, a national recompute + would answer the next county request from cache. + """ + other = {**SPM_SELECTION, "scenario": "zero_real"} + first_model = SPMModelVersion( + stored_config=other, stored_receipt=_spm_receipt(other) + ) + first = _make_spm_sim(first_model) + first.ensure() + assert first.artifact_outcome == ba.OUTCOME_INCOMPLETE + + second_model = SPMModelVersion( + stored_config=other, stored_receipt=_spm_receipt(other) + ) + second = _make_spm_sim(second_model, spm=other) + assert second.id == first.id + assert second.storage_id != first.storage_id + second.ensure() + # It read its own artifact rather than the entry the first request + # left behind, and both entries survive side by side. + assert second_model.calls == ["load"] + assert second.artifact_outcome == ba.OUTCOME_HIT + # The artifact class filed the recompute under the storage id, and + # nothing was ever filed under the bare simulation id. + assert fresh_cache.get(first.storage_id).spm_config == SPM_SELECTION + assert fresh_cache.get(second.storage_id) is not None + assert fresh_cache.get(first.id) is None + + def test_the_replaced_cache_entry_is_a_snapshot(self, fresh_cache): + """The replacement is a copy so that re-pointing this simulation at + another selection and running it again cannot overwrite the output + the earlier key names.""" + other = {**SPM_SELECTION, "scenario": "zero_real"} + model = SPMModelVersion(stored_config=other, stored_receipt=_spm_receipt(other)) + sim = _make_spm_sim(model) + sim.ensure() + cached = fresh_cache.get(sim.storage_id) + assert cached is not sim + + # A later request in the same process re-points the live object. + model.load(sim) + assert sim.spm_config == other + assert cached.spm_config == SPM_SELECTION + + @pytest.mark.skipif( + not _INSTALLED_WRAPPER_SUPPORTS_SPM, + reason=( + "The pinned policyengine is pre-canonical; this checks the double " + "above against the real surface as soon as an SPM-capable wrapper " + "is pinned, at which point the double should be retired for it." + ), + ) + def test_the_double_matches_the_installed_wrapper_surface(self): + from policyengine.core import Simulation + + assert "spm" in Simulation.model_fields + assert callable(getattr(Simulation, "spm_provenance", None)) + assert isinstance( + getattr(type(Simulation), "storage_id", None) + or Simulation.__dict__.get("storage_id"), + property, + ) diff --git a/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py new file mode 100644 index 000000000..45e9ed271 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_baseline_artifacts_spm_native.py @@ -0,0 +1,163 @@ +"""Real one-household SPM artifacts with deliberate local corruption. + +The producer always calculates authentic receipts with the installed country and +calculator. Corruption cases then damage only that temporary result artifact. +""" + +import json +import os +from pathlib import Path + +import pytest + +from native_spm_support import materialize_native_household + +pytestmark = pytest.mark.skipif( + not os.environ.get("SPM_NATIVE_SMOKE_SOURCE"), + reason="Requires the explicitly selected local native H5 and installed SPM bundle", +) + + +@pytest.fixture +def native_dataset(tmp_path): + return materialize_native_household(os.environ["SPM_NATIVE_SMOKE_SOURCE"], tmp_path) + + +def _baseline(dataset, *, geography="national", tax_only=False): + from policyengine_simulation_executor.baseline_artifacts import ( + ArtifactBaselineSimulation, + ) + from policyengine_simulation_executor.simulation_runtime import _build_simulation + + built = _build_simulation( + { + "country": "us", + "scope": "macro", + "time_period": "2024", + "data": "bounded-native-test", + "spm": {"geography_kind": geography}, + }, + dataset=dataset, + policy=None, + region_code="us", + ) + baseline = ArtifactBaselineSimulation.model_construct(**built.__dict__) + baseline.id = "bl1-native-receipt" + if tax_only: + # Produce a genuine lazy, tax-only output under the same selection. + model = baseline.tax_benefit_model_version + baseline.tax_benefit_model_version = model.model_copy( + update={ + "entity_variables": { + entity: ["income_tax"] if entity == "tax_unit" else [] + for entity in model.entity_variables + } + } + ) + return baseline + + +@pytest.mark.parametrize( + "damage, expected_outcome", + [ + ("tax_only_columns", "incomplete"), + ("empty_years", "incomplete"), + ("missing_receipt", "miss"), + ("malformed_receipt", "miss"), + ("different_selection", "miss"), + ("corrupt_hdf", "miss"), + ], +) +def test_incomplete_or_corrupt_spm_disk_artifact_recomputes( + native_dataset, monkeypatch, damage, expected_outcome +): + import h5py + from policyengine.core.simulation import _cache + from policyengine_simulation_executor.baseline_artifacts import ( + ArtifactBaselineSimulation, + ) + from policyengine_simulation_executor.spm import simulation_spm_result + + _cache._cache.clear() + producer = _baseline(native_dataset, tax_only=damage == "tax_only_columns") + producer.run() + producer.save() + path = Path(producer.output_dataset.filepath) + if damage == "tax_only_columns": + assert producer.spm_provenance()["years"] == {} + assert ( + "spm_unit_is_in_spm_poverty" + not in producer.output_dataset.data.spm_unit.columns + ) + elif damage == "corrupt_hdf": + path.write_bytes(b"deliberately corrupted test artifact") + else: + with h5py.File(path, "a") as stream: + recorded = json.loads(stream["policyengine_spm"].asstr()[()]) + del stream["policyengine_spm"] + if damage == "empty_years": + recorded["provenance"]["years"] = {} + elif damage == "malformed_receipt": + recorded["provenance"]["unexpected_test_field"] = True + elif damage == "different_selection": + recorded["config"]["geography_kind"] = "county" + if damage != "missing_receipt": + stream.create_dataset( + "policyengine_spm", + data=json.dumps(recorded), + dtype=h5py.string_dtype("utf-8"), + ) + + runs = [] + real_run = ArtifactBaselineSimulation.run + + def counted_run(simulation): + runs.append(simulation) + return real_run(simulation) + + monkeypatch.setattr(ArtifactBaselineSimulation, "run", counted_run) + consumer = _baseline(native_dataset) + requested = consumer.spm_config + consumer.ensure() + assert consumer.artifact_outcome == expected_outcome + assert runs == [consumer] + assert consumer.spm_config == requested + assert not consumer._missing_output_columns() + simulation_spm_result(consumer, consumer, requested, expected_year=2024) + # A second request uses the repaired cache; clearing it tests repaired disk. + cached = _baseline(native_dataset) + cached.ensure() + assert cached.artifact_outcome == "hit" + _cache._cache.clear() + reloaded = _baseline(native_dataset) + reloaded.ensure() + assert reloaded.artifact_outcome == "hit" + assert runs == [consumer] + assert reloaded.spm_provenance() == consumer.spm_provenance() + + +@pytest.mark.parametrize("damage", ["missing_receipt", "different_selection"]) +def test_invalid_spm_cache_entry_recomputes_requested_selection(native_dataset, damage): + from policyengine.core.simulation import _cache + from policyengine_simulation_executor.spm import simulation_spm_result + + _cache._cache.clear() + producer = _baseline( + native_dataset, + geography="county" if damage == "different_selection" else "national", + ) + producer.run() + if damage == "missing_receipt": + producer.spm_receipt = None + consumer = _baseline(native_dataset) + requested = consumer.spm_config + # Deliberate local cache corruption must never replace the request selection. + _cache.add(consumer.storage_id, producer) + consumer.ensure() + assert consumer.artifact_outcome == "incomplete" + assert consumer.spm_config == requested + simulation_spm_result(consumer, consumer, requested, expected_year=2024) + later = _baseline(native_dataset) + later.ensure() + assert later.artifact_outcome == "hit" + assert later.spm_provenance() == consumer.spm_provenance() diff --git a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py index 475c80049..4364f86d8 100644 --- a/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py +++ b/projects/policyengine-simulation-executor/tests/test_budget_window_scheduler.py @@ -17,9 +17,43 @@ import src.modal.budget_window_batch as batch_module import src.modal.budget_window_scheduler as scheduler_module import policyengine_simulation_contract.budget_window_state as state_module +from policyengine_simulation_contract.budget_window_state import ( + BUDGET_WINDOW_JOB_DICT_NAME, + BUDGET_WINDOW_JOB_SEED_DICT_NAME, +) +from policyengine_simulation_contract.spm import SPMInputError, SPMSelection from policyengine_simulation_gateway.testing import create_gateway_app from policyengine_simulation_gateway import endpoints +SPM_SELECTION = SPMSelection( + forecast_content_sha256="a" * 64, + scenario="ce_trend", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, +).model_dump(mode="json") + + +def spm_child_result(runtime, simulation_year, *, receipt_year): + """A child result carrying a canonical SPM receipt for ``receipt_year``.""" + receipt = { + "forecast_id": "test-only", + "forecast_sha256": SPM_SELECTION["forecast_content_sha256"], + "scenario": SPM_SELECTION["scenario"], + "geography_kind": SPM_SELECTION["geography_kind"], + "runtime_versions": {"policyengine-us": "test-only"}, + "years": {receipt_year: {"status": "forecast"}}, + "geographies": [], + "composition_method": "classified-inputs", + "storage_method": "formula", + } + return { + "spm_config": dict(SPM_SELECTION), + "spm_provenance": {"baseline": [receipt], "reform": [dict(receipt)]}, + **runtime.child_result_for_year(simulation_year), + } + @dataclass class SemiIntegrationRuntime: @@ -30,6 +64,15 @@ class SemiIntegrationRuntime: next_parent_call_id: str = "parent-batch-123" active_child_calls: set[str] = field(default_factory=set) max_active_child_calls: int = 0 + # Per-year injection seams: a child call raises ``child_errors[year]`` + # instead of returning, or returns ``child_results[year]`` verbatim. + child_errors: dict[str, BaseException] = field(default_factory=dict) + child_results: dict[str, dict] = field(default_factory=dict) + + def child_outcome_for_year(self, simulation_year: str) -> dict: + return self.child_results.get( + simulation_year, self.child_result_for_year(simulation_year) + ) def child_result_for_year(self, simulation_year: str) -> dict: offset = int(simulation_year) - 2025 @@ -69,15 +112,23 @@ def get(self, key: str, default=None): class MockChildCall: def __init__( - self, runtime: SemiIntegrationRuntime, *, object_id: str, result: dict + self, + runtime: SemiIntegrationRuntime, + *, + object_id: str, + result: dict, + error: BaseException | None = None, ): self.runtime = runtime self.object_id = object_id self.result = result + self.error = error self.runtime.child_started(object_id) def get(self, timeout: int = 0): self.runtime.child_finished(self.object_id) + if self.error is not None: + raise self.error return self.result @@ -124,7 +175,8 @@ def spawn(self, payload: dict): call = MockChildCall( self.runtime, object_id=f"child-{simulation_year}", - result=self.runtime.child_result_for_year(simulation_year), + result=self.runtime.child_outcome_for_year(simulation_year), + error=self.runtime.child_errors.get(simulation_year), ) self.runtime.calls[call.object_id] = call return call @@ -246,3 +298,183 @@ def test_budget_window_submit_and_poll_exercise_gateway_worker_seams( assert all("window_size" not in payload for payload in runtime.child_payloads) assert all("max_parallel" not in payload for payload in runtime.child_payloads) assert all("_metadata" not in payload for payload in runtime.child_payloads) + + +def submit_budget_window(client, *, window_size=3, max_parallel=1): + response = client.post( + "/simulate/economy/budget-window", + json={ + "country": "us", + "region": "us", + "scope": "macro", + "reform": {}, + "start_year": "2026", + "window_size": window_size, + "max_parallel": max_parallel, + }, + ) + assert response.status_code == 200, response.text + return response.json()["batch_job_id"] + + +def test_typed_child_call_failure_persists_errors_and_stops_the_batch( + budget_window_semi_integration_client, +): + """A child that raises ``SPMInputError`` is transported, not redacted.""" + client, runtime = budget_window_semi_integration_client + runtime.child_errors["2026"] = SPMInputError( + "SPM_GEOGRAPHY_REQUIRED", "County required" + ) + + batch_job_id = submit_budget_window(client) + assert client.get(f"/budget-window-jobs/{batch_job_id}").status_code == 202 + + failure = client.get(f"/budget-window-jobs/{batch_job_id}") + assert failure.status_code == 400 + body = failure.json() + expected = [{"code": "SPM_GEOGRAPHY_REQUIRED", "message": "County required"}] + assert body["status"] == "failed" + assert body["errors"] == expected + assert body["error"] == "County required" + # ``mark_child_failed`` replaces the child entry, so the scheduler has to + # re-attach the typed payload after it. + assert body["child_jobs"]["2026"]["errors"] == expected + assert body["child_jobs"]["2026"]["status"] == "failed" + assert body["failed_years"] == ["2026"] + # The scheduler returns early: the remaining years are never spawned. + assert body["queued_years"] == ["2027", "2028"] + assert [payload["time_period"] for payload in runtime.child_payloads] == ["2026"] + + +def test_typed_result_validation_failure_persists_errors_and_stops_the_batch( + budget_window_semi_integration_client, +): + """A mismatched SPM receipt fails the batch with the typed code.""" + client, runtime = budget_window_semi_integration_client + batch_job_id = submit_budget_window(client) + + # The gateway records the resolved selection on the seed whenever the + # route advertises a canonical SPM capability; the parent reads it back + # from the seed to validate each child receipt. Seeding it here keeps the + # test on the scheduler seam instead of the registry's. + seed = runtime.dicts[BUDGET_WINDOW_JOB_SEED_DICT_NAME][batch_job_id] + seed["request_payload"]["spm"] = dict(SPM_SELECTION) + runtime.child_results["2026"] = spm_child_result( + runtime, "2026", receipt_year="1999" + ) + + assert client.get(f"/budget-window-jobs/{batch_job_id}").status_code == 202 + + failure = client.get(f"/budget-window-jobs/{batch_job_id}") + assert failure.status_code == 400 + body = failure.json() + expected = [ + { + "code": "SPM_CONFIGURATION_UNAVAILABLE", + "message": "Result SPM provenance does not cover the requested year", + } + ] + assert body["status"] == "failed" + assert body["errors"] == expected + assert body["child_jobs"]["2026"]["errors"] == expected + assert body["failed_years"] == ["2026"] + assert body["queued_years"] == ["2027", "2028"] + assert [payload["time_period"] for payload in runtime.child_payloads] == ["2026"] + + +def test_a_typed_child_failure_cancels_the_siblings_already_running( + budget_window_semi_integration_client, +): + """The early return, where a single-parallel batch cannot show it. + + ``poll_running_children_once`` returns False on a typed failure instead + of continuing round the loop. With ``max_parallel=1`` there is never a + second running child, so replacing that return with ``continue`` changed + nothing observable. At two, the sibling that was already running must + be cancelled rather than harvested: the batch has failed, and a + completed 2027 alongside a failed 2026 would be a partial window nobody + asked for. + """ + client, runtime = budget_window_semi_integration_client + runtime.child_errors["2026"] = SPMInputError( + "SPM_GEOGRAPHY_REQUIRED", "County required" + ) + + batch_job_id = submit_budget_window(client, max_parallel=2) + assert client.get(f"/budget-window-jobs/{batch_job_id}").status_code == 202 + + body = client.get(f"/budget-window-jobs/{batch_job_id}").json() + + assert [payload["time_period"] for payload in runtime.child_payloads] == [ + "2026", + "2027", + ] + assert body["completed_years"] == [] + assert body["child_jobs"]["2027"]["status"] == "cancelled" + assert body["failed_years"] == ["2026"] + assert body["queued_years"] == ["2028"] + + +def test_a_typed_child_failure_is_persisted_not_only_returned( + budget_window_semi_integration_client, +): + """The poll body is serialized from the state run() returns in memory. + + Both typed branches write the re-attached errors back through + ``put_batch_job_state`` before returning, and nothing else re-reads the + batch-state dict, so deleting those writes left the suite green. A + later poll -- a different container, or the same one after a restart -- + reads the dict, so what is in it is what the client eventually sees. + """ + client, runtime = budget_window_semi_integration_client + runtime.child_errors["2026"] = SPMInputError( + "SPM_GEOGRAPHY_REQUIRED", "County required" + ) + + batch_job_id = submit_budget_window(client) + assert client.get(f"/budget-window-jobs/{batch_job_id}").status_code == 202 + assert client.get(f"/budget-window-jobs/{batch_job_id}").status_code == 400 + + persisted = runtime.dicts[BUDGET_WINDOW_JOB_DICT_NAME][batch_job_id] + expected = [{"code": "SPM_GEOGRAPHY_REQUIRED", "message": "County required"}] + assert persisted["status"] == "failed" + assert persisted["errors"] == expected + assert persisted["child_jobs"]["2026"]["errors"] == expected + + +def test_untyped_child_failure_is_redacted_and_carries_no_typed_errors( + budget_window_semi_integration_client, +): + """Only public typed errors escape redaction; others stay 500 with no code.""" + client, runtime = budget_window_semi_integration_client + runtime.child_errors["2026"] = RuntimeError("internal detail") + + batch_job_id = submit_budget_window(client) + assert client.get(f"/budget-window-jobs/{batch_job_id}").status_code == 202 + + failure = client.get(f"/budget-window-jobs/{batch_job_id}") + assert failure.status_code == 500 + body = failure.json() + assert "errors" not in body + assert "internal detail" not in body["error"] + assert "errors" not in body["child_jobs"]["2026"] + + +def test_untyped_result_parsing_failure_is_redacted_and_stops_the_batch( + budget_window_semi_integration_client, +): + """A malformed child result fails the batch without leaking the reason.""" + client, runtime = budget_window_semi_integration_client + runtime.child_results["2026"] = {"budget": "not-an-object"} + + batch_job_id = submit_budget_window(client) + assert client.get(f"/budget-window-jobs/{batch_job_id}").status_code == 202 + + failure = client.get(f"/budget-window-jobs/{batch_job_id}") + assert failure.status_code == 500 + body = failure.json() + assert "errors" not in body + assert "missing budget object" not in body["error"] + assert body["failed_years"] == ["2026"] + assert body["queued_years"] == ["2027", "2028"] + assert [payload["time_period"] for payload in runtime.child_payloads] == ["2026"] diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py new file mode 100644 index 000000000..cd990d087 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm.py @@ -0,0 +1,857 @@ +"""Bounded canonical worker contracts; no population simulation or Modal calls.""" + +import json +import pickle +import sys +from copy import deepcopy +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from policyengine_simulation_contract.gateway_models import ( + SimulationRequest, + BudgetWindowBatchRequest, +) +from policyengine_simulation_contract.spm import ( + SPMCapability, + SPMSelection, + SPMInputError, + resolve_spm_selection, + combine_spm_results, +) +from policyengine_simulation_executor import artifact_keys, simulation_runtime +from policyengine_simulation_executor.simulation import create_router +from src.modal.budget_window_results import ( + extract_annual_impact, + build_budget_window_result, +) +from src.modal.fanout import build_child_payload + +SELECTION = SPMSelection( + forecast_content_sha256="a" * 64, + scenario="ce_trend", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, +).model_dump() +CAPABILITY = SPMCapability(defaults=SELECTION).model_dump() + + +def result(selection=SELECTION, year="2026"): + receipt = dict( + forecast_id="test-only", + forecast_sha256=selection["forecast_content_sha256"], + scenario=selection["scenario"], + geography_kind=selection["geography_kind"], + runtime_versions={"policyengine-us": "test-only"}, + years={year: {"status": "forecast"}}, + geographies=[], + composition_method="classified-inputs", + storage_method="formula", + ) + return { + "spm_config": deepcopy(selection), + "spm_provenance": {"baseline": [receipt], "reform": [deepcopy(receipt)]}, + "budget": { + "tax_revenue_impact": 1, + "benefit_spending_impact": 2, + "budgetary_impact": -1, + }, + } + + +@pytest.mark.parametrize( + "model,extra", + [ + (SimulationRequest, {}), + ( + BudgetWindowBatchRequest, + {"region": "us", "start_year": "2026", "window_size": 2}, + ), + ], +) +def test_public_request_roundtrip(model, extra): + parsed = model(country="us", spm={"geography_kind": "national"}, **extra) + assert json.loads(parsed.model_dump_json())["spm"]["geography_kind"] == "national" + with pytest.raises(ValueError): + model(country="us", spm={"provider_path": "x"}, **extra) + with pytest.raises(ValueError): + model( + country="us", + spm={"geography_kind": "national", "geography_id": "12345"}, + **extra, + ) + + +def test_unknown_bundle_fails_closed_and_historical_uk_remain(): + for country, version, model in [ + ("us", "5.2.0", "1.764.6"), + ("us", "5.3.0", "1.764.6"), + ("uk", "future-test", "test"), + ]: + assert ( + resolve_spm_selection( + country, + None, + capability=None, + policyengine_version=version, + model_version=model, + ) + is None + ) + for selection in (None, {}, {"geography_kind": "national"}): + with pytest.raises(SPMInputError, match="no certified"): + resolve_spm_selection( + "us", + selection, + capability=None, + policyengine_version="future-test", + model_version="test", + ) + assert ( + resolve_spm_selection( + "us", + {}, + capability=CAPABILITY, + policyengine_version="test-only", + model_version="test", + ) + == SELECTION + ) + with pytest.raises(SPMInputError, match="hash"): + resolve_spm_selection( + "us", + {"forecast_content_sha256": "b" * 64}, + capability=CAPABILITY, + policyengine_version="test-only", + model_version="test", + ) + + +@pytest.mark.parametrize( + "kind,area", [("county", None), ("national", None), ("metro", "35620")] +) +def test_partial_selection_survives_ordinary_nested_json(kind, area): + capability = { + "contract_version": "canonical-spm-v1", + "defaults": { + **SELECTION, + "geography_kind": kind, + "geography_id": area, + "county_vintage": "2010", + "as_of": "2026-01-01", + }, + } + partial = {"scenario": "zero_real"} + request = SimulationRequest(country="us", spm=partial) + wire = json.loads(request.model_dump_json()) + assert wire["spm"] == partial + resolved = resolve_spm_selection( + "us", + wire["spm"], + capability=capability, + policyengine_version="test", + model_version="test", + ) + assert resolved == {**capability["defaults"], **partial} + assert SPMSelection.model_validate(resolved).model_dump() == resolved + explicit = {"geography_kind": "county", "county_vintage": "2020", "as_of": None} + assert resolve_spm_selection( + "us", + json.loads(SPMSelection(**explicit).model_dump_json()), + capability=capability, + policyengine_version="test", + model_version="test", + ) == {**capability["defaults"], **explicit, "geography_id": None} + + +def test_historical_allowance_does_not_certify_other_models(): + for version, model in [("5.3.0", "1.824.7"), ("5.3.1", "1.764.6")]: + with pytest.raises(SPMInputError, match="no certified"): + resolve_spm_selection( + "us", + None, + capability=None, + policyengine_version=version, + model_version=model, + ) + + +def test_annual_child_and_segment_receipts_must_cover_requested_year(): + wrong_year = result(year="2025") + with pytest.raises(SPMInputError, match="requested year"): + extract_annual_impact( + simulation_year="2026", child_result=wrong_year, spm=SELECTION + ) + with pytest.raises(SPMInputError, match="requested year"): + combine_spm_results([result(), wrong_year], SELECTION, expected_year=2026) + + +def test_result_transport_may_omit_nulls_but_not_resolved_options(): + transported = result() + transported["spm_config"] = { + key: value + for key, value in transported["spm_config"].items() + if value is not None + } + assert combine_spm_results([transported], SELECTION)["spm_config"] == SELECTION + for field in ( + "forecast_content_sha256", + "scenario", + "geography_kind", + "county_vintage", + ): + malformed = deepcopy(transported) + del malformed["spm_config"][field] + with pytest.raises(SPMInputError, match="complete resolved"): + combine_spm_results([malformed], SELECTION) + + +def test_sync_compatibility_endpoint_preserves_explicit_null(monkeypatch): + import policyengine_simulation_executor.simulation as simulation + + captured = [] + + def capture(params): + captured.append(params) + raise SPMInputError("SPM_GEOGRAPHY_REQUIRED", "intentional test stop") + + monkeypatch.setattr(simulation, "run_simulation_impl", capture) + app = FastAPI() + app.include_router(create_router()) + response = TestClient(app).post( + "/simulate/economy/comparison", + json={"country": "us", "spm": {"as_of": None}}, + ) + assert response.status_code == 400 + assert captured[0]["spm"] == {"as_of": None} + + +def test_year_alias_is_validated_before_dataset_loading(monkeypatch): + from policyengine_simulation_executor import release_bundle, spm + + monkeypatch.setattr( + release_bundle, + "get_country_release_bundle", + lambda country: SimpleNamespace( + policyengine_version="test", model_version="test" + ), + ) + monkeypatch.setattr(spm, "runtime_spm_capability", lambda: CAPABILITY) + year_metadata = Mock( + side_effect=SPMInputError("SPM_YEAR_UNAVAILABLE", "2040 is unavailable") + ) + monkeypatch.setattr(spm, "_forecast", lambda sha: object()) + monkeypatch.setitem( + sys.modules, + "spm_calculator.policyengine_adapter", + SimpleNamespace( + PolicyEngineSPMProvider=lambda *a, **k: SimpleNamespace( + year_metadata=year_metadata + ) + ), + ) + with pytest.raises(SPMInputError, match="2040 is unavailable") as error: + spm.normalize_runtime_spm({"country": "us", "year": "2040", "spm": SELECTION}) + assert error.value.code == "SPM_YEAR_UNAVAILABLE" + assert year_metadata.call_args.args == (2040,) + + +class TestResolutionIsNotRepeatedWork: + """One request resolves its selection about six times; pay once. + + ``_run_simulation_impl_core`` normalizes, ``_build_simulation`` + normalizes again, and the baseline and dataset identity collectors + normalize once each. Every one of those re-read the installed + capability and built a throwaway ``PolicyEngineSPMProvider`` to + prevalidate the years. The installed image cannot change under a + running container, and the prevalidation reads nothing but the + selection and the year range, so both are now memoized -- these cases + fix what that must and must not change. + """ + + @staticmethod + def _stub_bundle(monkeypatch): + from policyengine_simulation_executor import release_bundle + + monkeypatch.setattr( + release_bundle, + "get_country_release_bundle", + lambda country: SimpleNamespace( + policyengine_version="test", model_version="test" + ), + ) + + @staticmethod + def _stub_provider(monkeypatch, capability=CAPABILITY): + """Count the two pieces of work a resolution repeats.""" + import policyengine_simulation_executor.spm as spm + + year_metadata = Mock(return_value={}) + constructions = Mock( + side_effect=lambda *a, **k: SimpleNamespace(year_metadata=year_metadata) + ) + monkeypatch.setattr(spm, "runtime_spm_capability", lambda: capability) + monkeypatch.setattr(spm, "_forecast", lambda sha: object()) + monkeypatch.setitem( + sys.modules, + "spm_calculator.policyengine_adapter", + SimpleNamespace(PolicyEngineSPMProvider=constructions), + ) + return SimpleNamespace( + constructions=constructions, year_metadata=year_metadata, spm=spm + ) + + def test_one_selection_and_year_range_prevalidates_once(self, monkeypatch): + self._stub_bundle(monkeypatch) + stubs = self._stub_provider(monkeypatch) + request = {"country": "us", "time_period": "2026", "spm": SELECTION} + + resolutions = [stubs.spm.normalize_runtime_spm(request) for _ in range(6)] + + assert resolutions == [SELECTION] * 6 + assert stubs.constructions.call_count == 1 + assert stubs.year_metadata.call_count == 1 + + # The key is the selection and the range, so neither can be served + # the other's answer. + stubs.spm.normalize_runtime_spm({**request, "time_period": "2027"}) + assert stubs.constructions.call_count == 2 + stubs.spm.normalize_runtime_spm( + {**request, "spm": {**SELECTION, "scenario": "zero_real"}} + ) + assert stubs.constructions.call_count == 3 + stubs.spm.normalize_runtime_spm({**request, "window_size": 2}) + assert stubs.constructions.call_count == 4 + # One year each for the three single-year ranges, two for the window. + assert stubs.year_metadata.call_count == 5 + + stubs.spm.reset_spm_runtime_caches() + stubs.spm.normalize_runtime_spm(request) + assert stubs.constructions.call_count == 5 + + def test_a_rejected_selection_is_rejected_every_time(self, monkeypatch): + """Fail-closed cannot be weakened by a cache: exceptions are not kept.""" + self._stub_bundle(monkeypatch) + stubs = self._stub_provider(monkeypatch) + stubs.year_metadata.side_effect = SPMInputError( + "SPM_YEAR_UNAVAILABLE", "2040 is unavailable" + ) + request = {"country": "us", "time_period": "2040", "spm": SELECTION} + + for _ in range(3): + with pytest.raises(SPMInputError) as error: + stubs.spm.normalize_runtime_spm(request) + assert error.value.code == "SPM_YEAR_UNAVAILABLE" + assert stubs.constructions.call_count == 3 + + def test_an_unusable_year_does_not_mask_the_selection(self, monkeypatch): + """The order the memo had to preserve. + + The year range used to be parsed between the provider's construction + and its first use, so a selection the provider itself rejects was + reported as such even when the year was unusable. The memo is keyed + on the range, so the range has to be parsed first; an empty range + keeps the old order for the one request where it differs. + """ + self._stub_bundle(monkeypatch) + stubs = self._stub_provider(monkeypatch) + stubs.constructions.side_effect = ValueError( + "Unknown forecast scenario: made-up" + ) + unusable = {"country": "us", "time_period": "2026-01", "spm": SELECTION} + + with pytest.raises(SPMInputError) as error: + stubs.spm.normalize_runtime_spm(unusable) + assert error.value.code == "SPM_SCENARIO_UNAVAILABLE" + assert stubs.year_metadata.call_count == 0 + + with pytest.raises(SPMInputError) as error: + stubs.spm.normalize_runtime_spm({**unusable, "time_period": "2026"}) + assert error.value.code == "SPM_SCENARIO_UNAVAILABLE" + + def test_an_unusable_year_alone_is_a_settings_error(self, monkeypatch): + self._stub_bundle(monkeypatch) + stubs = self._stub_provider(monkeypatch) + + with pytest.raises(SPMInputError) as error: + stubs.spm.normalize_runtime_spm( + {"country": "us", "time_period": "2026-01", "spm": SELECTION} + ) + assert error.value.code == "SPM_SETTINGS_INVALID" + assert stubs.constructions.call_count == 1 + assert stubs.year_metadata.call_count == 0 + + def test_collecting_a_baseline_identity_prevalidates_once(self, monkeypatch): + """The repeat the review counted, through the real collectors. + + ``collect_baseline_identity`` normalizes the request's selection and + then ``collect_dataset_identity`` normalizes the bundle's default for + the same year. When both resolve to the same selection that is one + prevalidation for the pair, and repeating the collection -- as a + request does -- adds none. + """ + from fixtures.identity_stubs import install_identity_stubs + + install_identity_stubs(monkeypatch) + stubs = self._stub_provider(monkeypatch) + + first = artifact_keys.collect_baseline_identity( + "us", 2026, region="us", scope_key=None, spm={"geography_kind": "national"} + ) + assert first.spm == SELECTION + assert stubs.constructions.call_count == 1 + + second = artifact_keys.collect_baseline_identity( + "us", 2026, region="us", scope_key=None, spm={"geography_kind": "national"} + ) + assert second.storage_id == first.storage_id + assert stubs.constructions.call_count == 1 + + # A request that differs from the bundle default still resolves both. + county = artifact_keys.collect_baseline_identity( + "us", 2026, region="us", scope_key=None, spm={"geography_kind": "county"} + ) + assert county.storage_id != first.storage_id + assert stubs.constructions.call_count == 2 + + def test_the_installed_capability_is_read_once(self, monkeypatch): + import policyengine.bundle + import policyengine_simulation_executor.spm as spm + + reads = Mock(return_value={"measurements": {}}) + monkeypatch.setattr(policyengine.bundle, "get_current_bundle", reads) + + assert [spm.runtime_spm_capability() for _ in range(4)] == [None] * 4 + assert reads.call_count == 1 + + spm.reset_spm_runtime_caches() + assert spm.runtime_spm_capability() is None + assert reads.call_count == 2 + + def test_an_unavailable_capability_is_rederived_every_time(self, monkeypatch): + import policyengine.bundle + import policyengine_simulation_executor.spm as spm + + reads = Mock(side_effect=ValueError("bundle configuration unavailable")) + monkeypatch.setattr(policyengine.bundle, "get_current_bundle", reads) + + for _ in range(3): + with pytest.raises(SPMInputError) as error: + spm.runtime_spm_capability() + assert error.value.code == "SPM_CONFIGURATION_UNAVAILABLE" + assert reads.call_count == 3 + + +@pytest.mark.parametrize( + "code", + [ + "SPM_GEOGRAPHY_REQUIRED", + "SPM_GEOGRAPHY_UNAVAILABLE", + "SPM_COMPOSITION_REQUIRED", + "SPM_YEAR_UNAVAILABLE", + "SPM_SCENARIO_UNAVAILABLE", + ], +) +def test_actual_http_formula_error_contract(monkeypatch, code): + error = SPMInputError(code, "Explicit input is required") + assert pickle.loads(pickle.dumps(error)).to_dict() == error.to_dict() + monkeypatch.setattr( + "policyengine_simulation_executor.simulation.run_simulation_impl", + Mock(side_effect=error), + ) + app = FastAPI() + app.include_router(create_router()) + response = TestClient(app).post( + "/simulate/economy/comparison", + json={"country": "us", "spm": {"geography_kind": "national"}}, + ) + assert response.status_code == 400 + assert response.json()["errors"] == [error.to_dict()] + assert response.json()["result"] is None + + +@pytest.mark.parametrize( + "key,value", + [ + ("forecast_content_sha256", "b" * 64), + ("scenario", "zero_real"), + ("geography_kind", "county"), + ("geography_id", "12345"), + ("county_vintage", "2010"), + ("as_of", "2025-01-01"), + ], +) +def test_each_setting_rotates_artifact_identity(key, value): + common = dict( + country="us", + region="national", + scope_key=None, + dataset_digest="d" * 64, + model_version="test", + policyengine_version="test", + ) + changed = {**SELECTION, key: value} + assert artifact_keys.baseline_key( + **common, spm=SELECTION + ) != artifact_keys.baseline_key(**common, spm=changed) + ds = dict( + country="us", + dataset="test", + year=2026, + data_version="test", + data_artifact_revision="test", + source_sha256=None, + data_build_fingerprint=None, + model_version="test", + policyengine_version="test", + ) + assert artifact_keys.dataset_key(**ds, spm=SELECTION) != artifact_keys.dataset_key( + **ds, spm=changed + ) + + +def test_fanout_and_budget_window_preserve_settings_and_receipts(): + child = build_child_payload( + {"country": "us", "spm": SELECTION, "window_size": 2}, + strip_fields={"window_size"}, + overrides={"time_period": "2027"}, + ) + assert child["spm"] == SELECTION + rows = [ + extract_annual_impact( + simulation_year=year, child_result=result(year=year), spm=SELECTION + ) + for year in ("2026", "2027") + ] + window = build_budget_window_result( + start_year="2026", window_size=2, annual_impacts=rows + ) + dumped = json.loads(window.model_dump_json()) + assert dumped["annualImpacts"][1]["spm_provenance"]["baseline"][0]["years"] == { + "2027": {"status": "forecast"} + } + merged = combine_spm_results([result(), result()], SELECTION) + assert len(merged["spm_provenance"]["baseline"]) == 2 + changed = result({**SELECTION, "scenario": "zero_real"}) + for outputs in ([result(), changed], [result(), {"budget": {}}]): + with pytest.raises(SPMInputError): + combine_spm_results(outputs, SELECTION) + + +def test_worker_build_passes_selection_to_baseline_and_reform(monkeypatch): + import policyengine.core + from policyengine_simulation_executor import baseline_artifacts + import policyengine_simulation_executor.spm as spm + + calls = [] + monkeypatch.setattr(spm, "normalize_runtime_spm", lambda params: SELECTION) + monkeypatch.setattr( + simulation_runtime, "_country_module", lambda _: SimpleNamespace(model="test") + ) + monkeypatch.setattr( + baseline_artifacts, + "deterministic_baseline_id", + lambda *a, **k: "baseline-id" if k["policy"] is None else None, + ) + monkeypatch.setattr( + baseline_artifacts, "ArtifactBaselineSimulation", lambda **k: calls.append(k) + ) + monkeypatch.setattr(policyengine.core, "Simulation", lambda **k: calls.append(k)) + for policy in (None, {"test": 123}): + simulation_runtime._build_simulation( + {"country": "us", "spm": SELECTION}, + dataset="tiny", + policy=policy, + region_code="us", + ) + assert len(calls) == 2 + assert all(call["spm"] == SELECTION for call in calls) + assert calls[0]["id"] == "baseline-id" + + +def test_identity_errors_do_not_degrade_to_an_unidentified_baseline(monkeypatch): + from policyengine_simulation_executor import baseline_artifacts + + monkeypatch.setattr( + baseline_artifacts, + "qualifying_baseline_identity", + Mock(side_effect=SPMInputError("SPM_CONFIGURATION_UNAVAILABLE", "Uncertified")), + ) + with pytest.raises(SPMInputError): + baseline_artifacts.deterministic_baseline_id( + {}, + country="us", + policy=None, + region_code="us", + scoping_strategy=None, + year=2026, + ) + + +@pytest.mark.parametrize( + "code", + [ + "SPM_GEOGRAPHY_REQUIRED", + "SPM_GEOGRAPHY_UNAVAILABLE", + "SPM_COMPOSITION_REQUIRED", + "SPM_YEAR_UNAVAILABLE", + "SPM_SCENARIO_UNAVAILABLE", + ], +) +def test_gateway_poll_returns_structured_400(monkeypatch, code): + from policyengine_simulation_gateway import endpoints + from policyengine_simulation_gateway.testing import create_gateway_app + + monkeypatch.setattr(endpoints, "_job_metadata_store", lambda: {"job": {}}) + monkeypatch.setattr( + endpoints, + "modal", + SimpleNamespace( + FunctionCall=SimpleNamespace( + from_id=lambda _: SimpleNamespace( + get=Mock( + side_effect=SPMInputError(code, "Explicit input is required") + ) + ) + ) + ), + ) + response = TestClient(create_gateway_app()).get("/jobs/job") + assert response.status_code == 400 + assert response.json()["errors"] == [ + {"code": code, "message": "Explicit input is required"} + ] + assert response.json()["result"] is None + + +def test_gateway_submission_uses_registry_capability_before_spawn(monkeypatch): + from policyengine_simulation_gateway import endpoints + from policyengine_simulation_gateway.testing import create_gateway_app + from policyengine_simulation_contract.gateway_models import PolicyEngineBundle + + route = SimpleNamespace( + app_name="test-app", + response_version="test-only", + policyengine_version="test-only", + bundle_manifest={}, + country_model_version="test-only", + route_provenance=None, + ) + monkeypatch.setattr(endpoints, "resolve_route", lambda *args: route) + bundle = PolicyEngineBundle( + model_version="test-only", policyengine_version="test-only", spm=CAPABILITY + ) + monkeypatch.setattr(endpoints, "_build_policyengine_bundle", lambda *args: bundle) + spawn = Mock(return_value=SimpleNamespace(object_id="job")) + monkeypatch.setattr( + endpoints, + "modal", + SimpleNamespace( + Function=SimpleNamespace( + from_name=lambda *args: SimpleNamespace(spawn=spawn) + ) + ), + ) + monkeypatch.setattr(endpoints, "_job_metadata_store", lambda: {}) + response = TestClient(create_gateway_app()).post( + "/simulate/economy/comparison", + json={"country": "us", "spm": {"geography_kind": "national"}}, + ) + assert response.status_code == 200 + assert spawn.call_args.args[0]["spm"] == SELECTION + assert ( + SPMCapability.model_validate( + response.json()["policyengine_bundle"]["spm"] + ).model_dump() + == CAPABILITY + ) + bundle.spm = None + spawn.reset_mock() + response = TestClient(create_gateway_app()).post( + "/simulate/economy/comparison", json={"country": "us"} + ) + assert response.status_code == 400 + spawn.assert_not_called() + + +def test_segmented_child_input_error_is_not_redacted_or_retried(monkeypatch): + from src.modal.segmented_national import SegmentedNationalRunner + + child = SimpleNamespace( + get=Mock( + side_effect=SPMInputError("SPM_GEOGRAPHY_REQUIRED", "County required") + ), + cancel=Mock(), + ) + runner = object.__new__(SegmentedNationalRunner) + runner.poll_interval_initial_seconds = 0.001 + runner.poll_interval_max_seconds = 0.001 + with pytest.raises(SPMInputError) as error: + runner._collect([(["state/ca"], child)]) + assert error.value.code == "SPM_GEOGRAPHY_REQUIRED" + child.get.assert_called_once() + child.cancel.assert_called_once() + + +def test_budget_window_state_keeps_typed_failure_on_replay(): + from policyengine_simulation_contract.gateway_models import BudgetWindowBatchState + from policyengine_simulation_contract.budget_window_state import ( + build_batch_status_response, + ) + from policyengine_simulation_gateway.responses import batch_status_response + + state = BudgetWindowBatchState( + batch_job_id="job", + status="failed", + country="us", + region="us", + version="test", + resolved_app_name="app", + policyengine_bundle={"model_version": "test"}, + start_year="2026", + window_size=2, + max_parallel=2, + created_at="test", + updated_at="test", + error="County required", + errors=[{"code": "SPM_GEOGRAPHY_REQUIRED", "message": "County required"}], + ) + replay = BudgetWindowBatchState.model_validate_json(state.model_dump_json()) + response = batch_status_response(build_batch_status_response(replay)) + assert response.status_code == 400 + assert json.loads(response.body)["errors"][0]["code"] == "SPM_GEOGRAPHY_REQUIRED" + + +@pytest.mark.parametrize( + "code", + [ + "SPM_GEOGRAPHY_REQUIRED", + "SPM_GEOGRAPHY_UNAVAILABLE", + "SPM_COMPOSITION_REQUIRED", + "SPM_YEAR_UNAVAILABLE", + "SPM_SCENARIO_UNAVAILABLE", + ], +) +def test_country_error_is_transportable_without_country_package(monkeypatch, code): + from contextlib import nullcontext + + class CountryError(ValueError): + def __init__(self, code, message): + self.code = code + super().__init__(message) + + monkeypatch.setattr(simulation_runtime, "setup_gcp_credentials", nullcontext) + monkeypatch.setattr( + simulation_runtime, + "_run_simulation_impl_core", + Mock(side_effect=CountryError(code, "Observed input required")), + ) + with pytest.raises(SPMInputError) as caught: + simulation_runtime.run_simulation_impl({"country": "us"}) + transported = pickle.loads(pickle.dumps(caught.value)) + assert transported.to_dict() == {"code": code, "message": "Observed input required"} + + +@pytest.mark.parametrize( + "code", + [ + "SPM_GEOGRAPHY_REQUIRED", + "SPM_GEOGRAPHY_UNAVAILABLE", + "SPM_COMPOSITION_REQUIRED", + "SPM_YEAR_UNAVAILABLE", + "SPM_SCENARIO_UNAVAILABLE", + ], +) +def test_optional_analysis_does_not_swallow_spm_input_errors(code): + from policyengine_simulation_executor.simulation_output_common import ( + _try_compute_output, + ) + + with pytest.raises(SPMInputError): + _try_compute_output( + "optional impact", + Mock(side_effect=SPMInputError(code, "Explicit input required")), + ) + assert ( + _try_compute_output( + "optional impact", Mock(side_effect=RuntimeError("legacy optional output")) + ) + is None + ) + + +@pytest.mark.parametrize( + "date_fields,expected_as_of", + [({"as_of": None}, None), ({}, "2026-01-01")], +) +def test_as_of_presence_survives_budget_parent(date_fields, expected_as_of): + from policyengine_simulation_gateway.endpoints import ( + _resolve_request_spm, + _build_budget_window_parent_payload, + ) + from policyengine_simulation_contract.gateway_models import PolicyEngineBundle + + request = BudgetWindowBatchRequest( + country="us", + region="us", + start_year="2026", + window_size=2, + spm={"geography_kind": "national", **date_fields}, + ) + bundle = PolicyEngineBundle( + model_version="test", + policyengine_version="test", + spm={ + "contract_version": "canonical-spm-v1", + "defaults": {**SELECTION, "as_of": "2026-01-01"}, + }, + ) + selection = _resolve_request_spm( + request, + bundle, + SimpleNamespace( + bundle_manifest={}, country_model_version=None, route_provenance=None + ), + ) + assert selection["as_of"] == expected_as_of + parent = _build_budget_window_parent_payload( + request, resolved_version="test", resolved_app_name="test", bundle=bundle + ) + assert parent["spm"] == selection + + +def test_versions_http_retains_capability_contract_version(monkeypatch): + from policyengine_simulation_gateway import endpoints + from policyengine_simulation_gateway.testing import create_gateway_app + from src.modal.utils import update_version_registry as registry + + metadata = { + "app_name": "test-app", + "policyengine_version": "9.9.9", + "us": {"model_version": "test-us"}, + "uk": {"model_version": "test-uk"}, + "spm": CAPABILITY, + } + monkeypatch.setattr( + registry, "build_bundle_manifest_metadata", lambda **kwargs: metadata + ) + state = registry.build_next_routing_state( + current_state=None, + app_name="test-app", + policyengine_version="9.9.9", + us_version="test-us", + uk_version="test-uk", + ) + monkeypatch.setattr(endpoints, "_active_routing_state", lambda: state) + response = TestClient(create_gateway_app()).get("/versions") + assert response.status_code == 200 + payload = response.json() + assert payload["us"]["test-us"] == "test-app" + assert payload["policyengine"]["9.9.9"] == "test-app" + assert payload["spm_capabilities"] == {"9.9.9": CAPABILITY} + assert set(state["bundles"]) == {"9.9.9"} diff --git a/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py new file mode 100644 index 000000000..2d03c1d9d --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_canonical_spm_native.py @@ -0,0 +1,140 @@ +"""One-household real formula smoke, enabled explicitly with a local native H5. + +This test never loads or simulates the full population. It selects one household +and its native member records using the producer's HDF table indexes. +""" + +import json +import os + +import pytest + +from native_spm_support import materialize_native_household + +SOURCE = os.environ.get("SPM_NATIVE_SMOKE_SOURCE") +pytestmark = pytest.mark.skipif( + not SOURCE, + reason="Set SPM_NATIVE_SMOKE_SOURCE to the local native H5 for the bounded source integration smoke", +) + + +@pytest.fixture +def native_dataset(tmp_path): + return materialize_native_household(os.environ["SPM_NATIVE_SMOKE_SOURCE"], tmp_path) + + +def test_worker_baseline_reform_national_local_cache_and_receipts( + native_dataset, tmp_path +): + from policyengine.core.simulation import _cache + from policyengine_simulation_executor.simulation_runtime import _build_simulation + from policyengine_simulation_executor.spm import simulation_spm_result + + params = { + "country": "us", + "scope": "macro", + "time_period": "2024", + "data": str(tmp_path / "native.h5"), + "spm": {"geography_kind": "national"}, + } + baseline = _build_simulation( + params, dataset=native_dataset, policy=None, region_code="us" + ) + reform = _build_simulation( + params, + dataset=native_dataset, + policy={"gov.irs.credits.ctc.amount.base[0].amount": 3000}, + region_code="us", + ) + for simulation in (baseline, reform): + simulation.extra_variables = {"spm_unit": ["spm_unit_spm_threshold"]} + simulation.ensure() + assert ( + simulation.output_dataset.data.spm_unit["spm_unit_spm_threshold"] > 0 + ).all() + result = simulation_spm_result(baseline, reform, baseline.spm_config) + dumped = json.loads(json.dumps(result)) + assert dumped["spm_config"] == baseline.spm_config == reform.spm_config + assert dumped["spm_provenance"]["reform"][0]["years"]["2024"] + original = baseline.spm_provenance() + original["years"].clear() + assert baseline.spm_provenance()["years"] + # Disk replay uses the same identifier, selected config, and actual receipt. + _cache._cache.clear() + replay = _build_simulation( + params, dataset=native_dataset, policy=None, region_code="us" + ) + replay.id = baseline.id + replay.ensure() + assert replay.spm_provenance() == baseline.spm_provenance() + # A shared caller id cannot reuse national output for county selection. + local_params = {**params, "spm": {"geography_kind": "county"}} + local = _build_simulation( + local_params, + dataset=native_dataset, + policy=None, + region_code="us", + ) + local.id = baseline.id + assert local.storage_id != baseline.storage_id + local.ensure() + assert local.spm_provenance()["geography_kind"] == "county" + + # The precompute guard compares the planner's storage id against this + # wrapper property and refuses to publish on a mismatch, so agreement is + # what makes any canonical artifact publishable. Hermetic CI can only + # check our side against the wrapper's expression transcribed into + # test_artifact_keys; this is the same claim against the real wrapper. + from policyengine_simulation_executor.artifact_keys import canonical_digest + from policyengine_simulation_executor.spm import normalize_runtime_spm + + for simulation, request in ((baseline, params), (local, local_params)): + resolved = normalize_runtime_spm(request) + assert simulation.spm_config == resolved + assert simulation.storage_id == ( + f"{simulation.id}-spm-{canonical_digest(resolved)}" + ) + + +def test_native_state_only_worker_requires_geography_and_explicit_national_works( + native_dataset, +): + from policyengine_simulation_executor.simulation_runtime import _build_simulation + from policyengine_simulation_contract.spm import spm_error_detail + + native_dataset.data.household.drop(columns=["county_fips"], inplace=True) + params = { + "country": "us", + "scope": "macro", + "time_period": "2024", + "data": "test-only-custom-data", + } + simulation = _build_simulation( + params, dataset=native_dataset, policy=None, region_code="us" + ) + with pytest.raises(ValueError) as error: + simulation.run() + assert spm_error_detail(error.value).code == "SPM_GEOGRAPHY_REQUIRED" + national = _build_simulation( + {**params, "spm": {"geography_kind": "national"}}, + dataset=native_dataset, + policy=None, + region_code="us", + ) + national.run() + assert national.spm_provenance()["geography_kind"] == "national" + + +def test_actual_provider_unknown_metro_is_a_typed_input_error(): + from policyengine_simulation_executor.spm import normalize_runtime_spm + from policyengine_simulation_contract.spm import SPMInputError + + with pytest.raises(SPMInputError) as error: + normalize_runtime_spm( + { + "country": "us", + "time_period": "2024", + "spm": {"geography_kind": "metro", "geography_id": "00000"}, + } + ) + assert error.value.code == "SPM_GEOGRAPHY_UNAVAILABLE" diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index 172a3a57d..b96ed3c84 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -16,6 +16,7 @@ from pydantic import ValidationError from fixtures.identity_stubs import install_identity_stubs +from fixtures.wrapper_spm import wrapper_storage_id from policyengine_simulation_executor import precompute from policyengine_simulation_executor.artifact_keys import canonical_digest from policyengine_simulation_executor.precompute_models import ( @@ -390,9 +391,7 @@ def test_precompute_identity_equals_runtime_id(self, identity_stubs, fake_region year=2026, ) assert writer_identity.simulation_id == reader_id - assert writer_identity.store_path.endswith( - f"/{writer_identity.simulation_id}.h5" - ) + assert writer_identity.store_path.endswith(f"/{writer_identity.storage_id}.h5") def test_dataset_filename_matches_runtime_stem_lookup(self, monkeypatch): """Real-manifest coverage: the artifact filename equals the exact @@ -669,8 +668,34 @@ class SilentBaseline(ArtifactBaselineSimulation): def ensure(self): self._artifact_outcome = "miss" + class SPMStubBaseline(ArtifactBaselineSimulation): + """A baseline that names its artifact the way the wrapper does. + + The canonical wrapper derives ``storage_id`` from the resolved + selection and saves under it; ``compute_baseline_impl`` plans a + path from the identity's own derivation and refuses to act when + the two disagree. Only a stub that computes its side + independently can show the guard passing rather than aborting. + """ + + @property + def spm_config(self): + return state.spm + + @property + def storage_id(self): + return wrapper_storage_id(self.id, self.spm_config) + + def ensure(self): + self._artifact_outcome = "miss" + (tmp_path / f"{self.storage_id}.h5").write_bytes(b"artifact-bytes") + + state.spm = None state.baseline = StubBaseline.model_construct(id="bl1-cohort") state.make_silent = lambda: SilentBaseline.model_construct(id="bl1-cohort") + state.make_spm = lambda sim_id="bl1-cohort": SPMStubBaseline.model_construct( + id=sim_id + ) class FakeStore: def __init__(self, bucket): @@ -768,6 +793,116 @@ def test_refuses_to_act_under_a_mismatched_id(self, cohort_stubs): precompute.compute_baseline_impl("bucket-x", self._entry("bl1-other")) assert cohort_stubs.uploads == [] + def test_refuses_a_mismatched_planned_storage_id(self, cohort_stubs): + entry = self._entry() + entry.path = "baselines/us/bl-d/bl1-cohort-spm-planned.h5" + with pytest.raises(RuntimeError, match="storage ids disagree"): + precompute.compute_baseline_impl("bucket-x", entry) + assert cohort_stubs.configured == [] + assert cohort_stubs.uploads == [] + assert not (cohort_stubs.folder / "bl1-cohort.h5").exists() + + @staticmethod + def _planned_identity(selection): + """A real ``BaselineArtifactIdentity``, not a restated expression. + + The plan the deploy publishes comes from this object's properties, + so keying the test off them is what makes the write path a + writer==reader check: reformat or truncate ``storage_id``'s digest + and the container -- which derives its own the wrapper's way -- no + longer names the planned file. + """ + from policyengine_simulation_executor.artifact_keys import ( + BaselineArtifactIdentity, + DatasetArtifactIdentity, + ) + + return BaselineArtifactIdentity( + spm=selection, + country="us", + region="state/ca+state/wv", + scope_key="scoping", + dataset=DatasetArtifactIdentity( + spm=selection, + country="us", + dataset="populace_cps", + stem="populace", + year=2026, + data_version="1.2.3", + data_artifact_revision="rev-abc", + source_sha256=None, + data_build_fingerprint=None, + model_version="9.9.9", + policyengine_version="4.22.0", + ), + ) + + def test_publishes_a_selection_scoped_artifact_under_the_planned_path( + self, cohort_stubs + ): + """The canonical write path, end to end through the guard. + + Nothing else in hermetic CI runs ``compute_baseline_impl`` with a + selection: the identity's storage id was only ever compared to + itself, and the one case that reached the guard covered the abort. + Here the plan is the planner's own ``BaselineArtifactIdentity`` and + the container names its artifact the wrapper's way, so the guard + passing at all is the claim -- if the two derivations differed, no + canonical artifact could be published and the deploy would be + blocked. + """ + selection = { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, + } + planned = self._planned_identity(selection) + cohort_stubs.spm = selection + cohort_stubs.baseline = cohort_stubs.make_spm(planned.simulation_id) + + entry = self._entry(planned.simulation_id) + entry.digest = planned.digest + entry.path = planned.store_path + + result = precompute.compute_baseline_impl("bucket-x", entry) + + assert cohort_stubs.uploads == [ + ( + planned.store_path, + str(cohort_stubs.folder / f"{planned.storage_id}.h5"), + ) + ] + assert result.simulation_id == planned.simulation_id + assert result.uploaded is True + assert result.size_bytes == len(b"artifact-bytes") + # The digest the container computed is the one the planner planned. + assert planned.storage_id == wrapper_storage_id( + planned.simulation_id, selection + ) + # A selection-free plan for the same cohort is a different artifact. + assert planned.storage_id != planned.simulation_id + + def test_refuses_a_selection_scoped_plan_the_container_does_not_share( + self, cohort_stubs + ): + """The abort the review worried about, under a real selection: a + container that resolved a different selection names a different + artifact, and nothing is published.""" + cohort_stubs.spm = {"scenario": "zero_real"} + cohort_stubs.baseline = cohort_stubs.make_spm() + entry = self._entry() + entry.path = ( + f"baselines/us/bl-d/bl1-cohort-spm-" + f"{canonical_digest({'scenario': 'ce_trend'})}.h5" + ) + with pytest.raises(RuntimeError, match="storage ids disagree"): + precompute.compute_baseline_impl("bucket-x", entry) + assert cohort_stubs.configured == [] + assert cohort_stubs.uploads == [] + def test_refuses_a_plain_simulation(self, cohort_stubs): cohort_stubs.baseline = SimpleNamespace(id="bl1-cohort") with pytest.raises(RuntimeError, match="plain Simulation"): diff --git a/projects/policyengine-simulation-executor/tests/test_segmented_national.py b/projects/policyengine-simulation-executor/tests/test_segmented_national.py index 0ac6ea7ce..e5280a4c9 100644 --- a/projects/policyengine-simulation-executor/tests/test_segmented_national.py +++ b/projects/policyengine-simulation-executor/tests/test_segmented_national.py @@ -5,7 +5,9 @@ import pytest from src.modal import segmented_national as sn +from policyengine_simulation_contract.spm import SPMInputError from policyengine_simulation_executor import simulation_runtime as sr +from policyengine_simulation_executor import spm as executor_spm from policyengine_simulation_executor.national_partition import ( US_NATIONAL_REGION_GROUPS, ) @@ -363,6 +365,9 @@ def fake_segmented(params, *, app_name): return {"segmented": True} monkeypatch.setattr(sn, "run_segmented_national_impl", fake_segmented) + monkeypatch.setattr( + executor_spm, "normalize_runtime_spm", lambda params: None + ) monkeypatch.setattr( sr, "run_simulation_impl", @@ -388,9 +393,175 @@ def test__everything_else_takes_the_monolithic_path( "run_segmented_national_impl", lambda *a, **k: pytest.fail("segmented path must not run"), ) + monkeypatch.setattr( + executor_spm, "normalize_runtime_spm", lambda params: None + ) monkeypatch.setattr( sr, "run_simulation_impl", lambda params: {"monolithic": True} ) assert sn.dispatch_run_simulation(params, app_name="app-x") == { "monolithic": True } + + +SPM_SELECTION = { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, +} + + +def _spm_receipt(label, *, selection=SPM_SELECTION, year="2026"): + return { + "forecast_id": label, + "forecast_sha256": selection["forecast_content_sha256"], + "scenario": selection["scenario"], + "geography_kind": selection["geography_kind"], + "runtime_versions": {"policyengine-us": "test-only"}, + "years": {year: {"status": "forecast"}}, + "geographies": [], + "composition_method": "classified-inputs", + "storage_method": "formula", + } + + +def _spm_child(index, *, selection=SPM_SELECTION, year="2026"): + """One child's result carrying that segment's own SPM receipts.""" + return { + "child": index, + "spm_config": dict(selection), + "spm_provenance": { + "baseline": [ + _spm_receipt(f"baseline-{index}", selection=selection, year=year) + ], + "reform": [_spm_receipt(f"reform-{index}", selection=selection, year=year)], + }, + } + + +def _spm_child_missing(index, field): + """A child whose transported selection lost a resolved option.""" + child = _spm_child(index) + del child["spm_config"][field] + return child + + +@pytest.fixture +def stub_reduce(monkeypatch): + """The microdata reduce itself is covered by + test_segmented_national_reduce; these tests pin the SPM merge.""" + monkeypatch.setattr( + sn, "build_national_output", lambda children, **kwargs: {"budget": {}} + ) + + +class TestSegmentedNationalSPM: + def test__child_receipts_combine_into_one_national_receipt( + self, bare_country, stub_reduce + ): + fake = FakeModal([FakeCall(_spm_child(i)) for i in range(20)]) + runner = _runner(fake, params={**NATIONAL, "spm": SPM_SELECTION}) + + output = runner.run() + + # The resolved selection rides to every child unchanged and comes + # back on the parent as the one national selection. + assert all(p["spm"] == SPM_SELECTION for p in fake.spawned_payloads) + assert output["spm_config"] == SPM_SELECTION + # Every executed segment's receipts survive, baseline and reform + # concatenated in group order: a national SPM result must account + # for all 20 segments, not just the first. + provenance = output["spm_provenance"] + assert [r["forecast_id"] for r in provenance["baseline"]] == [ + f"baseline-{i}" for i in range(20) + ] + assert [r["forecast_id"] for r in provenance["reform"]] == [ + f"reform-{i}" for i in range(20) + ] + assert output["budget"] == {} + + @pytest.mark.parametrize( + "rogue,reason", + [ + ( + _spm_child(7, selection={**SPM_SELECTION, "scenario": "zero_real"}), + "selection differs from the request", + ), + (_spm_child(7, year="2025"), "does not cover the requested year"), + (_spm_child_missing(7, "county_vintage"), "complete resolved"), + ], + ) + def test__child_receipt_disagreeing_with_the_request_is_typed( + self, bare_country, stub_reduce, rogue, reason + ): + # One segment computed something other than what was requested: + # fail with the public typed code, never publish a mixed national + # result whose receipt does not describe every segment. + calls = [FakeCall(_spm_child(i)) for i in range(20)] + calls[7] = FakeCall(rogue) + runner = _runner(FakeModal(calls), params={**NATIONAL, "spm": SPM_SELECTION}) + + with pytest.raises(SPMInputError, match=reason) as error: + runner.run() + assert error.value.code == "SPM_CONFIGURATION_UNAVAILABLE" + + def test__unrequested_child_receipt_is_typed(self, bare_country, stub_reduce): + # A historical request must not inherit a child's stray receipt. + calls = [FakeCall({"child": i}) for i in range(20)] + calls[3] = FakeCall(_spm_child(3)) + runner = _runner(FakeModal(calls)) + + with pytest.raises(SPMInputError, match="Unexpected SPM receipt"): + runner.run() + + def test__a_run_without_spm_adds_no_receipt_keys(self, bare_country, stub_reduce): + runner = _runner(FakeModal([FakeCall({"child": i}) for i in range(20)])) + assert runner.run() == {"budget": {}} + + +class TestDispatchResolvesSPM: + def test__selection_is_resolved_before_the_fan_out(self, monkeypatch): + # Children and the parent's receipt check must both see the + # worker-resolved selection, not the caller's partial one. + monkeypatch.setattr( + executor_spm, "normalize_runtime_spm", lambda params: SPM_SELECTION + ) + seen = {} + + def fake_segmented(params, *, app_name): + seen["params"] = params + return {"segmented": True} + + monkeypatch.setattr(sn, "run_segmented_national_impl", fake_segmented) + + result = sn.dispatch_run_simulation( + {**NATIONAL, "spm": {"geography_kind": "national"}}, app_name="app-x" + ) + + assert result == {"segmented": True} + assert seen["params"]["spm"] == SPM_SELECTION + + def test__typed_spm_error_spawns_nothing(self, monkeypatch): + # Resolution failures are cheap only if they happen before 20 + # children are spawned. + def unavailable(params): + raise SPMInputError("SPM_GEOGRAPHY_REQUIRED", "County required") + + monkeypatch.setattr(executor_spm, "normalize_runtime_spm", unavailable) + monkeypatch.setattr( + sn, + "run_segmented_national_impl", + lambda *a, **k: pytest.fail("segmented path must not run"), + ) + monkeypatch.setattr( + sr, + "run_simulation_impl", + lambda params: pytest.fail("monolithic path must not run"), + ) + + with pytest.raises(SPMInputError) as error: + sn.dispatch_run_simulation(dict(NATIONAL), app_name="app-x") + assert error.value.code == "SPM_GEOGRAPHY_REQUIRED" diff --git a/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py b/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py new file mode 100644 index 000000000..a43aee2f9 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_spm_runtime_native.py @@ -0,0 +1,143 @@ +"""Worker prevalidation against the installed canonical calculator contracts. + +Enabled by the authenticated native qualification harness. These cases stop +before dataset loading or child submission and never run a population job. +""" + +import os +from unittest.mock import Mock + +import pytest + +pytestmark = pytest.mark.skipif( + not os.environ.get("SPM_NATIVE_SMOKE_SOURCE"), + reason="Requires the installed canonical SPM qualification environment", +) + + +@pytest.fixture +def forecast(): + from policyengine_simulation_executor.spm import _forecast, runtime_spm_capability + + return _forecast(runtime_spm_capability().defaults.forecast_content_sha256) + + +def _worker_entrypoint(params, monkeypatch): + from policyengine_simulation_executor import simulation_runtime + from src.modal import budget_window_batch + + if "start_year" in params: + entrypoint = budget_window_batch.run_budget_window_batch_impl + boundary = Mock(side_effect=AssertionError("Must reject before child setup")) + monkeypatch.setattr(budget_window_batch, "build_batch_context", boundary) + else: + entrypoint = simulation_runtime._run_simulation_impl_core + boundary = Mock(side_effect=AssertionError("Must reject before dataset lookup")) + monkeypatch.setattr(simulation_runtime, "_resolve_dataset_reference", boundary) + return entrypoint, boundary + + +@pytest.mark.parametrize( + "period,unavailable_year", + [ + ({"time_period": "2021"}, 2021), + ({"time_period": "2036"}, 2036), + ({"year": "2040"}, 2040), + ({"start_year": "2021", "window_size": 1}, 2021), + ({"start_year": "2036", "window_size": 1}, 2036), + ({"start_year": "2035", "window_size": 2}, 2036), + ], +) +def test_worker_year_boundary_matches_real_provider( + forecast, period, unavailable_year, monkeypatch +): + from spm_calculator.policyengine_adapter import PolicyEngineSPMProvider + from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail + + provider = PolicyEngineSPMProvider(forecast, geography_kind="national") + with pytest.raises(ValueError) as provider_error: + provider.year_metadata(unavailable_year) + expected = spm_error_detail(provider_error.value) + assert expected.code == "SPM_YEAR_UNAVAILABLE" + assert provider.provenance()["years"] == {} + params = {"country": "us", "spm": {"geography_kind": "national"}, **period} + entrypoint, boundary = _worker_entrypoint(params, monkeypatch) + with pytest.raises(SPMInputError) as error: + entrypoint(params) + assert error.value.to_dict() == expected.model_dump() + boundary.assert_not_called() + + +@pytest.mark.parametrize( + "period", + [ + {"time_period": "2024"}, + {"start_year": "2024", "window_size": 2}, + ], +) +def test_worker_scenario_boundary_matches_real_calculator_contract( + forecast, period, monkeypatch +): + from spm_calculator.axiom_adapter import export_build_spec + from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail + + scenario = "scenario-not-in-the-pinned-forecast" + with pytest.raises(ValueError) as adapter_error: + export_build_spec(forecast, years=[2024], scenario=scenario) + assert adapter_error.value.code == "SPM_SCENARIO_UNAVAILABLE" + expected = spm_error_detail(adapter_error.value) + assert expected is not None + params = { + "country": "us", + "spm": {"geography_kind": "national", "scenario": scenario}, + **period, + } + entrypoint, boundary = _worker_entrypoint(params, monkeypatch) + with pytest.raises(SPMInputError) as error: + entrypoint(params) + assert error.value.to_dict() == expected.model_dump() + boundary.assert_not_called() + + +def test_valid_worker_prevalidation_does_not_evaluate_amounts(forecast, monkeypatch): + from spm_calculator.policyengine_adapter import PolicyEngineSPMProvider + from policyengine_simulation_executor.spm import normalize_runtime_spm + + calculate = Mock(side_effect=AssertionError("Prevalidation is not measurement")) + monkeypatch.setattr(PolicyEngineSPMProvider, "calculate_unit", calculate) + resolved = normalize_runtime_spm( + { + "country": "us", + "start_year": "2034", + "window_size": 2, + "spm": {"geography_kind": "national"}, + } + ) + assert resolved["geography_kind"] == "national" + provider = PolicyEngineSPMProvider( + forecast=forecast, + geography_kind=resolved["geography_kind"], + scenario=resolved["scenario"], + as_of=resolved["as_of"], + ) + assert provider.provenance()["years"] == {} + calculate.assert_not_called() + + +@pytest.mark.parametrize( + "selection", + [{"as_of": "2020-01-01"}, {"county_vintage": "2010"}], +) +def test_other_invalid_settings_keep_the_settings_code(selection): + from policyengine_simulation_contract.spm import SPMInputError + from policyengine_simulation_executor.spm import normalize_runtime_spm + + with pytest.raises(SPMInputError) as error: + normalize_runtime_spm( + { + "country": "us", + "time_period": "2024", + "spm": {"geography_kind": "national", **selection}, + } + ) + assert error.value.code == "SPM_SETTINGS_INVALID" diff --git a/projects/policyengine-simulation-gateway/README.md b/projects/policyengine-simulation-gateway/README.md index 8372da0dc..619f6c6c4 100644 --- a/projects/policyengine-simulation-gateway/README.md +++ b/projects/policyengine-simulation-gateway/README.md @@ -4,6 +4,30 @@ The stable Modal gateway for the simulation service: routes simulation requests to versioned executor apps (`policyengine-simulation-py{X}`) via the routing state in `modal.Dict`, and serves the public API contract. +## Country-version routing + +A request that names a country model `version` resolves to the app the +routing state lists for it. When the registry ties that app to exactly one +wrapper bundle and the bundle's manifest states a *different* model version +for the country, the gateway refuses the request with a 400 instead of +routing it. + +That is a deliberate change from the pre-canonical gateway, which served +the route and returned a self-contradictory body — `version` from the stale +route, `policyengine_bundle.model_version` from the live manifest. It can +happen without anyone doing anything wrong: publishing only ever *adds* +country routes and overwrites the bundle manifest for the wrapper it +deploys, and nothing prunes the old country route, so re-publishing one +wrapper with an upgraded country model leaves the previous country version +pointing at an app that no longer serves it. Callers pinning an old +`version` (the `policyengine-apis-integ` suite pins one) should move to a +version the current bundle states, or pass `policyengine_version` +explicitly. + +A manifest that states *no* model version for the country — entry absent, +or the value absent, non-string, empty or whitespace — contradicts nothing +and still routes. + ## Image dependencies The Modal image installs with `uv_sync(frozen=True)` from this project's diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py index 235eaf154..19c1566d3 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/endpoints.py @@ -24,6 +24,12 @@ put_batch_job_state, ) from policyengine_simulation_gateway.auth import require_auth +from policyengine_simulation_contract.spm import ( + SPMCapability, + SPMInputError, + spm_error_detail, + resolve_spm_selection, +) from policyengine_simulation_observability.errors import log_and_redact_exception from policyengine_simulation_contract.gateway_models import ( BudgetWindowBatchRequest, @@ -76,6 +82,11 @@ class RouteResolution: response_version: str policyengine_version: str | None bundle_manifest: dict + route_provenance: str | None = None + # The country model version this route names, when it names one. A + # country route's key is that version; a policyengine-keyed route's is + # not, and echoing a wrapper version there reads as a model version. + country_model_version: str | None = None def _job_metadata_store(): @@ -338,22 +349,76 @@ def _bundle_manifest(state: dict, policyengine_version: str | None) -> dict: return manifest if isinstance(manifest, dict) else {} -def _policyengine_version_for_app(state: dict, app_name: str) -> str | None: - for policyengine_version, routed_app in _routing_state_routes( - state, "policyengine" - ).items(): - if routed_app == app_name: - return ( - policyengine_version if isinstance(policyengine_version, str) else None - ) - - for policyengine_version, manifest in _routing_state_bundles(state).items(): - if isinstance(manifest, dict) and manifest.get("app_name") == app_name: - return ( - policyengine_version if isinstance(policyengine_version, str) else None - ) - - return None +def _policyengine_version_for_app( + state: dict, app_name: str, *, country: str, model_version: str +) -> str | None: + """The wrapper version serving this country route, or None. + + Behaviour change from the pre-canonical gateway, deliberate: a country + route whose one candidate bundle states a *different* model version is + now rejected with a 400 instead of resolving. Publishing only ever adds + country routes and overwrites the bundle manifest for the wrapper + version it deploys (``update_version_registry``), so re-publishing one + wrapper with an upgraded country model leaves the old country route + pointing at an app whose manifest now states the new model. Before, that + route resolved and the 202 body contradicted itself -- ``version`` from + the stale route, ``policyengine_bundle.model_version`` from the live + manifest -- and with canonical SPM the caller would also have been + reading a capability the requested model never had. Nothing prunes the + stale route, so the gateway refuses it instead. + + "States a different model version" and "states none" are one + classification, computed once below: an absent country entry, an absent, + non-string, empty or whitespace ``model_version`` all contradict + nothing and still resolve. + """ + candidates = { + version + for version, routed_app in _routing_state_routes(state, "policyengine").items() + if isinstance(version, str) and version != "latest" and routed_app == app_name + } | { + version + for version, manifest in _routing_state_bundles(state).items() + if isinstance(version, str) + and version != "latest" + and isinstance(manifest, dict) + and manifest.get("app_name") == app_name + } + matching = [] + unclassified = [] + for version in candidates: + manifest = _bundle_manifest(state, version) + country_bundle = manifest.get(country) + candidate_model = ( + country_bundle.get("model_version") + if isinstance(country_bundle, dict) + else None + ) + if not isinstance(candidate_model, str) or not candidate_model.strip(): + unclassified.append(version) + elif candidate_model == model_version: + matching.append(version) + if len(matching) == 1 and not unclassified: + return matching[0] + if len(candidates) > 1: + raise ValueError( + f"Ambiguous bundle for {country} version {model_version}; pass " + "policyengine_version to select a bundle explicitly" + ) + if not candidates: + return None + version = next(iter(candidates)) + if not unclassified: + # Reuse the classification above rather than re-deriving it: the + # shared validator reads any string as a stated model version, so + # calling it unconditionally rejected a blank one here while the + # ambiguity check above treated blank as unstated. + _validate_legacy_version_matches_bundle( + country=country, + requested_version=model_version, + manifest=_bundle_manifest(state, version), + ) + return version def _policyengine_version_from_app_name(app_name: str) -> str | None: @@ -392,12 +457,30 @@ def _resolve_country_route( app_name = _routing_state_routes(state, country).get(version) if not isinstance(app_name, str): return None - policyengine_version = _policyengine_version_for_app(state, app_name) + policyengine_version = _policyengine_version_for_app( + state, app_name, country=country, model_version=version + ) + if policyengine_version is None: + # The legacy seed producer infers these same wrapper routes. Missing + # registry metadata must not erase a known future app version. + policyengine_version = _policyengine_version_from_app_name(app_name) return RouteResolution( app_name=app_name, response_version=version, policyengine_version=policyengine_version, bundle_manifest=_bundle_manifest(state, policyengine_version), + country_model_version=version, + # A country route the registry ties to no wrapper bundle at all is + # pre-canonical by construction: every publish writes the wrapper + # route and the bundle manifest for the app it deploys, and the + # legacy seed only infers wrapper routes for prefixed app names. The + # registry's ``generation`` marker cannot carry this — the next + # publish rewrites it (update_version_registry). + route_provenance=( + "legacy-country-route" + if policyengine_version is None and state.get("schema_version") == 1 + else None + ), ) @@ -534,6 +617,8 @@ def _resolve_from_legacy_dicts( response_version=resolved_version, policyengine_version=_policyengine_version_from_app_name(app_name), bundle_manifest={}, + country_model_version=resolved_version, + route_provenance="legacy-country-dict", ) @@ -568,6 +653,15 @@ def _build_policyengine_bundle( policyengine_version = app_bundle.get( "policyengine_version", resolution.policyengine_version ) + capability = app_bundle.get("spm") if country.lower() == "us" else None + if capability is not None: + try: + capability = SPMCapability.model_validate(capability) + except ValueError as exc: + raise SPMInputError( + "SPM_CONFIGURATION_UNAVAILABLE", + "This worker bundle has invalid canonical SPM capability metadata", + ) from exc return PolicyEngineBundle( model_version=str(model_version), policyengine_version=( @@ -575,6 +669,53 @@ def _build_policyengine_bundle( ), data_version=str(data_version) if isinstance(data_version, str) else None, dataset=resolved_dataset, + spm=capability, + ) + + +def _certified_model_version(country: str, route: RouteResolution) -> str | None: + """The country model version the registry states for this route. + + ``PolicyEngineBundle.model_version`` falls back to the routing response + version so the response always carries one. On a policyengine-keyed + route with no manifest that fallback is a wrapper version, and reading + it as a model version is how a legacy 5.2.0/5.3.0 route stopped + resolving. An unstated model version contradicts nothing. + """ + country_bundle = route.bundle_manifest.get(country.lower()) + if isinstance(country_bundle, dict): + stated = country_bundle.get("model_version") + if isinstance(stated, str) and stated.strip(): + return stated + return route.country_model_version + + +def _resolve_request_spm(request, bundle, route): + selection = resolve_spm_selection( + request.country, + request.spm, + capability=bundle.spm, + policyengine_version=bundle.policyengine_version, + model_version=_certified_model_version(request.country, route), + route_provenance=route.route_provenance, + ) + if selection is not None: + from policyengine_simulation_contract.spm import SPMSelection + + request.spm = SPMSelection.model_validate(selection) + return selection + + +def _bundle_payload(bundle: PolicyEngineBundle, **dump_kwargs) -> dict: + """Dump a bundle, omitting ``spm`` when the route has no capability. + + The 202 and 500 job bodies splat this dict in directly, bypassing the + routes' ``response_model_exclude_none``. Every other optional bundle + field was already emitted as an explicit null before canonical SPM, so + only the new key is dropped: a legacy no-SPM body stays byte-identical. + """ + return bundle.model_dump( + exclude=None if bundle.spm is not None else {"spm"}, **dump_kwargs ) @@ -585,7 +726,7 @@ def _serialize_job_metadata( ) -> dict: return { "resolved_app_name": resolved_app_name, - "policyengine_bundle": bundle.model_dump(), + "policyengine_bundle": _bundle_payload(bundle), "run_id": run_id, } @@ -602,13 +743,15 @@ def _build_budget_window_parent_payload( mode="json", exclude_none=True, ) + if request.spm is not None: + payload["spm"] = request.spm.model_dump(mode="json") payload["version"] = resolved_version if request.telemetry is not None: payload["_telemetry"] = request.telemetry.model_dump(mode="json") payload["_metadata"] = { "resolved_version": resolved_version, "resolved_app_name": resolved_app_name, - "policyengine_bundle": bundle.model_dump(mode="json"), + "policyengine_bundle": _bundle_payload(bundle, mode="json"), } return payload @@ -694,10 +837,19 @@ async def submit_simulation(request: SimulationRequest): route, payload, ) + _resolve_request_spm(request, bundle, route) except (ValueError, HuggingFaceDatasetReferenceError) as exc: + detail = spm_error_detail(exc) + if detail: + return failed_job_response( + error=detail.message, errors=[detail.model_dump()] + ) record_error(exc, handled=True, status_code=400, include_stack=False) raise HTTPException(status_code=400, detail=str(exc)) from exc + if request.spm is not None: + payload["spm"] = request.spm.model_dump(mode="json") + logger.info( "Routing %s:%s to app %s (run_id=%s)", request.country, @@ -766,7 +918,13 @@ async def submit_budget_window_batch(request: BudgetWindowBatchRequest): route, request.model_dump(mode="json"), ) + _resolve_request_spm(request, bundle, route) except (ValueError, HuggingFaceDatasetReferenceError) as exc: + detail = spm_error_detail(exc) + if detail: + return failed_job_response( + error=detail.message, errors=[detail.model_dump()] + ) record_error(exc, handled=True, status_code=400, include_stack=False) raise HTTPException(status_code=400, detail=str(exc)) from exc with segment(SegmentName.REQUEST_PARSE): @@ -854,6 +1012,13 @@ async def get_job_status(job_id: str): if _is_modal_job_not_found(exc): record_error(exc, handled=True, status_code=404, include_stack=False) raise HTTPException(status_code=404, detail=f"Job not found: {job_id}") + detail = spm_error_detail(exc) + if detail: + return failed_job_response( + error=detail.message, + errors=[detail.model_dump()], + job_metadata=job_metadata, + ) redacted = log_and_redact_exception( exc, scope="simulation_job_status", @@ -922,13 +1087,19 @@ async def get_budget_window_job_status(batch_job_id: str): # "submitted" status from the seed store (#448). We deliberately # overwrite the main job store entry as well as the seed so either # lookup path observes the terminal failed state. - redacted = log_and_redact_exception( - exc, - scope="budget_window_parent_call", - context={"batch_job_id": batch_job_id}, + detail = spm_error_detail(exc) + message = ( + detail.message + if detail + else log_and_redact_exception( + exc, + scope="budget_window_parent_call", + context={"batch_job_id": batch_job_id}, + ) ) seed_state.status = "failed" - seed_state.error = redacted + seed_state.errors = [detail] if detail else None + seed_state.error = message with segment(SegmentName.BUDGET_WINDOW_STATE_WRITE): put_batch_job_state(seed_state) put_batch_job_seed(seed_state) @@ -947,7 +1118,16 @@ async def list_versions() -> VersionsResponse: with segment(SegmentName.ROUTE_RESOLUTION): state = _active_routing_state() if state: + capabilities = {} + for name, bundle in _routing_state_bundles(state).items(): + if not isinstance(bundle, dict) or bundle.get("spm") is None: + continue + try: + capabilities[name] = SPMCapability.model_validate(bundle["spm"]) + except ValueError: + logger.warning("Omitting invalid SPM capability for bundle %s", name) return VersionsResponse( + spm_capabilities=capabilities, policyengine=_version_map_from_state(state, "policyengine"), us=_version_map_from_state(state, "us"), uk=_version_map_from_state(state, "uk"), diff --git a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py index 9aa2b2967..562d42b94 100644 --- a/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py +++ b/projects/policyengine-simulation-gateway/src/policyengine_simulation_gateway/responses.py @@ -25,6 +25,11 @@ def __init__(self, content: dict): def batch_status_payload(response: BudgetWindowBatchStatusResponse) -> dict: payload = response.model_dump(mode="json") + if payload.get("errors") is None: + payload.pop("errors", None) + for child in payload.get("child_jobs", {}).values(): + if child.get("errors") is None: + child.pop("errors", None) if response.policyengine_bundle is not None: payload["policyengine_bundle"] = response.policyengine_bundle.model_dump( mode="json", @@ -38,7 +43,9 @@ def batch_status_response(response: BudgetWindowBatchStatusResponse): if response.status in {"submitted", "running"}: return AcceptedResponse(payload) if response.status == "failed": - return ServerErrorResponse(payload) + return JSONResponse( + status_code=400 if response.errors else 500, content=payload + ) return response @@ -54,13 +61,15 @@ def running_job_response(job_metadata: dict | None = None) -> AcceptedResponse: def failed_job_response( - *, error: str, job_metadata: dict | None = None -) -> ServerErrorResponse: - return ServerErrorResponse( - { + *, error: str, job_metadata: dict | None = None, errors: list[dict] | None = None +) -> JSONResponse: + return JSONResponse( + status_code=400 if errors else 500, + content={ + **({"errors": errors} if errors else {}), "status": "failed", "result": None, "error": error, **(job_metadata or {}), - } + }, ) diff --git a/projects/policyengine-simulation-gateway/tests/golden/openapi.json b/projects/policyengine-simulation-gateway/tests/golden/openapi.json index 606d51a3d..2282cbf4b 100644 --- a/projects/policyengine-simulation-gateway/tests/golden/openapi.json +++ b/projects/policyengine-simulation-gateway/tests/golden/openapi.json @@ -370,6 +370,20 @@ ], "title": "Status" }, + "errors": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/SPMErrorDetail" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Errors" + }, "error": { "anyOf": [ { @@ -392,6 +406,26 @@ }, "BudgetWindowAnnualImpact": { "properties": { + "spm_config": { + "anyOf": [ + { + "$ref": "#/components/schemas/SPMSelection" + }, + { + "type": "null" + } + ] + }, + "spm_provenance": { + "anyOf": [ + { + "$ref": "#/components/schemas/SPMComparisonProvenance" + }, + { + "type": "null" + } + ] + }, "year": { "type": "string", "title": "Year" @@ -435,6 +469,16 @@ "type": "string", "title": "Country" }, + "spm": { + "anyOf": [ + { + "$ref": "#/components/schemas/SPMSelection" + }, + { + "type": "null" + } + ] + }, "version": { "anyOf": [ { @@ -696,6 +740,20 @@ } ] }, + "errors": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/SPMErrorDetail" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Errors" + }, "error": { "anyOf": [ { @@ -1262,6 +1320,20 @@ } ] }, + "errors": { + "anyOf": [ + { + "items": { + "$ref": "#/components/schemas/SPMErrorDetail" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Errors" + }, "error": { "anyOf": [ { @@ -1635,6 +1707,16 @@ } ], "title": "Dataset" + }, + "spm": { + "anyOf": [ + { + "$ref": "#/components/schemas/SPMCapability" + }, + { + "type": "null" + } + ] } }, "type": "object", @@ -1721,12 +1803,230 @@ ], "title": "RacePovertyOutput" }, + "SPMCapability": { + "properties": { + "contract_version": { + "type": "string", + "const": "canonical-spm-v1", + "title": "Contract Version", + "default": "canonical-spm-v1" + }, + "defaults": { + "$ref": "#/components/schemas/SPMSelection" + } + }, + "additionalProperties": false, + "type": "object", + "required": [ + "defaults" + ], + "title": "SPMCapability" + }, + "SPMComparisonProvenance": { + "properties": { + "baseline": { + "items": { + "$ref": "#/components/schemas/SPMProvenance" + }, + "type": "array", + "minItems": 1, + "title": "Baseline" + }, + "reform": { + "items": { + "$ref": "#/components/schemas/SPMProvenance" + }, + "type": "array", + "minItems": 1, + "title": "Reform" + } + }, + "additionalProperties": false, + "type": "object", + "required": [ + "baseline", + "reform" + ], + "title": "SPMComparisonProvenance", + "description": "One receipt per executed regional segment, separately for each policy." + }, + "SPMErrorDetail": { + "properties": { + "code": { + "type": "string", + "title": "Code" + }, + "message": { + "type": "string", + "title": "Message" + } + }, + "type": "object", + "required": [ + "code", + "message" + ], + "title": "SPMErrorDetail" + }, + "SPMProvenance": { + "properties": { + "forecast_id": { + "type": "string", + "title": "Forecast Id" + }, + "forecast_sha256": { + "type": "string", + "title": "Forecast Sha256" + }, + "scenario": { + "type": "string", + "title": "Scenario" + }, + "geography_kind": { + "type": "string", + "title": "Geography Kind" + }, + "runtime_versions": { + "additionalProperties": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ] + }, + "type": "object", + "title": "Runtime Versions" + }, + "years": { + "additionalProperties": { + "additionalProperties": true, + "type": "object" + }, + "type": "object", + "title": "Years" + }, + "geographies": { + "items": { + "additionalProperties": true, + "type": "object" + }, + "type": "array", + "title": "Geographies" + }, + "composition_method": { + "type": "string", + "title": "Composition Method" + }, + "storage_method": { + "type": "string", + "title": "Storage Method" + } + }, + "additionalProperties": false, + "type": "object", + "required": [ + "forecast_id", + "forecast_sha256", + "scenario", + "geography_kind", + "runtime_versions", + "years", + "geographies", + "composition_method", + "storage_method" + ], + "title": "SPMProvenance", + "description": "Detached calculation receipt; data certification is a separate claim." + }, + "SPMSelection": { + "properties": { + "forecast_content_sha256": { + "anyOf": [ + { + "type": "string", + "pattern": "^[0-9a-f]{64}$" + }, + { + "type": "null" + } + ], + "title": "Forecast Content Sha256" + }, + "scenario": { + "anyOf": [ + { + "type": "string", + "minLength": 1, + "pattern": "^\\S+$" + }, + { + "type": "null" + } + ], + "title": "Scenario" + }, + "geography_kind": { + "type": "string", + "enum": [ + "county", + "national", + "metro" + ], + "title": "Geography Kind" + }, + "geography_id": { + "anyOf": [ + { + "type": "string", + "minLength": 1 + }, + { + "type": "null" + } + ], + "title": "Geography Id" + }, + "county_vintage": { + "type": "string", + "pattern": "^[0-9]{4}$", + "title": "County Vintage" + }, + "as_of": { + "anyOf": [ + { + "type": "string" + }, + { + "type": "null" + } + ], + "title": "As Of" + } + }, + "additionalProperties": false, + "type": "object", + "title": "SPMSelection", + "description": "Select from the bundle's pinned artifact; national geography is explicit.\n\nCounty mode reads the household's observed ``county_fips``. A state alone\ndoes not identify an SPM area. These settings contain no provider or path." + }, "SimulationRequest": { "properties": { "country": { "type": "string", "title": "Country" }, + "spm": { + "anyOf": [ + { + "$ref": "#/components/schemas/SPMSelection" + }, + { + "type": "null" + } + ] + }, "version": { "anyOf": [ { @@ -1903,6 +2203,26 @@ }, "SingleYearMacroOutput": { "properties": { + "spm_config": { + "anyOf": [ + { + "$ref": "#/components/schemas/SPMSelection" + }, + { + "type": "null" + } + ] + }, + "spm_provenance": { + "anyOf": [ + { + "$ref": "#/components/schemas/SPMComparisonProvenance" + }, + { + "type": "null" + } + ] + }, "model_version": { "type": "string", "title": "Model Version" @@ -2195,6 +2515,14 @@ }, "VersionsResponse": { "properties": { + "spm_capabilities": { + "additionalProperties": { + "$ref": "#/components/schemas/SPMCapability" + }, + "type": "object", + "title": "Spm Capabilities", + "description": "Canonical SPM capabilities keyed by exact PolicyEngine wrapper version. Resolve a country-model route to its app in the routing maps, then find that app's exact version in the policyengine map. App names, country-model versions and latest aliases are not capability keys." + }, "policyengine": { "$ref": "#/components/schemas/VersionMap" }, diff --git a/projects/policyengine-simulation-gateway/tests/test_endpoints.py b/projects/policyengine-simulation-gateway/tests/test_endpoints.py index b29169cc7..244cc5d8a 100644 --- a/projects/policyengine-simulation-gateway/tests/test_endpoints.py +++ b/projects/policyengine-simulation-gateway/tests/test_endpoints.py @@ -1019,6 +1019,7 @@ def test__given_no_active_state__then_versions_fall_back_to_old_dicts( assert response.status_code == 200 assert response.json() == { + "spm_capabilities": {}, "policyengine": { "latest": "4.9.0", "4.9.0": "old-py-app", @@ -1236,7 +1237,9 @@ def test__given_budget_window_submission__then_initial_poll_returns_seed_state( def test__given_batch_state__then_poll_returns_completed_response( self, mock_modal, client: TestClient ): - from policyengine_simulation_contract.budget_window_state import put_batch_job_state + from policyengine_simulation_contract.budget_window_state import ( + put_batch_job_state, + ) from policyengine_simulation_contract.gateway_models import ( BudgetWindowAnnualImpact, BudgetWindowBatchState, diff --git a/projects/policyengine-simulation-gateway/tests/test_spm_routes.py b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py new file mode 100644 index 000000000..ad7122b27 --- /dev/null +++ b/projects/policyengine-simulation-gateway/tests/test_spm_routes.py @@ -0,0 +1,592 @@ +"""SPM capability resolution on historical and ambiguous registry routes.""" + +from copy import deepcopy + +import pytest + +from fixtures.gateway_endpoints import TEST_ROUTING_STATE + + +CAPABILITY = { + "contract_version": "canonical-spm-v1", + "defaults": {"forecast_content_sha256": "a" * 64, "scenario": "ce_trend"}, +} +ENDPOINTS = [ + ("/simulate/economy/comparison", {"scope": "macro"}), + ( + "/simulate/economy/budget-window", + {"region": "us", "start_year": "2026", "window_size": 2}, + ), +] + + +# ``generation`` is one global marker that every publish rewrites, so a +# seeded route has to resolve the same before and after a canonical deploy. +PUBLISHED_GENERATION = "5.4.0:policyengine-simulation-py5-4-0" +LEGACY_SOURCES = ["legacy-country-dict", "legacy-seed", PUBLISHED_GENERATION] + + +def legacy_route(mock_modal, source, model_version): + """Two shapes of pre-canonical route: the original per-country Modal + dicts, and a seeded routing-state country route the registry ties to no + wrapper bundle. For the latter, ``source`` is the state's ``generation`` + marker, which must not change how the route resolves.""" + app_name = "policyengine-simulation-us1-715-2-uk2-88-20" + if source == "legacy-country-dict": + app_name = "legacy-app" + del mock_modal["dicts"]["simulation-api-routing-state"] + mock_modal["dicts"]["simulation-api-us-versions"] = { + "latest": model_version, + model_version: app_name, + } + else: + mock_modal["dicts"]["simulation-api-routing-state"] = { + "active": { + "schema_version": 1, + "generation": source, + "latest": {"us": model_version}, + "routes": {"policyengine": {}, "us": {model_version: app_name}}, + "bundles": {}, + } + } + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +@pytest.mark.parametrize("source", LEGACY_SOURCES) +@pytest.mark.parametrize("version", [None, "1.715.2"]) +def test_historical_no_wrapper_route_submits_without_spm( + mock_modal, client, endpoint, extra, source, version +): + legacy_route(mock_modal, source, "1.715.2") + response = client.post( + endpoint, json={"country": "us", "version": version, **extra} + ) + assert response.status_code == 200, response.text + assert "spm" not in mock_modal["func"].last_payload + assert response.json()["policyengine_bundle"]["model_version"] == "1.715.2" + assert "policyengine_version" not in response.json()["policyengine_bundle"] + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +@pytest.mark.parametrize("source", LEGACY_SOURCES) +@pytest.mark.parametrize("selection", [{}, {"geography_kind": "national"}]) +def test_historical_route_rejects_any_explicit_spm( + mock_modal, client, endpoint, extra, source, selection +): + legacy_route(mock_modal, source, "1.715.2") + response = client.post(endpoint, json={"country": "us", "spm": selection, **extra}) + assert response.status_code == 400 + assert response.json()["errors"][0]["code"] == "SPM_CONFIGURATION_UNAVAILABLE" + assert mock_modal["func"].calls == [] + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +@pytest.mark.parametrize( + "source,model_version", + [ + ("legacy-country-dict", "1.824.7"), + ("legacy-seed", "1.824.7"), + (PUBLISHED_GENERATION, "1.824.7"), + ("legacy-seed", "future-model"), + (PUBLISHED_GENERATION, "future-model"), + ], +) +def test_missing_wrapper_does_not_certify_unknown_routes( + mock_modal, client, endpoint, extra, source, model_version +): + legacy_route(mock_modal, source, model_version) + response = client.post(endpoint, json={"country": "us", **extra}) + assert response.status_code == 400 + assert response.json()["errors"][0]["code"] == "SPM_CONFIGURATION_UNAVAILABLE" + assert mock_modal["func"].calls == [] + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +def test_legacy_seed_cannot_erase_future_wrapper_from_app_name( + mock_modal, client, endpoint, extra +): + legacy_route(mock_modal, "legacy-seed", "1.715.2") + state = mock_modal["dicts"]["simulation-api-routing-state"]["active"] + state["routes"]["us"]["1.715.2"] = "policyengine-simulation-py99-0-0" + response = client.post(endpoint, json={"country": "us", **extra}) + assert response.status_code == 400 + assert response.json()["errors"][0]["code"] == "SPM_CONFIGURATION_UNAVAILABLE" + assert mock_modal["func"].calls == [] + + +@pytest.mark.parametrize("schema_version", [None, 2]) +def test_legacy_seed_requires_supported_registry_schema( + mock_modal, client, schema_version +): + legacy_route(mock_modal, "legacy-seed", "1.715.2") + state = mock_modal["dicts"]["simulation-api-routing-state"]["active"] + state["schema_version"] = schema_version + response = client.post("/simulate/economy/comparison", json={"country": "us"}) + assert response.status_code == 400 + assert response.json()["errors"][0]["code"] == "SPM_CONFIGURATION_UNAVAILABLE" + assert mock_modal["func"].calls == [] + + +def shared_app_state(mock_modal, *, sibling_model): + state = deepcopy(TEST_ROUTING_STATE) + original = state["bundles"]["4.10.0"] + app_name = original["app_name"] + state["routes"]["policyengine"]["5.3.0"] = app_name + state["bundles"]["5.3.0"] = { + **deepcopy(original), + "policyengine_version": "5.3.0", + "us": {**original["us"], "model_version": sibling_model}, + "spm": deepcopy(CAPABILITY), + } + state["routes"]["us"][sibling_model] = app_name + mock_modal["dicts"]["simulation-api-routing-state"] = {"active": state} + return state + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +def test_shared_app_resolves_unique_country_model_bundle( + mock_modal, client, endpoint, extra +): + shared_app_state(mock_modal, sibling_model="1.824.7") + response = client.post( + endpoint, json={"country": "us", "version": "1.824.7", **extra} + ) + assert response.status_code == 200, response.text + assert response.json()["policyengine_bundle"]["policyengine_version"] == "5.3.0" + assert mock_modal["func"].last_payload["spm"]["scenario"] == "ce_trend" + + +def test_latest_route_alias_does_not_create_ambiguous_bundle(mock_modal, client): + state = shared_app_state(mock_modal, sibling_model="1.824.7") + state["routes"]["policyengine"]["latest"] = state["routes"]["policyengine"]["5.3.0"] + response = client.post( + "/simulate/economy/comparison", json={"country": "us", "version": "1.824.7"} + ) + assert response.status_code == 200 + assert response.json()["policyengine_bundle"]["policyengine_version"] == "5.3.0" + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +def test_shared_app_with_ambiguous_country_model_requires_explicit_bundle( + mock_modal, client, endpoint, extra +): + shared_app_state(mock_modal, sibling_model="1.500.0") + response = client.post( + endpoint, json={"country": "us", "version": "1.500.0", **extra} + ) + assert response.status_code == 400 + assert "policyengine_version" in response.json()["detail"] + assert mock_modal["func"].calls == [] + explicit = client.post( + endpoint, + json={"country": "us", "policyengine_version": "5.3.0", **extra}, + ) + assert explicit.status_code == 200 + assert mock_modal["func"].last_payload["spm"]["scenario"] == "ce_trend" + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +@pytest.mark.parametrize("sibling_model", [None, "", " ", 42]) +def test_shared_app_with_missing_model_metadata_is_ambiguous( + mock_modal, client, endpoint, extra, sibling_model +): + state = shared_app_state(mock_modal, sibling_model="1.824.7") + state["bundles"]["4.10.0"]["us"]["model_version"] = sibling_model + response = client.post( + endpoint, json={"country": "us", "version": "1.824.7", **extra} + ) + assert response.status_code == 400 + assert "policyengine_version" in response.json()["detail"] + assert mock_modal["func"].calls == [] + + +@pytest.mark.parametrize( + "invalid_capability", + [ + {"contract_version": "canonical-spm-v1", "defaults": {"scenario": "ce_trend"}}, + {**CAPABILITY, "contract_version": "unknown-contract"}, + {**CAPABILITY, "undeclared": True}, + {"defaults": {**CAPABILITY["defaults"], "geography_kind": "metro"}}, + ], +) +def test_versions_omits_malformed_capability_and_submission_rejects_it( + mock_modal, client, invalid_capability +): + state = shared_app_state(mock_modal, sibling_model="1.824.7") + state["bundles"]["4.10.0"]["spm"] = invalid_capability + response = client.get("/versions") + assert response.status_code == 200 + assert set(response.json()["spm_capabilities"]) == {"5.3.0"} + rejected = client.post( + "/simulate/economy/comparison", + json={"country": "us", "policyengine_version": "4.10.0"}, + ) + assert rejected.status_code == 400 + assert rejected.json()["errors"][0]["code"] == "SPM_CONFIGURATION_UNAVAILABLE" + assert mock_modal["func"].calls == [] + + +RESOLVED_SELECTION = { + "forecast_content_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, +} + + +def spm_receipt(year): + return { + "forecast_id": "ce-forecast", + "forecast_sha256": "a" * 64, + "scenario": "ce_trend", + "geography_kind": "national", + "runtime_versions": {"spm_calculator": "0.3.1"}, + "years": {str(year): {"entry": f"{year}-01-01"}}, + "geographies": [{"kind": "national"}], + "composition_method": "national", + "storage_method": "artifact", + } + + +def spm_provenance(year): + return {"baseline": [spm_receipt(year)], "reform": [spm_receipt(year)]} + + +def test_completed_result_body_keeps_resolved_nulls(mock_modal, client): + """A client must be able to replay the selection it was handed back. + + ``response_model_exclude_none`` applies to every poll body, and an + omitted option inherits the bundle default. Dropping a resolved null + would silently re-resolve ``as_of`` and ``geography_id`` on the next + request, changing the baseline key and the receipts. + """ + shared_app_state(mock_modal, sibling_model="1.824.7") + submitted = client.post( + "/simulate/economy/comparison", + json={ + "country": "us", + "version": "1.824.7", + "spm": {"geography_kind": "national", "as_of": None}, + }, + ) + assert submitted.status_code == 200, submitted.text + assert mock_modal["func"].last_payload["spm"] == RESOLVED_SELECTION + + job_id = submitted.json()["job_id"] + call = mock_modal["function_call"].registry[job_id] + call.result = { + **call.result, + "spm_config": RESOLVED_SELECTION, + "spm_provenance": spm_provenance(2026), + } + + polled = client.get(f"/jobs/{job_id}") + assert polled.status_code == 200, polled.text + assert polled.json()["result"]["spm_config"] == RESOLVED_SELECTION + + +def test_completed_budget_window_rows_keep_resolved_nulls(mock_modal, client): + """Each ``annualImpacts`` row carries the same replayable selection.""" + from policyengine_simulation_contract.budget_window_state import ( + put_batch_job_state, + ) + from policyengine_simulation_contract.gateway_models import ( + BudgetWindowAnnualImpact, + BudgetWindowBatchState, + BudgetWindowResult, + BudgetWindowTotals, + PolicyEngineBundle, + ) + + shared_app_state(mock_modal, sibling_model="1.824.7") + impact = BudgetWindowAnnualImpact( + spm_config=RESOLVED_SELECTION, + spm_provenance=spm_provenance(2026), + year="2026", + taxRevenueImpact=10, + federalTaxRevenueImpact=7, + stateTaxRevenueImpact=3, + benefitSpendingImpact=5, + budgetaryImpact=15, + ) + put_batch_job_state( + BudgetWindowBatchState( + batch_job_id="mock-batch-job-id-123", + status="complete", + country="us", + region="us", + version="1.824.7", + target="general", + resolved_app_name="policyengine-simulation-py4-10-0", + policyengine_bundle=PolicyEngineBundle(model_version="1.824.7"), + start_year="2026", + window_size=1, + max_parallel=1, + request_payload={"country": "us", "spm": RESOLVED_SELECTION}, + years=["2026"], + queued_years=[], + running_years=[], + completed_years=["2026"], + failed_years=[], + child_jobs={}, + partial_annual_impacts={}, + result=BudgetWindowResult( + startYear="2026", + endYear="2026", + windowSize=1, + annualImpacts=[impact], + totals=BudgetWindowTotals( + taxRevenueImpact=10, + federalTaxRevenueImpact=7, + stateTaxRevenueImpact=3, + benefitSpendingImpact=5, + budgetaryImpact=15, + ), + ), + error=None, + created_at="2026-01-01T00:00:00+00:00", + updated_at="2026-01-01T00:00:01+00:00", + run_id="batch-run-123", + ) + ) + + polled = client.get("/budget-window-jobs/mock-batch-job-id-123") + assert polled.status_code == 200, polled.text + row = polled.json()["result"]["annualImpacts"][0] + assert row["spm_config"] == RESOLVED_SELECTION + + +# Every optional field the pre-SPM gateway emitted in a raw 202/500 body. +LEGACY_BUNDLE_KEYS = { + "model_version", + "policyengine_version", + "data_version", + "dataset", +} + + +@pytest.mark.parametrize("source", LEGACY_SOURCES) +def test_legacy_poll_bodies_carry_no_spm_key(mock_modal, client, source): + """The 202 and 500 bodies splat job metadata in raw, bypassing the + routes' ``response_model_exclude_none``. A legacy no-SPM job must stay + byte-identical to the pre-SPM gateway.""" + legacy_route(mock_modal, source, "1.715.2") + submitted = client.post("/simulate/economy/comparison", json={"country": "us"}) + assert submitted.status_code == 200, submitted.text + job_id = submitted.json()["job_id"] + call = mock_modal["function_call"].registry[job_id] + + call.running = True + running = client.get(f"/jobs/{job_id}") + assert running.status_code == 202 + assert set(running.json()["policyengine_bundle"]) == LEGACY_BUNDLE_KEYS + + call.running = False + call.error = RuntimeError("worker exploded") + failed = client.get(f"/jobs/{job_id}") + assert failed.status_code == 500 + assert set(failed.json()["policyengine_bundle"]) == LEGACY_BUNDLE_KEYS + + +@pytest.mark.parametrize("source", LEGACY_SOURCES) +def test_legacy_budget_window_metadata_carries_no_spm_key(mock_modal, client, source): + """The same bundle rides to the worker in ``_metadata``.""" + legacy_route(mock_modal, source, "1.715.2") + submitted = client.post( + "/simulate/economy/budget-window", + json={"country": "us", "region": "us", "start_year": "2026", "window_size": 2}, + ) + assert submitted.status_code == 200, submitted.text + metadata = mock_modal["func"].last_payload["_metadata"] + assert set(metadata["policyengine_bundle"]) == LEGACY_BUNDLE_KEYS + + +def test_canonical_poll_bodies_still_carry_the_capability(mock_modal, client): + """Omission is scoped to routes with no capability, not to all routes.""" + shared_app_state(mock_modal, sibling_model="1.824.7") + submitted = client.post( + "/simulate/economy/comparison", json={"country": "us", "version": "1.824.7"} + ) + assert submitted.status_code == 200, submitted.text + job_id = submitted.json()["job_id"] + mock_modal["function_call"].registry[job_id].running = True + + running = client.get(f"/jobs/{job_id}") + assert running.status_code == 202 + bundle = running.json()["policyengine_bundle"] + assert set(bundle) == LEGACY_BUNDLE_KEYS | {"spm"} + assert bundle["spm"]["defaults"]["scenario"] == "ce_trend" + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +def test_publishing_does_not_revoke_a_seeded_country_route( + mock_modal, client, endpoint, extra +): + """The route is unchanged across a publish; only the global marker moved. + + Keying the allowance on ``generation`` meant the first canonical deploy + turned every seeded country route's bare submission into a 400. + """ + for generation in ("legacy-seed", PUBLISHED_GENERATION): + legacy_route(mock_modal, generation, "1.715.2") + response = client.post(endpoint, json={"country": "us", **extra}) + assert response.status_code == 200, f"{generation}: {response.text}" + assert "spm" not in mock_modal["func"].last_payload + + +def wrapper_route_without_manifest(mock_modal, wrapper): + """A seeded wrapper route the app-release snapshot had no bundle for.""" + state = deepcopy(TEST_ROUTING_STATE) + app_name = f"policyengine-simulation-py{wrapper.replace('.', '-')}" + state["routes"]["policyengine"][wrapper] = app_name + mock_modal["dicts"]["simulation-api-routing-state"] = {"active": state} + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +@pytest.mark.parametrize("wrapper", ["5.2.0", "5.3.0"]) +def test_pinned_wrapper_route_without_manifest_submits_without_spm( + mock_modal, client, endpoint, extra, wrapper +): + """With no manifest the response's ``model_version`` echoes the wrapper + version. Reading that back as a country model version rejected the two + pinned pre-canonical bundles.""" + wrapper_route_without_manifest(mock_modal, wrapper) + response = client.post( + endpoint, json={"country": "us", "policyengine_version": wrapper, **extra} + ) + assert response.status_code == 200, response.text + assert "spm" not in mock_modal["func"].last_payload + assert response.json()["policyengine_bundle"]["model_version"] == wrapper + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +@pytest.mark.parametrize("wrapper", ["5.2.0", "5.3.0"]) +def test_pinned_wrapper_through_legacy_dicts_submits_without_spm( + mock_modal, client, endpoint, extra, wrapper +): + """Same route shape through the pre-registry Modal dicts.""" + del mock_modal["dicts"]["simulation-api-routing-state"] + app_name = f"policyengine-simulation-py{wrapper.replace('.', '-')}" + mock_modal["dicts"]["simulation-api-policyengine-versions"] = {wrapper: app_name} + response = client.post( + endpoint, json={"country": "us", "policyengine_version": wrapper, **extra} + ) + assert response.status_code == 200, response.text + assert "spm" not in mock_modal["func"].last_payload + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +def test_unpinned_wrapper_route_without_manifest_still_fails_closed( + mock_modal, client, endpoint, extra +): + """Only the two pinned pre-canonical wrappers are vouched for without a + manifest; an unrecognized bundle proves nothing.""" + wrapper_route_without_manifest(mock_modal, "5.4.0") + response = client.post( + endpoint, json={"country": "us", "policyengine_version": "5.4.0", **extra} + ) + assert response.status_code == 400 + assert response.json()["errors"][0]["code"] == "SPM_CONFIGURATION_UNAVAILABLE" + assert mock_modal["func"].calls == [] + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +def test_stated_model_version_still_contradicts_the_wrapper_pin( + mock_modal, client, endpoint, extra +): + """A manifest that states a post-canonical model under a pinned wrapper + is a republish we cannot vouch for.""" + state = deepcopy(TEST_ROUTING_STATE) + original = state["bundles"]["4.10.0"] + state["routes"]["policyengine"]["5.3.0"] = original["app_name"] + state["bundles"]["5.3.0"] = { + **deepcopy(original), + "policyengine_version": "5.3.0", + "us": {**original["us"], "model_version": "1.824.7"}, + } + mock_modal["dicts"]["simulation-api-routing-state"] = {"active": state} + response = client.post( + endpoint, json={"country": "us", "policyengine_version": "5.3.0", **extra} + ) + assert response.status_code == 400 + assert response.json()["errors"][0]["code"] == "SPM_CONFIGURATION_UNAVAILABLE" + assert mock_modal["func"].calls == [] + + +def republished_wrapper_state(mock_modal, *, model_version): + """The registry a re-publish leaves behind. + + ``update_version_registry`` only ever adds country routes and it + overwrites the bundle manifest for the wrapper version it deploys, so + re-publishing one wrapper with an upgraded country model leaves + ``routes["us"]["1.459.0"]`` pointing at an app whose manifest now states + something else. Nothing prunes it. + """ + state = deepcopy(TEST_ROUTING_STATE) + state["bundles"]["3.9.0"]["us"] = { + **state["bundles"]["3.9.0"]["us"], + "model_version": model_version, + } + mock_modal["dicts"]["simulation-api-routing-state"] = {"active": state} + return state + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +def test_stale_country_route_is_refused_rather_than_served( + mock_modal, client, endpoint, extra +): + """The behaviour change. Before, this route resolved and the 202 body + contradicted itself: ``version`` from the stale route, + ``policyengine_bundle.model_version`` from the live manifest.""" + republished_wrapper_state(mock_modal, model_version="1.470.0") + response = client.post( + endpoint, json={"country": "us", "version": "1.459.0", **extra} + ) + assert response.status_code == 400 + detail = response.json()["detail"] + assert "1.459.0" in detail and "1.470.0" in detail + # A routing refusal, not an SPM one. + assert "errors" not in response.json() + assert mock_modal["func"].calls == [] + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +def test_a_country_route_its_bundle_still_states_is_served( + mock_modal, client, endpoint, extra +): + """The other half: a route the manifest agrees with keeps working, so + the refusal above is about disagreement and not about single-candidate + country routes in general.""" + republished_wrapper_state(mock_modal, model_version="1.459.0") + response = client.post( + endpoint, json={"country": "us", "version": "1.459.0", **extra} + ) + assert response.status_code == 200, response.text + bundle = response.json()["policyengine_bundle"] + assert bundle["policyengine_version"] == "3.9.0" + assert bundle["model_version"] == "1.459.0" + + +@pytest.mark.parametrize("endpoint,extra", ENDPOINTS) +@pytest.mark.parametrize("model_version", [None, "", " ", 42]) +def test_a_bundle_that_states_no_model_version_contradicts_nothing( + mock_modal, client, endpoint, extra, model_version +): + """Unstated is unstated, whichever way it is unstated. + + The ambiguity check classifies a blank or whitespace model version as + unstated; the shared validator reads any string as stated. Deriving the + same fact twice made a blank manifest entry a 400 here while + ``test_shared_app_with_missing_model_metadata_is_ambiguous`` treats the + identical value as missing metadata. + """ + state = republished_wrapper_state(mock_modal, model_version=model_version) + if model_version is None: + del state["bundles"]["3.9.0"]["us"]["model_version"] + response = client.post( + endpoint, json={"country": "us", "version": "1.459.0", **extra} + ) + assert response.status_code == 200, response.text + assert response.json()["policyengine_bundle"]["policyengine_version"] == "3.9.0"