diff --git a/.github/scripts/modal-record-deployment.sh b/.github/scripts/modal-record-deployment.sh new file mode 100755 index 000000000..87c6f70ba --- /dev/null +++ b/.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 +# 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" diff --git a/.github/workflows/modal-deploy.reusable.yml b/.github/workflows/modal-deploy.reusable.yml index 65ba662ac..1998b4def 100644 --- a/.github/workflows/modal-deploy.reusable.yml +++ b/.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] diff --git a/README.md b/README.md index 6e8302c5d..a85819f01 100644 --- a/README.md +++ b/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` diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 4c274f4b1..a5f64c066 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/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= 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/.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= + export POLICYENGINE_MANIFEST_DIGEST= + # 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=` 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 diff --git a/projects/policyengine-simulation-executor/src/modal/_image_setup.py b/projects/policyengine-simulation-executor/src/modal/_image_setup.py index 13a22899f..f8f7183c3 100644 --- a/projects/policyengine-simulation-executor/src/modal/_image_setup.py +++ b/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(): diff --git a/projects/policyengine-simulation-executor/src/modal/app.py b/projects/policyengine-simulation-executor/src/modal/app.py index 5c3822714..be83e0f7e 100644 --- a/projects/policyengine-simulation-executor/src/modal/app.py +++ b/projects/policyengine-simulation-executor/src/modal/app.py @@ -13,7 +13,7 @@ from policyengine_observability import operation, set_attribute -from src.modal._image_setup import prebuild_country_datasets, snapshot_models +from src.modal._image_setup import fetch_artifacts, snapshot_models from src.modal.dependency_pins import project_dependency_pin from policyengine_simulation_observability.logfire_legacy import ( configure_logfire, @@ -116,6 +116,48 @@ def get_app_name(policyengine_version: str) -> str: _UV_PROJECT_DIR = str(Path(__file__).resolve().parents[2]) if modal.is_local() else "." +def _deploy_time_artifact_inputs() -> tuple[str, dict | None]: + """Resolve the artifact bucket and deploy manifest on the deploying machine. + + The manifest content rides into the fetch layer's args, so Modal's + layer cache busts exactly when the deployed artifact set changes. + Resolution must never hard-fail at import: this module is imported + where no digest can exist (inside containers, by the precompute app + that produces the digest, by the PR image smoke), so a missing digest + yields a None manifest and fetch_artifacts fails at build time + instead — a digest-less deploy cannot silently ship an artifact-less + image. + + A resolved manifest is validated as ArtifactManifest and crosses into + the layer as its canonical payload: fetch_artifacts cannot import the + model by design (the layer is import-restricted), so the dict exists + only at that serialization boundary. + """ + if not modal.is_local(): + return "", None + digest = os.environ.get("POLICYENGINE_MANIFEST_DIGEST") + bucket = os.environ.get("POLICYENGINE_ARTIFACT_BUCKET", "") + if not digest: + return bucket, None + from policyengine_simulation_executor.artifact_store import ArtifactStore + from policyengine_simulation_executor.precompute_models import ArtifactManifest + + store = ArtifactStore(bucket or None) + payload = store.read_manifest(digest) + if payload is None: + raise RuntimeError( + f"Artifact manifest {digest} is not in the store: the deploy " + "must consume a digest published by the precompute run." + ) + # Validate on the runner so a corrupt or foreign manifest fails the + # deploy here, not mid-image-build. + manifest = ArtifactManifest.model_validate(payload) + return store.bucket_name, manifest.canonical_payload() + + +_ARTIFACT_BUCKET, _DEPLOY_MANIFEST = _deploy_time_artifact_inputs() + + def bundle_install_command(policyengine_version: str) -> str: return " ".join( [ @@ -146,10 +188,10 @@ def bundle_install_command(policyengine_version: str) -> str: def build_runtime_simulation_image() -> modal.Image: - """Image layers up to the version env — everything except the dataset - prebuild and model snapshot. + """Image layers up to the version env — everything except the artifact + fetch and model snapshot. - Shared by the deployed app, the prewarm app, and the image smoke app + Shared by the deployed app and the image smoke app (src/modal/smoke_app.py): all must construct these layers through this one code path so their definitions — and therefore Modal's content-addressed layer cache keys — are identical. @@ -175,25 +217,23 @@ def build_runtime_simulation_image() -> modal.Image: def build_base_simulation_image() -> modal.Image: - """Runtime image layers plus the dataset prebuild.""" + """Runtime image layers plus the artifact fetch.""" return ( build_runtime_simulation_image() - # TEMPORARY: remove once single-year datasets are published (issue - # #596). Prebuild US single-year datasets into the image so cold - # containers skip the slow runtime build. US only for now, to keep - # image build time low — UK requests still build at request time. - # This layer MUST stay before add_local_python_source — that layer - # is keyed on source file hashes, so anything after it rebuilds on - # every code change, and this layer takes hours. To force a rebuild - # of a cached layer (e.g. after a data re-release under the same - # revision), temporarily add force_build=True. + # Bake the precomputed store artifacts (single-year US datasets + # and cohort baselines) into the image. This layer MUST stay + # before add_local_python_source — that layer is keyed on source + # file hashes, so anything after it rebuilds on every code + # change. Its args carry the content-addressed manifest, so the + # cache busts exactly when the artifact set changes and never + # otherwise. UK requests still build datasets at request time. .run_function( - prebuild_country_datasets, - args=("us",), - secrets=[data_secret, hf_secret], - cpu=8.0, - memory=65536, - timeout=4 * 60 * 60, + fetch_artifacts, + args=(_ARTIFACT_BUCKET, _DEPLOY_MANIFEST), + secrets=[gcp_secret], + cpu=2.0, + memory=4096, + timeout=900, ) ) diff --git a/projects/policyengine-simulation-executor/src/modal/prewarm_app.py b/projects/policyengine-simulation-executor/src/modal/prewarm_app.py deleted file mode 100644 index 18e43684f..000000000 --- a/projects/policyengine-simulation-executor/src/modal/prewarm_app.py +++ /dev/null @@ -1,91 +0,0 @@ -"""TEMPORARY: minimum reproduction / pre-warm app for the dataset prebuild. - -Remove once single-year datasets are published (see issue #596). - -Builds ONLY the image layers up to and including the single-year dataset -prebuild (no local source, no model snapshot), then runs a container from -that image to verify the baked files exist and that the runtime load path -cache-hits. Because the layers are constructed via -``app.build_base_simulation_image()``, they are byte-identical to the real -app's layers, so a successful run leaves them in Modal's workspace image -cache and the next real deploy fast-forwards through them. - -Usage (from projects/policyengine-simulation-executor): - - uv run modal run --env=staging src/modal/prewarm_app.py - -The slow part is the image build (streamed before the function runs); the -verification function itself takes seconds. Interrupting the local client -mid-build does not deploy anything — the app here is ephemeral. -""" - -import json - -import modal - -from src.modal._image_setup import PREBUILD_DATASET_YEARS -from src.modal.app import build_base_simulation_image - -app = modal.App("policyengine-simulation-prebuild-prewarm") - - -@app.function( - image=build_base_simulation_image(), - cpu=4.0, - memory=32768, - timeout=1800, - serialized=True, -) -def verify_prebuilt_datasets(years: list[int]) -> dict: - """Verify baked datasets from inside a container on the built image. - - Serialized so the container never imports this module (the base image - deliberately lacks the local source packages). - """ - import os - import time - from importlib import import_module - from pathlib import Path - - os.environ.setdefault("POLICYENGINE_SKIP_COUNTRY_IMPORTS", "1") - from policyengine.provenance.manifest import ( - dataset_logical_name, - get_release_manifest, - resolve_dataset_reference, - ) - - data_folder = Path( - os.environ.get("POLICYENGINE_DATA_FOLDER", "/opt/policyengine/data") - ) - default_dataset = get_release_manifest("us").default_dataset - stem = dataset_logical_name(resolve_dataset_reference("us", default_dataset)) - - report: dict[str, str] = {} - missing = [] - for year in years: - path = data_folder / f"{stem}_year_{year}.h5" - if path.exists(): - report[path.name] = f"{path.stat().st_size / 1e6:.1f} MB" - else: - missing.append(path.name) - if missing: - raise RuntimeError(f"Baked datasets missing from image: {missing}") - - # Time the exact runtime load path: a cache hit loads the baked h5 in - # seconds; minutes would mean ensure_datasets fell back to a rebuild - # (i.e. the stems do not match what the runtime resolves). - start = time.monotonic() - country_module = import_module("policyengine.tax_benefit_models.us") - country_module.ensure_datasets( - datasets=[default_dataset], - years=[years[0]], - data_folder=str(data_folder), - ) - report["ensure_datasets_load_seconds"] = f"{time.monotonic() - start:.1f}" - return report - - -@app.local_entrypoint() -def main(): - report = verify_prebuilt_datasets.remote(years=PREBUILD_DATASET_YEARS) - print(json.dumps(report, indent=2)) diff --git a/projects/policyengine-simulation-executor/src/modal/smoke_app.py b/projects/policyengine-simulation-executor/src/modal/smoke_app.py index b0ecbe580..e05234009 100644 --- a/projects/policyengine-simulation-executor/src/modal/smoke_app.py +++ b/projects/policyengine-simulation-executor/src/modal/smoke_app.py @@ -2,9 +2,8 @@ Import parity, not data parity: the image here is the deployed image's layer prefix (pinned pip layer, policyengine bundle install, version -env) plus the source mounts — deliberately excluding the multi-hour -dataset-prebuild and model-snapshot layers, which add no Python -packages. Because layers are content-addressed and built through the +env) plus the source mounts — deliberately excluding the artifact-fetch +and model-snapshot layers, which add no Python packages. Because layers are content-addressed and built through the shared ``build_runtime_simulation_image()``, a warm cache makes this run take seconds; after a relock it pays only the bundle install. diff --git a/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py b/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py new file mode 100644 index 000000000..a3a8a0350 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/utils/record_deployment.py @@ -0,0 +1,67 @@ +"""Record a deployed environment's manifest in the artifact store. + +Runs on the deploy runner after the health check, writing the +``deployed/.json`` marker (last-writer-wins). The marker is +the artifact GC's liveness signal, so a deploy that cannot record itself +must fail the job even though Modal already serves it. +""" + +from __future__ import annotations + +import argparse +import os +from datetime import datetime, timezone +from typing import Mapping + +from policyengine_simulation_executor.artifact_store import ArtifactStore +from policyengine_simulation_executor.precompute_models import DeployedMarker + + +def build_marker( + environment: str, + manifest_digest: str, + env: Mapping[str, str], +) -> DeployedMarker: + """Pure marker assembly: the deploy identity that referenced this + manifest. Missing env vars become empty fields, never errors — the + marker write itself is the part that must not fail silently.""" + server = env.get("GITHUB_SERVER_URL", "") + repository = env.get("GITHUB_REPOSITORY", "") + run_id = env.get("GITHUB_RUN_ID", "") + run_url = ( + f"{server}/{repository}/actions/runs/{run_id}" + if server and repository and run_id + else "" + ) + return DeployedMarker( + environment=environment, + manifest_digest=manifest_digest, + policyengine_version=env.get("POLICYENGINE_VERSION", ""), + us_version=env.get("POLICYENGINE_US_VERSION", ""), + uk_version=env.get("POLICYENGINE_UK_VERSION", ""), + us_data_version=env.get("US_DATA_VERSION", ""), + uk_data_version=env.get("UK_DATA_VERSION", ""), + github_run_id=run_id, + github_run_url=run_url, + github_sha=env.get("GITHUB_SHA", ""), + deployed_at=datetime.now(timezone.utc).isoformat(), + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--environment", required=True) + parser.add_argument("--manifest-digest", required=True) + args = parser.parse_args() + + marker = build_marker(args.environment, args.manifest_digest, os.environ) + # The store speaks JSON-able mappings; models dump at the call site + # (same discipline as publish_manifest_impl). + ArtifactStore().write_deployed_marker(args.environment, marker.model_dump()) + print( + f"Recorded deployed/{args.environment}.json for manifest {args.manifest_digest}" + ) + + +if __name__ == "__main__": + main() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py index 56c249608..cb83156e5 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/artifact_keys.py @@ -294,7 +294,7 @@ def collect_dataset_identity(country: str, year: int) -> DatasetArtifactIdentity Reads the same sources the runtime trusts: the release bundle (versions), the bundle receipt (content sha), and the manifest certification (fingerprint). The stem comes through the same manifest helpers the - prebuild and runtime lookups use, so filename coupling is inherited, not + runtime lookup uses, so filename coupling is inherited, not re-derived. """ from policyengine.provenance.manifest import ( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py index 652189186..fd5332fd2 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/precompute_models.py @@ -1,8 +1,9 @@ """Typed schemas for the precompute pipeline's data flow. Every structure that crosses a function boundary in the precompute — the -plan, its entries, the bundle version identity, the deploy manifest, and -the worker results — is a strict Pydantic model (``extra="forbid"``), +plan, its entries, the bundle version identity, the deploy manifest, the +deployed-environment marker, and the worker results — is a strict +Pydantic model (``extra="forbid"``), matching the gateway contract's discipline. Plain dicts exist only at the Modal serialization boundary: wherever a structured value crosses, the app wrappers ``model_dump()`` on the way out of a container and @@ -112,6 +113,29 @@ def canonical_payload(self) -> dict[str, Any]: return self.model_dump(by_alias=True) +class DeployedMarker(_StrictModel): + """The ``deployed/.json`` record of what an environment runs. + + Written after a healthy deploy (last-writer-wins) by + ``src.modal.utils.record_deployment`` and read by artifact GC as its + liveness signal — the ``manifest_digest`` here is what keeps a live + artifact set from being collected. This is a persisted wire shape: + keep the keys stable. + """ + + environment: str + manifest_digest: str + policyengine_version: str + us_version: str + uk_version: str + us_data_version: str + uk_data_version: str + github_run_id: str + github_run_url: str + github_sha: str + deployed_at: str + + class DatasetBuildResult(_StrictModel): year: int path: str 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 d36366f0b..523ef82ba 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 @@ -449,13 +449,13 @@ def _load_dataset( from policyengine_simulation_executor.baseline_artifacts import uses_custom_data data_folder = resolve_data_folder() - # TEMPORARY: remove once single-year datasets are published (issue #596). - # The image bakes default-revision single-year files into - # POLICYENGINE_DATA_FOLDER, and ensure_datasets keys its cache on a - # revision-stripped filename stem — a custom dataset or revision whose - # stem matches the default would silently read the baked files. Only - # pure default requests may use the baked folder. (Region requests - # resolve their dataset without the "data" param, so they keep it.) + # The artifact fetch layer bakes default-revision single-year files + # into POLICYENGINE_DATA_FOLDER, and ensure_datasets keys its cache on + # a revision-stripped filename stem — a custom dataset or revision + # whose stem matches the default would silently read the baked files. + # Only pure default requests may use the baked folder. (Region + # requests resolve their dataset without the "data" param, so they + # keep it.) if uses_custom_data(params): data_folder = "/tmp/policyengine-data" diff --git a/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py new file mode 100644 index 000000000..2f4a71fe5 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_image_setup_fetch.py @@ -0,0 +1,278 @@ +"""Artifact-fetch layer tests against an injected fake GCS client. + +The load-bearing behaviors: every manifest artifact lands in the data +folder under its runtime filename, missing store objects fail the build +loudly before anything downloads, the freshness gate rejects manifests +computed for a different version-set, and the module stays importable +without the executor package (the layer runs before +add_local_python_source). +""" + +import ast +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +import src.modal._image_setup as image_setup +from policyengine_simulation_executor.precompute_models import ArtifactManifest +from src.modal._image_setup import fetch_artifacts + + +class FakeBlob: + def __init__(self, store, path): + self.store = store + self.path = path + + def exists(self): + return self.path in self.store.objects + + def download_to_filename(self, filename): + self.store.downloads.append(self.path) + Path(filename).write_bytes(self.store.objects[self.path]) + + +class FakeGcsClient: + def __init__(self): + self.objects = {} + self.downloads = [] + self.bucket_names = [] + + def bucket(self, name): + self.bucket_names.append(name) + return SimpleNamespace(blob=lambda path: FakeBlob(self, path)) + + +def _manifest(): + """The wire shape fetch_artifacts consumes, derived from the schema + (the layer itself cannot import the model, so it speaks the canonical + dict — building the fake through ArtifactManifest keeps it honest).""" + return ArtifactManifest.model_validate( + { + "schema": "mf1", + "country": "us", + "receipt": { + "policyengine_version": "4.22.0", + "model_version": "1.3.0", + "data_version": "1.2.3", + "data_artifact_revision": "rev-abc", + "default_dataset": "populace_cps", + }, + "artifacts": [ + { + "type": "dataset", + "path": "datasets/us/d1/populace_year_2026.h5", + "filename": "populace_year_2026.h5", + "year": 2026, + "digest": "d1", + }, + { + "type": "baseline", + "path": "baselines/us/b1/bl1-aaaa.h5", + "filename": "bl1-aaaa.h5", + "year": 2026, + "digest": "b1", + }, + ], + } + ).canonical_payload() + + +@pytest.fixture +def fake_client(): + client = FakeGcsClient() + for artifact in _manifest()["artifacts"]: + client.objects[artifact["path"]] = f"bytes-{artifact['digest']}".encode() + return client + + +@pytest.fixture +def data_folder(monkeypatch, tmp_path): + folder = tmp_path / "data" + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(folder)) + monkeypatch.setenv("POLICYENGINE_VERSION", "4.22.0") + monkeypatch.setenv("POLICYENGINE_US_VERSION", "1.3.0") + monkeypatch.delenv("POLICYENGINE_BUNDLE_RECEIPT", raising=False) + return folder + + +class TestFetchArtifacts: + def test_downloads_every_artifact_under_its_runtime_filename( + self, fake_client, data_folder + ): + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + assert "test-bucket" in fake_client.bucket_names + for artifact in _manifest()["artifacts"]: + target = data_folder / artifact["filename"] + assert target.read_bytes() == f"bytes-{artifact['digest']}".encode() + + def test_missing_store_object_fails_before_any_download( + self, fake_client, data_folder + ): + missing_path = _manifest()["artifacts"][1]["path"] + del fake_client.objects[missing_path] + with pytest.raises(RuntimeError, match=missing_path): + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + assert fake_client.downloads == [] + + def test_unresolved_manifest_fails_with_the_digest_env_var_named( + self, fake_client, data_folder + ): + with pytest.raises(RuntimeError, match="POLICYENGINE_MANIFEST_DIGEST"): + fetch_artifacts("test-bucket", None, client=fake_client) + assert fake_client.downloads == [] + + def test_overwrites_a_preexisting_same_named_file(self, fake_client, data_folder): + """No skip-if-exists: stale same-named bytes must never survive + into the image.""" + data_folder.mkdir(parents=True) + stale = data_folder / "populace_year_2026.h5" + stale.write_bytes(b"stale") + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + assert stale.read_bytes() == b"bytes-d1" + + +class TestFreshnessGate: + def test_rejects_a_policyengine_version_mismatch( + self, monkeypatch, fake_client, data_folder + ): + monkeypatch.setenv("POLICYENGINE_VERSION", "4.99.0") + with pytest.raises(RuntimeError, match=r"4\.99\.0.*4\.22\.0"): + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + assert fake_client.downloads == [] + + def test_rejects_a_model_version_mismatch( + self, monkeypatch, fake_client, data_folder + ): + monkeypatch.setenv("POLICYENGINE_US_VERSION", "9.9.9") + with pytest.raises(RuntimeError, match="POLICYENGINE_US_VERSION"): + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + + def test_rejects_an_installed_data_version_mismatch( + self, monkeypatch, tmp_path, fake_client, data_folder + ): + receipt_file = tmp_path / "bundle-receipt.json" + receipt_file.write_text('{"datasets": [{"country": "us", "version": "9.9.9"}]}') + monkeypatch.setenv("POLICYENGINE_BUNDLE_RECEIPT", str(receipt_file)) + with pytest.raises(RuntimeError, match=r"9\.9\.9.*1\.2\.3"): + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + + def test_passes_when_the_installed_data_version_matches( + self, monkeypatch, tmp_path, fake_client, data_folder + ): + receipt_file = tmp_path / "bundle-receipt.json" + receipt_file.write_text('{"datasets": [{"country": "us", "version": "1.2.3"}]}') + monkeypatch.setenv("POLICYENGINE_BUNDLE_RECEIPT", str(receipt_file)) + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + assert len(fake_client.downloads) == 2 + + def test_proceeds_without_an_installed_receipt_file( + self, monkeypatch, tmp_path, fake_client, data_folder + ): + """Lenient by design: local fault-drill runs have no bundle + receipt file.""" + monkeypatch.setenv( + "POLICYENGINE_BUNDLE_RECEIPT", str(tmp_path / "does-not-exist.json") + ) + fetch_artifacts("test-bucket", _manifest(), client=fake_client) + assert len(fake_client.downloads) == 2 + + +class TestCredentialShim: + """The client=None path: credentials are built in memory from the + secret blob via service_account.Credentials.from_service_account_info + (never written to disk — the layer's filesystem is committed into the + image), falling back to ambient ADC when no blob is present.""" + + @pytest.fixture + def patched_google(self, monkeypatch, fake_client): + from google.cloud import storage + from google.oauth2 import service_account + + recorder = SimpleNamespace(infos=[], client_calls=[], sentinel=object()) + + def fake_from_info(info): + recorder.infos.append(info) + return recorder.sentinel + + def fake_client_factory(**kwargs): + recorder.client_calls.append(kwargs) + return fake_client + + monkeypatch.setattr( + service_account.Credentials, "from_service_account_info", fake_from_info + ) + monkeypatch.setattr(storage, "Client", fake_client_factory) + return recorder + + def test_builds_in_memory_credentials_from_the_secret_blob( + self, monkeypatch, fake_client, data_folder, patched_google + ): + monkeypatch.setenv( + "GOOGLE_APPLICATION_CREDENTIALS_JSON", + '{"type": "service_account", "project_id": "pe-test"}', + ) + + fetch_artifacts("test-bucket", _manifest()) + + assert patched_google.infos == [ + {"type": "service_account", "project_id": "pe-test"} + ] + assert patched_google.client_calls == [ + {"credentials": patched_google.sentinel, "project": "pe-test"} + ] + assert len(fake_client.downloads) == 2 + + def test_unwraps_the_double_encoded_blob( + self, monkeypatch, fake_client, data_folder, patched_google + ): + clean = json.dumps({"type": "service_account", "project_id": "pe-test"}) + monkeypatch.setenv( + "GOOGLE_APPLICATION_CREDENTIALS_JSON", clean.replace('"', '\\"') + ) + + fetch_artifacts("test-bucket", _manifest()) + + assert patched_google.infos == [ + {"type": "service_account", "project_id": "pe-test"} + ] + + def test_falls_back_to_ambient_adc_without_a_blob( + self, monkeypatch, fake_client, data_folder, patched_google + ): + monkeypatch.delenv("GOOGLE_APPLICATION_CREDENTIALS_JSON", raising=False) + + fetch_artifacts("test-bucket", _manifest()) + + assert patched_google.infos == [] + assert patched_google.client_calls == [{}] + + +def test_fetch_never_writes_credentials_to_disk(): + """A run_function layer commits its container filesystem into the + image, so the shim must never materialize the key as a file.""" + source = Path(image_setup.__file__).read_text(encoding="utf-8") + assert "tempfile" not in source + + +def test_image_setup_module_stays_self_contained(): + """The fetch layer runs before add_local_python_source, so the module + must have no module-level imports and must never import the executor + package, even lazily inside a function body.""" + source = Path(image_setup.__file__).read_text(encoding="utf-8") + tree = ast.parse(source) + + module_level_imports = [ + node for node in tree.body if isinstance(node, (ast.Import, ast.ImportFrom)) + ] + assert module_level_imports == [] + + for node in ast.walk(tree): + if isinstance(node, ast.ImportFrom): + assert not (node.module or "").startswith( + "policyengine_simulation_executor" + ) + if isinstance(node, ast.Import): + for alias in node.names: + assert not alias.name.startswith("policyengine_simulation_executor") diff --git a/projects/policyengine-simulation-executor/tests/test_image_setup_prebuild.py b/projects/policyengine-simulation-executor/tests/test_image_setup_prebuild.py deleted file mode 100644 index 9f49ff3d3..000000000 --- a/projects/policyengine-simulation-executor/tests/test_image_setup_prebuild.py +++ /dev/null @@ -1,141 +0,0 @@ -"""Tests for the TEMPORARY single-year dataset prebuild (see issue #596). - -These pin the coupling between the image-build prebuild and the runtime -dataset lookup: both must resolve the same ``{stem}_year_{year}.h5`` cache -paths, otherwise the baked files are silently ignored and every cold -container falls back to the slow runtime build. -""" - -import sys -from pathlib import Path -from types import SimpleNamespace - -import pytest -from policyengine.provenance.manifest import ( - dataset_logical_name, - get_release_manifest, - resolve_dataset_reference, -) - -from policyengine_simulation_executor.release_bundle import ( - get_country_release_bundle, - resolve_bundle_dataset_name, -) -from policyengine_simulation_executor.simulation_runtime import DEFAULT_YEAR -from src.modal._image_setup import ( - PREBUILD_DATASET_YEARS, - prebuild_country_datasets, -) - - -def _expected_stem(country: str) -> str: - default_dataset = get_release_manifest(country).default_dataset - return dataset_logical_name(resolve_dataset_reference(country, default_dataset)) - - -def _install_country_stub(monkeypatch, country: str, ensure_datasets): - monkeypatch.setitem( - sys.modules, - f"policyengine.tax_benefit_models.{country}", - SimpleNamespace(ensure_datasets=ensure_datasets), - ) - - -def test_prebuild_stem_matches_runtime_default_lookup(monkeypatch): - """The critical coupling pin: the stem the prebuild writes must equal the - stem the runtime resolves for default requests. If bundle metadata ever - diverges from the packaged manifest default, the baked files would be - silently unused — this test turns that into a CI failure on every - policyengine version bump.""" - monkeypatch.delenv("POLICYENGINE_BUNDLE_RECEIPT", raising=False) - get_country_release_bundle.cache_clear() - - runtime_stem = dataset_logical_name( - resolve_dataset_reference("us", resolve_bundle_dataset_name("us", None)) - ) - - assert _expected_stem("us") == runtime_stem - - -def test_prebuild_years_cover_runtime_default_year(): - assert PREBUILD_DATASET_YEARS == [2025, 2026, 2027] - assert DEFAULT_YEAR in PREBUILD_DATASET_YEARS - - -def test_prebuild_passes_explicit_manifest_default(tmp_path, monkeypatch): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - stem = _expected_stem("us") - calls = [] - - def ensure_datasets(**kwargs): - calls.append(kwargs) - for year in kwargs["years"]: - Path(kwargs["data_folder"], f"{stem}_year_{year}.h5").write_bytes(b"h5") - return {} - - _install_country_stub(monkeypatch, "us", ensure_datasets) - - prebuild_country_datasets("us") - - assert calls == [ - { - "datasets": [get_release_manifest("us").default_dataset], - "years": [2025, 2026, 2027], - "data_folder": str(tmp_path), - } - ] - - -def test_prebuild_skips_years_that_already_exist(tmp_path, monkeypatch): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - stem = _expected_stem("us") - (tmp_path / f"{stem}_year_2025.h5").write_bytes(b"h5") - calls = [] - - def ensure_datasets(**kwargs): - calls.append(kwargs) - for year in kwargs["years"]: - Path(kwargs["data_folder"], f"{stem}_year_{year}.h5").write_bytes(b"h5") - return {} - - _install_country_stub(monkeypatch, "us", ensure_datasets) - - prebuild_country_datasets("us") - - assert calls == [ - { - "datasets": [get_release_manifest("us").default_dataset], - "years": [2026, 2027], - "data_folder": str(tmp_path), - } - ] - - -def test_prebuild_skips_entirely_when_all_years_exist(tmp_path, monkeypatch): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - stem = _expected_stem("us") - for year in PREBUILD_DATASET_YEARS: - (tmp_path / f"{stem}_year_{year}.h5").write_bytes(b"h5") - calls = [] - - def ensure_datasets(**kwargs): - calls.append(kwargs) - return {} - - _install_country_stub(monkeypatch, "us", ensure_datasets) - - prebuild_country_datasets("us") - - assert calls == [] - - -def test_prebuild_raises_when_files_not_produced(tmp_path, monkeypatch): - monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) - - def ensure_datasets(**kwargs): - return {} - - _install_country_stub(monkeypatch, "us", ensure_datasets) - - with pytest.raises(RuntimeError, match="did not produce"): - prebuild_country_datasets("us") diff --git a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py index de667d31d..63e62c2d7 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_bundle_image.py @@ -76,35 +76,84 @@ def test_modal_image_uses_policyengine_bundle_install(monkeypatch): ] -# TEMPORARY: remove once single-year datasets are published (issue #596). -def test_modal_image_prebuilds_datasets_between_env_and_local_source(monkeypatch): +def _fake_manifest(): + """A schema-valid store payload: app.py validates what it reads, so + the fake must satisfy ArtifactManifest — deriving it from the model + keeps the test from drifting off the wire shape.""" + from policyengine_simulation_executor.precompute_models import ArtifactManifest + + return ArtifactManifest.model_validate( + { + "schema": "mf1", + "country": "us", + "receipt": { + "policyengine_version": "4.19.1", + "model_version": "1.700.0", + "data_version": "1.2.3", + "data_artifact_revision": "rev-abc", + "default_dataset": "populace_cps", + }, + "artifacts": [ + { + "type": "dataset", + "path": "datasets/us/d1/populace_year_2026.h5", + "filename": "populace_year_2026.h5", + "year": 2026, + "digest": "d1", + }, + ], + } + ).canonical_payload() + + +def test_modal_image_fetches_artifacts_between_env_and_local_source(monkeypatch): install_fake_modal(monkeypatch) monkeypatch.setenv("POLICYENGINE_VERSION", "4.19.1") monkeypatch.setenv("POLICYENGINE_CORE_VERSION", "3.27.1") monkeypatch.setenv("POLICYENGINE_US_VERSION", "1.700.0") monkeypatch.setenv("POLICYENGINE_UK_VERSION", "2.90.0") + monkeypatch.setenv("POLICYENGINE_MANIFEST_DIGEST", "digest-abc") + monkeypatch.setenv("POLICYENGINE_ARTIFACT_BUCKET", "test-bucket") + + manifest = _fake_manifest() + + class FakeStore: + def __init__(self, bucket_name=None, **kwargs): + self.bucket_name = bucket_name + + def read_manifest(self, digest): + assert digest == "digest-abc" + return manifest + + # The app imports ArtifactStore lazily inside the digest-set branch, + # so patching the module attribute is enough — no GCS involved. + import policyengine_simulation_executor.artifact_store as artifact_store + + monkeypatch.setattr(artifact_store, "ArtifactStore", FakeStore) sys.modules.pop("src.modal.app", None) app = importlib.import_module("src.modal.app") calls = app.simulation_image.calls - prebuild_indices = [ + fetch_indices = [ index for index, call in enumerate(calls) - if call[0] == "run_function" and call[1] == "prebuild_country_datasets" + if call[0] == "run_function" and call[1] == "fetch_artifacts" ] - # US only — UK is deliberately not prebuilt (keeps image build short). - assert [calls[index][2]["args"] for index in prebuild_indices] == [("us",)] - prebuild_kwargs = calls[prebuild_indices[0]][2] - assert prebuild_kwargs["secrets"] == [app.data_secret, app.hf_secret] - assert prebuild_kwargs["timeout"] == 4 * 60 * 60 - assert prebuild_kwargs["memory"] == 65536 - - # The prebuild layer is a multi-hour build keyed only on upstream - # layers and its own definition. It must stay after the version env - # (so version bumps rebuild it) and before add_local_python_source - # (whose content-hash key would otherwise invalidate it on every - # source commit). + assert len(fetch_indices) == 1 + fetch_kwargs = calls[fetch_indices[0]][2] + # The manifest content is in the layer args: the layer cache key + # follows the artifact set, nothing else. + assert fetch_kwargs["args"] == ("test-bucket", manifest) + assert fetch_kwargs["secrets"] == [app.gcp_secret] + assert fetch_kwargs["cpu"] == 2.0 + assert fetch_kwargs["memory"] == 4096 + assert fetch_kwargs["timeout"] == 900 + + # The fetch layer must stay after the version env (its freshness gate + # reads the baked versions) and before add_local_python_source (whose + # content-hash key would otherwise invalidate it on every source + # commit). env_index = next(index for index, call in enumerate(calls) if call[0] == "env") local_source_index = next( index @@ -116,7 +165,7 @@ def test_modal_image_prebuilds_datasets_between_env_and_local_source(monkeypatch for index, call in enumerate(calls) if call[0] == "run_function" and call[1] == "snapshot_models" ) - assert env_index < prebuild_indices[0] < local_source_index < snapshot_index + assert env_index < fetch_indices[0] < local_source_index < snapshot_index # The shared libs ship into the image as mounted source; dropping one # from this tuple crashes workers at import time. @@ -128,6 +177,32 @@ def test_modal_image_prebuilds_datasets_between_env_and_local_source(monkeypatch ) +def test_modal_image_uses_failing_sentinel_without_manifest_digest(monkeypatch): + """Without a digest the app must still import — the precompute and + smoke apps import it where no digest can exist — but the fetch layer + args must carry a None manifest, which makes fetch_artifacts fail the + build loudly. A digest-less deploy cannot silently ship an + artifact-less image.""" + install_fake_modal(monkeypatch) + monkeypatch.setenv("POLICYENGINE_VERSION", "4.19.1") + monkeypatch.setenv("POLICYENGINE_CORE_VERSION", "3.27.1") + monkeypatch.setenv("POLICYENGINE_US_VERSION", "1.700.0") + monkeypatch.setenv("POLICYENGINE_UK_VERSION", "2.90.0") + monkeypatch.delenv("POLICYENGINE_MANIFEST_DIGEST", raising=False) + monkeypatch.delenv("POLICYENGINE_ARTIFACT_BUCKET", raising=False) + sys.modules.pop("src.modal.app", None) + + app = importlib.import_module("src.modal.app") + + fetch_calls = [ + call + for call in app.simulation_image.calls + if call[0] == "run_function" and call[1] == "fetch_artifacts" + ] + assert len(fetch_calls) == 1 + assert fetch_calls[0][2]["args"] == ("", None) + + def test_app_module_imports_at_container_entrypoint_path(monkeypatch): """Modal loads the deployed function's module as /root/app.py. @@ -143,6 +218,10 @@ def test_app_module_imports_at_container_entrypoint_path(monkeypatch): monkeypatch.setenv("POLICYENGINE_CORE_VERSION", "3.27.1") monkeypatch.setenv("POLICYENGINE_US_VERSION", "1.700.0") monkeypatch.setenv("POLICYENGINE_UK_VERSION", "2.90.0") + # Containers must stay inert even if a digest leaks into their env: + # the is_local() check comes before any env read or GCS work, so this + # value must be ignored entirely. + monkeypatch.setenv("POLICYENGINE_MANIFEST_DIGEST", "digest-must-be-ignored") sys.modules.pop("src.modal.app", None) source_path = Path(__file__).resolve().parents[1] / "src" / "modal" / "app.py" @@ -155,3 +234,9 @@ def test_app_module_imports_at_container_entrypoint_path(monkeypatch): exec(code, module.__dict__) assert module.APP_NAME.startswith("policyengine-simulation-py") + fetch_calls = [ + call + for call in module.simulation_image.calls + if call[0] == "run_function" and call[1] == "fetch_artifacts" + ] + assert fetch_calls[0][2]["args"] == ("", None) diff --git a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py index a81d4024b..ee99fc338 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py @@ -614,6 +614,38 @@ def test_deploy_workflow_threads_force_recompute_to_script(self): '"${{ inputs.force_recompute }}"' in reusable_workflow ) + def test_deploy_workflow_threads_manifest_and_store_credentials_to_deploy(self): + """The deploy step resolves the manifest from GCS on the runner, + so it needs the digest, the store bucket, and GCP credentials.""" + reusable_workflow = ( + REPO_ROOT / ".github" / "workflows" / "modal-deploy.reusable.yml" + ).read_text(encoding="utf-8") + + deploy_step = reusable_workflow[ + reusable_workflow.index( + "Deploy simulation API to Modal" + ) : reusable_workflow.index("Get deployed URL") + ] + assert ( + "POLICYENGINE_MANIFEST_DIGEST: " + "${{ needs.precompute.outputs.manifest_digest }}" in deploy_step + ) + assert ( + "POLICYENGINE_ARTIFACT_BUCKET: ${{ vars.POLICYENGINE_ARTIFACT_BUCKET }}" + in deploy_step + ) + assert "GCP_CREDENTIALS_JSON: ${{ secrets.GCP_CREDENTIALS_JSON }}" in ( + deploy_step + ) + # The precompute, deploy, and record-marker steps each carry the + # bucket var. + assert ( + reusable_workflow.count( + "POLICYENGINE_ARTIFACT_BUCKET: ${{ vars.POLICYENGINE_ARTIFACT_BUCKET }}" + ) + == 3 + ) + def test_precompute_job_syncs_its_own_secrets(self): """Precompute must not depend on a prior deploy's secret sync.""" reusable_workflow = ( @@ -629,6 +661,123 @@ def test_precompute_job_syncs_its_own_secrets(self): ) +class TestModalRecordDeployment: + """Tests for modal-record-deployment.sh""" + + script = SCRIPTS_DIR / "modal-record-deployment.sh" + + def _run_with_fake_uv(self, tmp_path, *args, bucket="policyengine-sim-artifacts"): + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + log_path = tmp_path / "uv-calls.log" + uv_path = bin_dir / "uv" + uv_path.write_text( + '#!/bin/bash\nprintf \'%s\\n\' "$*" >> "$UV_FAKE_LOG"\n', + encoding="utf-8", + ) + uv_path.chmod(0o755) + + env = os.environ.copy() + env.update( + { + "PATH": f"{bin_dir}{os.pathsep}{env['PATH']}", + "UV_FAKE_LOG": str(log_path), + } + ) + if bucket is None: + env.pop("POLICYENGINE_ARTIFACT_BUCKET", None) + else: + env["POLICYENGINE_ARTIFACT_BUCKET"] = bucket + + result = subprocess.run( + ["bash", str(self.script), *args], + capture_output=True, + text=True, + env=env, + ) + calls = ( + log_path.read_text(encoding="utf-8").splitlines() + if log_path.exists() + else [] + ) + return result, calls + + def test_script_exists(self): + """Script file should exist.""" + assert self.script.exists(), f"Script not found at {self.script}" + + def test_requires_environment_argument(self): + """Should fail with no arguments.""" + result = subprocess.run( + ["bash", str(self.script)], + capture_output=True, + text=True, + ) + assert result.returncode != 0 + + def test_requires_manifest_digest_argument(self, tmp_path): + """Should fail before invoking anything without a digest.""" + result, calls = self._run_with_fake_uv(tmp_path, "beta") + + assert result.returncode != 0 + assert calls == [] + + def test_requires_artifact_bucket(self, tmp_path): + """Should fail before invoking anything without the bucket env.""" + result, calls = self._run_with_fake_uv( + tmp_path, "beta", "digest-abc", bucket=None + ) + + assert result.returncode != 0 + assert "POLICYENGINE_ARTIFACT_BUCKET is required" in result.stderr + assert calls == [] + + def test_invokes_the_recorder_module(self, tmp_path): + """The exact python -m invocation, environment and digest threaded.""" + result, calls = self._run_with_fake_uv(tmp_path, "beta", "digest-abc") + + assert result.returncode == 0, result.stderr + assert calls == [ + "run python -m src.modal.utils.record_deployment " + "--environment beta --manifest-digest digest-abc" + ] + + def test_deploy_workflow_records_marker_after_health_check(self): + """Both legs write deployed/.json once the deploy is healthy.""" + reusable_workflow = ( + REPO_ROOT / ".github" / "workflows" / "modal-deploy.reusable.yml" + ).read_text(encoding="utf-8") + + invocation = ( + 'modal-record-deployment.sh "${{ inputs.environment }}" ' + '"${{ needs.precompute.outputs.manifest_digest }}"' + ) + assert invocation in reusable_workflow + # The marker means deployed AND healthy: it must come after the + # health check step. + assert reusable_workflow.index("Verify deployment health") < ( + reusable_workflow.index("modal-record-deployment.sh") + ) + + # The payload builder maps missing env vars to empty strings, so a + # silent rename here would empty the marker's fields — assert every + # env line the step feeds it. + marker_step = reusable_workflow[ + reusable_workflow.index("Record deployment marker") : + ] + marker_step = marker_step[: marker_step.index("integ_test:")] + for env_line in ( + "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 }}", + ): + assert env_line in marker_step, f"marker step is missing: {env_line}" + + class TestModalGetUrl: """Tests for modal-get-url.sh""" diff --git a/projects/policyengine-simulation-executor/tests/test_precompute.py b/projects/policyengine-simulation-executor/tests/test_precompute.py index bbf0628bf..172a3a57d 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute.py @@ -125,6 +125,14 @@ def test_force_selects_everything(self): def test_years_are_in_priority_order(self): assert precompute.PRECOMPUTE_YEARS == [2026, 2027, 2025] + def test_precompute_years_cover_runtime_default_year(self): + """The image bakes exactly the precomputed years, so the runtime + default year must be one of them or every default request pays a + runtime dataset build.""" + from policyengine_simulation_executor.simulation_runtime import DEFAULT_YEAR + + assert DEFAULT_YEAR in precompute.PRECOMPUTE_YEARS + def test_manifest_lists_every_artifact_with_runtime_filenames(self): manifest = precompute.build_manifest(_plan()) assert manifest.manifest_schema == precompute.MANIFEST_SCHEMA @@ -678,8 +686,8 @@ def upload_file(self, path, local_path): return True fake_analysis = ModuleType("policyengine.tax_benefit_models.us.analysis") - fake_analysis.configure_budgetary_impact_variables = ( - lambda baseline, reform: state.configured.append((baseline, reform)) + fake_analysis.configure_budgetary_impact_variables = lambda baseline, reform: ( + state.configured.append((baseline, reform)) ) monkeypatch.setitem( sys.modules, "policyengine.tax_benefit_models.us.analysis", fake_analysis diff --git a/projects/policyengine-simulation-executor/tests/test_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_precompute_app.py index cc6d21456..87f581de4 100644 --- a/projects/policyengine-simulation-executor/tests/test_precompute_app.py +++ b/projects/policyengine-simulation-executor/tests/test_precompute_app.py @@ -35,6 +35,9 @@ def precompute_module(monkeypatch): "POLICYENGINE_UK_VERSION", ): monkeypatch.setenv(env, "0.0.0-test") + # A developer following the local-deploy instructions may have the + # digest exported; importing src.modal.app must not reach for GCS. + monkeypatch.delenv("POLICYENGINE_MANIFEST_DIGEST", raising=False) for module in ("src.modal.app", "src.modal.precompute_app"): sys.modules.pop(module, None) module = importlib.import_module("src.modal.precompute_app") @@ -69,7 +72,7 @@ def test_compute_baseline_matches_segment_worker_shape(precompute_module): assert kwargs["timeout"] == 3600 -def test_dataset_builder_gets_prebuild_scale_resources(precompute_module): +def test_dataset_builder_gets_dataset_build_scale_resources(precompute_module): kwargs = _function_kwargs(precompute_module, "build_dataset") assert kwargs["cpu"] == 8.0 assert kwargs["memory"] == 65536 diff --git a/projects/policyengine-simulation-executor/tests/test_record_deployment.py b/projects/policyengine-simulation-executor/tests/test_record_deployment.py new file mode 100644 index 000000000..767e92dc9 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_record_deployment.py @@ -0,0 +1,96 @@ +"""Deployment-marker recording: typed marker assembly plus the store write. + +The deployed/.json marker is the artifact GC's liveness signal, so +main() writes through the real store client seam (faked here) and any +write failure propagates — a deploy that cannot record itself goes red. +""" + +from datetime import datetime + +import pytest + +from policyengine_simulation_executor.precompute_models import DeployedMarker +from src.modal.utils import record_deployment + + +def _github_env(): + return { + "POLICYENGINE_VERSION": "4.22.0", + "POLICYENGINE_US_VERSION": "1.3.0", + "POLICYENGINE_UK_VERSION": "2.9.0", + "US_DATA_VERSION": "1.2.3", + "UK_DATA_VERSION": "3.4.5", + "GITHUB_SERVER_URL": "https://github.com", + "GITHUB_REPOSITORY": "PolicyEngine/policyengine-sim-api", + "GITHUB_RUN_ID": "12345", + "GITHUB_SHA": "abc123", + } + + +class TestBuildMarker: + def test_records_the_deploy_identity(self): + marker = record_deployment.build_marker("beta", "digest-1", _github_env()) + + assert marker.environment == "beta" + assert marker.manifest_digest == "digest-1" + assert marker.policyengine_version == "4.22.0" + assert marker.us_version == "1.3.0" + assert marker.uk_version == "2.9.0" + assert marker.us_data_version == "1.2.3" + assert marker.uk_data_version == "3.4.5" + assert marker.github_run_id == "12345" + assert marker.github_run_url == ( + "https://github.com/PolicyEngine/policyengine-sim-api/actions/runs/12345" + ) + assert marker.github_sha == "abc123" + + def test_timestamp_is_utc_iso8601(self): + marker = record_deployment.build_marker("beta", "d", {}) + parsed = datetime.fromisoformat(marker.deployed_at) + assert parsed.utcoffset() is not None + assert parsed.utcoffset().total_seconds() == 0 + + def test_missing_env_yields_empty_fields_not_errors(self): + marker = record_deployment.build_marker("prod", "d", {}) + assert marker.github_run_url == "" + assert marker.policyengine_version == "" + + +class TestMain: + def test_writes_the_marker_through_the_store(self, monkeypatch): + writes = [] + + class FakeStore: + def __init__(self, bucket_name=None, **kwargs): + pass + + def write_deployed_marker(self, environment, payload): + writes.append((environment, payload)) + + monkeypatch.setattr(record_deployment, "ArtifactStore", FakeStore) + monkeypatch.setattr( + "sys.argv", + [ + "record_deployment", + "--environment", + "beta", + "--manifest-digest", + "digest-9", + ], + ) + + record_deployment.main() + + assert len(writes) == 1 + environment, payload = writes[0] + assert environment == "beta" + # The persisted payload is exactly the model's dump — the wire + # shape GC will read back. + assert set(payload) == set(DeployedMarker.model_fields) + assert payload["manifest_digest"] == "digest-9" + assert DeployedMarker.model_validate(payload).environment == "beta" + + def test_requires_both_arguments(self, monkeypatch): + monkeypatch.setattr("sys.argv", ["record_deployment", "--environment", "beta"]) + with pytest.raises(SystemExit): + record_deployment.main() 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 786f0b373..58170baf3 100644 --- a/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py +++ b/projects/policyengine-simulation-executor/tests/test_simulation_output_builder.py @@ -1145,8 +1145,7 @@ def ensure_datasets(**kwargs): ] -# TEMPORARY: remove once single-year datasets are published (issue #596). -# Pins the guard that keeps every non-default dataset request away from the +# Locks the guard that keeps every non-default dataset request away from the # baked default-revision single-year files in POLICYENGINE_DATA_FOLDER. # ensure_datasets keys its cache on a revision-stripped filename stem, so # even an explicit dataset name or foreign URI can collide with the baked