Bake per-state ACS-local year files into the main path's worker image

docs/generated/rules.md

@@ -68,7 +68,7 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf
 - `sim.ci.multiline-run-is-script`: 13 remaining in `rules/baselines/ci-multiline-run.txt`
 - `sim.microdata.weight-stripping-confined`: 1 remaining in `rules/baselines/sim-microdata-weight-stripping.txt`
 - `sim.quality.pyright-apis-integ`: 4 remaining in `rules/baselines/pyright-apis-integ.txt`
-- `sim.quality.pyright-executor`: 178 remaining in `rules/baselines/pyright-executor.txt`
+- `sim.quality.pyright-executor`: 124 remaining in `rules/baselines/pyright-executor.txt`
 - `sim.quality.pyright-gateway`: 0 remaining in `rules/baselines/pyright-gateway.txt`
 
 ### Exceptions and review dates

projects/policyengine-simulation-executor/README.md

@@ -82,9 +82,12 @@ durable plan without changing routing.
 ## Baseline artifact pipeline (precompute → store → image fetch)
 
 The precompute app fills a content-addressed GCS store with single-year US
-datasets (2026, 2027, 2025) and the 20 per-cohort national baseline
-simulations per year; the deploy bakes them into the image, and the
-runtime loads baselines instead of computing them. Library logic lives in
+datasets (2026, 2027, 2025), the 20 per-cohort national baseline
+simulations per year, and per-state datasets: for each state region whose
+requests load a dataset other than the bundle default (the ACS-local
+dataset since policyengine 6.7), that dataset's year file filtered to the
+state's households, one file per state and year. The deploy bakes them into
+the image, and the runtime loads baselines instead of computing them. Library logic lives in
 `src/policyengine_simulation_executor/precompute.py` (keys:
 `artifact_keys.py`, store client: `artifact_store.py`); the Modal app
 (`src/modal/precompute_app.py`) is plumbing only.
@@ -103,7 +106,10 @@ The deploy job consumes the precompute job's manifest digest:
 `src/modal/app.py` resolves the manifest from GCS on the deploying
 machine and passes its content as the `fetch_artifacts` image layer's
 args (`src/modal/_image_setup.py`). The layer downloads the full
-artifact set (~430 MB) into `POLICYENGINE_DATA_FOLDER`, sitting between
+artifact set into `POLICYENGINE_DATA_FOLDER`: about 430 MB of national
+datasets and baselines, plus several GB of per-state datasets, each placed
+in `state-datasets/<dataset>/<state>/` under the year file's own name. It
+sits between
 the version env and the source mount — do not reorder. Because the
 manifest is content-addressed and rides in the layer args, Modal's layer
 cache busts exactly when the artifact set changes and never otherwise;
@@ -150,6 +156,12 @@ Operational notes:
 - Only pure default requests read the baked folder: any request naming an
   explicit `data` dataset or pinning a `data_version` bypasses it (see
   `_load_dataset` in `simulation_runtime.py`).
+- A US state or congressional-district request whose region loads a
+  non-default dataset reads its state's baked file when the image has one for
+  the requested year and .py's year-file reuse checks accept it (see
+  `state_datasets.baked_state_data_folder`). Otherwise it prepares the full
+  regional dataset during the request, as before. The run records which in
+  its `dataset_source` attribute: `baked_state` or `on_demand`.
 - Known residual, inherited rather than introduced: a data re-release
   under the same revision labels leaves the cached bundle-install layer —
   and therefore the receipt every artifact key derives from — unchanged.

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

@@ -27,14 +27,21 @@
 def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace:
     """Stub the identity seams with the canonical synthetic version-set.
 
-    Returns mutable bundle, receipt, and SPM-capability state, plus a handle
-    to the installed capability resolver for tests that intentionally exercise
-    the installed release.
+    Returns mutable bundle, receipt, release-manifest dataset and
+    SPM-capability state, plus a handle to the installed capability resolver
+    for tests that intentionally exercise the installed release. The bundle
+    has one non-default dataset, ``populace_cps_local``, for the per-state
+    dataset identities.
     """
     # This module imports ``get_release_manifest`` by value while constructing
     # country models. Import it before replacing the manifest module's function
     # so full-suite ordering cannot bind this test's stub into production code.
-    import policyengine.tax_benefit_models.common.model_version  # noqa: F401  # must load before the manifest stubs are installed (see above)
+    import policyengine.tax_benefit_models.common.model_version
+
+    # The US dataset module also imports the manifest helpers by value, and its
+    # year-file reuse checks (``_expected_legacy_spm_profile``) must keep
+    # reading the installed release after a stubbed test has imported it.
+    import policyengine.tax_benefit_models.us.datasets  # noqa: F401  # both modules must load before the manifest stubs are installed (see above)
     from policyengine.provenance import (
         manifest as manifest_module,
     )
@@ -51,8 +58,12 @@ def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace:
         default_dataset="populace_cps",
         dataset_uris={
             "populace_cps": "hf://org/repo/populace_cps.h5@rev-abc",
+            "populace_cps_local": "hf://org/repo/populace_cps_local.h5@rev-local",
         },
     )
