Use ACS-local data for US regional simulations

.github/workflows/pr-image-smoke.yml

@@ -30,9 +30,10 @@ jobs:
     name: Image smoke (staging)
     runs-on: ubuntu-latest
     if: github.event.pull_request.head.repo.full_name == github.repository
-    # Warm cache: ~2 min. After a relock the executor smoke pays the
-    # certified dataset download layer (~15-20 min).
-    timeout-minutes: 45
+    # Image validation plus real US regional preparation/calculation. Prepared
+    # years are never reused; the 30-minute Modal calculation timeout and image
+    # construction must both fit inside this job's hard timeout.
+    timeout-minutes: 60
 
     steps:
     - uses: actions/checkout@v6

docs/us-regional-dataset-pr-validation.md

@@ -0,0 +1,56 @@
+# US regional dataset validation before merge
+
+The `PR image smoke` workflow checks more than package installation and the
+national dataset. Its executor command now also calls an ephemeral Modal
+function that:
+
+1. Resolves `state/ut` through the real executor region and dataset resolvers.
+2. Requires a non-default certified source with its revision and SHA-256 pin.
+3. Calls the same `ensure_datasets` API as production workers for 2026, using
+   a fresh temporary directory. No prepared year or baseline output is reused.
+4. Builds and runs a Utah baseline through the executor's simulation builder,
+   requesting WIC participation, WIC benefits, state FIPS, and net income.
+5. Requires complete boolean participation in prepared and calculated inputs,
+   nonempty Utah outputs, and finite net income and WIC benefits.
+6. Prints the installed `.py` version, selected source revision and hash, and
+   row counts. Calculation errors propagate and fail the PR check.
+
+The country year-preparation API currently prepares the shared ACS file before
+the simulation filters Utah. This is not just a small household-fixture test:
+the approximately 9.8-GB certified source must be downloaded and loaded. The
+function uses 8 CPUs, 64 GiB of memory, one container, and a hard 30-minute
+execution timeout. The workflow allows 60 minutes including image construction
+and its other checks. Timeouts and missing credentials are failures, not reasons
+to skip this validation.
+
+The function runs in an ephemeral Modal app using the same runtime image as
+the existing import check. It creates no scheduled process, public endpoint,
+artifact bucket, persistent volume, deployment, or manifest. Prepared files
+stay in its temporary directory; it does not write to production GCS or SQL.
+The existing US national artifact precompute remains unchanged.
+
+The source is the installed `.py` bundle's public US dataset. Existing GitHub
+Modal credentials launch the function; existing UK credentials continue to be
+used to construct and validate the shared image. No new environment variable,
+secret, permission, or provisioning step is introduced.
+
+The acceptance-condition unit tests use small DataFrames. They do not stand
+in for the real loading check. The real check is invoked by
+`.github/scripts/modal-image-smoke.sh` through `src/modal/smoke_app.py`, not by
+the local Docker integration command that excludes `beta_only` calculations.
+
+With the affected `.py` 6.2.2 release, real preparation is expected to fail on
+missing ACS WIC participation. Do not catch, skip, or mark that failure as
+expected in CI. The executor's two `.py` pins now target 6.2.5, including the
+narrowly scoped temporary compatibility fix in PolicyEngine/policyengine.py#566.
+
+## Pending publication and lockfile refresh
+
+Publish `.py` 6.2.5 before regenerating the executor's `uv.lock` with `uv lock`.
+Commit that generated lockfile, run `uv lock --check`, and rerun the real ACS
+validation before merging or deploying this change. Do not invent package
+URLs, hashes, or release metadata before publication.
+
+Until that refresh, the checked-in lockfile retains main's published 6.2.3
+release. Frozen installs therefore use 6.2.3, not the new target; this state
+is intentionally incomplete and must not be merged or deployed.

libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_bundle.py

@@ -36,12 +36,13 @@ class Stage12CountryBundle(StrictBundleModel):
     default_dataset: NonEmptyText
     default_dataset_uri: NonEmptyText
     datasets: tuple[Stage12Dataset, ...]
