Bake per-state ACS-local year files into Stage 12 and record the dataset state runs load

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/<dataset>/<state>/` 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

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):

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