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/README.md b/projects/policyengine-simulation-executor/README.md index 64f2f60f8..910c20a1b 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; @@ -150,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/fixtures/identity_stubs.py b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py index 307a6f5ec..bbd7b7f8c 100644 --- a/projects/policyengine-simulation-executor/fixtures/identity_stubs.py +++ b/projects/policyengine-simulation-executor/fixtures/identity_stubs.py @@ -27,14 +27,21 @@ 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 # 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, ) @@ -51,8 +58,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 +72,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 +90,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/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/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/src/modal/precompute_app.py b/projects/policyengine-simulation-executor/src/modal/precompute_app.py index b3b30a7db..3153b2056 100644 --- a/projects/policyengine-simulation-executor/src/modal/precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/precompute_app.py @@ -10,13 +10,17 @@ 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 +cohort, sized like the production ``run_simulation_segment`` workers). +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 +from typing import TYPE_CHECKING, cast + from src.modal.app import ( build_runtime_simulation_image, gcp_secret, @@ -25,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( @@ -86,6 +93,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, @@ -161,11 +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, - 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/artifact_keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py index ee9026334..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 @@ -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,14 +205,16 @@ def baseline_simulation_id(baseline_digest: str) -> str: # --- Store layout --------------------------------------------------------- # # datasets///_year_.h5 +# state-datasets///_year_.h5 # baselines///.h5 # manifests/.json # deployed/.json # # 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: @@ -187,6 +237,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 +319,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 +490,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/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute.py index 2fd267cd6..0e87f9dc8 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 @@ -51,16 +53,26 @@ DatasetPlanEntry, DeterminismVerdict, ManifestArtifact, + ManifestEntry, PrecomputePlan, RemoteFunction, + StateDatasetArtifact, + 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.core.scoping_strategy import RowFilterStrategy from policyengine_simulation_executor.artifact_store import ArtifactStore from policyengine_simulation_executor.simulation_runtime import RegionResolution @@ -152,17 +164,94 @@ 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``. + ``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 + state_code: str + state_fips: int + filter_variable: str + dataset: str + scoping_strategy: RowFilterStrategy | None = field( + default=None, compare=False, repr=False + ) + + +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, + scoping_strategy=strategy, + ) + ) + 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 + ], ) 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, @@ -171,7 +260,8 @@ def build_manifest(plan: PrecomputePlan) -> ArtifactManifest: digest=entry.digest, ) for entry in plan.datasets - ] + [ + ] + artifacts += [ ManifestArtifact( type="baseline", path=entry.path, @@ -181,6 +271,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, @@ -193,6 +295,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 +313,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 +341,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 +374,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: @@ -297,6 +431,203 @@ def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildR ) +def _checked_state_targets( + year: int, expected: list[StateDatasetPlanEntry] +) -> 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 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, + ) + + 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) + 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, + 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) + 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]: + """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 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 + district = household["congressional_district_geoid"] + in_another_state = district.isin(sorted(district_geoids)) & ( + (district // 100) != household["state_fips"] + ) + if in_another_state.any(): + msg = ( + "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) + + +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 [] + started = time.monotonic() + strategies = _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, _registered_district_geoids(country_module) + ) + + store = ArtifactStore(bucket) + results = [] + for entry in expected: + rows = strategies[entry.region].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, + build_seconds=round(time.monotonic() - started, 1), + ) + ) + return results + + def _prepare_cohort_baseline( bucket: str, expected: BaselinePlanEntry ) -> tuple[ArtifactStore, Path, Simulation]: @@ -480,10 +811,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 @@ -498,8 +830,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 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 @@ -512,7 +847,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)" ) @@ -520,6 +856,17 @@ def run_precompute( functions.build_dataset.spawn(bucket, entry.model_dump()) for entry in work.datasets ] + # 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_builds.setdefault((entry.year, entry.dataset), []).append( + entry.model_dump() + ) + state_handles = [ + (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()) echo( @@ -552,6 +899,18 @@ def run_precompute( ) echo("Determinism gate passed") + 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 {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()) echo(manifest_digest_line(manifest_digest)) return manifest_digest 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..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 @@ -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 @@ -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,10 +103,11 @@ class WorkSelection(_StrictModel): datasets: list[DatasetPlanEntry] baselines: list[BaselinePlanEntry] + state_datasets: list[StateDatasetPlanEntry] = Field(default_factory=list) 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 @@ -96,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).""" @@ -106,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.""" @@ -146,6 +187,21 @@ class DatasetBuildResult(_StrictModel): size_bytes: int +class StateDatasetBuildResult(_StrictModel): + """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): """Outcome of one cohort baseline computation and its write-once upload. 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..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,12 +23,14 @@ 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, ) 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, @@ -45,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 @@ -73,12 +75,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 +610,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, ) @@ -721,15 +727,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} @@ -753,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, @@ -819,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 ( @@ -848,24 +862,40 @@ 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, + 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, @@ -877,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"}, ): @@ -893,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/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..5b2adcfe1 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/state_datasets.py @@ -0,0 +1,112 @@ +"""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 + +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: + """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)) + + +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)) + + +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 + 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 + return folder 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_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_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_modal_bundle_image.py b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py index 2ad12baec..744db71c0 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 @@ -245,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 5bda67860..aca03f6d6 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] @@ -157,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. @@ -263,6 +357,48 @@ 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, build_seconds=30.0 +): + 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, + "build_seconds": build_seconds, + } + + class TestRunPrecompute: def _functions(self, verify_result=None): return { @@ -270,6 +406,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)} ), @@ -340,6 +477,103 @@ 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, build_seconds=40.5 + ), + ], + 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_each_builds_uploads_bytes_and_time(self): + _functions, _log, _digest, lines = self._state_run(uploaded_wyoming=False) + + 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) + assert functions["build_state_datasets"].spawn_calls == [] + def test_malformed_worker_result_fails_loudly(self): functions = self._functions() functions["build_dataset"] = FakeFunction( @@ -496,6 +730,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 +777,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 +874,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 @@ -676,6 +1091,346 @@ 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 + ) + + 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 + 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_each_state_file_holds_exactly_the_rows_its_filter_keeps(self, state_build): + from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES + + 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)) + + 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) + 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( + 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_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_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 == [] + assert state_build.ensure_calls == [] + + 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): diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_precompute_app.py index 16d710103..055ab06ee 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,61 @@ 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, + build_seconds=1.0, + ) + 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 +298,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, 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..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", @@ -1531,6 +1667,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, @@ -1618,6 +1815,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, 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..24d705fc6 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_state_datasets.py @@ -0,0 +1,203 @@ +"""The main path's baked per-state datasets: their folders and the runtime lookup.""" + +import contextlib +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(): + 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" + ) + + +@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_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 + + +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)] + + +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 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"