+    region_dataset_identities: dict[NonEmptyText, NonEmptyText]
 
 
 class Stage12BundleManifest(StrictBundleModel):
     """Stage 12's normalized subset of a PolicyEngine.py bundle."""
 
-    schema_version: Literal[1] = 1
+    schema_version: Literal[2] = 2
     source_bundle_schema_version: Literal[2] = 2
     policyengine_version: NonEmptyText
     policyengine_requirement: NonEmptyText
@@ -54,3 +55,43 @@ class Stage12BundleManifest(StrictBundleModel):
 class ResolvedStage12Bundle(StrictBundleModel):
     bundle: Stage12BundleManifest
     bundle_manifest_sha256: Sha256Digest
+
+
+def select_stage12_dataset(
+    country: Stage12CountryBundle,
+    region: str,
+) -> Stage12Dataset:
+    """Select the certified dataset declared for a concrete region type."""
+
+    normalized_region = region.strip().lower()
+    region_type = "national"
+    if normalized_region != country.country:
+        if (
+            country.country == "us"
+            and len(normalized_region) == 2
+            and normalized_region.isalpha()
+        ):
+            region_type = "state"
+        else:
+            requested_region_type, separator, _ = normalized_region.partition("/")
+            if separator:
+                region_type = requested_region_type
+    dataset_identity = country.region_dataset_identities.get(region_type)
+    if dataset_identity is None:
+        raise ValueError(
+            f"No certified dataset is declared for region type {region_type!r}."
+        )
+    selected = next(
+        (
+            dataset
+            for dataset in country.datasets
+            if dataset.identity == dataset_identity
+        ),
+        None,
+    )
+    if selected is None:
+        raise ValueError(
+            f"Certified region type {region_type!r} names absent dataset "
+            f"{dataset_identity!r}."
+        )
+    return selected

libs/policyengine-simulation-contract/tests/test_stage12_manifest.py

@@ -11,6 +11,7 @@
     Stage12BundleManifest,
     Stage12CountryBundle,
     Stage12Dataset,
