Cut the deploy over to the artifact store: image fetch, deployed markers, prewarm retired

.github/scripts/modal-record-deployment.sh

@@ -0,0 +1,21 @@
+#!/bin/bash
+# Record a deployed environment's manifest digest in the artifact store.
+# Usage: ./modal-record-deployment.sh <environment> <manifest-digest>
+# Run from projects/policyengine-simulation-executor (like modal-deploy-app.sh).
+# Required env vars: POLICYENGINE_ARTIFACT_BUCKET (GCS artifact store bucket),
+# plus GCP credentials (GCP_CREDENTIALS_JSON) for the store write.
+
+set -euo pipefail
+
+ENVIRONMENT="${1:?Environment required (e.g. beta or prod)}"
+MANIFEST_DIGEST="${2:?Manifest digest required (the precompute job output)}"
+
+if [ -z "${POLICYENGINE_ARTIFACT_BUCKET:-}" ]; then
+  echo "POLICYENGINE_ARTIFACT_BUCKET is required (the GCS artifact store bucket)." >&2
+  exit 1
+fi
+
+echo "Recording deployment marker for ${ENVIRONMENT} (manifest ${MANIFEST_DIGEST})"
+uv run python -m src.modal.utils.record_deployment \
+  --environment "$ENVIRONMENT" \
+  --manifest-digest "$MANIFEST_DIGEST"

.github/workflows/modal-deploy.reusable.yml

@@ -159,6 +159,12 @@ jobs:
         POLICYENGINE_CORE_VERSION: ${{ steps.versions.outputs.policyengine_core_version }}
         POLICYENGINE_US_VERSION: ${{ steps.versions.outputs.us_version }}
         POLICYENGINE_UK_VERSION: ${{ steps.versions.outputs.uk_version }}
+        # app.py resolves the deploy manifest from GCS on the runner at
+        # modal.is_local() time, so this step needs the digest, the store
+        # bucket, and GCP credentials alongside the version env vars.
+        POLICYENGINE_MANIFEST_DIGEST: ${{ needs.precompute.outputs.manifest_digest }}
+        POLICYENGINE_ARTIFACT_BUCKET: ${{ vars.POLICYENGINE_ARTIFACT_BUCKET }}
+        GCP_CREDENTIALS_JSON: ${{ secrets.GCP_CREDENTIALS_JSON }}
       run: ../../.github/scripts/modal-deploy-app.sh "${{ inputs.modal_environment }}" "${{ inputs.force_latest }}"
 
     - name: Get deployed URL
@@ -172,6 +178,21 @@ jobs:
     - name: Verify deployment health
       run: .github/scripts/modal-health-check.sh "${{ steps.get-url.outputs.simulation_api_url }}"
 
+    # After the health check, so the marker means deployed AND healthy.
+    # The marker is the artifact GC's liveness signal: a deploy that
+    # cannot record itself fails the job.
+    - name: Record deployment marker
+      working-directory: projects/policyengine-simulation-executor
+      env:
+        POLICYENGINE_ARTIFACT_BUCKET: ${{ vars.POLICYENGINE_ARTIFACT_BUCKET }}
+        GCP_CREDENTIALS_JSON: ${{ secrets.GCP_CREDENTIALS_JSON }}
+        POLICYENGINE_VERSION: ${{ steps.versions.outputs.policyengine_version }}
+        POLICYENGINE_US_VERSION: ${{ steps.versions.outputs.us_version }}
+        POLICYENGINE_UK_VERSION: ${{ steps.versions.outputs.uk_version }}
+        US_DATA_VERSION: ${{ steps.versions.outputs.us_data_version }}
+        UK_DATA_VERSION: ${{ steps.versions.outputs.uk_data_version }}
+      run: ../../.github/scripts/modal-record-deployment.sh "${{ inputs.environment }}" "${{ needs.precompute.outputs.manifest_digest }}"
+
   integ_test:
     name: Run integration tests
     needs: [deploy]

README.md

@@ -83,7 +83,7 @@ manual deploy step. On merge to `main`, the Modal deploy workflow
 
 The stable gateway app is `policyengine-simulation-gateway`; executors deploy as
 versioned `policyengine-simulation-py{version}` apps. For Modal image and deploy
