diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index fed0d09e4..953ca0d70 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -376,22 +376,56 @@ baselines, so running again after a failed check repeats the check rather than skipping it. A pass is recorded in the cache bucket, and the manifest is published only when that record exists. -Precompute publishes a canonical manifest containing all 63 files. The -deployment passes that manifest's SHA-256 digest into the US image build. The -image build downloads and verifies every listed file, then records the exact -cache release in worker-validation metadata and the v2 version manifest. The UK +Precompute also builds per-state datasets. For each US state region whose +requests load a dataset other than the bundle default (the ACS-local dataset +since policyengine 6.7), it prepares that dataset's year file for each +precomputed year once, keeps one state's households with the region's own row +filter, and stores each state's file under the year file's own name. It +refuses data in which a household's district code names another state's +district, because district runs read their state's file. + +Precompute publishes a canonical manifest containing the 63 national files and +the state datasets. The deployment passes that manifest's SHA-256 digest into +the US image build. The image build downloads and verifies every listed file, +placing each state dataset in `state-datasets///` under the +data folder, then records the exact cache release in worker-validation +metadata and the v2 version manifest. A state or district calculation reads +its state's file when the image holds it for the requested year and .py would +reuse it, and prepares the dataset during the calculation otherwise. The UK and coordinator images do not contain these files. Deployment validation runs one real 2026 region-group baseline and requires it to load the precomputed file without recalculation. Readers released before cache releases reject fields they do not know, and the manifest is published before the Cloud Run revision that reads it is -promoted. The publisher therefore writes the manifest under the -`active-with-cache` key and the same manifest without cache releases under the -`active` key, which those readers keep reading. Current readers read -`active-with-cache`, falling back to `active` only before the first publish -that writes both keys. Removing the `active` copy needs a later change, once -no deployed or rollback revision reads it. +promoted. The publisher therefore writes the whole manifest under the +`active-with-region-datasets` key, the same manifest without each country +bundle's `region_datasets` under the `active-with-cache` key, and without cache +releases either under the `active` key; readers released before each field keep +reading the key that lacks it. Current readers use the +`active-with-region-datasets` copy only while it, with its region datasets +removed, equals the manifest under `active-with-cache`; otherwise they read +`active-with-cache`, falling back to `active` before the first publish that +writes it. Code released before region datasets republishes only the two older +keys, so after a rollback the out-of-date region dataset copy stops matching and +is no longer read. Removing the older copies needs a later change, once no +deployed or rollback revision reads them. + +A rollback to code released before region datasets has one window: the older +worker is deployed before its publish job runs, and until that publish the +region dataset copy still matches, so updated entry services keep sending +regional provenance for state and district runs, which the older worker refuses. +Only Stage 12 comparison runs fail in that window, never production responses; +if the older deploy's validation fails, the window lasts until a later publish. +To avoid it, delete the `active-with-region-datasets` key from the +`simulation-api-v2-version-manifest` Modal Dict before deploying the older code. + +`region_datasets` names the dataset each region type loads. The entry service +records that dataset as a state or district run's dataset provenance, and the +worker refuses a calculation whose provenance names another dataset than the one +it loads, except that it still accepts the default dataset from entry services +released before regional provenance. The bundle digest leaves `region_datasets` +out, so adding it did not change the digest of a running application. Only an empty-policy baseline using the bundle's default dataset and an exact region-group scope can read a baseline cache file. Reform calculations always diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py index 5c406f742..2270497e4 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py @@ -10,6 +10,33 @@ Sha256Digest = Annotated[str, Field(pattern=r"^[0-9a-f]{64}$")] CountryId = Literal["us", "uk"] SUPPORTED_COUNTRIES: tuple[CountryId, ...] = ("us", "uk") +# A bare two-letter US region is a state code, such as "ut". +_US_STATE_CODE_LENGTH = 2 + + +def region_type_of(country: str, region: str | None) -> str | None: + """Return the region type a requested region code names, as bundles key it. + + Codes are read as the simulation runtime normalises them: surrounding + whitespace is ignored and the country itself, or no region, is + ``national``. A bare two-letter US code is a state and a bare UK name is + a UK country. Otherwise the type is the part before the first slash, in + lower case: ``State/UT`` is ``state`` and ``congressional_district/UT-01`` + is ``congressional_district``. None means the code names no region type. + """ + raw = (region or "").strip() + if not raw or raw.lower() == country: + return "national" + if raw.lower() in SUPPORTED_COUNTRIES: + # Another country's code names no region of this one. + return None + if "/" in raw: + return raw.split("/", 1)[0].lower() + if country == "us" and len(raw) == _US_STATE_CODE_LENGTH: + return "state" + if country == "uk": + return "country" + return None class StrictBundleModel(BaseModel): @@ -31,8 +58,23 @@ class Stage12Dataset(StrictBundleModel): repo_type: NonEmptyText +class Stage12RegionDataset(StrictBundleModel): + """The dataset one type of region loads, such as ``state`` or ``national``. + + ``dataset`` is a dataset identity in the country's bundle, or None when the + bundle gives a per-region file pattern instead of one dataset. + """ + + region_type: NonEmptyText + dataset: NonEmptyText | None + + class Stage12CountryBundle(StrictBundleModel): - """One country's pinned model package, data release and datasets within a bundle.""" + """One country's pinned model package, data release and datasets within a bundle. + + ``region_datasets`` names the dataset each type of region loads. It is None + in bundles published before it existed. + """ country: CountryId country_package_name: NonEmptyText @@ -45,6 +87,18 @@ class Stage12CountryBundle(StrictBundleModel): default_dataset: NonEmptyText default_dataset_uri: NonEmptyText datasets: tuple[Stage12Dataset, ...] + region_datasets: tuple[Stage12RegionDataset, ...] | None = None + + def dataset_for_region_type(self, region_type: str | None) -> str | None: + """Return the dataset identity ``region_type`` regions load, or None.""" + for item in self.region_datasets or (): + if item.region_type == region_type: + return item.dataset + return None + + def dataset_for_region(self, region: str | None) -> str | None: + """Return the dataset identity a requested region code loads, or None.""" + return self.dataset_for_region_type(region_type_of(self.country, region)) class Stage12BundleManifest(StrictBundleModel): diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index 75bd43790..a10d414cc 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -21,9 +21,14 @@ V1_ROUTING_STATE_NAME = "simulation-api-routing-state" V2_VERSION_MANIFEST_NAME = "simulation-api-v2-version-manifest" +# The complete manifest, including each country bundle's region datasets. +REGION_DATASETS_MANIFEST_KEY = "active-with-region-datasets" +# Readers released before region datasets forbid unknown fields, so they keep +# reading this key, which holds the same manifest without them. ACTIVE_MANIFEST_KEY = "active-with-cache" # Readers released before Stage 12 cache releases forbid unknown fields, so -# they keep reading this key, which holds the same manifest without them. +# they keep reading this key, which holds the same manifest without them or +# region datasets. PRE_CACHE_MANIFEST_KEY = "active" V2_APPLICATION_PREFIX = "policyengine-simulation-v2-py" V1_APPLICATION_PREFIX = "policyengine-simulation-py" @@ -219,20 +224,67 @@ def _is_newer(candidate: str, current: str | None) -> bool: return False +_REGION_DATASETS_FIELD = { + "bundle": {"countries": {"__all__": {"region_datasets"}}}, +} + + +def pre_region_datasets_manifest_payload( + manifest: V2VersionManifest, +) -> dict[str, object]: + """Return the manifest as readers released before region datasets accept it.""" + return manifest.model_dump( + mode="json", + exclude={"versions": {"__all__": _REGION_DATASETS_FIELD}}, + ) + + def pre_cache_manifest_payload(manifest: V2VersionManifest) -> dict[str, object]: """Return the manifest as readers released before cache releases accept it.""" return manifest.model_dump( mode="json", - exclude={"versions": {"__all__": {"cache": True, "validation": {"cache"}}}}, + exclude={ + "versions": { + "__all__": { + "cache": True, + "validation": {"cache"}, + **_REGION_DATASETS_FIELD, + } + } + }, ) +def _is_current_region_dataset_copy(value: object, cache_copy: object) -> bool: + """Return whether ``value`` is the copy the last publish wrote with ``cache_copy``. + + Code released before region datasets publishes only the older keys, for + example after a rollback, which leaves an out-of-date region dataset copy + behind. A current copy, with its region datasets removed, is exactly the + manifest under ``active-with-cache``. + """ + if value is None: + return False + if cache_copy is None: + return True + try: + manifest = V2VersionManifest.model_validate(value) + except (TypeError, ValueError): + return False + return pre_region_datasets_manifest_payload(manifest) == cache_copy + + def _current_manifest(store: V2ManifestStore) -> object: - current = store.get(ACTIVE_MANIFEST_KEY) - if current is None: - # Until the first publish that writes both keys. - current = store.get(PRE_CACHE_MANIFEST_KEY) - return current + # The newest key the last publish wrote: each older key holds the same + # manifest without the fields its readers do not accept. + cache_copy = store.get(ACTIVE_MANIFEST_KEY) + region_dataset_copy = store.get(REGION_DATASETS_MANIFEST_KEY) + if _is_current_region_dataset_copy(region_dataset_copy, cache_copy): + return region_dataset_copy + if cache_copy is not None: + return cache_copy + # Until the first publish that writes the cache copy. + return store.get(PRE_CACHE_MANIFEST_KEY) def build_next_v2_manifest( @@ -286,9 +338,11 @@ def publish_v2_manifest( worker=worker, force_default=force_default, ) - # Each key holds one entire validated document, and a reader reads only - # one key, so no reader observes a partial update assembled from several. - store[ACTIVE_MANIFEST_KEY] = next_manifest.model_dump(mode="json") + # Each key holds one entire validated document, and a reader uses one + # whole document (the region dataset copy only while it matches the cache + # copy), so no reader observes a partial update assembled from several. + store[REGION_DATASETS_MANIFEST_KEY] = next_manifest.model_dump(mode="json") + store[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(next_manifest) store[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(next_manifest) return next_manifest diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index 82bca08ba..78c1f908e 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -10,10 +10,13 @@ Stage12BundleManifest, Stage12CountryBundle, Stage12Dataset, + Stage12RegionDataset, + region_type_of, ) from policyengine_simulation_contract.stage12_manifest import ( ACTIVE_MANIFEST_KEY, PRE_CACHE_MANIFEST_KEY, + REGION_DATASETS_MANIFEST_KEY, V1_ROUTING_STATE_NAME, V2_VERSION_MANIFEST_NAME, Stage12CacheRelease, @@ -24,6 +27,7 @@ assert_separate_manifest_names, manifest_digest, pre_cache_manifest_payload, + pre_region_datasets_manifest_payload, publish_v2_manifest, v2_application_name, validate_manifest_has_no_credentials, @@ -159,13 +163,39 @@ def test_v2_application_name_rejects_versions_a_modal_name_cannot_hold( v2_application_name(version) +def _regional_worker(version: str = "5.2.0") -> V2WorkerVersion: + """Return a cached worker whose US bundle names each region type's dataset.""" + worker = _cached_worker(version) + countries = tuple( + country.model_copy( + update={ + "region_datasets": ( + Stage12RegionDataset( + region_type="national", dataset=country.default_dataset + ), + Stage12RegionDataset( + region_type="state", dataset=f"{country.default_dataset}_local" + ), + ) + } + ) + if country.country == "us" + else country + for country in worker.bundle.countries + ) + return worker.model_copy( + update={"bundle": worker.bundle.model_copy(update={"countries": countries})} + ) + + def test_publish_writes_one_complete_v2_document_per_key() -> None: store = FakeStore() - manifest = publish_v2_manifest(store=store, worker=_cached_worker()) + manifest = publish_v2_manifest(store=store, worker=_regional_worker()) assert store.writes == [ - (ACTIVE_MANIFEST_KEY, manifest.model_dump(mode="json")), + (REGION_DATASETS_MANIFEST_KEY, manifest.model_dump(mode="json")), + (ACTIVE_MANIFEST_KEY, pre_region_datasets_manifest_payload(manifest)), (PRE_CACHE_MANIFEST_KEY, pre_cache_manifest_payload(manifest)), ] assert manifest.default_version == "5.2.0" @@ -205,6 +235,75 @@ def test_readers_released_before_the_cache_get_only_fields_they_accept() -> None } +# The fields of a country bundle before region datasets. Readers released +# before them forbid any other field, so adding one would stop them loading. +_COUNTRY_BUNDLE_FIELDS_BEFORE_REGION_DATASETS = { + "country", + "country_package_name", + "country_package_version", + "country_package_requirement", + "data_package_name", + "data_package_version", + "data_release_version", + "data_artifact_revision", + "default_dataset", + "default_dataset_uri", + "datasets", +} + + +@pytest.mark.parametrize("key", [ACTIVE_MANIFEST_KEY, PRE_CACHE_MANIFEST_KEY]) +def test_readers_released_before_region_datasets_get_only_fields_they_accept( + key: str, +) -> None: + store = FakeStore() + publish_v2_manifest(store=store, worker=_regional_worker()) + + worker = store.values[key]["versions"]["5.2.0"] + for country in worker["bundle"]["countries"]: + assert set(country) == _COUNTRY_BUNDLE_FIELDS_BEFORE_REGION_DATASETS + + +def test_the_region_dataset_key_keeps_region_datasets_and_cache() -> None: + store = FakeStore() + publish_v2_manifest(store=store, worker=_regional_worker()) + + loaded = V2ManifestLoader(store).load().versions["5.2.0"] + us = next(country for country in loaded.bundle.countries if country.country == "us") + assert us.dataset_for_region_type("state") == "populace_us_2024_local" + assert us.dataset_for_region_type("national") == "populace_us_2024" + assert us.dataset_for_region_type("congressional_district") is None + assert loaded.cache is not None + + +def test_a_bundle_published_before_region_datasets_has_none() -> None: + worker = _cached_worker() + us = next(country for country in worker.bundle.countries if country.country == "us") + + assert us.region_datasets is None + assert us.dataset_for_region_type("state") is None + + +def test_first_publish_with_region_datasets_keeps_versions_from_the_cache_key() -> None: + earlier = _cached_worker("5.1.0").model_dump(mode="json") + store = FakeStore( + { + ACTIVE_MANIFEST_KEY: { + "schema_version": 1, + "generation": "5.1.0:earlier", + "default_version": "5.1.0", + "versions": {"5.1.0": earlier}, + } + } + ) + assert V2ManifestLoader(store).load().default_version == "5.1.0" + + manifest = publish_v2_manifest(store=store, worker=_regional_worker()) + + assert sorted(manifest.versions) == ["5.1.0", "5.2.0"] + assert V2ManifestLoader(store).load() == manifest + + def test_first_publish_with_cache_keeps_versions_from_the_pre_cache_key() -> None: earlier = _worker("5.1.0").model_dump( mode="json", exclude={"cache": True, "validation": {"cache"}} @@ -313,3 +412,70 @@ def test_storage_names_must_be_distinct() -> None: ) with pytest.raises(ValueError, match="different"): assert_separate_manifest_names(v1_name="same", v2_name="same") + + +@pytest.mark.parametrize( + ("country", "region", "region_type"), + [ + ("us", None, "national"), + ("us", "", "national"), + ("us", " US ", "national"), + ("us", "state/ut", "state"), + ("us", "State/UT", "state"), + ("us", "ut", "state"), + ("us", " ut ", "state"), + ("us", "congressional_district/ut-01", "congressional_district"), + ("us", "CONGRESSIONAL_DISTRICT/UT-01", "congressional_district"), + ("us", "place/UT-67000", "place"), + ("us", "utah", None), + ("us", "uk", None), + ("uk", "us", None), + ("uk", "uk", "national"), + ("uk", "country/scotland", "country"), + ("uk", "scotland", "country"), + ("uk", "constituency/Aldershot", "constituency"), + ], +) +def test_region_type_reads_codes_as_the_runtime_normalises_them( + country: str, region: str | None, region_type: str | None +) -> None: + assert region_type_of(country, region) == region_type + + +def test_a_country_bundle_gives_the_dataset_a_region_code_loads() -> None: + worker = _regional_worker() + us = next(country for country in worker.bundle.countries if country.country == "us") + uk = next(country for country in worker.bundle.countries if country.country == "uk") + + assert us.dataset_for_region("State/UT") == "populace_us_2024_local" + assert us.dataset_for_region("us") == "populace_us_2024" + assert us.dataset_for_region("place/UT-67000") is None + assert uk.dataset_for_region("country/scotland") is None + + +def test_a_region_dataset_copy_left_behind_by_an_older_publish_is_not_read() -> None: + """After a rollback, older code republishes only the keys it knows.""" + store = FakeStore() + publish_v2_manifest(store=store, worker=_regional_worker("5.2.0")) + stale = store.values[REGION_DATASETS_MANIFEST_KEY] + # A publish by code released before region datasets: it writes the two + # older keys only, leaving the region dataset copy as it was. + older = publish_v2_manifest(store=FakeStore(), worker=_cached_worker("5.3.0")) + store.values[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(older) + store.values[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(older) + + loaded = V2ManifestLoader(store).load() + + assert loaded.default_version == "5.3.0" + assert store.values[REGION_DATASETS_MANIFEST_KEY] == stale + + +def test_a_publish_after_a_rollback_keeps_the_versions_the_older_code_added() -> None: + store = FakeStore() + publish_v2_manifest(store=store, worker=_regional_worker("5.2.0")) + older = publish_v2_manifest(store=FakeStore(), worker=_cached_worker("5.3.0")) + store.values[ACTIVE_MANIFEST_KEY] = pre_region_datasets_manifest_payload(older) + + manifest = publish_v2_manifest(store=store, worker=_regional_worker("5.4.0")) + + assert sorted(manifest.versions) == ["5.3.0", "5.4.0"] diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py index 9f008fd77..c66929968 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py @@ -169,6 +169,20 @@ def _accepted_dataset_values(selected_dataset: Stage12Dataset) -> set[str | None } +def _region_dataset( + bundle_country: Stage12CountryBundle, region: str +) -> Stage12Dataset | None: + """Return the dataset the bundle says this region loads, if it names one. + + Workers published before region datasets name none, and their regions + load the default dataset. + """ + identity = bundle_country.dataset_for_region(region) + return next( + (item for item in bundle_country.datasets if item.identity == identity), None + ) + + def adapt_annual_comparison( payload: dict[str, Any], *, @@ -191,13 +205,18 @@ def adapt_annual_comparison( if skip_reason is not None: return _skip(skip_reason) - selected_dataset = next( + default_dataset = next( item for item in bundle_country.datasets if item.identity == bundle_country.default_dataset ) - if payload.get("data") not in _accepted_dataset_values(selected_dataset): + if payload.get("data") not in _accepted_dataset_values(default_dataset): return _skip(ComparisonSkipReason.UNSUPPORTED_DATASET) + # A state or district region loads the dataset the bundle names for its + # region type, so its provenance names that dataset. + selected_dataset = ( + _region_dataset(bundle_country, payload["region"]) or default_dataset + ) return AdaptedComparisonRun( report=_report_input( diff --git a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py index 46ebe97e7..fa4e83f6a 100644 --- a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py +++ b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py @@ -9,6 +9,7 @@ Stage12BundleManifest, Stage12CountryBundle, Stage12Dataset, + Stage12RegionDataset, ) from policyengine_simulation_contract.stage12_execution import ( ComparisonReportRecord, @@ -92,6 +93,56 @@ def worker() -> V2WorkerVersion: ) +REGIONAL_DATASET = "populace_us_2024_acs_local" +REGIONAL_REVISION = "acs-local-revision" +REGIONAL_URI = f"hf://policyengine/data/{REGIONAL_DATASET}.h5@{REGIONAL_REVISION}" + + +def regional_worker() -> V2WorkerVersion: + """Return ``worker()`` whose US bundle sends state and district regions elsewhere. + + State and congressional-district regions load the ACS-local dataset; + national regions keep the default. + """ + base = worker() + countries = [] + for country in base.bundle.countries: + if country.country != "us": + countries.append(country) + continue + regional = Stage12Dataset( + identity=REGIONAL_DATASET, + uri=REGIONAL_URI, + artifact_revision=REGIONAL_REVISION, + sha256="c" * 64, + repo_type="dataset", + ) + countries.append( + country.model_copy( + update={ + "datasets": (*country.datasets, regional), + "region_datasets": ( + Stage12RegionDataset( + region_type="national", dataset=country.default_dataset + ), + Stage12RegionDataset( + region_type="state", dataset=REGIONAL_DATASET + ), + Stage12RegionDataset( + region_type="congressional_district", + dataset=REGIONAL_DATASET, + ), + ), + } + ) + ) + return base.model_copy( + update={ + "bundle": base.bundle.model_copy(update={"countries": tuple(countries)}) + } + ) + + def eligible_payload() -> dict: return { "country": "us", diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index 401c15492..5751a9b24 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -311,3 +311,87 @@ def test_adapter_checks_eligibility_in_order(update, reason) -> None: assert result.skip_reason is reason assert result.eligible is (reason is None) + + +def _datasets_named(report) -> set[tuple[str, str, str, str, str]]: + return { + ( + simulation.bundle.dataset.identity, + simulation.bundle.dataset.uri, + simulation.bundle.dataset.artifact_revision, + simulation.population.artifact.uri, + simulation.population.artifact.content_sha256, + ) + for simulation in (report.baseline, report.reform) + } + + +@pytest.mark.parametrize( + "region", ["state/ut", "State/UT", "ut", " ut ", "congressional_district/UT-01"] +) +def test_state_and_district_requests_name_the_regional_dataset(region) -> None: + from stage12_fixtures import ( + REGIONAL_DATASET, + REGIONAL_REVISION, + REGIONAL_URI, + regional_worker, + ) + + result = adapt_annual_comparison( + {**eligible_payload(), "region": region}, + evaluation_id=EVALUATION_ID, + worker=regional_worker(), + ) + + assert result.skip_reason is None + assert _datasets_named(result.report) == { + (REGIONAL_DATASET, REGIONAL_URI, REGIONAL_REVISION, REGIONAL_URI, "c" * 64) + } + assert result.report.baseline.geography.region == region + + +def test_a_national_request_keeps_the_default_dataset_with_regional_workers() -> None: + from stage12_fixtures import regional_worker + + payload = {**eligible_payload(), "include_cliffs": True} + payload["spm"] = {"geography_kind": "national"} + + national = adapt_annual_comparison( + payload, evaluation_id=EVALUATION_ID, worker=regional_worker() + ) + today = adapt_annual_comparison( + payload, evaluation_id=EVALUATION_ID, worker=worker() + ) + + assert national.report.model_dump(mode="json") == today.report.model_dump( + mode="json" + ) + + +def test_a_worker_without_region_datasets_gets_the_default_dataset() -> None: + result = adapt_annual_comparison( + {**eligible_payload(), "region": "state/ut"}, + evaluation_id=EVALUATION_ID, + worker=worker(), + ) + + assert result.skip_reason is None + assert {item[0] for item in _datasets_named(result.report)} == {"populace_us_2024"} + + +def test_a_uk_request_keeps_the_default_dataset_with_regional_workers() -> None: + from stage12_fixtures import regional_worker + + result = adapt_annual_comparison( + { + **eligible_payload(), + "country": "uk", + "region": "country/scotland", + "version": "2.90.2", + }, + evaluation_id=EVALUATION_ID, + worker=regional_worker(), + ) + + assert result.skip_reason is None + assert {item[0] for item in _datasets_named(result.report)} == {"populace_uk_2023"} diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 910c20a1b..37f5e91cb 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -205,8 +205,10 @@ redacted cause (type, stage, simulation role and redacted message) followed by the same correlation ID. Stage 12 deployments run `src/modal/stage12_precompute_app.py` against the -dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets plus 60 -current-law US baseline files and emits +dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets, 60 +current-law US baseline files and, for each state region whose requests load +a dataset other than the bundle default, that dataset's year file filtered to +the state's households (one per state and year), and emits `STAGE12_CACHE_MANIFEST_DIGEST=`. Baseline files have the names .py saves and loads them under (`Simulation.storage_id`, which includes the SPM selection's digest). Stored files are never replaced: a run computes only @@ -217,7 +219,9 @@ workers, its image and functions mount only the read-only `HUGGING_FACE_TOKEN`, and `stage12-evaluation-gcp-credentials`; the US data it reads comes from the public `policyengine/populace-us` dataset repository. The US worker image -downloads and verifies the manifest's 63 files during image construction. The +downloads and verifies the manifest's files during image construction: the 63 +national files into the data folder, and each state dataset into +`state-datasets///` under the year file's own name. The release validation checks the recorded manifest, hashes a baseline file, and requires a real calculation to load that file without recomputing it. Worker validation also calculates a US household under a nonempty CTC reform using diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py b/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py index 423a14963..e6b2513f7 100644 --- a/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py +++ b/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py @@ -5,7 +5,7 @@ import json import os from hashlib import sha256 -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import TYPE_CHECKING if TYPE_CHECKING: @@ -20,16 +20,10 @@ def fetch_stage12_cache( if client is None: client = _storage_client() data_dir = Path(os.environ["POLICYENGINE_DATA_FOLDER"]) + destinations = _artifact_destinations(artifacts, data_dir) data_dir.mkdir(parents=True, exist_ok=True) - filenames: set[str] = set() - for artifact in artifacts: - _fetch_artifact( - artifact, - client=client, - bucket=bucket, - data_dir=data_dir, - filenames=filenames, - ) + for artifact, destination in zip(artifacts, destinations, strict=True): + _fetch_artifact(artifact, client=client, bucket=bucket, destination=destination) (data_dir / ".stage12-cache-manifest.json").write_bytes(canonical) @@ -75,23 +69,94 @@ def _require_installed_identity(manifest: dict) -> None: def _expected_artifacts(manifest: dict) -> list: - """Return the manifest's artifact list once it has 3 datasets and 60 baselines.""" + """Return the manifest's artifact list once its contents are as expected. + + It must have 3 datasets and 60 baselines, no artifact of another kind, and + either no state datasets or the same states' files for every year. + """ artifacts = manifest.get("artifacts") if not isinstance(artifacts, list): msg = "Stage 12 cache manifest has no artifact list" raise TypeError(msg) + entries = [item for item in artifacts if isinstance(item, dict)] type_counts = { - kind: sum( - item.get("type") == kind for item in artifacts if isinstance(item, dict) - ) + kind: sum(item.get("type") == kind for item in entries) for kind in ("dataset", "baseline") } if type_counts != {"dataset": 3, "baseline": 60}: msg = "Stage 12 cache manifest has an unexpected artifact count" raise RuntimeError(msg) + if any( + item.get("type") not in {"dataset", "baseline", "state_dataset"} + for item in entries + ): + msg = "Stage 12 cache manifest has an unexpected artifact type" + raise RuntimeError(msg) + states_by_year = {year: set() for year in manifest["years"]} + for item in entries: + if item.get("type") == "state_dataset": + states_by_year.setdefault(item.get("year"), set()).add( + item.get("state_code") + ) + if len({frozenset(states) for states in states_by_year.values()}) > 1: + msg = "Stage 12 cache manifest's state datasets differ between years" + raise RuntimeError(msg) return artifacts +def _relative_destination(artifact: dict) -> PurePosixPath | None: + """Return where a listed file lands, relative to the data folder, or None. + + Dataset and baseline files land directly in the data folder; a state + dataset lands in its ``folder``, a relative path without ``.`` or ``..`` + parts. None means the entry names an unsafe place. + """ + filename = artifact.get("filename") + if not isinstance(filename, str) or PurePosixPath(filename).name != filename: + return None + if "\\" in filename or filename in {"", ".", ".."}: + return None + if artifact.get("type") != "state_dataset": + return PurePosixPath(filename) + folder = artifact.get("folder") + if not isinstance(folder, str) or not folder or "\\" in folder: + return None + parts = PurePosixPath(folder).parts + if ( + PurePosixPath(folder).is_absolute() + or any(part in {".", ".."} for part in parts) + or "/".join(parts) != folder + ): + return None + return PurePosixPath(folder, filename) + + +def _artifact_destinations(artifacts: list, data_dir: Path) -> list[Path]: + """Check every listed file's metadata and return where each one lands. + + Runs before anything downloads, so a bad entry leaves no file behind. + """ + destinations = [] + seen: set[PurePosixPath] = set() + for artifact in artifacts: + if not isinstance(artifact, dict): + msg = "Stage 12 cache artifact entry is invalid" + raise TypeError(msg) + relative = _relative_destination(artifact) + if ( + relative is None + or relative in seen + or not isinstance(artifact.get("path"), str) + or not isinstance(artifact.get("content_sha256"), str) + or not isinstance(artifact.get("size_bytes"), int) + ): + msg = "Stage 12 cache artifact metadata is invalid" + raise RuntimeError(msg) + seen.add(relative) + destinations.append(data_dir.joinpath(*relative.parts)) + return destinations + + def _storage_client() -> Client: """Build a storage client from the deployment's service-account credentials. @@ -116,40 +181,23 @@ def _storage_client() -> Client: def _fetch_artifact( - artifact: object, + artifact: dict, *, client: Client, bucket: str, - data_dir: Path, - filenames: set[str], + destination: Path, ) -> None: - """Download one listed file into ``data_dir`` and verify its size and digest. + """Download one checked file to ``destination`` and verify its size and digest. The download goes to a partial file beside its destination and is moved into place only once verified, so a failed download leaves nothing. """ - if not isinstance(artifact, dict): - msg = "Stage 12 cache artifact entry is invalid" - raise TypeError(msg) - filename = artifact.get("filename") - path = artifact.get("path") - expected_digest = artifact.get("content_sha256") - expected_size = artifact.get("size_bytes") - if ( - not isinstance(filename, str) - or Path(filename).name != filename - or filename in filenames - or not isinstance(path, str) - or not isinstance(expected_digest, str) - or not isinstance(expected_size, int) - ): - msg = "Stage 12 cache artifact metadata is invalid" - raise RuntimeError(msg) - filenames.add(filename) - destination = data_dir / filename - partial = destination.with_name(filename + ".partial") + expected_digest = artifact["content_sha256"] + expected_size = artifact["size_bytes"] + destination.parent.mkdir(parents=True, exist_ok=True) + partial = destination.with_name(destination.name + ".partial") try: - client.bucket(bucket).blob(path).download_to_filename(str(partial)) + client.bucket(bucket).blob(artifact["path"]).download_to_filename(str(partial)) digest = sha256() with partial.open("rb") as stream: while chunk := stream.read(1024 * 1024): diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py index 414e5793b..2f0bc5160 100644 --- a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py +++ b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py @@ -14,10 +14,12 @@ BaselinePlanEntry, DatasetPlanEntry, PrecomputePlan, + StateDatasetPlanEntry, ) from policyengine_simulation_executor.stage12_cache.precompute import ( CachePrecomputeFunctions, build_dataset_impl, + build_state_datasets_impl, compute_baseline_impl, plan_artifacts_impl, publish_manifest_impl, @@ -104,6 +106,24 @@ def build_dataset(bucket: str, expected: dict) -> dict: ).model_dump() +@app.function( + image=precompute_image, + cpu=8.0, + # Preparing the regional dataset's full-country year file has used about + # 51 GB; the 51 filtered copies need room on top of that. + memory=98304, + timeout=7200, + secrets=worker_secrets, +) +def build_state_datasets(bucket: str, year: int, expected: list[dict]) -> list[dict]: + """Build and upload one year's planned state datasets; return their results.""" + entries = tuple(StateDatasetPlanEntry.model_validate(entry) for entry in expected) + return [ + result.model_dump() + for result in build_state_datasets_impl(bucket, year, entries) + ] + + @app.function( image=precompute_image, cpu=8.0, @@ -154,6 +174,7 @@ def main() -> None: functions=CachePrecomputeFunctions( plan_artifacts=plan_artifacts, build_dataset=build_dataset, + build_state_datasets=build_state_datasets, compute_baseline=compute_baseline, verify_determinism=verify_determinism, publish_manifest=publish_manifest, diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 0f98747e2..92642d673 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -195,7 +195,9 @@ def build_v2_image( secrets=[gcp_secret], cpu=2.0, memory=4096, - timeout=1800, + # The 63 national files and the state datasets (a file per state + # for each precomputed year), downloaded and hashed in turn. + timeout=3600, ) image = image.add_local_python_source( "src.modal", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py index 154f13002..f21114f4b 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/simulation_runtime.py @@ -722,6 +722,27 @@ def _resolve_dataset_selection( raise ValueError(msg) +def _selected_dataset_revision(country: str, selection: DatasetSelection) -> str | None: + """Return the release revision of the dataset a US state or district run loads. + + Such runs load a non-default dataset (the ACS-local dataset since .py 6.7), + so their ``data_version`` names that dataset's revision rather than the + default dataset's release. None leaves every other run reporting the + release's data version, as before. + """ + if country != "us" or selection.is_default: + return None + from policyengine.provenance.manifest import ( + get_release_manifest, + ) + + reference = getattr(get_release_manifest(country), "datasets", {}).get( + selection.name + ) + revision = getattr(reference, "revision", None) + return revision if isinstance(revision, str) and revision else None + + def _load_dataset( params: dict[str, Any], *, @@ -948,7 +969,7 @@ def _run_simulation_impl_core( dataset=dataset, baseline=baseline, reform=reform, - resolved_data_version=None, + resolved_data_version=_selected_dataset_revision(country, dataset_selection), resolved_region_code=region_resolution.code, runtime=runtime, uk_local_authority_metadata=uk_local_authority_metadata, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py index cf1a349e4..f80965a2a 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/manifest.py @@ -103,9 +103,19 @@ def normalize_stage12_bundle( def bundle_digest(bundle: Stage12BundleManifest) -> str: - """Return the sha256 of the manifest's canonical (key-sorted, compact) JSON.""" + """Return the sha256 of the manifest's canonical (key-sorted, compact) JSON. + + Each country's ``region_datasets`` is left out. It comes from the same .py + bundle, whose version the digest covers, and leaving it out keeps the + digest a running v2 application was validated with: a redeploy that adds + it replaces the application before the new version manifest is + published, and entry services send the digest they read from that + manifest. + """ payload = json.dumps( - bundle.model_dump(mode="json"), + bundle.model_dump( + mode="json", exclude={"countries": {"__all__": {"region_datasets"}}} + ), sort_keys=True, separators=(",", ":"), ensure_ascii=False, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py index b6975584c..047330999 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_bundle/parsing.py @@ -8,6 +8,7 @@ CountryId, Stage12CountryBundle, Stage12Dataset, + Stage12RegionDataset, ) from .errors import Stage12BundleError @@ -169,6 +170,7 @@ def normalize_country_bundle( f"data_releases.{country}.datasets", ) datasets: list[Stage12Dataset] = [] + identity_by_path: dict[str, str] = {} for identity in sorted(dataset_items): dataset = require_mapping( dataset_items[identity], @@ -194,6 +196,10 @@ def normalize_country_bundle( dataset.get("sha256"), f"data_releases.{country}.datasets.{identity}.sha256", ) + # The default dataset owns its path; otherwise the first identity + # in sorted order does. + if identity == default_dataset or path not in identity_by_path: + identity_by_path[path] = identity datasets.append( Stage12Dataset( identity=identity, @@ -256,4 +262,48 @@ def normalize_country_bundle( default_dataset=default_dataset, default_dataset_uri=default_dataset_uri, datasets=tuple(datasets), + region_datasets=_region_datasets( + release.get("region_datasets"), + identity_by_path=identity_by_path, + country=country, + ), ) + + +def _region_datasets( + value: object, + *, + identity_by_path: Mapping[str, str], + country: CountryId, +) -> tuple[Stage12RegionDataset, ...] | None: + """Return the dataset each region type loads, from ``region_datasets``. + + Each entry's ``path_template`` is matched to a dataset's ``path``, as + .py's ``resolve_region_dataset_path`` resolves it. A template with + placeholders names one file per region rather than one dataset, so it + maps to None. A bundle without ``region_datasets`` gives None. + """ + if value is None: + return None + entries = require_mapping(value, f"data_releases.{country}.region_datasets") + region_datasets = [] + for region_type in sorted(entries): + field_path = f"data_releases.{country}.region_datasets.{region_type}" + template = require_text( + require_mapping(entries[region_type], field_path).get("path_template"), + f"{field_path}.path_template", + ) + if "{" in template: + dataset = None + else: + dataset = identity_by_path.get(template) + if dataset is None: + msg = ( + f"PolicyEngine.py bundle {field_path} names {template!r}, " + "which is not one of its datasets" + ) + raise Stage12BundleError(msg) + region_datasets.append( + Stage12RegionDataset(region_type=region_type, dataset=dataset) + ) + return tuple(region_datasets) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py index f2d34788e..a31f41561 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py @@ -21,6 +21,7 @@ ) DATASET_SCHEMA = "stage12-ds1" +STATE_DATASET_SCHEMA = "stage12-sds1" BASELINE_SCHEMA = "stage12-bl1" CURRENT_LAW_POLICY = "current-law" _ID_DIGEST_CHARS = 16 @@ -95,6 +96,66 @@ def store_path(self) -> str: return f"datasets/{self.country}/{self.digest}/{self.filename}" +@dataclass(frozen=True) +class Stage12StateDatasetCacheIdentity: + """Identity of one state's cached year file, filtered from a regional dataset. + + The file is ``dataset``'s year file restricted to the households whose + ``filter_variable`` equals ``state_fips``. Its fields are those of a + dataset identity for the regional dataset, plus the state and filter. + """ + + country: CountryId + dataset: str + stem: str + year: int + source_sha256: str + data_artifact_revision: str + data_package_version: str + country_package_version: str + core_package_version: str + policyengine_version: str + bundle_manifest_sha256: str + state_code: str + state_fips: int + filter_variable: str + + @property + def payload(self) -> dict[str, Any]: + """The fields hashed into ``digest``.""" + return { + "schema": STATE_DATASET_SCHEMA, + "country": self.country, + "dataset": self.dataset, + "year": self.year, + "source_sha256": self.source_sha256, + "data_artifact_revision": self.data_artifact_revision, + "data_package_version": self.data_package_version, + "country_package_version": self.country_package_version, + "core_package_version": self.core_package_version, + "policyengine_version": self.policyengine_version, + "bundle_manifest_sha256": self.bundle_manifest_sha256, + "state_code": self.state_code, + "state_fips": self.state_fips, + "filter_variable": self.filter_variable, + } + + @property + def digest(self) -> str: + """The canonical digest of ``payload``; it names the cache directory.""" + return canonical_digest(self.payload) + + @property + def filename(self) -> str: + """The year file's own name, which every state's file shares.""" + return f"{self.stem}_year_{self.year}.h5" + + @property + def store_path(self) -> str: + """The object path in the cache bucket.""" + return f"state-datasets/{self.country}/{self.digest}/{self.filename}" + + @dataclass(frozen=True) class Stage12BaselineCacheIdentity: """Identity of one cached current-law baseline for a group of regions. @@ -193,6 +254,54 @@ def collect_dataset_identity( ) +def collect_state_dataset_identity( # noqa: PLR0913 # one keyword per input of the per-state identity + resolved: ResolvedStage12Bundle, + *, + country: CountryId, + year: int, + dataset: str, + state_code: str, + state_fips: int, + filter_variable: str, +) -> Stage12StateDatasetCacheIdentity: + """Collect one state's dataset identity from the installed Stage 12 bundle. + + ``dataset`` names a bundle dataset other than the default, such as the + regional dataset state regions load; its own hash and revision key the + file. + """ + from policyengine.provenance.manifest import ( + dataset_logical_name, + resolve_dataset_reference, + ) + + country_bundle = next( + item for item in resolved.bundle.countries if item.country == country + ) + entry = next( + (item for item in country_bundle.datasets if item.identity == dataset), None + ) + if entry is None: + msg = f"Dataset {dataset!r} is not in the installed {country} Stage 12 bundle" + raise ValueError(msg) + return Stage12StateDatasetCacheIdentity( + country=country, + dataset=entry.identity, + stem=dataset_logical_name(resolve_dataset_reference(country, dataset)), + year=year, + source_sha256=entry.sha256, + data_artifact_revision=entry.artifact_revision, + data_package_version=country_bundle.data_package_version, + country_package_version=country_bundle.country_package_version, + core_package_version=resolved.bundle.core_package_version, + policyengine_version=resolved.bundle.policyengine_version, + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + state_code=state_code.upper(), + state_fips=int(state_fips), + filter_variable=filter_variable, + ) + + def collect_baseline_identity( simulation: PlannedSimulationExecutionInput, *, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py index dcba07bdb..68d2aa875 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py @@ -35,6 +35,25 @@ class DatasetPlanEntry(StrictCacheModel): exists: bool +class StateDatasetPlanEntry(StrictCacheModel): + """One state's planned year file, filtered from a regional dataset. + + ``folder`` is where the worker image places the file, relative to its + data folder; every state's file has the same ``filename``. ``exists`` + records whether the bucket already holds it. + """ + + year: int + state_code: str + region: str + dataset: str + digest: str + path: str + filename: str + folder: str + exists: bool + + class BaselinePlanEntry(StrictCacheModel): """One planned baseline for a year and region group. @@ -67,13 +86,15 @@ class PrecomputePlan(StrictCacheModel): bundle: CacheBundleIdentity partition_sha256: str determinism_verified: bool + state_datasets: tuple[StateDatasetPlanEntry, ...] = () class WorkSelection(StrictCacheModel): - """The planned datasets and baselines a run still has to build.""" + """The planned datasets, state datasets and baselines a run still has to build.""" datasets: tuple[DatasetPlanEntry, ...] baselines: tuple[BaselinePlanEntry, ...] + state_datasets: tuple[StateDatasetPlanEntry, ...] = () class CacheObject(StrictCacheModel): @@ -92,6 +113,19 @@ class DatasetBuildResult(CacheObject): build_seconds: float +class StateDatasetBuildResult(CacheObject): + """A built state dataset file and whether this run uploaded it. + + ``build_seconds`` is how long the year's build had run when this file was + uploaded, so the last result of a build gives its whole duration. + """ + + year: int + state_code: str + uploaded: bool + build_seconds: float + + class BaselineBuildResult(CacheObject): """A computed baseline file and whether this run uploaded it.""" @@ -125,6 +159,28 @@ class ManifestArtifact(StrictCacheModel): size_bytes: int +class StateDatasetManifestArtifact(StrictCacheModel): + """One state's year file, which the worker image places in its state ``folder``. + + Its own model, so dataset and baseline entries keep their exact wire shape. + """ + + type: Literal["state_dataset"] + path: str + filename: str + folder: str + state_code: str + year: int + identity_digest: str + content_sha256: str + size_bytes: int + + +CacheManifestEntry = Annotated[ + ManifestArtifact | StateDatasetManifestArtifact, Field(discriminator="type") +] + + class CacheManifest(StrictCacheModel): """The published list of cache files for one bundle and region partition.""" @@ -137,7 +193,7 @@ class CacheManifest(StrictCacheModel): years: tuple[int, ...] partition_sha256: str bundle: CacheBundleIdentity - artifacts: tuple[ManifestArtifact, ...] + artifacts: tuple[CacheManifestEntry, ...] def canonical_payload(self) -> dict[str, Any]: """Return the JSON payload that is hashed and stored, keyed by ``schema``.""" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py index fd2b122c5..782647cf8 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -2,8 +2,10 @@ from __future__ import annotations +import gc +import tempfile import time -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any from uuid import NAMESPACE_URL, uuid5 @@ -26,6 +28,7 @@ from policyengine_simulation_executor import simulation_runtime from policyengine_simulation_executor.simulation_runtime import ( _country_module, + _resolve_dataset_selection, _resolve_region, ) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle @@ -49,26 +52,35 @@ canonical_digest, collect_baseline_identity, collect_dataset_identity, + collect_state_dataset_identity, ) from .models import ( BaselineBuildResult, BaselinePlanEntry, CacheBundleIdentity, CacheManifest, + CacheManifestEntry, DatasetBuildResult, DatasetPlanEntry, DeterminismVerdict, ManifestArtifact, PrecomputePlan, + StateDatasetBuildResult, + StateDatasetManifestArtifact, + StateDatasetPlanEntry, WorkSelection, ) +from .state_datasets import state_dataset_relative_folder, year_file_reuse_problem from .store import Stage12CacheStore if TYPE_CHECKING: from collections.abc import Callable, Mapping, Sequence + from types import ModuleType import modal import pandas as pd + from microdf import MicroDataFrame + from policyengine.core.scoping_strategy import RowFilterStrategy PRECOMPUTE_YEARS = (2026, 2027, 2025) MANIFEST_DIGEST_PREFIX = "STAGE12_CACHE_MANIFEST_DIGEST=" @@ -213,6 +225,69 @@ def determinism_record_path(baselines: Sequence[BaselinePlanEntry]) -> str: return f"determinism/us/{canonical_digest(determinism_record(baselines))}.json" +@dataclass(frozen=True) +class StateDatasetTarget: + """A US state whose requests load a non-default dataset, and its row filter. + + Stage 12 bakes this state's rows of ``dataset`` as their own year file: + the households whose ``filter_variable`` equals ``state_fips``. + ``scoping_strategy`` is the region's own .py strategy, the one a request + for the state applies. + """ + + region: str + state_code: str + state_fips: int + filter_variable: str + dataset: str + scoping_strategy: RowFilterStrategy | None = field( + default=None, compare=False, repr=False + ) + + +def state_dataset_targets() -> tuple[StateDatasetTarget, ...]: + """Return every US state region that loads a non-default dataset, by state code. + + Each state is resolved as a Stage 12 worker resolves a request for it + (``_resolve_region`` then ``_resolve_dataset_selection``), so a bundle + that sends states back to the default dataset gets no state files. + """ + country_module = _country_module("us") + targets = [] + for region in country_module.model.region_registry.get_by_type("state"): + params = { + "country": "us", + "region": region.code, + "time_period": PRECOMPUTE_YEARS[0], + } + resolution = _resolve_region( + country_module=country_module, country="us", params=params + ) + selection = _resolve_dataset_selection(params, region_resolution=resolution) + if selection.is_default: + continue + strategy = resolution.scoping_strategy + variable_name = getattr(strategy, "variable_name", None) + variable_value = getattr(strategy, "variable_value", None) + if not variable_name or variable_value is None or not resolution.state_code: + msg = ( + f"State region {resolution.code} loads {selection.name} but has no " + "row filter and state code to build its state file from" + ) + raise RuntimeError(msg) + targets.append( + StateDatasetTarget( + region=resolution.code, + state_code=resolution.state_code, + state_fips=int(variable_value), + filter_variable=variable_name, + dataset=selection.name, + scoping_strategy=strategy, + ) + ) + return tuple(sorted(targets, key=lambda target: target.state_code)) + + def plan_artifacts_impl(bucket: str) -> PrecomputePlan: """Plan every dataset and baseline file and record which the bucket holds. @@ -221,8 +296,10 @@ def plan_artifacts_impl(bucket: str) -> PrecomputePlan: """ resolved = load_stage12_bundle() store = Stage12CacheStore(bucket) + targets = state_dataset_targets() datasets = [] baselines = [] + state_datasets = [] for year in PRECOMPUTE_YEARS: dataset = collect_dataset_identity(resolved, country="us", year=year) datasets.append( @@ -251,10 +328,36 @@ def plan_artifacts_impl(bucket: str) -> PrecomputePlan: exists=store.exists(identity.store_path), ) ) + for target in targets: + state_identity = collect_state_dataset_identity( + resolved, + country="us", + year=year, + dataset=target.dataset, + state_code=target.state_code, + state_fips=target.state_fips, + filter_variable=target.filter_variable, + ) + state_datasets.append( + StateDatasetPlanEntry( + year=year, + state_code=target.state_code, + region=target.region, + dataset=target.dataset, + digest=state_identity.digest, + path=state_identity.store_path, + filename=state_identity.filename, + folder=state_dataset_relative_folder( + target.dataset, target.state_code + ), + exists=store.exists(state_identity.store_path), + ) + ) return PrecomputePlan( years=PRECOMPUTE_YEARS, datasets=tuple(datasets), baselines=tuple(baselines), + state_datasets=tuple(state_datasets), bundle=_bundle_identity(), partition_sha256=canonical_digest( {"groups": [list(group) for group in US_REGION_GROUPS]} @@ -268,6 +371,9 @@ def select_work(plan: PrecomputePlan) -> WorkSelection: return WorkSelection( datasets=tuple(entry for entry in plan.datasets if not entry.exists), baselines=tuple(entry for entry in plan.baselines if not entry.exists), + state_datasets=tuple( + entry for entry in plan.state_datasets if not entry.exists + ), ) @@ -298,6 +404,189 @@ def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildR ) +def _checked_state_targets( + year: int, expected: Sequence[StateDatasetPlanEntry] +) -> dict[str, RowFilterStrategy]: + """Return each planned state's row filter, after re-deriving its key here. + + Raises ``RuntimeError`` when the installed bundle gives another identity + than the plan, so nothing is built under a key that does not describe it, + and when a state has no row filter, before anything is prepared. + """ + resolved = load_stage12_bundle() + targets = {target.region: target for target in state_dataset_targets()} + 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( + resolved, + country="us", + year=year, + dataset=target.dataset, + state_code=target.state_code, + state_fips=target.state_fips, + filter_variable=target.filter_variable, + ) + folder = state_dataset_relative_folder(target.dataset, target.state_code) + if (identity.store_path, identity.filename, folder) != ( + entry.path, + entry.filename, + entry.folder, + ): + msg = ( + f"planned and runtime Stage 12 state dataset identities differ " + f"for {entry.region} {year}" + ) + raise RuntimeError(msg) + datasets.add(target.dataset) + if len(datasets) > 1: + msg = f"One state dataset build covers one dataset, not {sorted(datasets)}" + raise ValueError(msg) + 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 code that names no district, such as an unassigned 0, matches + no district request and changes no result. Both columns are labels; + nothing here is weighted. + """ + if "congressional_district_geoid" not in household.columns: + return + 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: Sequence[StateDatasetPlanEntry], + *, + work_folder: Path | None = None, +) -> tuple[StateDatasetBuildResult, ...]: + """Build, check and upload one year's Stage 12 state datasets. + + Has .py prepare the regional dataset's full-country year file once, then + for each planned state keeps the rows its region's own scoping strategy + keeps, writes them with .py's ``save()``, copies the year-file stamp, + confirms that .py's reuse checks accept the file, and uploads it. + """ + from policyengine.tax_benefit_models.us.datasets import ( + PolicyEngineUSDataset, + USYearData, + ) + from policyengine.tax_benefit_models.us.legacy_inputs import ( + read_year_file_stamp, + write_year_file_stamp, + ) + from policyengine.tax_benefit_models.us.model import US_GROUP_ENTITIES + + if not expected: + return () + started = time.monotonic() + strategies = _checked_state_targets(year, expected) + dataset = expected[0].dataset + work_folder = work_folder or Path( + tempfile.mkdtemp(prefix="policyengine-stage12-state-datasets-") + ) + source_folder = str(work_folder / "source") + country_module = _country_module("us") + + # Preparing the year file runs a Microsimulation over every state, whose + # objects hold reference cycles; release them before filtering 51 copies. + country_module.ensure_datasets( + datasets=[dataset], years=[year], data_folder=source_folder + ) + gc.collect() + (source,) = country_module.load_datasets( + datasets=[dataset], years=[year], data_folder=source_folder + ).values() + stamp = read_year_file_stamp(source.filepath) + if stamp is None: + msg = f"The prepared year file {source.filepath} has no year-file stamp" + raise RuntimeError(msg) + _require_districts_within_their_states( + source.data.household, _registered_district_geoids(country_module) + ) + + store = Stage12CacheStore(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) + uploaded, stored = store.upload_file(entry.path, local_file) + results.append( + StateDatasetBuildResult( + **stored.model_dump(), + year=year, + state_code=entry.state_code, + uploaded=uploaded, + build_seconds=round(time.monotonic() - started, 1), + ) + ) + return tuple(results) + + def _ensure_dataset_file(bucket: str, year: int) -> None: identity = collect_dataset_identity(load_stage12_bundle(), country="us", year=year) @@ -399,8 +688,12 @@ def verify_determinism_impl( def build_manifest(plan: PrecomputePlan, store: Stage12CacheStore) -> CacheManifest: - """List every planned file with the digest and size the bucket reports.""" - artifacts = [] + """List every planned file with the digest and size the bucket reports. + + Datasets, then baselines, then state datasets, each in plan order, so a + plan without state datasets keeps the manifest it always had. + """ + artifacts: list[CacheManifestEntry] = [] for entry in plan.datasets: stored = store.describe(entry.path) artifacts.append( @@ -427,6 +720,21 @@ def build_manifest(plan: PrecomputePlan, store: Stage12CacheStore) -> CacheManif size_bytes=stored.size_bytes, ) ) + for entry in plan.state_datasets: + stored = store.describe(entry.path) + artifacts.append( + StateDatasetManifestArtifact( + type="state_dataset", + path=entry.path, + filename=entry.filename, + folder=entry.folder, + state_code=entry.state_code, + year=entry.year, + identity_digest=entry.digest, + content_sha256=stored.content_sha256, + size_bytes=stored.size_bytes, + ) + ) return CacheManifest( years=plan.years, partition_sha256=plan.partition_sha256, @@ -452,10 +760,11 @@ def publish_manifest_impl(bucket: str, plan: PrecomputePlan) -> str: @dataclass(frozen=True) class CachePrecomputeFunctions: - """The five Modal function handles the cache precompute orchestration calls.""" + """The six Modal function handles the cache precompute orchestration calls.""" plan_artifacts: modal.Function build_dataset: modal.Function + build_state_datasets: modal.Function compute_baseline: modal.Function verify_determinism: modal.Function publish_manifest: modal.Function @@ -470,8 +779,11 @@ def run_precompute( """Build the missing cache files, check determinism, publish the manifest. Missing datasets are built first, then missing baselines, in parallel - Modal calls. Prints the manifest digest through ``echo`` and returns it. - Exits with ``SystemExit`` when the determinism check fails. + Modal calls. Missing state datasets are built alongside both, one call per + year and dataset, and waited for only before the manifest is published, so + the baselines never wait for them. Prints the manifest digest through + ``echo`` and returns it. Exits with ``SystemExit`` when the determinism + check fails. """ plan = PrecomputePlan.model_validate(functions.plan_artifacts.remote(bucket)) work = select_work(plan) @@ -479,6 +791,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[str, Any]]] = {} + for entry in work.state_datasets: + state_builds.setdefault((entry.year, entry.dataset), []).append( + entry.model_dump() + ) + state_calls = [ + (year, dataset, functions.build_state_datasets.spawn(bucket, year, entries)) + for (year, dataset), entries in state_builds.items() + ] for call in dataset_calls: DatasetBuildResult.model_validate(call.get()) baseline_calls = [ @@ -502,6 +825,17 @@ def run_precompute( "Stage 12 cache determinism validation failed: " + "; ".join(verdict.differences) ) + for year, dataset, call in state_calls: + results = [ + StateDatasetBuildResult.model_validate(result) for result in call.get() + ] + uploaded = [result for result in results if result.uploaded] + echo( + f"State datasets {year} {dataset}: {len(results)} built, " + f"{len(uploaded)} uploaded " + f"({sum(result.size_bytes for result in uploaded)} bytes) in " + f"{max((result.build_seconds for result in results), default=0.0)} s" + ) digest = functions.publish_manifest.remote(bucket, plan.model_dump()) if not isinstance(digest, str) or not digest: msg = "Stage 12 cache manifest publisher returned no digest" diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py new file mode 100644 index 000000000..7633041f2 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/state_datasets.py @@ -0,0 +1,102 @@ +"""Where the Stage 12 US worker image keeps baked per-state datasets. + +Stage 12's own copy of the main path's per-state dataset layout (see +``policyengine_simulation_executor.state_datasets``): Stage 12 does not +import the main path's precompute modules. A state's year files sit at +``/state-datasets///`` under the +basename the dataset's own year file has, so ``ensure_datasets`` called on +that folder finds the state's file where it would look for the full one. +""" + +from __future__ import annotations + +import logging +from pathlib import Path +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + +logger = logging.getLogger(__name__) + +STATE_DATASETS_DIRECTORY = "state-datasets" + + +def state_dataset_relative_folder(dataset: str, state_code: str) -> str: + """Return a state's folder relative to the Stage 12 data folder.""" + return f"{STATE_DATASETS_DIRECTORY}/{dataset}/{state_code.upper()}" + + +def state_dataset_folder(data_folder: str, dataset: str, state_code: str) -> str: + """Return the absolute folder of a state's baked year files.""" + return str(Path(data_folder) / state_dataset_relative_folder(dataset, state_code)) + + +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 calculation should load. + + A calculation qualifies when it loads a non-default dataset for a US + state (a district passes its state's code) and the image holds that + state's year file for ``year``, which must also pass .py's reuse checks: + otherwise ``ensure_datasets`` would rebuild the full regional dataset + inside the baked folder. Returns None in every other case, and the + calculation prepares the dataset as before. + """ + if country.lower() != "us" or not state_code or selection.is_default: + return None + from policyengine.provenance.manifest import ( + dataset_logical_name, + resolve_dataset_reference, + ) + + folder = state_dataset_folder(data_folder, selection.name, state_code) + stem = dataset_logical_name(resolve_dataset_reference("us", selection.name)) + year_file = Path(folder) / f"{stem}_year_{year}.h5" + if not year_file.is_file(): + return None + try: + problem = year_file_reuse_problem(year_file, selection.name) + except ImportError: + # The reuse checks read private .py helpers. If an upgrade moves one, + # the calculation prepares its dataset instead of failing. + logger.warning( + "Not reading the baked state file %s: .py's year-file reuse " + "checks could not be imported", + year_file, + ) + return None + if problem is not None: + logger.warning("Not reading the baked state file %s: %s", year_file, problem) + return None + return folder diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index b1a5695c9..3b2bb5b56 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -187,12 +187,30 @@ def _require_installed_bundle(simulation: SimulationExecutionInput) -> None: for item in resolved.bundle.countries if item.country == simulation.geography.country ) + # A region loads the dataset the bundle names for its region type. Entry + # services released before regional provenance name the default dataset + # for every region, so the default stays accepted until they are updated. + regional = country.dataset_for_region(simulation.geography.region) + named = simulation.bundle.dataset.identity + regional_entry = next( + (item for item in country.datasets if item.identity == regional), None + ) + if regional_entry is not None and named == regional_entry.identity: + expected_dataset = { + "dataset_identity": regional_entry.identity, + "dataset_uri": regional_entry.uri, + "dataset_revision": regional_entry.artifact_revision, + } + else: + expected_dataset = { + "dataset_identity": country.default_dataset, + "dataset_uri": country.default_dataset_uri, + "dataset_revision": country.data_artifact_revision, + } expected = { "country_package_name": country.country_package_name, "country_package_version": country.country_package_version, - "dataset_identity": country.default_dataset, - "dataset_uri": country.default_dataset_uri, - "dataset_revision": country.data_artifact_revision, + **expected_dataset, } actual = { "country_package_name": simulation.bundle.country_package_name, @@ -206,6 +224,34 @@ def _require_installed_bundle(simulation: SimulationExecutionInput) -> None: raise RuntimeError(msg) +def _require_loaded_dataset_matches_provenance( + simulation: SimulationExecutionInput, selection: DatasetSelection +) -> None: + """Require a calculation to load the dataset its provenance names. + + The provenance is recorded with the comparison run, so it must name the + dataset the region actually loads. The one accepted difference is a + state or district run whose provenance names the default dataset while it + loads a regional one: entry services released before regional provenance + send that until they are updated. + """ + from policyengine_simulation_executor.simulation_runtime import ( # noqa: PLC0415 # resolved per call so tests can patch simulation_runtime helpers + _resolve_dataset_reference, + ) + + named = simulation.bundle.dataset + if (selection.name, selection.uri) == (named.identity, named.uri): + return + if not selection.is_default and named.identity == _resolve_dataset_reference( + simulation.geography.country + ): + return + msg = ( + "simulation provenance names a dataset other than the dataset its region loads" + ) + raise RuntimeError(msg) + + def calculate_simulation_frames( simulation: PlannedSimulationExecutionInput, *, @@ -236,8 +282,12 @@ def calculate_simulation_frames( _require_certified_uk_weight_matrix, _resolve_dataset_selection, _resolve_region, + resolve_data_folder, setup_gcp_credentials, ) + from policyengine_simulation_executor.stage12_cache import ( # noqa: PLC0415 # resolved per call so tests can patch the baked-folder lookup + state_datasets, + ) with ( stage_span(Stage.CREDENTIAL_SETUP), @@ -256,11 +306,19 @@ def calculate_simulation_frames( params, region_resolution=region, ) + _require_loaded_dataset_matches_provenance(simulation, dataset_selection) with stage_span(Stage.DATASET_LOAD): dataset = _load_dataset( params, selection=dataset_selection, country_module=country_module, + data_folder=state_datasets.baked_state_data_folder( + country=country, + selection=dataset_selection, + state_code=region.state_code, + year=simulation.year, + data_folder=resolve_data_folder(), + ), ) uk_local_authority_metadata = detect_uk_local_authority_metadata( country, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index d861c427d..102593c99 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -18,6 +18,9 @@ resolve_data_folder, ) from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle +from policyengine_simulation_executor.stage12_cache import ( + precompute as cache_precompute, +) from policyengine_simulation_executor.stage12_cache.models import CacheManifest from policyengine_simulation_executor.stage12_cache.precompute import ( PRECOMPUTE_YEARS, @@ -147,6 +150,31 @@ def _check_uk_local_authority_dataset(dataset_path: str) -> None: STAGE12_CACHE_PROBE_YEAR = PRECOMPUTE_YEARS[0] +def expected_state_dataset_count() -> int: + """Return how many state dataset files the installed bundle's cache holds. + + One per precomputed year for each US state region on a non-default + dataset; none when every state uses the default dataset. + """ + return len(PRECOMPUTE_YEARS) * len(cache_precompute.state_dataset_targets()) + + +def _require_installed_artifacts( + manifest: CacheManifest, data_folder: Path, *, state_dataset_count: int +) -> None: + """Require the expected number of listed files and every state file on disk.""" + if len(manifest.artifacts) != STAGE12_CACHE_ARTIFACT_COUNT + state_dataset_count: + msg = "installed Stage 12 cache artifact count differs" + raise RuntimeError(msg) + for item in manifest.artifacts: + if ( + item.type == "state_dataset" + and not (data_folder / item.folder / item.filename).is_file() + ): + msg = f"installed Stage 12 state dataset file is missing: {item.path}" + raise RuntimeError(msg) + + def _check_installed_cache(expected_manifest_sha256: str) -> dict[str, Any]: """Validate the baked manifest and require one real baseline cache hit.""" data_folder = Path(resolve_data_folder()) @@ -159,9 +187,11 @@ def _check_installed_cache(expected_manifest_sha256: str) -> dict[str, Any]: if manifest.years != PRECOMPUTE_YEARS: msg = "installed Stage 12 cache years differ" raise RuntimeError(msg) - if len(manifest.artifacts) != STAGE12_CACHE_ARTIFACT_COUNT: - msg = "installed Stage 12 cache artifact count differs" - raise RuntimeError(msg) + _require_installed_artifacts( + manifest, + data_folder, + state_dataset_count=expected_state_dataset_count(), + ) probe = next( item for item in manifest.artifacts @@ -175,8 +205,19 @@ def _check_installed_cache(expected_manifest_sha256: str) -> dict[str, Any]: if calculation.cache_outcome != "hit": msg = "deployed Stage 12 cache probe did not load a cache hit" raise RuntimeError(msg) + return _cache_release_summary(manifest, expected_manifest_sha256) + + +def _cache_release_summary( + manifest: CacheManifest, manifest_sha256: str +) -> dict[str, Any]: + """Return the cache release a worker reports, in ``Stage12CacheRelease`` form. + + Only national datasets and baselines are counted: the release model pins + those counts and forbids other fields, so state datasets are not reported. + """ return { - "manifest_sha256": expected_manifest_sha256, + "manifest_sha256": manifest_sha256, "partition_sha256": manifest.partition_sha256, "years": list(manifest.years), "segment_count": len(US_REGION_GROUPS), diff --git a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py index 162c7620a..f4abd82e8 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1368,6 +1368,98 @@ def test_runtime_loads_a_national_run_from_the_default_folder(monkeypatch): assert "dataset_source" not in loaded.context +def _builder_data_version(monkeypatch, *, region_resolution, dataset_selection): + """Run the annual runtime and return the data version it gives the builder.""" + from policyengine_simulation_executor import state_datasets + + selection = installed_spm_selection() + simulations = [spm_capable_simulation(selection), spm_capable_simulation(selection)] + builder_calls = [] + + class FakeSimulationOutputBuilder: + def __init__(self, **kwargs): + builder_calls.append(kwargs) + + def serialize(self): + return dict(CURRENT_SINGLE_YEAR_MACRO_RESULT) + + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._country_module", + lambda _country: SimpleNamespace(model=SimpleNamespace(version="2.2.1")), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._resolve_region", + lambda **_kwargs: region_resolution, + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._resolve_dataset_selection", + lambda _params, *, region_resolution: dataset_selection, + ) + monkeypatch.setattr( + state_datasets, "baked_state_data_folder", lambda **_kwargs: None + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._load_dataset", + lambda _params, country_module, selection, data_folder=None: object(), + ) + build_iter = iter(simulations) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime._build_simulation", + lambda *_args, **_kwargs: next(build_iter), + ) + monkeypatch.setattr( + "policyengine_simulation_executor.simulation_runtime.SimulationOutputBuilder", + FakeSimulationOutputBuilder, + ) + + _run_simulation_impl_core( + { + "country": "us", + "region": region_resolution.code, + "baseline": {}, + "reform": {}, + }, + runtime=_TrackingRuntime(), + ) + return builder_calls[0]["resolved_data_version"] + + +def test_a_state_run_reports_the_revision_of_the_dataset_it_loads(monkeypatch): + from policyengine.provenance.manifest import get_release_manifest + + bundle = get_country_release_bundle("us") + regional = "populace_us_2024_acs_local" + + data_version = _builder_data_version( + monkeypatch, + region_resolution=RegionResolution( + code="state/ut", dataset_reference="ref", state_code="UT" + ), + dataset_selection=DatasetSelection( + name=regional, uri=bundle.dataset_uris[regional], is_default=False + ), + ) + + assert data_version == get_release_manifest("us").datasets[regional].revision + assert data_version != bundle.data_version + + +def test_a_national_run_keeps_the_release_data_version(monkeypatch): + bundle = get_country_release_bundle("us") + + data_version = _builder_data_version( + monkeypatch, + region_resolution=RegionResolution(code="us"), + dataset_selection=DatasetSelection( + name=bundle.default_dataset, + uri=bundle.default_dataset_uri, + is_default=True, + ), + ) + + assert data_version is None + + def test_run_simulation_impl_core_passes_detected_uk_boundary_to_builder(monkeypatch): dataset = SimpleNamespace( data=SimpleNamespace( diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py index 45d98f3f8..8bd6a03f6 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_bundle.py @@ -69,6 +69,9 @@ def test_complete_packaged_bundle_selects_every_worker_dependency_and_dataset() lambda bundle: bundle["data_releases"]["us"]["certified_data_artifact"].update( {"sha256": "0" * 64} ), + lambda bundle: bundle["data_releases"]["us"]["region_datasets"].update( + {"state": {"path_template": "no_such_dataset.h5"}} + ), ], ids=[ "missing-data-releases", @@ -78,6 +81,7 @@ def test_complete_packaged_bundle_selects_every_worker_dependency_and_dataset() "dataset-uri-revision-mismatch", "malformed-dataset-digest", "certified-artifact-mismatch", + "region-dataset-names-no-dataset", ], ) def test_missing_or_internally_inconsistent_bundle_values_fail_closed(mutate) -> None: @@ -130,3 +134,126 @@ def test_only_bundle_schema_version_2_is_supported(schema_version) -> None: Stage12BundleError, match=r"^PolicyEngine\.py bundle schema_version must be 2$" ): normalize_stage12_bundle(bundle) + + +def _us(bundle): + return next(country for country in bundle.countries if country.country == "us") + + +def test_the_installed_bundle_names_the_dataset_each_us_region_type_loads() -> None: + us = _us(normalize_stage12_bundle(_bundle())) + + assert us.dataset_for_region_type("national") == "populace_us_2024" + assert us.dataset_for_region_type("state") == "populace_us_2024_acs_local" + assert ( + us.dataset_for_region_type("congressional_district") + == "populace_us_2024_acs_local" + ) + + +def test_a_per_region_file_pattern_names_no_single_dataset() -> None: + bundle = _bundle() + bundle["data_releases"]["us"]["region_datasets"]["state"] = { + "path_template": "states/{state_code}.h5" + } + + us = _us(normalize_stage12_bundle(bundle)) + + assert any(item.region_type == "state" for item in us.region_datasets) + assert us.dataset_for_region_type("state") is None + + +def test_a_bundle_without_region_datasets_records_none() -> None: + bundle = _bundle() + bundle["data_releases"]["us"].pop("region_datasets") + + assert _us(normalize_stage12_bundle(bundle)).region_datasets is None + + +def test_region_datasets_leave_the_bundle_digest_unchanged() -> None: + """A worker redeployed with region datasets keeps the digest it had. + + The redeploy replaces the running v2 application before the new version + manifest is published, and entry services send the digest they read from + that manifest, so a changed digest would refuse every run in between. + Region datasets come from the same .py bundle, whose version the digest + already covers. + """ + bundle = normalize_stage12_bundle(_bundle()) + without_region_datasets = bundle.model_copy( + update={ + "countries": tuple( + country.model_copy(update={"region_datasets": None}) + for country in bundle.countries + ) + } + ) + + assert _us(bundle).region_datasets is not None + assert bundle_digest(bundle) == bundle_digest(without_region_datasets) + + +def test_a_dataset_sharing_the_default_path_does_not_replace_the_default() -> None: + bundle = _bundle() + release = bundle["data_releases"]["us"] + release["datasets"]["a_alias_of_the_default"] = deepcopy( + release["datasets"][release["default_dataset"]] + ) + + us = _us(normalize_stage12_bundle(bundle)) + + assert us.dataset_for_region_type("national") == us.default_dataset + + +def test_the_bundle_digest_of_a_fixed_bundle_is_pinned() -> None: + """Any change to what the digest covers is deliberate. + + A changed digest refuses comparison runs while a redeployed v2 application + waits for its version manifest, and rebuilds the Stage 12 cache, so a new + bundle field must not reach it by accident. The pinned value is the one + the code before region datasets gives the same bundle. + """ + from policyengine_simulation_contract.stage12_bundle import ( + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, + Stage12RegionDataset, + ) + + dataset = Stage12Dataset( + identity="populace_us_2024", + uri="hf://policyengine/populace-us/populace_us_2024.h5@revision", + artifact_revision="revision", + sha256="a" * 64, + repo_type="dataset", + ) + bundle = Stage12BundleManifest( + policyengine_version="6.7.5", + policyengine_requirement="policyengine==6.7.5", + core_package_version="3.33.0", + core_package_requirement="policyengine-core==3.33.0", + countries=( + Stage12CountryBundle( + country="us", + country_package_name="policyengine-us", + country_package_version="2.2.1", + country_package_requirement="policyengine-us==2.2.1", + data_package_name="populace-us", + data_package_version="1.0.0", + data_release_version="revision", + data_artifact_revision="revision", + default_dataset=dataset.identity, + default_dataset_uri=dataset.uri, + datasets=(dataset,), + region_datasets=( + Stage12RegionDataset( + region_type="national", dataset=dataset.identity + ), + ), + ), + ), + ) + + assert bundle_digest(bundle) == ( + "47f3f7c69ea3b99005ffee585092074862bf4f55a91a95511f12924a79aac71c" + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py index d737a9d54..388e54ebd 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py @@ -308,3 +308,133 @@ def build_client(**arguments): {"credentials": "service-account-credentials", "project": "cache-project"} ] assert len(list(tmp_path.glob("*.h5"))) == 63 + + +REGIONAL_DATASET = "populace_us_2024_acs_local" + + +def _state_artifact(year: int, state_code: str, **changes) -> tuple[dict, bytes]: + filename = f"{REGIONAL_DATASET}_year_{year}.h5" + payload = f"state-{year}-{state_code}".encode() + artifact = { + "type": "state_dataset", + "path": f"state-datasets/us/{year}-{state_code}/{filename}", + "filename": filename, + "folder": f"state-datasets/{REGIONAL_DATASET}/{state_code}", + "state_code": state_code, + "year": year, + "identity_digest": f"identity-state-{year}-{state_code}", + "content_sha256": sha256(payload).hexdigest(), + "size_bytes": len(payload), + **changes, + } + return artifact, payload + + +def _manifest_with_state_datasets(states=("UT", "WY")) -> tuple[dict, dict]: + manifest, payloads = _manifest() + for year in (2026, 2027, 2025): + for state_code in states: + artifact, payload = _state_artifact(year, state_code) + manifest["artifacts"].append(artifact) + payloads[artifact["path"]] = payload + return manifest, payloads + + +def test_places_each_state_file_in_its_state_folder(monkeypatch, tmp_path) -> None: + manifest, payloads = _manifest_with_state_datasets() + _use_deployment(monkeypatch, tmp_path, manifest) + + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + state_folder = tmp_path / "state-datasets" / REGIONAL_DATASET + for year in (2026, 2027, 2025): + for state_code in ("UT", "WY"): + stored = state_folder / state_code / f"{REGIONAL_DATASET}_year_{year}.h5" + assert stored.read_bytes() == f"state-{year}-{state_code}".encode() + assert len(list(tmp_path.glob("*.h5"))) == 63 + + +def _drop_one_state_file(manifest: dict) -> None: + manifest["artifacts"] = [ + item + for item in manifest["artifacts"] + if not (item["type"] == "state_dataset" and item["year"] == 2027) + or item["state_code"] != "WY" + ] + + +def _add_unknown_type(manifest: dict) -> None: + manifest["artifacts"].append({**manifest["artifacts"][0], "type": "mystery"}) + + +@pytest.mark.parametrize( + ("change", "message"), + [ + pytest.param( + _drop_one_state_file, + "^Stage 12 cache manifest's state datasets differ between years$", + id="state-missing-in-one-year", + ), + pytest.param( + _add_unknown_type, + "^Stage 12 cache manifest has an unexpected artifact type$", + id="unknown-type", + ), + ], +) +def test_rejects_inconsistent_state_datasets_before_downloading( + monkeypatch, tmp_path, change, message +) -> None: + manifest, payloads = _manifest_with_state_datasets() + change(manifest) + _use_deployment(monkeypatch, tmp_path, manifest) + + with pytest.raises(RuntimeError, match=message): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + assert list(tmp_path.iterdir()) == [] + + +@pytest.mark.parametrize( + "folder", + [ + "/absolute/UT", + "state-datasets/../../escape", + "state-datasets/./UT", + "state-datasets//UT", + "state-datasets\\UT", + "", + None, + ], +) +def test_rejects_an_unsafe_state_folder_before_downloading( + monkeypatch, tmp_path, folder +) -> None: + manifest, payloads = _manifest_with_state_datasets() + manifest["artifacts"][-1]["folder"] = folder + _use_deployment(monkeypatch, tmp_path, manifest) + + with pytest.raises( + RuntimeError, match=r"^Stage 12 cache artifact metadata is invalid$" + ): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + assert list(tmp_path.iterdir()) == [] + + +def test_rejects_two_entries_for_one_state_file_before_downloading( + monkeypatch, tmp_path +) -> None: + manifest, payloads = _manifest_with_state_datasets() + repeated = dict(manifest["artifacts"][-1]) + manifest["artifacts"].append({**repeated, "path": "state-datasets/us/other/x.h5"}) + payloads["state-datasets/us/other/x.h5"] = payloads[repeated["path"]] + _use_deployment(monkeypatch, tmp_path, manifest) + + with pytest.raises( + RuntimeError, match=r"^Stage 12 cache artifact metadata is invalid$" + ): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + assert list(tmp_path.iterdir()) == [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py index 6b5bb337b..3291fb84d 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -2,6 +2,7 @@ from __future__ import annotations +from pathlib import Path from types import SimpleNamespace import pandas as pd @@ -15,6 +16,7 @@ CacheObject, DatasetPlanEntry, PrecomputePlan, + StateDatasetPlanEntry, ) from policyengine_simulation_executor.stage12_cache.store import Stage12CacheStore from policyengine_simulation_executor.stage12_runtime import ( @@ -76,6 +78,36 @@ def _plan(*, exists: bool = False, verified: bool | None = None) -> PrecomputePl ) +REGIONAL_DATASET = "populace_us_2024_acs_local" + + +def _state_entry(year: int, state_code: str, *, exists: bool) -> StateDatasetPlanEntry: + digest = f"state-{year}-{state_code}" + return StateDatasetPlanEntry( + year=year, + state_code=state_code, + region=f"state/{state_code.lower()}", + dataset=REGIONAL_DATASET, + digest=digest, + path=f"state-datasets/us/{digest}/{REGIONAL_DATASET}_year_{year}.h5", + filename=f"{REGIONAL_DATASET}_year_{year}.h5", + folder=f"state-datasets/{REGIONAL_DATASET}/{state_code}", + exists=exists, + ) + + +def _plan_with_state_datasets(*, exists: bool = False) -> PrecomputePlan: + return _plan(exists=exists).model_copy( + update={ + "state_datasets": tuple( + _state_entry(year, state_code, exists=exists) + for year in precompute.PRECOMPUTE_YEARS + for state_code in ("UT", "WY") + ) + } + ) + + class _FakeBlob: """A GCS blob over a dict, honouring create-only uploads.""" @@ -98,6 +130,14 @@ def upload_from_string(self, payload, *, content_type, if_generation_match): raise PreconditionFailed(self.path) self.objects[self.path] = {"payload": payload, "metadata": self.metadata} + def upload_from_filename(self, filename, *, if_generation_match): + assert if_generation_match == 0 + if self.path in self.objects: + raise PreconditionFailed(self.path) + with open(filename, "rb") as source: # noqa: PTH123 # mirrors the GCS client's own file read + payload = source.read() + self.objects[self.path] = {"payload": payload, "metadata": self.metadata} + def download_as_bytes(self) -> bytes: if self.path not in self.objects: raise NotFound(self.path) @@ -158,6 +198,186 @@ def test_work_selection_computes_only_missing_files() -> None: assert warm.baselines == () +def test_work_selection_builds_only_missing_state_datasets() -> None: + plan = _plan_with_state_datasets() + stored = plan.state_datasets[0].model_copy(update={"exists": True}) + plan = plan.model_copy( + update={"state_datasets": (stored, *plan.state_datasets[1:])} + ) + + work = precompute.select_work(plan) + + assert [(entry.year, entry.state_code) for entry in work.state_datasets] == [ + (2026, "WY"), + (2027, "UT"), + (2027, "WY"), + (2025, "UT"), + (2025, "WY"), + ] + assert precompute.select_work(_plan()).state_datasets == () + + +def test_plan_with_state_datasets_round_trips_the_modal_boundary() -> None: + plan = _plan_with_state_datasets() + + assert PrecomputePlan.model_validate(plan.model_dump()) == plan + + +@pytest.fixture +def stubbed_planning(monkeypatch): + """Stub everything plan_artifacts_impl reads except the state targets.""" + client = _FakeClient() + _use_store(monkeypatch, client) + monkeypatch.setattr(precompute, "load_stage12_bundle", lambda: "resolved") + monkeypatch.setattr(precompute, "_bundle_identity", _bundle) + monkeypatch.setattr( + precompute, + "collect_dataset_identity", + lambda _resolved, *, country, year: SimpleNamespace( + digest=f"dataset-{year}", + store_path=f"datasets/{country}/dataset-{year}/populace_year_{year}.h5", + filename=f"populace_year_{year}.h5", + ), + ) + monkeypatch.setattr(precompute, "US_REGION_GROUPS", (("state/ca",),)) + monkeypatch.setattr( + precompute, + "_baseline_identity", + lambda year, _group: ( + SimpleNamespace(output_plan=None), + SimpleNamespace( + digest=f"baseline-{year}", + store_path=f"baselines/us/baseline-{year}/stage12-bl1-{year}.h5", + filename=f"stage12-bl1-{year}.h5", + simulation_id=f"stage12-bl1-{year}", + ), + ), + ) + monkeypatch.setattr(precompute, "stage12_output_plan_sha256", lambda _: "b" * 64) + + def fake_state_identity(resolved, *, country, year, dataset, **state_filter): + assert resolved == "resolved" + digest = ( + f"state-{year}-{state_filter['state_code']}-" + f"{state_filter['state_fips']}-{state_filter['filter_variable']}" + ) + return SimpleNamespace( + digest=digest, + store_path=f"state-datasets/{country}/{digest}/{dataset}_year_{year}.h5", + filename=f"{dataset}_year_{year}.h5", + ) + + monkeypatch.setattr( + precompute, "collect_state_dataset_identity", fake_state_identity + ) + targets = [ + precompute.StateDatasetTarget( + region="state/ut", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + dataset=REGIONAL_DATASET, + ), + precompute.StateDatasetTarget( + region="state/wy", + state_code="WY", + state_fips=56, + filter_variable="state_fips", + dataset=REGIONAL_DATASET, + ), + ] + monkeypatch.setattr(precompute, "state_dataset_targets", lambda: tuple(targets)) + return SimpleNamespace(client=client, targets=targets) + + +def test_plan_lists_every_state_dataset_year_major_with_presence( + stubbed_planning, +) -> None: + stubbed_planning.client.objects[ + f"state-datasets/us/state-2026-UT-49-state_fips/{REGIONAL_DATASET}_year_2026.h5" + ] = {"payload": b"h5", "metadata": {"sha256": "f" * 64, "size_bytes": "2"}} + + plan = precompute.plan_artifacts_impl("cache-bucket") + + assert [(entry.year, entry.state_code) for entry in plan.state_datasets] == [ + (2026, "UT"), + (2026, "WY"), + (2027, "UT"), + (2027, "WY"), + (2025, "UT"), + (2025, "WY"), + ] + assert [entry.exists for entry in plan.state_datasets] == [ + True, + False, + False, + False, + False, + False, + ] + first = plan.state_datasets[0] + assert first.region == "state/ut" + assert first.dataset == REGIONAL_DATASET + assert first.folder == f"state-datasets/{REGIONAL_DATASET}/UT" + assert first.filename == f"{REGIONAL_DATASET}_year_2026.h5" + assert len(plan.datasets) == 3 + + +def test_states_on_the_default_dataset_add_no_state_datasets( + stubbed_planning, monkeypatch +) -> None: + monkeypatch.setattr(precompute, "state_dataset_targets", tuple) + + assert precompute.plan_artifacts_impl("cache-bucket").state_datasets == () + + +def test_state_targets_follow_how_a_request_for_the_state_resolves( + monkeypatch, +) -> None: + from policyengine.core.scoping_strategy import RowFilterStrategy + from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + ) + + bundle = get_country_release_bundle("us") + + def region(state_code, fips, dataset_path): + return SimpleNamespace( + code=f"state/{state_code.lower()}", + dataset_path=dataset_path, + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=fips + ), + parent_code="us", + state_code=state_code, + ) + + regions = [ + region("WY", 56, bundle.dataset_uris[REGIONAL_DATASET]), + region("UT", 49, bundle.dataset_uris[REGIONAL_DATASET]), + region("ID", 16, bundle.default_dataset_uri), + ] + by_code = {item.code: item for item in regions} + country_module = SimpleNamespace( + model=SimpleNamespace( + get_region=by_code.get, + region_registry=SimpleNamespace( + get_by_type=lambda kind: list(regions) if kind == "state" else [] + ), + ) + ) + monkeypatch.setattr(precompute, "_country_module", lambda _country: country_module) + + targets = precompute.state_dataset_targets() + + assert [(target.region, target.state_fips) for target in targets] == [ + ("state/ut", 49), + ("state/wy", 56), + ] + assert {target.dataset for target in targets} == {REGIONAL_DATASET} + assert targets[0].scoping_strategy is by_code["state/ut"].scoping_strategy + + class _Call: def __init__(self, result, events, label): self.result = result @@ -228,6 +448,7 @@ def test_precompute_finishes_dataset_wave_before_baseline_wave() -> None: functions=precompute.CachePrecomputeFunctions( plan_artifacts=planner, build_dataset=dataset, + build_state_datasets=_Function(lambda *_: [], events, "state"), compute_baseline=baseline, verify_determinism=verify, publish_manifest=publish, @@ -250,6 +471,138 @@ def test_precompute_finishes_dataset_wave_before_baseline_wave() -> None: assert lines[-1] == "STAGE12_CACHE_MANIFEST_DIGEST=manifest-digest" +def _state_run_handles(events, plan): + def baseline_result(_bucket, entry): + return { + "path": entry["path"], + "content_sha256": "e" * 64, + "size_bytes": 1, + "year": entry["year"], + "segment_index": entry["segment_index"], + "simulation_id": entry["simulation_id"], + "cache_outcome": "miss", + "uploaded": True, + "compute_seconds": 1.0, + } + + def dataset_result(_bucket, entry): + return { + "path": entry["path"], + "content_sha256": "d" * 64, + "size_bytes": 1, + "year": entry["year"], + "uploaded": True, + "build_seconds": 1.0, + } + + def state_results(_bucket, year, entries): + return [ + { + "path": entry["path"], + "content_sha256": "c" * 64, + "size_bytes": 100, + "year": year, + "state_code": entry["state_code"], + "uploaded": entry["state_code"] != "WY", + "build_seconds": 12.5 if entry["state_code"] == "WY" else 6.0, + } + for entry in entries + ] + + return { + "plan_artifacts": _Function(lambda _: plan.model_dump(), events, "plan"), + "build_dataset": _Function(dataset_result, events, "dataset"), + "build_state_datasets": _Function(state_results, events, "state"), + "compute_baseline": _Function(baseline_result, events, "baseline"), + "verify_determinism": _Function( + lambda *_: {"equal": True, "differences": []}, events, "verify" + ), + "publish_manifest": _Function(lambda *_: "manifest-digest", events, "publish"), + } + + +def test_precompute_builds_each_year_of_state_datasets_in_one_call() -> None: + events = [] + plan = _plan_with_state_datasets() + handles = _state_run_handles(events, plan) + + precompute.run_precompute( + "cache-bucket", + functions=precompute.CachePrecomputeFunctions(**handles), + echo=lambda _line: None, + ) + + assert handles["build_state_datasets"].spawn_calls == [ + ( + "cache-bucket", + year, + [entry.model_dump() for entry in plan.state_datasets if entry.year == year], + ) + for year in precompute.PRECOMPUTE_YEARS + ] + + +def test_precompute_waits_for_state_datasets_only_before_publishing() -> None: + events = [] + handles = _state_run_handles(events, _plan_with_state_datasets()) + + precompute.run_precompute( + "cache-bucket", + functions=precompute.CachePrecomputeFunctions(**handles), + echo=lambda _line: None, + ) + + first_state_get = events.index("get:state") + assert events.index("spawn:state") < events.index("spawn:baseline") + assert events.index("remote:verify") < first_state_get + assert first_state_get < events.index("remote:publish") + + +def test_precompute_reports_each_builds_uploads_bytes_and_time() -> None: + events = [] + lines = [] + handles = _state_run_handles(events, _plan_with_state_datasets()) + + precompute.run_precompute( + "cache-bucket", + functions=precompute.CachePrecomputeFunctions(**handles), + echo=lines.append, + ) + + for year in precompute.PRECOMPUTE_YEARS: + assert ( + f"State datasets {year} {REGIONAL_DATASET}: 2 built, 1 uploaded " + "(100 bytes) in 12.5 s" + ) in lines + assert lines[-1] == "STAGE12_CACHE_MANIFEST_DIGEST=manifest-digest" + + +def test_precompute_builds_each_year_and_dataset_in_its_own_call() -> None: + events = [] + plan = _plan_with_state_datasets() + other = _state_entry(2026, "ID", exists=False).model_copy( + update={"dataset": "populace_other", "folder": "state-datasets/other/ID"} + ) + plan = plan.model_copy(update={"state_datasets": (*plan.state_datasets, other)}) + handles = _state_run_handles(events, plan) + + precompute.run_precompute( + "cache-bucket", + functions=precompute.CachePrecomputeFunctions(**handles), + echo=lambda _line: None, + ) + + assert [ + (call[1], {entry["dataset"] for entry in call[2]}) + for call in handles["build_state_datasets"].spawn_calls + ] == [ + (2026, {REGIONAL_DATASET}), + (2027, {REGIONAL_DATASET}), + (2025, {REGIONAL_DATASET}), + (2026, {"populace_other"}), + ] + + def test_warm_verified_precompute_skips_builds_and_determinism_check() -> None: events = [] plan = _plan(exists=True) @@ -265,6 +618,7 @@ def test_warm_verified_precompute_skips_builds_and_determinism_check() -> None: functions=precompute.CachePrecomputeFunctions( plan_artifacts=planner, build_dataset=unused, + build_state_datasets=unused, compute_baseline=unused, verify_determinism=verify, publish_manifest=publish, @@ -319,6 +673,7 @@ def test_failed_determinism_validation_prevents_manifest_publication() -> None: functions=precompute.CachePrecomputeFunctions( plan_artifacts=planner, build_dataset=dataset, + build_state_datasets=dataset, compute_baseline=baseline, verify_determinism=verify, publish_manifest=publish, @@ -345,6 +700,71 @@ def test_manifest_records_content_hashes_and_sizes() -> None: assert all(item.size_bytes == 123 for item in manifest.artifacts) +def test_manifest_lists_state_datasets_after_the_others_in_plan_order() -> None: + plan = _plan_with_state_datasets(exists=True) + store = SimpleNamespace( + describe=lambda path: CacheObject( + path=path, content_sha256="f" * 64, size_bytes=123 + ) + ) + + manifest = precompute.build_manifest(plan, store) + + assert len(manifest.artifacts) == 63 + len(plan.state_datasets) + state_artifacts = [ + artifact.model_dump() + for artifact in manifest.artifacts + if artifact.type == "state_dataset" + ] + assert state_artifacts == [ + { + "type": "state_dataset", + "path": entry.path, + "filename": entry.filename, + "folder": entry.folder, + "state_code": entry.state_code, + "year": entry.year, + "identity_digest": entry.digest, + "content_sha256": "f" * 64, + "size_bytes": 123, + } + for entry in plan.state_datasets + ] + assert [artifact.type for artifact in manifest.artifacts[:63]] == ( + ["dataset"] * 3 + ["baseline"] * 60 + ) + + +def test_manifest_without_state_datasets_keeps_its_published_shape() -> None: + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + + store = SimpleNamespace( + describe=lambda path: CacheObject( + path=path, content_sha256="f" * 64, size_bytes=123 + ) + ) + payload = precompute.build_manifest(_plan(exists=True), store).canonical_payload() + + assert all(set(item) != {"folder"} for item in payload["artifacts"]) + assert all("folder" not in item for item in payload["artifacts"]) + assert CacheManifest.model_validate(payload).canonical_payload() == payload + + +def test_manifest_with_state_datasets_round_trips_its_wire_shape() -> None: + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + + store = SimpleNamespace( + describe=lambda path: CacheObject( + path=path, content_sha256="f" * 64, size_bytes=123 + ) + ) + payload = precompute.build_manifest( + _plan_with_state_datasets(exists=True), store + ).canonical_payload() + + assert CacheManifest.model_validate(payload).canonical_payload() == payload + + @pytest.mark.parametrize("passes", [True, False]) def test_a_run_with_no_new_baselines_still_checks_unverified_ones( *, @@ -370,6 +790,7 @@ def run() -> None: functions=precompute.CachePrecomputeFunctions( plan_artifacts=planner, build_dataset=unused, + build_state_datasets=unused, compute_baseline=unused, verify_determinism=verify, publish_manifest=publish, @@ -549,6 +970,26 @@ def write_json(self, *_): assert downloads == expected_downloads +def test_installed_bundle_plans_a_regional_file_per_state_and_year( + monkeypatch, +) -> None: + """Uses the real bundle and US region registry.""" + client = _FakeClient() + _use_store(monkeypatch, client) + + plan = precompute.plan_artifacts_impl("cache-bucket") + + assert len(plan.state_datasets) == 51 * len(precompute.PRECOMPUTE_YEARS) + assert {entry.dataset for entry in plan.state_datasets} == {REGIONAL_DATASET} + utah = next( + entry + for entry in plan.state_datasets + if entry.year == 2026 and entry.state_code == "UT" + ) + assert utah.folder == f"state-datasets/{REGIONAL_DATASET}/UT" + assert utah.filename == f"{REGIONAL_DATASET}_year_2026.h5" + + def test_planned_baseline_file_is_the_one_the_worker_saves_and_loads( monkeypatch, ) -> None: @@ -597,3 +1038,240 @@ def test_planned_baseline_file_is_the_one_the_worker_saves_and_loads( # A US simulation always resolves an SPM selection, which .py names. assert entry.filename != f"{entry.simulation_id}.h5" assert plan.determinism_verified is False + + +class TestBuildStateDatasetsImpl: + """Building, checking and uploading one year's Stage 12 state datasets.""" + + YEAR = 2026 + + @pytest.fixture + def state_build(self, monkeypatch, tmp_path): + from fixtures import state_dataset_sources as sources + from policyengine.core.scoping_strategy import RowFilterStrategy + from policyengine.tax_benefit_models import us as us_module + + client = _FakeClient() + _use_store(monkeypatch, client) + state = SimpleNamespace( + client=client, + ensure_calls=[], + entity_frames=sources.source_entity_frames(), + work_folder=tmp_path / "work", + sources=sources, + ) + + def target(state_code, fips): + return precompute.StateDatasetTarget( + region=f"state/{state_code.lower()}", + state_code=state_code, + state_fips=fips, + filter_variable="state_fips", + dataset=sources.REGIONAL_DATASET, + scoping_strategy=RowFilterStrategy( + variable_name="state_fips", variable_value=fips + ), + ) + + state.targets = [ + target("UT", sources.UTAH_FIPS), + target("WY", sources.WYOMING_FIPS), + ] + monkeypatch.setattr( + precompute, "state_dataset_targets", lambda: tuple(state.targets) + ) + monkeypatch.setattr(precompute, "load_stage12_bundle", lambda: "resolved") + + def fake_state_identity(resolved, *, country, year, dataset, **state_filter): + digest = f"state-{year}-{state_filter['state_code']}" + return SimpleNamespace( + digest=digest, + store_path=f"state-datasets/{country}/{digest}/{dataset}_year_{year}.h5", + filename=f"{dataset}_year_{year}.h5", + ) + + monkeypatch.setattr( + precompute, "collect_state_dataset_identity", fake_state_identity + ) + + def fake_ensure_datasets(*, datasets, years, data_folder): + state.ensure_calls.append((datasets, years, data_folder)) + folder = Path(data_folder) + folder.mkdir(parents=True, exist_ok=True) + for year in years: + sources.write_year_file( + folder, year=year, entity_frames=state.entity_frames + ) + return us_module.load_datasets( + datasets=datasets, years=years, data_folder=data_folder + ) + + districts = [ + SimpleNamespace( + scoping_strategy=RowFilterStrategy( + variable_name="congressional_district_geoid", + variable_value=geoid, + ) + ) + for geoid in (4901, 4902, 4903, 4904, 5600) + ] + monkeypatch.setattr( + precompute, + "_country_module", + lambda _country: SimpleNamespace( + ensure_datasets=fake_ensure_datasets, + load_datasets=us_module.load_datasets, + model=SimpleNamespace( + region_registry=SimpleNamespace( + get_by_type=lambda kind: ( + districts if kind == "congressional_district" else [] + ) + ) + ), + ), + ) + state.identity_for = fake_state_identity + return state + + def _entry(self, state_build, state_code): + dataset = state_build.sources.REGIONAL_DATASET + identity = state_build.identity_for( + "resolved", + country="us", + year=self.YEAR, + dataset=dataset, + state_code=state_code, + ) + return StateDatasetPlanEntry( + year=self.YEAR, + state_code=state_code, + region=f"state/{state_code.lower()}", + dataset=dataset, + digest=identity.digest, + path=identity.store_path, + filename=identity.filename, + folder=f"state-datasets/{dataset}/{state_code}", + exists=False, + ) + + def _build(self, state_build, entries): + return precompute.build_state_datasets_impl( + "cache-bucket", self.YEAR, entries, work_folder=state_build.work_folder + ) + + def test_uploads_each_state_file_with_only_its_households(self, state_build): + entries = [self._entry(state_build, "UT"), self._entry(state_build, "WY")] + + results = self._build(state_build, entries) + + sources = state_build.sources + utah = state_build.work_folder / entries[0].folder / entries[0].filename + wyoming = state_build.work_folder / entries[1].folder / entries[1].filename + assert sources.stored_ids(utah, "household") == [1, 2] + assert sources.stored_ids(utah, "person") == [1, 2, 3] + assert sources.stored_weights(utah, "person") == [10.0, 11.0, 20.0] + assert sources.stored_ids(wyoming, "household") == [3] + assert [result.path for result in results] == [entry.path for entry in entries] + assert all(result.uploaded for result in results) + assert all(result.size_bytes > 0 for result in results) + assert state_build.client.objects[entries[0].path]["payload"] == ( + utah.read_bytes() + ) + + def test_dotpy_reuses_a_state_file_without_rebuilding( + self, state_build, monkeypatch + ): + from policyengine.tax_benefit_models.us import datasets as us_datasets + + entry = self._entry(state_build, "UT") + self._build(state_build, [entry]) + + def refuse_to_rebuild(**_kwargs): + msg = "ensure_datasets rebuilt a baked state file" + raise AssertionError(msg) + + monkeypatch.setattr(us_datasets, "create_datasets", refuse_to_rebuild) + (dataset,) = us_datasets.ensure_datasets( + datasets=[state_build.sources.REGIONAL_DATASET], + years=[self.YEAR], + data_folder=str(state_build.work_folder / entry.folder), + ).values() + + assert sorted(dataset.data.household["household_id"]) == [1, 2] + + def test_prepares_the_full_country_year_file_once(self, state_build): + self._build( + state_build, + [self._entry(state_build, "UT"), self._entry(state_build, "WY")], + ) + + assert len(state_build.ensure_calls) == 1 + datasets, years, _folder = state_build.ensure_calls[0] + assert datasets == [state_build.sources.REGIONAL_DATASET] + assert years == [self.YEAR] + + def test_refuses_a_district_that_lies_outside_its_state(self, state_build): + state_build.entity_frames = state_build.sources.source_entity_frames( + wyoming_district_geoid=4903 + ) + + with pytest.raises(RuntimeError, match="congressional district"): + self._build(state_build, [self._entry(state_build, "UT")]) + assert state_build.client.objects == {} + + def test_a_household_without_a_district_code_does_not_stop_the_build( + self, state_build + ): + state_build.entity_frames = state_build.sources.source_entity_frames( + wyoming_district_geoid=0 + ) + + results = self._build( + state_build, + [self._entry(state_build, "UT"), self._entry(state_build, "WY")], + ) + + assert [result.state_code for result in results] == ["UT", "WY"] + + def test_refuses_a_target_without_a_row_filter(self, state_build): + from dataclasses import replace + + state_build.targets[0] = replace(state_build.targets[0], scoping_strategy=None) + + with pytest.raises(RuntimeError, match="no row filter"): + self._build(state_build, [self._entry(state_build, "UT")]) + assert state_build.client.objects == {} + assert state_build.ensure_calls == [] + + def test_reports_how_long_the_years_build_had_run(self, state_build): + results = self._build( + state_build, + [self._entry(state_build, "UT"), self._entry(state_build, "WY")], + ) + + assert 0 <= results[0].build_seconds <= results[1].build_seconds + + def test_refuses_a_plan_entry_whose_key_differs_in_the_container(self, state_build): + entry = self._entry(state_build, "UT").model_copy( + update={"path": "state-datasets/us/other/x.h5"} + ) + + with pytest.raises(RuntimeError, match="identities differ"): + self._build(state_build, [entry]) + assert state_build.ensure_calls == [] + + def test_refuses_a_state_the_container_does_not_target(self, state_build): + state_build.targets = state_build.targets[:1] + + with pytest.raises(RuntimeError, match="state/wy"): + self._build(state_build, [self._entry(state_build, "WY")]) + assert state_build.ensure_calls == [] + + def test_refuses_entries_for_another_year(self, state_build): + with pytest.raises(ValueError, match="2027"): + precompute.build_state_datasets_impl( + "cache-bucket", + 2027, + [self._entry(state_build, "UT")], + work_folder=state_build.work_folder, + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index 597ab967d..00ae66afc 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -223,6 +223,17 @@ def test_only_us_worker_image_fetches_stage12_cache(monkeypatch) -> None: assert not any(call[0] == "run_function" for call in module.coordinator_image.calls) +def test_stage12_cache_fetch_has_room_for_the_state_datasets(monkeypatch) -> None: + module = _load(monkeypatch) + + (fetch,) = [ + call for call in module.us_worker_image.calls if call[0] == "run_function" + ] + # The 63 national files plus a file per state for each precomputed year, + # downloaded and hashed one after another; the main path allows the same. + assert fetch[2]["timeout"] == 3600 + + def test_v2_app_declares_validation_workers_and_non_http_coordinator( monkeypatch, ) -> None: diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py index 160be263a..c260afda0 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py @@ -21,6 +21,7 @@ def test_precompute_app_declares_independent_worker_resources(monkeypatch) -> No assert set(functions) == { "plan_artifacts", "build_dataset", + "build_state_datasets", "compute_baseline", "verify_determinism", "publish_manifest", @@ -28,6 +29,11 @@ def test_precompute_app_declares_independent_worker_resources(monkeypatch) -> No assert functions["build_dataset"]["cpu"] == 8.0 assert functions["build_dataset"]["memory"] == 65536 assert functions["build_dataset"]["timeout"] == 7200 + assert functions["build_state_datasets"]["cpu"] == 8.0 + # Preparing the regional dataset's full-country year file has used about + # 51 GB; the 51 filtered copies need room on top of that. + assert functions["build_state_datasets"]["memory"] == 98304 + assert functions["build_state_datasets"]["timeout"] == 7200 assert functions["compute_baseline"]["cpu"] == 8.0 assert functions["compute_baseline"]["memory"] == 32768 assert functions["compute_baseline"]["timeout"] == 3600 @@ -86,3 +92,50 @@ def test_precompute_app_mounts_only_the_read_only_hugging_face_secret( assert bundle_install[2]["secrets"] == [read_only_secret] for function_name, declaration in module.app.function_calls: assert declaration["secrets"] == [gcp_secret, read_only_secret], function_name + + +def test_state_dataset_wrapper_validates_entries_in_and_dumps_results_out( + monkeypatch, +) -> None: + import pytest + from policyengine_simulation_executor.stage12_cache.models import ( + StateDatasetBuildResult, + StateDatasetPlanEntry, + ) + from pydantic import ValidationError + + module = _load(monkeypatch) + entry = StateDatasetPlanEntry( + year=2026, + state_code="UT", + region="state/ut", + dataset="populace_us_2024_acs_local", + digest="s1", + path="state-datasets/us/s1/populace_us_2024_acs_local_year_2026.h5", + filename="populace_us_2024_acs_local_year_2026.h5", + folder="state-datasets/populace_us_2024_acs_local/UT", + exists=False, + ) + result = StateDatasetBuildResult( + path=entry.path, + content_sha256="c" * 64, + size_bytes=5, + year=2026, + state_code="UT", + uploaded=True, + build_seconds=1.0, + ) + seen = {} + + def fake_build(bucket, year, expected): + seen["call"] = (bucket, year, expected) + return (result,) + + monkeypatch.setattr(module, "build_state_datasets_impl", fake_build) + + assert module.build_state_datasets("cache-bucket", 2026, [entry.model_dump()]) == [ + result.model_dump() + ] + assert seen["call"] == ("cache-bucket", 2026, (entry,)) + with pytest.raises(ValidationError): + module.build_state_datasets("cache-bucket", 2026, [{"year": 2026}]) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index e91f4d759..83ee55cf0 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -539,11 +539,17 @@ def _frames(value=100.0): def test_calculator_uses_the_current_dataset_selection_contract(monkeypatch) -> None: from policyengine_simulation_executor import simulation_runtime + from policyengine_simulation_executor.simulation_runtime import ( + DatasetSelection, + RegionResolution, + ) from policyengine_simulation_executor.stage12_runtime import simulation as worker country_module = object() - region = type("Region", (), {"code": "us", "scoping_strategy": None})() - selection = object() + region = RegionResolution(code="us") + selection = DatasetSelection( + name="populace_us_2024", uri=_bundle().dataset.uri, is_default=True + ) dataset = object() received: dict[str, object] = {} @@ -587,9 +593,10 @@ def resolve_dataset(params, *, region_resolution): received["region"] = region_resolution return selection - def load_dataset(_params, *, selection, country_module): + def load_dataset(_params, *, selection, country_module, data_folder=None): received["load_selection"] = selection received["load_country_module"] = country_module + received["load_data_folder"] = data_folder return dataset def build_simulation(_params, *, resolved, **_kwargs): @@ -614,6 +621,7 @@ def build_simulation(_params, *, resolved, **_kwargs): assert received["region"] is region assert received["load_selection"] is selection assert received["load_country_module"] is country_module + assert received["load_data_folder"] is None assert received["build_dataset"] is dataset assert received["build_selection"] is selection assert received["ensured"] is True @@ -621,6 +629,343 @@ def build_simulation(_params, *, resolved, **_kwargs): assert set(result.frames) == {"household", "person"} +class _StoppedAfterLoadingError(Exception): + """Raised by a stubbed simulation build once the dataset has loaded.""" + + +def _stage12_load(monkeypatch, *, region, selection, baked_folder, planned=None): + """Run a Stage 12 calculation up to its dataset load; return what it did.""" + from policyengine_simulation_executor import simulation_runtime + from policyengine_simulation_executor.stage12_cache import state_datasets + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + loads = [] + lookups = [] + + def load_dataset(_params, *, selection, country_module, data_folder=None): + loads.append({"selection": selection, "data_folder": data_folder}) + return object() + + def lookup(**kwargs): + lookups.append(kwargs) + return baked_folder + + def stop_after_loading(*_args, **_kwargs): + raise _StoppedAfterLoadingError + + monkeypatch.setattr(worker, "_require_installed_bundle", lambda _: None) + monkeypatch.setattr(simulation_runtime, "setup_gcp_credentials", nullcontext) + monkeypatch.setattr( + simulation_runtime, "_country_module", lambda _country: object() + ) + monkeypatch.setattr(simulation_runtime, "_resolve_region", lambda **_: region) + monkeypatch.setattr( + simulation_runtime, + "_resolve_dataset_selection", + lambda _params, *, region_resolution: selection, + ) + monkeypatch.setattr(simulation_runtime, "_load_dataset", load_dataset) + monkeypatch.setattr(state_datasets, "baked_state_data_folder", lookup) + monkeypatch.setattr( + worker, "detect_uk_local_authority_metadata", lambda *_a, **_k: None + ) + monkeypatch.setattr( + simulation_runtime, "_require_certified_uk_weight_matrix", lambda *_: None + ) + monkeypatch.setattr(worker, "build_stage12_simulation", stop_after_loading) + + planned = planned or _planned_simulation(SimulationRole.BASELINE) + with pytest.raises(_StoppedAfterLoadingError): + worker.calculate_simulation_frames(planned) + return SimpleNamespace(loads=loads, lookups=lookups, year=planned.year) + + +def _regional_selection(): + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + + return DatasetSelection( + name="populace_us_2024_acs_local", + uri="hf://policyengine/populace-us/populace_us_2024_acs_local.h5@rev", + is_default=False, + ) + + +@pytest.mark.parametrize("region_code", ["state/ut", "congressional_district/UT-01"]) +def test_stage12_state_and_district_runs_read_their_states_baked_folder( + monkeypatch, region_code +) -> None: + from policyengine_simulation_executor.simulation_runtime import ( + RegionResolution, + resolve_data_folder, + ) + + loaded = _stage12_load( + monkeypatch, + region=RegionResolution(code=region_code, state_code="UT"), + selection=_regional_selection(), + baked_folder="/opt/policyengine/stage12-data/state-datasets/acs/UT", + ) + + assert loaded.loads == [ + { + "selection": _regional_selection(), + "data_folder": "/opt/policyengine/stage12-data/state-datasets/acs/UT", + } + ] + assert loaded.lookups == [ + { + "country": "us", + "selection": _regional_selection(), + "state_code": "UT", + "year": loaded.year, + "data_folder": resolve_data_folder(), + } + ] + + +def test_stage12_state_run_without_a_baked_file_loads_on_demand(monkeypatch) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + loaded = _stage12_load( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_regional_selection(), + baked_folder=None, + ) + + assert loaded.loads[0]["data_folder"] is None + + +REGIONAL_URI = "hf://policyengine/data/populace_us_2024_acs_local.h5@regional-revision" + + +def _regional_provenance() -> BundleProvenance: + return _bundle().model_copy( + update={ + "dataset": DatasetProvenance( + identity="populace_us_2024_acs_local", + uri=REGIONAL_URI, + artifact_revision="regional-revision", + data_package_name="populace-data", + data_package_version="0.1.0", + ) + } + ) + + +def _installed_bundle_with_regions(): + from policyengine_simulation_contract.stage12_bundle import ( + ResolvedStage12Bundle, + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, + Stage12RegionDataset, + ) + + national = Stage12Dataset( + identity="populace_us_2024", + uri=_bundle().dataset.uri, + artifact_revision="revision", + sha256="a" * 64, + repo_type="dataset", + ) + regional = Stage12Dataset( + identity="populace_us_2024_acs_local", + uri=REGIONAL_URI, + artifact_revision="regional-revision", + sha256="c" * 64, + repo_type="dataset", + ) + return ResolvedStage12Bundle( + bundle=Stage12BundleManifest( + policyengine_version="5.2.0", + policyengine_requirement="policyengine==5.2.0", + core_package_version="3.0.0", + core_package_requirement="policyengine-core==3.0.0", + countries=( + Stage12CountryBundle( + country="us", + country_package_name="policyengine-us", + country_package_version="1.764.6", + country_package_requirement="policyengine-us==1.764.6", + data_package_name="populace-data", + data_package_version="0.1.0", + data_release_version="revision", + data_artifact_revision="revision", + default_dataset=national.identity, + default_dataset_uri=national.uri, + datasets=(national, regional), + region_datasets=( + Stage12RegionDataset( + region_type="national", dataset=national.identity + ), + Stage12RegionDataset( + region_type="state", dataset=regional.identity + ), + Stage12RegionDataset( + region_type="congressional_district", + dataset=regional.identity, + ), + ), + ), + ), + ), + bundle_manifest_sha256="b" * 64, + ) + + +def _simulation_for(region: str, provenance: BundleProvenance): + simulation = _planned_simulation(SimulationRole.BASELINE) + return simulation.model_copy( + update={ + "geography": GeographySelection(country="us", region=region), + "bundle": provenance, + "population": DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=provenance.dataset.uri, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256="a" * 64, + ) + ), + } + ) + + +@pytest.mark.parametrize( + ("region", "provenance"), + [ + ("state/ut", _regional_provenance), + ("State/UT", _regional_provenance), + ("ut", _regional_provenance), + ("congressional_district/UT-01", _regional_provenance), + ("us", _bundle), + # Entry services released before regional provenance send the default + # dataset for state runs until they are updated. + ("state/ut", _bundle), + ], +) +def test_installed_bundle_accepts_the_dataset_a_region_loads( + monkeypatch, region, provenance +) -> None: + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + monkeypatch.setattr(worker, "load_stage12_bundle", _installed_bundle_with_regions) + + worker._require_installed_bundle(_simulation_for(region, provenance())) + + +def test_uk_regions_keep_the_default_dataset_when_the_bundle_names_region_types( + monkeypatch, +) -> None: + from policyengine_simulation_contract.stage12_bundle import Stage12RegionDataset + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + resolved = _installed_bundle_with_regions() + us = resolved.bundle.countries[0] + uk = us.model_copy( + update={ + "country": "uk", + "region_datasets": ( + Stage12RegionDataset( + region_type="national", dataset=us.default_dataset + ), + ), + } + ) + resolved = resolved.model_copy( + update={"bundle": resolved.bundle.model_copy(update={"countries": (us, uk)})} + ) + monkeypatch.setattr(worker, "load_stage12_bundle", lambda: resolved) + simulation = _simulation_for("us", _bundle()).model_copy( + update={"geography": GeographySelection(country="uk", region="country/wales")} + ) + + worker._require_installed_bundle(simulation) + with pytest.raises(RuntimeError, match="differs from the installed bundle"): + worker._require_installed_bundle( + simulation.model_copy(update={"bundle": _regional_provenance()}) + ) + + +@pytest.mark.parametrize( + ("region", "dataset"), + [ + ("us", "populace_us_2024_acs_local"), + ("state/ut", "source_enrichment"), + ], +) +def test_installed_bundle_refuses_a_dataset_the_region_does_not_load( + monkeypatch, region, dataset +) -> None: + from policyengine_simulation_executor.stage12_runtime import simulation as worker + + monkeypatch.setattr(worker, "load_stage12_bundle", _installed_bundle_with_regions) + provenance = _regional_provenance() + provenance = provenance.model_copy( + update={"dataset": provenance.dataset.model_copy(update={"identity": dataset})} + ) + + with pytest.raises(RuntimeError, match="differs from the installed bundle"): + worker._require_installed_bundle(_simulation_for(region, provenance)) + + +def _selection(name: str, uri: str, *, is_default: bool): + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + + return DatasetSelection(name=name, uri=uri, is_default=is_default) + + +def test_a_calculation_loads_the_dataset_its_provenance_names(monkeypatch) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + loaded = _stage12_load( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_selection( + "populace_us_2024_acs_local", REGIONAL_URI, is_default=False + ), + baked_folder=None, + planned=_simulation_for("state/ut", _regional_provenance()), + ) + + assert loaded.loads[0]["selection"].name == "populace_us_2024_acs_local" + + +def test_a_calculation_refuses_provenance_naming_another_dataset( + monkeypatch, +) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + with pytest.raises(RuntimeError, match="dataset its region loads"): + _stage12_load( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_selection( + "populace_us_2024", _bundle().dataset.uri, is_default=True + ), + baked_folder=None, + planned=_simulation_for("state/ut", _regional_provenance()), + ) + + +def test_default_provenance_from_an_older_entry_service_is_still_loaded( + monkeypatch, +) -> None: + from policyengine_simulation_executor.simulation_runtime import RegionResolution + + loaded = _stage12_load( + monkeypatch, + region=RegionResolution(code="state/ut", state_code="UT"), + selection=_selection( + "populace_us_2024_acs_local", REGIONAL_URI, is_default=False + ), + baked_folder=None, + planned=_simulation_for("state/ut", _bundle()), + ) + + assert loaded.loads[0]["selection"].name == "populace_us_2024_acs_local" + + def test_uk_boundary_version_is_detected_from_the_unscoped_dataset() -> None: dataset = type( "Dataset", @@ -4165,3 +4510,44 @@ def setup_that_uses_the_whole_window(_report): assert [len(call.probed_at) for call in invoker.calls.values()] == [1, 1] assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] assert store.parent.error_code == "simulation_wait_timeout" + + +def _runtime_region_type(country: str, region: str | None) -> str | None: + """Return the region type the runtime's own region normalisation implies.""" + from policyengine_simulation_executor.simulation_runtime import ( + _normalise_region_code, + ) + + code = _normalise_region_code(country, region) + if code == country: + return "national" + return code.split("/", 1)[0] if "/" in code else None + + +@pytest.mark.parametrize( + ("country", "region"), + [ + ("us", None), + ("us", ""), + ("us", " US "), + ("us", "uk"), + ("us", "state/ut"), + ("us", "State/UT"), + ("us", "ut"), + ("us", " ut "), + ("us", "congressional_district/ut-01"), + ("us", "place/UT-67000"), + ("us", "utah"), + ("uk", "uk"), + ("uk", "us"), + ("uk", "country/scotland"), + ("uk", "scotland"), + ("uk", "Northern Ireland"), + ("uk", "constituency/Aldershot"), + ("uk", "local_authority/E06000001"), + ], +) +def test_the_contract_reads_region_codes_as_the_runtime_does(country, region) -> None: + from policyengine_simulation_contract.stage12_bundle import region_type_of + + assert region_type_of(country, region) == _runtime_region_type(country, region) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py b/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py new file mode 100644 index 000000000..5b6eef5a2 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_state_datasets.py @@ -0,0 +1,238 @@ +"""Stage 12's per-state datasets: their cache identity and image folders. + +Stage 12 keeps its own copy of the main path's per-state dataset machinery, +keyed from the published Stage 12 bundle. +""" + +from __future__ import annotations + +import dataclasses +from pathlib import Path + +import pytest +from policyengine_simulation_contract.stage12_bundle import ( + ResolvedStage12Bundle, + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, +) +from policyengine_simulation_executor.stage12_cache import state_datasets +from policyengine_simulation_executor.stage12_cache.keys import ( + Stage12StateDatasetCacheIdentity, + collect_dataset_identity, + collect_state_dataset_identity, +) + +REGIONAL_DATASET = "populace_us_2024_acs_local" + + +def _resolved() -> ResolvedStage12Bundle: + national = Stage12Dataset( + identity="populace_us_2024", + uri="hf://policyengine/populace-us/populace_us_2024.h5@national-revision", + artifact_revision="national-revision", + sha256="a" * 64, + repo_type="dataset", + ) + regional = Stage12Dataset( + identity=REGIONAL_DATASET, + uri=f"hf://policyengine/populace-us/{REGIONAL_DATASET}.h5@regional-revision", + artifact_revision="regional-revision", + sha256="c" * 64, + repo_type="dataset", + ) + country = Stage12CountryBundle( + country="us", + country_package_name="policyengine-us", + country_package_version="1.0.0", + country_package_requirement="policyengine-us==1.0.0", + data_package_name="policyengine-us-data", + data_package_version="2.0.0", + data_release_version="2.0.0", + data_artifact_revision="national-revision", + default_dataset=national.identity, + default_dataset_uri=national.uri, + datasets=(national, regional), + ) + return ResolvedStage12Bundle( + bundle=Stage12BundleManifest( + policyengine_version="5.2.0", + policyengine_requirement="policyengine==5.2.0", + core_package_version="3.0.0", + core_package_requirement="policyengine-core==3.0.0", + countries=(country,), + ), + bundle_manifest_sha256="b" * 64, + ) + + +def _utah_2026() -> Stage12StateDatasetCacheIdentity: + return collect_state_dataset_identity( + _resolved(), + country="us", + year=2026, + dataset=REGIONAL_DATASET, + state_code="ut", + state_fips=49, + filter_variable="state_fips", + ) + + +def test_identity_is_taken_from_the_regional_dataset_entry(): + identity = _utah_2026() + + assert identity.dataset == REGIONAL_DATASET + assert identity.source_sha256 == "c" * 64 + assert identity.data_artifact_revision == "regional-revision" + assert identity.state_code == "UT" + assert identity.payload == { + "schema": "stage12-sds1", + "country": "us", + "dataset": REGIONAL_DATASET, + "year": 2026, + "source_sha256": "c" * 64, + "data_artifact_revision": "regional-revision", + "data_package_version": "2.0.0", + "country_package_version": "1.0.0", + "core_package_version": "3.0.0", + "policyengine_version": "5.2.0", + "bundle_manifest_sha256": "b" * 64, + "state_code": "UT", + "state_fips": 49, + "filter_variable": "state_fips", + } + + +def test_file_keeps_the_year_file_name_under_its_own_store_path(): + identity = _utah_2026() + + assert identity.filename == f"{REGIONAL_DATASET}_year_2026.h5" + assert identity.store_path == ( + f"state-datasets/us/{identity.digest}/{REGIONAL_DATASET}_year_2026.h5" + ) + + +@pytest.mark.parametrize( + "field", + [ + field.name + for field in dataclasses.fields(Stage12StateDatasetCacheIdentity) + if field.name not in {"stem", "country"} + ], +) +def test_every_identity_field_changes_the_digest(field): + identity = _utah_2026() + replacement = {"year": 2027, "state_fips": 56}.get(field, "other") + + assert dataclasses.replace(identity, **{field: replacement}).digest != ( + identity.digest + ) + + +def test_refuses_a_dataset_missing_from_the_bundle(): + with pytest.raises(ValueError, match="populace_missing"): + collect_state_dataset_identity( + _resolved(), + country="us", + year=2026, + dataset="populace_missing", + state_code="UT", + state_fips=49, + filter_variable="state_fips", + ) + + +def test_the_national_dataset_identity_is_unchanged(): + identity = collect_dataset_identity(_resolved(), country="us", year=2026) + + assert identity.dataset == "populace_us_2024" + assert identity.source_sha256 == "a" * 64 + assert identity.payload["schema"] == "stage12-ds1" + + +def test_state_folder_sits_under_the_stage12_data_folder(): + assert state_datasets.state_dataset_relative_folder(REGIONAL_DATASET, "ut") == ( + f"state-datasets/{REGIONAL_DATASET}/UT" + ) + assert state_datasets.state_dataset_folder( + "/opt/policyengine/stage12-data", REGIONAL_DATASET, "UT" + ) == (f"/opt/policyengine/stage12-data/state-datasets/{REGIONAL_DATASET}/UT") + + +def _regional_selection(): + from policyengine_simulation_executor.release_bundle import ( + get_country_release_bundle, + ) + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + + bundle = get_country_release_bundle("us") + return DatasetSelection( + name=REGIONAL_DATASET, + uri=bundle.dataset_uris[REGIONAL_DATASET], + is_default=False, + ) + + +def _bake_utah(data_folder: Path, *, stamp: bool = True) -> Path: + from fixtures.state_dataset_sources import write_year_file + + folder = Path( + state_datasets.state_dataset_folder(str(data_folder), REGIONAL_DATASET, "UT") + ) + folder.mkdir(parents=True) + write_year_file(folder, year=2026, stamp=stamp) + return folder + + +def _lookup(data_folder, selection, *, state_code="UT", year=2026, country="us"): + return state_datasets.baked_state_data_folder( + country=country, + selection=selection, + state_code=state_code, + year=year, + data_folder=str(data_folder), + ) + + +def test_a_state_reads_its_baked_folder(tmp_path): + folder = _bake_utah(tmp_path) + + assert _lookup(tmp_path, _regional_selection()) == str(folder) + + +def test_a_year_or_state_without_a_baked_file_reads_nothing_baked(tmp_path): + _bake_utah(tmp_path) + + assert _lookup(tmp_path, _regional_selection(), year=2028) is None + assert _lookup(tmp_path, _regional_selection(), state_code="WY") is None + + +def test_an_unstamped_file_reads_nothing_baked(tmp_path): + _bake_utah(tmp_path, stamp=False) + + assert _lookup(tmp_path, _regional_selection()) is None + + +def test_default_datasets_and_other_regions_read_nothing_baked(tmp_path): + from policyengine_simulation_executor.simulation_runtime import DatasetSelection + + _bake_utah(tmp_path) + default = DatasetSelection( + name="populace_us_2024", uri="hf://default", is_default=True + ) + + assert _lookup(tmp_path, default) is None + assert _lookup(tmp_path, _regional_selection(), state_code=None) is None + assert _lookup(tmp_path, _regional_selection(), country="uk") is None + + +def test_a_renamed_dotpy_reuse_check_falls_back_to_on_demand_loading( + monkeypatch, tmp_path +): + """A .py upgrade moving the private SPM profile helper must not fail runs.""" + from policyengine.tax_benefit_models.us import datasets as us_datasets + + _bake_utah(tmp_path) + monkeypatch.delattr(us_datasets, "_expected_legacy_spm_profile") + + assert _lookup(tmp_path, _regional_selection()) is None diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index 4e02df25e..eac3dde27 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -411,3 +411,126 @@ def test_digest_mismatch_is_checked_before_any_dataset_access() -> None: ) assert events == [] + + +def _installed_manifest(*, state_codes=("UT", "WY")): + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + from policyengine_simulation_executor.stage12_cache.precompute import ( + PRECOMPUTE_YEARS, + ) + from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, + ) + + def entry(kind, year, index): + return { + "type": kind, + "path": f"{kind}s/us/{year}-{index}/{kind}-{year}-{index}.h5", + "filename": f"{kind}-{year}-{index}.h5", + "year": year, + "identity_digest": f"{kind}-{year}-{index}", + "content_sha256": "f" * 64, + "size_bytes": 1, + } + + artifacts = [entry("dataset", year, 0) for year in PRECOMPUTE_YEARS] + artifacts += [ + entry("baseline", year, index) + for year in PRECOMPUTE_YEARS + for index in range(len(US_REGION_GROUPS)) + ] + artifacts += [ + { + "type": "state_dataset", + "path": f"state-datasets/us/{year}-{state}/acs_year_{year}.h5", + "filename": f"acs_year_{year}.h5", + "folder": f"state-datasets/acs/{state}", + "state_code": state, + "year": year, + "identity_digest": f"state-{year}-{state}", + "content_sha256": "f" * 64, + "size_bytes": 1, + } + for year in PRECOMPUTE_YEARS + for state in state_codes + ] + return CacheManifest.model_validate( + { + "schema": "stage12-mf1", + "country": "us", + "years": list(PRECOMPUTE_YEARS), + "partition_sha256": "d" * 64, + "bundle": { + "bundle_manifest_sha256": "b" * 64, + "policyengine_version": "5.2.0", + "core_package_version": "3.0.0", + "country_package_version": "1.0.0", + "data_package_version": "2.0.0", + "data_artifact_revision": "revision", + "default_dataset": "populace_us_2024", + }, + "artifacts": artifacts, + } + ) + + +def _install_state_files(data_folder, manifest): + for item in manifest.artifacts: + if item.type == "state_dataset": + folder = data_folder / item.folder + folder.mkdir(parents=True, exist_ok=True) + (folder / item.filename).write_bytes(b"h5") + + +def test_installed_cache_accepts_the_expected_state_files(tmp_path) -> None: + from policyengine_simulation_executor import stage12_worker_validation as module + + manifest = _installed_manifest() + _install_state_files(tmp_path, manifest) + + module._require_installed_artifacts(manifest, tmp_path, state_dataset_count=6) + + +def test_installed_cache_refuses_a_state_file_count_the_bundle_does_not_give( + tmp_path, +) -> None: + from policyengine_simulation_executor import stage12_worker_validation as module + + manifest = _installed_manifest(state_codes=("UT",)) + _install_state_files(tmp_path, manifest) + + with pytest.raises(RuntimeError, match="artifact count differs"): + module._require_installed_artifacts(manifest, tmp_path, state_dataset_count=6) + + +def test_installed_cache_refuses_a_missing_state_file(tmp_path) -> None: + from policyengine_simulation_executor import stage12_worker_validation as module + + manifest = _installed_manifest() + + with pytest.raises(RuntimeError, match="state dataset file is missing"): + module._require_installed_artifacts(manifest, tmp_path, state_dataset_count=6) + + +def test_expected_state_file_count_follows_the_installed_bundle(monkeypatch) -> None: + from policyengine_simulation_executor import stage12_worker_validation as module + from policyengine_simulation_executor.stage12_cache import precompute + + monkeypatch.setattr(precompute, "state_dataset_targets", lambda: ("UT", "WY")) + + assert module.expected_state_dataset_count() == 2 * len(precompute.PRECOMPUTE_YEARS) + + +def test_the_reported_cache_release_still_matches_its_contract_with_state_files( + tmp_path, +) -> None: + from policyengine_simulation_contract.stage12_manifest import Stage12CacheRelease + from policyengine_simulation_executor import stage12_worker_validation as module + + manifest = _installed_manifest() + + summary = module._cache_release_summary(manifest, "c" * 64) + + release = Stage12CacheRelease.model_validate(summary) + assert release.dataset_count == 3 + assert release.baseline_count == 60