From 6e2dcd4ee3202ffead3f5cf9679e57a638f961b4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:43:23 +0400 Subject: [PATCH 01/25] Key and lay out Stage 12's per-state datasets Stage 12 gets its own copy of the per-state dataset identity, without importing the main path's precompute modules. Stage12StateDatasetCacheIdentity (schema stage12-sds1) holds a dataset identity's fields for the regional dataset (its own SHA-256 and artifact revision from the published Stage 12 bundle, the package versions and the bundle manifest digest) plus the state code and FIPS and the filter variable. Its store path is state-datasets///_year_.h5. stage12_cache/state_datasets.py places each state's files at /state-datasets//. The national dataset and baseline identities are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../stage12_cache/keys.py | 109 ++++++++++++ .../stage12_cache/state_datasets.py | 25 +++ .../tests/test_stage12_state_datasets.py | 158 ++++++++++++++++++ 3 files changed, 292 insertions(+) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py create mode 100644 projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py index f2d34788e..a31f41561 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py @@ -21,6 +21,7 @@ ) DATASET_SCHEMA = "stage12-ds1" +STATE_DATASET_SCHEMA = "stage12-sds1" BASELINE_SCHEMA = "stage12-bl1" CURRENT_LAW_POLICY = "current-law" _ID_DIGEST_CHARS = 16 @@ -95,6 +96,66 @@ def store_path(self) -> str: return f"datasets/{self.country}/{self.digest}/{self.filename}" +@dataclass(frozen=True) +class Stage12StateDatasetCacheIdentity: + """Identity of one state's cached year file, filtered from a regional dataset. + + The file is ``dataset``'s year file restricted to the households whose + ``filter_variable`` equals ``state_fips``. Its fields are those of a + dataset identity for the regional dataset, plus the state and filter. + """ + + country: CountryId + dataset: str + stem: str + year: int + source_sha256: str + data_artifact_revision: str + data_package_version: str + country_package_version: str + core_package_version: str + policyengine_version: str + bundle_manifest_sha256: str + state_code: str + state_fips: int + filter_variable: str + + @property + def payload(self) -> dict[str, Any]: + """The fields hashed into ``digest``.""" + return { + "schema": STATE_DATASET_SCHEMA, + "country": self.country, + "dataset": self.dataset, + "year": self.year, + "source_sha256": self.source_sha256, + "data_artifact_revision": self.data_artifact_revision, + "data_package_version": self.data_package_version, + "country_package_version": self.country_package_version, + "core_package_version": self.core_package_version, + "policyengine_version": self.policyengine_version, + "bundle_manifest_sha256": self.bundle_manifest_sha256, + "state_code": self.state_code, + "state_fips": self.state_fips, + "filter_variable": self.filter_variable, + } + + @property + def digest(self) -> str: + """The canonical digest of ``payload``; it names the cache directory.""" + return canonical_digest(self.payload) + + @property + def filename(self) -> str: + """The year file's own name, which every state's file shares.""" + return f"{self.stem}_year_{self.year}.h5" + + @property + def store_path(self) -> str: + """The object path in the cache bucket.""" + return f"state-datasets/{self.country}/{self.digest}/{self.filename}" + + @dataclass(frozen=True) class Stage12BaselineCacheIdentity: """Identity of one cached current-law baseline for a group of regions. @@ -193,6 +254,54 @@ def collect_dataset_identity( ) +def collect_state_dataset_identity( # noqa: PLR0913 # one keyword per input of the per-state identity + resolved: ResolvedStage12Bundle, + *, + country: CountryId, + year: int, + dataset: str, + state_code: str, + state_fips: int, + filter_variable: str, +) -> Stage12StateDatasetCacheIdentity: + """Collect one state's dataset identity from the installed Stage 12 bundle. + + ``dataset`` names a bundle dataset other than the default, such as the + regional dataset state regions load; its own hash and revision key the + file. + """ + from policyengine.provenance.manifest import ( + dataset_logical_name, + resolve_dataset_reference, + ) + + country_bundle = next( + item for item in resolved.bundle.countries if item.country == country + ) + entry = next( + (item for item in country_bundle.datasets if item.identity == dataset), None + ) + if entry is None: + msg = f"Dataset {dataset!r} is not in the installed {country} Stage 12 bundle" + raise ValueError(msg) + return Stage12StateDatasetCacheIdentity( + country=country, + dataset=entry.identity, + stem=dataset_logical_name(resolve_dataset_reference(country, dataset)), + year=year, + source_sha256=entry.sha256, + data_artifact_revision=entry.artifact_revision, + data_package_version=country_bundle.data_package_version, + country_package_version=country_bundle.country_package_version, + core_package_version=resolved.bundle.core_package_version, + policyengine_version=resolved.bundle.policyengine_version, + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + state_code=state_code.upper(), + state_fips=int(state_fips), + filter_variable=filter_variable, + ) + + def collect_baseline_identity( simulation: PlannedSimulationExecutionInput, *, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py new file mode 100644 index 000000000..9646704af --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py @@ -0,0 +1,25 @@ +"""Where the Stage 12 US worker image keeps baked per-state datasets. + +Stage 12's own copy of the main path's per-state dataset layout (see +``policyengine_simulation_executor.state_datasets``): Stage 12 does not +import the main path's precompute modules. A state's year files sit at +``/state-datasets///`` under the +basename the dataset's own year file has, so ``ensure_datasets`` called on +that folder finds the state's file where it would look for the full one. +""" + +from __future__ import annotations + +from pathlib import Path + +STATE_DATASETS_DIRECTORY = "state-datasets" + + +def state_dataset_relative_folder(dataset: str, state_code: str) -> str: + """Return a state's folder relative to the Stage 12 data folder.""" + return f"{STATE_DATASETS_DIRECTORY}/{dataset}/{state_code.upper()}" + + +def state_dataset_folder(data_folder: str, dataset: str, state_code: str) -> str: + """Return the absolute folder of a state's baked year files.""" + return str(Path(data_folder) / state_dataset_relative_folder(dataset, state_code)) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py b/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py new file mode 100644 index 000000000..4432b205a --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py @@ -0,0 +1,158 @@ +"""Stage 12's per-state datasets: their cache identity and image folders. + +Stage 12 keeps its own copy of the main path's per-state dataset machinery, +keyed from the published Stage 12 bundle. +""" + +from __future__ import annotations + +import dataclasses + +import pytest +from policyengine_simulation_contract.stage12_bundle import ( + ResolvedStage12Bundle, + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, +) +from policyengine_simulation_executor.stage12_cache import state_datasets +from policyengine_simulation_executor.stage12_cache.keys import ( + Stage12StateDatasetCacheIdentity, + collect_dataset_identity, + collect_state_dataset_identity, +) + +REGIONAL_DATASET = "populace_us_2024_acs_local" + + +def _resolved() -> ResolvedStage12Bundle: + national = Stage12Dataset( + identity="populace_us_2024", + uri="hf://policyengine/populace-us/populace_us_2024.h5@national-revision", + artifact_revision="national-revision", + sha256="a" * 64, + repo_type="dataset", + ) + regional = Stage12Dataset( + identity=REGIONAL_DATASET, + uri=f"hf://policyengine/populace-us/{REGIONAL_DATASET}.h5@regional-revision", + artifact_revision="regional-revision", + sha256="c" * 64, + repo_type="dataset", + ) + country = Stage12CountryBundle( + country="us", + country_package_name="policyengine-us", + country_package_version="1.0.0", + country_package_requirement="policyengine-us==1.0.0", + data_package_name="policyengine-us-data", + data_package_version="2.0.0", + data_release_version="2.0.0", + data_artifact_revision="national-revision", + default_dataset=national.identity, + default_dataset_uri=national.uri, + datasets=(national, regional), + ) + return ResolvedStage12Bundle( + bundle=Stage12BundleManifest( + policyengine_version="5.2.0", + policyengine_requirement="policyengine==5.2.0", + core_package_version="3.0.0", + core_package_requirement="policyengine-core==3.0.0", + countries=(country,), + ), + bundle_manifest_sha256="b" * 64, + ) + + +def _utah_2026() -> Stage12StateDatasetCacheIdentity: + return collect_state_dataset_identity( + _resolved(), + country="us", + year=2026, + dataset=REGIONAL_DATASET, + state_code="ut", + state_fips=49, + filter_variable="state_fips", + ) + + +def test_identity_is_taken_from_the_regional_dataset_entry(): + identity = _utah_2026() + + assert identity.dataset == REGIONAL_DATASET + assert identity.source_sha256 == "c" * 64 + assert identity.data_artifact_revision == "regional-revision" + assert identity.state_code == "UT" + assert identity.payload == { + "schema": "stage12-sds1", + "country": "us", + "dataset": REGIONAL_DATASET, + "year": 2026, + "source_sha256": "c" * 64, + "data_artifact_revision": "regional-revision", + "data_package_version": "2.0.0", + "country_package_version": "1.0.0", + "core_package_version": "3.0.0", + "policyengine_version": "5.2.0", + "bundle_manifest_sha256": "b" * 64, + "state_code": "UT", + "state_fips": 49, + "filter_variable": "state_fips", + } + + +def test_file_keeps_the_year_file_name_under_its_own_store_path(): + identity = _utah_2026() + + assert identity.filename == f"{REGIONAL_DATASET}_year_2026.h5" + assert identity.store_path == ( + f"state-datasets/us/{identity.digest}/{REGIONAL_DATASET}_year_2026.h5" + ) + + +@pytest.mark.parametrize( + "field", + [ + field.name + for field in dataclasses.fields(Stage12StateDatasetCacheIdentity) + if field.name not in {"stem", "country"} + ], +) +def test_every_identity_field_changes_the_digest(field): + identity = _utah_2026() + replacement = {"year": 2027, "state_fips": 56}.get(field, "other") + + assert dataclasses.replace(identity, **{field: replacement}).digest != ( + identity.digest + ) + + +def test_refuses_a_dataset_missing_from_the_bundle(): + with pytest.raises(ValueError, match="populace_missing"): + collect_state_dataset_identity( + _resolved(), + country="us", + year=2026, + dataset="populace_missing", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + ) + + +def test_the_national_dataset_identity_is_unchanged(): + identity = collect_dataset_identity(_resolved(), country="us", year=2026) + + assert identity.dataset == "populace_us_2024" + assert identity.source_sha256 == "a" * 64 + assert identity.payload["schema"] == "stage12-ds1" + + +def test_state_folder_sits_under_the_stage12_data_folder(): + assert state_datasets.state_dataset_relative_folder(REGIONAL_DATASET, "ut") == ( + f"state-datasets/{REGIONAL_DATASET}/UT" + ) + assert state_datasets.state_dataset_folder( + "/opt/policyengine/stage12-data", REGIONAL_DATASET, "UT" + ) == (f"/opt/policyengine/stage12-data/state-datasets/{REGIONAL_DATASET}/UT") From 98afd38e0c506f9d70630cca3f877a1f08421eb6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:46:53 +0400 Subject: [PATCH 02/25] Plan Stage 12's per-state datasets for each state on a regional dataset Stage 12's precompute gets its own state_dataset_targets, which resolves every US state region the way a Stage 12 worker resolves a request for it and keeps the states whose dataset is not the bundle default, with their row filter and state code. The plan lists one StateDatasetPlanEntry per precompute year and target state, year-major and by state code, with its key, store path, file name, image folder and presence in the bucket; select_work keeps the missing ones. On the installed bundle that is 51 states for each of the three years, all on the ACS-local dataset. Nothing builds or publishes these entries yet. Co-Authored-By: Claude Opus 5.5 --- .../stage12_cache/models.py | 23 +- .../stage12_cache/precompute.py | 98 +++++++- .../tests/test_stage12_cache_precompute.py | 231 ++++++++++++++++++ 3 files changed, 350 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py index dcba07bdb..40a7b442a 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py @@ -35,6 +35,25 @@ class DatasetPlanEntry(StrictCacheModel): exists: bool +class StateDatasetPlanEntry(StrictCacheModel): + """One state's planned year file, filtered from a regional dataset. + + ``folder`` is where the worker image places the file, relative to its + data folder; every state's file has the same ``filename``. ``exists`` + records whether the bucket already holds it. + """ + + year: int + state_code: str + region: str + dataset: str + digest: str + path: str + filename: str + folder: str + exists: bool + + class BaselinePlanEntry(StrictCacheModel): """One planned baseline for a year and region group. @@ -67,13 +86,15 @@ class PrecomputePlan(StrictCacheModel): bundle: CacheBundleIdentity partition_sha256: str determinism_verified: bool + state_datasets: tuple[StateDatasetPlanEntry, ...] = () class WorkSelection(StrictCacheModel): - """The planned datasets and baselines a run still has to build.""" + """The planned datasets, state datasets and baselines a run still has to build.""" datasets: tuple[DatasetPlanEntry, ...] baselines: tuple[BaselinePlanEntry, ...] + state_datasets: tuple[StateDatasetPlanEntry, ...] = () class CacheObject(StrictCacheModel): diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index fd2b122c5..1dffd2c71 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -3,7 +3,7 @@ from __future__ import annotations import time -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any from uuid import NAMESPACE_URL, uuid5 @@ -26,6 +26,7 @@ from policyengine_simulation_executor import simulation_runtime from policyengine_simulation_executor.simulation_runtime import ( _country_module, + _resolve_dataset_selection, _resolve_region, ) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle @@ -49,6 +50,7 @@ canonical_digest, collect_baseline_identity, collect_dataset_identity, + collect_state_dataset_identity, ) from .models import ( BaselineBuildResult, @@ -60,8 +62,10 @@ DeterminismVerdict, ManifestArtifact, PrecomputePlan, + StateDatasetPlanEntry, WorkSelection, ) +from .state_datasets import state_dataset_relative_folder from .store import Stage12CacheStore if TYPE_CHECKING: @@ -213,6 +217,67 @@ def determinism_record_path(baselines: Sequence[BaselinePlanEntry]) -> str: return f"determinism/us/{canonical_digest(determinism_record(baselines))}.json" +@dataclass(frozen=True) +class StateDatasetTarget: + """A US state whose requests load a non-default dataset, and its row filter. + + Stage 12 bakes this state's rows of ``dataset`` as their own year file: + the households whose ``filter_variable`` equals ``state_fips``. + ``scoping_strategy`` is the region's own .py strategy, the one a request + for the state applies. + """ + + region: str + state_code: str + state_fips: int + filter_variable: str + dataset: str + scoping_strategy: Any = field(default=None, compare=False, repr=False) + + +def state_dataset_targets() -> tuple[StateDatasetTarget, ...]: + """Return every US state region that loads a non-default dataset, by state code. + + Each state is resolved as a Stage 12 worker resolves a request for it + (``_resolve_region`` then ``_resolve_dataset_selection``), so a bundle + that sends states back to the default dataset gets no state files. + """ + country_module = _country_module("us") + targets = [] + for region in country_module.model.region_registry.get_by_type("state"): + params = { + "country": "us", + "region": region.code, + "time_period": PRECOMPUTE_YEARS[0], + } + resolution = _resolve_region( + country_module=country_module, country="us", params=params + ) + selection = _resolve_dataset_selection(params, region_resolution=resolution) + if selection.is_default: + continue + strategy = resolution.scoping_strategy + variable_name = getattr(strategy, "variable_name", None) + variable_value = getattr(strategy, "variable_value", None) + if not variable_name or variable_value is None or not resolution.state_code: + msg = ( + f"State region {resolution.code} loads {selection.name} but has no " + "row filter and state code to build its state file from" + ) + raise RuntimeError(msg) + targets.append( + StateDatasetTarget( + region=resolution.code, + state_code=resolution.state_code, + state_fips=int(variable_value), + filter_variable=variable_name, + dataset=selection.name, + scoping_strategy=strategy, + ) + ) + return tuple(sorted(targets, key=lambda target: target.state_code)) + + def plan_artifacts_impl(bucket: str) -> PrecomputePlan: """Plan every dataset and baseline file and record which the bucket holds. @@ -221,8 +286,10 @@ def plan_artifacts_impl(bucket: str) -> PrecomputePlan: """ resolved = load_stage12_bundle() store = Stage12CacheStore(bucket) + targets = state_dataset_targets() datasets = [] baselines = [] + state_datasets = [] for year in PRECOMPUTE_YEARS: dataset = collect_dataset_identity(resolved, country="us", year=year) datasets.append( @@ -251,10 +318,36 @@ def plan_artifacts_impl(bucket: str) -> PrecomputePlan: exists=store.exists(identity.store_path), ) ) + for target in targets: + state_identity = collect_state_dataset_identity( + resolved, + country="us", + year=year, + dataset=target.dataset, + state_code=target.state_code, + state_fips=target.state_fips, + filter_variable=target.filter_variable, + ) + state_datasets.append( + StateDatasetPlanEntry( + year=year, + state_code=target.state_code, + region=target.region, + dataset=target.dataset, + digest=state_identity.digest, + path=state_identity.store_path, + filename=state_identity.filename, + folder=state_dataset_relative_folder( + target.dataset, target.state_code + ), + exists=store.exists(state_identity.store_path), + ) + ) return PrecomputePlan( years=PRECOMPUTE_YEARS, datasets=tuple(datasets), baselines=tuple(baselines), + state_datasets=tuple(state_datasets), bundle=_bundle_identity(), partition_sha256=canonical_digest( {"groups": [list(group) for group in US_REGION_GROUPS]} @@ -268,6 +361,9 @@ def select_work(plan: PrecomputePlan) -> WorkSelection: return WorkSelection( datasets=tuple(entry for entry in plan.datasets if not entry.exists), baselines=tuple(entry for entry in plan.baselines if not entry.exists), + state_datasets=tuple( + entry for entry in plan.state_datasets if not entry.exists + ), ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py index 6b5bb337b..82a0ebbed 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -15,6 +15,7 @@ CacheObject, DatasetPlanEntry, PrecomputePlan, + StateDatasetPlanEntry, ) from policyengine_simulation_executor.stage12_cache.store import Stage12CacheStore from policyengine_simulation_executor.stage12_runtime import ( @@ -76,6 +77,36 @@ def _plan(*, exists: bool = False, verified: bool | None = None) -> PrecomputePl ) +REGIONAL_DATASET = "populace_us_2024_acs_local" + + +def _state_entry(year: int, state_code: str, *, exists: bool) -> StateDatasetPlanEntry: + digest = f"state-{year}-{state_code}" + return StateDatasetPlanEntry( + year=year, + state_code=state_code, + region=f"state/{state_code.lower()}", + dataset=REGIONAL_DATASET, + digest=digest, + path=f"state-datasets/us/{digest}/{REGIONAL_DATASET}_year_{year}.h5", + filename=f"{REGIONAL_DATASET}_year_{year}.h5", + folder=f"state-datasets/{REGIONAL_DATASET}/{state_code}", + exists=exists, + ) + + +def _plan_with_state_datasets(*, exists: bool = False) -> PrecomputePlan: + return _plan(exists=exists).model_copy( + update={ + "state_datasets": tuple( + _state_entry(year, state_code, exists=exists) + for year in precompute.PRECOMPUTE_YEARS + for state_code in ("UT", "WY") + ) + } + ) + + class _FakeBlob: """A GCS blob over a dict, honouring create-only uploads.""" @@ -158,6 +189,186 @@ def test_work_selection_computes_only_missing_files() -> None: assert warm.baselines == () +def test_work_selection_builds_only_missing_state_datasets() -> None: + plan = _plan_with_state_datasets() + stored = plan.state_datasets[0].model_copy(update={"exists": True}) + plan = plan.model_copy( + update={"state_datasets": (stored, *plan.state_datasets[1:])} + ) + + work = precompute.select_work(plan) + + assert [(entry.year, entry.state_code) for entry in work.state_datasets] == [ + (2026, "WY"), + (2027, "UT"), + (2027, "WY"), + (2025, "UT"), + (2025, "WY"), + ] + assert precompute.select_work(_plan()).state_datasets == () + + +def test_plan_with_state_datasets_round_trips_the_modal_boundary() -> None: + plan = _plan_with_state_datasets() + + assert PrecomputePlan.model_validate(plan.model_dump()) == plan + + +@pytest.fixture +def stubbed_planning(monkeypatch): + """Stub everything plan_artifacts_impl reads except the state targets.""" + client = _FakeClient() + _use_store(monkeypatch, client) + monkeypatch.setattr(precompute, "load_stage12_bundle", lambda: "resolved") + monkeypatch.setattr(precompute, "_bundle_identity", _bundle) + monkeypatch.setattr( + precompute, + "collect_dataset_identity", + lambda _resolved, *, country, year: SimpleNamespace( + digest=f"dataset-{year}", + store_path=f"datasets/{country}/dataset-{year}/populace_year_{year}.h5", + filename=f"populace_year_{year}.h5", + ), + ) + monkeypatch.setattr(precompute, "US_REGION_GROUPS", (("state/ca",),)) + monkeypatch.setattr( + precompute, + "_baseline_identity", + lambda year, _group: ( + SimpleNamespace(output_plan=None), + SimpleNamespace( + digest=f"baseline-{year}", + store_path=f"baselines/us/baseline-{year}/stage12-bl1-{year}.h5", + filename=f"stage12-bl1-{year}.h5", + simulation_id=f"stage12-bl1-{year}", + ), + ), + ) + monkeypatch.setattr(precompute, "stage12_output_plan_sha256", lambda _: "b" * 64) + + def fake_state_identity(resolved, *, country, year, dataset, **state_filter): + assert resolved == "resolved" + digest = ( + f"state-{year}-{state_filter['state_code']}-" + f"{state_filter['state_fips']}-{state_filter['filter_variable']}" + ) + return SimpleNamespace( + digest=digest, + store_path=f"state-datasets/{country}/{digest}/{dataset}_year_{year}.h5", + filename=f"{dataset}_year_{year}.h5", + ) + + monkeypatch.setattr( + precompute, "collect_state_dataset_identity", fake_state_identity + ) + targets = [ + precompute.StateDatasetTarget( + region="state/ut", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + dataset=REGIONAL_DATASET, + ), + precompute.StateDatasetTarget( + region="state/wy", + state_code="WY", + state_fips=56, + filter_variable="state_fips", + dataset=REGIONAL_DATASET, + ), + ] + monkeypatch.setattr(precompute, "state_dataset_targets", lambda: tuple(targets)) + return SimpleNamespace(client=client, targets=targets) + + +def test_plan_lists_every_state_dataset_year_major_with_presence( + stubbed_planning, +) -> None: + stubbed_planning.client.objects[ + f"state-datasets/us/state-2026-UT-49-state_fips/{REGIONAL_DATASET}_year_2026.h5" + ] = {"payload": b"h5", "metadata": {"sha256": "f" * 64, "size_bytes": "2"}} + + plan = precompute.plan_artifacts_impl("cache-bucket") + + assert [(entry.year, entry.state_code) for entry in plan.state_datasets] == [ + (2026, "UT"), + (2026, "WY"), + (2027, "UT"), + (2027, "WY"), + (2025, "UT"), + (2025, "WY"), + ] + assert [entry.exists for entry in plan.state_datasets] == [ + True, + False, + False, + False, + False, + False, + ] + first = plan.state_datasets[0] + assert first.region == "state/ut" + assert first.dataset == REGIONAL_DATASET + assert first.folder == f"state-datasets/{REGIONAL_DATASET}/UT" + assert first.filename == f"{REGIONAL_DATASET}_year_2026.h5" + assert len(plan.datasets) == 3 + + +def test_states_on_the_default_dataset_add_no_state_datasets( + stubbed_planning, monkeypatch +) -> None: + monkeypatch.setattr(precompute, "state_dataset_targets", tuple) + + assert precompute.plan_artifacts_impl("cache-bucket").state_datasets == () + + +def test_state_targets_follow_how_a_request_for_the_state_resolves( + monkeypatch, +) -> None: + from policyengine.core.scoping_strategy import RowFilterStrategy + from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + ) + + bundle = get_country_release_bundle("us") + + def region(state_code, fips, dataset_path): + return SimpleNamespace( + code=f"state/{state_code.lower()}", + dataset_path=dataset_path, + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=fips + ), + parent_code="us", + state_code=state_code, + ) + + regions = [ + region("WY", 56, bundle.dataset_uris[REGIONAL_DATASET]), + region("UT", 49, bundle.dataset_uris[REGIONAL_DATASET]), + region("ID", 16, bundle.default_dataset_uri), + ] + by_code = {item.code: item for item in regions} + country_module = SimpleNamespace( + model=SimpleNamespace( + get_region=by_code.get, + region_registry=SimpleNamespace( + get_by_type=lambda kind: list(regions) if kind == "state" else [] + ), + ) + ) + monkeypatch.setattr(precompute, "_country_module", lambda _country: country_module) + + targets = precompute.state_dataset_targets() + + assert [(target.region, target.state_fips) for target in targets] == [ + ("state/ut", 49), + ("state/wy", 56), + ] + assert {target.dataset for target in targets} == {REGIONAL_DATASET} + assert targets[0].scoping_strategy is by_code["state/ut"].scoping_strategy + + class _Call: def __init__(self, result, events, label): self.result = result @@ -549,6 +760,26 @@ def write_json(self, *_): assert downloads == expected_downloads +def test_installed_bundle_plans_a_regional_file_per_state_and_year( + monkeypatch, +) -> None: + """Uses the real bundle and US region registry.""" + client = _FakeClient() + _use_store(monkeypatch, client) + + plan = precompute.plan_artifacts_impl("cache-bucket") + + assert len(plan.state_datasets) == 51 * len(precompute.PRECOMPUTE_YEARS) + assert {entry.dataset for entry in plan.state_datasets} == {REGIONAL_DATASET} + utah = next( + entry + for entry in plan.state_datasets + if entry.year == 2026 and entry.state_code == "UT" + ) + assert utah.folder == f"state-datasets/{REGIONAL_DATASET}/UT" + assert utah.filename == f"{REGIONAL_DATASET}_year_2026.h5" + + def test_planned_baseline_file_is_the_one_the_worker_saves_and_loads( monkeypatch, ) -> None: From ba19a03f34d02be6236d41712e255b13102dc1b7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:49:39 +0400 Subject: [PATCH 03/25] Build, check and upload one year's Stage 12 state datasets Stage 12's own build_state_datasets_impl mirrors the main path's: it re-derives each planned state's key from the installed Stage 12 bundle and refuses to build on a mismatch, has .py prepare the regional dataset's full-country year file once, collects garbage and reloads the year, refuses data in which a household's district code names another state, then for each state writes the rows its region's own scoping strategy keeps with .py's save(), copies the year-file stamp, confirms .py's reuse checks accept the file, and uploads it to the write-once Stage 12 cache, which checks the stored digest and size. Stage 12 keeps its own copy of the reuse check. The fake GCS blob gains upload_from_filename, the first test of a file upload on the Stage 12 side. Co-Authored-By: Claude Opus 5.5 --- .../stage12_cache/models.py | 8 + .../stage12_cache/precompute.py | 160 +++++++++++++- .../stage12_cache/state_datasets.py | 24 +++ .../tests/test_stage12_cache_precompute.py | 198 ++++++++++++++++++ 4 files changed, 389 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py index 40a7b442a..a5bf6591f 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py @@ -113,6 +113,14 @@ class DatasetBuildResult(CacheObject): build_seconds: float +class StateDatasetBuildResult(CacheObject): + """A built state dataset file and whether this run uploaded it.""" + + year: int + state_code: str + uploaded: bool + + class BaselineBuildResult(CacheObject): """A computed baseline file and whether this run uploaded it.""" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index 1dffd2c71..623ad6442 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -2,6 +2,8 @@ from __future__ import annotations +import gc +import tempfile import time from dataclasses import dataclass, field from pathlib import Path @@ -62,10 +64,11 @@ DeterminismVerdict, ManifestArtifact, PrecomputePlan, + StateDatasetBuildResult, StateDatasetPlanEntry, WorkSelection, ) -from .state_datasets import state_dataset_relative_folder +from .state_datasets import state_dataset_relative_folder, year_file_reuse_problem from .store import Stage12CacheStore if TYPE_CHECKING: @@ -73,6 +76,7 @@ import modal import pandas as pd + from microdf import MicroDataFrame PRECOMPUTE_YEARS = (2026, 2027, 2025) MANIFEST_DIGEST_PREFIX = "STAGE12_CACHE_MANIFEST_DIGEST=" @@ -394,6 +398,160 @@ def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildR ) +def _checked_state_targets( + year: int, expected: Sequence[StateDatasetPlanEntry] +) -> dict[str, StateDatasetTarget]: + """Return each planned state's target, after re-deriving its key here. + + Raises ``RuntimeError`` when the installed bundle gives another identity + than the plan, so nothing is built under a key that does not describe it. + """ + resolved = load_stage12_bundle() + targets = {target.region: target for target in state_dataset_targets()} + datasets = set() + for entry in expected: + if entry.year != year: + msg = ( + f"State dataset entry for {entry.state_code} is for year " + f"{entry.year}, not {year}" + ) + raise ValueError(msg) + target = targets.get(entry.region) + if target is None: + msg = f"No state dataset target for {entry.region} in this container" + raise RuntimeError(msg) + identity = collect_state_dataset_identity( + resolved, + country="us", + year=year, + dataset=target.dataset, + state_code=target.state_code, + state_fips=target.state_fips, + filter_variable=target.filter_variable, + ) + folder = state_dataset_relative_folder(target.dataset, target.state_code) + if (identity.store_path, identity.filename, folder) != ( + entry.path, + entry.filename, + entry.folder, + ): + msg = ( + f"planned and runtime Stage 12 state dataset identities differ " + f"for {entry.region} {year}" + ) + raise RuntimeError(msg) + datasets.add(target.dataset) + if len(datasets) > 1: + msg = f"One state dataset build covers one dataset, not {sorted(datasets)}" + raise ValueError(msg) + return {entry.region: targets[entry.region] for entry in expected} + + +def _require_districts_within_their_states(household: MicroDataFrame) -> None: + """Refuse a dataset whose district codes name a different state. + + District requests read their state's file and filter it by + ``congressional_district_geoid`` (state FIPS times 100 plus the district + number). Both columns are labels; nothing here is weighted. + """ + if "congressional_district_geoid" not in household.columns: + return + outside = (household["congressional_district_geoid"] // 100) != household[ + "state_fips" + ] + if outside.any(): + msg = ( + "Some households have a congressional district code outside their " + "state; their state's file would not match the district's results" + ) + raise RuntimeError(msg) + + +def build_state_datasets_impl( + bucket: str, + year: int, + expected: Sequence[StateDatasetPlanEntry], + *, + work_folder: Path | None = None, +) -> tuple[StateDatasetBuildResult, ...]: + """Build, check and upload one year's Stage 12 state datasets. + + Has .py prepare the regional dataset's full-country year file once, then + for each planned state keeps the rows its region's own scoping strategy + keeps, writes them with .py's ``save()``, copies the year-file stamp, + confirms that .py's reuse checks accept the file, and uploads it. + """ + from policyengine.tax_benefit_models.us.datasets import ( + PolicyEngineUSDataset, + USYearData, + ) + from policyengine.tax_benefit_models.us.legacy_inputs import ( + read_year_file_stamp, + write_year_file_stamp, + ) + from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES + + if not expected: + return () + targets = _checked_state_targets(year, expected) + dataset = expected[0].dataset + work_folder = work_folder or Path( + tempfile.mkdtemp(prefix="policyengine-stage12-state-datasets-") + ) + source_folder = str(work_folder / "source") + country_module = _country_module("us") + + # Preparing the year file runs a Microsimulation over every state, whose + # objects hold reference cycles; release them before filtering 51 copies. + country_module.ensure_datasets( + datasets=[dataset], years=[year], data_folder=source_folder + ) + gc.collect() + (source,) = country_module.load_datasets( + datasets=[dataset], years=[year], data_folder=source_folder + ).values() + stamp = read_year_file_stamp(source.filepath) + if stamp is None: + msg = f"The prepared year file {source.filepath} has no year-file stamp" + raise RuntimeError(msg) + _require_districts_within_their_states(source.data.household) + + store = Stage12CacheStore(bucket) + results = [] + for entry in expected: + rows = targets[entry.region].scoping_strategy.apply( + entity_data=source.data.entity_data, + group_entities=US_GROUP_ENTITIES, + year=year, + ) + local_file = work_folder / entry.folder / entry.filename + PolicyEngineUSDataset( + name=f"{dataset}-{entry.state_code.lower()}-year-{year}", + description=f"{entry.region} rows of {dataset} for {year}", + filepath=str(local_file), + year=year, + metadata=dict(source.metadata), + data=USYearData(**rows), + ).save() + # save() rewrites the file without the stamp, so stamp it afterwards, + # as create_datasets does. + write_year_file_stamp(local_file, stamp) + problem = year_file_reuse_problem(local_file, dataset) + if problem is not None: + msg = f"ensure_datasets would not reuse {local_file}: {problem}" + raise RuntimeError(msg) + uploaded, stored = store.upload_file(entry.path, local_file) + results.append( + StateDatasetBuildResult( + **stored.model_dump(), + year=year, + state_code=entry.state_code, + uploaded=uploaded, + ) + ) + return tuple(results) + + def _ensure_dataset_file(bucket: str, year: int) -> None: identity = collect_dataset_identity(load_stage12_bundle(), country="us", year=year) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py index 9646704af..a6ff52d42 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py @@ -23,3 +23,27 @@ def state_dataset_relative_folder(dataset: str, state_code: str) -> str: def state_dataset_folder(data_folder: str, dataset: str, state_code: str) -> str: """Return the absolute folder of a state's baked year files.""" return str(Path(data_folder) / state_dataset_relative_folder(dataset, state_code)) + + +def year_file_reuse_problem(path: str | Path, dataset: str) -> str | None: + """Return why ``ensure_datasets`` would rebuild this year file, or None. + + These are the two checks .py's ``ensure_datasets`` and ``load_datasets`` + run before reusing a year file of ``dataset``: the year-file stamp and + renames record, then the legacy SPM record against the profile this + release expects for the dataset. + """ + from policyengine.tax_benefit_models.us.datasets import ( + _expected_legacy_spm_profile, + ) + from policyengine.tax_benefit_models.us.legacy_inputs import ( + year_file_stamp_problem, + ) + from policyengine.tax_benefit_models.us.legacy_spm.record import ( + year_file_problem, + ) + + problem = year_file_stamp_problem(path) + if problem is not None: + return problem + return year_file_problem(path, _expected_legacy_spm_profile(dataset)) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py index 82a0ebbed..f3c4fa7d4 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -2,6 +2,7 @@ from __future__ import annotations +from pathlib import Path from types import SimpleNamespace import pandas as pd @@ -129,6 +130,14 @@ def upload_from_string(self, payload, *, content_type, if_generation_match): raise PreconditionFailed(self.path) self.objects[self.path] = {"payload": payload, "metadata": self.metadata} + def upload_from_filename(self, filename, *, if_generation_match): + assert if_generation_match == 0 + if self.path in self.objects: + raise PreconditionFailed(self.path) + with open(filename, "rb") as source: # noqa: PTH123 # mirrors the GCS client's own file read + payload = source.read() + self.objects[self.path] = {"payload": payload, "metadata": self.metadata} + def download_as_bytes(self) -> bytes: if self.path not in self.objects: raise NotFound(self.path) @@ -828,3 +837,192 @@ def test_planned_baseline_file_is_the_one_the_worker_saves_and_loads( # A US simulation always resolves an SPM selection, which .py names. assert entry.filename != f"{entry.simulation_id}.h5" assert plan.determinism_verified is False + + +class TestBuildStateDatasetsImpl: + """Building, checking and uploading one year's Stage 12 state datasets.""" + + YEAR = 2026 + + @pytest.fixture + def state_build(self, monkeypatch, tmp_path): + from fixtures import state_dataset_sources as sources + from policyengine.core.scoping_strategy import RowFilterStrategy + from policyengine.tax_benefit_models import us as us_module + + client = _FakeClient() + _use_store(monkeypatch, client) + state = SimpleNamespace( + client=client, + ensure_calls=[], + entity_frames=sources.source_entity_frames(), + work_folder=tmp_path / "work", + sources=sources, + ) + + def target(state_code, fips): + return precompute.StateDatasetTarget( + region=f"state/{state_code.lower()}", + state_code=state_code, + state_fips=fips, + filter_variable="state_fips", + dataset=sources.REGIONAL_DATASET, + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=fips + ), + ) + + state.targets = [ + target("UT", sources.UTAH_FIPS), + target("WY", sources.WYOMING_FIPS), + ] + monkeypatch.setattr( + precompute, "state_dataset_targets", lambda: tuple(state.targets) + ) + monkeypatch.setattr(precompute, "load_stage12_bundle", lambda: "resolved") + + def fake_state_identity(resolved, *, country, year, dataset, **state_filter): + digest = f"state-{year}-{state_filter['state_code']}" + return SimpleNamespace( + digest=digest, + store_path=f"state-datasets/{country}/{digest}/{dataset}_year_{year}.h5", + filename=f"{dataset}_year_{year}.h5", + ) + + monkeypatch.setattr( + precompute, "collect_state_dataset_identity", fake_state_identity + ) + + def fake_ensure_datasets(*, datasets, years, data_folder): + state.ensure_calls.append((datasets, years, data_folder)) + folder = Path(data_folder) + folder.mkdir(parents=True, exist_ok=True) + for year in years: + sources.write_year_file( + folder, year=year, entity_frames=state.entity_frames + ) + return us_module.load_datasets( + datasets=datasets, years=years, data_folder=data_folder + ) + + monkeypatch.setattr( + precompute, + "_country_module", + lambda _country: SimpleNamespace( + ensure_datasets=fake_ensure_datasets, + load_datasets=us_module.load_datasets, + ), + ) + state.identity_for = fake_state_identity + return state + + def _entry(self, state_build, state_code): + dataset = state_build.sources.REGIONAL_DATASET + identity = state_build.identity_for( + "resolved", + country="us", + year=self.YEAR, + dataset=dataset, + state_code=state_code, + ) + return StateDatasetPlanEntry( + year=self.YEAR, + state_code=state_code, + region=f"state/{state_code.lower()}", + dataset=dataset, + digest=identity.digest, + path=identity.store_path, + filename=identity.filename, + folder=f"state-datasets/{dataset}/{state_code}", + exists=False, + ) + + def _build(self, state_build, entries): + return precompute.build_state_datasets_impl( + "cache-bucket", self.YEAR, entries, work_folder=state_build.work_folder + ) + + def test_uploads_each_state_file_with_only_its_households(self, state_build): + entries = [self._entry(state_build, "UT"), self._entry(state_build, "WY")] + + results = self._build(state_build, entries) + + sources = state_build.sources + utah = state_build.work_folder / entries[0].folder / entries[0].filename + wyoming = state_build.work_folder / entries[1].folder / entries[1].filename + assert sources.stored_ids(utah, "household") == [1, 2] + assert sources.stored_ids(utah, "person") == [1, 2, 3] + assert sources.stored_weights(utah, "person") == [10.0, 11.0, 20.0] + assert sources.stored_ids(wyoming, "household") == [3] + assert [result.path for result in results] == [entry.path for entry in entries] + assert all(result.uploaded for result in results) + assert all(result.size_bytes > 0 for result in results) + assert state_build.client.objects[entries[0].path]["payload"] == ( + utah.read_bytes() + ) + + def test_dotpy_reuses_a_state_file_without_rebuilding( + self, state_build, monkeypatch + ): + from policyengine.tax_benefit_models.us import datasets as us_datasets + + entry = self._entry(state_build, "UT") + self._build(state_build, [entry]) + + def refuse_to_rebuild(**_kwargs): + msg = "ensure_datasets rebuilt a baked state file" + raise AssertionError(msg) + + monkeypatch.setattr(us_datasets, "create_datasets", refuse_to_rebuild) + (dataset,) = us_datasets.ensure_datasets( + datasets=[state_build.sources.REGIONAL_DATASET], + years=[self.YEAR], + data_folder=str(state_build.work_folder / entry.folder), + ).values() + + assert sorted(dataset.data.household["household_id"]) == [1, 2] + + def test_prepares_the_full_country_year_file_once(self, state_build): + self._build( + state_build, + [self._entry(state_build, "UT"), self._entry(state_build, "WY")], + ) + + assert len(state_build.ensure_calls) == 1 + datasets, years, _folder = state_build.ensure_calls[0] + assert datasets == [state_build.sources.REGIONAL_DATASET] + assert years == [self.YEAR] + + def test_refuses_a_district_that_lies_outside_its_state(self, state_build): + state_build.entity_frames = state_build.sources.source_entity_frames( + wyoming_district_geoid=4903 + ) + + with pytest.raises(RuntimeError, match="congressional district"): + self._build(state_build, [self._entry(state_build, "UT")]) + assert state_build.client.objects == {} + + def test_refuses_a_plan_entry_whose_key_differs_in_the_container(self, state_build): + entry = self._entry(state_build, "UT").model_copy( + update={"path": "state-datasets/us/other/x.h5"} + ) + + with pytest.raises(RuntimeError, match="identities differ"): + self._build(state_build, [entry]) + assert state_build.ensure_calls == [] + + def test_refuses_a_state_the_container_does_not_target(self, state_build): + state_build.targets = state_build.targets[:1] + + with pytest.raises(RuntimeError, match="state/wy"): + self._build(state_build, [self._entry(state_build, "WY")]) + assert state_build.ensure_calls == [] + + def test_refuses_entries_for_another_year(self, state_build): + with pytest.raises(ValueError, match="2027"): + precompute.build_state_datasets_impl( + "cache-bucket", + 2027, + [self._entry(state_build, "UT")], + work_folder=state_build.work_folder, + ) From 1a728e2dea26f53866170542f9b4cff36fb76219 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:52:29 +0400 Subject: [PATCH 04/25] Run the Stage 12 state dataset builds in Stage 12's precompute The Stage 12 precompute app gains build_state_datasets (8 CPUs, 96 GiB, two hours), one call per year, validating plan entries in and dumping results out. run_precompute starts the state builds with the dataset builds and waits for them only after the determinism check, before publishing the manifest, so baselines never wait for them, then reports how many were built and uploaded and their total size. The determinism record still covers baselines only. The manifest does not list the state datasets yet. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/stage12_precompute_app.py | 21 ++++ .../stage12_cache/precompute.py | 28 ++++- .../tests/test_stage12_cache_precompute.py | 105 ++++++++++++++++++ .../tests/test_stage12_precompute_app.py | 52 +++++++++ 4 files changed, 203 insertions(+), 3 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py index 414e5793b..2f0bc5160 100644 --- a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py @@ -14,10 +14,12 @@ BaselinePlanEntry, DatasetPlanEntry, PrecomputePlan, + StateDatasetPlanEntry, ) from policyengine_simulation_executor.stage12_cache.precompute import ( CachePrecomputeFunctions, build_dataset_impl, + build_state_datasets_impl, compute_baseline_impl, plan_artifacts_impl, publish_manifest_impl, @@ -104,6 +106,24 @@ def build_dataset(bucket: str, expected: dict) -> dict: ).model_dump() +@app.function( + image=precompute_image, + cpu=8.0, + # Preparing the regional dataset's full-country year file has used about + # 51 GB; the 51 filtered copies need room on top of that. + memory=98304, + timeout=7200, + secrets=worker_secrets, +) +def build_state_datasets(bucket: str, year: int, expected: list[dict]) -> list[dict]: + """Build and upload one year's planned state datasets; return their results.""" + entries = tuple(StateDatasetPlanEntry.model_validate(entry) for entry in expected) + return [ + result.model_dump() + for result in build_state_datasets_impl(bucket, year, entries) + ] + + @app.function( image=precompute_image, cpu=8.0, @@ -154,6 +174,7 @@ def main() -> None: functions=CachePrecomputeFunctions( plan_artifacts=plan_artifacts, build_dataset=build_dataset, + build_state_datasets=build_state_datasets, compute_baseline=compute_baseline, verify_determinism=verify_determinism, publish_manifest=publish_manifest, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index 623ad6442..7c49f4980 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -706,10 +706,11 @@ def publish_manifest_impl(bucket: str, plan: PrecomputePlan) -> str: @dataclass(frozen=True) class CachePrecomputeFunctions: - """The five Modal function handles the cache precompute orchestration calls.""" + """The six Modal function handles the cache precompute orchestration calls.""" plan_artifacts: modal.Function build_dataset: modal.Function + build_state_datasets: modal.Function compute_baseline: modal.Function verify_determinism: modal.Function publish_manifest: modal.Function @@ -724,8 +725,11 @@ def run_precompute( """Build the missing cache files, check determinism, publish the manifest. Missing datasets are built first, then missing baselines, in parallel - Modal calls. Prints the manifest digest through ``echo`` and returns it. - Exits with ``SystemExit`` when the determinism check fails. + Modal calls. Missing state datasets are built alongside both, one call per + year, and waited for only before the manifest is published, so the + baselines never wait for them. Prints the manifest digest through + ``echo`` and returns it. Exits with ``SystemExit`` when the determinism + check fails. """ plan = PrecomputePlan.model_validate(functions.plan_artifacts.remote(bucket)) work = select_work(plan) @@ -733,6 +737,13 @@ def run_precompute( functions.build_dataset.spawn(bucket, entry.model_dump()) for entry in work.datasets ] + state_entries_by_year: dict[int, list[dict[str, Any]]] = {} + for entry in work.state_datasets: + state_entries_by_year.setdefault(entry.year, []).append(entry.model_dump()) + state_calls = [ + functions.build_state_datasets.spawn(bucket, year, entries) + for year, entries in state_entries_by_year.items() + ] for call in dataset_calls: DatasetBuildResult.model_validate(call.get()) baseline_calls = [ @@ -756,6 +767,17 @@ def run_precompute( "Stage 12 cache determinism validation failed: " + "; ".join(verdict.differences) ) + if state_calls: + state_results = [ + StateDatasetBuildResult.model_validate(result) + for call in state_calls + for result in call.get() + ] + echo( + f"State datasets: {len(state_results)} built, " + f"{sum(result.uploaded for result in state_results)} uploaded, " + f"{sum(result.size_bytes for result in state_results)} bytes" + ) digest = functions.publish_manifest.remote(bucket, plan.model_dump()) if not isinstance(digest, str) or not digest: msg = "Stage 12 cache manifest publisher returned no digest" diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py index f3c4fa7d4..22f3c2025 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -448,6 +448,7 @@ def test_precompute_finishes_dataset_wave_before_baseline_wave() -> None: functions=precompute.CachePrecomputeFunctions( plan_artifacts=planner, build_dataset=dataset, + build_state_datasets=_Function(lambda *_: [], events, "state"), compute_baseline=baseline, verify_determinism=verify, publish_manifest=publish, @@ -470,6 +471,107 @@ def test_precompute_finishes_dataset_wave_before_baseline_wave() -> None: assert lines[-1] == "STAGE12_CACHE_MANIFEST_DIGEST=manifest-digest" +def _state_run_handles(events, plan): + def baseline_result(_bucket, entry): + return { + "path": entry["path"], + "content_sha256": "e" * 64, + "size_bytes": 1, + "year": entry["year"], + "segment_index": entry["segment_index"], + "simulation_id": entry["simulation_id"], + "cache_outcome": "miss", + "uploaded": True, + "compute_seconds": 1.0, + } + + def dataset_result(_bucket, entry): + return { + "path": entry["path"], + "content_sha256": "d" * 64, + "size_bytes": 1, + "year": entry["year"], + "uploaded": True, + "build_seconds": 1.0, + } + + def state_results(_bucket, year, entries): + return [ + { + "path": entry["path"], + "content_sha256": "c" * 64, + "size_bytes": 100, + "year": year, + "state_code": entry["state_code"], + "uploaded": entry["state_code"] != "WY", + } + for entry in entries + ] + + return { + "plan_artifacts": _Function(lambda _: plan.model_dump(), events, "plan"), + "build_dataset": _Function(dataset_result, events, "dataset"), + "build_state_datasets": _Function(state_results, events, "state"), + "compute_baseline": _Function(baseline_result, events, "baseline"), + "verify_determinism": _Function( + lambda *_: {"equal": True, "differences": []}, events, "verify" + ), + "publish_manifest": _Function(lambda *_: "manifest-digest", events, "publish"), + } + + +def test_precompute_builds_each_year_of_state_datasets_in_one_call() -> None: + events = [] + plan = _plan_with_state_datasets() + handles = _state_run_handles(events, plan) + + precompute.run_precompute( + "cache-bucket", + functions=precompute.CachePrecomputeFunctions(**handles), + echo=lambda _line: None, + ) + + assert handles["build_state_datasets"].spawn_calls == [ + ( + "cache-bucket", + year, + [entry.model_dump() for entry in plan.state_datasets if entry.year == year], + ) + for year in precompute.PRECOMPUTE_YEARS + ] + + +def test_precompute_waits_for_state_datasets_only_before_publishing() -> None: + events = [] + handles = _state_run_handles(events, _plan_with_state_datasets()) + + precompute.run_precompute( + "cache-bucket", + functions=precompute.CachePrecomputeFunctions(**handles), + echo=lambda _line: None, + ) + + first_state_get = events.index("get:state") + assert events.index("spawn:state") < events.index("spawn:baseline") + assert events.index("remote:verify") < first_state_get + assert first_state_get < events.index("remote:publish") + + +def test_precompute_reports_the_state_dataset_bytes_uploaded() -> None: + events = [] + lines = [] + handles = _state_run_handles(events, _plan_with_state_datasets()) + + precompute.run_precompute( + "cache-bucket", + functions=precompute.CachePrecomputeFunctions(**handles), + echo=lines.append, + ) + + assert "State datasets: 6 built, 3 uploaded, 600 bytes" in lines + assert lines[-1] == "STAGE12_CACHE_MANIFEST_DIGEST=manifest-digest" + + def test_warm_verified_precompute_skips_builds_and_determinism_check() -> None: events = [] plan = _plan(exists=True) @@ -485,6 +587,7 @@ def test_warm_verified_precompute_skips_builds_and_determinism_check() -> None: functions=precompute.CachePrecomputeFunctions( plan_artifacts=planner, build_dataset=unused, + build_state_datasets=unused, compute_baseline=unused, verify_determinism=verify, publish_manifest=publish, @@ -539,6 +642,7 @@ def test_failed_determinism_validation_prevents_manifest_publication() -> None: functions=precompute.CachePrecomputeFunctions( plan_artifacts=planner, build_dataset=dataset, + build_state_datasets=dataset, compute_baseline=baseline, verify_determinism=verify, publish_manifest=publish, @@ -590,6 +694,7 @@ def run() -> None: functions=precompute.CachePrecomputeFunctions( plan_artifacts=planner, build_dataset=unused, + build_state_datasets=unused, compute_baseline=unused, verify_determinism=verify, publish_manifest=publish, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py index 160be263a..2118148c9 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py @@ -21,6 +21,7 @@ def test_precompute_app_declares_independent_worker_resources(monkeypatch) -> No assert set(functions) == { "plan_artifacts", "build_dataset", + "build_state_datasets", "compute_baseline", "verify_determinism", "publish_manifest", @@ -28,6 +29,11 @@ def test_precompute_app_declares_independent_worker_resources(monkeypatch) -> No assert functions["build_dataset"]["cpu"] == 8.0 assert functions["build_dataset"]["memory"] == 65536 assert functions["build_dataset"]["timeout"] == 7200 + assert functions["build_state_datasets"]["cpu"] == 8.0 + # Preparing the regional dataset's full-country year file has used about + # 51 GB; the 51 filtered copies need room on top of that. + assert functions["build_state_datasets"]["memory"] == 98304 + assert functions["build_state_datasets"]["timeout"] == 7200 assert functions["compute_baseline"]["cpu"] == 8.0 assert functions["compute_baseline"]["memory"] == 32768 assert functions["compute_baseline"]["timeout"] == 3600 @@ -86,3 +92,49 @@ def test_precompute_app_mounts_only_the_read_only_hugging_face_secret( assert bundle_install[2]["secrets"] == [read_only_secret] for function_name, declaration in module.app.function_calls: assert declaration["secrets"] == [gcp_secret, read_only_secret], function_name + + +def test_state_dataset_wrapper_validates_entries_in_and_dumps_results_out( + monkeypatch, +) -> None: + import pytest + from policyengine_simulation_executor.stage12_cache.models import ( + StateDatasetBuildResult, + StateDatasetPlanEntry, + ) + from pydantic import ValidationError + + module = _load(monkeypatch) + entry = StateDatasetPlanEntry( + year=2026, + state_code="UT", + region="state/ut", + dataset="populace_us_2024_acs_local", + digest="s1", + path="state-datasets/us/s1/populace_us_2024_acs_local_year_2026.h5", + filename="populace_us_2024_acs_local_year_2026.h5", + folder="state-datasets/populace_us_2024_acs_local/UT", + exists=False, + ) + result = StateDatasetBuildResult( + path=entry.path, + content_sha256="c" * 64, + size_bytes=5, + year=2026, + state_code="UT", + uploaded=True, + ) + seen = {} + + def fake_build(bucket, year, expected): + seen["call"] = (bucket, year, expected) + return (result,) + + monkeypatch.setattr(module, "build_state_datasets_impl", fake_build) + + assert module.build_state_datasets("cache-bucket", 2026, [entry.model_dump()]) == [ + result.model_dump() + ] + assert seen["call"] == ("cache-bucket", 2026, (entry,)) + with pytest.raises(ValidationError): + module.build_state_datasets("cache-bucket", 2026, [{"year": 2026}]) From 7b4ce739740cd66b9210301f9335725891eedd2e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:56:37 +0400 Subject: [PATCH 05/25] Let Stage 12's image and worker checks accept state dataset files The Stage 12 cache manifest model accepts state_dataset entries, with their folder and state code, in their own model chosen by "type", so dataset and baseline entries keep their wire shape. The image fetch now checks every entry before downloading anything: it refuses unknown types, state datasets that do not cover the same states in every year, unsafe folders and file names, and two entries for one file (by folder and name, so every state's same-named file can coexist), and places each state file in its folder. The worker check expects the 63 national files plus one file per year for each state the installed bundle maps to a non-default dataset, and requires every listed state file on disk; the release it reports keeps today's fields. No manifest lists state datasets yet. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/stage12_cache_image.py | 126 +++++++++++------ .../stage12_cache/models.py | 24 +++- .../stage12_worker_validation.py | 36 ++++- .../tests/test_stage12_cache_image.py | 130 ++++++++++++++++++ .../tests/test_stage12_worker_validation.py | 108 +++++++++++++++ 5 files changed, 381 insertions(+), 43 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py b/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py index 423a14963..e6b2513f7 100644 --- a/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py +++ b/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py @@ -5,7 +5,7 @@ import json import os from hashlib import sha256 -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -20,16 +20,10 @@ def fetch_stage12_cache( if client is None: client = _storage_client() data_dir = Path(os.environ["POLICYENGINE_DATA_FOLDER"]) + destinations = _artifact_destinations(artifacts, data_dir) data_dir.mkdir(parents=True, exist_ok=True) - filenames: set[str] = set() - for artifact in artifacts: - _fetch_artifact( - artifact, - client=client, - bucket=bucket, - data_dir=data_dir, - filenames=filenames, - ) + for artifact, destination in zip(artifacts, destinations, strict=True): + _fetch_artifact(artifact, client=client, bucket=bucket, destination=destination) (data_dir / ".stage12-cache-manifest.json").write_bytes(canonical) @@ -75,23 +69,94 @@ def _require_installed_identity(manifest: dict) -> None: def _expected_artifacts(manifest: dict) -> list: - """Return the manifest's artifact list once it has 3 datasets and 60 baselines.""" + """Return the manifest's artifact list once its contents are as expected. + + It must have 3 datasets and 60 baselines, no artifact of another kind, and + either no state datasets or the same states' files for every year. + """ artifacts = manifest.get("artifacts") if not isinstance(artifacts, list): msg = "Stage 12 cache manifest has no artifact list" raise TypeError(msg) + entries = [item for item in artifacts if isinstance(item, dict)] type_counts = { - kind: sum( - item.get("type") == kind for item in artifacts if isinstance(item, dict) - ) + kind: sum(item.get("type") == kind for item in entries) for kind in ("dataset", "baseline") } if type_counts != {"dataset": 3, "baseline": 60}: msg = "Stage 12 cache manifest has an unexpected artifact count" raise RuntimeError(msg) + if any( + item.get("type") not in {"dataset", "baseline", "state_dataset"} + for item in entries + ): + msg = "Stage 12 cache manifest has an unexpected artifact type" + raise RuntimeError(msg) + states_by_year = {year: set() for year in manifest["years"]} + for item in entries: + if item.get("type") == "state_dataset": + states_by_year.setdefault(item.get("year"), set()).add( + item.get("state_code") + ) + if len({frozenset(states) for states in states_by_year.values()}) > 1: + msg = "Stage 12 cache manifest's state datasets differ between years" + raise RuntimeError(msg) return artifacts +def _relative_destination(artifact: dict) -> PurePosixPath | None: + """Return where a listed file lands, relative to the data folder, or None. + + Dataset and baseline files land directly in the data folder; a state + dataset lands in its ``folder``, a relative path without ``.`` or ``..`` + parts. None means the entry names an unsafe place. + """ + filename = artifact.get("filename") + if not isinstance(filename, str) or PurePosixPath(filename).name != filename: + return None + if "\\" in filename or filename in {"", ".", ".."}: + return None + if artifact.get("type") != "state_dataset": + return PurePosixPath(filename) + folder = artifact.get("folder") + if not isinstance(folder, str) or not folder or "\\" in folder: + return None + parts = PurePosixPath(folder).parts + if ( + PurePosixPath(folder).is_absolute() + or any(part in {".", ".."} for part in parts) + or "/".join(parts) != folder + ): + return None + return PurePosixPath(folder, filename) + + +def _artifact_destinations(artifacts: list, data_dir: Path) -> list[Path]: + """Check every listed file's metadata and return where each one lands. + + Runs before anything downloads, so a bad entry leaves no file behind. + """ + destinations = [] + seen: set[PurePosixPath] = set() + for artifact in artifacts: + if not isinstance(artifact, dict): + msg = "Stage 12 cache artifact entry is invalid" + raise TypeError(msg) + relative = _relative_destination(artifact) + if ( + relative is None + or relative in seen + or not isinstance(artifact.get("path"), str) + or not isinstance(artifact.get("content_sha256"), str) + or not isinstance(artifact.get("size_bytes"), int) + ): + msg = "Stage 12 cache artifact metadata is invalid" + raise RuntimeError(msg) + seen.add(relative) + destinations.append(data_dir.joinpath(*relative.parts)) + return destinations + + def _storage_client() -> Client: """Build a storage client from the deployment's service-account credentials. @@ -116,40 +181,23 @@ def _storage_client() -> Client: def _fetch_artifact( - artifact: object, + artifact: dict, *, client: Client, bucket: str, - data_dir: Path, - filenames: set[str], + destination: Path, ) -> None: - """Download one listed file into ``data_dir`` and verify its size and digest. + """Download one checked file to ``destination`` and verify its size and digest. The download goes to a partial file beside its destination and is moved into place only once verified, so a failed download leaves nothing. """ - if not isinstance(artifact, dict): - msg = "Stage 12 cache artifact entry is invalid" - raise TypeError(msg) - filename = artifact.get("filename") - path = artifact.get("path") - expected_digest = artifact.get("content_sha256") - expected_size = artifact.get("size_bytes") - if ( - not isinstance(filename, str) - or Path(filename).name != filename - or filename in filenames - or not isinstance(path, str) - or not isinstance(expected_digest, str) - or not isinstance(expected_size, int) - ): - msg = "Stage 12 cache artifact metadata is invalid" - raise RuntimeError(msg) - filenames.add(filename) - destination = data_dir / filename - partial = destination.with_name(filename + ".partial") + expected_digest = artifact["content_sha256"] + expected_size = artifact["size_bytes"] + destination.parent.mkdir(parents=True, exist_ok=True) + partial = destination.with_name(destination.name + ".partial") try: - client.bucket(bucket).blob(path).download_to_filename(str(partial)) + client.bucket(bucket).blob(artifact["path"]).download_to_filename(str(partial)) digest = sha256() with partial.open("rb") as stream: while chunk := stream.read(1024 * 1024): diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py index a5bf6591f..9a5ffad52 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py @@ -154,6 +154,28 @@ class ManifestArtifact(StrictCacheModel): size_bytes: int +class StateDatasetManifestArtifact(StrictCacheModel): + """One state's year file, which the worker image places in its state ``folder``. + + Its own model, so dataset and baseline entries keep their exact wire shape. + """ + + type: Literal["state_dataset"] + path: str + filename: str + folder: str + state_code: str + year: int + identity_digest: str + content_sha256: str + size_bytes: int + + +CacheManifestEntry = Annotated[ + ManifestArtifact | StateDatasetManifestArtifact, Field(discriminator="type") +] + + class CacheManifest(StrictCacheModel): """The published list of cache files for one bundle and region partition.""" @@ -166,7 +188,7 @@ class CacheManifest(StrictCacheModel): years: tuple[int, ...] partition_sha256: str bundle: CacheBundleIdentity - artifacts: tuple[ManifestArtifact, ...] + artifacts: tuple[CacheManifestEntry, ...] def canonical_payload(self) -> dict[str, Any]: """Return the JSON payload that is hashed and stored, keyed by ``schema``.""" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index d861c427d..992d1d049 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -18,6 +18,9 @@ resolve_data_folder, ) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle +from policyengine_simulation_executor.stage12_cache import ( + precompute as cache_precompute, +) from policyengine_simulation_executor.stage12_cache.models import CacheManifest from policyengine_simulation_executor.stage12_cache.precompute import ( PRECOMPUTE_YEARS, @@ -147,6 +150,31 @@ def _check_uk_local_authority_dataset(dataset_path: str) -> None: STAGE12_CACHE_PROBE_YEAR = PRECOMPUTE_YEARS[0] +def expected_state_dataset_count() -> int: + """Return how many state dataset files the installed bundle's cache holds. + + One per precomputed year for each US state region on a non-default + dataset; none when every state uses the default dataset. + """ + return len(PRECOMPUTE_YEARS) * len(cache_precompute.state_dataset_targets()) + + +def _require_installed_artifacts( + manifest: CacheManifest, data_folder: Path, *, state_dataset_count: int +) -> None: + """Require the expected number of listed files and every state file on disk.""" + if len(manifest.artifacts) != STAGE12_CACHE_ARTIFACT_COUNT + state_dataset_count: + msg = "installed Stage 12 cache artifact count differs" + raise RuntimeError(msg) + for item in manifest.artifacts: + if ( + item.type == "state_dataset" + and not (data_folder / item.folder / item.filename).is_file() + ): + msg = f"installed Stage 12 state dataset file is missing: {item.path}" + raise RuntimeError(msg) + + def _check_installed_cache(expected_manifest_sha256: str) -> dict[str, Any]: """Validate the baked manifest and require one real baseline cache hit.""" data_folder = Path(resolve_data_folder()) @@ -159,9 +187,11 @@ def _check_installed_cache(expected_manifest_sha256: str) -> dict[str, Any]: if manifest.years != PRECOMPUTE_YEARS: msg = "installed Stage 12 cache years differ" raise RuntimeError(msg) - if len(manifest.artifacts) != STAGE12_CACHE_ARTIFACT_COUNT: - msg = "installed Stage 12 cache artifact count differs" - raise RuntimeError(msg) + _require_installed_artifacts( + manifest, + data_folder, + state_dataset_count=expected_state_dataset_count(), + ) probe = next( item for item in manifest.artifacts diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py index d737a9d54..388e54ebd 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py @@ -308,3 +308,133 @@ def build_client(**arguments): {"credentials": "service-account-credentials", "project": "cache-project"} ] assert len(list(tmp_path.glob("*.h5"))) == 63 + + +REGIONAL_DATASET = "populace_us_2024_acs_local" + + +def _state_artifact(year: int, state_code: str, **changes) -> tuple[dict, bytes]: + filename = f"{REGIONAL_DATASET}_year_{year}.h5" + payload = f"state-{year}-{state_code}".encode() + artifact = { + "type": "state_dataset", + "path": f"state-datasets/us/{year}-{state_code}/{filename}", + "filename": filename, + "folder": f"state-datasets/{REGIONAL_DATASET}/{state_code}", + "state_code": state_code, + "year": year, + "identity_digest": f"identity-state-{year}-{state_code}", + "content_sha256": sha256(payload).hexdigest(), + "size_bytes": len(payload), + **changes, + } + return artifact, payload + + +def _manifest_with_state_datasets(states=("UT", "WY")) -> tuple[dict, dict]: + manifest, payloads = _manifest() + for year in (2026, 2027, 2025): + for state_code in states: + artifact, payload = _state_artifact(year, state_code) + manifest["artifacts"].append(artifact) + payloads[artifact["path"]] = payload + return manifest, payloads + + +def test_places_each_state_file_in_its_state_folder(monkeypatch, tmp_path) -> None: + manifest, payloads = _manifest_with_state_datasets() + _use_deployment(monkeypatch, tmp_path, manifest) + + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + state_folder = tmp_path / "state-datasets" / REGIONAL_DATASET + for year in (2026, 2027, 2025): + for state_code in ("UT", "WY"): + stored = state_folder / state_code / f"{REGIONAL_DATASET}_year_{year}.h5" + assert stored.read_bytes() == f"state-{year}-{state_code}".encode() + assert len(list(tmp_path.glob("*.h5"))) == 63 + + +def _drop_one_state_file(manifest: dict) -> None: + manifest["artifacts"] = [ + item + for item in manifest["artifacts"] + if not (item["type"] == "state_dataset" and item["year"] == 2027) + or item["state_code"] != "WY" + ] + + +def _add_unknown_type(manifest: dict) -> None: + manifest["artifacts"].append({**manifest["artifacts"][0], "type": "mystery"}) + + +@pytest.mark.parametrize( + ("change", "message"), + [ + pytest.param( + _drop_one_state_file, + "^Stage 12 cache manifest's state datasets differ between years$", + id="state-missing-in-one-year", + ), + pytest.param( + _add_unknown_type, + "^Stage 12 cache manifest has an unexpected artifact type$", + id="unknown-type", + ), + ], +) +def test_rejects_inconsistent_state_datasets_before_downloading( + monkeypatch, tmp_path, change, message +) -> None: + manifest, payloads = _manifest_with_state_datasets() + change(manifest) + _use_deployment(monkeypatch, tmp_path, manifest) + + with pytest.raises(RuntimeError, match=message): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + assert list(tmp_path.iterdir()) == [] + + +@pytest.mark.parametrize( + "folder", + [ + "/absolute/UT", + "state-datasets/../../escape", + "state-datasets/./UT", + "state-datasets//UT", + "state-datasets\\UT", + "", + None, + ], +) +def test_rejects_an_unsafe_state_folder_before_downloading( + monkeypatch, tmp_path, folder +) -> None: + manifest, payloads = _manifest_with_state_datasets() + manifest["artifacts"][-1]["folder"] = folder + _use_deployment(monkeypatch, tmp_path, manifest) + + with pytest.raises( + RuntimeError, match=r"^Stage 12 cache artifact metadata is invalid$" + ): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + assert list(tmp_path.iterdir()) == [] + + +def test_rejects_two_entries_for_one_state_file_before_downloading( + monkeypatch, tmp_path +) -> None: + manifest, payloads = _manifest_with_state_datasets() + repeated = dict(manifest["artifacts"][-1]) + manifest["artifacts"].append({**repeated, "path": "state-datasets/us/other/x.h5"}) + payloads["state-datasets/us/other/x.h5"] = payloads[repeated["path"]] + _use_deployment(monkeypatch, tmp_path, manifest) + + with pytest.raises( + RuntimeError, match=r"^Stage 12 cache artifact metadata is invalid$" + ): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + assert list(tmp_path.iterdir()) == [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index 4e02df25e..fc104c9df 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -411,3 +411,111 @@ def test_digest_mismatch_is_checked_before_any_dataset_access() -> None: ) assert events == [] + + +def _installed_manifest(*, state_codes=("UT", "WY")): + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + from policyengine_simulation_executor.stage12_cache.precompute import ( + PRECOMPUTE_YEARS, + ) + from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, + ) + + def entry(kind, year, index): + return { + "type": kind, + "path": f"{kind}s/us/{year}-{index}/{kind}-{year}-{index}.h5", + "filename": f"{kind}-{year}-{index}.h5", + "year": year, + "identity_digest": f"{kind}-{year}-{index}", + "content_sha256": "f" * 64, + "size_bytes": 1, + } + + artifacts = [entry("dataset", year, 0) for year in PRECOMPUTE_YEARS] + artifacts += [ + entry("baseline", year, index) + for year in PRECOMPUTE_YEARS + for index in range(len(US_REGION_GROUPS)) + ] + artifacts += [ + { + "type": "state_dataset", + "path": f"state-datasets/us/{year}-{state}/acs_year_{year}.h5", + "filename": f"acs_year_{year}.h5", + "folder": f"state-datasets/acs/{state}", + "state_code": state, + "year": year, + "identity_digest": f"state-{year}-{state}", + "content_sha256": "f" * 64, + "size_bytes": 1, + } + for year in PRECOMPUTE_YEARS + for state in state_codes + ] + return CacheManifest.model_validate( + { + "schema": "stage12-mf1", + "country": "us", + "years": list(PRECOMPUTE_YEARS), + "partition_sha256": "p" * 64, + "bundle": { + "bundle_manifest_sha256": "b" * 64, + "policyengine_version": "5.2.0", + "core_package_version": "3.0.0", + "country_package_version": "1.0.0", + "data_package_version": "2.0.0", + "data_artifact_revision": "revision", + "default_dataset": "populace_us_2024", + }, + "artifacts": artifacts, + } + ) + + +def _install_state_files(data_folder, manifest): + for item in manifest.artifacts: + if item.type == "state_dataset": + folder = data_folder / item.folder + folder.mkdir(parents=True, exist_ok=True) + (folder / item.filename).write_bytes(b"h5") + + +def test_installed_cache_accepts_the_expected_state_files(tmp_path) -> None: + from policyengine_simulation_executor import stage12_worker_validation as module + + manifest = _installed_manifest() + _install_state_files(tmp_path, manifest) + + module._require_installed_artifacts(manifest, tmp_path, state_dataset_count=6) + + +def test_installed_cache_refuses_a_state_file_count_the_bundle_does_not_give( + tmp_path, +) -> None: + from policyengine_simulation_executor import stage12_worker_validation as module + + manifest = _installed_manifest(state_codes=("UT",)) + _install_state_files(tmp_path, manifest) + + with pytest.raises(RuntimeError, match="artifact count differs"): + module._require_installed_artifacts(manifest, tmp_path, state_dataset_count=6) + + +def test_installed_cache_refuses_a_missing_state_file(tmp_path) -> None: + from policyengine_simulation_executor import stage12_worker_validation as module + + manifest = _installed_manifest() + + with pytest.raises(RuntimeError, match="state dataset file is missing"): + module._require_installed_artifacts(manifest, tmp_path, state_dataset_count=6) + + +def test_expected_state_file_count_follows_the_installed_bundle(monkeypatch) -> None: + from policyengine_simulation_executor import stage12_worker_validation as module + from policyengine_simulation_executor.stage12_cache import precompute + + monkeypatch.setattr(precompute, "state_dataset_targets", lambda: ("UT", "WY")) + + assert module.expected_state_dataset_count() == 2 * len(precompute.PRECOMPUTE_YEARS) From 61e622ae0aba031749e3e9fedf8f752fe37a5a15 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:59:54 +0400 Subject: [PATCH 06/25] List the state datasets in Stage 12's cache manifest build_manifest appends a state_dataset entry, with its folder, state code and the stored digest and size, for every planned state file after the datasets and baselines. A plan without state datasets publishes the manifest it always did. The Stage 12 image fetch places each state file in its folder, and the worker check expects them. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 10 ++- .../stage12_cache/precompute.py | 25 ++++++- .../tests/test_stage12_cache_precompute.py | 65 +++++++++++++++++++ 3 files changed, 95 insertions(+), 5 deletions(-) diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 910c20a1b..37f5e91cb 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -205,8 +205,10 @@ redacted cause (type, stage, simulation role and redacted message) followed by the same correlation ID. Stage 12 deployments run `src/modal/stage12_precompute_app.py` against the -dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets plus 60 -current-law US baseline files and emits +dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets, 60 +current-law US baseline files and, for each state region whose requests load +a dataset other than the bundle default, that dataset's year file filtered to +the state's households (one per state and year), and emits `STAGE12_CACHE_MANIFEST_DIGEST=`. Baseline files have the names .py saves and loads them under (`Simulation.storage_id`, which includes the SPM selection's digest). Stored files are never replaced: a run computes only @@ -217,7 +219,9 @@ workers, its image and functions mount only the read-only `HUGGING_FACE_TOKEN`, and `stage12-evaluation-gcp-credentials`; the US data it reads comes from the public `policyengine/populace-us` dataset repository. The US worker image -downloads and verifies the manifest's 63 files during image construction. The +downloads and verifies the manifest's files during image construction: the 63 +national files into the data folder, and each state dataset into +`state-datasets///` under the year file's own name. The release validation checks the recorded manifest, hashes a baseline file, and requires a real calculation to load that file without recomputing it. Worker validation also calculates a US household under a nonempty CTC reform using diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index 7c49f4980..ac591dc49 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -59,12 +59,14 @@ BaselinePlanEntry, CacheBundleIdentity, CacheManifest, + CacheManifestEntry, DatasetBuildResult, DatasetPlanEntry, DeterminismVerdict, ManifestArtifact, PrecomputePlan, StateDatasetBuildResult, + StateDatasetManifestArtifact, StateDatasetPlanEntry, WorkSelection, ) @@ -653,8 +655,12 @@ def verify_determinism_impl( def build_manifest(plan: PrecomputePlan, store: Stage12CacheStore) -> CacheManifest: - """List every planned file with the digest and size the bucket reports.""" - artifacts = [] + """List every planned file with the digest and size the bucket reports. + + Datasets, then baselines, then state datasets, each in plan order, so a + plan without state datasets keeps the manifest it always had. + """ + artifacts: list[CacheManifestEntry] = [] for entry in plan.datasets: stored = store.describe(entry.path) artifacts.append( @@ -681,6 +687,21 @@ def build_manifest(plan: PrecomputePlan, store: Stage12CacheStore) -> CacheManif size_bytes=stored.size_bytes, ) ) + for entry in plan.state_datasets: + stored = store.describe(entry.path) + artifacts.append( + StateDatasetManifestArtifact( + type="state_dataset", + path=entry.path, + filename=entry.filename, + folder=entry.folder, + state_code=entry.state_code, + year=entry.year, + identity_digest=entry.digest, + content_sha256=stored.content_sha256, + size_bytes=stored.size_bytes, + ) + ) return CacheManifest( years=plan.years, partition_sha256=plan.partition_sha256, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py index 22f3c2025..978bd8c00 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -669,6 +669,71 @@ def test_manifest_records_content_hashes_and_sizes() -> None: assert all(item.size_bytes == 123 for item in manifest.artifacts) +def test_manifest_lists_state_datasets_after_the_others_in_plan_order() -> None: + plan = _plan_with_state_datasets(exists=True) + store = SimpleNamespace( + describe=lambda path: CacheObject( + path=path, content_sha256="f" * 64, size_bytes=123 + ) + ) + + manifest = precompute.build_manifest(plan, store) + + assert len(manifest.artifacts) == 63 + len(plan.state_datasets) + state_artifacts = [ + artifact.model_dump() + for artifact in manifest.artifacts + if artifact.type == "state_dataset" + ] + assert state_artifacts == [ + { + "type": "state_dataset", + "path": entry.path, + "filename": entry.filename, + "folder": entry.folder, + "state_code": entry.state_code, + "year": entry.year, + "identity_digest": entry.digest, + "content_sha256": "f" * 64, + "size_bytes": 123, + } + for entry in plan.state_datasets + ] + assert [artifact.type for artifact in manifest.artifacts[:63]] == ( + ["dataset"] * 3 + ["baseline"] * 60 + ) + + +def test_manifest_without_state_datasets_keeps_its_published_shape() -> None: + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + + store = SimpleNamespace( + describe=lambda path: CacheObject( + path=path, content_sha256="f" * 64, size_bytes=123 + ) + ) + payload = precompute.build_manifest(_plan(exists=True), store).canonical_payload() + + assert all(set(item) != {"folder"} for item in payload["artifacts"]) + assert all("folder" not in item for item in payload["artifacts"]) + assert CacheManifest.model_validate(payload).canonical_payload() == payload + + +def test_manifest_with_state_datasets_round_trips_its_wire_shape() -> None: + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + + store = SimpleNamespace( + describe=lambda path: CacheObject( + path=path, content_sha256="f" * 64, size_bytes=123 + ) + ) + payload = precompute.build_manifest( + _plan_with_state_datasets(exists=True), store + ).canonical_payload() + + assert CacheManifest.model_validate(payload).canonical_payload() == payload + + @pytest.mark.parametrize("passes", [True, False]) def test_a_run_with_no_new_baselines_still_checks_unverified_ones( *, From 5537991b5c599fe48bd311807850d0a081438a01 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 21:04:32 +0400 Subject: [PATCH 07/25] Read a state's baked dataset in Stage 12 state and district calculations Stage 12's worker asks its own copy of the baked-folder lookup, in stage12_cache/state_datasets.py, for the Stage 12 image folder of the calculation's state before loading a non-default dataset, and passes it to _load_dataset. The lookup returns the folder only when the image holds that state's year file for the year and .py's reuse checks accept it, so ensure_datasets never rebuilds the regional dataset inside a baked folder; otherwise the calculation prepares the dataset as before. District calculations read their state's file and keep their district filter. National, region-group and UK calculations are unchanged. The selection-contract test now uses real RegionResolution and DatasetSelection values. Co-Authored-By: Claude Opus 5.5 --- .../stage12_cache/state_datasets.py | 43 ++++++ .../stage12_runtime/simulation.py | 11 ++ .../tests/test_stage12_runtime.py | 123 +++++++++++++++++- .../tests/test_stage12_state_datasets.py | 68 ++++++++++ 4 files changed, 242 insertions(+), 3 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py index a6ff52d42..98bc21010 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py @@ -10,7 +10,14 @@ from __future__ import annotations +import logging from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + +logger = logging.getLogger(__name__) STATE_DATASETS_DIRECTORY = "state-datasets" @@ -47,3 +54,39 @@ def year_file_reuse_problem(path: str | Path, dataset: str) -> str | None: if problem is not None: return problem return year_file_problem(path, _expected_legacy_spm_profile(dataset)) + + +def baked_state_data_folder( + *, + country: str, + selection: DatasetSelection, + state_code: str | None, + year: int, + data_folder: str, +) -> str | None: + """Return the baked folder a US state or district calculation should load. + + A calculation qualifies when it loads a non-default dataset for a US + state (a district passes its state's code) and the image holds that + state's year file for ``year``, which must also pass .py's reuse checks: + otherwise ``ensure_datasets`` would rebuild the full regional dataset + inside the baked folder. Returns None in every other case, and the + calculation prepares the dataset as before. + """ + if country.lower() != "us" or not state_code or selection.is_default: + return None + from policyengine.provenance.manifest import ( + dataset_logical_name, + resolve_dataset_reference, + ) + + folder = state_dataset_folder(data_folder, selection.name, state_code) + stem = dataset_logical_name(resolve_dataset_reference("us", selection.name)) + year_file = Path(folder) / f"{stem}_year_{year}.h5" + if not year_file.is_file(): + return None + problem = year_file_reuse_problem(year_file, selection.name) + if problem is not None: + logger.warning("Not reading the baked state file %s: %s", year_file, problem) + return None + return folder diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index b1a5695c9..09c67bb42 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -236,8 +236,12 @@ def calculate_simulation_frames( _require_certified_uk_weight_matrix, _resolve_dataset_selection, _resolve_region, + resolve_data_folder, setup_gcp_credentials, ) + from policyengine_simulation_executor.stage12_cache import ( # noqa: PLC0415 # resolved per call so tests can patch the baked-folder lookup + state_datasets, + ) with ( stage_span(Stage.CREDENTIAL_SETUP), @@ -261,6 +265,13 @@ def calculate_simulation_frames( params, selection=dataset_selection, country_module=country_module, + data_folder=state_datasets.baked_state_data_folder( + country=country, + selection=dataset_selection, + state_code=region.state_code, + year=simulation.year, + data_folder=resolve_data_folder(), + ), ) uk_local_authority_metadata = detect_uk_local_authority_metadata( country, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index e91f4d759..8019886a4 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -539,11 +539,19 @@ def _frames(value=100.0): def test_calculator_uses_the_current_dataset_selection_contract(monkeypatch) -> None: from policyengine_simulation_executor import simulation_runtime + from policyengine_simulation_executor.simulation_runtime import ( + DatasetSelection, + RegionResolution, + ) from policyengine_simulation_executor.stage12_runtime import simulation as worker country_module = object() - region = type("Region", (), {"code": "us", "scoping_strategy": None})() - selection = object() + region = RegionResolution(code="us") + selection = DatasetSelection( + name="populace_us_2024", + uri="hf://policyengine/populace-us/populace_us_2024.h5@revision", + is_default=True, + ) dataset = object() received: dict[str, object] = {} @@ -587,9 +595,10 @@ def resolve_dataset(params, *, region_resolution): received["region"] = region_resolution return selection - def load_dataset(_params, *, selection, country_module): + def load_dataset(_params, *, selection, country_module, data_folder=None): received["load_selection"] = selection received["load_country_module"] = country_module + received["load_data_folder"] = data_folder return dataset def build_simulation(_params, *, resolved, **_kwargs): @@ -614,6 +623,7 @@ def build_simulation(_params, *, resolved, **_kwargs): assert received["region"] is region assert received["load_selection"] is selection assert received["load_country_module"] is country_module + assert received["load_data_folder"] is None assert received["build_dataset"] is dataset assert received["build_selection"] is selection assert received["ensured"] is True @@ -621,6 +631,113 @@ def build_simulation(_params, *, resolved, **_kwargs): assert set(result.frames) == {"household", "person"} +class _StoppedAfterLoadingError(Exception): + """Raised by a stubbed simulation build once the dataset has loaded.""" + + +def _stage12_load(monkeypatch, *, region, selection, baked_folder): + """Run a Stage 12 calculation up to its dataset load; return what it did.""" + from policyengine_simulation_executor import simulation_runtime + from policyengine_simulation_executor.stage12_cache import state_datasets + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + loads = [] + lookups = [] + + def load_dataset(_params, *, selection, country_module, data_folder=None): + loads.append({"selection": selection, "data_folder": data_folder}) + return object() + + def lookup(**kwargs): + lookups.append(kwargs) + return baked_folder + + def stop_after_loading(*_args, **_kwargs): + raise _StoppedAfterLoadingError + + monkeypatch.setattr(worker, "_require_installed_bundle", lambda _: None) + monkeypatch.setattr(simulation_runtime, "setup_gcp_credentials", nullcontext) + monkeypatch.setattr( + simulation_runtime, "_country_module", lambda _country: object() + ) + monkeypatch.setattr(simulation_runtime, "_resolve_region", lambda **_: region) + monkeypatch.setattr( + simulation_runtime, + "_resolve_dataset_selection", + lambda _params, *, region_resolution: selection, + ) + monkeypatch.setattr(simulation_runtime, "_load_dataset", load_dataset) + monkeypatch.setattr(state_datasets, "baked_state_data_folder", lookup) + monkeypatch.setattr( + worker, "detect_uk_local_authority_metadata", lambda *_a, **_k: None + ) + monkeypatch.setattr( + simulation_runtime, "_require_certified_uk_weight_matrix", lambda *_: None + ) + monkeypatch.setattr(worker, "build_stage12_simulation", stop_after_loading) + + planned = _planned_simulation(SimulationRole.BASELINE) + with pytest.raises(_StoppedAfterLoadingError): + worker.calculate_simulation_frames(planned) + return SimpleNamespace(loads=loads, lookups=lookups, year=planned.year) + + +def _regional_selection(): + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + + return DatasetSelection( + name="populace_us_2024_acs_local", + uri="hf://policyengine/populace-us/populace_us_2024_acs_local.h5@rev", + is_default=False, + ) + + +@pytest.mark.parametrize("region_code", ["state/ut", "congressional_district/UT-01"]) +def test_stage12_state_and_district_runs_read_their_states_baked_folder( + monkeypatch, region_code +) -> None: + from policyengine_simulation_executor.simulation_runtime import ( + RegionResolution, + resolve_data_folder, + ) + + loaded = _stage12_load( + monkeypatch, + region=RegionResolution(code=region_code, state_code="UT"), + selection=_regional_selection(), + baked_folder="/opt/policyengine/stage12-data/state-datasets/acs/UT", + ) + + assert loaded.loads == [ + { + "selection": _regional_selection(), + "data_folder": "/opt/policyengine/stage12-data/state-datasets/acs/UT", + } + ] + assert loaded.lookups == [ + { + "country": "us", + "selection": _regional_selection(), + "state_code": "UT", + "year": loaded.year, + "data_folder": resolve_data_folder(), + } + ] + + +def test_stage12_state_run_without_a_baked_file_loads_on_demand(monkeypatch) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + loaded = _stage12_load( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_regional_selection(), + baked_folder=None, + ) + + assert loaded.loads[0]["data_folder"] is None + + def test_uk_boundary_version_is_detected_from_the_unscoped_dataset() -> None: dataset = type( "Dataset", diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py b/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py index 4432b205a..ba5b26c96 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py @@ -7,6 +7,7 @@ from __future__ import annotations import dataclasses +from pathlib import Path import pytest from policyengine_simulation_contract.stage12_bundle import ( @@ -156,3 +157,70 @@ def test_state_folder_sits_under_the_stage12_data_folder(): assert state_datasets.state_dataset_folder( "/opt/policyengine/stage12-data", REGIONAL_DATASET, "UT" ) == (f"/opt/policyengine/stage12-data/state-datasets/{REGIONAL_DATASET}/UT") + + +def _regional_selection(): + from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + ) + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + + bundle = get_country_release_bundle("us") + return DatasetSelection( + name=REGIONAL_DATASET, + uri=bundle.dataset_uris[REGIONAL_DATASET], + is_default=False, + ) + + +def _bake_utah(data_folder: Path, *, stamp: bool = True) -> Path: + from fixtures.state_dataset_sources import write_year_file + + folder = Path( + state_datasets.state_dataset_folder(str(data_folder), REGIONAL_DATASET, "UT") + ) + folder.mkdir(parents=True) + write_year_file(folder, year=2026, stamp=stamp) + return folder + + +def _lookup(data_folder, selection, *, state_code="UT", year=2026, country="us"): + return state_datasets.baked_state_data_folder( + country=country, + selection=selection, + state_code=state_code, + year=year, + data_folder=str(data_folder), + ) + + +def test_a_state_reads_its_baked_folder(tmp_path): + folder = _bake_utah(tmp_path) + + assert _lookup(tmp_path, _regional_selection()) == str(folder) + + +def test_a_year_or_state_without_a_baked_file_reads_nothing_baked(tmp_path): + _bake_utah(tmp_path) + + assert _lookup(tmp_path, _regional_selection(), year=2028) is None + assert _lookup(tmp_path, _regional_selection(), state_code="WY") is None + + +def test_an_unstamped_file_reads_nothing_baked(tmp_path): + _bake_utah(tmp_path, stamp=False) + + assert _lookup(tmp_path, _regional_selection()) is None + + +def test_default_datasets_and_other_regions_read_nothing_baked(tmp_path): + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + + _bake_utah(tmp_path) + default = DatasetSelection( + name="populace_us_2024", uri="hf://default", is_default=True + ) + + assert _lookup(tmp_path, default) is None + assert _lookup(tmp_path, _regional_selection(), state_code=None) is None + assert _lookup(tmp_path, _regional_selection(), country="uk") is None From 45f20a8565a978cc6e8f6a2aaa5aa573a89f7b5b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 21:09:29 +0400 Subject: [PATCH 08/25] Report the loaded dataset's revision as a US state run's data_version Main-path US state and district runs load a non-default dataset (the ACS-local dataset since .py 6.7), but their results named the default dataset's release as data_version, because the runtime passed no resolved data version and the year file carries none. The runtime now passes the revision the .py release manifest gives that dataset (populace-us-2024-buildo-acs-local-... for 6.7.5). National runs and every run on the default dataset report the release's data version as before. Stage 12 takes data_version from its dataset provenance; once state runs record ACS-local there, both paths report the same value for them. Co-Authored-By: Claude Opus 5.5 --- .../simulation_runtime.py | 23 ++++- .../tests/test_simulation_output_builder.py | 92 +++++++++++++++++++ 2 files changed, 114 insertions(+), 1 deletion(-) 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 154f13002..f21114f4b 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 @@ -722,6 +722,27 @@ def _resolve_dataset_selection( raise ValueError(msg) +def _selected_dataset_revision(country: str, selection: DatasetSelection) -> str | None: + """Return the release revision of the dataset a US state or district run loads. + + Such runs load a non-default dataset (the ACS-local dataset since .py 6.7), + so their ``data_version`` names that dataset's revision rather than the + default dataset's release. None leaves every other run reporting the + release's data version, as before. + """ + if country != "us" or selection.is_default: + return None + from policyengine.provenance.manifest import ( + get_release_manifest, + ) + + reference = getattr(get_release_manifest(country), "datasets", {}).get( + selection.name + ) + revision = getattr(reference, "revision", None) + return revision if isinstance(revision, str) and revision else None + + def _load_dataset( params: dict[str, Any], *, @@ -948,7 +969,7 @@ def _run_simulation_impl_core( dataset=dataset, baseline=baseline, reform=reform, - resolved_data_version=None, + resolved_data_version=_selected_dataset_revision(country, dataset_selection), resolved_region_code=region_resolution.code, runtime=runtime, uk_local_authority_metadata=uk_local_authority_metadata, diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index 162c7620a..f4abd82e8 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1368,6 +1368,98 @@ def test_runtime_loads_a_national_run_from_the_default_folder(monkeypatch): assert "dataset_source" not in loaded.context +def _builder_data_version(monkeypatch, *, region_resolution, dataset_selection): + """Run the annual runtime and return the data version it gives the builder.""" + from policyengine_simulation_executor import state_datasets + + selection = installed_spm_selection() + simulations = [spm_capable_simulation(selection), spm_capable_simulation(selection)] + builder_calls = [] + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + builder_calls.append(kwargs) + + def serialize(self): + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._country_module", + lambda _country: SimpleNamespace(model=SimpleNamespace(version="2.2.1")), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._resolve_region", + lambda **_kwargs: region_resolution, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._resolve_dataset_selection", + lambda _params, *, region_resolution: dataset_selection, + ) + monkeypatch.setattr( + state_datasets, "baked_state_data_folder", lambda **_kwargs: None + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + lambda _params, country_module, selection, data_folder=None: object(), + ) + build_iter = iter(simulations) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + lambda *_args, **_kwargs: next(build_iter), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", + FakeSimulationOutputBuilder, + ) + + _run_simulation_impl_core( + { + "country": "us", + "region": region_resolution.code, + "baseline": {}, + "reform": {}, + }, + runtime=_TrackingRuntime(), + ) + return builder_calls[0]["resolved_data_version"] + + +def test_a_state_run_reports_the_revision_of_the_dataset_it_loads(monkeypatch): + from policyengine.provenance.manifest import get_release_manifest + + bundle = get_country_release_bundle("us") + regional = "populace_us_2024_acs_local" + + data_version = _builder_data_version( + monkeypatch, + region_resolution=RegionResolution( + code="state/ut", dataset_reference="ref", state_code="UT" + ), + dataset_selection=DatasetSelection( + name=regional, uri=bundle.dataset_uris[regional], is_default=False + ), + ) + + assert data_version == get_release_manifest("us").datasets[regional].revision + assert data_version != bundle.data_version + + +def test_a_national_run_keeps_the_release_data_version(monkeypatch): + bundle = get_country_release_bundle("us") + + data_version = _builder_data_version( + monkeypatch, + region_resolution=RegionResolution(code="us"), + dataset_selection=DatasetSelection( + name=bundle.default_dataset, + uri=bundle.default_dataset_uri, + is_default=True, + ), + ) + + assert data_version is None + + def test_run_simulation_impl_core_passes_detected_uk_boundary_to_builder(monkeypatch): dataset = SimpleNamespace( data=SimpleNamespace( From a81f6e078d0cd3e8263220d6269dee29e1f55edc Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 21:15:01 +0400 Subject: [PATCH 09/25] Publish which dataset each region type loads in the Stage 12 bundle Stage12CountryBundle gains region_datasets: for each region type in .py's region_datasets, the identity of the bundle dataset its path names, or None for a per-region file pattern such as states/{state_code}.h5; a path that names no bundle dataset is refused. Bundles published before the field have None. On the installed 6.7.5 bundle, state and congressional_district regions load populace_us_2024_acs_local and national regions populace_us_2024. Entry services already running forbid unknown fields in the manifest they read, so the v2 version manifest now goes under three keys: active-with-region-datasets holds it all, active-with-cache keeps its current shape without region datasets, and active keeps its shape without cache metadata either. Loaders read the newest key a publish has written, and a first publish keeps the versions stored under the older keys. Adding the field changes the Stage 12 bundle digest, which keys every Stage 12 cache file, so the first deploy rebuilds the Stage 12 cache. Co-Authored-By: Claude Opus 5.5 --- .../stage12_bundle.py | 25 ++++- .../stage12_manifest.py | 51 +++++++-- .../tests/test_stage12_manifest.py | 102 +++++++++++++++++- .../stage12_bundle/parsing.py | 47 ++++++++ .../tests/test_stage12_bundle.py | 38 +++++++ 5 files changed, 252 insertions(+), 11 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py index 5c406f742..4d2b48db1 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py @@ -31,8 +31,23 @@ class Stage12Dataset(StrictBundleModel): repo_type: NonEmptyText +class Stage12RegionDataset(StrictBundleModel): + """The dataset one type of region loads, such as ``state`` or ``national``. + + ``dataset`` is a dataset identity in the country's bundle, or None when the + bundle gives a per-region file pattern instead of one dataset. + """ + + region_type: NonEmptyText + dataset: NonEmptyText | None + + class Stage12CountryBundle(StrictBundleModel): - """One country's pinned model package, data release and datasets within a bundle.""" + """One country's pinned model package, data release and datasets within a bundle. + + ``region_datasets`` names the dataset each type of region loads. It is None + in bundles published before it existed. + """ country: CountryId country_package_name: NonEmptyText @@ -45,6 +60,14 @@ class Stage12CountryBundle(StrictBundleModel): default_dataset: NonEmptyText default_dataset_uri: NonEmptyText datasets: tuple[Stage12Dataset, ...] + region_datasets: tuple[Stage12RegionDataset, ...] | None = None + + def dataset_for_region_type(self, region_type: str) -> str | None: + """Return the dataset identity ``region_type`` regions load, or None.""" + for item in self.region_datasets or (): + if item.region_type == region_type: + return item.dataset + return None class Stage12BundleManifest(StrictBundleModel): diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index 75bd43790..a924a24ba 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -21,9 +21,14 @@ V1_ROUTING_STATE_NAME = "simulation-api-routing-state" V2_VERSION_MANIFEST_NAME = "simulation-api-v2-version-manifest" +# The complete manifest, including each country bundle's region datasets. +REGION_DATASETS_MANIFEST_KEY = "active-with-region-datasets" +# Readers released before region datasets forbid unknown fields, so they keep +# reading this key, which holds the same manifest without them. ACTIVE_MANIFEST_KEY = "active-with-cache" # Readers released before Stage 12 cache releases forbid unknown fields, so -# they keep reading this key, which holds the same manifest without them. +# they keep reading this key, which holds the same manifest without them or +# region datasets. PRE_CACHE_MANIFEST_KEY = "active" V2_APPLICATION_PREFIX = "policyengine-simulation-v2-py" V1_APPLICATION_PREFIX = "policyengine-simulation-py" @@ -219,20 +224,49 @@ def _is_newer(candidate: str, current: str | None) -> bool: return False +_REGION_DATASETS_FIELD = { + "bundle": {"countries": {"__all__": {"region_datasets"}}}, +} + + +def pre_region_datasets_manifest_payload( + manifest: V2VersionManifest, +) -> dict[str, object]: + """Return the manifest as readers released before region datasets accept it.""" + return manifest.model_dump( + mode="json", + exclude={"versions": {"__all__": _REGION_DATASETS_FIELD}}, + ) + + def pre_cache_manifest_payload(manifest: V2VersionManifest) -> dict[str, object]: """Return the manifest as readers released before cache releases accept it.""" return manifest.model_dump( mode="json", - exclude={"versions": {"__all__": {"cache": True, "validation": {"cache"}}}}, + exclude={ + "versions": { + "__all__": { + "cache": True, + "validation": {"cache"}, + **_REGION_DATASETS_FIELD, + } + } + }, ) def _current_manifest(store: V2ManifestStore) -> object: - current = store.get(ACTIVE_MANIFEST_KEY) - if current is None: - # Until the first publish that writes both keys. - current = store.get(PRE_CACHE_MANIFEST_KEY) - return current + # The newest key a publish has written yet: each older key holds the same + # manifest without the fields its readers do not accept. + for key in ( + REGION_DATASETS_MANIFEST_KEY, + ACTIVE_MANIFEST_KEY, + PRE_CACHE_MANIFEST_KEY, + ): + current = store.get(key) + if current is not None: + return current + return None def build_next_v2_manifest( @@ -288,7 +322,8 @@ def publish_v2_manifest( ) # Each key holds one entire validated document, and a reader reads only # one key, so no reader observes a partial update assembled from several. - store[ACTIVE_MANIFEST_KEY] = next_manifest.model_dump(mode="json") + store[REGION_DATASETS_MANIFEST_KEY] = next_manifest.model_dump(mode="json") + store[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(next_manifest) store[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(next_manifest) return next_manifest diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 82bca08ba..ce9cd8daa 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -10,10 +10,12 @@ Stage12BundleManifest, Stage12CountryBundle, Stage12Dataset, + Stage12RegionDataset, ) from policyengine_simulation_contract.stage12_manifest import ( ACTIVE_MANIFEST_KEY, PRE_CACHE_MANIFEST_KEY, + REGION_DATASETS_MANIFEST_KEY, V1_ROUTING_STATE_NAME, V2_VERSION_MANIFEST_NAME, Stage12CacheRelease, @@ -24,6 +26,7 @@ assert_separate_manifest_names, manifest_digest, pre_cache_manifest_payload, + pre_region_datasets_manifest_payload, publish_v2_manifest, v2_application_name, validate_manifest_has_no_credentials, @@ -159,13 +162,39 @@ def test_v2_application_name_rejects_versions_a_modal_name_cannot_hold( v2_application_name(version) +def _regional_worker(version: str = "5.2.0") -> V2WorkerVersion: + """Return a cached worker whose US bundle names each region type's dataset.""" + worker = _cached_worker(version) + countries = tuple( + country.model_copy( + update={ + "region_datasets": ( + Stage12RegionDataset( + region_type="national", dataset=country.default_dataset + ), + Stage12RegionDataset( + region_type="state", dataset=f"{country.default_dataset}_local" + ), + ) + } + ) + if country.country == "us" + else country + for country in worker.bundle.countries + ) + return worker.model_copy( + update={"bundle": worker.bundle.model_copy(update={"countries": countries})} + ) + + def test_publish_writes_one_complete_v2_document_per_key() -> None: store = FakeStore() - manifest = publish_v2_manifest(store=store, worker=_cached_worker()) + manifest = publish_v2_manifest(store=store, worker=_regional_worker()) assert store.writes == [ - (ACTIVE_MANIFEST_KEY, manifest.model_dump(mode="json")), + (REGION_DATASETS_MANIFEST_KEY, manifest.model_dump(mode="json")), + (ACTIVE_MANIFEST_KEY, pre_region_datasets_manifest_payload(manifest)), (PRE_CACHE_MANIFEST_KEY, pre_cache_manifest_payload(manifest)), ] assert manifest.default_version == "5.2.0" @@ -205,6 +234,75 @@ def test_readers_released_before_the_cache_get_only_fields_they_accept() -> None } +# The fields of a country bundle before region datasets. Readers released +# before them forbid any other field, so adding one would stop them loading. +_COUNTRY_BUNDLE_FIELDS_BEFORE_REGION_DATASETS = { + "country", + "country_package_name", + "country_package_version", + "country_package_requirement", + "data_package_name", + "data_package_version", + "data_release_version", + "data_artifact_revision", + "default_dataset", + "default_dataset_uri", + "datasets", +} + + +@pytest.mark.parametrize("key", [ACTIVE_MANIFEST_KEY, PRE_CACHE_MANIFEST_KEY]) +def test_readers_released_before_region_datasets_get_only_fields_they_accept( + key: str, +) -> None: + store = FakeStore() + publish_v2_manifest(store=store, worker=_regional_worker()) + + worker = store.values[key]["versions"]["5.2.0"] + for country in worker["bundle"]["countries"]: + assert set(country) == _COUNTRY_BUNDLE_FIELDS_BEFORE_REGION_DATASETS + + +def test_the_region_dataset_key_keeps_region_datasets_and_cache() -> None: + store = FakeStore() + publish_v2_manifest(store=store, worker=_regional_worker()) + + loaded = V2ManifestLoader(store).load().versions["5.2.0"] + us = next(country for country in loaded.bundle.countries if country.country == "us") + assert us.dataset_for_region_type("state") == "populace_us_2024_local" + assert us.dataset_for_region_type("national") == "populace_us_2024" + assert us.dataset_for_region_type("congressional_district") is None + assert loaded.cache is not None + + +def test_a_bundle_published_before_region_datasets_has_none() -> None: + worker = _cached_worker() + us = next(country for country in worker.bundle.countries if country.country == "us") + + assert us.region_datasets is None + assert us.dataset_for_region_type("state") is None + + +def test_first_publish_with_region_datasets_keeps_versions_from_the_cache_key() -> None: + earlier = _cached_worker("5.1.0").model_dump(mode="json") + store = FakeStore( + { + ACTIVE_MANIFEST_KEY: { + "schema_version": 1, + "generation": "5.1.0:earlier", + "default_version": "5.1.0", + "versions": {"5.1.0": earlier}, + } + } + ) + assert V2ManifestLoader(store).load().default_version == "5.1.0" + + manifest = publish_v2_manifest(store=store, worker=_regional_worker()) + + assert sorted(manifest.versions) == ["5.1.0", "5.2.0"] + assert V2ManifestLoader(store).load() == manifest + + def test_first_publish_with_cache_keeps_versions_from_the_pre_cache_key() -> None: earlier = _worker("5.1.0").model_dump( mode="json", exclude={"cache": True, "validation": {"cache"}} diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py index b6975584c..942fa7ef3 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py @@ -8,6 +8,7 @@ CountryId, Stage12CountryBundle, Stage12Dataset, + Stage12RegionDataset, ) from .errors import Stage12BundleError @@ -169,6 +170,7 @@ def normalize_country_bundle( f"data_releases.{country}.datasets", ) datasets: list[Stage12Dataset] = [] + identity_by_path: dict[str, str] = {} for identity in sorted(dataset_items): dataset = require_mapping( dataset_items[identity], @@ -194,6 +196,7 @@ def normalize_country_bundle( dataset.get("sha256"), f"data_releases.{country}.datasets.{identity}.sha256", ) + identity_by_path.setdefault(path, identity) datasets.append( Stage12Dataset( identity=identity, @@ -256,4 +259,48 @@ def normalize_country_bundle( default_dataset=default_dataset, default_dataset_uri=default_dataset_uri, datasets=tuple(datasets), + region_datasets=_region_datasets( + release.get("region_datasets"), + identity_by_path=identity_by_path, + country=country, + ), ) + + +def _region_datasets( + value: object, + *, + identity_by_path: Mapping[str, str], + country: CountryId, +) -> tuple[Stage12RegionDataset, ...] | None: + """Return the dataset each region type loads, from ``region_datasets``. + + Each entry's ``path_template`` is matched to a dataset's ``path``, as + .py's ``resolve_region_dataset_path`` resolves it. A template with + placeholders names one file per region rather than one dataset, so it + maps to None. A bundle without ``region_datasets`` gives None. + """ + if value is None: + return None + entries = require_mapping(value, f"data_releases.{country}.region_datasets") + region_datasets = [] + for region_type in sorted(entries): + field_path = f"data_releases.{country}.region_datasets.{region_type}" + template = require_text( + require_mapping(entries[region_type], field_path).get("path_template"), + f"{field_path}.path_template", + ) + if "{" in template: + dataset = None + else: + dataset = identity_by_path.get(template) + if dataset is None: + msg = ( + f"PolicyEngine.py bundle {field_path} names {template!r}, " + "which is not one of its datasets" + ) + raise Stage12BundleError(msg) + region_datasets.append( + Stage12RegionDataset(region_type=region_type, dataset=dataset) + ) + return tuple(region_datasets) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py index 45d98f3f8..f12988ae5 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py @@ -69,6 +69,9 @@ def test_complete_packaged_bundle_selects_every_worker_dependency_and_dataset() lambda bundle: bundle["data_releases"]["us"]["certified_data_artifact"].update( {"sha256": "0" * 64} ), + lambda bundle: bundle["data_releases"]["us"]["region_datasets"].update( + {"state": {"path_template": "no_such_dataset.h5"}} + ), ], ids=[ "missing-data-releases", @@ -78,6 +81,7 @@ def test_complete_packaged_bundle_selects_every_worker_dependency_and_dataset() "dataset-uri-revision-mismatch", "malformed-dataset-digest", "certified-artifact-mismatch", + "region-dataset-names-no-dataset", ], ) def test_missing_or_internally_inconsistent_bundle_values_fail_closed(mutate) -> None: @@ -130,3 +134,37 @@ def test_only_bundle_schema_version_2_is_supported(schema_version) -> None: Stage12BundleError, match=r"^PolicyEngine\.py bundle schema_version must be 2$" ): normalize_stage12_bundle(bundle) + + +def _us(bundle): + return next(country for country in bundle.countries if country.country == "us") + + +def test_the_installed_bundle_names_the_dataset_each_us_region_type_loads() -> None: + us = _us(normalize_stage12_bundle(_bundle())) + + assert us.dataset_for_region_type("national") == "populace_us_2024" + assert us.dataset_for_region_type("state") == "populace_us_2024_acs_local" + assert ( + us.dataset_for_region_type("congressional_district") + == "populace_us_2024_acs_local" + ) + + +def test_a_per_region_file_pattern_names_no_single_dataset() -> None: + bundle = _bundle() + bundle["data_releases"]["us"]["region_datasets"]["state"] = { + "path_template": "states/{state_code}.h5" + } + + us = _us(normalize_stage12_bundle(bundle)) + + assert any(item.region_type == "state" for item in us.region_datasets) + assert us.dataset_for_region_type("state") is None + + +def test_a_bundle_without_region_datasets_records_none() -> None: + bundle = _bundle() + bundle["data_releases"]["us"].pop("region_datasets") + + assert _us(normalize_stage12_bundle(bundle)).region_datasets is None From e559fdd3a97b12e7673387382744f532824d7c16 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 23:39:44 +0400 Subject: [PATCH 10/25] Accept regional dataset provenance in Stage 12 and check it is loaded A Stage 12 worker now accepts a simulation whose provenance names the dataset the installed bundle gives its region type (state and congressional_district regions: the ACS-local dataset), with that dataset's URI and revision, as well as the default dataset. The default stays accepted for state and district runs because entry services released before regional provenance send it until they are updated; any other dataset, and a regional dataset for a national run, are refused as before. After resolving the region's dataset, the worker requires the provenance to name the dataset it loads, by name and URI, so a comparison record can no longer name one dataset while the calculation used another. The one accepted difference is default provenance on a run that loads a regional dataset, from those older entry services. Co-Authored-By: Claude Opus 5.5 --- .../stage12_runtime/simulation.py | 75 ++++++- .../tests/test_stage12_runtime.py | 204 +++++++++++++++++- 2 files changed, 271 insertions(+), 8 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 09c67bb42..2953b2e88 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -76,6 +76,9 @@ logger = logging.getLogger(__name__) +# A bare two-letter US region is a state code, such as "ut". +_US_STATE_CODE_LENGTH = 2 + @dataclass(frozen=True) class SimulationCalculation: @@ -174,6 +177,23 @@ def _require_context( raise ValueError(msg) +def _region_type(country: str, region_code: str | None) -> str | None: + """Return the bundle region type of a requested region code. + + ``state/ut`` gives ``state`` and ``congressional_district/UT-01`` gives + ``congressional_district``; the country itself (or no region) is + ``national``, and a bare two-letter US code is a state, as the gateway + reads it. None means the code names no region type. + """ + if not region_code or region_code.lower() == country: + return "national" + if "/" in region_code: + return region_code.split("/", 1)[0] + if country == "us" and len(region_code) == _US_STATE_CODE_LENGTH: + return "state" + return None + + def _require_installed_bundle(simulation: SimulationExecutionInput) -> None: resolved = load_stage12_bundle() if resolved.bundle_manifest_sha256 != simulation.bundle.bundle_manifest_sha256: @@ -187,12 +207,32 @@ def _require_installed_bundle(simulation: SimulationExecutionInput) -> None: for item in resolved.bundle.countries if item.country == simulation.geography.country ) + # A region loads the dataset the bundle names for its region type. Entry + # services released before regional provenance name the default dataset + # for every region, so the default stays accepted until they are updated. + regional = country.dataset_for_region_type( + _region_type(simulation.geography.country, simulation.geography.region) + ) + named = simulation.bundle.dataset.identity + regional_entry = next( + (item for item in country.datasets if item.identity == regional), None + ) + if regional_entry is not None and named == regional_entry.identity: + expected_dataset = { + "dataset_identity": regional_entry.identity, + "dataset_uri": regional_entry.uri, + "dataset_revision": regional_entry.artifact_revision, + } + else: + expected_dataset = { + "dataset_identity": country.default_dataset, + "dataset_uri": country.default_dataset_uri, + "dataset_revision": country.data_artifact_revision, + } expected = { "country_package_name": country.country_package_name, "country_package_version": country.country_package_version, - "dataset_identity": country.default_dataset, - "dataset_uri": country.default_dataset_uri, - "dataset_revision": country.data_artifact_revision, + **expected_dataset, } actual = { "country_package_name": simulation.bundle.country_package_name, @@ -206,6 +246,34 @@ def _require_installed_bundle(simulation: SimulationExecutionInput) -> None: raise RuntimeError(msg) +def _require_loaded_dataset_matches_provenance( + simulation: SimulationExecutionInput, selection: DatasetSelection +) -> None: + """Require a calculation to load the dataset its provenance names. + + The provenance is recorded with the comparison run, so it must name the + dataset the region actually loads. The one accepted difference is a + state or district run whose provenance names the default dataset while it + loads a regional one: entry services released before regional provenance + send that until they are updated. + """ + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _resolve_dataset_reference, + ) + + named = simulation.bundle.dataset + if (selection.name, selection.uri) == (named.identity, named.uri): + return + if not selection.is_default and named.identity == _resolve_dataset_reference( + simulation.geography.country + ): + return + msg = ( + "simulation provenance names a dataset other than the dataset its region loads" + ) + raise RuntimeError(msg) + + def calculate_simulation_frames( simulation: PlannedSimulationExecutionInput, *, @@ -260,6 +328,7 @@ def calculate_simulation_frames( params, region_resolution=region, ) + _require_loaded_dataset_matches_provenance(simulation, dataset_selection) with stage_span(Stage.DATASET_LOAD): dataset = _load_dataset( params, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 8019886a4..04d3e6bba 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -548,9 +548,7 @@ def test_calculator_uses_the_current_dataset_selection_contract(monkeypatch) -> country_module = object() region = RegionResolution(code="us") selection = DatasetSelection( - name="populace_us_2024", - uri="hf://policyengine/populace-us/populace_us_2024.h5@revision", - is_default=True, + name="populace_us_2024", uri=_bundle().dataset.uri, is_default=True ) dataset = object() received: dict[str, object] = {} @@ -635,7 +633,7 @@ class _StoppedAfterLoadingError(Exception): """Raised by a stubbed simulation build once the dataset has loaded.""" -def _stage12_load(monkeypatch, *, region, selection, baked_folder): +def _stage12_load(monkeypatch, *, region, selection, baked_folder, planned=None): """Run a Stage 12 calculation up to its dataset load; return what it did.""" from policyengine_simulation_executor import simulation_runtime from policyengine_simulation_executor.stage12_cache import state_datasets @@ -676,7 +674,7 @@ def stop_after_loading(*_args, **_kwargs): ) monkeypatch.setattr(worker, "build_stage12_simulation", stop_after_loading) - planned = _planned_simulation(SimulationRole.BASELINE) + planned = planned or _planned_simulation(SimulationRole.BASELINE) with pytest.raises(_StoppedAfterLoadingError): worker.calculate_simulation_frames(planned) return SimpleNamespace(loads=loads, lookups=lookups, year=planned.year) @@ -738,6 +736,202 @@ def test_stage12_state_run_without_a_baked_file_loads_on_demand(monkeypatch) -> assert loaded.loads[0]["data_folder"] is None +REGIONAL_URI = "hf://policyengine/data/populace_us_2024_acs_local.h5@regional-revision" + + +def _regional_provenance() -> BundleProvenance: + return _bundle().model_copy( + update={ + "dataset": DatasetProvenance( + identity="populace_us_2024_acs_local", + uri=REGIONAL_URI, + artifact_revision="regional-revision", + data_package_name="populace-data", + data_package_version="0.1.0", + ) + } + ) + + +def _installed_bundle_with_regions(): + from policyengine_simulation_contract.stage12_bundle import ( + ResolvedStage12Bundle, + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, + Stage12RegionDataset, + ) + + national = Stage12Dataset( + identity="populace_us_2024", + uri=_bundle().dataset.uri, + artifact_revision="revision", + sha256="a" * 64, + repo_type="dataset", + ) + regional = Stage12Dataset( + identity="populace_us_2024_acs_local", + uri=REGIONAL_URI, + artifact_revision="regional-revision", + sha256="c" * 64, + repo_type="dataset", + ) + return ResolvedStage12Bundle( + bundle=Stage12BundleManifest( + policyengine_version="5.2.0", + policyengine_requirement="policyengine==5.2.0", + core_package_version="3.0.0", + core_package_requirement="policyengine-core==3.0.0", + countries=( + Stage12CountryBundle( + country="us", + country_package_name="policyengine-us", + country_package_version="1.764.6", + country_package_requirement="policyengine-us==1.764.6", + data_package_name="populace-data", + data_package_version="0.1.0", + data_release_version="revision", + data_artifact_revision="revision", + default_dataset=national.identity, + default_dataset_uri=national.uri, + datasets=(national, regional), + region_datasets=( + Stage12RegionDataset( + region_type="national", dataset=national.identity + ), + Stage12RegionDataset( + region_type="state", dataset=regional.identity + ), + Stage12RegionDataset( + region_type="congressional_district", + dataset=regional.identity, + ), + ), + ), + ), + ), + bundle_manifest_sha256="b" * 64, + ) + + +def _simulation_for(region: str, provenance: BundleProvenance): + simulation = _planned_simulation(SimulationRole.BASELINE) + return simulation.model_copy( + update={ + "geography": GeographySelection(country="us", region=region), + "bundle": provenance, + "population": DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=provenance.dataset.uri, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256="a" * 64, + ) + ), + } + ) + + +@pytest.mark.parametrize( + ("region", "provenance"), + [ + ("state/ut", _regional_provenance), + ("ut", _regional_provenance), + ("congressional_district/UT-01", _regional_provenance), + ("us", _bundle), + # Entry services released before regional provenance send the default + # dataset for state runs until they are updated. + ("state/ut", _bundle), + ], +) +def test_installed_bundle_accepts_the_dataset_a_region_loads( + monkeypatch, region, provenance +) -> None: + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + monkeypatch.setattr(worker, "load_stage12_bundle", _installed_bundle_with_regions) + + worker._require_installed_bundle(_simulation_for(region, provenance())) + + +@pytest.mark.parametrize( + ("region", "dataset"), + [ + ("us", "populace_us_2024_acs_local"), + ("state/ut", "source_enrichment"), + ], +) +def test_installed_bundle_refuses_a_dataset_the_region_does_not_load( + monkeypatch, region, dataset +) -> None: + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + monkeypatch.setattr(worker, "load_stage12_bundle", _installed_bundle_with_regions) + provenance = _regional_provenance() + provenance = provenance.model_copy( + update={"dataset": provenance.dataset.model_copy(update={"identity": dataset})} + ) + + with pytest.raises(RuntimeError, match="differs from the installed bundle"): + worker._require_installed_bundle(_simulation_for(region, provenance)) + + +def _selection(name: str, uri: str, *, is_default: bool): + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + + return DatasetSelection(name=name, uri=uri, is_default=is_default) + + +def test_a_calculation_loads_the_dataset_its_provenance_names(monkeypatch) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + loaded = _stage12_load( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_selection( + "populace_us_2024_acs_local", REGIONAL_URI, is_default=False + ), + baked_folder=None, + planned=_simulation_for("state/ut", _regional_provenance()), + ) + + assert loaded.loads[0]["selection"].name == "populace_us_2024_acs_local" + + +def test_a_calculation_refuses_provenance_naming_another_dataset( + monkeypatch, +) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + with pytest.raises(RuntimeError, match="dataset its region loads"): + _stage12_load( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_selection( + "populace_us_2024", _bundle().dataset.uri, is_default=True + ), + baked_folder=None, + planned=_simulation_for("state/ut", _regional_provenance()), + ) + + +def test_default_provenance_from_an_older_entry_service_is_still_loaded( + monkeypatch, +) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + loaded = _stage12_load( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_selection( + "populace_us_2024_acs_local", REGIONAL_URI, is_default=False + ), + baked_folder=None, + planned=_simulation_for("state/ut", _bundle()), + ) + + assert loaded.loads[0]["selection"].name == "populace_us_2024_acs_local" + + def test_uk_boundary_version_is_detected_from_the_unscoped_dataset() -> None: dataset = type( "Dataset", From 2524ce4f11f60b93d2d282b1ef433f281f34d0c3 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 23:42:36 +0400 Subject: [PATCH 11/25] Send the regional dataset's provenance for Stage 12 state and district runs The entry adapter now records, for a state or congressional-district request, the dataset the worker's published bundle names for that region type (the ACS-local dataset): its identity, URI and revision in the dataset and bundle provenance, and its URI and SHA-256 as the population artifact. A bare two-letter US region counts as a state, as the gateway reads it. The entry service keeps its own copy of the worker's region-type reading. National and UK requests, and requests to workers whose manifest predates region datasets, still name the default dataset, and the request's data value is still checked against the default. Workers accept regional provenance since the previous commit, so they never refuse what this sends. Co-Authored-By: Claude Opus 5.5 --- .../stage12_adapter.py | 45 +++++++++- .../tests/stage12_fixtures.py | 51 ++++++++++++ .../tests/test_stage12_adapter.py | 82 +++++++++++++++++++ 3 files changed, 176 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py index 9f008fd77..de13ee099 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py @@ -86,6 +86,8 @@ def eligible(self) -> bool: ) _AGGREGATES = tuple(ReportAggregate) _YEAR_DIGIT_COUNT = 4 +# A bare two-letter US region is a state code, such as "ut". +_US_STATE_CODE_LENGTH = 2 def _skip(reason: ComparisonSkipReason) -> AdaptedComparisonRun: @@ -169,6 +171,40 @@ def _accepted_dataset_values(selected_dataset: Stage12Dataset) -> set[str | None } +def _region_type(country: str, region: str) -> str | None: + """Return the bundle region type of a requested region code. + + The same reading as the Stage 12 worker's: ``state/ut`` is ``state``, + ``congressional_district/UT-01`` is ``congressional_district``, the + country itself is ``national``, and a bare two-letter US code is a state, + as the gateway reads it. The entry service cannot import the worker's + code, so it keeps this copy. + """ + if region.lower() == country: + return "national" + if "/" in region: + return region.split("/", 1)[0] + if country == "us" and len(region) == _US_STATE_CODE_LENGTH: + return "state" + return None + + +def _region_dataset( + bundle_country: Stage12CountryBundle, region: str +) -> Stage12Dataset | None: + """Return the dataset the bundle says this region loads, if it names one. + + Workers published before region datasets name none, and their regions + load the default dataset. + """ + identity = bundle_country.dataset_for_region_type( + _region_type(bundle_country.country, region) + ) + return next( + (item for item in bundle_country.datasets if item.identity == identity), None + ) + + def adapt_annual_comparison( payload: dict[str, Any], *, @@ -191,13 +227,18 @@ def adapt_annual_comparison( if skip_reason is not None: return _skip(skip_reason) - selected_dataset = next( + default_dataset = next( item for item in bundle_country.datasets if item.identity == bundle_country.default_dataset ) - if payload.get("data") not in _accepted_dataset_values(selected_dataset): + if payload.get("data") not in _accepted_dataset_values(default_dataset): return _skip(ComparisonSkipReason.UNSUPPORTED_DATASET) + # A state or district region loads the dataset the bundle names for its + # region type, so its provenance names that dataset. + selected_dataset = ( + _region_dataset(bundle_country, payload["region"]) or default_dataset + ) return AdaptedComparisonRun( report=_report_input( diff --git a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py index 46ebe97e7..fa4e83f6a 100644 --- a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py +++ b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py @@ -9,6 +9,7 @@ Stage12BundleManifest, Stage12CountryBundle, Stage12Dataset, + Stage12RegionDataset, ) from policyengine_simulation_contract.stage12_execution import ( ComparisonReportRecord, @@ -92,6 +93,56 @@ def worker() -> V2WorkerVersion: ) +REGIONAL_DATASET = "populace_us_2024_acs_local" +REGIONAL_REVISION = "acs-local-revision" +REGIONAL_URI = f"hf://policyengine/data/{REGIONAL_DATASET}.h5@{REGIONAL_REVISION}" + + +def regional_worker() -> V2WorkerVersion: + """Return ``worker()`` whose US bundle sends state and district regions elsewhere. + + State and congressional-district regions load the ACS-local dataset; + national regions keep the default. + """ + base = worker() + countries = [] + for country in base.bundle.countries: + if country.country != "us": + countries.append(country) + continue + regional = Stage12Dataset( + identity=REGIONAL_DATASET, + uri=REGIONAL_URI, + artifact_revision=REGIONAL_REVISION, + sha256="c" * 64, + repo_type="dataset", + ) + countries.append( + country.model_copy( + update={ + "datasets": (*country.datasets, regional), + "region_datasets": ( + Stage12RegionDataset( + region_type="national", dataset=country.default_dataset + ), + Stage12RegionDataset( + region_type="state", dataset=REGIONAL_DATASET + ), + Stage12RegionDataset( + region_type="congressional_district", + dataset=REGIONAL_DATASET, + ), + ), + } + ) + ) + return base.model_copy( + update={ + "bundle": base.bundle.model_copy(update={"countries": tuple(countries)}) + } + ) + + def eligible_payload() -> dict: return { "country": "us", diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index 401c15492..8a2d0dc10 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -311,3 +311,85 @@ def test_adapter_checks_eligibility_in_order(update, reason) -> None: assert result.skip_reason is reason assert result.eligible is (reason is None) + + +def _datasets_named(report) -> set[tuple[str, str, str, str, str]]: + return { + ( + simulation.bundle.dataset.identity, + simulation.bundle.dataset.uri, + simulation.bundle.dataset.artifact_revision, + simulation.population.artifact.uri, + simulation.population.artifact.content_sha256, + ) + for simulation in (report.baseline, report.reform) + } + + +@pytest.mark.parametrize("region", ["state/ut", "ut", "congressional_district/UT-01"]) +def test_state_and_district_requests_name_the_regional_dataset(region) -> None: + from stage12_fixtures import ( + REGIONAL_DATASET, + REGIONAL_REVISION, + REGIONAL_URI, + regional_worker, + ) + + result = adapt_annual_comparison( + {**eligible_payload(), "region": region}, + evaluation_id=EVALUATION_ID, + worker=regional_worker(), + ) + + assert result.skip_reason is None + assert _datasets_named(result.report) == { + (REGIONAL_DATASET, REGIONAL_URI, REGIONAL_REVISION, REGIONAL_URI, "c" * 64) + } + assert result.report.baseline.geography.region == region + + +def test_a_national_request_keeps_the_default_dataset_with_regional_workers() -> None: + from stage12_fixtures import regional_worker + + payload = {**eligible_payload(), "include_cliffs": True} + payload["spm"] = {"geography_kind": "national"} + + national = adapt_annual_comparison( + payload, evaluation_id=EVALUATION_ID, worker=regional_worker() + ) + today = adapt_annual_comparison( + payload, evaluation_id=EVALUATION_ID, worker=worker() + ) + + assert national.report.model_dump(mode="json") == today.report.model_dump( + mode="json" + ) + + +def test_a_worker_without_region_datasets_gets_the_default_dataset() -> None: + result = adapt_annual_comparison( + {**eligible_payload(), "region": "state/ut"}, + evaluation_id=EVALUATION_ID, + worker=worker(), + ) + + assert result.skip_reason is None + assert {item[0] for item in _datasets_named(result.report)} == {"populace_us_2024"} + + +def test_a_uk_request_keeps_the_default_dataset_with_regional_workers() -> None: + from stage12_fixtures import regional_worker + + result = adapt_annual_comparison( + { + **eligible_payload(), + "country": "uk", + "region": "country/scotland", + "version": "2.90.2", + }, + evaluation_id=EVALUATION_ID, + worker=regional_worker(), + ) + + assert result.skip_reason is None + assert {item[0] for item in _datasets_named(result.report)} == {"populace_uk_2023"} From be3f5554ff76f3f969fee0912980843925a2cee4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:08:15 +0400 Subject: [PATCH 12/25] Keep region datasets out of the Stage 12 bundle digest Adding region_datasets to the bundle changed bundle_digest for every bundle. A deploy replaces the running v2 application, whose name depends only on the .py version, before the new version manifest is published, and entry services send the digest they read from that manifest; until the publish, every comparison run would have been refused with "installed bundle digest differs", and the Stage 12 cache, keyed by the digest, would have been rebuilt. region_datasets comes from the same .py bundle whose version the digest already covers, so the digest now leaves it out. The installed 6.7.5 bundle keeps the digest it had before this change. Co-Authored-By: Claude Opus 5.5 --- .../stage12_bundle/manifest.py | 14 +++++++++-- .../tests/test_stage12_bundle.py | 23 +++++++++++++++++++ 2 files changed, 35 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py index cf1a349e4..f80965a2a 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py @@ -103,9 +103,19 @@ def normalize_stage12_bundle( def bundle_digest(bundle: Stage12BundleManifest) -> str: - """Return the sha256 of the manifest's canonical (key-sorted, compact) JSON.""" + """Return the sha256 of the manifest's canonical (key-sorted, compact) JSON. + + Each country's ``region_datasets`` is left out. It comes from the same .py + bundle, whose version the digest covers, and leaving it out keeps the + digest a running v2 application was validated with: a redeploy that adds + it replaces the application before the new version manifest is + published, and entry services send the digest they read from that + manifest. + """ payload = json.dumps( - bundle.model_dump(mode="json"), + bundle.model_dump( + mode="json", exclude={"countries": {"__all__": {"region_datasets"}}} + ), sort_keys=True, separators=(",", ":"), ensure_ascii=False, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py index f12988ae5..7fc624d6e 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py @@ -168,3 +168,26 @@ def test_a_bundle_without_region_datasets_records_none() -> None: bundle["data_releases"]["us"].pop("region_datasets") assert _us(normalize_stage12_bundle(bundle)).region_datasets is None + + +def test_region_datasets_leave_the_bundle_digest_unchanged() -> None: + """A worker redeployed with region datasets keeps the digest it had. + + The redeploy replaces the running v2 application before the new version + manifest is published, and entry services send the digest they read from + that manifest, so a changed digest would refuse every run in between. + Region datasets come from the same .py bundle, whose version the digest + already covers. + """ + bundle = normalize_stage12_bundle(_bundle()) + without_region_datasets = bundle.model_copy( + update={ + "countries": tuple( + country.model_copy(update={"region_datasets": None}) + for country in bundle.countries + ) + } + ) + + assert _us(bundle).region_datasets is not None + assert bundle_digest(bundle) == bundle_digest(without_region_datasets) From 4ed83ddc0554ef9779f7d33e5131985558093103 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:12:23 +0400 Subject: [PATCH 13/25] Read region codes once, as the runtime does, for regional provenance The Stage 12 worker and the entry adapter each kept a copy of the region-type reading, and neither ignored surrounding whitespace or upper case as the runtime's region normalisation does: for a region of State/UT the entry sent default provenance while the worker loaded the ACS-local dataset, hidden only by the temporary acceptance of the default. The contract library now has one region_type_of, matching the runtime's reading, and Stage12CountryBundle.dataset_for_region, which both use. Co-Authored-By: Claude Opus 5.5 --- .../stage12_bundle.py | 30 ++++++++++++++- .../tests/test_stage12_manifest.py | 38 +++++++++++++++++++ .../stage12_adapter.py | 24 +----------- .../tests/test_stage12_adapter.py | 4 +- .../stage12_runtime/simulation.py | 24 +----------- .../tests/test_stage12_runtime.py | 1 + 6 files changed, 73 insertions(+), 48 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py index 4d2b48db1..90dac9bb3 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py @@ -10,6 +10,30 @@ Sha256Digest = Annotated[str, Field(pattern=r"^[0-9a-f]{64}$")] CountryId = Literal["us", "uk"] SUPPORTED_COUNTRIES: tuple[CountryId, ...] = ("us", "uk") +# A bare two-letter US region is a state code, such as "ut". +_US_STATE_CODE_LENGTH = 2 + + +def region_type_of(country: str, region: str | None) -> str | None: + """Return the region type a requested region code names, as bundles key it. + + Codes are read as the simulation runtime normalises them: surrounding + whitespace is ignored and the country itself, or no region, is + ``national``. A bare two-letter US code is a state and a bare UK name is + a UK country. Otherwise the type is the part before the first slash, in + lower case: ``State/UT`` is ``state`` and ``congressional_district/UT-01`` + is ``congressional_district``. None means the code names no region type. + """ + raw = (region or "").strip() + if not raw or raw.lower() == country: + return "national" + if "/" in raw: + return raw.split("/", 1)[0].lower() + if country == "us" and len(raw) == _US_STATE_CODE_LENGTH: + return "state" + if country == "uk" and raw.lower() not in SUPPORTED_COUNTRIES: + return "country" + return None class StrictBundleModel(BaseModel): @@ -62,13 +86,17 @@ class Stage12CountryBundle(StrictBundleModel): datasets: tuple[Stage12Dataset, ...] region_datasets: tuple[Stage12RegionDataset, ...] | None = None - def dataset_for_region_type(self, region_type: str) -> str | None: + def dataset_for_region_type(self, region_type: str | None) -> str | None: """Return the dataset identity ``region_type`` regions load, or None.""" for item in self.region_datasets or (): if item.region_type == region_type: return item.dataset return None + def dataset_for_region(self, region: str | None) -> str | None: + """Return the dataset identity a requested region code loads, or None.""" + return self.dataset_for_region_type(region_type_of(self.country, region)) + class Stage12BundleManifest(StrictBundleModel): """Stage 12's normalized subset of a PolicyEngine.py bundle.""" diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index ce9cd8daa..0337a8cc2 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -11,6 +11,7 @@ Stage12CountryBundle, Stage12Dataset, Stage12RegionDataset, + region_type_of, ) from policyengine_simulation_contract.stage12_manifest import ( ACTIVE_MANIFEST_KEY, @@ -411,3 +412,40 @@ def test_storage_names_must_be_distinct() -> None: ) with pytest.raises(ValueError, match="different"): assert_separate_manifest_names(v1_name="same", v2_name="same") + + +@pytest.mark.parametrize( + ("country", "region", "region_type"), + [ + ("us", None, "national"), + ("us", "", "national"), + ("us", " US ", "national"), + ("us", "state/ut", "state"), + ("us", "State/UT", "state"), + ("us", "ut", "state"), + ("us", " ut ", "state"), + ("us", "congressional_district/ut-01", "congressional_district"), + ("us", "CONGRESSIONAL_DISTRICT/UT-01", "congressional_district"), + ("us", "place/UT-67000", "place"), + ("us", "utah", None), + ("uk", "uk", "national"), + ("uk", "country/scotland", "country"), + ("uk", "scotland", "country"), + ("uk", "constituency/Aldershot", "constituency"), + ], +) +def test_region_type_reads_codes_as_the_runtime_normalises_them( + country: str, region: str | None, region_type: str | None +) -> None: + assert region_type_of(country, region) == region_type + + +def test_a_country_bundle_gives_the_dataset_a_region_code_loads() -> None: + worker = _regional_worker() + us = next(country for country in worker.bundle.countries if country.country == "us") + uk = next(country for country in worker.bundle.countries if country.country == "uk") + + assert us.dataset_for_region("State/UT") == "populace_us_2024_local" + assert us.dataset_for_region("us") == "populace_us_2024" + assert us.dataset_for_region("place/UT-67000") is None + assert uk.dataset_for_region("country/scotland") is None diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py index de13ee099..c66929968 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py @@ -86,8 +86,6 @@ def eligible(self) -> bool: ) _AGGREGATES = tuple(ReportAggregate) _YEAR_DIGIT_COUNT = 4 -# A bare two-letter US region is a state code, such as "ut". -_US_STATE_CODE_LENGTH = 2 def _skip(reason: ComparisonSkipReason) -> AdaptedComparisonRun: @@ -171,24 +169,6 @@ def _accepted_dataset_values(selected_dataset: Stage12Dataset) -> set[str | None } -def _region_type(country: str, region: str) -> str | None: - """Return the bundle region type of a requested region code. - - The same reading as the Stage 12 worker's: ``state/ut`` is ``state``, - ``congressional_district/UT-01`` is ``congressional_district``, the - country itself is ``national``, and a bare two-letter US code is a state, - as the gateway reads it. The entry service cannot import the worker's - code, so it keeps this copy. - """ - if region.lower() == country: - return "national" - if "/" in region: - return region.split("/", 1)[0] - if country == "us" and len(region) == _US_STATE_CODE_LENGTH: - return "state" - return None - - def _region_dataset( bundle_country: Stage12CountryBundle, region: str ) -> Stage12Dataset | None: @@ -197,9 +177,7 @@ def _region_dataset( Workers published before region datasets name none, and their regions load the default dataset. """ - identity = bundle_country.dataset_for_region_type( - _region_type(bundle_country.country, region) - ) + identity = bundle_country.dataset_for_region(region) return next( (item for item in bundle_country.datasets if item.identity == identity), None ) diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index 8a2d0dc10..5751a9b24 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -326,7 +326,9 @@ def _datasets_named(report) -> set[tuple[str, str, str, str, str]]: } -@pytest.mark.parametrize("region", ["state/ut", "ut", "congressional_district/UT-01"]) +@pytest.mark.parametrize( + "region", ["state/ut", "State/UT", "ut", " ut ", "congressional_district/UT-01"] +) def test_state_and_district_requests_name_the_regional_dataset(region) -> None: from stage12_fixtures import ( REGIONAL_DATASET, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index 2953b2e88..3b2bb5b56 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -76,9 +76,6 @@ logger = logging.getLogger(__name__) -# A bare two-letter US region is a state code, such as "ut". -_US_STATE_CODE_LENGTH = 2 - @dataclass(frozen=True) class SimulationCalculation: @@ -177,23 +174,6 @@ def _require_context( raise ValueError(msg) -def _region_type(country: str, region_code: str | None) -> str | None: - """Return the bundle region type of a requested region code. - - ``state/ut`` gives ``state`` and ``congressional_district/UT-01`` gives - ``congressional_district``; the country itself (or no region) is - ``national``, and a bare two-letter US code is a state, as the gateway - reads it. None means the code names no region type. - """ - if not region_code or region_code.lower() == country: - return "national" - if "/" in region_code: - return region_code.split("/", 1)[0] - if country == "us" and len(region_code) == _US_STATE_CODE_LENGTH: - return "state" - return None - - def _require_installed_bundle(simulation: SimulationExecutionInput) -> None: resolved = load_stage12_bundle() if resolved.bundle_manifest_sha256 != simulation.bundle.bundle_manifest_sha256: @@ -210,9 +190,7 @@ def _require_installed_bundle(simulation: SimulationExecutionInput) -> None: # A region loads the dataset the bundle names for its region type. Entry # services released before regional provenance name the default dataset # for every region, so the default stays accepted until they are updated. - regional = country.dataset_for_region_type( - _region_type(simulation.geography.country, simulation.geography.region) - ) + regional = country.dataset_for_region(simulation.geography.region) named = simulation.bundle.dataset.identity regional_entry = next( (item for item in country.datasets if item.identity == regional), None diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 04d3e6bba..fde76eb48 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -835,6 +835,7 @@ def _simulation_for(region: str, provenance: BundleProvenance): ("region", "provenance"), [ ("state/ut", _regional_provenance), + ("State/UT", _regional_provenance), ("ut", _regional_provenance), ("congressional_district/UT-01", _regional_provenance), ("us", _bundle), From 52407662764f20b3d8cc43b16678f13d529cf6b6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:17:09 +0400 Subject: [PATCH 14/25] Give the Stage 12 cache fetch time for the state datasets The US worker image's Stage 12 cache fetch kept its 1,800 s limit while gaining a file per state for each precomputed year, downloaded and hashed one after another on two CPUs. It now allows 3,600 s, as the main path's fetch does for the same files. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/v2_app.py | 4 +++- .../tests/test_stage12_modal_app.py | 11 +++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 0f98747e2..92642d673 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -195,7 +195,9 @@ def build_v2_image( secrets=[gcp_secret], cpu=2.0, memory=4096, - timeout=1800, + # The 63 national files and the state datasets (a file per state + # for each precomputed year), downloaded and hashed in turn. + timeout=3600, ) image = image.add_local_python_source( "src.modal", diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index 597ab967d..00ae66afc 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -223,6 +223,17 @@ def test_only_us_worker_image_fetches_stage12_cache(monkeypatch) -> None: assert not any(call[0] == "run_function" for call in module.coordinator_image.calls) +def test_stage12_cache_fetch_has_room_for_the_state_datasets(monkeypatch) -> None: + module = _load(monkeypatch) + + (fetch,) = [ + call for call in module.us_worker_image.calls if call[0] == "run_function" + ] + # The 63 national files plus a file per state for each precomputed year, + # downloaded and hashed one after another; the main path allows the same. + assert fetch[2]["timeout"] == 3600 + + def test_v2_app_declares_validation_workers_and_non_http_coordinator( monkeypatch, ) -> None: From 54c6ed4f3fc2fe2587ce6650d3d12d2fddd3d71c Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:17:09 +0400 Subject: [PATCH 15/25] Apply the per-state build review fixes to Stage 12's copy Stage 12's copy of the per-state build gets the fixes made to the main path's: - it refuses only district codes that name a registered district of a different state, so an unassigned code such as 0 no longer stops the precompute and the deploy waiting on it; - a target without a row filter is refused before anything is written, and the target's strategy is typed as an optional RowFilterStrategy; - builds are grouped by year and dataset, since one build prepares one dataset's year file; - each result records how long the year's build had run, and the log reports per build the files built and uploaded, the bytes actually uploaded, and the build time. Co-Authored-By: Claude Opus 5.5 --- .../stage12_cache/models.py | 7 +- .../stage12_cache/precompute.py | 80 ++++++++++++------ .../tests/test_stage12_cache_precompute.py | 82 ++++++++++++++++++- .../tests/test_stage12_precompute_app.py | 1 + 4 files changed, 143 insertions(+), 27 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py index 9a5ffad52..68d2aa875 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py @@ -114,11 +114,16 @@ class DatasetBuildResult(CacheObject): class StateDatasetBuildResult(CacheObject): - """A built state dataset file and whether this run uploaded it.""" + """A built state dataset file and whether this run uploaded it. + + ``build_seconds`` is how long the year's build had run when this file was + uploaded, so the last result of a build gives its whole duration. + """ year: int state_code: str uploaded: bool + build_seconds: float class BaselineBuildResult(CacheObject): diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index ac591dc49..c44b23498 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -75,10 +75,12 @@ if TYPE_CHECKING: from collections.abc import Callable, Mapping, Sequence + from types import ModuleType import modal import pandas as pd from microdf import MicroDataFrame + from policyengine.core.scoping_strategy import RowFilterStrategy PRECOMPUTE_YEARS = (2026, 2027, 2025) MANIFEST_DIGEST_PREFIX = "STAGE12_CACHE_MANIFEST_DIGEST=" @@ -238,7 +240,9 @@ class StateDatasetTarget: state_fips: int filter_variable: str dataset: str - scoping_strategy: Any = field(default=None, compare=False, repr=False) + scoping_strategy: RowFilterStrategy | None = field( + default=None, compare=False, repr=False + ) def state_dataset_targets() -> tuple[StateDatasetTarget, ...]: @@ -449,22 +453,38 @@ def _checked_state_targets( return {entry.region: targets[entry.region] for entry in expected} -def _require_districts_within_their_states(household: MicroDataFrame) -> None: - """Refuse a dataset whose district codes name a different state. +def _registered_district_geoids(country_module: ModuleType) -> set[int]: + """Return the district codes the US congressional-district regions filter on.""" + return { + int(region.scoping_strategy.variable_value) + for region in country_module.model.region_registry.get_by_type( + "congressional_district" + ) + } + + +def _require_districts_within_their_states( + household: MicroDataFrame, district_geoids: set[int] +) -> None: + """Refuse a dataset in which a household sits in another state's district. District requests read their state's file and filter it by ``congressional_district_geoid`` (state FIPS times 100 plus the district - number). Both columns are labels; nothing here is weighted. + number). A code that names no district, such as an unassigned 0, matches + no district request and changes no result. Both columns are labels; + nothing here is weighted. """ if "congressional_district_geoid" not in household.columns: return - outside = (household["congressional_district_geoid"] // 100) != household[ - "state_fips" - ] - if outside.any(): + district = household["congressional_district_geoid"] + in_another_state = district.isin(district_geoids) & ( + (district // 100) != household["state_fips"] + ) + if in_another_state.any(): msg = ( - "Some households have a congressional district code outside their " - "state; their state's file would not match the district's results" + "Some households have a congressional district code that names " + "another state's district; their state's file would not match " + "that district's results" ) raise RuntimeError(msg) @@ -495,6 +515,7 @@ def build_state_datasets_impl( if not expected: return () + started = time.monotonic() targets = _checked_state_targets(year, expected) dataset = expected[0].dataset work_folder = work_folder or Path( @@ -516,12 +537,18 @@ def build_state_datasets_impl( if stamp is None: msg = f"The prepared year file {source.filepath} has no year-file stamp" raise RuntimeError(msg) - _require_districts_within_their_states(source.data.household) + _require_districts_within_their_states( + source.data.household, _registered_district_geoids(country_module) + ) store = Stage12CacheStore(bucket) results = [] for entry in expected: - rows = targets[entry.region].scoping_strategy.apply( + strategy = targets[entry.region].scoping_strategy + if strategy is None: + msg = f"{entry.region} has no row filter to build its state file with" + raise RuntimeError(msg) + rows = strategy.apply( entity_data=source.data.entity_data, group_entities=US_GROUP_ENTITIES, year=year, @@ -549,6 +576,7 @@ def build_state_datasets_impl( year=year, state_code=entry.state_code, uploaded=uploaded, + build_seconds=round(time.monotonic() - started, 1), ) ) return tuple(results) @@ -758,12 +786,16 @@ def run_precompute( functions.build_dataset.spawn(bucket, entry.model_dump()) for entry in work.datasets ] - state_entries_by_year: dict[int, list[dict[str, Any]]] = {} + # One build prepares one dataset's year file, so builds are per year and + # dataset. + state_builds: dict[tuple[int, str], list[dict[str, Any]]] = {} for entry in work.state_datasets: - state_entries_by_year.setdefault(entry.year, []).append(entry.model_dump()) + state_builds.setdefault((entry.year, entry.dataset), []).append( + entry.model_dump() + ) state_calls = [ - functions.build_state_datasets.spawn(bucket, year, entries) - for year, entries in state_entries_by_year.items() + (year, dataset, functions.build_state_datasets.spawn(bucket, year, entries)) + for (year, dataset), entries in state_builds.items() ] for call in dataset_calls: DatasetBuildResult.model_validate(call.get()) @@ -788,16 +820,16 @@ def run_precompute( "Stage 12 cache determinism validation failed: " + "; ".join(verdict.differences) ) - if state_calls: - state_results = [ - StateDatasetBuildResult.model_validate(result) - for call in state_calls - for result in call.get() + for year, dataset, call in state_calls: + results = [ + StateDatasetBuildResult.model_validate(result) for result in call.get() ] + uploaded = [result for result in results if result.uploaded] echo( - f"State datasets: {len(state_results)} built, " - f"{sum(result.uploaded for result in state_results)} uploaded, " - f"{sum(result.size_bytes for result in state_results)} bytes" + f"State datasets {year} {dataset}: {len(results)} built, " + f"{len(uploaded)} uploaded " + f"({sum(result.size_bytes for result in uploaded)} bytes) in " + f"{max((result.build_seconds for result in results), default=0.0)} s" ) digest = functions.publish_manifest.remote(bucket, plan.model_dump()) if not isinstance(digest, str) or not digest: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py index 978bd8c00..7db8837b1 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -504,6 +504,7 @@ def state_results(_bucket, year, entries): "year": year, "state_code": entry["state_code"], "uploaded": entry["state_code"] != "WY", + "build_seconds": 12.5 if entry["state_code"] == "WY" else 6.0, } for entry in entries ] @@ -557,7 +558,7 @@ def test_precompute_waits_for_state_datasets_only_before_publishing() -> None: assert first_state_get < events.index("remote:publish") -def test_precompute_reports_the_state_dataset_bytes_uploaded() -> None: +def test_precompute_reports_each_builds_uploads_bytes_and_time() -> None: events = [] lines = [] handles = _state_run_handles(events, _plan_with_state_datasets()) @@ -568,10 +569,40 @@ def test_precompute_reports_the_state_dataset_bytes_uploaded() -> None: echo=lines.append, ) - assert "State datasets: 6 built, 3 uploaded, 600 bytes" in lines + for year in precompute.PRECOMPUTE_YEARS: + assert ( + f"State datasets {year} {REGIONAL_DATASET}: 2 built, 1 uploaded " + "(100 bytes) in 12.5 s" + ) in lines assert lines[-1] == "STAGE12_CACHE_MANIFEST_DIGEST=manifest-digest" +def test_precompute_builds_each_year_and_dataset_in_its_own_call() -> None: + events = [] + plan = _plan_with_state_datasets() + other = _state_entry(2026, "ID", exists=False).model_copy( + update={"dataset": "populace_other", "folder": "state-datasets/other/ID"} + ) + plan = plan.model_copy(update={"state_datasets": (*plan.state_datasets, other)}) + handles = _state_run_handles(events, plan) + + precompute.run_precompute( + "cache-bucket", + functions=precompute.CachePrecomputeFunctions(**handles), + echo=lambda _line: None, + ) + + assert [ + (call[1], {entry["dataset"] for entry in call[2]}) + for call in handles["build_state_datasets"].spawn_calls + ] == [ + (2026, {REGIONAL_DATASET}), + (2027, {REGIONAL_DATASET}), + (2025, {REGIONAL_DATASET}), + (2026, {"populace_other"}), + ] + + def test_warm_verified_precompute_skips_builds_and_determinism_check() -> None: events = [] plan = _plan(exists=True) @@ -1075,12 +1106,28 @@ def fake_ensure_datasets(*, datasets, years, data_folder): datasets=datasets, years=years, data_folder=data_folder ) + districts = [ + SimpleNamespace( + scoping_strategy=RowFilterStrategy( + variable_name="congressional_district_geoid", + variable_value=geoid, + ) + ) + for geoid in (4901, 4902, 4903, 4904, 5600) + ] monkeypatch.setattr( precompute, "_country_module", lambda _country: SimpleNamespace( ensure_datasets=fake_ensure_datasets, load_datasets=us_module.load_datasets, + model=SimpleNamespace( + region_registry=SimpleNamespace( + get_by_type=lambda kind: ( + districts if kind == "congressional_district" else [] + ) + ) + ), ), ) state.identity_for = fake_state_identity @@ -1172,6 +1219,37 @@ def test_refuses_a_district_that_lies_outside_its_state(self, state_build): self._build(state_build, [self._entry(state_build, "UT")]) assert state_build.client.objects == {} + def test_a_household_without_a_district_code_does_not_stop_the_build( + self, state_build + ): + state_build.entity_frames = state_build.sources.source_entity_frames( + wyoming_district_geoid=0 + ) + + results = self._build( + state_build, + [self._entry(state_build, "UT"), self._entry(state_build, "WY")], + ) + + assert [result.state_code for result in results] == ["UT", "WY"] + + def test_refuses_a_target_without_a_row_filter(self, state_build): + from dataclasses import replace + + state_build.targets[0] = replace(state_build.targets[0], scoping_strategy=None) + + with pytest.raises(RuntimeError, match="no row filter"): + self._build(state_build, [self._entry(state_build, "UT")]) + assert state_build.client.objects == {} + + def test_reports_how_long_the_years_build_had_run(self, state_build): + results = self._build( + state_build, + [self._entry(state_build, "UT"), self._entry(state_build, "WY")], + ) + + assert 0 <= results[0].build_seconds <= results[1].build_seconds + def test_refuses_a_plan_entry_whose_key_differs_in_the_container(self, state_build): entry = self._entry(state_build, "UT").model_copy( update={"path": "state-datasets/us/other/x.h5"} diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py index 2118148c9..c260afda0 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py @@ -123,6 +123,7 @@ def test_state_dataset_wrapper_validates_entries_in_and_dumps_results_out( year=2026, state_code="UT", uploaded=True, + build_seconds=1.0, ) seen = {} From 2a1c884d8962c6c698b611d930ed95ebffa748bd Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:24:44 +0400 Subject: [PATCH 16/25] Describe Stage 12 state datasets and region provenance in the migration doc The Stage 12 worker foundation document said the cache manifest held all 63 files and described two manifest keys. It now describes the per-state datasets the precompute builds and the image places in state folders, the third manifest key that carries region datasets, the dataset provenance the entry service and worker now agree on, and why the bundle digest leaves region datasets out. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 40 ++++++++++++++----- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index fed0d09e4..86f1f2ac8 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -376,22 +376,42 @@ baselines, so running again after a failed check repeats the check rather than skipping it. A pass is recorded in the cache bucket, and the manifest is published only when that record exists. -Precompute publishes a canonical manifest containing all 63 files. The -deployment passes that manifest's SHA-256 digest into the US image build. The -image build downloads and verifies every listed file, then records the exact -cache release in worker-validation metadata and the v2 version manifest. The UK +Precompute also builds per-state datasets. For each US state region whose +requests load a dataset other than the bundle default (the ACS-local dataset +since policyengine 6.7), it prepares that dataset's year file for each +precomputed year once, keeps one state's households with the region's own row +filter, and stores each state's file under the year file's own name. It +refuses data in which a household's district code names another state's +district, because district runs read their state's file. + +Precompute publishes a canonical manifest containing the 63 national files and +the state datasets. The deployment passes that manifest's SHA-256 digest into +the US image build. The image build downloads and verifies every listed file, +placing each state dataset in `state-datasets///` under the +data folder, then records the exact cache release in worker-validation +metadata and the v2 version manifest. A state or district calculation reads +its state's file when the image holds it for the requested year and .py would +reuse it, and prepares the dataset during the calculation otherwise. The UK and coordinator images do not contain these files. Deployment validation runs one real 2026 region-group baseline and requires it to load the precomputed file without recalculation. Readers released before cache releases reject fields they do not know, and the manifest is published before the Cloud Run revision that reads it is -promoted. The publisher therefore writes the manifest under the -`active-with-cache` key and the same manifest without cache releases under the -`active` key, which those readers keep reading. Current readers read -`active-with-cache`, falling back to `active` only before the first publish -that writes both keys. Removing the `active` copy needs a later change, once -no deployed or rollback revision reads it. +promoted. The publisher therefore writes the whole manifest under the +`active-with-region-datasets` key, the same manifest without each country +bundle's `region_datasets` under the `active-with-cache` key, and without cache +releases either under the `active` key; readers released before each field keep +reading the key that lacks it. Current readers read the newest key a publish has +written. Removing the older copies needs a later change, once no deployed or +rollback revision reads them. + +`region_datasets` names the dataset each region type loads. The entry service +records that dataset as a state or district run's dataset provenance, and the +worker refuses a calculation whose provenance names another dataset than the one +it loads, except that it still accepts the default dataset from entry services +released before regional provenance. The bundle digest leaves `region_datasets` +out, so adding it did not change the digest of a running application. Only an empty-policy baseline using the bundle's default dataset and an exact region-group scope can read a baseline cache file. Reform calculations always From ded3296e1fa80826ba61d857580900795caa7af1 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:24:44 +0400 Subject: [PATCH 17/25] Test the worker's cache release report and UK provenance with region datasets The cache release a US worker reports is now built by _cache_release_summary, and a test shows it still validates against Stage12CacheRelease, which pins the national counts and forbids other fields, when the installed manifest also lists state datasets. Another test shows UK regions keep requiring the default dataset when the bundle names dataset per region type. Co-Authored-By: Claude Opus 5.5 --- .../stage12_worker_validation.py | 13 +++++++- .../tests/test_stage12_runtime.py | 33 +++++++++++++++++++ .../tests/test_stage12_worker_validation.py | 17 +++++++++- 3 files changed, 61 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index 992d1d049..102593c99 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -205,8 +205,19 @@ def _check_installed_cache(expected_manifest_sha256: str) -> dict[str, Any]: if calculation.cache_outcome != "hit": msg = "deployed Stage 12 cache probe did not load a cache hit" raise RuntimeError(msg) + return _cache_release_summary(manifest, expected_manifest_sha256) + + +def _cache_release_summary( + manifest: CacheManifest, manifest_sha256: str +) -> dict[str, Any]: + """Return the cache release a worker reports, in ``Stage12CacheRelease`` form. + + Only national datasets and baselines are counted: the release model pins + those counts and forbids other fields, so state datasets are not reported. + """ return { - "manifest_sha256": expected_manifest_sha256, + "manifest_sha256": manifest_sha256, "partition_sha256": manifest.partition_sha256, "years": list(manifest.years), "segment_count": len(US_REGION_GROUPS), diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index fde76eb48..fc67261b3 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -854,6 +854,39 @@ def test_installed_bundle_accepts_the_dataset_a_region_loads( worker._require_installed_bundle(_simulation_for(region, provenance())) +def test_uk_regions_keep_the_default_dataset_when_the_bundle_names_region_types( + monkeypatch, +) -> None: + from policyengine_simulation_contract.stage12_bundle import Stage12RegionDataset + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + resolved = _installed_bundle_with_regions() + us = resolved.bundle.countries[0] + uk = us.model_copy( + update={ + "country": "uk", + "region_datasets": ( + Stage12RegionDataset( + region_type="national", dataset=us.default_dataset + ), + ), + } + ) + resolved = resolved.model_copy( + update={"bundle": resolved.bundle.model_copy(update={"countries": (us, uk)})} + ) + monkeypatch.setattr(worker, "load_stage12_bundle", lambda: resolved) + simulation = _simulation_for("us", _bundle()).model_copy( + update={"geography": GeographySelection(country="uk", region="country/wales")} + ) + + worker._require_installed_bundle(simulation) + with pytest.raises(RuntimeError, match="differs from the installed bundle"): + worker._require_installed_bundle( + simulation.model_copy(update={"bundle": _regional_provenance()}) + ) + + @pytest.mark.parametrize( ("region", "dataset"), [ diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index fc104c9df..eac3dde27 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -459,7 +459,7 @@ def entry(kind, year, index): "schema": "stage12-mf1", "country": "us", "years": list(PRECOMPUTE_YEARS), - "partition_sha256": "p" * 64, + "partition_sha256": "d" * 64, "bundle": { "bundle_manifest_sha256": "b" * 64, "policyengine_version": "5.2.0", @@ -519,3 +519,18 @@ def test_expected_state_file_count_follows_the_installed_bundle(monkeypatch) -> monkeypatch.setattr(precompute, "state_dataset_targets", lambda: ("UT", "WY")) assert module.expected_state_dataset_count() == 2 * len(precompute.PRECOMPUTE_YEARS) + + +def test_the_reported_cache_release_still_matches_its_contract_with_state_files( + tmp_path, +) -> None: + from policyengine_simulation_contract.stage12_manifest import Stage12CacheRelease + from policyengine_simulation_executor import stage12_worker_validation as module + + manifest = _installed_manifest() + + summary = module._cache_release_summary(manifest, "c" * 64) + + release = Stage12CacheRelease.model_validate(summary) + assert release.dataset_count == 3 + assert release.baseline_count == 60 From a79487ea02dcc43021636b152f1e88010590f73d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:40:32 +0400 Subject: [PATCH 18/25] Read the region dataset manifest copy only while it is current Code released before region datasets publishes only the active-with-cache and active keys, so after a rollback the active-with-region-datasets copy is left out of date. Entry services from this change would have kept reading it, sending ACS-local provenance to the older workers, which refuse it, and the next publish would have dropped the worker versions the older code added. Readers and the publisher now use that copy only when, with its region datasets removed, it equals the manifest under active-with-cache; otherwise they read active-with-cache. Co-Authored-By: Claude Opus 5.5 --- .../stage12_manifest.py | 38 ++++++++++++++----- .../tests/test_stage12_manifest.py | 28 ++++++++++++++ 2 files changed, 56 insertions(+), 10 deletions(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index a924a24ba..afeddea98 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -255,18 +255,36 @@ def pre_cache_manifest_payload(manifest: V2VersionManifest) -> dict[str, object] ) +def _is_current_region_dataset_copy(value: object, cache_copy: object) -> bool: + """Return whether ``value`` is the copy the last publish wrote with ``cache_copy``. + + Code released before region datasets publishes only the older keys, for + example after a rollback, which leaves an out-of-date region dataset copy + behind. A current copy, with its region datasets removed, is exactly the + manifest under ``active-with-cache``. + """ + if value is None: + return False + if cache_copy is None: + return True + try: + manifest = V2VersionManifest.model_validate(value) + except (TypeError, ValueError): + return False + return pre_region_datasets_manifest_payload(manifest) == cache_copy + + def _current_manifest(store: V2ManifestStore) -> object: - # The newest key a publish has written yet: each older key holds the same + # The newest key the last publish wrote: each older key holds the same # manifest without the fields its readers do not accept. - for key in ( - REGION_DATASETS_MANIFEST_KEY, - ACTIVE_MANIFEST_KEY, - PRE_CACHE_MANIFEST_KEY, - ): - current = store.get(key) - if current is not None: - return current - return None + cache_copy = store.get(ACTIVE_MANIFEST_KEY) + region_dataset_copy = store.get(REGION_DATASETS_MANIFEST_KEY) + if _is_current_region_dataset_copy(region_dataset_copy, cache_copy): + return region_dataset_copy + if cache_copy is not None: + return cache_copy + # Until the first publish that writes the cache copy. + return store.get(PRE_CACHE_MANIFEST_KEY) def build_next_v2_manifest( diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 0337a8cc2..f3bf51167 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -449,3 +449,31 @@ def test_a_country_bundle_gives_the_dataset_a_region_code_loads() -> None: assert us.dataset_for_region("us") == "populace_us_2024" assert us.dataset_for_region("place/UT-67000") is None assert uk.dataset_for_region("country/scotland") is None + + +def test_a_region_dataset_copy_left_behind_by_an_older_publish_is_not_read() -> None: + """After a rollback, older code republishes only the keys it knows.""" + store = FakeStore() + publish_v2_manifest(store=store, worker=_regional_worker("5.2.0")) + stale = store.values[REGION_DATASETS_MANIFEST_KEY] + # A publish by code released before region datasets: it writes the two + # older keys only, leaving the region dataset copy as it was. + older = publish_v2_manifest(store=FakeStore(), worker=_cached_worker("5.3.0")) + store.values[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(older) + store.values[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(older) + + loaded = V2ManifestLoader(store).load() + + assert loaded.default_version == "5.3.0" + assert store.values[REGION_DATASETS_MANIFEST_KEY] == stale + + +def test_a_publish_after_a_rollback_keeps_the_versions_the_older_code_added() -> None: + store = FakeStore() + publish_v2_manifest(store=store, worker=_regional_worker("5.2.0")) + older = publish_v2_manifest(store=FakeStore(), worker=_cached_worker("5.3.0")) + store.values[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(older) + + manifest = publish_v2_manifest(store=store, worker=_regional_worker("5.4.0")) + + assert sorted(manifest.versions) == ["5.3.0", "5.4.0"] From 0613893687d05882e5371093e7454d9b5d1f9dfc Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:41:48 +0400 Subject: [PATCH 19/25] Test that the contract and the runtime read region codes alike region_type_of restates the runtime's region normalisation, so an executor test now runs both on the same codes (US and UK, bare codes, mixed case, surrounding spaces, prefixed codes). It found one difference: a bare "uk" requested as a US region was read as a state; another country's code now names no region type, as at run time. Co-Authored-By: Claude Opus 5.5 --- .../stage12_bundle.py | 5 ++- .../tests/test_stage12_manifest.py | 2 + .../tests/test_stage12_runtime.py | 41 +++++++++++++++++++ 3 files changed, 47 insertions(+), 1 deletion(-) diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py index 90dac9bb3..2270497e4 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py @@ -27,11 +27,14 @@ def region_type_of(country: str, region: str | None) -> str | None: raw = (region or "").strip() if not raw or raw.lower() == country: return "national" + if raw.lower() in SUPPORTED_COUNTRIES: + # Another country's code names no region of this one. + return None if "/" in raw: return raw.split("/", 1)[0].lower() if country == "us" and len(raw) == _US_STATE_CODE_LENGTH: return "state" - if country == "uk" and raw.lower() not in SUPPORTED_COUNTRIES: + if country == "uk": return "country" return None diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index f3bf51167..78c1f908e 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -428,6 +428,8 @@ def test_storage_names_must_be_distinct() -> None: ("us", "CONGRESSIONAL_DISTRICT/UT-01", "congressional_district"), ("us", "place/UT-67000", "place"), ("us", "utah", None), + ("us", "uk", None), + ("uk", "us", None), ("uk", "uk", "national"), ("uk", "country/scotland", "country"), ("uk", "scotland", "country"), diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index fc67261b3..83ee55cf0 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -4510,3 +4510,44 @@ def setup_that_uses_the_whole_window(_report): assert [len(call.probed_at) for call in invoker.calls.values()] == [1, 1] assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] assert store.parent.error_code == "simulation_wait_timeout" + + +def _runtime_region_type(country: str, region: str | None) -> str | None: + """Return the region type the runtime's own region normalisation implies.""" + from policyengine_simulation_executor.simulation_runtime import ( + _normalise_region_code, + ) + + code = _normalise_region_code(country, region) + if code == country: + return "national" + return code.split("/", 1)[0] if "/" in code else None + + +@pytest.mark.parametrize( + ("country", "region"), + [ + ("us", None), + ("us", ""), + ("us", " US "), + ("us", "uk"), + ("us", "state/ut"), + ("us", "State/UT"), + ("us", "ut"), + ("us", " ut "), + ("us", "congressional_district/ut-01"), + ("us", "place/UT-67000"), + ("us", "utah"), + ("uk", "uk"), + ("uk", "us"), + ("uk", "country/scotland"), + ("uk", "scotland"), + ("uk", "Northern Ireland"), + ("uk", "constituency/Aldershot"), + ("uk", "local_authority/E06000001"), + ], +) +def test_the_contract_reads_region_codes_as_the_runtime_does(country, region) -> None: + from policyengine_simulation_contract.stage12_bundle import region_type_of + + assert region_type_of(country, region) == _runtime_region_type(country, region) From bde5b00d34f7323b853c764b2634049f0b5592f6 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:42:42 +0400 Subject: [PATCH 20/25] Map a region type to the default dataset when another dataset shares its file When two bundle datasets have the same file path, region_datasets took the identity first in sorted order, so an alias listed for the default dataset's file would have become the national region's dataset and workers would then refuse national runs naming the default. The default dataset now owns its path; other collisions still take the first identity in sorted order. Co-Authored-By: Claude Opus 5.5 --- .../stage12_bundle/parsing.py | 5 ++++- .../tests/test_stage12_bundle.py | 12 ++++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py index 942fa7ef3..047330999 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py @@ -196,7 +196,10 @@ def normalize_country_bundle( dataset.get("sha256"), f"data_releases.{country}.datasets.{identity}.sha256", ) - identity_by_path.setdefault(path, identity) + # The default dataset owns its path; otherwise the first identity + # in sorted order does. + if identity == default_dataset or path not in identity_by_path: + identity_by_path[path] = identity datasets.append( Stage12Dataset( identity=identity, diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py index 7fc624d6e..033d06030 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py @@ -191,3 +191,15 @@ def test_region_datasets_leave_the_bundle_digest_unchanged() -> None: assert _us(bundle).region_datasets is not None assert bundle_digest(bundle) == bundle_digest(without_region_datasets) + + +def test_a_dataset_sharing_the_default_path_does_not_replace_the_default() -> None: + bundle = _bundle() + release = bundle["data_releases"]["us"] + release["datasets"]["a_alias_of_the_default"] = deepcopy( + release["datasets"][release["default_dataset"]] + ) + + us = _us(normalize_stage12_bundle(bundle)) + + assert us.dataset_for_region_type("national") == us.default_dataset From 086d0dfcd3c8ca2f4fab8156bc9970dbaf1aee91 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:43:20 +0400 Subject: [PATCH 21/25] Pin the Stage 12 bundle digest of a fixed bundle A test only showed that the digest ignores region_datasets, so the next field added to the bundle could change every digest unnoticed, refusing comparison runs during a deploy and rebuilding the Stage 12 cache. The digest of a fixed bundle is now pinned to the value the code before region datasets gives it. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_stage12_bundle.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py index 033d06030..8bd6a03f6 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py @@ -203,3 +203,57 @@ def test_a_dataset_sharing_the_default_path_does_not_replace_the_default() -> No us = _us(normalize_stage12_bundle(bundle)) assert us.dataset_for_region_type("national") == us.default_dataset + + +def test_the_bundle_digest_of_a_fixed_bundle_is_pinned() -> None: + """Any change to what the digest covers is deliberate. + + A changed digest refuses comparison runs while a redeployed v2 application + waits for its version manifest, and rebuilds the Stage 12 cache, so a new + bundle field must not reach it by accident. The pinned value is the one + the code before region datasets gives the same bundle. + """ + from policyengine_simulation_contract.stage12_bundle import ( + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, + Stage12RegionDataset, + ) + + dataset = Stage12Dataset( + identity="populace_us_2024", + uri="hf://policyengine/populace-us/populace_us_2024.h5@revision", + artifact_revision="revision", + sha256="a" * 64, + repo_type="dataset", + ) + bundle = Stage12BundleManifest( + policyengine_version="6.7.5", + policyengine_requirement="policyengine==6.7.5", + core_package_version="3.33.0", + core_package_requirement="policyengine-core==3.33.0", + countries=( + Stage12CountryBundle( + country="us", + country_package_name="policyengine-us", + country_package_version="2.2.1", + country_package_requirement="policyengine-us==2.2.1", + data_package_name="populace-us", + data_package_version="1.0.0", + data_release_version="revision", + data_artifact_revision="revision", + default_dataset=dataset.identity, + default_dataset_uri=dataset.uri, + datasets=(dataset,), + region_datasets=( + Stage12RegionDataset( + region_type="national", dataset=dataset.identity + ), + ), + ), + ), + ) + + assert bundle_digest(bundle) == ( + "47f3f7c69ea3b99005ffee585092074862bf4f55a91a95511f12924a79aac71c" + ) From 8babec84e5e9182dce5080e15aaac7b3349dcbe7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:48:11 +0400 Subject: [PATCH 22/25] Load on demand in Stage 12 when .py's reuse checks cannot be imported Stage 12's copy of the baked-folder lookup now logs a warning and returns no folder when .py's year-file reuse checks, which import a private .py helper, cannot be imported after an upgrade, so state and district calculations prepare their dataset instead of failing. Co-Authored-By: Claude Opus 5.5 --- .../stage12_cache/state_datasets.py | 12 +++++++++++- .../tests/test_stage12_state_datasets.py | 12 ++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py index 98bc21010..7633041f2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py @@ -85,7 +85,17 @@ def baked_state_data_folder( year_file = Path(folder) / f"{stem}_year_{year}.h5" if not year_file.is_file(): return None - problem = year_file_reuse_problem(year_file, selection.name) + try: + problem = year_file_reuse_problem(year_file, selection.name) + except ImportError: + # The reuse checks read private .py helpers. If an upgrade moves one, + # the calculation prepares its dataset instead of failing. + logger.warning( + "Not reading the baked state file %s: .py's year-file reuse " + "checks could not be imported", + year_file, + ) + return None if problem is not None: logger.warning("Not reading the baked state file %s: %s", year_file, problem) return None diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py b/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py index ba5b26c96..5b6eef5a2 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py @@ -224,3 +224,15 @@ def test_default_datasets_and_other_regions_read_nothing_baked(tmp_path): assert _lookup(tmp_path, default) is None assert _lookup(tmp_path, _regional_selection(), state_code=None) is None assert _lookup(tmp_path, _regional_selection(), country="uk") is None + + +def test_a_renamed_dotpy_reuse_check_falls_back_to_on_demand_loading( + monkeypatch, tmp_path +): + """A .py upgrade moving the private SPM profile helper must not fail runs.""" + from policyengine.tax_benefit_models.us import datasets as us_datasets + + _bake_utah(tmp_path) + monkeypatch.delattr(us_datasets, "_expected_legacy_spm_profile") + + assert _lookup(tmp_path, _regional_selection()) is None From b1dd135a57815a8974fe376d765d8ef4202950f4 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:48:11 +0400 Subject: [PATCH 23/25] Refuse a Stage 12 state without a row filter before preparing anything Stage 12's copy of the in-container checks now refuses a state target without a row filter before the full-country year file is prepared, and returns each state's row filter. The orchestration docstring says builds run one call per year and dataset. Co-Authored-By: Claude Opus 5.5 --- .../stage12_cache/precompute.py | 29 +++++++++++-------- .../tests/test_stage12_cache_precompute.py | 1 + 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index c44b23498..f0fe1a5ed 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -406,11 +406,12 @@ def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildR def _checked_state_targets( year: int, expected: Sequence[StateDatasetPlanEntry] -) -> dict[str, StateDatasetTarget]: - """Return each planned state's target, after re-deriving its key here. +) -> dict[str, RowFilterStrategy]: + """Return each planned state's row filter, after re-deriving its key here. Raises ``RuntimeError`` when the installed bundle gives another identity - than the plan, so nothing is built under a key that does not describe it. + than the plan, so nothing is built under a key that does not describe it, + and when a state has no row filter, before anything is prepared. """ resolved = load_stage12_bundle() targets = {target.region: target for target in state_dataset_targets()} @@ -426,6 +427,9 @@ def _checked_state_targets( if target is None: msg = f"No state dataset target for {entry.region} in this container" raise RuntimeError(msg) + if target.scoping_strategy is None: + msg = f"{entry.region} has no row filter to build its state file with" + raise RuntimeError(msg) identity = collect_state_dataset_identity( resolved, country="us", @@ -450,7 +454,12 @@ def _checked_state_targets( if len(datasets) > 1: msg = f"One state dataset build covers one dataset, not {sorted(datasets)}" raise ValueError(msg) - return {entry.region: targets[entry.region] for entry in expected} + strategies = {} + for entry in expected: + strategy = targets[entry.region].scoping_strategy + if strategy is not None: + strategies[entry.region] = strategy + return strategies def _registered_district_geoids(country_module: ModuleType) -> set[int]: @@ -516,7 +525,7 @@ def build_state_datasets_impl( if not expected: return () started = time.monotonic() - targets = _checked_state_targets(year, expected) + strategies = _checked_state_targets(year, expected) dataset = expected[0].dataset work_folder = work_folder or Path( tempfile.mkdtemp(prefix="policyengine-stage12-state-datasets-") @@ -544,11 +553,7 @@ def build_state_datasets_impl( store = Stage12CacheStore(bucket) results = [] for entry in expected: - strategy = targets[entry.region].scoping_strategy - if strategy is None: - msg = f"{entry.region} has no row filter to build its state file with" - raise RuntimeError(msg) - rows = strategy.apply( + rows = strategies[entry.region].apply( entity_data=source.data.entity_data, group_entities=US_GROUP_ENTITIES, year=year, @@ -775,8 +780,8 @@ def run_precompute( Missing datasets are built first, then missing baselines, in parallel Modal calls. Missing state datasets are built alongside both, one call per - year, and waited for only before the manifest is published, so the - baselines never wait for them. Prints the manifest digest through + year and dataset, and waited for only before the manifest is published, so + the baselines never wait for them. Prints the manifest digest through ``echo`` and returns it. Exits with ``SystemExit`` when the determinism check fails. """ diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py index 7db8837b1..3291fb84d 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -1241,6 +1241,7 @@ def test_refuses_a_target_without_a_row_filter(self, state_build): with pytest.raises(RuntimeError, match="no row filter"): self._build(state_build, [self._entry(state_build, "UT")]) assert state_build.client.objects == {} + assert state_build.ensure_calls == [] def test_reports_how_long_the_years_build_had_run(self, state_build): results = self._build( From 077bf4282a020300fbf1e2b034cfb1b02b7327d8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:01:13 +0400 Subject: [PATCH 24/25] Document how manifest readers pick a key and the rollback step The Stage 12 worker foundation document said current readers read the newest key a publish has written. It now gives the actual rule (the region dataset copy only while it matches active-with-cache once region datasets are removed) and the one window it leaves on a rollback to code from before region datasets: until the older code publishes, entry services keep sending regional provenance, which the older worker refuses, failing Stage 12 comparison runs for state and district requests but no production response. The rollback step that avoids it is to delete the active-with-region-datasets key before deploying the older code. A comment in publish_v2_manifest now gives the right reason readers never combine parts of two documents. Co-Authored-By: Claude Opus 5.5 --- .../stage-12-modal-worker-foundation.md | 20 ++++++++++++++++--- .../stage12_manifest.py | 5 +++-- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index 86f1f2ac8..953ca0d70 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -402,9 +402,23 @@ promoted. The publisher therefore writes the whole manifest under the `active-with-region-datasets` key, the same manifest without each country bundle's `region_datasets` under the `active-with-cache` key, and without cache releases either under the `active` key; readers released before each field keep -reading the key that lacks it. Current readers read the newest key a publish has -written. Removing the older copies needs a later change, once no deployed or -rollback revision reads them. +reading the key that lacks it. Current readers use the +`active-with-region-datasets` copy only while it, with its region datasets +removed, equals the manifest under `active-with-cache`; otherwise they read +`active-with-cache`, falling back to `active` before the first publish that +writes it. Code released before region datasets republishes only the two older +keys, so after a rollback the out-of-date region dataset copy stops matching and +is no longer read. Removing the older copies needs a later change, once no +deployed or rollback revision reads them. + +A rollback to code released before region datasets has one window: the older +worker is deployed before its publish job runs, and until that publish the +region dataset copy still matches, so updated entry services keep sending +regional provenance for state and district runs, which the older worker refuses. +Only Stage 12 comparison runs fail in that window, never production responses; +if the older deploy's validation fails, the window lasts until a later publish. +To avoid it, delete the `active-with-region-datasets` key from the +`simulation-api-v2-version-manifest` Modal Dict before deploying the older code. `region_datasets` names the dataset each region type loads. The entry service records that dataset as a state or district run's dataset provenance, and the diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index afeddea98..a10d414cc 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -338,8 +338,9 @@ def publish_v2_manifest( worker=worker, force_default=force_default, ) - # Each key holds one entire validated document, and a reader reads only - # one key, so no reader observes a partial update assembled from several. + # Each key holds one entire validated document, and a reader uses one + # whole document (the region dataset copy only while it matches the cache + # copy), so no reader observes a partial update assembled from several. store[REGION_DATASETS_MANIFEST_KEY] = next_manifest.model_dump(mode="json") store[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(next_manifest) store[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(next_manifest) From 6b23b26997da233d2f6305c4f97fb8e66645545d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:51:47 +0400 Subject: [PATCH 25/25] Pass isin a list in Stage 12's district-code check Pyright reads `Series.isin` as taking a sequence, mapping, Series or DataFrame, so it refused the set of registered district codes. The check now passes them sorted, which keeps the same matches. Co-Authored-By: Claude Opus 5.5 --- .../stage12_cache/precompute.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index f0fe1a5ed..782647cf8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -486,7 +486,7 @@ def _require_districts_within_their_states( if "congressional_district_geoid" not in household.columns: return district = household["congressional_district_geoid"] - in_another_state = district.isin(district_geoids) & ( + in_another_state = district.isin(sorted(district_geoids)) & ( (district // 100) != household["state_fips"] ) if in_another_state.any():