-specifics (image dependency pinning, dataset prebuild, observability), see the
+specifics (image dependency pinning, artifact fetch, observability), see the
 service READMEs:
 
 - `projects/policyengine-simulation-gateway/README.md`

projects/policyengine-simulation-executor/README.md

@@ -20,46 +20,16 @@ To change image dependencies, edit the `modal-simulation-image`
 dependency group and run `uv lock`. PRs touching image inputs run an
 in-image import smoke (`src/modal/smoke_app.py` via
 `.github/workflows/pr-image-smoke.yml`). Note that any change to the
-group or lock invalidates the image layer cache, including the dataset
-prebuild layer below.
+group or lock invalidates the image layer cache, including the artifact
+fetch layer below.
 
-## Temporary: prebuilt single-year datasets in the Modal image
-
-The Modal image prebuilds single-year datasets (2025–2027, US national
-default only — UK deliberately excluded to keep image builds short) into
-`POLICYENGINE_DATA_FOLDER` at image build time
-(`src/modal/_image_setup.py:prebuild_country_datasets`), so cold containers
-skip the slow runtime `ensure_datasets()` build. This is temporary until
-Populace publishes single-year datasets to Hugging Face — see issue #596 for
-the removal checklist; all code sites are greppable via `TEMPORARY`.
-
-Operational notes:
-
-- The prebuild layer runs on Modal's cloud image builder during
-  `modal deploy`, and only rebuilds when `POLICYENGINE_*` versions change or
-  the prebuild function itself is edited (even comment changes). It sits
-  before `add_local_python_source` on purpose — do not reorder.
-- The first deploy after a version bump pays the multi-hour build. Pre-warm
-  it off the deploy path with the minimum-reproduction app, which builds
-  only the layers up to the prebuild and then verifies the baked files and
-  load path from inside a container:
-  `uv run modal run --env=staging src/modal/prewarm_app.py`.
-  Because it shares the layer construction with `app.py`, a successful run
-  leaves the layers cached and the next real deploy fast-forwards.
-- To force a rebuild of a cached layer (e.g. a data re-release under the
-  same revision label), temporarily add `force_build=True` to the affected
-  `run_function` call in `src/modal/app.py` for one deploy.
-- 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`).
-
-## Artifact precompute (baseline artifact pipeline)
+## 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, so a later deploy can bake them into the image and
-the runtime can load baselines instead of computing them. Library logic
-lives in `src/policyengine_simulation_executor/precompute.py` (keys:
+simulations per 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.
 
@@ -73,6 +43,32 @@ Manual runs remain valid for ad-hoc warming:
     export POLICYENGINE_ARTIFACT_BUCKET=<bucket-name>
     uv run modal run --env=staging src/modal/precompute_app.py
 
+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
+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;
+there is no force-rebuild ritual. A freshness gate inside the layer
+refuses to bake artifacts computed for versions other than the image's
+own, and any missing store object fails the build loudly.
+
+After the health check, each deploy leg writes a
+`deployed/<environment>.json` marker (`deployed/beta.json`,
+`deployed/prod.json`) recording the manifest digest, versions, and run
+identity — the liveness signal for artifact garbage collection.
+
+Local `modal deploy` runs need the same inputs the CI deploy job has:
+
+    export POLICYENGINE_ARTIFACT_BUCKET=<bucket-name>
+    export POLICYENGINE_MANIFEST_DIGEST=<from a precompute run's MANIFEST_DIGEST= line>
+    # plus GCP credentials (GCP_CREDENTIALS_JSON, or gcloud ADC)
+
+Without a digest the deploy fails at the fetch layer, by design: a
+digest-less deploy can never silently ship an artifact-less image.
+
 Operational notes:
 
 - Idempotent by construction: artifact keys digest the full input closure
@@ -87,12 +83,22 @@ Operational notes:
   bytes can never drift), and protection against stale-code runners
   clobbering trusted artifacts. The heal procedure for a bad artifact is
   therefore always: delete its object from the bucket, then re-run the
-  precompute — deletion turns the key back into an ordinary miss.
+  precompute — deletion turns the key back into an ordinary miss, and the
+  next deploy's precompute job recomputes it before the deploy leg builds
+  the image.
 - A determinism gate runs whenever baselines were computed: one cohort's
   uploaded artifact is compared frame-by-frame against an independent
   fresh run. The run fails if they differ.
 - The final stdout line `MANIFEST_DIGEST=<digest>` names the published
   deploy manifest; the deploy pipeline consumes exactly that line.
+- 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`).
+- 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.
+  That staleness class predates the artifact pipeline and lives in the
+  bundle install, not the fetch.
 
 ## Observability
 

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