+    release_datasets = {
+        "populace_cps_local": SimpleNamespace(sha256="c0ffee00" * 8),
+    }
     receipt_entry = {
         "country": "us",
         "version": "1.2.3",
@@ -61,6 +72,7 @@ def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace:
     state = SimpleNamespace(
         bundle=bundle,
         receipt_entry=receipt_entry,
+        release_datasets=release_datasets,
         spm_capability=None,
     )
 
@@ -78,13 +90,14 @@ def install_identity_stubs(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace:
     monkeypatch.setattr(
         manifest_module,
         "dataset_logical_name",
-        lambda _reference: state.bundle.default_dataset,
+        lambda reference: reference.rsplit("/", 1)[-1].split(".h5", 1)[0],
     )
     monkeypatch.setattr(
         manifest_module,
         "get_release_manifest",
         lambda _country: SimpleNamespace(
-            certification=SimpleNamespace(data_build_fingerprint="fp-123")
+            certification=SimpleNamespace(data_build_fingerprint="fp-123"),
+            datasets=state.release_datasets,
         ),
     )
 

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

@@ -0,0 +1,141 @@
+"""Small US year files for the per-state dataset tests.
+
+The files are written with policyengine.py's own ``PolicyEngineUSDataset.save()``
+and year-file stamp, and carry the SPM record the installed release expects
+for the regional dataset, so .py's reuse checks treat them as it would a year
+file ``create_datasets`` wrote. Three households: two in Utah (state FIPS 49)
+and one in Wyoming (56), with uneven weights on every entity.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pandas as pd
+
+from fixtures.weighted_frames import weighted_entity_frame
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from microdf import MicroDataFrame
+
+REGIONAL_DATASET = "populace_us_2024_acs_local"
+UTAH_FIPS = 49
+WYOMING_FIPS = 56
+
+
+def source_entity_frames(
+    *, wyoming_district_geoid: int = 5600
+) -> dict[str, MicroDataFrame]:
+    """Return the six US entity frames of the three-household sample."""
+    return {
+        "household": weighted_entity_frame(
+            {
+                "household_id": [1, 2, 3],
+                "household_weight": [10.0, 20.0, 30.0],
+                "state_fips": [UTAH_FIPS, UTAH_FIPS, WYOMING_FIPS],
+                "congressional_district_geoid": [4901, 4902, wyoming_district_geoid],
+            }
+        ),
+        "person": weighted_entity_frame(
+            {
+                "person_id": [1, 2, 3, 4],
+                "person_household_id": [1, 1, 2, 3],
+                "person_tax_unit_id": [1, 1, 2, 3],
+                "person_spm_unit_id": [1, 1, 2, 3],
+                "person_family_id": [1, 1, 2, 3],
+                "person_marital_unit_id": [1, 2, 3, 4],
+                "person_weight": [10.0, 11.0, 20.0, 30.0],
+                "age": [40, 12, 70, 33],
+            },
+            entity="person",
+        ),
+        "tax_unit": weighted_entity_frame(
+            {"tax_unit_id": [1, 2, 3], "tax_unit_weight": [10.0, 20.0, 30.0]},
+            entity="tax_unit",
+        ),
+        "spm_unit": weighted_entity_frame(
+            {
+                "spm_unit_id": [1, 2, 3],
+                "spm_unit_weight": [10.0, 20.0, 30.0],
+                "spm_unit_spm_universe_status": ["INCLUDED"] * 3,
+            },
+            entity="spm_unit",
+        ),
+        "family": weighted_entity_frame(
+            {"family_id": [1, 2, 3], "family_weight": [10.0, 20.0, 30.0]},
+            entity="family",
+        ),
+        "marital_unit": weighted_entity_frame(
+            {
+                "marital_unit_id": [1, 2, 3, 4],
+                "marital_unit_weight": [10.0, 11.0, 20.0, 30.0],
+            },
+            entity="marital_unit",
+        ),
+    }
+
+
+def year_file_metadata(dataset: str = REGIONAL_DATASET) -> dict:
+    """Return the metadata a ``create_datasets`` year file of ``dataset`` carries."""
+    from policyengine.tax_benefit_models.us.datasets import (
+        _expected_legacy_spm_profile,
+    )
+    from policyengine.tax_benefit_models.us.legacy_inputs import RENAMES_RECORD_KEY
+    from policyengine.tax_benefit_models.us.legacy_spm.record import (
+        RECORD_KEY,
+        LegacySPMRecord,
+    )
+
+    metadata: dict = {RENAMES_RECORD_KEY: {}}
+    profile = _expected_legacy_spm_profile(dataset)
+    if profile is not None:
+        metadata[RECORD_KEY] = LegacySPMRecord(
+            profile=profile, outside_spm_units=0, included_spm_units=3
+        ).model_dump(mode="json")
+    return metadata
+
+
+def write_year_file(
+    folder: Path,
+    *,
+    year: int,
+    dataset: str = REGIONAL_DATASET,
+    entity_frames: dict[str, MicroDataFrame] | None = None,
+    stamp: bool = True,
+) -> Path:
+    """Write ``<folder>/<dataset>_year_<year>.h5`` as ``create_datasets`` would."""
+    from policyengine.tax_benefit_models.us.datasets import (
+        PolicyEngineUSDataset,
+        USYearData,
+    )
+    from policyengine.tax_benefit_models.us.legacy_inputs import (
+        write_year_file_stamp,
+        year_file_stamp,
+    )
+
+    path = folder / f"{dataset}_year_{year}.h5"
+    PolicyEngineUSDataset(
+        name=f"{dataset}-year-{year}",
+        description=f"Test year file for {year}",
+        filepath=str(path),
+        year=year,
+        metadata=year_file_metadata(dataset),
+        data=USYearData(**(entity_frames or source_entity_frames())),
+    ).save()
+    if stamp:
+        write_year_file_stamp(path, year_file_stamp())
+    return path
+
+
+def stored_ids(path: Path, entity: str) -> list[int]:
+    """Return an entity's ids as stored in a year file, in row order."""
+    with pd.HDFStore(path, mode="r") as store:
+        return [int(value) for value in store[entity][f"{entity}_id"]]
+
+
+def stored_weights(path: Path, entity: str) -> list[float]:
+    """Return an entity's stored weights, in row order."""
+    with pd.HDFStore(path, mode="r") as store:
+        return [float(value) for value in store[entity][f"{entity}_weight"]]

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

@@ -15,7 +15,7 @@
 import json
 import logging
 import os
-from pathlib import Path
+from pathlib import Path, PurePosixPath
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
@@ -178,17 +178,58 @@ def _service_account_credentials(blob: str) -> "dict[str, Any]":
     }
 
 
+def _artifact_target(data_folder: "Path", artifact: "Mapping[str, Any]") -> "Path":
+    """Return where an artifact lands: its ``folder`` (if any) under the data folder.
+
+    Per-state datasets share one ``filename`` and differ by ``folder``, a
+    relative path such as ``state-datasets/<dataset>/<state>``. Anything that
+    could leave the data folder, or name a different file than intended, is
+    refused.
+    """
+    filename = artifact.get("filename")
+    if (
+        not isinstance(filename, str)
+        or filename in {"", ".", ".."}
+        or "/" in filename
+        or "\\" in filename
+    ):
+        msg = f"Artifact {artifact.get('path')!r} has an unsafe filename {filename!r}"
+        raise RuntimeError(msg)
+    folder = artifact.get("folder")
+    if folder is None:
+        return data_folder / filename
+    parts = (
+        PurePosixPath(folder).parts
+        if isinstance(folder, str) and folder and "\\" not in folder
+        else ()
+    )
+    if (
+        not parts
+        or PurePosixPath(folder).is_absolute()
+        or any(part in {".", ".."} for part in parts)
+        or "/".join(parts) != folder
+    ):
+        msg = f"Artifact {artifact.get('path')!r} has an unsafe folder {folder!r}"
+        raise RuntimeError(msg)
+    return data_folder.joinpath(*parts, filename)
+
+
 def _download_manifest_artifacts(
     bucket_handle: "storage.Bucket",
     artifacts: "Sequence[Mapping[str, Any]]",
     data_folder: "Path",
 ) -> None:
-    """Download every artifact under its runtime filename, failing on any gap.
+    """Download every artifact to its runtime location, failing on any gap.
 
-    Every store object is checked before anything downloads, and every
-    expected file is checked after.
+    Every target and store object is checked before anything downloads, and
+    every expected file is checked after.
     """
     logger = logging.getLogger(__name__)
+    targets = [_artifact_target(data_folder, artifact) for artifact in artifacts]
+    repeated = sorted({str(target) for target in targets if targets.count(target) > 1})
+    if repeated:
+        msg = f"The manifest lists these files more than once: {repeated}"
+        raise RuntimeError(msg)
     missing = [
         artifact["path"]
         for artifact in artifacts
@@ -206,16 +247,12 @@ def _download_manifest_artifacts(
     # Unconditional downloads: a same-named stale file must never
     # survive into the image, and the layer starts empty of these
     # files anyway.
-    for artifact in artifacts:
-        target = data_folder / artifact["filename"]
+    for artifact, target in zip(artifacts, targets, strict=True):
+        target.parent.mkdir(parents=True, exist_ok=True)
         bucket_handle.blob(artifact["path"]).download_to_filename(str(target))
         logger.info("Fetched %s (%.1f MB)", target, target.stat().st_size / 1e6)
 
-    absent = [
-        str(data_folder / artifact["filename"])
-        for artifact in artifacts
-        if not (data_folder / artifact["filename"]).exists()
-    ]
+    absent = [str(target) for target in targets if not target.exists()]
     if absent:
         msg = f"Artifact fetch did not produce expected files: {absent}"
         raise RuntimeError(msg)

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

@@ -235,7 +235,9 @@ def build_base_simulation_image() -> modal.Image:
             secrets=[gcp_secret],
             cpu=2.0,
             memory=4096,
-            timeout=900,
+            # The national files and baselines are about 430 MB; the
+            # per-state datasets (51 states for each year) add several GB.
+            timeout=3600,
         )
     )