From f5aab4d6b3f3525033ff850ee178cb8ccc08fd0b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 19:56:09 +0400 Subject: [PATCH 01/16] Report the US state code when resolving a state or district region RegionResolution gains state_code: the two-letter code of the US state a state or congressional-district region lies in, taken from the .py region. National, UK and region-group resolutions leave it None. The baked per-state datasets added next are looked up by this code. Co-Authored-By: Claude Opus 5.5 --- .../simulation_runtime.py | 6 +- .../tests/test_simulation_output_builder.py | 129 ++++++++++++++++++ 2 files changed, 134 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 39a30ae20..23529c82a 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 @@ -73,12 +73,15 @@ class RegionResolution: """A request's normalized region code, the dataset to load, and how to scope it. ``scoping_strategy`` narrows the simulation to the region; None simulates - the whole dataset, as a national run does. + the whole dataset, as a national run does. ``state_code`` is the two-letter + code of the US state a state or congressional-district region lies in, and + None for every other region. """ code: str dataset_reference: str | None = None scoping_strategy: Any | None = None + state_code: str | None = None @dataclass(frozen=True) @@ -605,6 +608,7 @@ def _resolve_region( code=region_code, dataset_reference=dataset_reference, scoping_strategy=getattr(region, "scoping_strategy", None), + state_code=getattr(region, "state_code", None) if country == "us" else None, ) 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 9d95eab0b..578a92546 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1618,6 +1618,135 @@ def test_resolve_region_scopes_us_congressional_district_from_national_dataset() assert resolution.scoping_strategy is scoping_strategy +def _us_country_module_with_regions(regions): + return SimpleNamespace(model=SimpleNamespace(get_region=regions.get)) + + +@pytest.mark.parametrize( + ("requested_region", "state_code"), + [ + ("state/ut", "UT"), + ("state/UT", "UT"), + ("ut", "UT"), + ("congressional_district/ut-01", "UT"), + ("congressional_district/WY-00", "WY"), + ], +) +def test_resolve_region_reports_the_us_state_code(requested_region, state_code): + utah = SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="us", + state_code="UT", + ) + regions = { + "us": SimpleNamespace( + dataset_path=get_country_release_bundle("us").default_dataset_uri, + scoping_strategy=None, + parent_code=None, + state_code=None, + ), + "state/ut": utah, + "congressional_district/UT-01": SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="state/ut", + state_code="UT", + ), + "congressional_district/WY-00": SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="state/wy", + state_code="WY", + ), + "state/wy": SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="us", + state_code="WY", + ), + } + + resolution = _resolve_region( + country_module=_us_country_module_with_regions(regions), + country="us", + params={"region": requested_region}, + ) + + assert resolution.state_code == state_code + + +def test_resolve_region_reports_no_state_code_for_a_national_us_run(): + resolution = _resolve_region( + country_module=_us_country_module_with_regions({}), + country="us", + params={"region": "us"}, + ) + + assert resolution.state_code is None + + +def test_resolve_region_reports_no_state_code_for_a_uk_region(): + scotland = SimpleNamespace( + dataset_path=None, + scoping_strategy=object(), + parent_code="uk", + state_code=None, + ) + national = SimpleNamespace( + dataset_path=get_country_release_bundle("uk").default_dataset_uri, + scoping_strategy=None, + parent_code=None, + state_code=None, + ) + country_module = SimpleNamespace( + model=SimpleNamespace( + get_region={"country/scotland": scotland, "uk": national}.get + ) + ) + + resolution = _resolve_region( + country_module=country_module, + country="uk", + params={"region": "country/scotland"}, + ) + + assert resolution.state_code is None + + +def test_resolve_region_reports_no_state_code_for_a_region_group(): + from policyengine.core.scoping_strategy import RowFilterStrategy + + regions = { + code: SimpleNamespace( + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=fips + ), + state_code=code.removeprefix("state/").upper(), + ) + for code, fips in (("state/ut", 49), ("state/wy", 56)) + } + + resolution = _resolve_region( + country_module=_us_country_module_with_regions(regions), + country="us", + params={"region": "us", "region_group": ["state/ut", "state/wy"]}, + ) + + assert resolution.state_code is None + + +def test_every_installed_us_state_and_district_region_names_its_state(): + from policyengine.countries.us.regions import us_region_registry + + for region_type in ("state", "congressional_district"): + regions = us_region_registry.get_by_type(region_type) + assert regions + for region in regions: + assert region.state_code + assert region.state_code == region.state_code.upper() + + def test_resolve_region_rejects_unscoped_us_place_region(): place = SimpleNamespace( dataset_path=None, From c4a72a824f500b7246a9606d04d0f13834670ef2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 19:58:10 +0400 Subject: [PATCH 02/16] Let callers choose the folder _load_dataset reads a year file from _load_dataset takes an optional data_folder that replaces the folder the dataset selection would otherwise use (the image's data folder for the default dataset, a temporary folder for any other). Callers are unchanged; the main path and Stage 12 will each pass the folder of a baked per-state dataset when one exists. Co-Authored-By: Claude Opus 5.5 --- .../simulation_runtime.py | 18 ++++-- .../tests/test_simulation_output_builder.py | 61 +++++++++++++++++++ 2 files changed, 74 insertions(+), 5 deletions(-) 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 23529c82a..add789dfe 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 @@ -725,15 +725,23 @@ def _load_dataset( *, selection: DatasetSelection, country_module: ModuleType | None = None, + data_folder: str | None = None, ) -> Dataset: + """Load the selected dataset's year file, preparing it if the folder lacks it. + + ``data_folder`` replaces the folder the selection would otherwise use: the + image's data folder for the default dataset, or a temporary folder for any + other dataset. + """ country = params.get("country", "us").lower() year = _parse_year(params) country_module = country_module or _country_module(country) - data_folder = ( - resolve_data_folder() - if selection.is_default - else _nondefault_data_folder(country, selection.name, selection.uri) - ) + if data_folder is None: + data_folder = ( + resolve_data_folder() + if selection.is_default + else _nondefault_data_folder(country, selection.name, selection.uri) + ) start = time.monotonic() load_options = {"years": [year], "data_folder": data_folder} 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 578a92546..9306148e4 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1531,6 +1531,67 @@ def test_nondefault_dataset_cache_is_distinct_for_shared_filename_stem(): assert first != second +def test_load_dataset_reads_a_nondefault_dataset_from_the_folder_it_is_given( + tmp_path, + monkeypatch, +): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path / "image-data")) + bundle = get_country_release_bundle("us") + nondefault = next( + name for name in bundle.dataset_uris if name != bundle.default_dataset + ) + selection = DatasetSelection( + name=nondefault, uri=bundle.dataset_uris[nondefault], is_default=False + ) + state_folder = str(tmp_path / "state-datasets" / nondefault / "UT") + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + _load_selected_dataset( + {"country": "us", "time_period": "2026"}, + selection=selection, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + data_folder=state_folder, + ) + + assert ensure_calls == [ + {"datasets": [nondefault], "years": [2026], "data_folder": state_folder} + ] + + +def test_load_dataset_without_a_folder_keeps_the_temporary_nondefault_folder( + tmp_path, + monkeypatch, +): + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + bundle = get_country_release_bundle("us") + nondefault = next( + name for name in bundle.dataset_uris if name != bundle.default_dataset + ) + selection = DatasetSelection( + name=nondefault, uri=bundle.dataset_uris[nondefault], is_default=False + ) + ensure_calls = [] + + def ensure_datasets(**kwargs): + ensure_calls.append(kwargs) + return {"dataset": SimpleNamespace()} + + _load_selected_dataset( + {"country": "us", "time_period": "2026"}, + selection=selection, + country_module=SimpleNamespace(ensure_datasets=ensure_datasets), + data_folder=None, + ) + + assert ensure_calls[0]["data_folder"] == _nondefault_data_folder( + "us", nondefault, bundle.dataset_uris[nondefault] + ) + + def test_load_dataset_uses_baked_folder_for_default_requests( tmp_path, monkeypatch, From c4ca8dd71f1cd2c3b9e65fc9c9365381a2f32035 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:02:47 +0400 Subject: [PATCH 03/16] Key and lay out per-state datasets for the main path's precompute A per-state dataset is a regional dataset's year file filtered to one state's households. Its content key (schema sds1) covers the regional dataset's name, revision-pinned URI and certified SHA-256, the year, the state code and FIPS, the filter variable, the model and .py versions and the SPM selection. Its store path is state-datasets///_year_.h5, keeping the basename ensure_datasets resolves. state_datasets.py gives each state's folder under the image data folder, state-datasets//, so the manifest builder and the runtime lookup agree on where a state's files live. The default dataset's key and identity are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../fixtures/identity_stubs.py | 18 +- .../artifact_keys.py | 176 +++++++++++++++++- .../state_datasets.py | 28 +++ .../tests/test_artifact_keys.py | 84 +++++++++ .../tests/test_state_datasets.py | 19 ++ 5 files changed, 318 insertions(+), 7 deletions(-) create mode 100644 projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py create mode 100644 projects/policyengine-simulation-executor/tests/test_state_datasets.py diff --git a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py index 307a6f5ec..0de437c35 100644 --- a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py +++ b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py @@ -27,9 +27,11 @@ def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: """Stub the identity seams with the canonical synthetic version-set. - Returns mutable bundle, receipt, and SPM-capability state, plus a handle - to the installed capability resolver for tests that intentionally exercise - the installed release. + Returns mutable bundle, receipt, release-manifest dataset and + SPM-capability state, plus a handle to the installed capability resolver + for tests that intentionally exercise the installed release. The bundle + has one non-default dataset, ``populace_cps_local``, for the per-state + dataset identities. """ # This module imports ``get_release_manifest`` by value while constructing # country models. Import it before replacing the manifest module's function @@ -51,8 +53,12 @@ def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: default_dataset="populace_cps", dataset_uris={ "populace_cps": "hf://org/repo/populace_cps.h5@rev-abc", + "populace_cps_local": "hf://org/repo/populace_cps_local.h5@rev-local", }, ) + release_datasets = { + "populace_cps_local": SimpleNamespace(sha256="c0ffee00" * 8), + } receipt_entry = { "country": "us", "version": "1.2.3", @@ -61,6 +67,7 @@ def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: state = SimpleNamespace( bundle=bundle, receipt_entry=receipt_entry, + release_datasets=release_datasets, spm_capability=None, ) @@ -78,13 +85,14 @@ def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: monkeypatch.setattr( manifest_module, "dataset_logical_name", - lambda _reference: state.bundle.default_dataset, + lambda reference: reference.rsplit("/", 1)[-1].split(".h5", 1)[0], ) monkeypatch.setattr( manifest_module, "get_release_manifest", lambda _country: SimpleNamespace( - certification=SimpleNamespace(data_build_fingerprint="fp-123") + certification=SimpleNamespace(data_build_fingerprint="fp-123"), + datasets=state.release_datasets, ), ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py index ee9026334..fdfe2fefd 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py @@ -1,7 +1,9 @@ """Content-addressed identity for precomputed simulation artifacts. -The artifact pipeline caches two artifact types in a shared store: single-year -dataset files and precomputed baseline simulation outputs. An artifact's key is +The artifact pipeline caches three artifact types in a shared store: single-year +dataset files, per-state single-year dataset files (a regional dataset's +year file filtered to one state's households) and precomputed baseline +simulation outputs. An artifact's key is a sha256 digest over the FULL closure of inputs that determine its bytes, so "is it cached?" and "is it stale?" are the same question, answered by key equality: same key implies a byte-identical artifact; any input change rotates @@ -39,6 +41,7 @@ from typing import Any DATASET_KEY_SCHEMA = "ds1" +STATE_DATASET_KEY_SCHEMA = "sds1" BASELINE_KEY_SCHEMA = "bl1" # Length of the digest prefix used in simulation ids and store paths. 16 hex @@ -106,6 +109,51 @@ def dataset_key(inputs: DatasetKeyInputs) -> str: ) +@dataclass(frozen=True) +class StateDatasetKeyInputs: + """Every input that can change one per-state single-year dataset's bytes. + + The file is the regional dataset's year file (built by a full + Microsimulation pass, so model and .py versions matter) filtered to the + households whose ``filter_variable`` equals ``state_fips``. + ``dataset_uri`` pins the source to its Hugging Face revision and + ``source_sha256`` is the bundle's certified hash of that source, which may + be None when the bundle has none. + """ + + country: str + dataset: str + dataset_uri: str + source_sha256: str | None + year: int + state_code: str + state_fips: int + filter_variable: str + model_version: str + policyengine_version: str + spm: dict | None = None + + +def state_dataset_key(inputs: StateDatasetKeyInputs) -> str: + """Digest identifying one per-state single-year dataset artifact.""" + return canonical_digest( + { + "schema": STATE_DATASET_KEY_SCHEMA, + "spm": inputs.spm, + "country": inputs.country.lower(), + "dataset": inputs.dataset, + "dataset_uri": inputs.dataset_uri, + "source_sha256": inputs.source_sha256, + "year": int(inputs.year), + "state_code": inputs.state_code.upper(), + "state_fips": int(inputs.state_fips), + "filter_variable": inputs.filter_variable, + "model_version": inputs.model_version, + "policyengine_version": inputs.policyengine_version, + } + ) + + @dataclass(frozen=True) class BaselineKeyInputs: """Every input that can change one precomputed baseline artifact's bytes. @@ -157,6 +205,7 @@ def baseline_simulation_id(baseline_digest: str) -> str: # --- Store layout --------------------------------------------------------- # # datasets///_year_.h5 +# state-datasets///_year_.h5 # baselines///.h5 # manifests/.json # deployed/.json @@ -187,6 +236,20 @@ def dataset_artifact_path( ) +def state_dataset_artifact_path( + country: str, state_dataset_digest: str, stem: str, year: int +) -> str: + """Return the store key ``state-datasets///_year_.h5``. + + Every state's file has the same basename, the one the runtime's + ``ensure_datasets`` resolves; the image keeps them apart by folder. + """ + return ( + f"state-datasets/{country.lower()}/{state_dataset_digest}/" + f"{dataset_artifact_filename(stem, year)}" + ) + + def baseline_artifact_path( country: str, baseline_digest: str, simulation_id: str ) -> str: @@ -255,6 +318,55 @@ def store_path(self) -> str: return dataset_artifact_path(self.country, self.digest, self.stem, self.year) +@dataclass(frozen=True) +class StateDatasetArtifactIdentity: + """Everything needed to key and name one per-state single-year dataset.""" + + country: str + dataset: str + dataset_uri: str + source_sha256: str | None + stem: str + year: int + state_code: str + state_fips: int + filter_variable: str + model_version: str + policyengine_version: str + spm: dict | None = None + + @property + def digest(self) -> str: + """Content key over every input that can change this file's bytes.""" + return state_dataset_key( + StateDatasetKeyInputs( + country=self.country, + dataset=self.dataset, + dataset_uri=self.dataset_uri, + source_sha256=self.source_sha256, + year=self.year, + state_code=self.state_code, + state_fips=self.state_fips, + filter_variable=self.filter_variable, + model_version=self.model_version, + policyengine_version=self.policyengine_version, + spm=self.spm, + ) + ) + + @property + def filename(self) -> str: + """Basename the runtime resolves inside the state's folder.""" + return dataset_artifact_filename(self.stem, self.year) + + @property + def store_path(self) -> str: + """Store key of this per-state dataset artifact.""" + return state_dataset_artifact_path( + self.country, self.digest, self.stem, self.year + ) + + @dataclass(frozen=True) class BaselineArtifactIdentity: """Everything needed to key, id, and name one baseline artifact.""" @@ -377,6 +489,66 @@ def collect_dataset_identity(country: str, year: int) -> DatasetArtifactIdentity ) +def _release_dataset_sha256(country: str, dataset: str) -> str | None: + """Certified SHA-256 of a bundle dataset, from the .py release manifest.""" + from policyengine.provenance.manifest import ( + get_release_manifest, + ) + + reference = getattr(get_release_manifest(country), "datasets", {}).get(dataset) + digest = getattr(reference, "sha256", None) + return digest.lower() if isinstance(digest, str) and digest else None + + +def collect_state_dataset_identity( # noqa: PLR0913 # one keyword per input of the per-state key + country: str, + year: int, + *, + dataset: str, + state_code: str, + state_fips: int, + filter_variable: str, +) -> StateDatasetArtifactIdentity: + """Identity of one state's year file filtered from a bundle dataset. + + ``dataset`` names a non-default bundle dataset, such as the regional + dataset state regions load. Versions and the revision-pinned URI come from + the release bundle, the source hash from the .py release manifest, and the + stem through the same manifest helpers the runtime lookup uses. + """ + from policyengine.provenance.manifest import ( + dataset_logical_name, + resolve_dataset_reference, + ) + + from policyengine_simulation_executor.release_bundle import ( # noqa: PLC0415 # identity seam: resolved per call so tests can patch release_bundle + get_country_release_bundle, + ) + from policyengine_simulation_executor.spm import ( + normalize_runtime_spm, + ) + + bundle = get_country_release_bundle(country) + dataset_uri = bundle.dataset_uris.get(dataset) + if dataset_uri is None: + msg = f"Dataset {dataset!r} is not in the {country} release bundle" + raise ValueError(msg) + return StateDatasetArtifactIdentity( + country=bundle.country, + dataset=dataset, + dataset_uri=dataset_uri, + source_sha256=_release_dataset_sha256(bundle.country, dataset), + stem=dataset_logical_name(resolve_dataset_reference(bundle.country, dataset)), + year=int(year), + state_code=state_code.upper(), + state_fips=int(state_fips), + filter_variable=filter_variable, + model_version=bundle.model_version, + policyengine_version=bundle.policyengine_version, + spm=normalize_runtime_spm({"country": country, "time_period": year}), + ) + + def collect_baseline_identity( country: str, year: int, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py new file mode 100644 index 000000000..e43f9929a --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py @@ -0,0 +1,28 @@ +"""Where the main path's worker image keeps baked per-state datasets. + +State and congressional-district regions load a regional dataset and keep +one state's rows. The precompute writes each state's filtered year file and +the image places it at ``/state-datasets///``, +under the basename the dataset's own year file has. Every state therefore +gets its own folder, and ``ensure_datasets`` called on that folder finds the +state's file where it would look for the full year file. + +The manifest builder and the runtime lookup both take folders from here, so +the folder an artifact is copied into is the folder a request reads. +""" + +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 image 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_artifact_keys.py b/projects/policyengine-simulation-executor/tests/test_artifact_keys.py index 2e4a9dc24..a0cbadddf 100644 --- a/projects/policyengine-simulation-executor/tests/test_artifact_keys.py +++ b/projects/policyengine-simulation-executor/tests/test_artifact_keys.py @@ -49,6 +49,23 @@ _BASELINE_GOLDEN = "f9cac05d94509895eb48ea33ea5f72d40eb9b954989df05ca6506009d8dedae2" +_STATE_DATASET_KWARGS = { + "country": "us", + "dataset": "populace_cps_local", + "dataset_uri": "hf://org/repo/populace_cps_local.h5@rev-local", + "source_sha256": "c0ffee00" * 8, + "year": 2026, + "state_code": "UT", + "state_fips": 49, + "filter_variable": "state_fips", + "model_version": "9.9.9", + "policyengine_version": "4.22.0", +} + +_STATE_DATASET_GOLDEN = ( + "e1f1f9f7573eb090ef887967a49314c1cbf0fb20c8021779bf3804b26b24117f" +) + class TestDigests: def test_dataset_key_golden(self): @@ -82,6 +99,27 @@ def test_canonical_digest_ignores_insertion_order(self): backward = {"b": {"d": None, "c": 2}, "a": 1} assert ak.canonical_digest(forward) == ak.canonical_digest(backward) + def test_state_dataset_key_golden(self): + assert ( + ak.state_dataset_key(ak.StateDatasetKeyInputs(**_STATE_DATASET_KWARGS)) + == _STATE_DATASET_GOLDEN + ) + + @pytest.mark.parametrize("field", sorted(_STATE_DATASET_KWARGS)) + def test_every_state_dataset_field_perturbs_digest(self, field): + replacement = {"year": 2027, "state_fips": 56}.get(field, "other") + perturbed = {**_STATE_DATASET_KWARGS, field: replacement} + assert ( + ak.state_dataset_key(ak.StateDatasetKeyInputs(**perturbed)) + != _STATE_DATASET_GOLDEN + ) + + def test_state_dataset_spm_selection_is_part_of_the_key(self): + with_spm = ak.StateDatasetKeyInputs( + **_STATE_DATASET_KWARGS, spm={"forecast": "x"} + ) + assert ak.state_dataset_key(with_spm) != _STATE_DATASET_GOLDEN + def test_simulation_id_shape(self): sim_id = ak.baseline_simulation_id(_BASELINE_GOLDEN) assert sim_id == "bl1-f9cac05d94509895" @@ -99,6 +137,12 @@ def test_baseline_path(self): path = ak.baseline_artifact_path("us", "b" * 64, "bl1-abc") assert path == f"baselines/us/{'b' * 64}/bl1-abc.h5" + def test_state_dataset_path(self): + path = ak.state_dataset_artifact_path( + "US", "s" * 64, "populace_cps_local", 2026 + ) + assert path == (f"state-datasets/us/{'s' * 64}/populace_cps_local_year_2026.h5") + def test_manifest_and_marker_paths(self): assert ak.manifest_path("m" * 64) == f"manifests/{'m' * 64}.json" assert ak.deployed_marker_path("beta") == "deployed/beta.json" @@ -365,3 +409,43 @@ def test_spm_arm_matches_the_installed_wrapper( tax_benefit_model_version=SimpleNamespace(country_code="us"), ) assert wrapper.storage_id == built.storage_id + + +class TestStateDatasetIdentityCollection: + def test_collected_identity_matches_golden(self, stub_identity_sources): + identity = ak.collect_state_dataset_identity( + "us", + 2026, + dataset="populace_cps_local", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + ) + + assert identity.digest == _STATE_DATASET_GOLDEN + assert identity.filename == "populace_cps_local_year_2026.h5" + assert identity.store_path == ( + f"state-datasets/us/{_STATE_DATASET_GOLDEN}/populace_cps_local_year_2026.h5" + ) + + def test_a_dataset_without_a_certified_sha_is_a_distinct_identity( + self, stub_identity_sources + ): + stub_identity_sources.release_datasets["populace_cps_local"] = SimpleNamespace( + sha256=None + ) + + identity = ak.collect_state_dataset_identity( + "us", + 2026, + dataset="populace_cps_local", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + ) + + assert identity.source_sha256 is None + assert identity.digest != _STATE_DATASET_GOLDEN + + def test_the_default_dataset_identity_is_unchanged(self, stub_identity_sources): + assert ak.collect_dataset_identity("us", 2026).digest == _DATASET_GOLDEN diff --git a/projects/policyengine-simulation-executor/tests/test_state_datasets.py b/projects/policyengine-simulation-executor/tests/test_state_datasets.py new file mode 100644 index 000000000..dae2ea7d5 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_state_datasets.py @@ -0,0 +1,19 @@ +"""The main path's folder layout for baked per-state datasets.""" + +from policyengine_simulation_executor import state_datasets + + +def test_state_dataset_relative_folder_names_the_dataset_and_the_state(): + assert ( + state_datasets.state_dataset_relative_folder("populace_us_2024_acs_local", "ut") + == "state-datasets/populace_us_2024_acs_local/UT" + ) + + +def test_state_dataset_folder_sits_under_the_data_folder(): + assert ( + state_datasets.state_dataset_folder( + "/opt/policyengine/data", "populace_us_2024_acs_local", "UT" + ) + == "/opt/policyengine/data/state-datasets/populace_us_2024_acs_local/UT" + ) From 842b58e428576a1d48459d151756831ebc0c29a8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:06:57 +0400 Subject: [PATCH 04/16] Plan a per-state dataset for each state that loads a regional dataset state_dataset_targets resolves every US state region exactly as a request for it is resolved and keeps the states whose dataset is not the bundle default, with their row filter (state_fips) 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, runtime filename, image folder and store presence; select_work keeps the missing ones. A bundle that sends states back to the default dataset plans none. Nothing builds or publishes these entries yet. Co-Authored-By: Claude Opus 5.5 --- .../precompute.py | 104 +++++++- .../precompute_models.py | 20 ++ .../tests/test_precompute.py | 229 +++++++++++++++++- 3 files changed, 351 insertions(+), 2 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 2fd267cd6..671c44105 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -53,8 +53,12 @@ ManifestArtifact, PrecomputePlan, RemoteFunction, + StateDatasetPlanEntry, WorkSelection, ) +from policyengine_simulation_executor.state_datasets import ( + state_dataset_relative_folder, +) if TYPE_CHECKING: from collections.abc import Callable @@ -152,11 +156,77 @@ def cohort_identity(year: int, group: list[str]) -> BaselineArtifactIdentity: return identity +@dataclass(frozen=True) +class StateDatasetTarget: + """A US state whose requests load a non-default dataset, and its row filter. + + The precompute bakes this state's rows of ``dataset`` as their own year + file: the households whose ``filter_variable`` equals ``state_fips``. + """ + + region: str + state_code: str + state_fips: int + filter_variable: str + dataset: str + + +def state_dataset_targets() -> list[StateDatasetTarget]: + """Return every US state region that loads a non-default dataset, by state code. + + Each state is resolved exactly as a request for it is (``_resolve_region`` + then ``_resolve_dataset_selection``), so a bundle that sends states back to + the default dataset gets no state files. + """ + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _country_module, + _resolve_dataset_selection, + _resolve_region, + ) + + country_module = _country_module(PRECOMPUTE_COUNTRY) + targets = [] + for region in country_module.model.region_registry.get_by_type("state"): + params = { + "country": PRECOMPUTE_COUNTRY, + "region": region.code, + "time_period": PRECOMPUTE_YEARS[0], + } + resolution = _resolve_region( + country_module=country_module, country=PRECOMPUTE_COUNTRY, 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, + ) + ) + return sorted(targets, key=lambda target: target.state_code) + + def select_work(plan: PrecomputePlan, *, force: bool) -> WorkSelection: """Which plan entries need computing.""" return WorkSelection( datasets=[entry for entry in plan.datasets if force or not entry.exists], baselines=[entry for entry in plan.baselines if force or not entry.exists], + state_datasets=[ + entry for entry in plan.state_datasets if force or not entry.exists + ], ) @@ -193,6 +263,7 @@ def plan_artifacts_impl(bucket: str) -> PrecomputePlan: """Compute every expected artifact identity and its store presence.""" from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_dataset_identity collect_dataset_identity, + collect_state_dataset_identity, ) from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore ArtifactStore, @@ -210,8 +281,10 @@ def plan_artifacts_impl(bucket: str) -> PrecomputePlan: msg = f"No national partition for {PRECOMPUTE_COUNTRY!r}" raise RuntimeError(msg) + targets = state_dataset_targets() datasets: list[DatasetPlanEntry] = [] baselines: list[BaselinePlanEntry] = [] + state_datasets: list[StateDatasetPlanEntry] = [] for year in PRECOMPUTE_YEARS: dataset_identity = collect_dataset_identity(PRECOMPUTE_COUNTRY, year) datasets.append( @@ -236,6 +309,30 @@ def plan_artifacts_impl(bucket: str) -> PrecomputePlan: exists=store.exists(identity.store_path), ) ) + for target in targets: + state_identity = collect_state_dataset_identity( + PRECOMPUTE_COUNTRY, + 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), + ) + ) bundle = get_country_release_bundle(PRECOMPUTE_COUNTRY) receipt = BundleVersionIdentity( @@ -245,7 +342,12 @@ def plan_artifacts_impl(bucket: str) -> PrecomputePlan: data_artifact_revision=bundle.data_artifact_revision, default_dataset=bundle.default_dataset, ) - return PrecomputePlan(datasets=datasets, baselines=baselines, receipt=receipt) + return PrecomputePlan( + datasets=datasets, + baselines=baselines, + receipt=receipt, + state_datasets=state_datasets, + ) def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildResult: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py index 1807b7244..c736ede88 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py @@ -59,6 +59,24 @@ class DatasetPlanEntry(_StrictModel): exists: bool +class StateDatasetPlanEntry(_StrictModel): + """One state's year file filtered from a regional dataset: identity plus presence. + + ``folder`` is where the image places the file, relative to its data folder; + every state's file has the same ``filename``. + """ + + year: int + state_code: str + region: str + dataset: str + digest: str + path: str + filename: str + folder: str + exists: bool + + class BaselinePlanEntry(_StrictModel): """One cohort baseline artifact: identity plus store presence.""" @@ -77,6 +95,7 @@ class PrecomputePlan(_StrictModel): datasets: list[DatasetPlanEntry] baselines: list[BaselinePlanEntry] receipt: BundleVersionIdentity + state_datasets: list[StateDatasetPlanEntry] = Field(default_factory=list) class WorkSelection(_StrictModel): @@ -84,6 +103,7 @@ class WorkSelection(_StrictModel): datasets: list[DatasetPlanEntry] baselines: list[BaselinePlanEntry] + state_datasets: list[StateDatasetPlanEntry] = Field(default_factory=list) class ManifestArtifact(_StrictModel): diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index 5bda67860..c9c8f78c1 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -24,6 +24,7 @@ BundleVersionIdentity, DatasetPlanEntry, PrecomputePlan, + StateDatasetPlanEntry, ) from pydantic import ValidationError @@ -80,6 +81,33 @@ def _plan() -> PrecomputePlan: ) +def _state_dataset_entry(year: int, state_code: str, *, exists: bool): + digest = f"sd-{year}-{state_code}" + return StateDatasetPlanEntry( + year=year, + state_code=state_code, + region=f"state/{state_code.lower()}", + dataset="populace_local", + digest=digest, + path=f"state-datasets/us/{digest}/populace_local_year_{year}.h5", + filename=f"populace_local_year_{year}.h5", + folder=f"state-datasets/populace_local/{state_code}", + exists=exists, + ) + + +def _plan_with_state_datasets() -> PrecomputePlan: + return _plan().model_copy( + update={ + "state_datasets": [ + _state_dataset_entry(2026, "UT", exists=True), + _state_dataset_entry(2026, "WY", exists=False), + _state_dataset_entry(2027, "UT", exists=False), + ] + } + ) + + class TestSchemas: def test_plan_round_trips_through_the_modal_boundary(self): plan = _plan() @@ -126,6 +154,24 @@ def test_force_selects_everything(self): assert len(work.datasets) == 2 assert len(work.baselines) == 2 + def test_select_work_computes_only_missing_state_datasets(self): + work = precompute.select_work(_plan_with_state_datasets(), force=False) + assert [(entry.year, entry.state_code) for entry in work.state_datasets] == [ + (2026, "WY"), + (2027, "UT"), + ] + + def test_force_selects_every_state_dataset(self): + work = precompute.select_work(_plan_with_state_datasets(), force=True) + assert len(work.state_datasets) == 3 + + def test_a_plan_without_state_datasets_selects_none(self): + assert precompute.select_work(_plan(), force=False).state_datasets == [] + + def test_plan_with_state_datasets_round_trips_through_the_modal_boundary(self): + plan = _plan_with_state_datasets() + assert PrecomputePlan.model_validate(plan.model_dump()) == plan + def test_years_are_in_priority_order(self): assert precompute.PRECOMPUTE_YEARS == [2026, 2027, 2025] @@ -496,6 +542,42 @@ def fake_cohort_identity(year, group): ) monkeypatch.setattr(precompute, "cohort_identity", fake_cohort_identity) + targets = [ + precompute.StateDatasetTarget( + region="state/ut", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + dataset="populace_local", + ), + precompute.StateDatasetTarget( + region="state/wy", + state_code="WY", + state_fips=56, + filter_variable="state_fips", + dataset="populace_local", + ), + ] + monkeypatch.setattr(precompute, "state_dataset_targets", lambda: list(targets)) + existing.add("state-datasets/us/sd-2026-UT/populace_local_year_2026.h5") + + def fake_state_identity( + _country, year, *, dataset, state_code, state_fips, filter_variable + ): + assert (state_fips, filter_variable) in { + (49, "state_fips"), + (56, "state_fips"), + } + digest = f"sd-{year}-{state_code}" + return SimpleNamespace( + digest=digest, + store_path=f"state-datasets/us/{digest}/{dataset}_year_{year}.h5", + filename=f"{dataset}_year_{year}.h5", + ) + + monkeypatch.setattr( + artifact_keys, "collect_state_dataset_identity", fake_state_identity + ) monkeypatch.setattr( release_bundle, "get_country_release_bundle", @@ -507,7 +589,47 @@ def fake_cohort_identity(year, group): default_dataset="populace_cps", ), ) - return SimpleNamespace(existing=existing) + return SimpleNamespace(existing=existing, targets=targets) + + @pytest.mark.usefixtures("planning_stubs") + def test_plan_lists_every_state_dataset_year_major_with_presence(self): + plan = precompute.plan_artifacts_impl("bucket-x") + + 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 == "populace_local" + assert first.digest == "sd-2026-UT" + assert first.path == "state-datasets/us/sd-2026-UT/populace_local_year_2026.h5" + assert first.filename == "populace_local_year_2026.h5" + assert first.folder == "state-datasets/populace_local/UT" + work = precompute.select_work(plan, force=False) + assert len(work.state_datasets) == 5 + + def test_states_on_the_default_dataset_add_no_state_datasets( + self, planning_stubs, monkeypatch + ): + monkeypatch.setattr(precompute, "state_dataset_targets", list) + + plan = precompute.plan_artifacts_impl("bucket-x") + + assert plan.state_datasets == [] + assert len(plan.datasets) == 3 @pytest.mark.usefixtures("planning_stubs") def test_plan_enumerates_years_and_cohorts_with_presence(self): @@ -564,6 +686,111 @@ def test_empty_partition_fails_loudly(self, monkeypatch): precompute.plan_artifacts_impl("bucket-x") +def _us_state_region(state_code: str, state_fips: int, dataset_path: str): + from policyengine.core.scoping_strategy import RowFilterStrategy + + return SimpleNamespace( + code=f"state/{state_code.lower()}", + dataset_path=dataset_path, + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=state_fips + ), + parent_code="us", + state_code=state_code, + ) + + +class TestStateDatasetTargets: + """Which states get baked files, resolved the way a request resolves them.""" + + @pytest.fixture + def regional_bundle(self): + from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + ) + + bundle = get_country_release_bundle("us") + regional = next( + name for name in bundle.dataset_uris if name.endswith("_acs_local") + ) + return SimpleNamespace( + default_uri=bundle.default_dataset_uri, + regional=regional, + regional_uri=bundle.dataset_uris[regional], + ) + + def _install_regions(self, monkeypatch, regions): + from policyengine_simulation_executor import simulation_runtime + + by_code = {region.code: region for region in regions} + country_module = SimpleNamespace( + model=SimpleNamespace( + get_region=by_code.get, + region_registry=SimpleNamespace( + get_by_type=lambda region_type: ( + list(regions) if region_type == "state" else [] + ) + ), + ) + ) + monkeypatch.setattr( + simulation_runtime, "_country_module", lambda _country: country_module + ) + + def test_targets_states_whose_requests_load_a_regional_dataset( + self, monkeypatch, regional_bundle + ): + self._install_regions( + monkeypatch, + [ + _us_state_region("WY", 56, regional_bundle.regional_uri), + _us_state_region("UT", 49, regional_bundle.regional_uri), + ], + ) + + targets = precompute.state_dataset_targets() + + assert targets == [ + precompute.StateDatasetTarget( + region="state/ut", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + dataset=regional_bundle.regional, + ), + precompute.StateDatasetTarget( + region="state/wy", + state_code="WY", + state_fips=56, + filter_variable="state_fips", + dataset=regional_bundle.regional, + ), + ] + + def test_skips_states_whose_requests_load_the_default_dataset( + self, monkeypatch, regional_bundle + ): + self._install_regions( + monkeypatch, + [ + _us_state_region("UT", 49, regional_bundle.default_uri), + _us_state_region("WY", 56, regional_bundle.regional_uri), + ], + ) + + assert [target.state_code for target in precompute.state_dataset_targets()] == [ + "WY" + ] + + def test_refuses_a_state_without_a_row_filter(self, monkeypatch, regional_bundle): + unfiltered = _us_state_region("UT", 49, regional_bundle.regional_uri) + unfiltered.scoping_strategy = None + self._install_regions(monkeypatch, [unfiltered]) + + with pytest.raises(RuntimeError, match="state/ut"): + precompute.state_dataset_targets() + + class TestCohortIdentityGuard: def test_diverged_predicates_fail_loudly(self, monkeypatch): from policyengine_simulation_executor import baseline_artifacts as ba From cab957a20d2fc3deda5259398d8c7484caab94d8 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:18:04 +0400 Subject: [PATCH 05/16] Build, check and upload one year's per-state datasets build_state_datasets_impl re-derives each planned state's key in its own container and refuses to build when the planner disagreed, then has .py's ensure_datasets prepare the regional dataset's full-country year file once, collects garbage, and reloads that one year. It refuses data in which a household's congressional district code names another state, since district requests read their state's file. For each state it keeps the rows the state region's own scoping strategy keeps, writes them with .py's save(), copies the year-file stamp that save() drops, confirms that .py's reuse checks (stamp and renames record, then the legacy SPM record) accept the file, and uploads it write-once. Weights are unchanged. The tests write a three-household year file with .py's own save, stamp and SPM record, and show that .py's real ensure_datasets reuses a state file without rebuilding it. The identity stubs now import .py's US dataset module before patching the manifest, so a stubbed test cannot bind the stub into it. Co-Authored-By: Claude Opus 5.5 --- .../fixtures/identity_stubs.py | 7 +- .../fixtures/state_dataset_sources.py | 141 ++++++++ .../precompute.py | 180 +++++++++- .../precompute_models.py | 10 + .../state_datasets.py | 24 ++ .../tests/test_precompute.py | 308 ++++++++++++++++++ 6 files changed, 668 insertions(+), 2 deletions(-) create mode 100644 projects/policyengine-simulation-executor/fixtures/state_dataset_sources.py diff --git a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py index 0de437c35..bbd7b7f8c 100644 --- a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py +++ b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py @@ -36,7 +36,12 @@ def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: # This module imports ``get_release_manifest`` by value while constructing # country models. Import it before replacing the manifest module's function # so full-suite ordering cannot bind this test's stub into production code. - import policyengine.tax_benefit_models.common.model_version # noqa: F401 # must load before the manifest stubs are installed (see above) + import policyengine.tax_benefit_models.common.model_version + + # The US dataset module also imports the manifest helpers by value, and its + # year-file reuse checks (``_expected_legacy_spm_profile``) must keep + # reading the installed release after a stubbed test has imported it. + import policyengine.tax_benefit_models.us.datasets # noqa: F401 # both modules must load before the manifest stubs are installed (see above) from policyengine.provenance import ( manifest as manifest_module, ) diff --git a/projects/policyengine-simulation-executor/fixtures/state_dataset_sources.py b/projects/policyengine-simulation-executor/fixtures/state_dataset_sources.py new file mode 100644 index 000000000..3050c62dd --- /dev/null +++ b/projects/policyengine-simulation-executor/fixtures/state_dataset_sources.py @@ -0,0 +1,141 @@ +"""Small US year files for the per-state dataset tests. + +The files are written with policyengine.py's own ``PolicyEngineUSDataset.save()`` +and year-file stamp, and carry the SPM record the installed release expects +for the regional dataset, so .py's reuse checks treat them as it would a year +file ``create_datasets`` wrote. Three households: two in Utah (state FIPS 49) +and one in Wyoming (56), with uneven weights on every entity. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pandas as pd + +from fixtures.weighted_frames import weighted_entity_frame + +if TYPE_CHECKING: + from pathlib import Path + + from microdf import MicroDataFrame + +REGIONAL_DATASET = "populace_us_2024_acs_local" +UTAH_FIPS = 49 +WYOMING_FIPS = 56 + + +def source_entity_frames( + *, wyoming_district_geoid: int = 5600 +) -> dict[str, MicroDataFrame]: + """Return the six US entity frames of the three-household sample.""" + return { + "household": weighted_entity_frame( + { + "household_id": [1, 2, 3], + "household_weight": [10.0, 20.0, 30.0], + "state_fips": [UTAH_FIPS, UTAH_FIPS, WYOMING_FIPS], + "congressional_district_geoid": [4901, 4902, wyoming_district_geoid], + } + ), + "person": weighted_entity_frame( + { + "person_id": [1, 2, 3, 4], + "person_household_id": [1, 1, 2, 3], + "person_tax_unit_id": [1, 1, 2, 3], + "person_spm_unit_id": [1, 1, 2, 3], + "person_family_id": [1, 1, 2, 3], + "person_marital_unit_id": [1, 2, 3, 4], + "person_weight": [10.0, 11.0, 20.0, 30.0], + "age": [40, 12, 70, 33], + }, + entity="person", + ), + "tax_unit": weighted_entity_frame( + {"tax_unit_id": [1, 2, 3], "tax_unit_weight": [10.0, 20.0, 30.0]}, + entity="tax_unit", + ), + "spm_unit": weighted_entity_frame( + { + "spm_unit_id": [1, 2, 3], + "spm_unit_weight": [10.0, 20.0, 30.0], + "spm_unit_spm_universe_status": ["INCLUDED"] * 3, + }, + entity="spm_unit", + ), + "family": weighted_entity_frame( + {"family_id": [1, 2, 3], "family_weight": [10.0, 20.0, 30.0]}, + entity="family", + ), + "marital_unit": weighted_entity_frame( + { + "marital_unit_id": [1, 2, 3, 4], + "marital_unit_weight": [10.0, 11.0, 20.0, 30.0], + }, + entity="marital_unit", + ), + } + + +def year_file_metadata(dataset: str = REGIONAL_DATASET) -> dict: + """Return the metadata a ``create_datasets`` year file of ``dataset`` carries.""" + from policyengine.tax_benefit_models.us.datasets import ( + _expected_legacy_spm_profile, + ) + from policyengine.tax_benefit_models.us.legacy_inputs import RENAMES_RECORD_KEY + from policyengine.tax_benefit_models.us.legacy_spm.record import ( + RECORD_KEY, + LegacySPMRecord, + ) + + metadata: dict = {RENAMES_RECORD_KEY: {}} + profile = _expected_legacy_spm_profile(dataset) + if profile is not None: + metadata[RECORD_KEY] = LegacySPMRecord( + profile=profile, outside_spm_units=0, included_spm_units=3 + ).model_dump(mode="json") + return metadata + + +def write_year_file( + folder: Path, + *, + year: int, + dataset: str = REGIONAL_DATASET, + entity_frames: dict[str, MicroDataFrame] | None = None, + stamp: bool = True, +) -> Path: + """Write ``/_year_.h5`` as ``create_datasets`` would.""" + from policyengine.tax_benefit_models.us.datasets import ( + PolicyEngineUSDataset, + USYearData, + ) + from policyengine.tax_benefit_models.us.legacy_inputs import ( + write_year_file_stamp, + year_file_stamp, + ) + + path = folder / f"{dataset}_year_{year}.h5" + PolicyEngineUSDataset( + name=f"{dataset}-year-{year}", + description=f"Test year file for {year}", + filepath=str(path), + year=year, + metadata=year_file_metadata(dataset), + data=USYearData(**(entity_frames or source_entity_frames())), + ).save() + if stamp: + write_year_file_stamp(path, year_file_stamp()) + return path + + +def stored_ids(path: Path, entity: str) -> list[int]: + """Return an entity's ids as stored in a year file, in row order.""" + with pd.HDFStore(path, mode="r") as store: + return [int(value) for value in store[entity][f"{entity}_id"]] + + +def stored_weights(path: Path, entity: str) -> list[float]: + """Return an entity's stored weights, in row order.""" + with pd.HDFStore(path, mode="r") as store: + return [float(value) for value in store[entity][f"{entity}_weight"]] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 671c44105..933a5b733 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -32,9 +32,11 @@ from __future__ import annotations +import gc import logging +import tempfile import time -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any @@ -53,17 +55,20 @@ ManifestArtifact, PrecomputePlan, RemoteFunction, + StateDatasetBuildResult, StateDatasetPlanEntry, WorkSelection, ) from policyengine_simulation_executor.state_datasets import ( state_dataset_relative_folder, + year_file_reuse_problem, ) if TYPE_CHECKING: from collections.abc import Callable from types import ModuleType + from microdf import MicroDataFrame from policyengine.core import Simulation from policyengine_simulation_executor.artifact_store import ArtifactStore @@ -162,6 +167,9 @@ class StateDatasetTarget: The precompute 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, so a baked file holds exactly the rows the request + would keep. """ region: str @@ -169,6 +177,7 @@ class StateDatasetTarget: state_fips: int filter_variable: str dataset: str + scoping_strategy: Any = field(default=None, compare=False, repr=False) def state_dataset_targets() -> list[StateDatasetTarget]: @@ -214,6 +223,7 @@ def state_dataset_targets() -> list[StateDatasetTarget]: state_fips=int(variable_value), filter_variable=variable_name, dataset=selection.name, + scoping_strategy=strategy, ) ) return sorted(targets, key=lambda target: target.state_code) @@ -399,6 +409,174 @@ def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildR ) +def _checked_state_targets( + year: int, expected: list[StateDatasetPlanEntry] +) -> dict[str, StateDatasetTarget]: + """Return each planned state's target, after re-deriving its key in this container. + + The planner and this builder run in different containers; both derive a + state's key from the installed bundle. A disagreement means they see + different bundles, so nothing is built rather than uploading a file under + a key that does not describe it. + """ + from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_state_dataset_identity + collect_state_dataset_identity, + ) + + 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( + PRECOMPUTE_COUNTRY, + 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 in-container identities of {entry.region} {year} " + f"disagree ({entry.path} != {identity.store_path}); refusing to " + "build under a mismatched key." + ) + 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). A household whose district code names another state would be + missing from that district's results, so the build stops instead. 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: list[StateDatasetPlanEntry], + *, + work_folder: Path | None = None, +) -> list[StateDatasetBuildResult]: + """Build, check and upload one year's per-state datasets. + + Prepares 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 + + from policyengine_simulation_executor.artifact_store import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_store.ArtifactStore + ArtifactStore, + ) + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _country_module, + ) + + if not expected: + return [] + targets = _checked_state_targets(year, expected) + dataset = expected[0].dataset + work_folder = work_folder or Path( + tempfile.mkdtemp(prefix="policyengine-state-datasets-") + ) + source_folder = str(work_folder / "source") + country_module = _country_module(PRECOMPUTE_COUNTRY) + + # 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 = ArtifactStore(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) + results.append( + StateDatasetBuildResult( + year=year, + state_code=entry.state_code, + path=entry.path, + uploaded=store.upload_file(entry.path, local_file), + size_bytes=local_file.stat().st_size, + ) + ) + return results + + def _prepare_cohort_baseline( bucket: str, expected: BaselinePlanEntry ) -> tuple[ArtifactStore, Path, Simulation]: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py index c736ede88..d3fe2d1ba 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py @@ -166,6 +166,16 @@ class DatasetBuildResult(_StrictModel): size_bytes: int +class StateDatasetBuildResult(_StrictModel): + """Result of one state's dataset build; ``uploaded`` is False if the key existed.""" + + year: int + state_code: str + path: str + uploaded: bool + size_bytes: int + + class BaselineComputeResult(_StrictModel): """Outcome of one cohort baseline computation and its write-once upload. diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py index e43f9929a..a513bb7ff 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py @@ -26,3 +26,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_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index c9c8f78c1..a474ab2a6 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -903,6 +903,314 @@ def test_fails_loudly_when_no_file_is_produced(self, dataset_stubs): assert dataset_stubs.uploads == [] +class TestBuildStateDatasetsImpl: + """Building, checking and uploading one year's per-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 + from policyengine_simulation_executor import ( + artifact_keys, + artifact_store, + ) + from policyengine_simulation_executor import simulation_runtime as sr + + state = SimpleNamespace( + uploads=[], + existing=set(), + ensure_calls=[], + entity_frames=sources.source_entity_frames(), + work_folder=tmp_path / "work", + ) + + 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: list(state.targets) + ) + + def fake_state_identity( + _country, year, *, dataset, state_code, state_fips, filter_variable + ): + digest = f"sd-{year}-{state_code}-{state_fips}-{filter_variable}" + return SimpleNamespace( + digest=digest, + store_path=f"state-datasets/us/{digest}/{dataset}_year_{year}.h5", + filename=f"{dataset}_year_{year}.h5", + ) + + monkeypatch.setattr( + artifact_keys, "collect_state_dataset_identity", fake_state_identity + ) + + class FakeStore: + def __init__(self, bucket): + self.bucket = bucket + + def upload_file(self, path, local_path): + state.uploads.append((path, Path(local_path))) + return path not in state.existing + + monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) + + 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( + sr, + "_country_module", + lambda _country: SimpleNamespace( + ensure_datasets=fake_ensure_datasets, + load_datasets=us_module.load_datasets, + ), + ) + state.sources = sources + state.entry = lambda state_code, fips: self._entry( + fake_state_identity, state_code, fips + ) + return state + + def _entry(self, identity_for, state_code, fips): + from fixtures.state_dataset_sources import REGIONAL_DATASET + + identity = identity_for( + "us", + self.YEAR, + dataset=REGIONAL_DATASET, + state_code=state_code, + state_fips=fips, + filter_variable="state_fips", + ) + return StateDatasetPlanEntry( + year=self.YEAR, + state_code=state_code, + region=f"state/{state_code.lower()}", + dataset=REGIONAL_DATASET, + digest=identity.digest, + path=identity.store_path, + filename=identity.filename, + folder=f"state-datasets/{REGIONAL_DATASET}/{state_code}", + exists=False, + ) + + def _build(self, state_build, entries): + return precompute.build_state_datasets_impl( + "bucket-x", self.YEAR, entries, work_folder=state_build.work_folder + ) + + def _both_states(self, state_build): + return [ + state_build.entry("UT", state_build.sources.UTAH_FIPS), + state_build.entry("WY", state_build.sources.WYOMING_FIPS), + ] + + def test_writes_each_state_file_with_only_its_households(self, state_build): + sources = state_build.sources + results = self._build(state_build, self._both_states(state_build)) + + files = dict(state_build.uploads) + utah, wyoming = (files[entry.path] for entry in self._both_states(state_build)) + assert sources.stored_ids(utah, "household") == [1, 2] + assert sources.stored_ids(utah, "person") == [1, 2, 3] + assert sources.stored_ids(utah, "tax_unit") == [1, 2] + assert sources.stored_ids(utah, "marital_unit") == [1, 2, 3] + assert sources.stored_ids(wyoming, "household") == [3] + assert sources.stored_ids(wyoming, "person") == [4] + assert sources.stored_ids(wyoming, "spm_unit") == [3] + assert [(result.state_code, result.year) for result in results] == [ + ("UT", 2026), + ("WY", 2026), + ] + + def test_keeps_every_weight_unchanged(self, state_build): + sources = state_build.sources + self._build(state_build, self._both_states(state_build)) + + utah = state_build.uploads[0][1] + assert sources.stored_weights(utah, "household") == [10.0, 20.0] + assert sources.stored_weights(utah, "person") == [10.0, 11.0, 20.0] + assert sources.stored_weights(utah, "marital_unit") == [10.0, 11.0, 20.0] + + def test_state_files_pass_dotpy_reuse_checks(self, state_build): + from policyengine.tax_benefit_models.us.datasets import ( + _expected_legacy_spm_profile, + ) + from policyengine.tax_benefit_models.us.legacy_inputs import ( + read_year_file_stamp, + year_file_stamp, + year_file_stamp_problem, + ) + from policyengine.tax_benefit_models.us.legacy_spm.record import ( + year_file_problem, + ) + + self._build(state_build, self._both_states(state_build)) + + for _path, local in state_build.uploads: + assert read_year_file_stamp(local) == year_file_stamp() + assert year_file_stamp_problem(local) is None + assert ( + year_file_problem( + local, + _expected_legacy_spm_profile(state_build.sources.REGIONAL_DATASET), + ) + is None + ) + + def test_dotpy_reuses_a_state_file_without_rebuilding( + self, state_build, monkeypatch + ): + from policyengine.tax_benefit_models.us import datasets as us_datasets + + self._build(state_build, self._both_states(state_build)) + + 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) + utah_folder = state_build.uploads[0][1].parent + loaded = us_datasets.ensure_datasets( + datasets=[state_build.sources.REGIONAL_DATASET], + years=[self.YEAR], + data_folder=str(utah_folder), + ) + + (dataset,) = loaded.values() + assert sorted(dataset.data.household["household_id"]) == [1, 2] + + def test_filtering_a_state_file_again_keeps_the_same_rows(self, state_build): + from policyengine.tax_benefit_models.us import load_datasets + from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES + + self._build(state_build, self._both_states(state_build)) + utah_target = state_build.targets[0] + utah_folder = state_build.uploads[0][1].parent + (state_file,) = load_datasets( + datasets=[state_build.sources.REGIONAL_DATASET], + years=[self.YEAR], + data_folder=str(utah_folder), + ).values() + + filtered_again = utah_target.scoping_strategy.apply( + entity_data=state_file.data.entity_data, + group_entities=US_GROUP_ENTITIES, + year=self.YEAR, + ) + + for entity, frame in state_file.data.entity_data.items(): + assert list(filtered_again[entity][f"{entity}_id"]) == list( + frame[f"{entity}_id"] + ) + + def test_prepares_the_full_country_year_file_once(self, state_build): + self._build(state_build, self._both_states(state_build)) + + 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] + assert Path(folder).is_relative_to(state_build.work_folder) + + def test_uploads_each_file_under_its_planned_path(self, state_build): + entries = self._both_states(state_build) + state_build.existing.add(entries[1].path) + + results = self._build(state_build, entries) + + assert [path for path, _local in state_build.uploads] == [ + entry.path for entry in entries + ] + assert [local.name for _path, local in state_build.uploads] == [ + entry.filename for entry in entries + ] + assert [result.uploaded for result in results] == [True, False] + assert [result.path for result in results] == [entry.path for entry in entries] + assert all(result.size_bytes > 0 for result in results) + + 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._both_states(state_build)) + assert state_build.uploads == [] + + def test_refuses_a_state_with_no_households(self, state_build): + from policyengine.core.scoping_strategy import RowFilterStrategy + + state_build.targets.append( + precompute.StateDatasetTarget( + region="state/id", + state_code="ID", + state_fips=16, + filter_variable="state_fips", + dataset=state_build.sources.REGIONAL_DATASET, + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=16 + ), + ) + ) + + with pytest.raises(ValueError, match="No households"): + self._build(state_build, [state_build.entry("ID", 16)]) + assert state_build.uploads == [] + + def test_refuses_a_plan_entry_whose_key_differs_in_the_container(self, state_build): + entry = state_build.entry("UT", state_build.sources.UTAH_FIPS) + mismatched = entry.model_copy(update={"path": "state-datasets/us/OTHER/x.h5"}) + + with pytest.raises(RuntimeError, match="refusing to build"): + self._build(state_build, [mismatched]) + 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._both_states(state_build)) + assert state_build.ensure_calls == [] + + def test_refuses_entries_for_another_year(self, state_build): + entry = state_build.entry("UT", state_build.sources.UTAH_FIPS) + + with pytest.raises(ValueError, match="2027"): + precompute.build_state_datasets_impl( + "bucket-x", + 2027, + [entry], + work_folder=state_build.work_folder, + ) + + class TestComputeBaselineImpl: @pytest.fixture def cohort_stubs(self, monkeypatch, tmp_path): From abd995e7f7070e8789e5cfcaf12fcf3d022f9ab7 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:21:03 +0400 Subject: [PATCH 06/16] Place manifest artifacts that name a folder in that folder of the image The image's artifact fetch now puts an artifact with a "folder" at //, so per-state datasets, which all share the year file's basename, land in their own state folders. Artifacts without a folder still go straight into the data folder. Before anything downloads, it refuses absolute folders, folders with "." or ".." parts, empty parts or backslashes, filenames that are not bare names, and two entries for the same file. The fetch layer's timeout rises from 900 s to 3,600 s for the several GB of per-state datasets. No manifest lists such artifacts yet. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/_image_setup.py | 59 ++++++++--- .../src/modal/app.py | 4 +- .../tests/test_image_setup_fetch.py | 97 +++++++++++++++++++ .../tests/test_modal_bundle_image.py | 4 +- 4 files changed, 151 insertions(+), 13 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/_image_setup.py b/projects/policyengine-simulation-executor/src/modal/_image_setup.py index 220e8b3a7..011bfc1e9 100644 --- a/projects/policyengine-simulation-executor/src/modal/_image_setup.py +++ b/projects/policyengine-simulation-executor/src/modal/_image_setup.py @@ -15,7 +15,7 @@ import json import logging import os -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -178,17 +178,58 @@ def _service_account_credentials(blob: str) -> "dict[str, Any]": } +def _artifact_target(data_folder: "Path", artifact: "Mapping[str, Any]") -> "Path": + """Return where an artifact lands: its ``folder`` (if any) under the data folder. + + Per-state datasets share one ``filename`` and differ by ``folder``, a + relative path such as ``state-datasets//``. Anything that + could leave the data folder, or name a different file than intended, is + refused. + """ + filename = artifact.get("filename") + if ( + not isinstance(filename, str) + or filename in {"", ".", ".."} + or "/" in filename + or "\\" in filename + ): + msg = f"Artifact {artifact.get('path')!r} has an unsafe filename {filename!r}" + raise RuntimeError(msg) + folder = artifact.get("folder") + if folder is None: + return data_folder / filename + parts = ( + PurePosixPath(folder).parts + if isinstance(folder, str) and folder and "\\" not in folder + else () + ) + if ( + not parts + or PurePosixPath(folder).is_absolute() + or any(part in {".", ".."} for part in parts) + or "/".join(parts) != folder + ): + msg = f"Artifact {artifact.get('path')!r} has an unsafe folder {folder!r}" + raise RuntimeError(msg) + return data_folder.joinpath(*parts, filename) + + def _download_manifest_artifacts( bucket_handle: "storage.Bucket", artifacts: "Sequence[Mapping[str, Any]]", data_folder: "Path", ) -> None: - """Download every artifact under its runtime filename, failing on any gap. + """Download every artifact to its runtime location, failing on any gap. - Every store object is checked before anything downloads, and every - expected file is checked after. + Every target and store object is checked before anything downloads, and + every expected file is checked after. """ logger = logging.getLogger(__name__) + targets = [_artifact_target(data_folder, artifact) for artifact in artifacts] + repeated = sorted({str(target) for target in targets if targets.count(target) > 1}) + if repeated: + msg = f"The manifest lists these files more than once: {repeated}" + raise RuntimeError(msg) missing = [ artifact["path"] for artifact in artifacts @@ -206,16 +247,12 @@ def _download_manifest_artifacts( # Unconditional downloads: a same-named stale file must never # survive into the image, and the layer starts empty of these # files anyway. - for artifact in artifacts: - target = data_folder / artifact["filename"] + for artifact, target in zip(artifacts, targets, strict=True): + target.parent.mkdir(parents=True, exist_ok=True) bucket_handle.blob(artifact["path"]).download_to_filename(str(target)) logger.info("Fetched %s (%.1f MB)", target, target.stat().st_size / 1e6) - absent = [ - str(data_folder / artifact["filename"]) - for artifact in artifacts - if not (data_folder / artifact["filename"]).exists() - ] + absent = [str(target) for target in targets if not target.exists()] if absent: msg = f"Artifact fetch did not produce expected files: {absent}" raise RuntimeError(msg) diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index 08c5e4c1c..eb52bab63 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -235,7 +235,9 @@ def build_base_simulation_image() -> modal.Image: secrets=[gcp_secret], cpu=2.0, memory=4096, - timeout=900, + # The national files and baselines are about 430 MB; the + # per-state datasets (51 states for each year) add several GB. + timeout=3600, ) ) diff --git a/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py index c61522820..0cb9ee80a 100644 --- a/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py +++ b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py @@ -132,6 +132,103 @@ def test_overwrites_a_preexisting_same_named_file(self, fake_client, data_folder assert stale.read_bytes() == b"bytes-d1" +def _state_dataset_artifact(state_code, *, folder=None, filename=None): + digest = f"s-{state_code}" + return { + "type": "state_dataset", + "path": f"state-datasets/us/{digest}/populace_local_year_2026.h5", + "filename": filename or "populace_local_year_2026.h5", + "folder": folder or f"state-datasets/populace_local/{state_code}", + "state_code": state_code, + "year": 2026, + "digest": digest, + } + + +def _manifest_with_state_datasets(*state_artifacts): + manifest = _manifest() + manifest["artifacts"] = [*manifest["artifacts"], *state_artifacts] + return manifest + + +def _store_every_artifact(client, manifest): + for artifact in manifest["artifacts"]: + client.objects[artifact["path"]] = f"bytes-{artifact['digest']}".encode() + + +class TestStateDatasetPlacement: + """Entries with a ``folder`` land in it; every state shares one filename.""" + + def test_places_each_state_file_in_its_own_folder(self, fake_client, data_folder): + manifest = _manifest_with_state_datasets( + _state_dataset_artifact("UT"), _state_dataset_artifact("WY") + ) + _store_every_artifact(fake_client, manifest) + + fetch_artifacts("test-bucket", manifest, client=fake_client) + + state_folder = data_folder / "state-datasets" / "populace_local" + assert (state_folder / "UT" / "populace_local_year_2026.h5").read_bytes() == ( + b"bytes-s-UT" + ) + assert (state_folder / "WY" / "populace_local_year_2026.h5").read_bytes() == ( + b"bytes-s-WY" + ) + assert (data_folder / "populace_year_2026.h5").read_bytes() == b"bytes-d1" + assert not (data_folder / "populace_local_year_2026.h5").exists() + + @pytest.mark.parametrize( + "folder", + [ + "/abs/state-datasets/UT", + "state-datasets/../../escape", + "state-datasets/./UT", + "state-datasets//UT", + "state-datasets\\UT", + "", + ], + ) + @pytest.mark.usefixtures("data_folder") + def test_refuses_an_unsafe_folder_before_any_download(self, fake_client, folder): + manifest = _manifest_with_state_datasets( + _state_dataset_artifact("UT", folder=folder or None) + ) + manifest["artifacts"][-1]["folder"] = folder + _store_every_artifact(fake_client, manifest) + + with pytest.raises(RuntimeError, match="folder"): + fetch_artifacts("test-bucket", manifest, client=fake_client) + assert fake_client.downloads == [] + + @pytest.mark.parametrize( + "filename", ["../populace_local_year_2026.h5", "UT/x.h5", "..", "."] + ) + @pytest.mark.usefixtures("data_folder") + def test_refuses_an_unsafe_filename_before_any_download( + self, fake_client, filename + ): + manifest = _manifest_with_state_datasets( + _state_dataset_artifact("UT", filename=filename) + ) + _store_every_artifact(fake_client, manifest) + + with pytest.raises(RuntimeError, match="filename"): + fetch_artifacts("test-bucket", manifest, client=fake_client) + assert fake_client.downloads == [] + + @pytest.mark.usefixtures("data_folder") + def test_refuses_two_entries_for_one_file_before_any_download(self, fake_client): + repeated = _state_dataset_artifact("UT") + manifest = _manifest_with_state_datasets( + repeated, {**repeated, "path": "state-datasets/us/other/x.h5"} + ) + _store_every_artifact(fake_client, manifest) + + with pytest.raises(RuntimeError, match="more than once"): + fetch_artifacts("test-bucket", manifest, client=fake_client) + assert fake_client.downloads == [] + + class TestFreshnessGate: @pytest.mark.usefixtures("data_folder") def test_rejects_a_policyengine_version_mismatch(self, monkeypatch, fake_client): diff --git a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py index 2ad12baec..52e538a83 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py @@ -196,7 +196,9 @@ def read_manifest(self, digest): assert fetch_kwargs["secrets"] == [app.gcp_secret] assert fetch_kwargs["cpu"] == 2.0 assert fetch_kwargs["memory"] == 4096 - assert fetch_kwargs["timeout"] == 900 + # Room for the per-state datasets (51 states for each precomputed year) + # on top of the national files and baselines. + assert fetch_kwargs["timeout"] == 3600 # The fetch layer must stay after the version env (its freshness gate # reads the baked versions) and before add_local_python_source (whose From 6f24fa31e99343ac4a8609779dd62a8bf73e002e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:25:00 +0400 Subject: [PATCH 07/16] Run the per-state dataset builds in the main path's precompute The precompute app gains build_state_datasets (8 CPUs, 96 GiB, two hours), one container per year, validating plan entries in and dumping build results out like the other wrappers. run_precompute spawns the state builds with the dataset builds and waits for them only before publishing the manifest, so baseline computation never waits for them, then reports how many state datasets were built and uploaded and their total size. The manifest does not list the state datasets yet. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/precompute_app.py | 37 ++++++- .../precompute.py | 30 +++++- .../tests/test_precompute.py | 99 +++++++++++++++++++ .../tests/test_precompute_app.py | 52 ++++++++++ 4 files changed, 212 insertions(+), 6 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/precompute_app.py b/projects/policyengine-simulation-executor/src/modal/precompute_app.py index b3b30a7db..48d048861 100644 --- a/projects/policyengine-simulation-executor/src/modal/precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/precompute_app.py @@ -10,9 +10,10 @@ Wave 1 builds the single-year US datasets (one container per year), wave 2 computes the 20 per-cohort national baselines per year (one container per -cohort, sized like the production ``run_simulation_segment`` workers). See -the library module for the idempotency, ``--force``, and determinism-gate -semantics. +cohort, sized like the production ``run_simulation_segment`` workers). +Alongside both waves, one container per year filters the regional dataset +that state regions load into one year file per state. See the library module +for the idempotency, ``--force``, and determinism-gate semantics. """ from __future__ import annotations @@ -86,6 +87,35 @@ 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=2 * 60 * 60, + secrets=_worker_secrets, +) +def build_state_datasets(bucket: str, year: int, expected: list[dict]) -> list[dict]: + """Build and upload one year's per-state datasets. + + Raises if a state's key resolved in this container differs from the + planned one, so planner/worker drift cannot upload under the wrong key. + """ + from policyengine_simulation_executor.precompute import ( # noqa: PLC0415 # Modal function body: loaded in the precompute worker, not at deploy time + build_state_datasets_impl, + ) + from policyengine_simulation_executor.precompute_models import ( # noqa: PLC0415 # Modal function body: loaded in the precompute worker, not at deploy time + StateDatasetPlanEntry, + ) + + entries = [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, @@ -164,6 +194,7 @@ def main(*, force: bool = False) -> None: functions=PrecomputeFunctions( 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/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 933a5b733..58a4eb513 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -760,10 +760,11 @@ def publish_manifest_impl(bucket: str, plan: PrecomputePlan) -> str: @dataclass(frozen=True) class PrecomputeFunctions: - """The five Modal function handles the precompute orchestration calls.""" + """The six Modal function handles the precompute orchestration calls.""" plan_artifacts: RemoteFunction build_dataset: RemoteFunction + build_state_datasets: RemoteFunction compute_baseline: RemoteFunction verify_determinism: RemoteFunction publish_manifest: RemoteFunction @@ -778,8 +779,11 @@ def run_precompute( ) -> str: """Orchestrate the local entrypoint: plan, spawn misses, gate, publish. - Takes the five Modal function handles as ``functions`` (each offering + Takes the six Modal function handles as ``functions`` (each offering ``.remote``/``.spawn``) so the flow is testable with plain fakes. + Per-state dataset builds (one container per year) start with the dataset + builds and are waited for only before the manifest is published, so the + baselines never wait for them. Everything crossing a handle is a plain dict (Modal's serialization boundary); it is validated back into models immediately on return. The determinism gate runs only when baselines were actually computed @@ -792,7 +796,8 @@ def run_precompute( work = select_work(plan, force=force) echo( f"Plan: {len(plan.datasets)} datasets ({len(work.datasets)} to " - f"build), {len(plan.baselines)} baselines " + f"build), {len(plan.state_datasets)} state datasets " + f"({len(work.state_datasets)} to build), {len(plan.baselines)} baselines " f"({len(work.baselines)} to compute)" ) @@ -800,6 +805,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]] = {} + for entry in work.state_datasets: + state_entries_by_year.setdefault(entry.year, []).append(entry.model_dump()) + state_handles = [ + functions.build_state_datasets.spawn(bucket, year, entries) + for year, entries in state_entries_by_year.items() + ] for handle in dataset_handles: result = DatasetBuildResult.model_validate(handle.get()) echo( @@ -832,6 +844,18 @@ def run_precompute( ) echo("Determinism gate passed") + if state_handles: + state_results = [ + StateDatasetBuildResult.model_validate(result) + for handle in state_handles + for result in handle.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" + ) + manifest_digest = functions.publish_manifest.remote(bucket, plan.model_dump()) echo(manifest_digest_line(manifest_digest)) return manifest_digest diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index a474ab2a6..0b1d2bc6e 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -309,6 +309,45 @@ def _baseline_result(year): } +class RecordingFunction: + """A Modal handle fake that logs, in one shared list, when it is used. + + ``results`` maps the spawned year to the result returned. The year is the + second argument itself (a state build) or that plan entry's ``year``. + """ + + def __init__(self, name, log, results): + self.name = name + self.log = log + self.results = results + self.spawn_calls = [] + + def spawn(self, *args): + self.spawn_calls.append(args) + year = args[1] if isinstance(args[1], int) else args[1]["year"] + self.log.append(("spawn", self.name, year)) + log, name, result = self.log, self.name, self.results[year] + + class Handle: + def get(self): + log.append(("get", name, year)) + return result + + return Handle() + + +def _state_build_result(year, state_code, *, uploaded=True, size_bytes=100): + return { + "year": year, + "state_code": state_code, + "path": ( + f"state-datasets/us/sd-{year}-{state_code}/populace_local_year_{year}.h5" + ), + "uploaded": uploaded, + "size_bytes": size_bytes, + } + + class TestRunPrecompute: def _functions(self, verify_result=None): return { @@ -316,6 +355,7 @@ def _functions(self, verify_result=None): "build_dataset": FakeFunction( results_by_year={2027: _dataset_result(2027)} ), + "build_state_datasets": RecordingFunction("build_state_datasets", [], {}), "compute_baseline": FakeFunction( results_by_year={2027: _baseline_result(2027)} ), @@ -386,6 +426,65 @@ def test_failed_gate_aborts_before_publishing(self): self._run(functions) assert functions["publish_manifest"].remote_calls == [] + def _state_run(self, *, uploaded_wyoming=True): + log = [] + functions = self._functions() + functions["plan_artifacts"] = FakeFunction( + result=_plan_with_state_datasets().model_dump() + ) + functions["build_state_datasets"] = RecordingFunction( + "build_state_datasets", + log, + { + 2026: [ + _state_build_result(2026, "WY", uploaded=uploaded_wyoming), + ], + 2027: [_state_build_result(2027, "UT", size_bytes=250)], + }, + ) + functions["compute_baseline"] = RecordingFunction( + "compute_baseline", log, {2027: _baseline_result(2027)} + ) + publish = functions["publish_manifest"] + + def publish_remote(*args): + log.append(("remote", "publish_manifest", None)) + return FakeFunction.remote(publish, *args) + + publish.remote = publish_remote + digest, lines = self._run(functions) + return functions, log, digest, lines + + def test_spawns_one_state_build_per_year_with_its_missing_entries(self): + functions, _log, _digest, _lines = self._state_run() + + plan = _plan_with_state_datasets() + assert functions["build_state_datasets"].spawn_calls == [ + ("bucket-x", 2026, [plan.state_datasets[1].model_dump()]), + ("bucket-x", 2027, [plan.state_datasets[2].model_dump()]), + ] + + def test_waits_for_state_builds_only_before_publishing(self): + _functions, log, _digest, _lines = self._state_run() + + first_state_get = log.index(("get", "build_state_datasets", 2026)) + assert log.index(("spawn", "build_state_datasets", 2026)) < log.index( + ("spawn", "compute_baseline", 2027) + ) + assert log.index(("get", "compute_baseline", 2027)) < first_state_get + assert first_state_get < log.index(("remote", "publish_manifest", None)) + + def test_reports_the_state_dataset_bytes_uploaded(self): + _functions, _log, _digest, lines = self._state_run(uploaded_wyoming=False) + + assert "State datasets: 2 built, 1 uploaded, 350 bytes" in lines + assert lines[-1] == "MANIFEST_DIGEST=digest-123" + + def test_a_run_without_state_datasets_spawns_no_state_build(self): + functions = self._functions() + self._run(functions) + assert functions["build_state_datasets"].spawn_calls == [] + def test_malformed_worker_result_fails_loudly(self): functions = self._functions() functions["build_dataset"] = FakeFunction( diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_precompute_app.py index 16d710103..fabd0ec96 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute_app.py @@ -18,6 +18,7 @@ _WORKER_FUNCTIONS = ( "plan_artifacts", "build_dataset", + "build_state_datasets", "compute_baseline", "verify_determinism", "publish_manifest", @@ -86,6 +87,56 @@ def test_dataset_builder_gets_dataset_build_scale_resources(precompute_module): assert kwargs["memory"] == 65536 +def test_state_dataset_builder_has_room_for_the_regional_year_file( + precompute_module, +): + kwargs = _function_kwargs(precompute_module, "build_state_datasets") + assert kwargs["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 kwargs["memory"] == 98304 + assert kwargs["timeout"] == 2 * 60 * 60 + + +def test_state_dataset_wrapper_validates_entries_in_and_dumps_results_out( + precompute_module, monkeypatch +): + from policyengine_simulation_executor import precompute as lib + from policyengine_simulation_executor.precompute_models import ( + StateDatasetBuildResult, + StateDatasetPlanEntry, + ) + + entry = StateDatasetPlanEntry( + year=2026, + state_code="UT", + region="state/ut", + dataset="populace_local", + digest="sd1", + path="state-datasets/us/sd1/populace_local_year_2026.h5", + filename="populace_local_year_2026.h5", + folder="state-datasets/populace_local/UT", + exists=False, + ) + result = StateDatasetBuildResult( + year=2026, state_code="UT", path=entry.path, uploaded=True, size_bytes=5 + ) + seen = {} + + def fake_build_state_datasets(bucket, year, expected): + seen["call"] = (bucket, year, expected) + return [result] + + monkeypatch.setattr(lib, "build_state_datasets_impl", fake_build_state_datasets) + + assert precompute_module.build_state_datasets( + "bucket-x", 2026, [entry.model_dump()] + ) == [result.model_dump()] + assert seen["call"] == ("bucket-x", 2026, [entry]) + with pytest.raises(ValidationError): + precompute_module.build_state_datasets("bucket-x", 2026, [{"year": 2026}]) + + def test_image_includes_local_source(precompute_module): calls = precompute_module.precompute_image.calls local_source = [call for call in calls if call[0] == "add_local_python_source"] @@ -242,6 +293,7 @@ def fake_run_precompute(bucket, *, force, functions): assert captured["handles"] == { "plan_artifacts": precompute_module.plan_artifacts, "build_dataset": precompute_module.build_dataset, + "build_state_datasets": precompute_module.build_state_datasets, "compute_baseline": precompute_module.compute_baseline, "verify_determinism": precompute_module.verify_determinism, "publish_manifest": precompute_module.publish_manifest, From 6386db4d55382ed36e017e5c81446a962d735c72 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:28:12 +0400 Subject: [PATCH 08/16] List the per-state datasets in the main path's deploy manifest build_manifest appends a state_dataset entry for every planned state file, with its folder and state code, after the datasets and baselines. The new entries use their own StateDatasetArtifact model, chosen by "type", so a plan without per-state datasets publishes byte-for-byte the manifest it always did (the pinned wire-shape test is unchanged). The deploy-time manifest validation accepts the new entries and hands them to the fetch layer, which places each in its state's folder. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 14 ++++-- .../precompute.py | 22 ++++++++- .../precompute_models.py | 27 +++++++++-- .../tests/test_modal_bundle_image.py | 46 ++++++++++++++++++ .../tests/test_precompute.py | 48 +++++++++++++++++++ 5 files changed, 148 insertions(+), 9 deletions(-) diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 64f2f60f8..72087b956 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -82,9 +82,12 @@ durable plan without changing routing. ## Baseline artifact pipeline (precompute → store → image fetch) The precompute app fills a content-addressed GCS store with single-year US -datasets (2026, 2027, 2025) and the 20 per-cohort national baseline -simulations per year; the deploy bakes them into the image, and the -runtime loads baselines instead of computing them. Library logic lives in +datasets (2026, 2027, 2025), the 20 per-cohort national baseline +simulations per year, and per-state datasets: for each state region whose +requests load a dataset other than the bundle default (the ACS-local +dataset since policyengine 6.7), that dataset's year file filtered to the +state's households, one file per state and year. The deploy bakes them into +the image, and the runtime loads baselines instead of computing them. Library logic lives in `src/policyengine_simulation_executor/precompute.py` (keys: `artifact_keys.py`, store client: `artifact_store.py`); the Modal app (`src/modal/precompute_app.py`) is plumbing only. @@ -103,7 +106,10 @@ The deploy job consumes the precompute job's manifest digest: `src/modal/app.py` resolves the manifest from GCS on the deploying machine and passes its content as the `fetch_artifacts` image layer's args (`src/modal/_image_setup.py`). The layer downloads the full -artifact set (~430 MB) into `POLICYENGINE_DATA_FOLDER`, sitting between +artifact set into `POLICYENGINE_DATA_FOLDER`: about 430 MB of national +datasets and baselines, plus several GB of per-state datasets, each placed +in `state-datasets///` under the year file's own name. It +sits between the version env and the source mount — do not reorder. Because the manifest is content-addressed and rides in the layer args, Modal's layer cache busts exactly when the artifact set changes and never otherwise; diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 58a4eb513..622a89300 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -53,8 +53,10 @@ DatasetPlanEntry, DeterminismVerdict, ManifestArtifact, + ManifestEntry, PrecomputePlan, RemoteFunction, + StateDatasetArtifact, StateDatasetBuildResult, StateDatasetPlanEntry, WorkSelection, @@ -241,8 +243,12 @@ def select_work(plan: PrecomputePlan, *, force: bool) -> WorkSelection: def build_manifest(plan: PrecomputePlan) -> ArtifactManifest: - """Build the deploy manifest: the receipt and each artifact the image must fetch.""" - artifacts = [ + """Build the deploy manifest: the receipt and each artifact the image must fetch. + + Datasets, then baselines, then per-state datasets, each in plan order, so + a plan without per-state datasets keeps the manifest it always had. + """ + artifacts: list[ManifestEntry] = [ ManifestArtifact( type="dataset", path=entry.path, @@ -261,6 +267,18 @@ def build_manifest(plan: PrecomputePlan) -> ArtifactManifest: ) for entry in plan.baselines ] + artifacts += [ + StateDatasetArtifact( + type="state_dataset", + path=entry.path, + filename=entry.filename, + folder=entry.folder, + state_code=entry.state_code, + year=entry.year, + digest=entry.digest, + ) + for entry in plan.state_datasets + ] return ArtifactManifest( manifest_schema=MANIFEST_SCHEMA, country=PRECOMPUTE_COUNTRY, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py index d3fe2d1ba..f05a4c417 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py @@ -20,7 +20,7 @@ from __future__ import annotations -from typing import Any, Literal, Protocol +from typing import Annotated, Any, Literal, Protocol from pydantic import BaseModel, ConfigDict, Field @@ -107,7 +107,7 @@ class WorkSelection(_StrictModel): class ManifestArtifact(_StrictModel): - """One artifact the deploy image must fetch.""" + """One artifact the deploy image must fetch into its data folder.""" type: Literal["dataset", "baseline"] path: str @@ -116,6 +116,27 @@ class ManifestArtifact(_StrictModel): digest: str +class StateDatasetArtifact(_StrictModel): + """One state's year file, which the image fetches into the state's ``folder``. + + Its own model, rather than optional fields on ``ManifestArtifact``, so the + dataset and baseline entries keep exactly their published wire shape. + """ + + type: Literal["state_dataset"] + path: str + filename: str + folder: str + state_code: str + year: int + digest: str + + +ManifestEntry = Annotated[ + ManifestArtifact | StateDatasetArtifact, Field(discriminator="type") +] + + class ArtifactManifest(_StrictModel): """The published deploy manifest (content-addressed in the store).""" @@ -126,7 +147,7 @@ class ArtifactManifest(_StrictModel): manifest_schema: str = Field(alias="schema") country: str receipt: BundleVersionIdentity - artifacts: list[ManifestArtifact] + artifacts: list[ManifestEntry] def canonical_payload(self) -> dict[str, Any]: """Return the exact dict that gets digested and stored, keyed by wire alias.""" diff --git a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py index 52e538a83..744db71c0 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py @@ -247,6 +247,52 @@ def read_manifest(self, digest): ) +def test_a_manifest_listing_state_datasets_reaches_the_fetch_layer(monkeypatch): + install_fake_modal(monkeypatch) + monkeypatch.setenv("POLICYENGINE_VERSION", "4.19.1") + monkeypatch.setenv("POLICYENGINE_CORE_VERSION", "3.27.1") + monkeypatch.setenv("POLICYENGINE_US_VERSION", "1.700.0") + monkeypatch.setenv("POLICYENGINE_UK_VERSION", "2.90.0") + monkeypatch.setenv("POLICYENGINE_MANIFEST_DIGEST", "digest-abc") + monkeypatch.setenv("POLICYENGINE_ARTIFACT_BUCKET", "test-bucket") + + manifest = _fake_manifest() + manifest["artifacts"].append( + { + "type": "state_dataset", + "path": "state-datasets/us/s1/populace_local_year_2026.h5", + "filename": "populace_local_year_2026.h5", + "folder": "state-datasets/populace_local/UT", + "state_code": "UT", + "year": 2026, + "digest": "s1", + } + ) + + class FakeStore: + def __init__(self, bucket_name=None, **_kwargs): + self.bucket_name = bucket_name + + def read_manifest(self, digest): + assert digest == "digest-abc" + return manifest + + from policyengine_simulation_executor import artifact_store + + monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) + sys.modules.pop("src.modal.app", None) + + app = importlib.import_module("src.modal.app") + + (fetch_call,) = [ + call + for call in app.simulation_image.calls + if call[0] == "run_function" and call[1] == "fetch_artifacts" + ] + assert fetch_call[2]["args"] == ("test-bucket", manifest) + sys.modules.pop("src.modal.app", None) + + def test_modal_image_uses_failing_sentinel_without_manifest_digest(monkeypatch): """Without a digest the app imports, but its fetch layer fails the build. diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index 0b1d2bc6e..dd0c78c66 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -203,6 +203,54 @@ def test_manifest_lists_every_artifact_with_runtime_filenames(self): "bl1-bbbb.h5", ] + def test_manifest_lists_state_datasets_after_the_others_in_plan_order(self): + plan = _plan_with_state_datasets() + + manifest = precompute.build_manifest(plan) + + assert [artifact.type for artifact in manifest.artifacts] == [ + "dataset", + "dataset", + "baseline", + "baseline", + "state_dataset", + "state_dataset", + "state_dataset", + ] + 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, + "digest": entry.digest, + } + for entry in plan.state_datasets + ] + + def test_manifest_with_state_datasets_round_trips_its_wire_shape(self): + payload = precompute.build_manifest( + _plan_with_state_datasets() + ).canonical_payload() + + assert ArtifactManifest.model_validate(payload).canonical_payload() == payload + + def test_a_state_dataset_entry_must_name_its_folder_and_state(self): + payload = precompute.build_manifest( + _plan_with_state_datasets() + ).canonical_payload() + del payload["artifacts"][-1]["folder"] + + with pytest.raises(ValidationError): + ArtifactManifest.model_validate(payload) + def test_manifest_payload_is_wire_compatible(self): """The typed manifest must digest identically to the published dict shape. From 794e8f4c8345cbd09face4e6ad8f4ffe56a18845 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 20:33:17 +0400 Subject: [PATCH 09/16] Read a state's baked dataset for US state and district requests Before loading a non-default dataset, the main path's runtime asks state_datasets.baked_state_data_folder for the image folder of the request's state (a district passes its state's code). It returns the folder only when the image holds that state's year file for the requested year and .py's reuse checks (stamp and renames record, legacy SPM record) accept it; a file .py would rebuild is skipped, so ensure_datasets never rebuilds the full regional dataset inside a baked folder. _load_dataset then reads from that folder, and the district filter still applies at run time. Otherwise the request prepares the regional dataset on demand, as before. Runs whose region loads a non-default dataset record dataset_source as baked_state or on_demand. National and default-dataset runs are unchanged, and Stage 12 still loads without a folder. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 6 + .../simulation_runtime.py | 17 ++ .../state_datasets.py | 50 ++++++ .../tests/test_job_progress.py | 2 +- .../tests/test_simulation_output_builder.py | 146 +++++++++++++++++- .../tests/test_state_datasets.py | 128 ++++++++++++++- 6 files changed, 342 insertions(+), 7 deletions(-) diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 72087b956..910c20a1b 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -156,6 +156,12 @@ Operational notes: - Only pure default requests read the baked folder: any request naming an explicit `data` dataset or pinning a `data_version` bypasses it (see `_load_dataset` in `simulation_runtime.py`). +- A US state or congressional-district request whose region loads a + non-default dataset reads its state's baked file when the image has one for + the requested year and .py's year-file reuse checks accept it (see + `state_datasets.baked_state_data_folder`). Otherwise it prepares the full + regional dataset during the request, as before. The run records which in + its `dataset_source` attribute: `baked_state` or `on_demand`. - Known residual, inherited rather than introduced: a data re-release under the same revision labels leaves the cached bundle-install layer — and therefore the receipt every artifact key derives from — unchanged. 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 add789dfe..0f5bfb5b6 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 @@ -29,6 +29,7 @@ ) from policyengine_simulation_observability.telemetry import split_internal_payload +from policyengine_simulation_executor import state_datasets from policyengine_simulation_executor.job_progress import ( NULL_JOB_PROGRESS, JobProgressReporter, @@ -874,10 +875,26 @@ def _run_simulation_impl_core( simulation_params, region_resolution=region_resolution ) with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD)): + baked_folder = state_datasets.baked_state_data_folder( + country=country, + selection=dataset_selection, + state_code=region_resolution.state_code, + year=_parse_year(simulation_params), + data_folder=resolve_data_folder(), + ) + if not dataset_selection.is_default: + runtime.set_context( + dataset_source=( + state_datasets.BAKED_STATE_DATASET_SOURCE + if baked_folder + else state_datasets.ON_DEMAND_DATASET_SOURCE + ) + ) dataset = _load_dataset( simulation_params, selection=dataset_selection, country_module=country_module, + data_folder=baked_folder, ) from policyengine_simulation_executor.uk_local_authority_metadata import ( # noqa: PLC0415 # the UK metadata module is not loaded at import; loaded per call detect_uk_local_authority_metadata, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py index a513bb7ff..ff6199bbd 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py @@ -13,9 +13,23 @@ from __future__ import annotations +import logging from pathlib import Path +from typing import TYPE_CHECKING + +from policyengine_simulation_executor.artifact_keys import dataset_artifact_filename + +if TYPE_CHECKING: + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + +logger = logging.getLogger(__name__) STATE_DATASETS_DIRECTORY = "state-datasets" +# Values of the run's ``dataset_source`` attribute when its region loads a +# non-default dataset: read from the image's baked state file, or prepared +# from the full regional dataset during the request. +BAKED_STATE_DATASET_SOURCE = "baked_state" +ON_DEMAND_DATASET_SOURCE = "on_demand" def state_dataset_relative_folder(dataset: str, state_code: str) -> str: @@ -50,3 +64,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 request should load, or None. + + A request 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``. The file must also pass .py's reuse checks: + otherwise ``ensure_datasets`` would rebuild the full regional dataset + inside the baked folder, so the request goes to the on-demand path + instead. + """ + if country.lower() != "us" or selection.is_default or not state_code: + 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) / dataset_artifact_filename(stem, year) + 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/tests/test_job_progress.py b/projects/policyengine-simulation-executor/tests/test_job_progress.py index 6ad82f818..2a172362e 100644 --- a/projects/policyengine-simulation-executor/tests/test_job_progress.py +++ b/projects/policyengine-simulation-executor/tests/test_job_progress.py @@ -599,7 +599,7 @@ def serialize(self): monkeypatch.setattr( sr, "_load_dataset", - lambda _params, *, country_module, selection: object(), + lambda _params, *, country_module, selection, data_folder=None: object(), ) monkeypatch.setattr( sr, "_build_simulation", lambda *_args, **_kwargs: next(simulations) 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 9306148e4..162c7620a 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -74,6 +74,7 @@ _resolve_dataset_selection, _resolve_region, _run_simulation_impl_core, + resolve_data_folder, run_simulation_impl, ) from policyengine_simulation_executor.simulation_runtime import ( @@ -574,7 +575,7 @@ def serialize(self): ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime._load_dataset", - lambda _params, country_module, selection: dataset, + lambda _params, country_module, selection, data_folder=None: dataset, ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime._build_simulation", @@ -710,7 +711,7 @@ def serialize(self): ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime._load_dataset", - lambda _params, country_module, selection: dataset, + lambda _params, country_module, selection, data_folder=None: dataset, ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime._build_simulation", @@ -1104,7 +1105,7 @@ def serialize(self): ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime._load_dataset", - lambda _params, country_module, selection: dataset, + lambda _params, country_module, selection, data_folder=None: dataset, ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime._build_simulation", @@ -1182,7 +1183,7 @@ def fake_build_simulation( ) return baseline_simulation if len(build_calls) == 1 else reform_simulation - def fake_load_dataset(_params, *, country_module, selection): + def fake_load_dataset(_params, *, country_module, selection, data_folder=None): load_selections.append(selection) return dataset @@ -1232,6 +1233,141 @@ def serialize(self): assert build_calls[1][4] is load_selections[0] +class _StoppedAfterLoadingError(Exception): + """Raised by a stubbed simulation build once the dataset has loaded.""" + + +def _load_through_runtime( + monkeypatch, *, region_resolution, selection, baked_folder, time_period=2026 +): + """Run the annual runtime up to its dataset load and return what it did.""" + from policyengine_simulation_executor import state_datasets + + load_calls = [] + lookups = [] + + def fake_load_dataset(_params, *, country_module, selection, data_folder=None): + load_calls.append({"selection": selection, "data_folder": data_folder}) + return object() + + def fake_baked_state_data_folder(**kwargs): + lookups.append(kwargs) + return baked_folder + + def stop_after_loading(*_args, **_kwargs): + raise _StoppedAfterLoadingError + + 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: selection, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + fake_load_dataset, + ) + monkeypatch.setattr( + state_datasets, "baked_state_data_folder", fake_baked_state_data_folder + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + stop_after_loading, + ) + runtime = _TrackingRuntime() + with pytest.raises(_StoppedAfterLoadingError): + _run_simulation_impl_core( + { + "country": "us", + "region": region_resolution.code, + "time_period": time_period, + "baseline": {}, + "reform": {}, + }, + runtime=runtime, + ) + return SimpleNamespace( + load_calls=load_calls, lookups=lookups, context=dict(runtime.context) + ) + + +def _regional_selection(): + return DatasetSelection( + name="populace_us_2024_acs_local", + uri="hf://policyengine/populace-us/populace_us_2024_acs_local.h5@rev", + is_default=False, + ) + + +def test_runtime_loads_a_state_from_its_baked_folder(monkeypatch): + loaded = _load_through_runtime( + monkeypatch, + region_resolution=RegionResolution( + code="state/ut", dataset_reference="ref", state_code="UT" + ), + selection=_regional_selection(), + baked_folder="/opt/policyengine/data/state-datasets/acs/UT", + ) + + assert loaded.load_calls == [ + { + "selection": _regional_selection(), + "data_folder": "/opt/policyengine/data/state-datasets/acs/UT", + } + ] + assert loaded.lookups == [ + { + "country": "us", + "selection": _regional_selection(), + "state_code": "UT", + "year": 2026, + "data_folder": resolve_data_folder(), + } + ] + assert loaded.context["dataset_source"] == "baked_state" + + +def test_runtime_prepares_a_state_on_demand_without_a_baked_file(monkeypatch): + loaded = _load_through_runtime( + monkeypatch, + region_resolution=RegionResolution( + code="congressional_district/UT-01", + dataset_reference="ref", + state_code="UT", + ), + selection=_regional_selection(), + baked_folder=None, + time_period=2028, + ) + + assert loaded.load_calls[0]["data_folder"] is None + assert loaded.lookups[0]["year"] == 2028 + assert loaded.context["dataset_source"] == "on_demand" + + +def test_runtime_loads_a_national_run_from_the_default_folder(monkeypatch): + bundle = get_country_release_bundle("us") + loaded = _load_through_runtime( + monkeypatch, + region_resolution=RegionResolution(code="us"), + selection=DatasetSelection( + name=bundle.default_dataset, + uri=bundle.default_dataset_uri, + is_default=True, + ), + baked_folder=None, + ) + + assert loaded.load_calls[0]["data_folder"] is None + assert "dataset_source" not in loaded.context + + def test_run_simulation_impl_core_passes_detected_uk_boundary_to_builder(monkeypatch): dataset = SimpleNamespace( data=SimpleNamespace( @@ -1262,7 +1398,7 @@ def serialize(self): ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime._load_dataset", - lambda _params, country_module, selection: dataset, + lambda _params, country_module, selection, data_folder=None: dataset, ) monkeypatch.setattr( "policyengine_simulation_executor.simulation_runtime._build_simulation", diff --git a/projects/policyengine-simulation-executor/tests/test_state_datasets.py b/projects/policyengine-simulation-executor/tests/test_state_datasets.py index dae2ea7d5..a0bcd4911 100644 --- a/projects/policyengine-simulation-executor/tests/test_state_datasets.py +++ b/projects/policyengine-simulation-executor/tests/test_state_datasets.py @@ -1,6 +1,12 @@ -"""The main path's folder layout for baked per-state datasets.""" +"""The main path's baked per-state datasets: their folders and the runtime lookup.""" +from pathlib import Path +from types import SimpleNamespace + +import pytest from policyengine_simulation_executor import state_datasets +from policyengine_simulation_executor.release_bundle import get_country_release_bundle +from policyengine_simulation_executor.simulation_runtime import DatasetSelection def test_state_dataset_relative_folder_names_the_dataset_and_the_state(): @@ -17,3 +23,123 @@ def test_state_dataset_folder_sits_under_the_data_folder(): ) == "/opt/policyengine/data/state-datasets/populace_us_2024_acs_local/UT" ) + + +@pytest.fixture +def regional_selection(): + from fixtures.state_dataset_sources import REGIONAL_DATASET + + bundle = get_country_release_bundle("us") + return DatasetSelection( + name=REGIONAL_DATASET, + uri=bundle.dataset_uris[REGIONAL_DATASET], + is_default=False, + ) + + +@pytest.fixture +def baked_utah(tmp_path): + from fixtures.state_dataset_sources import REGIONAL_DATASET, write_year_file + + folder = Path( + state_datasets.state_dataset_folder(str(tmp_path), REGIONAL_DATASET, "UT") + ) + folder.mkdir(parents=True) + return SimpleNamespace( + data_folder=str(tmp_path), + folder=folder, + year_file=write_year_file(folder, year=2026), + ) + + +def _lookup(baked, 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=baked.data_folder, + ) + + +def test_a_state_reads_its_baked_folder(baked_utah, regional_selection): + assert _lookup(baked_utah, regional_selection) == str(baked_utah.folder) + + +def test_a_district_reads_its_state_folder(baked_utah, regional_selection): + # A district resolves to its state's two-letter code, so it shares the + # state's file and keeps applying its own district filter at run time. + assert _lookup(baked_utah, regional_selection, state_code="UT") == str( + baked_utah.folder + ) + + +def test_a_year_without_a_baked_file_reads_nothing_baked( + baked_utah, regional_selection +): + assert _lookup(baked_utah, regional_selection, year=2028) is None + + +def test_a_state_without_a_baked_folder_reads_nothing_baked( + baked_utah, regional_selection +): + assert _lookup(baked_utah, regional_selection, state_code="WY") is None + + +def test_an_unstamped_file_reads_nothing_baked(tmp_path, regional_selection): + from fixtures.state_dataset_sources import REGIONAL_DATASET, write_year_file + + folder = Path( + state_datasets.state_dataset_folder(str(tmp_path), REGIONAL_DATASET, "UT") + ) + folder.mkdir(parents=True) + write_year_file(folder, year=2026, stamp=False) + + assert ( + state_datasets.baked_state_data_folder( + country="us", + selection=regional_selection, + state_code="UT", + year=2026, + data_folder=str(tmp_path), + ) + is None + ) + + +def test_a_file_with_another_spm_profile_reads_nothing_baked( + baked_utah, regional_selection +): + from policyengine.tax_benefit_models.us.legacy_spm.record import ( + LegacySPMProfile, + LegacySPMRecord, + read_record_value, + write_record, + ) + + stored = LegacySPMRecord.model_validate(read_record_value(baked_utah.year_file)) + write_record( + baked_utah.year_file, + LegacySPMRecord( + profile=LegacySPMProfile( + adapter_revision=stored.profile.adapter_revision + 1, + source=stored.profile.source, + ) + ), + ) + + assert _lookup(baked_utah, regional_selection) is None + + +def test_the_default_dataset_reads_nothing_baked(baked_utah): + bundle = get_country_release_bundle("us") + default = DatasetSelection( + name=bundle.default_dataset, uri=bundle.default_dataset_uri, is_default=True + ) + + assert _lookup(baked_utah, default) is None + + +def test_a_region_without_a_state_reads_nothing_baked(baked_utah, regional_selection): + assert _lookup(baked_utah, regional_selection, state_code=None) is None + assert _lookup(baked_utah, regional_selection, country="uk") is None From e600d36cac2bc8f1a0349573059f13d11257db5f Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 23:55:44 +0400 Subject: [PATCH 10/16] Refuse only district codes that name another state's district The per-state build refused the whole year whenever any household's congressional_district_geoid did not start with its own state FIPS, including codes that name no district, such as an unassigned 0. Such a household matches no district request on either path, so it cannot change a result, yet it would have stopped the precompute and the deploy waiting on it. The build now refuses only codes that name a registered district of a different state. Co-Authored-By: Claude Opus 5.5 --- .../precompute.py | 42 +++++++++++++------ .../tests/test_precompute.py | 27 ++++++++++++ 2 files changed, 57 insertions(+), 12 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 622a89300..1913b280b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -481,24 +481,40 @@ 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). A household whose district code names another state would be - missing from that district's results, so the build stops instead. Both - columns are labels; nothing here is weighted. + number). A household whose code names a district of another state would + be missing from that district's results, so the build stops instead. A + code that names no district, such as an unassigned 0, matches no district + request on either path 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) @@ -557,7 +573,9 @@ 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 = ArtifactStore(bucket) results = [] diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index dd0c78c66..bcdb5fd59 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -1130,12 +1130,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( sr, "_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.sources = sources @@ -1311,6 +1327,17 @@ def test_refuses_a_district_that_lies_outside_its_state(self, state_build): self._build(state_build, self._both_states(state_build)) assert state_build.uploads == [] + 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._both_states(state_build)) + + assert [result.state_code for result in results] == ["UT", "WY"] + def test_refuses_a_state_with_no_households(self, state_build): from policyengine.core.scoping_strategy import RowFilterStrategy From 21089c7f2e30cae32f28efeeb13e5cb4cafb9c0a Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Wed, 21 Oct 2026 23:58:14 +0400 Subject: [PATCH 11/16] Build state datasets per year and dataset, and report what was uploaded run_precompute grouped state entries by year only, while one build prepares a single dataset's year file and refuses entries for two. A bundle sending different states to different datasets would have failed every year's build; builds are now grouped by year and dataset. Each state build result records how long the year's build had run when the file was uploaded, and the precompute log now reports, per build, how many files were built and uploaded, the bytes actually uploaded (files already stored no longer count), and the build time, for tuning the two-hour limit. Co-Authored-By: Claude Opus 5.5 --- .../precompute.py | 30 ++++++----- .../precompute_models.py | 7 ++- .../tests/test_precompute.py | 50 +++++++++++++++++-- .../tests/test_precompute_app.py | 7 ++- 4 files changed, 76 insertions(+), 18 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 1913b280b..ef70dec30 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -552,6 +552,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( @@ -608,6 +609,7 @@ def build_state_datasets_impl( path=entry.path, uploaded=store.upload_file(entry.path, local_file), size_bytes=local_file.stat().st_size, + build_seconds=round(time.monotonic() - started, 1), ) ) return results @@ -841,12 +843,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]] = {} + # One build prepares one dataset's year file, so builds are per year and + # dataset. + state_builds: dict[tuple[int, str], list[dict]] = {} 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_handles = [ - 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 handle in dataset_handles: result = DatasetBuildResult.model_validate(handle.get()) @@ -880,16 +886,16 @@ def run_precompute( ) echo("Determinism gate passed") - if state_handles: - state_results = [ - StateDatasetBuildResult.model_validate(result) - for handle in state_handles - for result in handle.get() + for year, dataset, handle in state_handles: + results = [ + StateDatasetBuildResult.model_validate(result) for result in handle.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" ) manifest_digest = functions.publish_manifest.remote(bucket, plan.model_dump()) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py index f05a4c417..282639668 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py @@ -188,13 +188,18 @@ class DatasetBuildResult(_StrictModel): class StateDatasetBuildResult(_StrictModel): - """Result of one state's dataset build; ``uploaded`` is False if the key existed.""" + """Result of one state's dataset build; ``uploaded`` is False if the key existed. + + ``build_seconds`` is how long the year's build had run when this state's + file was uploaded, so the last result of a build gives its whole duration. + """ year: int state_code: str path: str uploaded: bool size_bytes: int + build_seconds: float class BaselineComputeResult(_StrictModel): diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index bcdb5fd59..b3abea77f 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -384,7 +384,9 @@ def get(self): return Handle() -def _state_build_result(year, state_code, *, uploaded=True, size_bytes=100): +def _state_build_result( + year, state_code, *, uploaded=True, size_bytes=100, build_seconds=30.0 +): return { "year": year, "state_code": state_code, @@ -393,6 +395,7 @@ def _state_build_result(year, state_code, *, uploaded=True, size_bytes=100): ), "uploaded": uploaded, "size_bytes": size_bytes, + "build_seconds": build_seconds, } @@ -485,7 +488,9 @@ def _state_run(self, *, uploaded_wyoming=True): log, { 2026: [ - _state_build_result(2026, "WY", uploaded=uploaded_wyoming), + _state_build_result( + 2026, "WY", uploaded=uploaded_wyoming, build_seconds=40.5 + ), ], 2027: [_state_build_result(2027, "UT", size_bytes=250)], }, @@ -522,12 +527,48 @@ def test_waits_for_state_builds_only_before_publishing(self): assert log.index(("get", "compute_baseline", 2027)) < first_state_get assert first_state_get < log.index(("remote", "publish_manifest", None)) - def test_reports_the_state_dataset_bytes_uploaded(self): + def test_reports_each_builds_uploads_bytes_and_time(self): _functions, _log, _digest, lines = self._state_run(uploaded_wyoming=False) - assert "State datasets: 2 built, 1 uploaded, 350 bytes" in lines + assert ( + "State datasets 2026 populace_local: 1 built, 0 uploaded " + "(0 bytes) in 40.5 s" + ) in lines + assert ( + "State datasets 2027 populace_local: 1 built, 1 uploaded " + "(250 bytes) in 30.0 s" + ) in lines assert lines[-1] == "MANIFEST_DIGEST=digest-123" + def test_spawns_one_state_build_per_year_and_dataset(self): + log = [] + functions = self._functions() + other = _state_dataset_entry(2026, "ID", exists=False).model_copy( + update={"dataset": "populace_other", "folder": "state-datasets/o/ID"} + ) + plan = _plan_with_state_datasets() + plan = plan.model_copy(update={"state_datasets": [*plan.state_datasets, other]}) + functions["plan_artifacts"] = FakeFunction(result=plan.model_dump()) + functions["build_state_datasets"] = RecordingFunction( + "build_state_datasets", + log, + { + 2026: [_state_build_result(2026, "WY")], + 2027: [_state_build_result(2027, "UT")], + }, + ) + + self._run(functions) + + assert [ + (args[1], [entry["dataset"] for entry in args[2]]) + for args in functions["build_state_datasets"].spawn_calls + ] == [ + (2026, ["populace_local"]), + (2027, ["populace_local"]), + (2026, ["populace_other"]), + ] + def test_a_run_without_state_datasets_spawns_no_state_build(self): functions = self._functions() self._run(functions) @@ -1317,6 +1358,7 @@ def test_uploads_each_file_under_its_planned_path(self, state_build): assert [result.uploaded for result in results] == [True, False] assert [result.path for result in results] == [entry.path for entry in entries] assert all(result.size_bytes > 0 for result in results) + assert 0 <= results[0].build_seconds <= results[1].build_seconds def test_refuses_a_district_that_lies_outside_its_state(self, state_build): state_build.entity_frames = state_build.sources.source_entity_frames( diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_precompute_app.py index fabd0ec96..055ab06ee 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute_app.py @@ -119,7 +119,12 @@ def test_state_dataset_wrapper_validates_entries_in_and_dumps_results_out( exists=False, ) result = StateDatasetBuildResult( - year=2026, state_code="UT", path=entry.path, uploaded=True, size_bytes=5 + year=2026, + state_code="UT", + path=entry.path, + uploaded=True, + size_bytes=5, + build_seconds=1.0, ) seen = {} From 54294edec5749d52d564421bfbf9915a28107d5d Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:00:17 +0400 Subject: [PATCH 12/16] Test district loading end to end and state file rows against the source A district request is now followed from region resolution, with the real dataset selection and baked-folder lookup over a stamped Utah file, to the folder _load_dataset is given; the old test only repeated a state lookup. Each state file's rows are now compared with .py's filter applied to the full source frames, instead of with the file filtered again, which would not have noticed another state's rows. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_precompute.py | 34 +++++----- .../tests/test_state_datasets.py | 63 ++++++++++++++++--- 2 files changed, 69 insertions(+), 28 deletions(-) diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index b3abea77f..d31285ae1 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -1310,29 +1310,23 @@ def refuse_to_rebuild(**_kwargs): (dataset,) = loaded.values() assert sorted(dataset.data.household["household_id"]) == [1, 2] - def test_filtering_a_state_file_again_keeps_the_same_rows(self, state_build): - from policyengine.tax_benefit_models.us import load_datasets + def test_each_state_file_holds_exactly_the_rows_its_filter_keeps(self, state_build): from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES - self._build(state_build, self._both_states(state_build)) - utah_target = state_build.targets[0] - utah_folder = state_build.uploads[0][1].parent - (state_file,) = load_datasets( - datasets=[state_build.sources.REGIONAL_DATASET], - years=[self.YEAR], - data_folder=str(utah_folder), - ).values() - - filtered_again = utah_target.scoping_strategy.apply( - entity_data=state_file.data.entity_data, - group_entities=US_GROUP_ENTITIES, - year=self.YEAR, - ) - - for entity, frame in state_file.data.entity_data.items(): - assert list(filtered_again[entity][f"{entity}_id"]) == list( - frame[f"{entity}_id"] + entries = self._both_states(state_build) + self._build(state_build, entries) + stored = dict(state_build.uploads) + + for target, entry in zip(state_build.targets, entries, strict=True): + expected = target.scoping_strategy.apply( + entity_data=state_build.sources.source_entity_frames(), + group_entities=US_GROUP_ENTITIES, + year=self.YEAR, ) + for entity, frame in expected.items(): + assert state_build.sources.stored_ids(stored[entry.path], entity) == [ + int(value) for value in frame[f"{entity}_id"] + ], (target.state_code, entity) def test_prepares_the_full_country_year_file_once(self, state_build): self._build(state_build, self._both_states(state_build)) diff --git a/projects/policyengine-simulation-executor/tests/test_state_datasets.py b/projects/policyengine-simulation-executor/tests/test_state_datasets.py index a0bcd4911..dfd762a83 100644 --- a/projects/policyengine-simulation-executor/tests/test_state_datasets.py +++ b/projects/policyengine-simulation-executor/tests/test_state_datasets.py @@ -1,5 +1,6 @@ """The main path's baked per-state datasets: their folders and the runtime lookup.""" +import contextlib from pathlib import Path from types import SimpleNamespace @@ -66,14 +67,6 @@ def test_a_state_reads_its_baked_folder(baked_utah, regional_selection): assert _lookup(baked_utah, regional_selection) == str(baked_utah.folder) -def test_a_district_reads_its_state_folder(baked_utah, regional_selection): - # A district resolves to its state's two-letter code, so it shares the - # state's file and keeps applying its own district filter at run time. - assert _lookup(baked_utah, regional_selection, state_code="UT") == str( - baked_utah.folder - ) - - def test_a_year_without_a_baked_file_reads_nothing_baked( baked_utah, regional_selection ): @@ -143,3 +136,57 @@ def test_the_default_dataset_reads_nothing_baked(baked_utah): def test_a_region_without_a_state_reads_nothing_baked(baked_utah, regional_selection): assert _lookup(baked_utah, regional_selection, state_code=None) is None assert _lookup(baked_utah, regional_selection, country="uk") is None + + +def test_a_district_request_loads_its_states_baked_file( + monkeypatch, baked_utah, regional_selection +): + """A district resolves to its state's code and reads the state's file.""" + from policyengine.core.scoping_strategy import RowFilterStrategy + from policyengine_simulation_executor import simulation_runtime + + district = SimpleNamespace( + dataset_path=regional_selection.uri, + scoping_strategy=RowFilterStrategy( + variable_name="congressional_district_geoid", variable_value=4901 + ), + parent_code="state/ut", + state_code="UT", + ) + country_module = SimpleNamespace( + model=SimpleNamespace( + version="2.2.1", + get_region={"congressional_district/UT-01": district}.get, + ) + ) + loads = [] + + class _StoppedAfterLoadingError(Exception): + pass + + def load_dataset(_params, *, country_module, selection, data_folder=None): + loads.append(data_folder) + raise _StoppedAfterLoadingError + + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", baked_utah.data_folder) + monkeypatch.setattr( + simulation_runtime, "_country_module", lambda _country: country_module + ) + monkeypatch.setattr(simulation_runtime, "_load_dataset", load_dataset) + + with pytest.raises(_StoppedAfterLoadingError): + simulation_runtime._run_simulation_impl_core( + { + "country": "us", + "region": "congressional_district/UT-01", + "time_period": 2026, + "baseline": {}, + "reform": {}, + }, + runtime=SimpleNamespace( + span=lambda *_args, **_kwargs: contextlib.nullcontext(), + set_context=lambda **_attributes: None, + ), + ) + + assert loads == [str(baked_utah.folder)] From fdaba75c8b5187d67c530c52e32f597897e72e00 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:03:13 +0400 Subject: [PATCH 13/16] Type the state target's row filter and refuse a target without one StateDatasetTarget.scoping_strategy is typed as an optional RowFilterStrategy instead of Any, and a build whose target has none raises a clear error before writing anything, instead of failing on None. A comment in artifact_keys now says that per-state datasets are fetched into their state's folder rather than straight into the data folder. Co-Authored-By: Claude Opus 5.5 --- .../policyengine_simulation_executor/artifact_keys.py | 5 +++-- .../policyengine_simulation_executor/precompute.py | 11 +++++++++-- .../tests/test_precompute.py | 9 +++++++++ 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py index fdfe2fefd..f80b4920b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py @@ -212,8 +212,9 @@ def baseline_simulation_id(baseline_digest: str) -> str: # # The filename inside a dataset object's path is EXACTLY the name the # runtime's ensure_datasets existence check resolves — the fetch layer -# downloads objects to POLICYENGINE_DATA_FOLDER under their basenames, so -# the object basename IS the runtime filename contract. +# downloads each object under its basename, into POLICYENGINE_DATA_FOLDER or, +# for a per-state dataset, into its state's folder there, so the object +# basename IS the runtime filename contract. def dataset_artifact_filename(stem: str, year: int) -> str: diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index ef70dec30..27d8dad01 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -72,6 +72,7 @@ from microdf import MicroDataFrame from policyengine.core import Simulation + from policyengine.core.scoping_strategy import RowFilterStrategy from policyengine_simulation_executor.artifact_store import ArtifactStore from policyengine_simulation_executor.simulation_runtime import RegionResolution @@ -179,7 +180,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() -> list[StateDatasetTarget]: @@ -581,7 +584,11 @@ def build_state_datasets_impl( store = ArtifactStore(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, diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index d31285ae1..35702c926 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -1374,6 +1374,15 @@ def test_a_household_without_a_district_code_does_not_stop_the_build( 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._both_states(state_build)) + assert state_build.uploads == [] + def test_refuses_a_state_with_no_households(self, state_build): from policyengine.core.scoping_strategy import RowFilterStrategy From 3c153872190baa8909e1d853bf7095691402822e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:46:04 +0400 Subject: [PATCH 14/16] Load on demand when .py's year-file reuse checks cannot be imported The baked-folder lookup runs .py's reuse checks, one of which imports a private .py helper. If a .py upgrade moved it, every state and district request would have failed with ImportError while loading its dataset. The lookup now logs a warning and returns no folder, so the request prepares its dataset on demand as it would without a baked file. Co-Authored-By: Claude Opus 5.5 --- .../state_datasets.py | 12 +++++++++++- .../tests/test_state_datasets.py | 11 +++++++++++ 2 files changed, 22 insertions(+), 1 deletion(-) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py index ff6199bbd..5b2adcfe1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py @@ -95,7 +95,17 @@ def baked_state_data_folder( year_file = Path(folder) / dataset_artifact_filename(stem, year) 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 request prepares its dataset on demand 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_state_datasets.py b/projects/policyengine-simulation-executor/tests/test_state_datasets.py index dfd762a83..24d705fc6 100644 --- a/projects/policyengine-simulation-executor/tests/test_state_datasets.py +++ b/projects/policyengine-simulation-executor/tests/test_state_datasets.py @@ -190,3 +190,14 @@ def load_dataset(_params, *, country_module, selection, data_folder=None): ) assert loads == [str(baked_utah.folder)] + + +def test_a_renamed_dotpy_reuse_check_falls_back_to_on_demand_loading( + monkeypatch, baked_utah, regional_selection +): + """A .py upgrade moving the private SPM profile helper must not fail requests.""" + from policyengine.tax_benefit_models.us import datasets as us_datasets + + monkeypatch.delattr(us_datasets, "_expected_legacy_spm_profile") + + assert _lookup(baked_utah, regional_selection) is None From b7d54a55017a63646aeaedc654bbb4142fca1b40 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 00:46:04 +0400 Subject: [PATCH 15/16] Refuse a state without a row filter before preparing the year file The check for a state target without a row filter ran inside the per-state loop, after the full-country year file was prepared and earlier states were uploaded. _checked_state_targets now refuses it with the other in-container checks, before anything is prepared, and returns each state's row filter. Docstrings now say builds run one container per year and dataset. Co-Authored-By: Claude Opus 5.5 --- .../src/modal/precompute_app.py | 7 +++-- .../precompute.py | 31 +++++++++++-------- .../tests/test_precompute.py | 1 + 3 files changed, 23 insertions(+), 16 deletions(-) diff --git a/projects/policyengine-simulation-executor/src/modal/precompute_app.py b/projects/policyengine-simulation-executor/src/modal/precompute_app.py index 48d048861..3e63cda39 100644 --- a/projects/policyengine-simulation-executor/src/modal/precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/precompute_app.py @@ -11,9 +11,10 @@ Wave 1 builds the single-year US datasets (one container per year), wave 2 computes the 20 per-cohort national baselines per year (one container per cohort, sized like the production ``run_simulation_segment`` workers). -Alongside both waves, one container per year filters the regional dataset -that state regions load into one year file per state. See the library module -for the idempotency, ``--force``, and determinism-gate semantics. +Alongside both waves, one container per year and dataset filters the +regional dataset that state regions load into one year file per state. See +the library module for the idempotency, ``--force``, and determinism-gate +semantics. """ from __future__ import annotations diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 27d8dad01..f6bbce204 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -432,13 +432,14 @@ def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildR def _checked_state_targets( year: int, expected: list[StateDatasetPlanEntry] -) -> dict[str, StateDatasetTarget]: - """Return each planned state's target, after re-deriving its key in this container. +) -> dict[str, RowFilterStrategy]: + """Return each planned state's row filter, after re-deriving its key here. The planner and this builder run in different containers; both derive a state's key from the installed bundle. A disagreement means they see different bundles, so nothing is built rather than uploading a file under - a key that does not describe it. + a key that does not describe it. A state without a row filter is refused + here too, before anything is prepared. """ from policyengine_simulation_executor.artifact_keys import ( # noqa: PLC0415 # resolved per call so tests can patch artifact_keys.collect_state_dataset_identity collect_state_dataset_identity, @@ -457,6 +458,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( PRECOMPUTE_COUNTRY, year, @@ -481,7 +485,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]: @@ -556,7 +565,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-state-datasets-") @@ -584,11 +593,7 @@ def build_state_datasets_impl( store = ArtifactStore(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, @@ -826,9 +831,9 @@ def run_precompute( Takes the six Modal function handles as ``functions`` (each offering ``.remote``/``.spawn``) so the flow is testable with plain fakes. - Per-state dataset builds (one container per year) start with the dataset - builds and are waited for only before the manifest is published, so the - baselines never wait for them. + Per-state dataset builds (one container per year and dataset) start with + the dataset builds and are waited for only before the manifest is + published, so the baselines never wait for them. Everything crossing a handle is a plain dict (Modal's serialization boundary); it is validated back into models immediately on return. The determinism gate runs only when baselines were actually computed diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index 35702c926..aca03f6d6 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -1382,6 +1382,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._both_states(state_build)) assert state_build.uploads == [] + assert state_build.ensure_calls == [] def test_refuses_a_state_with_no_households(self, state_build): from policyengine.core.scoping_strategy import RowFilterStrategy From aaa97c8df1d3398374a7c24fc74f5dd814c480c2 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 22 Oct 2026 01:47:52 +0400 Subject: [PATCH 16/16] Fix the executor Pyright errors this branch adds Main now runs Pyright on every package (#864). Against the executor's baseline, this branch added seven errors: - The run's new reads of `dataset_selection`, `region_resolution` and `simulation_params` were reported as possibly unbound, because `runtime.span(...)` is typed as possibly suppressing exceptions. The run's spans now open with `scopes.span(runtime, ...)`, as docs/engineering/skills/type-checking.md says. `_build_simulation` is annotated with the `StageReportingSimulation` it returns, which the progress callback needs. - `build_state_datasets` couldn't be passed as a `RemoteFunction`, because Pyright doesn't resolve `modal` in src/modal to the Modal package. The precompute app casts its six Modal functions to that protocol. - The manifest's artifact list is built with `ManifestEntry` elements, and the district-code check passes `isin` a sorted list. 54 baseline entries no longer occur, so the baseline drops from 178 to 124 entries. Co-Authored-By: Claude Opus 5.5 --- docs/generated/rules.md | 2 +- .../src/modal/precompute_app.py | 20 ++++--- .../precompute.py | 5 +- .../simulation_runtime.py | 23 ++++---- rules/baselines/pyright-executor.txt | 54 ------------------- rules/registry.toml | 2 +- 6 files changed, 32 insertions(+), 74 deletions(-) diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 5223e0d95..52a3fa432 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt` - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt` - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt` -- `sim.quality.pyright-executor`: 178 remaining in `rules/baselines/pyright-executor.txt` +- `sim.quality.pyright-executor`: 124 remaining in `rules/baselines/pyright-executor.txt` - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt` ### Exceptions and review dates diff --git a/projects/policyengine-simulation-executor/src/modal/precompute_app.py b/projects/policyengine-simulation-executor/src/modal/precompute_app.py index 3e63cda39..3153b2056 100644 --- a/projects/policyengine-simulation-executor/src/modal/precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/precompute_app.py @@ -19,6 +19,8 @@ from __future__ import annotations +from typing import TYPE_CHECKING, cast + from src.modal.app import ( build_runtime_simulation_image, gcp_secret, @@ -27,6 +29,9 @@ import modal +if TYPE_CHECKING: + from policyengine_simulation_executor.precompute_models import RemoteFunction + app = modal.App("policyengine-simulation-precompute") precompute_image = build_runtime_simulation_image().add_local_python_source( @@ -192,12 +197,15 @@ def main(*, force: bool = False) -> None: run_precompute( resolve_bucket_name(), force=force, + # Pyright doesn't resolve `modal` here to the Modal package (it reports + # `modal.App` as unknown), so it types each decorated function as the + # plain function it wraps. At run time each is a Modal function. functions=PrecomputeFunctions( - 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, + plan_artifacts=cast("RemoteFunction", plan_artifacts), + build_dataset=cast("RemoteFunction", build_dataset), + build_state_datasets=cast("RemoteFunction", build_state_datasets), + compute_baseline=cast("RemoteFunction", compute_baseline), + verify_determinism=cast("RemoteFunction", verify_determinism), + publish_manifest=cast("RemoteFunction", publish_manifest), ), ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index f6bbce204..0e87f9dc8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py @@ -260,7 +260,8 @@ def build_manifest(plan: PrecomputePlan) -> ArtifactManifest: digest=entry.digest, ) for entry in plan.datasets - ] + [ + ] + artifacts += [ ManifestArtifact( type="baseline", path=entry.path, @@ -519,7 +520,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(): 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 0f5bfb5b6..154f13002 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 @@ -23,6 +23,7 @@ from policyengine_simulation_contract.dataset_uri import runtime_dataset_uri from policyengine_simulation_contract.job_progress import JobStage from policyengine_simulation_contract.spm import SPMInputError, spm_error_detail +from policyengine_simulation_observability import scopes from policyengine_simulation_observability.stages import ( ANNUAL_IMPACT_STAGES, Stage, @@ -46,7 +47,7 @@ from collections.abc import Iterator from types import ModuleType - from policyengine.core import Dataset, Policy, Region, Simulation + from policyengine.core import Dataset, Policy, Region from policyengine.core.scoping_strategy import ScopingStrategy from policyengine.provenance.manifest import ArtifactPathReference from policyengine_observability import ObservabilityRuntime @@ -766,7 +767,7 @@ def _build_simulation( dataset_selection: DatasetSelection, policy: dict[str, Any] | None, region: RegionResolution, -) -> Simulation: +) -> StageReportingSimulation: from policyengine_simulation_executor.baseline_artifacts import ( ArtifactBaselineSimulation, BaselineIdentityRequest, @@ -832,7 +833,7 @@ def _run_simulation_impl_core( runtime: ObservabilityRuntime, progress: JobProgressReporter = NULL_JOB_PROGRESS, ) -> dict: - with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.REQUEST_PARSE)): + with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.REQUEST_PARSE)): simulation_params, telemetry, metadata = split_internal_payload(params) metadata = metadata or {} from policyengine_simulation_executor.spm import ( @@ -861,20 +862,20 @@ def _run_simulation_impl_core( country=country, ) - with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): + with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.COUNTRY_MODULE_LOAD)): country_module = _country_module(country) - with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.REGION_RESOLUTION)): + with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.REGION_RESOLUTION)): region_resolution = _resolve_region( country_module=country_module, country=country, params=simulation_params, ) runtime.set_context(region=region_resolution.code) - with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.DATASET_RESOLUTION)): + with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.DATASET_RESOLUTION)): dataset_selection = _resolve_dataset_selection( simulation_params, region_resolution=region_resolution ) - with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD)): + with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.DATASET_LOAD)): baked_folder = state_datasets.baked_state_data_folder( country=country, selection=dataset_selection, @@ -906,12 +907,13 @@ def _run_simulation_impl_core( _require_certified_uk_weight_matrix( region_resolution.scoping_strategy, dataset, dataset_selection ) - with runtime.span(ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): + with scopes.span(runtime, ANNUAL_IMPACT_STAGES.name(Stage.POLICY_NORMALIZATION)): baseline_policy = _normalise_policy(simulation_params.get("baseline")) reform_policy = _normalise_policy(simulation_params.get("reform")) logger.info("Initialising baseline and reform simulations") - with runtime.span( + with scopes.span( + runtime, ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_BUILD), attributes={"simulation_kind": "baseline"}, ): @@ -922,7 +924,8 @@ def _run_simulation_impl_core( policy=baseline_policy, region=region_resolution, ) - with runtime.span( + with scopes.span( + runtime, ANNUAL_IMPACT_STAGES.name(Stage.SIMULATION_BUILD), attributes={"simulation_kind": "reform"}, ): diff --git a/rules/baselines/pyright-executor.txt b/rules/baselines/pyright-executor.txt index a1f792020..181b6e689 100644 --- a/rules/baselines/pyright-executor.txt +++ b/rules/baselines/pyright-executor.txt @@ -37,11 +37,6 @@ projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py : projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "state" is possibly unbound projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportPossiblyUnboundVariable :: "state" is possibly unbound projects/policyengine-simulation-executor/src/modal/budget_window_scheduler.py :: reportReturnType :: Function with declared return type "dict[str, Any]" must return value on all code paths -projects/policyengine-simulation-executor/src/modal/precompute_app.py :: reportArgumentType :: Argument of type "(bucket: str) -> dict[Unknown, Unknown]" cannot be assigned to parameter "plan_artifacts" of type "RemoteFunction" in function "__init__" -projects/policyengine-simulation-executor/src/modal/precompute_app.py :: reportArgumentType :: Argument of type "(bucket: str, expected: dict[Unknown, Unknown]) -> dict[Unknown, Unknown]" cannot be assigned to parameter "build_dataset" of type "RemoteFunction" in function "__init__" -projects/policyengine-simulation-executor/src/modal/precompute_app.py :: reportArgumentType :: Argument of type "(bucket: str, expected: dict[Unknown, Unknown]) -> dict[Unknown, Unknown]" cannot be assigned to parameter "compute_baseline" of type "RemoteFunction" in function "__init__" -projects/policyengine-simulation-executor/src/modal/precompute_app.py :: reportArgumentType :: Argument of type "(bucket: str, expected: dict[Unknown, Unknown]) -> dict[Unknown, Unknown]" cannot be assigned to parameter "verify_determinism" of type "RemoteFunction" in function "__init__" -projects/policyengine-simulation-executor/src/modal/precompute_app.py :: reportArgumentType :: Argument of type "(bucket: str, plan: dict[Unknown, Unknown]) -> str" cannot be assigned to parameter "publish_manifest" of type "RemoteFunction" in function "__init__" projects/policyengine-simulation-executor/src/modal/precompute_app.py :: reportAttributeAccessIssue :: "App" is not a known attribute of module "modal" projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportAttributeAccessIssue :: "Function" is not a known attribute of module "modal" projects/policyengine-simulation-executor/src/modal/segmented_national.py :: reportAttributeAccessIssue :: "FunctionCall" is not a known attribute of module "modal" @@ -109,60 +104,11 @@ projects/policyengine-simulation-executor/src/policyengine_simulation_executor/s projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportAttributeAccessIssue :: Cannot access attribute "household" for class "BaseModel" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "data" is not a known attribute of "None" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_output_geographic.py :: reportOptionalMemberAccess :: "household" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportArgumentType :: Argument of type "Unbound | Simulation" cannot be assigned to parameter "simulation" of type "StageReportingSimulation" in function "_report_calculation_start" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportArgumentType :: Argument of type "Unbound | Simulation" cannot be assigned to parameter "simulation" of type "StageReportingSimulation" in function "_report_calculation_start" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportArgumentType :: Argument of type "str" cannot be assigned to parameter "region_type" of type "RegionType" in function "__init__" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAssignmentType :: Type "dict[str, Any] | Policy | None" is not assignable to declared type "dict[str, Any] | None" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportAttributeAccessIssue :: Cannot access attribute "entity_data" for class "BaseModel" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportOptionalMemberAccess :: "entity_data" is not a known attribute of "None" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportOptionalMemberAccess :: "items" is not a known attribute of "None" -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "baseline" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "baseline" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "baseline" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "baseline" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "baseline" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "baseline_policy" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "country_module" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "country_module" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "country_module" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "dataset" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "dataset" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "dataset" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "dataset" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "dataset" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "dataset_selection" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "dataset_selection" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "dataset_selection" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "dataset_selection" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "metadata" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "reform" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "reform" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "reform" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "reform" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "reform_policy" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "region_resolution" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "region_resolution" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "region_resolution" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "region_resolution" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "region_resolution" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "region_resolution" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "region_resolution" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "simulation_params" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "telemetry" is possibly unbound -projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py :: reportPossiblyUnboundVariable :: "telemetry" is possibly unbound projects/policyengine-simulation-executor/src/policyengine_simulation_executor/spm.py :: reportArgumentType :: Argument of type "object" cannot be assigned to parameter "defaults" of type "SPMSelection" in function "__init__" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py :: reportAttributeAccessIssue :: "Function" is not a known attribute of module "modal" projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py :: reportAttributeAccessIssue :: "Function" is not a known attribute of module "modal" diff --git a/rules/registry.toml b/rules/registry.toml index b3ab3490b..618222fb7 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -1139,7 +1139,7 @@ job = "pyright" invocation = "python -m scripts.rules pyright" [rule.baseline] file = "rules/baselines/pyright-executor.txt" -count = 178 +count = 124 [rule.proof] bad_fixture = "rules/fixtures/sim.quality.pyright-executor/bad" good_fixture = "rules/fixtures/sim.quality.pyright-executor/good"