@@ -3,83 +3,156 @@
 
 These functions are executed during Modal image build and must not
 import any other modules from this package to avoid dependency issues.
-The dataset prebuild additionally runs BEFORE add_local_python_source,
+The artifact fetch additionally runs BEFORE add_local_python_source,
 so policyengine_simulation_executor is not importable there at all.
 """
 
-# TEMPORARY: remove once single-year datasets are published (see issue #596).
-# These years are baked into the image; uncovered years still build at runtime.
-# NOTE: this constant is a referenced global of prebuild_country_datasets, so
-# changing it (like editing the function body, even comments) invalidates the
-# cached multi-hour image layer.
-PREBUILD_DATASET_YEARS = [2025, 2026, 2027]
 
+def fetch_artifacts(bucket: str, manifest, *, client=None):
+    """Download every manifest-listed artifact into the image data folder.
 
-def prebuild_country_datasets(country: str):
-    """TEMPORARY: bake single-year national datasets into the image layer.
+    Runs as an image-build layer between ``.env(VERSION_ENV)`` and
+    ``add_local_python_source``, so it must stay self-contained (no
+    imports from this package; everything lazy in-body). The manifest
+    rides in the layer args: it is content-addressed, so the layer cache
+    busts exactly when the deployed artifact set changes.
 
-    Remove once single-year datasets are published (see issue #596).
+    ``manifest`` is None when the deploying machine resolved no manifest
+    (POLICYENGINE_MANIFEST_DIGEST unset) — importing the app module must
+    stay safe for the precompute and smoke apps, so that case fails here,
+    at build time, instead of at import time.
 
-    policyengine.py's ensure_datasets() rebuilds missing
-    ``{stem}_year_{year}.h5`` files at runtime with a full Microsimulation
-    pass, which dominates cold-container latency. Building them here bakes
-    the files into the image at POLICYENGINE_DATA_FOLDER so the runtime
-    existence check hits immediately.
+    ``client`` is a test seam; real builds construct a storage client from
+    the layer's GCP secret (or ambient ADC for local drill runs).
     """
+    import json
     import logging
     import os
-    from importlib import import_module
     from pathlib import Path
 
     logging.basicConfig(level=logging.INFO)
     logger = logging.getLogger(__name__)
 
-    os.environ.setdefault("POLICYENGINE_SKIP_COUNTRY_IMPORTS", "1")
-
-    from policyengine.provenance.manifest import (
-        dataset_logical_name,
-        get_release_manifest,
-        resolve_dataset_reference,
-    )
+    if not manifest:
+        raise RuntimeError(
+            "No artifact manifest was resolved on the deploying machine. "
+            "Set POLICYENGINE_MANIFEST_DIGEST (from the precompute run's "
+            "MANIFEST_DIGEST= output line) and POLICYENGINE_ARTIFACT_BUCKET, "
+            "with GCP credentials available, then redeploy."
+        )
 