+    select_stage12_dataset,
 )
 from policyengine_simulation_contract.stage12_manifest import (
     ACTIVE_MANIFEST_KEY,
@@ -41,13 +42,37 @@ def __setitem__(self, key: str, value: object) -> None:
         self.values[key] = value
 
 
+def _dataset_fixture(identity: str, sha256: str) -> Stage12Dataset:
+    """Build synthetic provenance without reading deployed dataset metadata."""
+    revision = f"{identity}-revision"
+    return Stage12Dataset(
+        identity=identity,
+        uri=f"hf://policyengine/data/{identity}.h5@{revision}",
+        artifact_revision=revision,
+        sha256=sha256,
+        repo_type="dataset",
+    )
+
+
 def _resolved_bundle(version: str = "5.2.0") -> ResolvedStage12Bundle:
-    countries = []
+    countries: list[Stage12CountryBundle] = []
     for country, package, package_version, dataset in (
         ("us", "policyengine-us", "1.764.6", "populace_us_2024"),
         ("uk", "policyengine-uk", "2.90.2", "populace_uk_2023"),
     ):
-        revision = f"{dataset}-revision"
+        national_dataset = _dataset_fixture(dataset, "a" * 64)
+        datasets = [national_dataset]
+        region_dataset_identities = {"national": dataset}
+        if country == "us":
+            regional_dataset = _dataset_fixture("populace_us_2024_acs_local", "c" * 64)
+            datasets.append(regional_dataset)
+            region_dataset_identities.update(
+                {
+                    "state": regional_dataset.identity,
+                    "congressional_district": regional_dataset.identity,
+                }
+            )
+        revision = national_dataset.artifact_revision
         countries.append(
             Stage12CountryBundle(
                 country=country,
@@ -59,16 +84,9 @@ def _resolved_bundle(version: str = "5.2.0") -> ResolvedStage12Bundle:
                 data_release_version=revision,
                 data_artifact_revision=revision,
                 default_dataset=dataset,
-                default_dataset_uri=f"hf://policyengine/data/{dataset}.h5@{revision}",
-                datasets=(
-                    Stage12Dataset(
-                        identity=dataset,
-                        uri=f"hf://policyengine/data/{dataset}.h5@{revision}",
-                        artifact_revision=revision,
-                        sha256="a" * 64,
-                        repo_type="dataset",
-                    ),
-                ),
+                default_dataset_uri=national_dataset.uri,
+                datasets=tuple(datasets),
+                region_dataset_identities=region_dataset_identities,
             )
         )
     bundle = Stage12BundleManifest(
@@ -131,6 +149,43 @@ def test_publish_writes_one_complete_v2_document() -> None:
     assert manifest.versions["5.2.0"].validation.validated is True
 
 
+@pytest.mark.parametrize(
+    "region, expected_dataset",
+    [
+        ("us", "populace_us_2024"),
+        ("CA", "populace_us_2024_acs_local"),
+        ("DC", "populace_us_2024_acs_local"),
+        ("state/ca", "populace_us_2024_acs_local"),
+        ("state/DC", "populace_us_2024_acs_local"),
+        ("congressional_district/CA-01", "populace_us_2024_acs_local"),
+        ("congressional_district/DC-01", "populace_us_2024_acs_local"),
+    ],
+)
+def test_manifest_round_trip_preserves_regional_dataset_selection(
+    region: str, expected_dataset: str
+) -> None:
+    store = FakeStore()
+    worker = _worker()
+    publish_v2_manifest(store=store, worker=worker)
+
+    _, loaded_worker, _ = V2ManifestLoader(store).resolve()
+    us = next(
+        country for country in loaded_worker.bundle.countries if country.country == "us"
+    )
+    selected = select_stage12_dataset(us, region)
+
+    assert selected.identity == expected_dataset
+    original_us = next(
+        country for country in worker.bundle.countries if country.country == "us"
+    )
+    assert selected == next(
+        dataset
+        for dataset in original_us.datasets
+        if dataset.identity == expected_dataset
+    )
+    assert us.region_dataset_identities == original_us.region_dataset_identities
+
+
 def test_loader_retains_its_last_valid_v2_document() -> None:
     store = FakeStore()
     expected = publish_v2_manifest(store=store, worker=_worker())

projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py

@@ -7,7 +7,10 @@
 from typing import Any, cast
 from uuid import UUID, uuid5
 
-from policyengine_simulation_contract.stage12_bundle import CountryId
+from policyengine_simulation_contract.stage12_bundle import (
+    CountryId,
+    select_stage12_dataset,
+)
 from policyengine_simulation_contract.stage12_execution import (
     BundleProvenance,
     DatasetArtifactMediaType,
@@ -134,11 +137,7 @@ def adapt_annual_comparison(
     if not isinstance(region, str) or not region:
         return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE)
 
-    selected_dataset = next(
-        item
-        for item in bundle_country.datasets
-        if item.identity == bundle_country.default_dataset
-    )
+    selected_dataset = select_stage12_dataset(bundle_country, region)
     requested_dataset = payload.get("data")
     accepted_dataset_values = {
         None,

projects/policyengine-simulation-entry/tests/stage12_fixtures.py

@@ -36,6 +36,33 @@ def worker() -> V2WorkerVersion:
         ("uk", "policyengine-uk", "2.90.2", "populace_uk_2023"),
     ):
         revision = f"{dataset}-revision"
+        datasets = [
+            Stage12Dataset(
+                identity=dataset,
+                uri=f"hf://policyengine/data/{dataset}.h5@{revision}",
+                artifact_revision=revision,
+                sha256="a" * 64,
+                repo_type="dataset",
+            )
+        ]
+        region_dataset_identities = {"national": dataset}
+        if country == "us":
+            local_dataset = "populace_us_2024_acs_local"
+            local_revision = f"{local_dataset}-revision"
+            datasets.append(
+                Stage12Dataset(
+                    identity=local_dataset,
+                    uri=(f"hf://policyengine/data/{local_dataset}.h5@{local_revision}"),
+                    artifact_revision=local_revision,
+                    sha256="c" * 64,
+                    repo_type="dataset",
+                )
+            )
+            region_dataset_identities = {
+                "national": dataset,
+                "state": local_dataset,
+                "congressional_district": local_dataset,
+            }
         countries.append(
             Stage12CountryBundle(
                 country=country,
@@ -48,15 +75,8 @@ def worker() -> V2WorkerVersion:
                 data_artifact_revision=revision,
                 default_dataset=dataset,
                 default_dataset_uri=f"hf://policyengine/data/{dataset}.h5@{revision}",
-                datasets=(
-                    Stage12Dataset(
-                        identity=dataset,
-                        uri=f"hf://policyengine/data/{dataset}.h5@{revision}",
-                        artifact_revision=revision,
-                        sha256="a" * 64,
-                        repo_type="dataset",
-                    ),
-                ),
+                datasets=tuple(datasets),
+                region_dataset_identities=region_dataset_identities,
             )
         )
         country_workers.append(

projects/policyengine-simulation-entry/tests/test_stage12_adapter.py

@@ -56,6 +56,26 @@ def test_absent_data_resolves_exact_bundle_default() -> None:
     assert result.report.baseline.bundle.dataset.artifact_revision.endswith("-revision")
 
 
+@pytest.mark.parametrize(
+    "region",
+    ["DC", "state/DC", "congressional_district/DC-01"],
+)
+def test_regional_request_resolves_declared_acs_local_dataset(region: str) -> None:
+    payload = {**eligible_payload(), "region": region}
+
+    result = adapt_annual_comparison(
+        payload,
+        evaluation_id=EVALUATION_ID,
+        worker=worker(),
+    )
+
+    assert result.report is not None
+    baseline = result.report.baseline
+    assert baseline.bundle.dataset.identity == "populace_us_2024_acs_local"
+    assert baseline.population.artifact.uri == baseline.bundle.dataset.uri
+    assert baseline.population.artifact.content_sha256 == "c" * 64
+
+
 def test_absent_package_versions_resolve_from_the_selected_bundle() -> None:
     payload = eligible_payload()
     payload.pop("version")

projects/policyengine-simulation-executor/fixtures/identity_stubs.py

@@ -44,6 +44,9 @@ def install_identity_stubs(monkeypatch):
         dataset_uris={
             "populace_cps": "hf://org/repo/populace_cps.h5@rev-abc",
         },
+        dataset_data_versions={"populace_cps": "1.2.3"},
+        dataset_revisions={"populace_cps": "rev-abc"},
+        dataset_sha256s={"populace_cps": "feedbead" * 8},
     )
     receipt_entry = {
         "country": "us",

projects/policyengine-simulation-executor/pyproject.toml

@@ -21,7 +21,7 @@ dependencies = [
     "policyengine-stage12-persistence",
     # PolicyEngine.py's models extra is the only selector for core, country
     # models, and the SPM calculator. Do not duplicate those versions here.
-    "policyengine[models]==6.2.3",
+    "policyengine[models]==6.2.5",
     "tables>=3.10.2",
     "modal>=0.73.0",
     "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.2,<4",
@@ -46,7 +46,7 @@ packages = ["src/policyengine_simulation_executor"]
 # certified datasets and writes their receipt.
 [dependency-groups]
 modal-simulation-image = [
-    "policyengine[models]==6.2.3",
+    "policyengine[models]==6.2.5",
     "fastapi>=0.115.0",
     "tables>=3.10.2",
     "policyengine-observability[fastapi,google,otlp-grpc]>=3.0.2,<4",

projects/policyengine-simulation-executor/src/modal/app.py

@@ -261,7 +261,7 @@ def _set_modal_call_attributes(runtime) -> None:
 @app.function(
     image=simulation_image,
     cpu=8.0,
-    memory=32768,
+    memory=65536,
     timeout=3600,
     retries=0,
     max_containers=100,