-    data_folder = os.environ.get("POLICYENGINE_DATA_FOLDER", "/opt/policyengine/data")
-    # Pass the dataset name explicitly to mirror the runtime call shape in
-    # simulation_runtime._load_dataset (not every country module accepts
-    # datasets=None).
-    default_dataset = get_release_manifest(country).default_dataset
-    stem = dataset_logical_name(resolve_dataset_reference(country, default_dataset))
-    targets = {
-        year: Path(data_folder) / f"{stem}_year_{year}.h5"
-        for year in PREBUILD_DATASET_YEARS
+    # Freshness gate: refuse to bake artifacts computed for a different
+    # version-set than the one this image installs.
+    country = str(manifest["country"])
+    receipt = manifest["receipt"]
+    expected_env = {
+        "POLICYENGINE_VERSION": receipt["policyengine_version"],
+        f"POLICYENGINE_{country.upper()}_VERSION": receipt["model_version"],
     }
-    missing_years = [year for year, path in targets.items() if not path.exists()]
-    if not missing_years:
-        logger.info("All %s single-year datasets already present", country)
-        return
-
-    logger.info(
-        "Prebuilding %s dataset %s for years %s into %s",
-        country,
-        default_dataset,
-        missing_years,
-        data_folder,
-    )
-    country_module = import_module(f"policyengine.tax_benefit_models.{country}")
-    country_module.ensure_datasets(
-        datasets=[default_dataset],
-        years=missing_years,
-        data_folder=data_folder,
-    )
+    for env_var, manifest_value in expected_env.items():
+        image_value = os.environ.get(env_var)
+        if image_value != manifest_value:
+            raise RuntimeError(
+                f"Stale artifact manifest: the image has "
+                f"{env_var}={image_value!r} but the manifest was computed "
+                f"for {manifest_value!r}. Re-run the precompute for the "
+                "current version-set."
+            )
+
+    installed_data_version = None
+    receipt_path = os.environ.get("POLICYENGINE_BUNDLE_RECEIPT")
+    if receipt_path and Path(receipt_path).exists():
+        try:
+            installed = json.loads(Path(receipt_path).read_text(encoding="utf-8"))
+        except (OSError, json.JSONDecodeError):
+            installed = None
+        if isinstance(installed, dict):
+            for dataset in installed.get("datasets", []):
+                if isinstance(dataset, dict) and dataset.get("country") == country:
+                    installed_data_version = dataset.get("version")
+    if installed_data_version is None:
+        # Lenient by design: local drill runs have no bundle receipt file.
+        logger.warning(
+            "No installed bundle receipt entry for %s; "
+            "skipping the data-version freshness check",
+            country,
+        )
+    elif installed_data_version != receipt["data_version"]:
+        raise RuntimeError(
+            f"Stale artifact manifest: installed {country} data version is "
+            f"{installed_data_version!r} but the manifest was computed for "
+            f"{receipt['data_version']!r}. Re-run the precompute for the "
+            "current version-set."
+        )
 
-    still_missing = [str(path) for path in targets.values() if not path.exists()]
-    if still_missing:
-        # Fail the image build loudly rather than shipping an image that
-        # silently falls back to per-request dataset builds.
+    if client is None:
+        creds_kwargs = {}
+        blob = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS_JSON")
+        if blob:
+            # The secret payload is sometimes double-encoded (escaped
+            # interior quotes, with or without outer quotes) — same
+            # unwrap as the runtime's _normalize_credentials_blob.
+            try:
+                parsed = json.loads(blob)
+            except json.JSONDecodeError:
+                if blob.lstrip().startswith('"') or '\\"' in blob:
+                    parsed = json.loads(f'"{blob}"')
+                else:
+                    raise
+            info = json.loads(parsed) if isinstance(parsed, str) else parsed
+
+            from google.oauth2 import service_account
+
+            # Credentials built in memory, never written to disk: this
+            # layer's container filesystem is committed into the image,
+            # so a credentials file would ship inside it.
+            creds_kwargs = {
+                "credentials": service_account.Credentials.from_service_account_info(
+                    info
+                ),
+                "project": info.get("project_id"),
+            }
+        # No blob: ambient ADC (the local fault-drill path).
+        from google.cloud import storage
+
+        client = storage.Client(**creds_kwargs)
+
+    data_folder = Path(
+        os.environ.get("POLICYENGINE_DATA_FOLDER", "/opt/policyengine/data")
+    )
+    data_folder.mkdir(parents=True, exist_ok=True)
+    bucket_handle = client.bucket(bucket)
+    artifacts = manifest["artifacts"]
+
+    missing = [
+        artifact["path"]
+        for artifact in artifacts
+        if not bucket_handle.blob(artifact["path"]).exists()
+    ]
+    if missing:
+        # Fail the image build loudly rather than shipping an image
+        # with a partial artifact set. Heal: re-run the precompute
+        # (a deleted object is an ordinary miss to it).
         raise RuntimeError(
-            f"Prebuild did not produce expected dataset files: {still_missing}"
+            "Artifact store is missing objects the manifest lists: "
+            + ", ".join(missing)
         )
-    for year, path in targets.items():
-        logger.info("Prebuilt %s (%.1f MB)", path, path.stat().st_size / 1e6)
+
+    # 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"]
+        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()
+    ]
+    if absent:
+        raise RuntimeError(f"Artifact fetch did not produce expected files: {absent}")
 
 
 def snapshot_models():