diff --git a/.github/scripts/stage12-cache-precompute.sh b/.github/scripts/stage12-cache-precompute.sh new file mode 100755 index 000000000..2b94ad263 --- /dev/null +++ b/.github/scripts/stage12-cache-precompute.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# Build the independent Stage 12 cache and export its manifest digest. +# Cache files are immutable and addressed by identity, so there is no force +# option: a run computes only missing files. + +set -euo pipefail + +modal_environment="${1:?Modal environment is required}" +: "${STAGE12_CACHE_BUCKET:?STAGE12_CACHE_BUCKET is required}" +: "${GITHUB_OUTPUT:?GITHUB_OUTPUT is required}" + +output_file="$(mktemp)" +trap 'rm -f "${output_file}"' EXIT +uv run modal run --env="${modal_environment}" src/modal/stage12_precompute_app.py 2>&1 | tee "${output_file}" +digest_line="$(tr -d '\r' <"${output_file}" | grep '^STAGE12_CACHE_MANIFEST_DIGEST=' | tail -n 1 || true)" +if [[ -z "${digest_line}" ]]; then + echo "Stage 12 precompute returned no cache manifest digest" >&2 + exit 1 +fi +digest="${digest_line#STAGE12_CACHE_MANIFEST_DIGEST=}" +if [[ ! "${digest}" =~ ^[0-9a-f]{64}$ ]]; then + echo "Stage 12 precompute returned an invalid cache manifest digest" >&2 + exit 1 +fi +echo "manifest_digest=${digest}" >>"${GITHUB_OUTPUT}" diff --git a/.github/scripts/stage12-sync-modal-secrets.sh b/.github/scripts/stage12-sync-modal-secrets.sh index c6c97b180..5954fd68a 100755 --- a/.github/scripts/stage12-sync-modal-secrets.sh +++ b/.github/scripts/stage12-sync-modal-secrets.sh @@ -9,6 +9,7 @@ modal_environment="${1:?Modal environment is required}" : "${STAGE12_DATABASE_URL_SECRET_NAME:?STAGE12_DATABASE_URL_SECRET_NAME is required}" : "${STAGE12_GCP_CREDENTIALS_SECRET_NAME:?STAGE12_GCP_CREDENTIALS_SECRET_NAME is required}" : "${STAGE12_ARTIFACT_BUCKET:?STAGE12_ARTIFACT_BUCKET is required}" +: "${STAGE12_CACHE_BUCKET:?STAGE12_CACHE_BUCKET is required}" : "${MODAL_TOKEN_ID:?MODAL_TOKEN_ID is required}" : "${MODAL_TOKEN_SECRET:?MODAL_TOKEN_SECRET is required}" @@ -34,7 +35,8 @@ gcloud secrets versions access latest \ jq -n \ --rawfile database_url "${database_url_file}" \ --arg artifact_bucket "${STAGE12_ARTIFACT_BUCKET}" \ - '{STAGE12_DATABASE_URL: ($database_url | rtrimstr("\n")), STAGE12_ARTIFACT_BUCKET: $artifact_bucket}' \ + --arg cache_bucket "${STAGE12_CACHE_BUCKET}" \ + '{STAGE12_DATABASE_URL: ($database_url | rtrimstr("\n")), STAGE12_ARTIFACT_BUCKET: $artifact_bucket, STAGE12_CACHE_BUCKET: $cache_bucket}' \ >"${runtime_modal_file}" jq -n \ --rawfile credentials "${credentials_file}" \ diff --git a/.github/scripts/stage12-validate-infrastructure.sh b/.github/scripts/stage12-validate-infrastructure.sh index d12daabad..c64413204 100755 --- a/.github/scripts/stage12-validate-infrastructure.sh +++ b/.github/scripts/stage12-validate-infrastructure.sh @@ -7,6 +7,7 @@ set +x : "${STAGE12_ENVIRONMENT:?STAGE12_ENVIRONMENT is required}" : "${STAGE12_GCP_PROJECT_ID:?STAGE12_GCP_PROJECT_ID is required}" : "${STAGE12_ARTIFACT_BUCKET:?STAGE12_ARTIFACT_BUCKET is required}" +: "${STAGE12_CACHE_BUCKET:?STAGE12_CACHE_BUCKET is required}" : "${STAGE12_ENTRYPOINT_SERVICE_ACCOUNT:?STAGE12_ENTRYPOINT_SERVICE_ACCOUNT is required}" : "${STAGE12_DEPLOY_SERVICE_ACCOUNT:?STAGE12_DEPLOY_SERVICE_ACCOUNT is required}" : "${STAGE12_MODAL_SERVICE_ACCOUNT:?STAGE12_MODAL_SERVICE_ACCOUNT is required}" @@ -28,6 +29,11 @@ if [[ ! "${STAGE12_ARTIFACT_BUCKET}" =~ ^[a-z0-9][a-z0-9._-]{2,61}[a-z0-9]$ ]] | echo "STAGE12_ARTIFACT_BUCKET must be valid and identify the environment" >&2 exit 1 fi +if [[ ! "${STAGE12_CACHE_BUCKET}" =~ ^[a-z0-9][a-z0-9._-]{2,61}[a-z0-9]$ ]] || + [[ ! "${STAGE12_CACHE_BUCKET}" =~ ${STAGE12_ENVIRONMENT} ]]; then + echo "STAGE12_CACHE_BUCKET must be valid and identify the environment" >&2 + exit 1 +fi for account in \ "${STAGE12_ENTRYPOINT_SERVICE_ACCOUNT}" \ "${STAGE12_DEPLOY_SERVICE_ACCOUNT}" \ @@ -68,6 +74,7 @@ canary_source_file="$(mktemp)" canary_download_file="$(mktemp)" runtime_gcloud_config="$(mktemp -d)" canary_object="" +cache_canary_object="" runtime_gcloud() { # The GitHub auth action exports credentials for the deployment identity. # Remove them only here so the activated Modal worker account controls this check. @@ -83,6 +90,10 @@ cleanup() { runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ "${canary_object}" --quiet >/dev/null 2>&1 || true fi + if [[ -n "${cache_canary_object}" ]]; then + runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ + "${cache_canary_object}" --quiet >/dev/null 2>&1 || true + fi rm -f \ "${database_url_file}" \ "${credentials_file}" \ @@ -164,6 +175,19 @@ runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ "${canary_object}" --quiet canary_object="" +cache_canary_object="gs://${STAGE12_CACHE_BUCKET}/_deployment-validation/${STAGE12_ENVIRONMENT}/${canary_id}.txt" +runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage cp \ + "${canary_source_file}" "${cache_canary_object}" --quiet +runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage cp \ + "${cache_canary_object}" "${canary_download_file}" --quiet +cmp "${canary_source_file}" "${canary_download_file}" +gcloud storage cp \ + "${cache_canary_object}" "${canary_download_file}" --quiet +cmp "${canary_source_file}" "${canary_download_file}" +runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage rm \ + "${cache_canary_object}" --quiet +cache_canary_object="" + for lookup_object in "${uk_geography_lookup_objects[@]}"; do runtime_gcloud --account="${STAGE12_MODAL_SERVICE_ACCOUNT}" storage cp \ "${lookup_object}" "${canary_download_file}" --quiet @@ -173,4 +197,4 @@ for lookup_object in "${uk_geography_lookup_objects[@]}"; do fi done -echo "Pre-provisioned Stage 12 database, secret, artifact storage, and UK geography lookup access is verified." +echo "Pre-provisioned Stage 12 database, secret, report storage, cache storage, and UK geography lookup access is verified." diff --git a/.github/scripts/stage12-validate-v2-workers.sh b/.github/scripts/stage12-validate-v2-workers.sh index ad0ed6e7b..df813a8be 100755 --- a/.github/scripts/stage12-validate-v2-workers.sh +++ b/.github/scripts/stage12-validate-v2-workers.sh @@ -1,7 +1,8 @@ #!/usr/bin/env bash -# Validate the installed bundles, secrets, datasets and calculations of the -# deployed v2 workers, and keep the result for manifest publication. -# Usage: stage12-validate-v2-workers.sh +# Validate the installed bundles, secrets, datasets, calculations and Stage 12 +# cache of the deployed v2 workers, and keep the result for manifest +# publication. +# Usage: stage12-validate-v2-workers.sh # Run from projects/policyengine-simulation-executor. # Required env vars: RUNNER_TEMP, GITHUB_OUTPUT. # Outputs: file (the validation result path) in GITHUB_OUTPUT. @@ -9,9 +10,11 @@ set -euo pipefail MODAL_ENV="${1:?Modal environment required}" +CACHE_MANIFEST_SHA256="${2:?Stage 12 cache manifest digest required}" output="${RUNNER_TEMP}/stage12-worker-validation.json" uv run python -m src.modal.utils.validate_v2_worker_release \ --environment "${MODAL_ENV}" \ + --cache-manifest-sha256 "${CACHE_MANIFEST_SHA256}" \ --output "${output}" echo "file=${output}" >> "$GITHUB_OUTPUT" diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index 11d47de71..3d9914111 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -547,6 +547,7 @@ jobs: STAGE12_ENVIRONMENT: ${{ inputs.deployment_environment }} STAGE12_GCP_PROJECT_ID: ${{ vars.SIMULATION_ENTRYPOINT_GCP_PROJECT_ID }} STAGE12_ARTIFACT_BUCKET: ${{ vars.STAGE12_ARTIFACT_BUCKET }} + STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} STAGE12_ENTRYPOINT_SERVICE_ACCOUNT: ${{ vars.SIMULATION_ENTRYPOINT_RUNTIME_SERVICE_ACCOUNT }} STAGE12_DEPLOY_SERVICE_ACCOUNT: ${{ secrets.SIMULATION_ENTRYPOINT_GCP_DEPLOY_SERVICE_ACCOUNT }} STAGE12_MODAL_SERVICE_ACCOUNT: ${{ vars.STAGE12_MODAL_SERVICE_ACCOUNT }} @@ -568,6 +569,7 @@ jobs: outputs: application_name: ${{ steps.application.outputs.name }} validation_file: ${{ steps.validation.outputs.file }} + cache_manifest_digest: ${{ steps.stage12_cache.outputs.manifest_digest }} steps: - uses: actions/checkout@v6 @@ -603,11 +605,21 @@ jobs: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} STAGE12_ARTIFACT_BUCKET: ${{ vars.STAGE12_ARTIFACT_BUCKET }} + STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} STAGE12_DATABASE_URL_SECRET_NAME: ${{ vars.STAGE12_DATABASE_URL_SECRET_NAME }} STAGE12_GCP_PROJECT_ID: ${{ vars.SIMULATION_ENTRYPOINT_GCP_PROJECT_ID }} STAGE12_GCP_CREDENTIALS_SECRET_NAME: ${{ vars.STAGE12_GCP_CREDENTIALS_SECRET_NAME }} run: ../../.github/scripts/stage12-sync-modal-secrets.sh "${{ inputs.modal_environment }}" + - name: Build the independent Stage 12 cache + id: stage12_cache + working-directory: projects/policyengine-simulation-executor + env: + MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} + MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} + STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} + run: ../../.github/scripts/stage12-cache-precompute.sh "${{ inputs.modal_environment }}" + - name: Deploy the additional versioned v2 Modal application working-directory: projects/policyengine-simulation-executor env: @@ -618,6 +630,8 @@ jobs: STAGE12_EXPECT_US_VERSION: ${{ needs.prepare.outputs.us_version }} STAGE12_EXPECT_UK_VERSION: ${{ needs.prepare.outputs.uk_version }} STAGE12_ARTIFACT_BUCKET: ${{ vars.STAGE12_ARTIFACT_BUCKET }} + STAGE12_CACHE_BUCKET: ${{ vars.STAGE12_CACHE_BUCKET }} + STAGE12_CACHE_MANIFEST_DIGEST: ${{ steps.stage12_cache.outputs.manifest_digest }} OBSERVABILITY_SERVICE_NAMESPACE: ${{ vars.OBSERVABILITY_SERVICE_NAMESPACE }} OBSERVABILITY_TRACE_PROJECT_ID: ${{ vars.OBSERVABILITY_TRACE_PROJECT_ID }} OBSERVABILITY_LOGGING_PROJECT_ID: ${{ vars.OBSERVABILITY_LOGGING_PROJECT_ID }} @@ -634,7 +648,7 @@ jobs: env: MODAL_TOKEN_ID: ${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET: ${{ secrets.MODAL_TOKEN_SECRET }} - run: ../../.github/scripts/stage12-validate-v2-workers.sh "${{ inputs.modal_environment }}" + run: ../../.github/scripts/stage12-validate-v2-workers.sh "${{ inputs.modal_environment }}" "${{ steps.stage12_cache.outputs.manifest_digest }}" - name: Retain worker validation for manifest publication uses: actions/upload-artifact@v7 diff --git a/.github/workflows/stage12-deploy.yml b/.github/workflows/stage12-deploy.yml index c3fb9d1ee..cecb6b097 100644 --- a/.github/workflows/stage12-deploy.yml +++ b/.github/workflows/stage12-deploy.yml @@ -54,7 +54,9 @@ jobs: entrypoint_min_instances: 0 entrypoint_max_instances: 2 promote_entrypoint: false - run_full_integration: true + # This workflow deploys and validates only the independent Stage 12 path. + # The main API deployment continues to run the existing-stack suite. + run_full_integration: false deploy_existing_stack: false deploy_stage12_v2: true secrets: inherit diff --git a/AGENTS.md b/AGENTS.md index 7d6016fe9..16838dede 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -93,7 +93,7 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. ### Stage 12 persistence -- `sim.stage12.deployment-scope`: Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources; the Stage 12 worker can request a Google OIDC token. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) +- `sim.stage12.deployment-scope`: Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources or runs the existing stack's complete integration suite; the Stage 12 worker can request a Google OIDC token. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) - `sim.stage12.deployment-uses-api-schema`: Stage 12 deployment uses the API-owned schema directly, and the Stage 12 runtime it deploys uses typed DML without raw SQL or schema changes. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) - `sim.stage12.no-migration-chain`: This repository has no schema migration chain (no Alembic); Stage 12 tables are migrated by policyengine-api. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py`) - `sim.stage12.no-production-user-associations`: Stage 12 runtime modules never create production user associations (user_simulations or user_reports). (blocking: pytest `projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py`) diff --git a/changelog_entry.yaml b/changelog_entry.yaml index 83047a7a5..163ded3ed 100644 --- a/changelog_entry.yaml +++ b/changelog_entry.yaml @@ -3,6 +3,7 @@ added: - Made CLAUDE.md and the Copilot instructions thin pointers to AGENTS.md, required workflow logic longer than one line to live in tested and shellchecked scripts, and registered the tracking-issue rule - Added a rules registry (rules/registry.toml) with a standard-library tool that validates it, checks every rule is enforced where it claims, and renders docs/generated/rules.md and the Rules block in AGENTS.md + - Ran Stage 12 US national simulations as 20 region-group calls each, and built an independent, immutable Stage 12 cache of three annual US datasets and 60 current-law baselines before each deployment, publishing it only after a recorded determinism check and recording its release in the v2 version manifest under a new key, while the existing key keeps a copy that earlier readers accept changed: - Ran ruff check in CI with the first group of rule families (imports, pyupgrade and other safe auto-fixes) and fixed every finding - Moved the multi-line `run` blocks of the simulation deployment workflows into single-line steps and tested scripts without changing their behaviour or outputs, leaving only the release-sentinel and routing-rollback blocks in the one-line `run` rule's baseline @@ -23,3 +24,4 @@ - Released each Stage 12 database session after its operation instead of holding pooled sessions through calculations, retried opening each Stage 12 coordinator and worker database connection that the pooler refused without ever retrying a statement or commit, and reported Stage 12 failures to telemetry before recording them in the database so a failed write cannot hide the original error - Accepted a null SPM provenance in Stage 12 aggregate artifacts, so UK and other reports without SPM calculations are saved, while every supplied provenance record is still validated - Recorded a redacted, bounded Stage 12 failure cause naming the exception, failing stage, simulation role and message instead of only the exception class, and kept a failed worker's recorded cause on its simulation and report + - Gave Stage 12 baseline and reform simulations one shared wait window that opens when the coordinator starts and fits inside the coordinator's Modal timeout, and cancelled every unfinished simulation, recording it as failed with simulation_cancelled, when a simulation fails, the window closes, or the coordinator fails after starting a simulation diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 48e1975ed..7ebba899a 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -89,7 +89,7 @@ flowchart LR | Rule | Statement | Tier | Status | Blocking | Enforced by | CI | Proof | | --- | --- | --- | --- | --- | --- | --- | --- | -| `sim.stage12.deployment-scope` | Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources; the Stage 12 worker can request a Google OIDC token. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py` | `pr.yml:test` | **unproven** | +| `sim.stage12.deployment-scope` | Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources or runs the existing stack's complete integration suite; the Stage 12 worker can request a Google OIDC token. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py` | `pr.yml:test` | **unproven** | | `sim.stage12.deployment-uses-api-schema` | Stage 12 deployment uses the API-owned schema directly, and the Stage 12 runtime it deploys uses typed DML without raw SQL or schema changes. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py` | `pr.yml:test` | **unproven** | | `sim.stage12.no-migration-chain` | This repository has no schema migration chain (no Alembic); Stage 12 tables are migrated by policyengine-api. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py` | `pr.yml:test` | **unproven** | | `sim.stage12.no-production-user-associations` | Stage 12 runtime modules never create production user associations (user_simulations or user_reports). | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py` | `pr.yml:test` | **unproven** | diff --git a/docs/migration/stage-12-modal-worker-foundation.md b/docs/migration/stage-12-modal-worker-foundation.md index f5a802179..0afbea3f2 100644 --- a/docs/migration/stage-12-modal-worker-foundation.md +++ b/docs/migration/stage-12-modal-worker-foundation.md @@ -19,10 +19,19 @@ simulation submission or polling contracts. simulations therefore resolve in Modal `main` without relabeling durable records as `main`. - One v2 report-coordinator invocation starts baseline and reform as distinct - single-simulation calls before awaiting either result. -- Each country-specific single-simulation function and the report-coordinator - function has a Modal `max_containers` value of 10. This is a per-function - limit; the functions do not share one application-wide container quota. + single-simulation calls, then checks both calls with nonblocking result probes + on every polling pass. +- An eligible US national single-simulation call starts 20 region-group + calculation calls and merges their typed Parquet results. Because baseline + and reform are separate and start before either is awaited, an eligible + report starts 40 region-group calculation calls. These are invocations; + Modal may reuse a warm container for more than one invocation. +- Country-specific single-simulation functions and the report coordinator have + a Modal `max_containers` value of 10. The separate US region-group function + has a value of 300 so calculation work cannot wait behind the parent calls + that are collecting it. These are per-function limits, not one shared quota. +- Stage 12 owns its partition, request and result models, orchestration, merge, + and cache implementation. It does not import the corresponding v1 modules. - The v2 functions write canonical private artifacts and use a shared SQLAlchemy Core package to persist temporary comparison state directly. They do not create production simulations, reports, report runs, or user @@ -67,6 +76,10 @@ with the live API-owned schema. validation downloads that lookup file using the worker credential. Stage 12 local-authority output uses packaged metadata and does not require access to `local_authorities_2021.csv`. +- Stage 12 uses a dedicated environment-specific cache bucket named by + `STAGE12_CACHE_BUCKET`. Its object keys and manifest namespace are separate + from the v1 precompute store. The same restricted worker credential can + create, read, and delete objects in this cache bucket. - `STAGE12_DATABASE_URL` is delivered from the environment-specific Secret Manager resource named by `STAGE12_DATABASE_URL_SECRET_NAME`. It authenticates as the existing shared `policyengine_v2_runtime` account; Stage 12 does not @@ -98,9 +111,10 @@ automation or infrastructure configuration: | Operation | Maintained owner | | --- | --- | | Build and deploy each versioned v2 Modal application | Simulation deployment workflow | +| Build and verify the Stage 12 dataset and baseline cache | Simulation deployment workflow | | Create or update the separate v2 manifest storage object | V2 manifest publisher | | Configure the v2 manifest reader and `STAGE12_ENABLED` | Cloud Run deployment workflow | -| Create the private artifact namespace and its retention policy | Infrastructure configuration | +| Create the private report-artifact and cache buckets and their policies | Infrastructure configuration | | Ensure the existing API v2 runtime identity has comparison-row database access | API infrastructure configuration | | Deliver the existing shared runtime database secret to Cloud Run and Modal | Simulation deployment workflow | | Grant private-object access only to the Modal runtime | Infrastructure configuration | @@ -160,9 +174,10 @@ Before deploying, the simulation workflow obtains the shared runtime URL from Secret Manager, verifies the expected `policyengine_v2_runtime` identity and permissions, inspects the complete temporary schema against the declared SQLAlchemy mappings, and runs a rolled-back insert/read/update/delete canary. It -also verifies required Secret Manager access and performs a create/read/delete -canary in the private artifact bucket. The storage canary is deleted in the -normal path and by exit cleanup after a failure. +also verifies required Secret Manager access and performs separate +create/read/delete checks in the private report-artifact and cache buckets. The +temporary objects are deleted in the normal path and by exit cleanup after a +failure. The current Public API omits `data` for the certified default dataset. The v2 adapter therefore resolves an absent `data` field to the exact default dataset @@ -255,6 +270,136 @@ move to the front end, and Stage 14 or later work should remove these temporary executor resources when the authoritative v2 report architecture supersedes Stage 12. +## US segmentation and baseline cache + +Segmentation applies by default when both simulation inputs describe a US +national society-wide calculation and the shared output plan does not require +cliff analysis or labor-supply responses. `segmented=false` selects the +single-container calculation explicitly. UK calculations, subnational US +calculations, cliff analysis, and labor-supply-response calculations also use +the single-container path. `segmented=true` is accepted, but it does not +override those numerical eligibility conditions. + +The Stage 12 coordinator still persists only two logical simulation records: +one baseline and one reform. Each eligible logical simulation call performs +the following work independently: + +```text +logical baseline or reform simulation + | + +---- start 20 US region-group calculation calls + | + +---- collect and validate all 20 results + | + +---- concatenate entity rows and reject duplicate identifiers + | + +---- write one logical simulation artifact and update one record +``` + +Region-group calls do not write temporary database rows or private report +artifacts. They return compressed, dtype-preserving data to their logical +simulation parent. The parent validates the result role, partition index, +region membership, content digest, entity schemas, column dtypes, and unique +entity identifiers before merging. A missing or failed region-group call fails +that logical simulation. The parent makes a best-effort cancellation request +for every calculation call that it already started. + +The report coordinator observes baseline and reform together rather than +blocking on one role first. On the first failure, it copies the failed +simulation's `error_code` and safe `error_summary` to the report, records the +unresolved peer as `failed` with `simulation_cancelled` and the summary +"Cancelled after the baseline simulation failed." (or reform), and completes +the report as failed without waiting for the peer calculation. An unsegmented +peer is also cancelled through Modal. A segmented peer checks its persisted +simulation status during bounded waits; after seeing the cancellation, it +cancels all region-group calls it owns and exits. Modal does not cancel +spawned calls recursively, so the coordinator cancels a segmented peer through +Modal only if it could not write the cancellation record. Only a `running` +record changes, so a simulation that already succeeded or failed keeps its +outcome, and atomic status replacement prevents a late calculation result from +replacing the cancellation record. + +Both simulations share one 3,900-second wait window, the longest simulation +worker timeout. It opens when the coordinator starts, so the coordinator's own +setup, including loading the country model and dispatching both simulations, +comes out of that window and not out of the time it needs afterwards. When the +window closes, the report fails with `simulation_wait_timeout`, and each +unfinished simulation is cancelled in the same way. The coordinator's +5,400-second Modal timeout is that window plus up to 900 seconds for the +production-result wait plus 600 seconds to aggregate, write artifacts, and +record the outcome. + +If the coordinator itself fails after starting a simulation, for example while +recording a later dispatch or attaching a started call's invocation +identifier, it cancels every started, unfinished simulation in the same way, +with the summary "Cancelled after the coordinator failed.". The report keeps +the coordinator's original failure. + +Errors while reading or updating simulation records never replace the failure +that caused them. If the coordinator cannot read or update the failed +simulation's own record, or record a peer's cancellation, it logs that error +with the report and simulation identifiers and the error type, but not the +error's message. The report still records the simulation's failure, the +failed simulation is never recorded as cancelled, and its peers are still +cancelled with the reason "Cancelled after the baseline simulation failed." +(or reform). + +Typed input failures, including invalid policy periods and SPM selection +errors, retain their code and bounded message through the segment, simulation, +and report layers. Unexpected exceptions are logged with their complete stack +trace and persisted as their redacted cause followed by the correlation ID, +for example `OperationalError in stage12_child_state_create: ... +(correlation_id=...)`. Redaction is described in the executor's +`stage12_runtime.failure_causes` module; the cause is shortened, never the +correlation ID, to fit the 512-character summary. The coordinator and +simulation workers report each failure once, before any database write, and +that report carries the same correlation ID. + +Before each v2 deployment, the independent Stage 12 precompute application +builds three annual US dataset files and 60 current-law baseline files: one for +each of 20 region groups in 2025, 2026, and 2027. Cache identities include the +bundle manifest digest, package and data versions, source-data digest, year, +region membership, scoping identity, output-plan digest, and the resolved SPM +selection the worker builds the simulation with. A baseline file has the name +.py's `Simulation.storage_id` gives the file it saves and loads: the identity's +simulation ID followed by the SPM selection's digest. Objects are +content-addressed and written with a create-only condition; an existing object +must have the expected digest and size. A stored file is never replaced, so +precompute computes only missing files and has no force option; the v1 +`force_recompute` input does not reach it. To replace a bad file, delete it and +deploy again: the run recomputes it and repeats the determinism check. + +Before publishing, precompute loads one stored baseline from the cache and +compares it with a fresh calculation. The check runs whenever a run computed +baselines, and also whenever no earlier run recorded a pass for the planned +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 +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. + +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 +compute. Missing cache files fall back to calculation; a loaded file that does +not satisfy the requested output plan is recalculated. Each region-group log +records whether it loaded a complete file, recalculated an incomplete file, or +computed because no file existed. + ## Temporary direct runner endpoint > **Temporary Stage 12 interface:** The authenticated routes in this section @@ -299,9 +444,9 @@ Modal report coordinator | atomically create or resolve the temporary parent row | stop if another invocation already owns the same logical run | - +------> baseline single-simulation call + +------> baseline logical simulation call ------> 20 region-group calls | - +------> reform single-simulation call + +------> reform logical simulation call --------> 20 region-group calls | +------> write aggregate artifact and durable state @@ -333,11 +478,25 @@ curl -X POST "${SIMULATION_ENTRYPOINT_URL}/internal/stage12/reports" \ "scope": "macro", "region": "us", "time_period": "2026", + "segmented": true, "baseline": {}, - "reform": {} + "reform": { + "gov.irs.credits.ctc.amount.base[0].amount": { + "2026-01-01.2100-12-31": 3000 + } + } }' ``` +Policy values may be scalars, ISO effective-date mappings, or frontend interval +mappings in `YYYY-MM-DD.YYYY-MM-DD` form. A bare year such as `"2026"` is not +a valid effective-period key. + +The `segmented` field is optional. Omitting it uses segmentation when the +request meets the eligibility conditions above. `false` explicitly selects the +single-container path. `true` requests segmentation but still falls back to +the single-container path when the output requirements are not eligible. + An accepted request returns `202` and `Retry-After: 1` without waiting for the calculations or for the coordinator's first database write: @@ -359,6 +518,10 @@ While the report is pending or running, GET returns `202` and `Retry-After: 5`; completed, failed, incomplete, or skipped states return `200`. The response contains the complete temporary parent and child execution metadata, including private artifact references and digests, but never returns artifact contents. +For a failed report, the parent record contains the originating safe failure; +the failed simulation contains the same values, and a peer stopped because of +that failure is `failed` with `simulation_cancelled`; its `error_summary` +gives the reason. Cloud Run therefore requires no additional object-storage permission for this interface. @@ -392,7 +555,7 @@ starting child simulations. Unsupported automatic inputs are logged with a bounded reason and are not persisted. The report coordinator starts baseline and reform as independent Modal calls, -waits for them, and derives the Stage 12 aggregate. It then restores the +observes both until they succeed, and derives the Stage 12 aggregate. It then restores the production Modal function call by the retained production job identifier and waits up to 15 minutes for that production call to finish. It compares both complete aggregate result objects exactly. The private @@ -434,3 +597,9 @@ aggregate, and compares every aggregate field. Differences fail qualification with only a bounded code and field path; calculation values are not included in the diagnostic or receipt. Numerical tolerances default to zero and can be introduced only for explicit field paths through a reviewed JSON file. +For an eligible US national input, the Stage 12 side of qualification executes +and merges all 20 region groups separately for both baseline and reform. This +checks the partition and merge numerically in addition to comparing the final +entity tables and aggregate report with the existing combined executor. Inputs +that are not eligible for segmentation continue through the single-container +Stage 12 calculation. diff --git a/docs/operations/api-v1-observability.md b/docs/operations/api-v1-observability.md index 07cacc8b6..77ed93d53 100644 --- a/docs/operations/api-v1-observability.md +++ b/docs/operations/api-v1-observability.md @@ -54,8 +54,13 @@ The simulation entry service dispatches the current production run and the Stage 12 report with the same `observability_id`. Stage 12 creates separate functional identifiers for the report and its baseline and reform simulations. The coordinator passes captured context to both simulation workers. When Stage -12 becomes the authoritative runner, the same report identifier flow applies; -only the selected functional result changes. +12 runs an eligible US national report, each simulation worker passes that same +context to 20 region-group calculation calls. Baseline and reform therefore +produce 40 region-group invocations with one `observability_id`, two +`simulation_execution_id` values, and no additional persisted simulation +records. The region-group index and cache outcome distinguish their structured +logs. When Stage 12 becomes the authoritative runner, the same report +identifier flow applies; only the selected functional result changes. ## Query model diff --git a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py index a048b8f3a..f44db5ed3 100644 --- a/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py +++ b/libs/policyengine-simulation-contract/src/policyengine_simulation_contract/stage12_manifest.py @@ -21,7 +21,10 @@ V1_ROUTING_STATE_NAME = "simulation-api-routing-state" V2_VERSION_MANIFEST_NAME = "simulation-api-v2-version-manifest" -ACTIVE_MANIFEST_KEY = "active" +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. +PRE_CACHE_MANIFEST_KEY = "active" V2_APPLICATION_PREFIX = "policyengine-simulation-v2-py" V1_APPLICATION_PREFIX = "policyengine-simulation-py" _CALLABLE_PATTERN = re.compile(r"^[a-z][a-z0-9_]{0,127}$") @@ -39,6 +42,15 @@ class V2CountryWorker(StrictManifestModel): single_simulation_callable: CallableName +class Stage12CacheRelease(StrictManifestModel): + manifest_sha256: Sha256Digest + partition_sha256: Sha256Digest + years: tuple[int, ...] + segment_count: Literal[20] = 20 + dataset_count: Literal[3] = 3 + baseline_count: Literal[60] = 60 + + class V2WorkerValidation(StrictManifestModel): validated: Literal[True] application_name: ManifestText @@ -46,6 +58,7 @@ class V2WorkerValidation(StrictManifestModel): validated_at: datetime validation_invocation_id: ManifestText country_validation_invocation_ids: dict[CountryId, ManifestText] + cache: Stage12CacheRelease | None = None class V2WorkerVersion(StrictManifestModel): @@ -55,6 +68,7 @@ class V2WorkerVersion(StrictManifestModel): bundle: Stage12BundleManifest bundle_manifest_sha256: Sha256Digest validation: V2WorkerValidation + cache: Stage12CacheRelease | None = None @model_validator(mode="after") def validate_worker(self) -> V2WorkerVersion: @@ -69,6 +83,8 @@ def validate_worker(self) -> V2WorkerVersion: raise ValueError("worker validation names another application") if self.validation.bundle_manifest_sha256 != self.bundle_manifest_sha256: raise ValueError("worker validation names another bundle digest") + if self.validation.cache != self.cache: + raise ValueError("worker cache metadata differs from worker validation") countries = tuple(worker.country for worker in self.countries) bundle_countries = tuple(country.country for country in self.bundle.countries) if len(countries) != len(set(countries)): @@ -137,6 +153,23 @@ def _is_newer(candidate: str, current: str | None) -> bool: return False +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"}}}}, + ) + + +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 + + def build_next_v2_manifest( *, current: object, @@ -173,13 +206,14 @@ def publish_v2_manifest( force_default: bool = False, ) -> V2VersionManifest: next_manifest = build_next_v2_manifest( - current=store.get(ACTIVE_MANIFEST_KEY), + current=_current_manifest(store), worker=worker, force_default=force_default, ) - # The entire validated document replaces one store value. Readers cannot - # observe a partial update assembled from several keys. + # 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") + store[PRE_CACHE_MANIFEST_KEY] = pre_cache_manifest_payload(next_manifest) return next_manifest @@ -191,7 +225,7 @@ def __init__(self, store: V2ManifestStore): self._last_valid: V2VersionManifest | None = None def load(self) -> V2VersionManifest: - value = self._store.get(ACTIVE_MANIFEST_KEY) + value = _current_manifest(self._store) try: manifest = V2VersionManifest.model_validate(value) except (TypeError, ValueError): diff --git a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py index ddf024aff..7ab548a54 100644 --- a/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py +++ b/libs/policyengine-simulation-contract/tests/test_stage12_manifest.py @@ -2,7 +2,7 @@ from __future__ import annotations -from datetime import datetime, timezone +from datetime import UTC, datetime import pytest from policyengine_simulation_contract.stage12_bundle import ( @@ -13,14 +13,17 @@ ) from policyengine_simulation_contract.stage12_manifest import ( ACTIVE_MANIFEST_KEY, + PRE_CACHE_MANIFEST_KEY, V1_ROUTING_STATE_NAME, V2_VERSION_MANIFEST_NAME, + Stage12CacheRelease, V2CountryWorker, V2ManifestLoader, V2WorkerValidation, V2WorkerVersion, assert_separate_manifest_names, manifest_digest, + pre_cache_manifest_payload, publish_v2_manifest, v2_application_name, validate_manifest_has_no_credentials, @@ -105,7 +108,7 @@ def _worker(version: str = "5.2.0") -> V2WorkerVersion: validated=True, application_name=application_name, bundle_manifest_sha256=resolved.bundle_manifest_sha256, - validated_at=datetime(2026, 9, 14, tzinfo=timezone.utc), + validated_at=datetime(2026, 9, 14, tzinfo=UTC), validation_invocation_id="validation-123", country_validation_invocation_ids={ "us": "validation-us-123", @@ -120,14 +123,87 @@ def test_v2_application_name_cannot_collide_with_v1() -> None: assert not v2_application_name("5.2.0").startswith("policyengine-simulation-py") -def test_publish_writes_one_complete_v2_document() -> None: +def _cached_worker(version: str = "5.2.0") -> V2WorkerVersion: + worker = _worker(version) + cache = Stage12CacheRelease( + manifest_sha256="c" * 64, + partition_sha256="d" * 64, + years=(2026, 2027, 2025), + ) + return worker.model_copy( + update={ + "cache": cache, + "validation": worker.validation.model_copy(update={"cache": cache}), + } + ) + + +def test_publish_writes_one_complete_v2_document_per_key() -> None: store = FakeStore() - manifest = publish_v2_manifest(store=store, worker=_worker()) + manifest = publish_v2_manifest(store=store, worker=_cached_worker()) - assert store.writes == [(ACTIVE_MANIFEST_KEY, manifest.model_dump(mode="json"))] + assert store.writes == [ + (ACTIVE_MANIFEST_KEY, manifest.model_dump(mode="json")), + (PRE_CACHE_MANIFEST_KEY, pre_cache_manifest_payload(manifest)), + ] assert manifest.default_version == "5.2.0" assert manifest.versions["5.2.0"].validation.validated is True + assert V2ManifestLoader(store).load().versions["5.2.0"].cache is not None + + +def test_readers_released_before_the_cache_get_only_fields_they_accept() -> None: + store = FakeStore() + publish_v2_manifest(store=store, worker=_cached_worker()) + + # The fields of each model before Stage 12 cache releases. Those readers + # forbid any other field, so adding one here would stop them loading. + pre_cache = store.values[PRE_CACHE_MANIFEST_KEY] + assert set(pre_cache) == { + "default_version", + "generation", + "schema_version", + "versions", + } + worker = pre_cache["versions"]["5.2.0"] + assert set(worker) == { + "application_name", + "bundle", + "bundle_manifest_sha256", + "countries", + "report_coordinator_callable", + "validation", + } + assert set(worker["validation"]) == { + "application_name", + "bundle_manifest_sha256", + "country_validation_invocation_ids", + "validated", + "validated_at", + "validation_invocation_id", + } + + +def test_first_publish_with_cache_keeps_versions_from_the_pre_cache_key() -> None: + earlier = _worker("5.1.0").model_dump( + mode="json", exclude={"cache": True, "validation": {"cache"}} + ) + store = FakeStore( + { + PRE_CACHE_MANIFEST_KEY: { + "schema_version": 1, + "generation": "5.1.0:earlier", + "default_version": "5.1.0", + "versions": {"5.1.0": earlier}, + } + } + ) + assert V2ManifestLoader(store).load().default_version == "5.1.0" + + manifest = publish_v2_manifest(store=store, worker=_cached_worker()) + + assert sorted(manifest.versions) == ["5.1.0", "5.2.0"] + assert V2ManifestLoader(store).load() == manifest def test_loader_retains_its_last_valid_v2_document() -> None: @@ -177,6 +253,21 @@ def test_manifest_rejects_incomplete_or_unvalidated_worker() -> None: with pytest.raises(ValueError): V2WorkerVersion.model_validate(value) + +def test_cache_metadata_is_optional_for_old_workers_and_must_match_validation() -> None: + old = _worker() + assert old.cache is None + + cache = Stage12CacheRelease( + manifest_sha256="c" * 64, + partition_sha256="d" * 64, + years=(2026, 2027, 2025), + ) + value = old.model_dump(mode="json") + value["cache"] = cache.model_dump(mode="json") + with pytest.raises(ValueError, match="cache metadata differs"): + V2WorkerVersion.model_validate(value) + value = _worker().model_dump(mode="json") value["validation"]["validated"] = False with pytest.raises(ValueError): diff --git a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py index 6b8d6afbc..93f4e2627 100644 --- a/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py +++ b/libs/policyengine-simulation-observability/src/policyengine_simulation_observability/stages.py @@ -20,6 +20,7 @@ class RunConfiguration(StrEnum): STAGE12_SHADOW_REPORT = "stage12_shadow_report" STAGE12_CANONICAL_REPORT = "stage12_canonical_report" STAGE12_SIMULATION = "stage12_simulation" + STAGE12_SEGMENT = "stage12_segment" class Stage(StrEnum): @@ -91,6 +92,11 @@ class Stage(StrEnum): STAGE12_CHILD_WAIT = "stage12_child_wait" STAGE12_INPUT_WRITE = "stage12_input_write" STAGE12_SIMULATION_EXECUTION = "stage12_simulation_execution" + STAGE12_SEGMENT_EXECUTION = "stage12_segment_execution" + STAGE12_SEGMENT_DISPATCH = "stage12_segment_dispatch" + STAGE12_SEGMENT_WAIT = "stage12_segment_wait" + STAGE12_SEGMENT_MERGE = "stage12_segment_merge" + STAGE12_CACHE_LOOKUP = "stage12_cache_lookup" STAGE12_CALCULATION = "stage12_calculation" STAGE12_SIMULATION_ARTIFACT_WRITE = "stage12_simulation_artifact_write" STAGE12_ARTIFACT_READ = "stage12_artifact_read" @@ -226,6 +232,9 @@ def name(self, stage: Stage) -> str: ( Stage.STAGE12_SIMULATION_EXECUTION, Stage.STAGE12_INPUT_WRITE, + Stage.STAGE12_SEGMENT_DISPATCH, + Stage.STAGE12_SEGMENT_WAIT, + Stage.STAGE12_SEGMENT_MERGE, Stage.CREDENTIAL_SETUP, Stage.COUNTRY_MODULE_LOAD, Stage.REGION_RESOLUTION, @@ -233,10 +242,26 @@ def name(self, stage: Stage) -> str: Stage.DATASET_LOAD, Stage.POLICY_NORMALIZATION, Stage.SIMULATION_BUILD, + Stage.STAGE12_CACHE_LOOKUP, Stage.STAGE12_CALCULATION, Stage.STAGE12_SIMULATION_ARTIFACT_WRITE, ), ), + RunConfiguration.STAGE12_SEGMENT: StagePlan( + RunConfiguration.STAGE12_SEGMENT, + ( + Stage.STAGE12_SEGMENT_EXECUTION, + Stage.CREDENTIAL_SETUP, + Stage.COUNTRY_MODULE_LOAD, + Stage.REGION_RESOLUTION, + Stage.DATASET_RESOLUTION, + Stage.DATASET_LOAD, + Stage.POLICY_NORMALIZATION, + Stage.SIMULATION_BUILD, + Stage.STAGE12_CACHE_LOOKUP, + Stage.STAGE12_CALCULATION, + ), + ), } ) @@ -250,3 +275,4 @@ def name(self, stage: Stage) -> str: RunConfiguration.STAGE12_CANONICAL_REPORT ] STAGE12_SIMULATION_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.STAGE12_SIMULATION] +STAGE12_SEGMENT_STAGES = RUN_STAGE_REGISTRY[RunConfiguration.STAGE12_SEGMENT] diff --git a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py index f463d0a32..b9e0ce05d 100644 --- a/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py +++ b/libs/policyengine-stage12-persistence/src/policyengine_stage12_persistence/store.py @@ -492,6 +492,55 @@ def replace_simulation( ) return _simulation(updated) + def replace_simulation_if_status( + self, + record: ComparisonSimulationRecord, + *, + expected_status: ComparisonRunLifecycleStatus, + ) -> tuple[ComparisonSimulationRecord, bool]: + """Replace a simulation only while it remains in ``expected_status``.""" + + values = _record_values(record) + with self._transaction("replace_simulation_if_status") as connection: + row = ( + connection.execute( + select(comparison_simulations) + .where( + comparison_simulations.c.simulation_execution_id + == record.simulation_execution_id + ) + .with_for_update() + ) + .mappings() + .one_or_none() + ) + if row is None: + raise LookupError( + "comparison simulation " + f"{record.simulation_execution_id} does not exist" + ) + existing = _simulation(row) + require_simulation_identity(existing, record) + if existing.status is not expected_status: + return existing, False + require_lifecycle_transition(existing.status, record.status) + updated = ( + connection.execute( + update(comparison_simulations) + .where( + comparison_simulations.c.simulation_execution_id + == record.simulation_execution_id + ) + .values( + **{name: values[name] for name in SIMULATION_MUTABLE_COLUMNS} + ) + .returning(*comparison_simulations.c) + ) + .mappings() + .one() + ) + return _simulation(updated), True + def attach_simulation_invocation( self, simulation_execution_id: UUID, diff --git a/libs/policyengine-stage12-persistence/tests/test_store.py b/libs/policyengine-stage12-persistence/tests/test_store.py index 9df7938b2..7cf41b236 100644 --- a/libs/policyengine-stage12-persistence/tests/test_store.py +++ b/libs/policyengine-stage12-persistence/tests/test_store.py @@ -370,6 +370,74 @@ def test_invocation_attachment_is_compare_and_set() -> None: assert result == attached +def test_simulation_status_replacement_is_compare_and_set() -> None: + running = _simulation().model_copy( + update={"status": ComparisonRunLifecycleStatus.RUNNING} + ) + failed = running.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_execution_failed", + "completed_at": NOW, + } + ) + engine = FakeEngine( + [ + FakeResult(row=_row(running)), + FakeResult(row=_row(failed)), + ] + ) + opener = RetryingConnectionOpener() + store = Stage12PersistenceStore( + engine=engine, # type: ignore[arg-type] + open_connection=opener, + ) + + result, replaced = store.replace_simulation_if_status( + failed, + expected_status=ComparisonRunLifecycleStatus.RUNNING, + ) + + assert replaced is True + assert result == failed + assert opener.operation_names == ["replace_simulation_if_status"] + assert len(engine.connection.statements) == 2 + + +def test_simulation_status_replacement_preserves_a_concurrent_cancellation() -> None: + cancelled = _simulation().model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_cancelled", + "error_summary": "Cancelled after the reform simulation failed.", + "completed_at": NOW, + } + ) + succeeded = _simulation().model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": "gs://private/baseline.parquet", + "output_sha256": "c" * 64, + "output_schema_version": 1, + "row_identity_columns": ("household.household_id",), + "row_count": 1, + "row_identity_sha256": "d" * 64, + "completed_at": NOW, + } + ) + engine = FakeEngine([FakeResult(row=_row(cancelled))]) + store = Stage12PersistenceStore(engine=engine) # type: ignore[arg-type] + + result, replaced = store.replace_simulation_if_status( + succeeded, + expected_status=ComparisonRunLifecycleStatus.RUNNING, + ) + + assert replaced is False + assert result == cancelled + assert len(engine.connection.statements) == 1 + + def test_database_url_uses_one_canonical_secret_format() -> None: parsed = normalized_database_url( "postgresql://policyengine_v2_runtime:secret@db.example/postgres" diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py index c51e7aa2e..f0873ac4b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/stage12_adapter.py @@ -99,8 +99,9 @@ def adapt_annual_comparison( include_cliffs = payload.get("include_cliffs") if include_cliffs is not None and not isinstance(include_cliffs, bool): return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE) - if payload.get("segmented") not in {None, False}: - return _skip(ComparisonSkipReason.UNSUPPORTED_OPTIONS) + segmented = payload.get("segmented") + if segmented is not None and not isinstance(segmented, bool): + return _skip(ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE) country = payload.get("country") if not isinstance(country, str) or _country_worker(worker, country) is None: return _skip(ComparisonSkipReason.UNSUPPORTED_COUNTRY) @@ -175,7 +176,11 @@ def adapt_annual_comparison( country=cast(CountryId, country), region=region, ) - options = {key: payload[key] for key in ("spm",) if payload.get(key) is not None} + options = { + key: payload[key] + for key in ("spm", "segmented") + if payload.get(key) is not None + } if include_cliffs is True: options["include_cliffs"] = True diff --git a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py index 4133ac0c5..46ebe97e7 100644 --- a/projects/policyengine-simulation-entry/tests/stage12_fixtures.py +++ b/projects/policyengine-simulation-entry/tests/stage12_fixtures.py @@ -97,7 +97,9 @@ def eligible_payload() -> dict: "country": "us", "scope": "macro", "baseline": {}, - "reform": {"gov.irs.credits.ctc.amount.base[0].amount": {"2026": 3000}}, + "reform": { + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01.2100-12-31": 3000} + }, "time_period": "2026", "region": "us", "include_cliffs": False, diff --git a/projects/policyengine-simulation-entry/tests/test_app.py b/projects/policyengine-simulation-entry/tests/test_app.py index d50e4e7f9..a7539de32 100644 --- a/projects/policyengine-simulation-entry/tests/test_app.py +++ b/projects/policyengine-simulation-entry/tests/test_app.py @@ -487,6 +487,54 @@ def test_temporary_stage12_poll_returns_terminal_metadata_with_200(backend): ) +def test_temporary_stage12_poll_returns_propagated_failure_values(backend): + report = comparison_report(status=ComparisonRunLifecycleStatus.FAILED).model_copy( + update={ + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + } + ) + baseline = comparison_simulation(SimulationRole.BASELINE).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_cancelled", + "error_summary": "Cancelled after the reform simulation failed.", + "completed_at": report.completed_at, + } + ) + reform = comparison_simulation(SimulationRole.REFORM).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + "completed_at": report.completed_at, + } + ) + comparison = TemporaryComparisonBackend( + report=report, + simulations=(baseline, reform), + ) + app = create_app( + settings=make_settings(), + backend=backend, + auth_dependency=lambda: None, + comparison_backend=comparison, + ) + + from fastapi.testclient import TestClient + + with TestClient(app) as test_client: + result = test_client.get(f"/internal/stage12/reports/{EVALUATION_ID}") + + assert result.status_code == 200 + payload = result.json() + assert payload["report"]["error_code"] == "invalid_policy_period" + assert payload["report"]["error_summary"] == "The reform period is invalid." + assert payload["simulations"][0]["status"] == "failed" + assert payload["simulations"][0]["error_code"] == "simulation_cancelled" + assert payload["simulations"][1]["error_code"] == "invalid_policy_period" + + def test_temporary_stage12_routes_are_excluded_from_openapi(backend): app = create_app( settings=make_settings(), diff --git a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py index 302133ae2..1fda38a7c 100644 --- a/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py +++ b/projects/policyengine-simulation-entry/tests/test_deploy_workflow_scripts.py @@ -293,6 +293,7 @@ def test_candidate_fails_when_the_tag_is_not_deployed(fakes, tmp_path): V2_APPLICATION_NAME = "stage12-v2-application-name.sh" VALIDATE_V2_WORKERS = "stage12-validate-v2-workers.sh" +CACHE_MANIFEST_SHA256 = "a" * 64 STAGE12_READY = "stage12-deployment-ready.sh" V2_APPLICATION_CODE = ( "from policyengine_simulation_contract.stage12_manifest import " @@ -320,9 +321,9 @@ def test_deploy_workflow_runs_the_stage12_worker_scripts(): validation ) assert ( - f'run: ../../.github/scripts/{VALIDATE_V2_WORKERS} "${{{{ inputs.modal_environment }}}}"' - in validation - ) + f'run: ../../.github/scripts/{VALIDATE_V2_WORKERS} "${{{{ inputs.modal_environment }}}}" ' + '"${{ steps.stage12_cache.outputs.manifest_digest }}"\n' + ) in validation assert "application_name: ${{ steps.application.outputs.name }}" in workflow assert "validation_file: ${{ steps.validation.outputs.file }}" in workflow @@ -367,7 +368,9 @@ def _validate_v2_workers(fakes: FakeCommands, tmp_path: Path, *args: str): def test_validate_v2_workers_records_the_validation_file(fakes, tmp_path): fakes.add("uv") - result, output, validation_file = _validate_v2_workers(fakes, tmp_path, "main") + result, output, validation_file = _validate_v2_workers( + fakes, tmp_path, "main", CACHE_MANIFEST_SHA256 + ) assert result.returncode == 0, result.stderr assert output == f"file={validation_file}\n" @@ -380,6 +383,8 @@ def test_validate_v2_workers_records_the_validation_file(fakes, tmp_path): "src.modal.utils.validate_v2_worker_release", "--environment", "main", + "--cache-manifest-sha256", + CACHE_MANIFEST_SHA256, "--output", str(validation_file), ] @@ -391,7 +396,9 @@ def test_validate_v2_workers_fails_without_writing_when_validation_fails( ): fakes.add("uv", "exit 1\n") - result, output, _ = _validate_v2_workers(fakes, tmp_path, "main") + result, output, _ = _validate_v2_workers( + fakes, tmp_path, "main", CACHE_MANIFEST_SHA256 + ) assert result.returncode == 1 assert output == "" @@ -408,6 +415,17 @@ def test_validate_v2_workers_requires_the_modal_environment(fakes, tmp_path): assert fakes.calls() == [] +def test_validate_v2_workers_requires_the_cache_manifest_digest(fakes, tmp_path): + fakes.add("uv") + + result, output, _ = _validate_v2_workers(fakes, tmp_path, "main") + + assert result.returncode != 0 + assert "Stage 12 cache manifest digest required" in result.stderr + assert output == "" + assert fakes.calls() == [] + + def test_deploy_workflow_checks_stage12_results_with_the_tested_script(): workflow = REUSABLE_DEPLOY_WORKFLOW.read_text(encoding="utf-8") step = _step(workflow, "Check Stage 12 deployment result") diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py index 5c7b0c080..5cc6cb484 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py @@ -712,6 +712,13 @@ def test_complete_integration_suite_is_configured_for_beta_only(): assert "run_full_integration:" in reusable_workflow +def test_stage12_only_deployment_skips_existing_stack_integration_suite(): + stage12_workflow = STAGE12_DEPLOY_WORKFLOW.read_text(encoding="utf-8") + + assert stage12_workflow.count("run_full_integration: false") == 2 + assert "run_full_integration: true" not in stage12_workflow + + def test_main_deployment_automatically_deploys_stage12_in_both_environments(): deploy_workflow = DEPLOY_WORKFLOW.read_text(encoding="utf-8") reusable_workflow = REUSABLE_DEPLOY_WORKFLOW.read_text(encoding="utf-8") diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py index c1b8ef4a9..53a449ec9 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_adapter.py @@ -106,6 +106,21 @@ def test_cliff_analysis_is_forwarded_to_both_simulations() -> None: assert result.report.reform.options == {"include_cliffs": True} +@pytest.mark.parametrize("segmented", [True, False]) +def test_segmentation_selection_is_forwarded_to_both_simulations(segmented) -> None: + payload = {**eligible_payload(), "segmented": segmented} + + result = adapt_annual_comparison( + payload, + evaluation_id=EVALUATION_ID, + worker=worker(), + ) + + assert result.report is not None + assert result.report.baseline.options == {"segmented": segmented} + assert result.report.reform.options == {"segmented": segmented} + + @pytest.mark.parametrize( ("update", "reason"), [ @@ -123,7 +138,7 @@ def test_cliff_analysis_is_forwarded_to_both_simulations() -> None: ), ({"time_period": "2026-01"}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), ({"region": None}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), - ({"segmented": True}, ComparisonSkipReason.UNSUPPORTED_OPTIONS), + ({"segmented": "true"}, ComparisonSkipReason.UNSUPPORTED_REQUEST_SHAPE), ({"region_group": ["state/ca"]}, ComparisonSkipReason.UNSUPPORTED_OPTIONS), ], ) diff --git a/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py b/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py index 4eb899072..2fd05a021 100644 --- a/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py +++ b/projects/policyengine-simulation-entry/tests/test_stage12_infrastructure.py @@ -29,9 +29,10 @@ def test_infrastructure_validation_is_bounded_and_valid_shell() -> None: source = SCRIPT.read_text(encoding="utf-8") assert "STAGE12_ENVIRONMENT" in source assert "STAGE12_ARTIFACT_BUCKET" in source + assert "STAGE12_CACHE_BUCKET" in source assert ( - "Pre-provisioned Stage 12 database, secret, artifact storage, and UK " - "geography lookup access" in source + "Pre-provisioned Stage 12 database, secret, report storage, cache storage, " + "and UK geography lookup access" in source ) assert "STAGE12_DATABASE_URL_SECRET_NAME" in source assert "policyengine_v2_runtime" in source @@ -60,6 +61,7 @@ def test_infrastructure_validation_is_bounded_and_valid_shell() -> None: assert "stage12-evaluation-runtime" in modal_sync assert "stage12-evaluation-gcp-credentials" in modal_sync assert "STAGE12_DATABASE_URL:" in modal_sync + assert "STAGE12_CACHE_BUCKET:" in modal_sync assert "GOOGLE_APPLICATION_CREDENTIALS_JSON:" in modal_sync @@ -75,7 +77,7 @@ def test_infrastructure_validation_rejects_missing_configuration() -> None: assert "STAGE12_ENVIRONMENT is required" in result.stderr -def test_storage_validation_uses_only_the_modal_worker_credentials() -> None: +def test_storage_validation_checks_runtime_and_deployment_credentials() -> None: source = SCRIPT.read_text(encoding="utf-8") helper_start = source.index("runtime_gcloud()") helper_end = source.index("\n}\n", helper_start) @@ -90,8 +92,13 @@ def test_storage_validation_uses_only_the_modal_worker_credentials() -> None: assert "runtime_gcloud auth activate-service-account" in source modal_account = '--account="${STAGE12_MODAL_SERVICE_ACCOUNT}"' - assert source.count(f"runtime_gcloud {modal_account} storage cp") == 3 - assert source.count(f"runtime_gcloud {modal_account} storage rm") == 2 + assert source.count(f"runtime_gcloud {modal_account} storage cp") == 5 + assert source.count(f"runtime_gcloud {modal_account} storage rm") == 4 + assert ( + 'gcloud storage cp \\\n "${cache_canary_object}" ' + '"${canary_download_file}" --quiet' + ) in source + assert "gs://${STAGE12_CACHE_BUCKET}/_deployment-validation" in source assert "gs://policyengine-uk-data-private/constituencies_2024.csv" in source assert "gs://policyengine-uk-data-private/local_authorities_2021.csv" not in source assert "STAGE12_UK_GEOGRAPHY_BUCKET" not in source diff --git a/projects/policyengine-simulation-executor/README.md b/projects/policyengine-simulation-executor/README.md index 83170934c..64f2f60f8 100644 --- a/projects/policyengine-simulation-executor/README.md +++ b/projects/policyengine-simulation-executor/README.md @@ -156,6 +156,68 @@ Operational notes: That staleness class predates the artifact pipeline and lives in the bundle install, not the fetch. +## Independent Stage 12 segmentation and cache + +The separately named Stage 12 v2 application has its own implementation under +`stage12_runtime/` and `stage12_cache/`. It does not call the v1 segmented-run +or precompute modules described above. + +For an eligible US national report, the coordinator starts baseline and reform +as separate logical simulations. Each logical simulation starts 20 calls to +`run_single_simulation_segment_us`, then validates and concatenates their +dtype-preserving entity tables. This produces 40 region-group calculation +invocations per report while retaining two simulation records and two private +simulation artifacts. `segmented=false`, US subnational requests, UK requests, +cliff analysis, and labor-supply-response calculations use the +single-container Stage 12 path. + +After dispatching both logical simulations, the coordinator checks both Modal +calls with nonblocking result probes on every polling pass. It does not wait +for baseline before checking reform. Both simulations share one 3,900-second +wait window, the longest simulation worker timeout, that opens when the +coordinator starts, so coordinator setup comes out of that window. The first +failed simulation supplies the report's `error_code` and safe `error_summary`. +When a simulation fails, when the window closes, or when the coordinator +itself fails after starting a simulation, the coordinator records every +started, unfinished simulation as `failed` with `simulation_cancelled` and a +summary that gives the reason, then stops it. Only a `running` record changes, +so a simulation that already succeeded or failed keeps its outcome. An +unsegmented simulation is cancelled through Modal. A segmented simulation +reads its cancelled record, cancels the region-group calls it owns, and exits +without overwriting the record; it is cancelled through Modal only if that +record could not be written. An error while reading or updating a simulation +record is logged with its identifiers and never replaces the failure that +caused it. Expected input failures retain their typed code and message. +Unexpected exceptions are reported once, in full, and persisted as their +redacted cause (type, stage, simulation role and redacted message) followed +by the same correlation ID. + +Stage 12 deployments run `src/modal/stage12_precompute_app.py` against the +dedicated `STAGE12_CACHE_BUCKET`. It publishes three annual datasets plus 60 +current-law US baseline files and emits +`STAGE12_CACHE_MANIFEST_DIGEST=`. Baseline files have the names .py +saves and loads them under (`Simulation.storage_id`, which includes the SPM +selection's digest). Stored files are never replaced: a run computes only +missing files, takes no `--force`, and publishes the manifest only once a +determinism check of the planned baselines has a recorded pass. Like the v2 +workers, its image and functions mount only the read-only +`pe-uk-private-hf-read-token` Hugging Face secret, which supplies +`HUGGING_FACE_TOKEN`, and +`stage12-evaluation-gcp-credentials`; the US data it reads comes from the +public `policyengine/populace-us` dataset repository. The US worker image +downloads and verifies the manifest's 63 files during image construction. The +release validation checks the recorded manifest, hashes a baseline file, and +requires a real calculation to load that file without recomputing it. Worker +validation also calculates a US household under a nonempty CTC reform using +the frontend interval format `2026-01-01.2100-12-31`. The cache manifest +digest is then stored in the separate v2 version manifest. + +The cache is an optimization, not required state for a calculation: an absent +file causes a normal calculation, and a file missing planned output columns is +recalculated. Only empty-policy baseline region-group calculations against the +bundle's default dataset are eligible to read it. Reform calculations never +read a baseline file. + ## Observability `policyengine-observability` 3.x emits structured request, operation, diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py b/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py new file mode 100644 index 000000000..5e69a2c6c --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/stage12_cache_image.py @@ -0,0 +1,100 @@ +"""Self-contained Stage 12 cache fetch used during Modal image construction.""" + +from __future__ import annotations + + +def fetch_stage12_cache(bucket: str, manifest: dict | None, *, client=None) -> None: + """Download and verify every manifest-listed cache file into the image.""" + + import json + import os + from hashlib import sha256 + from pathlib import Path + + if not bucket or not manifest: + raise RuntimeError("Stage 12 cache deployment has no bucket or manifest") + if manifest.get("schema") != "stage12-mf1": + raise RuntimeError("Stage 12 cache manifest schema is unsupported") + bundle = manifest.get("bundle") + if not isinstance(bundle, dict) or bundle.get( + "bundle_manifest_sha256" + ) != os.getenv("STAGE12_BUNDLE_MANIFEST_SHA256"): + raise RuntimeError("Stage 12 cache manifest names another installed bundle") + if manifest.get("years") != [2026, 2027, 2025]: + raise RuntimeError("Stage 12 cache manifest has an unexpected year set") + artifacts = manifest.get("artifacts") + if not isinstance(artifacts, list): + raise RuntimeError("Stage 12 cache manifest has no artifact list") + type_counts = { + kind: sum( + item.get("type") == kind for item in artifacts if isinstance(item, dict) + ) + for kind in ("dataset", "baseline") + } + if type_counts != {"dataset": 3, "baseline": 60}: + raise RuntimeError("Stage 12 cache manifest has an unexpected artifact count") + canonical = json.dumps( + manifest, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + manifest_digest = sha256(canonical).hexdigest() + if manifest_digest != os.getenv("STAGE12_CACHE_MANIFEST_DIGEST"): + raise RuntimeError("Stage 12 cache manifest digest differs from deployment") + + if client is None: + credentials = {} + raw_credentials = os.getenv("GOOGLE_APPLICATION_CREDENTIALS_JSON") + if raw_credentials: + decoded = json.loads(raw_credentials) + info = json.loads(decoded) if isinstance(decoded, str) else decoded + from google.oauth2 import service_account + + credentials = { + "credentials": service_account.Credentials.from_service_account_info( + info + ), + "project": info.get("project_id"), + } + from google.cloud import storage + + client = storage.Client(**credentials) + + data_dir = Path(os.environ["POLICYENGINE_DATA_FOLDER"]) + data_dir.mkdir(parents=True, exist_ok=True) + filenames: set[str] = set() + for artifact in artifacts: + if not isinstance(artifact, dict): + raise RuntimeError("Stage 12 cache artifact entry is invalid") + filename = artifact.get("filename") + path = artifact.get("path") + expected_digest = artifact.get("content_sha256") + expected_size = artifact.get("size_bytes") + if ( + not isinstance(filename, str) + or Path(filename).name != filename + or filename in filenames + or not isinstance(path, str) + or not isinstance(expected_digest, str) + or not isinstance(expected_size, int) + ): + raise RuntimeError("Stage 12 cache artifact metadata is invalid") + filenames.add(filename) + destination = data_dir / filename + partial = destination.with_name(filename + ".partial") + try: + client.bucket(bucket).blob(path).download_to_filename(str(partial)) + digest = sha256() + with partial.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + if partial.stat().st_size != expected_size: + raise RuntimeError("Stage 12 cache artifact size differs") + if digest.hexdigest() != expected_digest: + raise RuntimeError("Stage 12 cache artifact digest differs") + partial.replace(destination) + finally: + partial.unlink(missing_ok=True) + (data_dir / ".stage12-cache-manifest.json").write_bytes(canonical) diff --git a/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py new file mode 100644 index 000000000..f78cce17e --- /dev/null +++ b/projects/policyengine-simulation-executor/src/modal/stage12_precompute_app.py @@ -0,0 +1,178 @@ +# pyright: reportAttributeAccessIssue=false +"""Ephemeral Modal application that builds the independent Stage 12 cache.""" + +from __future__ import annotations + +import shlex +from pathlib import Path + +from policyengine_simulation_executor.hf_access_validation import ( + HF_MODAL_SECRET_NAME, + HF_RUNTIME_ENV_NAME, +) +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle + +import modal + +STAGE12_DATA_DIR = "/opt/policyengine/stage12-data" +_UV_PROJECT_DIR = str(Path(__file__).resolve().parents[2]) if modal.is_local() else "." +RESOLVED_BUNDLE = load_stage12_bundle() + +app = modal.App("policyengine-stage12-cache-precompute") +gcp_secret = modal.Secret.from_name("stage12-evaluation-gcp-credentials") +# The read-only Hugging Face credential the other simulation apps mount. The +# US cache reads only the public policyengine/populace-us dataset repository. +hf_secret = modal.Secret.from_name( + HF_MODAL_SECRET_NAME, + required_keys=[HF_RUNTIME_ENV_NAME], +) +worker_secrets = [gcp_secret, hf_secret] + + +def _bundle_install_command() -> str: + parts = [ + "uvx", + "--from", + RESOLVED_BUNDLE.bundle.policyengine_requirement, + "policyengine", + "bundle", + "install", + RESOLVED_BUNDLE.bundle.policyengine_version, + "--venv", + "/.uv/.venv", + "--country", + "us", + "--data-dir", + STAGE12_DATA_DIR, + "--yes", + ] + return " ".join(shlex.quote(part) for part in parts) + + +precompute_image = ( + modal.Image.debian_slim(python_version="3.13") + .uv_sync( + uv_project_dir=_UV_PROJECT_DIR, + frozen=True, + extra_options="--only-group modal-simulation-image", + ) + .run_commands(_bundle_install_command(), secrets=[hf_secret]) + .env({"POLICYENGINE_DATA_FOLDER": STAGE12_DATA_DIR}) + .add_local_python_source( + "src.modal", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_simulation_contract", + "policyengine_stage12_persistence", + copy=True, + ) +) + + +@app.function( + image=precompute_image, + cpu=4.0, + memory=16384, + timeout=1800, + secrets=worker_secrets, +) +def plan_artifacts(bucket: str) -> dict: + from policyengine_simulation_executor.stage12_cache.precompute import ( + plan_artifacts_impl, + ) + + return plan_artifacts_impl(bucket).model_dump() + + +@app.function( + image=precompute_image, + cpu=8.0, + memory=65536, + timeout=7200, + secrets=worker_secrets, +) +def build_dataset(bucket: str, expected: dict) -> dict: + from policyengine_simulation_executor.stage12_cache.models import DatasetPlanEntry + from policyengine_simulation_executor.stage12_cache.precompute import ( + build_dataset_impl, + ) + + return build_dataset_impl( + bucket, DatasetPlanEntry.model_validate(expected) + ).model_dump() + + +@app.function( + image=precompute_image, + cpu=8.0, + memory=32768, + timeout=3600, + secrets=worker_secrets, +) +def compute_baseline(bucket: str, expected: dict) -> dict: + from policyengine_simulation_executor.stage12_cache.models import BaselinePlanEntry + from policyengine_simulation_executor.stage12_cache.precompute import ( + compute_baseline_impl, + ) + + return compute_baseline_impl( + bucket, BaselinePlanEntry.model_validate(expected) + ).model_dump() + + +@app.function( + image=precompute_image, + cpu=8.0, + memory=32768, + timeout=3600, + secrets=worker_secrets, +) +def verify_determinism(bucket: str, plan: dict, expected: dict) -> dict: + from policyengine_simulation_executor.stage12_cache.models import ( + BaselinePlanEntry, + PrecomputePlan, + ) + from policyengine_simulation_executor.stage12_cache.precompute import ( + verify_determinism_impl, + ) + + return verify_determinism_impl( + bucket, + PrecomputePlan.model_validate(plan), + BaselinePlanEntry.model_validate(expected), + ).model_dump() + + +@app.function( + image=precompute_image, + cpu=2.0, + memory=8192, + timeout=600, + secrets=worker_secrets, +) +def publish_manifest(bucket: str, plan: dict) -> str: + from policyengine_simulation_executor.stage12_cache.models import PrecomputePlan + from policyengine_simulation_executor.stage12_cache.precompute import ( + publish_manifest_impl, + ) + + return publish_manifest_impl(bucket, PrecomputePlan.model_validate(plan)) + + +@app.local_entrypoint() +def main(): + from policyengine_simulation_executor.stage12_cache.precompute import ( + run_precompute, + ) + from policyengine_simulation_executor.stage12_cache.store import ( + resolve_cache_bucket, + ) + + run_precompute( + resolve_cache_bucket(), + plan_artifacts=plan_artifacts, + build_dataset=build_dataset, + compute_baseline=compute_baseline, + verify_determinism=verify_determinism, + publish_manifest=publish_manifest, + ) diff --git a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py index 4717bb628..7bf387f65 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/update_v2_version_manifest.py @@ -37,6 +37,7 @@ def build_worker(validation_payload: object) -> V2WorkerVersion: bundle=resolved.bundle, bundle_manifest_sha256=resolved.bundle_manifest_sha256, validation=validation, + cache=validation.cache, ) validate_manifest_has_no_credentials(worker.model_dump(mode="json")) return worker diff --git a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py index fc088bd2e..8c2a9dd8a 100644 --- a/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py +++ b/projects/policyengine-simulation-executor/src/modal/utils/validate_v2_worker_release.py @@ -13,6 +13,7 @@ from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_manifest import ( ManifestText, + Stage12CacheRelease, V2WorkerValidation, v2_application_name, ) @@ -38,10 +39,13 @@ def _spawn_validation( return result, invocation_id -def validate_release(*, environment: str) -> V2WorkerValidation: +def validate_release( + *, environment: str, expected_cache_manifest_sha256: str +) -> V2WorkerValidation: resolved = load_stage12_bundle() application_name = v2_application_name(resolved.bundle.policyengine_version) invocation_ids: dict[CountryId, ManifestText] = {} + cache_release: Stage12CacheRelease | None = None for country_bundle in resolved.bundle.countries: function_name = f"validate_worker_{country_bundle.country}" result, invocation_id = _spawn_validation( @@ -70,6 +74,13 @@ def validate_release(*, environment: str) -> V2WorkerValidation: raise RuntimeError( f"{function_name} reported bundle values that differ from deployment" ) + if country_bundle.country == "us": + reported_cache = result.get("cache") + cache_release = Stage12CacheRelease.model_validate(reported_cache) + if cache_release.manifest_sha256 != expected_cache_manifest_sha256: + raise RuntimeError(f"{function_name} reported another cache manifest") + elif result.get("cache") is not None: + raise RuntimeError(f"{function_name} unexpectedly reported a cache") invocation_ids[country_bundle.country] = invocation_id return V2WorkerValidation( validated=True, @@ -78,6 +89,7 @@ def validate_release(*, environment: str) -> V2WorkerValidation: validated_at=datetime.now(timezone.utc), validation_invocation_id=f"stage12-validation-{uuid4()}", country_validation_invocation_ids=invocation_ids, + cache=cache_release, ) @@ -87,8 +99,12 @@ def main() -> None: ) parser.add_argument("--environment", required=True) parser.add_argument("--output", required=True, type=Path) + parser.add_argument("--cache-manifest-sha256", required=True) args = parser.parse_args() - validation = validate_release(environment=args.environment) + validation = validate_release( + environment=args.environment, + expected_cache_manifest_sha256=args.cache_manifest_sha256, + ) args.output.write_text( json.dumps( validation.model_dump(mode="json"), diff --git a/projects/policyengine-simulation-executor/src/modal/v2_app.py b/projects/policyengine-simulation-executor/src/modal/v2_app.py index 076f9dce0..ea09f3c26 100644 --- a/projects/policyengine-simulation-executor/src/modal/v2_app.py +++ b/projects/policyengine-simulation-executor/src/modal/v2_app.py @@ -22,12 +22,18 @@ assertion_values, load_stage12_bundle, ) +from policyengine_simulation_executor.stage12_timeouts import ( + COORDINATOR_TIMEOUT_SECONDS, + UK_SIMULATION_TIMEOUT_SECONDS, + US_SIMULATION_TIMEOUT_SECONDS, +) from policyengine_simulation_observability.observability import ( init_process_observability, modal_image_environment, ) from policyengine_simulation_observability.stages import ( STAGE12_CANONICAL_REPORT_STAGES, + STAGE12_SEGMENT_STAGES, STAGE12_SHADOW_REPORT_STAGES, STAGE12_SIMULATION_STAGES, Stage, @@ -36,6 +42,7 @@ normalize_observability_context, ) from src.modal.bundle_data import bundle_data_install_command +from src.modal.stage12_cache_image import fetch_stage12_cache from src.modal.static_runtime_files import add_static_runtime_files import modal @@ -47,6 +54,30 @@ APP_NAME = v2_application_name(RESOLVED_BUNDLE.bundle.policyengine_version) +def _deploy_time_cache_inputs() -> tuple[str, dict | None]: + """Resolve and validate the cache manifest on the deploying machine.""" + + if not modal.is_local(): + return "", None + bucket = os.environ.get("STAGE12_CACHE_BUCKET", "") + digest = os.environ.get("STAGE12_CACHE_MANIFEST_DIGEST") + if not digest: + return bucket, None + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + from policyengine_simulation_executor.stage12_cache.store import Stage12CacheStore + + payload = Stage12CacheStore(bucket).read_manifest(digest) + if payload is None: + raise RuntimeError("Stage 12 cache manifest is unavailable") + manifest = CacheManifest.model_validate(payload) + if manifest.bundle.bundle_manifest_sha256 != RESOLVED_BUNDLE.bundle_manifest_sha256: + raise RuntimeError("Stage 12 cache manifest names another bundle") + return bucket, manifest.canonical_payload() + + +STAGE12_CACHE_BUCKET, STAGE12_CACHE_MANIFEST = _deploy_time_cache_inputs() + + def _external_assertions(environment: dict[str, str]) -> dict[str, str]: mapping = { "STAGE12_EXPECT_POLICYENGINE_VERSION": "policyengine_version", @@ -90,7 +121,9 @@ def _country_bundle(country: CountryId): ) -def build_v2_image(countries: tuple[CountryId, ...]) -> modal.Image: +def build_v2_image( + countries: tuple[CountryId, ...], *, include_us_cache: bool = False +) -> modal.Image: country_values = { country: _country_bundle(country).model_dump(mode="json") for country in countries @@ -122,21 +155,33 @@ def build_v2_image(countries: tuple[CountryId, ...]) -> modal.Image: sort_keys=True, separators=(",", ":"), ), + "STAGE12_CACHE_MANIFEST_DIGEST": os.environ.get( + "STAGE12_CACHE_MANIFEST_DIGEST", "" + ), } ) - .add_local_python_source( - "src.modal", - "policyengine_simulation_executor", - "policyengine_simulation_observability", - "policyengine_simulation_contract", - "policyengine_stage12_persistence", - copy=True, + ) + if include_us_cache: + image = image.run_function( + fetch_stage12_cache, + args=(STAGE12_CACHE_BUCKET, STAGE12_CACHE_MANIFEST), + secrets=[gcp_secret], + cpu=2.0, + memory=4096, + timeout=1800, ) + image = image.add_local_python_source( + "src.modal", + "policyengine_simulation_executor", + "policyengine_simulation_observability", + "policyengine_simulation_contract", + "policyengine_stage12_persistence", + copy=True, ) return add_static_runtime_files(image, uv_project_dir=_UV_PROJECT_DIR) -us_worker_image = build_v2_image(("us",)) +us_worker_image = build_v2_image(("us",), include_us_cache=True) uk_worker_image = build_v2_image(("uk",)) coordinator_image = build_v2_image(("us", "uk")) @@ -149,6 +194,9 @@ def _validate(country: CountryId) -> dict: return validate_country_worker( country=country, expected_bundle_manifest_sha256=os.environ["STAGE12_BUNDLE_MANIFEST_SHA256"], + expected_cache_manifest_sha256=( + os.environ.get("STAGE12_CACHE_MANIFEST_DIGEST") if country == "us" else None + ), ) @@ -180,7 +228,52 @@ def validate_worker_uk() -> dict: image=us_worker_image, cpu=8.0, memory=32768, - timeout=3000, + timeout=3600, + retries=0, + max_containers=300, + secrets=worker_secrets, +) +def run_single_simulation_segment_us( + payload: dict, + *, + observability_context: dict | None = None, +) -> dict: + from policyengine_simulation_executor.stage12_runtime import calculate_segment + + runtime = init_process_observability( + service_name="policyengine-stage12-us-segment-worker", + service_role="stage12_segment_worker", + platform="modal", + environment=os.getenv("MODAL_ENVIRONMENT", "local"), + ) + + try: + propagated = normalize_observability_context(observability_context) + result = None + with runtime.operation( + STAGE12_SEGMENT_STAGES.name(Stage.STAGE12_SEGMENT_EXECUTION), + attributes={ + "runner_name": "stage12", + "simulation_role": payload.get("simulation", {}).get("role"), + "segment_index": payload.get("segment_index"), + }, + remote_context=propagated, + ): + result = calculate_segment(payload, runtime=runtime).model_dump( + mode="python" + ) + if result is None: + raise RuntimeError("Stage 12 segment calculation returned no result") + return result + finally: + runtime.shutdown() + + +@app.function( + image=us_worker_image, + cpu=8.0, + memory=32768, + timeout=US_SIMULATION_TIMEOUT_SECONDS, retries=0, max_containers=10, secrets=worker_secrets, @@ -192,6 +285,7 @@ def run_single_simulation_us( observability_context: dict | None = None, ) -> dict: from policyengine_simulation_executor.stage12_runtime import ( + run_segmented_simulation, run_single_simulation, ) @@ -201,8 +295,18 @@ def run_single_simulation_us( platform="modal", environment=os.getenv("MODAL_ENVIRONMENT", "local"), ) + + def segmented_calculator(simulation, cancellation_requested): + return run_segmented_simulation( + simulation, + app_name=APP_NAME, + runtime=runtime, + cancellation_requested=cancellation_requested, + ) + try: propagated = normalize_observability_context(observability_context) + result = None with runtime.operation( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), attributes={ @@ -211,12 +315,16 @@ def run_single_simulation_us( }, remote_context=propagated, ): - return run_single_simulation( + result = run_single_simulation( payload, context, required_country="us", + segmented_calculator=segmented_calculator, runtime=runtime, ) + if result is None: + raise RuntimeError("Stage 12 US simulation returned no result") + return result finally: runtime.shutdown() @@ -225,7 +333,7 @@ def run_single_simulation_us( image=uk_worker_image, cpu=8.0, memory=32768, - timeout=3000, + timeout=UK_SIMULATION_TIMEOUT_SECONDS, retries=0, max_containers=10, secrets=worker_secrets, @@ -248,6 +356,7 @@ def run_single_simulation_uk( ) try: propagated = normalize_observability_context(observability_context) + result = None with runtime.operation( STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SIMULATION_EXECUTION), attributes={ @@ -256,12 +365,15 @@ def run_single_simulation_uk( }, remote_context=propagated, ): - return run_single_simulation( + result = run_single_simulation( payload, context, required_country="uk", runtime=runtime, ) + if result is None: + raise RuntimeError("Stage 12 UK simulation returned no result") + return result finally: runtime.shutdown() @@ -270,9 +382,9 @@ def run_single_simulation_uk( image=coordinator_image, cpu=2.0, memory=8192, - # Allow one 50-minute child-calculation window, the subsequent 15-minute - # production-result wait, and bounded aggregation/persistence overhead. - timeout=4500, + # The shared child-wait window (the longest simulation worker timeout), + # the production-result wait, and time to finish and record the report. + timeout=COORDINATOR_TIMEOUT_SECONDS, retries=0, max_containers=10, secrets=worker_secrets, @@ -304,6 +416,7 @@ def coordinate_report( ) try: propagated = normalize_observability_context(observability_context) + result = None with runtime.operation( stage_plan.name(Stage.STAGE12_COORDINATOR_EXECUTION), attributes={ @@ -316,7 +429,7 @@ def coordinate_report( }, remote_context=propagated, ): - return run_report_coordinator( + result = run_report_coordinator( payload, context, parent, @@ -324,5 +437,8 @@ def coordinate_report( coordinator_invocation_id=coordinator_invocation_id, runtime=runtime, ) + if result is None: + raise RuntimeError("Stage 12 report coordinator returned no result") + return result finally: runtime.shutdown() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/__init__.py new file mode 100644 index 000000000..b68b75e70 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/__init__.py @@ -0,0 +1,27 @@ +"""Independent cache identities and runtime loading for Stage 12.""" + +from .keys import ( + Stage12BaselineCacheIdentity, + Stage12DatasetCacheIdentity, + collect_baseline_identity, + collect_dataset_identity, +) +from .runtime import ( + CACHE_HIT, + CACHE_INCOMPLETE, + CACHE_MISS, + Stage12CachedSimulation, + qualifying_cache_identity, +) + +__all__ = [ + "CACHE_HIT", + "CACHE_INCOMPLETE", + "CACHE_MISS", + "Stage12BaselineCacheIdentity", + "Stage12CachedSimulation", + "Stage12DatasetCacheIdentity", + "collect_baseline_identity", + "collect_dataset_identity", + "qualifying_cache_identity", +] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py new file mode 100644 index 000000000..b5535b406 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/keys.py @@ -0,0 +1,193 @@ +"""Content-derived Stage 12 cache keys and object paths.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping +from dataclasses import dataclass +from hashlib import sha256 +from typing import Any + +from policyengine_simulation_contract.stage12_bundle import ( + CountryId, + ResolvedStage12Bundle, +) +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + stage12_output_plan_sha256, +) + +DATASET_SCHEMA = "stage12-ds1" +BASELINE_SCHEMA = "stage12-bl1" +CURRENT_LAW_POLICY = "current-law" +_ID_DIGEST_CHARS = 16 + + +def canonical_digest(payload: Mapping[str, Any]) -> str: + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + return sha256(encoded).hexdigest() + + +@dataclass(frozen=True) +class Stage12DatasetCacheIdentity: + country: CountryId + dataset: str + stem: str + year: int + source_sha256: str + data_artifact_revision: str + data_package_version: str + country_package_version: str + core_package_version: str + policyengine_version: str + bundle_manifest_sha256: str + + @property + def payload(self) -> dict[str, Any]: + return { + "schema": DATASET_SCHEMA, + "country": self.country, + "dataset": self.dataset, + "year": self.year, + "source_sha256": self.source_sha256, + "data_artifact_revision": self.data_artifact_revision, + "data_package_version": self.data_package_version, + "country_package_version": self.country_package_version, + "core_package_version": self.core_package_version, + "policyengine_version": self.policyengine_version, + "bundle_manifest_sha256": self.bundle_manifest_sha256, + } + + @property + def digest(self) -> str: + return canonical_digest(self.payload) + + @property + def filename(self) -> str: + return f"{self.stem}_year_{self.year}.h5" + + @property + def store_path(self) -> str: + return f"datasets/{self.country}/{self.digest}/{self.filename}" + + +@dataclass(frozen=True) +class Stage12BaselineCacheIdentity: + dataset: Stage12DatasetCacheIdentity + region_codes: tuple[str, ...] + scope_key: str + output_plan_sha256: str + spm: Mapping[str, Any] | None = None + + @property + def payload(self) -> dict[str, Any]: + return { + "schema": BASELINE_SCHEMA, + "country": self.dataset.country, + "dataset_digest": self.dataset.digest, + "region_codes": list(self.region_codes), + "scope_key": self.scope_key, + "output_plan_sha256": self.output_plan_sha256, + "spm": dict(self.spm) if self.spm is not None else None, + "policy": CURRENT_LAW_POLICY, + } + + @property + def digest(self) -> str: + return canonical_digest(self.payload) + + @property + def simulation_id(self) -> str: + return f"{BASELINE_SCHEMA}-{self.digest[:_ID_DIGEST_CHARS]}" + + @property + def storage_id(self) -> str: + """The name .py saves and loads this baseline under. + + Mirrors ``policyengine.core.Simulation.storage_id``: a US simulation + always has a resolved SPM selection, so .py appends its digest. + """ + + if self.spm is None: + return self.simulation_id + return f"{self.simulation_id}-spm-{canonical_digest(self.spm)}" + + @property + def filename(self) -> str: + return f"{self.storage_id}.h5" + + @property + def store_path(self) -> str: + return f"baselines/{self.dataset.country}/{self.digest}/{self.filename}" + + +def collect_dataset_identity( + resolved: ResolvedStage12Bundle, + *, + country: CountryId, + year: int, +) -> Stage12DatasetCacheIdentity: + """Collect dataset identity only from the installed Stage 12 bundle.""" + + from policyengine.provenance.manifest import ( + dataset_logical_name, + resolve_dataset_reference, + ) + + country_bundle = next( + item for item in resolved.bundle.countries if item.country == country + ) + dataset = next( + item + for item in country_bundle.datasets + if item.identity == country_bundle.default_dataset + ) + stem = dataset_logical_name( + resolve_dataset_reference(country, country_bundle.default_dataset) + ) + return Stage12DatasetCacheIdentity( + country=country, + dataset=dataset.identity, + stem=stem, + year=year, + source_sha256=dataset.sha256, + data_artifact_revision=dataset.artifact_revision, + data_package_version=country_bundle.data_package_version, + country_package_version=country_bundle.country_package_version, + core_package_version=resolved.bundle.core_package_version, + policyengine_version=resolved.bundle.policyengine_version, + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + ) + + +def collect_baseline_identity( + simulation: PlannedSimulationExecutionInput, + *, + resolved: ResolvedStage12Bundle, + region_codes: tuple[str, ...], + scope_key: str, + spm: Mapping[str, Any] | None, +) -> Stage12BaselineCacheIdentity: + """Build the identity shared by the Stage 12 writer and reader. + + ``spm`` is the resolved selection the simulation is built with, so the + identity's file name matches the one .py saves and loads. + """ + + return Stage12BaselineCacheIdentity( + dataset=collect_dataset_identity( + resolved, + country=simulation.geography.country, + year=simulation.year, + ), + region_codes=region_codes, + scope_key=scope_key, + output_plan_sha256=stage12_output_plan_sha256(simulation.output_plan), + spm=spm, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py new file mode 100644 index 000000000..e13591fca --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/models.py @@ -0,0 +1,114 @@ +"""Strict models crossing Stage 12 precompute function boundaries.""" + +from __future__ import annotations + +from typing import Annotated, Any, Literal, Protocol + +from pydantic import BaseModel, ConfigDict, Field + + +class StrictCacheModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class CacheBundleIdentity(StrictCacheModel): + bundle_manifest_sha256: str + policyengine_version: str + core_package_version: str + country_package_version: str + data_package_version: str + data_artifact_revision: str + default_dataset: str + + +class DatasetPlanEntry(StrictCacheModel): + year: int + digest: str + path: str + filename: str + exists: bool + + +class BaselinePlanEntry(StrictCacheModel): + year: int + segment_index: Annotated[int, Field(ge=0, lt=20)] + region_codes: Annotated[tuple[str, ...], Field(min_length=1)] + digest: str + path: str + filename: str + simulation_id: str + output_plan_sha256: str + exists: bool + + +class PrecomputePlan(StrictCacheModel): + schema_version: Literal[1] = 1 + years: tuple[int, ...] + datasets: tuple[DatasetPlanEntry, ...] + baselines: tuple[BaselinePlanEntry, ...] + bundle: CacheBundleIdentity + partition_sha256: str + determinism_verified: bool + + +class WorkSelection(StrictCacheModel): + datasets: tuple[DatasetPlanEntry, ...] + baselines: tuple[BaselinePlanEntry, ...] + + +class CacheObject(StrictCacheModel): + path: str + content_sha256: str + size_bytes: Annotated[int, Field(ge=0)] + + +class DatasetBuildResult(CacheObject): + year: int + uploaded: bool + build_seconds: float + + +class BaselineBuildResult(CacheObject): + year: int + segment_index: int + simulation_id: str + cache_outcome: Literal["hit", "incomplete", "miss"] | None + uploaded: bool + compute_seconds: float + + +class DeterminismVerdict(StrictCacheModel): + equal: bool + differences: tuple[str, ...] + + +class ManifestArtifact(StrictCacheModel): + type: Literal["dataset", "baseline"] + path: str + filename: str + year: int + identity_digest: str + content_sha256: str + size_bytes: int + + +class CacheManifest(StrictCacheModel): + manifest_schema: Literal["stage12-mf1"] = Field( + default="stage12-mf1", + alias="schema", + serialization_alias="schema", + ) + country: Literal["us"] = "us" + years: tuple[int, ...] + partition_sha256: str + bundle: CacheBundleIdentity + artifacts: tuple[ManifestArtifact, ...] + + def canonical_payload(self) -> dict[str, Any]: + return self.model_dump(mode="json", by_alias=True) + + +class RemoteFunction(Protocol): + def remote(self, *args: Any) -> Any: ... + + def spawn(self, *args: Any) -> Any: ... diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py new file mode 100644 index 000000000..dca28a9d4 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/precompute.py @@ -0,0 +1,460 @@ +"""Build and publish the independent Stage 12 US cache.""" + +from __future__ import annotations + +import time +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any +from uuid import NAMESPACE_URL, uuid5 + +from policyengine_simulation_contract.stage12_execution import ( + BundleProvenance, + DatasetArtifactMediaType, + DatasetArtifactReference, + DatasetPopulationInput, + DatasetProvenance, + GeographySelection, + PlannedSimulationExecutionInput, + ReportAggregate, + ReportExecutionInput, + SimulationExecutionInput, + SimulationRole, + stage12_output_plan_sha256, +) + +from policyengine_simulation_executor.stage12_bundle import load_stage12_bundle +from policyengine_simulation_executor.stage12_runtime.output_planning import ( + plan_simulation_input, + resolve_report_output_plan, +) +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, +) + +from .keys import ( + Stage12BaselineCacheIdentity, + canonical_digest, + collect_baseline_identity, + collect_dataset_identity, +) +from .models import ( + BaselineBuildResult, + BaselinePlanEntry, + CacheBundleIdentity, + CacheManifest, + DatasetBuildResult, + DatasetPlanEntry, + DeterminismVerdict, + ManifestArtifact, + PrecomputePlan, + RemoteFunction, + WorkSelection, +) +from .store import Stage12CacheStore + +PRECOMPUTE_YEARS = (2026, 2027, 2025) +MANIFEST_DIGEST_PREFIX = "STAGE12_CACHE_MANIFEST_DIGEST=" +DETERMINISM_SCHEMA = "stage12-dv1" + + +def manifest_digest_line(digest: str) -> str: + return f"{MANIFEST_DIGEST_PREFIX}{digest}" + + +def _bundle_identity() -> CacheBundleIdentity: + resolved = load_stage12_bundle() + country = next(item for item in resolved.bundle.countries if item.country == "us") + return CacheBundleIdentity( + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + policyengine_version=resolved.bundle.policyengine_version, + core_package_version=resolved.bundle.core_package_version, + country_package_version=country.country_package_version, + data_package_version=country.data_package_version, + data_artifact_revision=country.data_artifact_revision, + default_dataset=country.default_dataset, + ) + + +def _planned_pair(year: int): + resolved = load_stage12_bundle() + country = next(item for item in resolved.bundle.countries if item.country == "us") + dataset = next( + item for item in country.datasets if item.identity == country.default_dataset + ) + provenance = BundleProvenance( + policyengine_version=resolved.bundle.policyengine_version, + country_package_name=country.country_package_name, + country_package_version=country.country_package_version, + dataset=DatasetProvenance( + identity=dataset.identity, + uri=dataset.uri, + artifact_revision=dataset.artifact_revision, + data_package_name=country.data_package_name, + data_package_version=country.data_package_version, + ), + bundle_manifest_sha256=resolved.bundle_manifest_sha256, + ) + population = DatasetPopulationInput( + artifact=DatasetArtifactReference( + uri=dataset.uri, + media_type=DatasetArtifactMediaType.HDF5, + content_sha256=dataset.sha256, + ) + ) + evaluation_id = uuid5(NAMESPACE_URL, f"stage12-cache:{year}") + + def simulation(role: SimulationRole) -> SimulationExecutionInput: + return SimulationExecutionInput( + evaluation_id=evaluation_id, + simulation_execution_id=uuid5(evaluation_id, role.value), + role=role, + policy={}, + population=population, + year=year, + geography=GeographySelection(country="us", region="us"), + bundle=provenance, + ) + + report = ReportExecutionInput( + evaluation_id=evaluation_id, + baseline=simulation(SimulationRole.BASELINE), + reform=simulation(SimulationRole.REFORM), + requested_aggregates=tuple(ReportAggregate), + ) + plan = resolve_report_output_plan(report) + return ( + plan_simulation_input(report.baseline, plan), + plan_simulation_input(report.reform, plan), + ) + + +def _scoped_planned_simulation( + year: int, group: tuple[str, ...], *, reform: bool = False +): + baseline, reform_input = _planned_pair(year) + selected = reform_input if reform else baseline + return selected.model_copy(update={"options": {"region_group": list(group)}}) + + +def _scope_key(year: int, group: tuple[str, ...]) -> str: + from policyengine_simulation_executor.simulation_runtime import ( + _country_module, + _resolve_region, + ) + + params = { + "country": "us", + "scope": "macro", + "time_period": str(year), + "region_group": list(group), + } + resolution = _resolve_region( + country_module=_country_module("us"), + country="us", + params=params, + ) + value = getattr(resolution.scoping_strategy, "cache_key", None) + if not isinstance(value, str) or not value: + raise RuntimeError("Stage 12 region group produced no cache identity") + return value + + +def _baseline_identity( + year: int, group: tuple[str, ...] +) -> tuple[PlannedSimulationExecutionInput, Stage12BaselineCacheIdentity]: + """Return a planned baseline and the identity its worker reads it under.""" + + from policyengine_simulation_executor.spm import normalize_runtime_spm + from policyengine_simulation_executor.stage12_runtime.simulation import ( + simulation_runtime_params, + ) + + simulation = _scoped_planned_simulation(year, group) + identity = collect_baseline_identity( + simulation, + resolved=load_stage12_bundle(), + region_codes=group, + scope_key=_scope_key(year, group), + # The selection the worker builds the simulation with, which names + # the file .py saves and loads. + spm=normalize_runtime_spm(simulation_runtime_params(simulation)), + ) + return simulation, identity + + +def determinism_record(baselines: Sequence[BaselinePlanEntry]) -> dict[str, Any]: + """Return the record a passing determinism check writes for these baselines.""" + + return { + "schema": DETERMINISM_SCHEMA, + "baselines": [entry.digest for entry in baselines], + } + + +def determinism_record_path(baselines: Sequence[BaselinePlanEntry]) -> str: + return f"determinism/us/{canonical_digest(determinism_record(baselines))}.json" + + +def plan_artifacts_impl(bucket: str) -> PrecomputePlan: + resolved = load_stage12_bundle() + store = Stage12CacheStore(bucket) + datasets = [] + baselines = [] + for year in PRECOMPUTE_YEARS: + dataset = collect_dataset_identity(resolved, country="us", year=year) + datasets.append( + DatasetPlanEntry( + year=year, + digest=dataset.digest, + path=dataset.store_path, + filename=dataset.filename, + exists=store.exists(dataset.store_path), + ) + ) + for index, group in enumerate(US_REGION_GROUPS): + simulation, identity = _baseline_identity(year, group) + baselines.append( + BaselinePlanEntry( + year=year, + segment_index=index, + region_codes=group, + digest=identity.digest, + path=identity.store_path, + filename=identity.filename, + simulation_id=identity.simulation_id, + output_plan_sha256=stage12_output_plan_sha256( + simulation.output_plan + ), + exists=store.exists(identity.store_path), + ) + ) + return PrecomputePlan( + years=PRECOMPUTE_YEARS, + datasets=tuple(datasets), + baselines=tuple(baselines), + bundle=_bundle_identity(), + partition_sha256=canonical_digest( + {"groups": [list(group) for group in US_REGION_GROUPS]} + ), + determinism_verified=store.exists(determinism_record_path(baselines)), + ) + + +def select_work(plan: PrecomputePlan) -> WorkSelection: + """Select only missing files: a stored cache file is never replaced.""" + + return WorkSelection( + datasets=tuple(entry for entry in plan.datasets if not entry.exists), + baselines=tuple(entry for entry in plan.baselines if not entry.exists), + ) + + +def build_dataset_impl(bucket: str, expected: DatasetPlanEntry) -> DatasetBuildResult: + from policyengine_simulation_executor.simulation_runtime import ( + _country_module, + resolve_data_folder, + ) + + resolved = load_stage12_bundle() + identity = collect_dataset_identity(resolved, country="us", year=expected.year) + if identity.store_path != expected.path: + raise RuntimeError("planned and runtime Stage 12 dataset identities differ") + started = time.monotonic() + _country_module("us").ensure_datasets( + datasets=[identity.dataset], + years=[expected.year], + data_folder=resolve_data_folder(), + ) + local_path = Path(resolve_data_folder()) / identity.filename + uploaded, stored = Stage12CacheStore(bucket).upload_file(expected.path, local_path) + return DatasetBuildResult( + **stored.model_dump(), + year=expected.year, + uploaded=uploaded, + build_seconds=round(time.monotonic() - started, 1), + ) + + +def _ensure_dataset_file(bucket: str, year: int) -> None: + from policyengine_simulation_executor.simulation_runtime import resolve_data_folder + + identity = collect_dataset_identity(load_stage12_bundle(), country="us", year=year) + local_path = Path(resolve_data_folder()) / identity.filename + if not local_path.exists(): + Stage12CacheStore(bucket).download_file(identity.store_path, local_path) + + +def compute_baseline_impl( + bucket: str, expected: BaselinePlanEntry +) -> BaselineBuildResult: + from policyengine_simulation_executor.simulation_runtime import resolve_data_folder + from policyengine_simulation_executor.stage12_runtime.simulation import ( + calculate_simulation_frames, + ) + + simulation, identity = _baseline_identity( + expected.year, tuple(expected.region_codes) + ) + if identity.store_path != expected.path: + raise RuntimeError("planned and runtime Stage 12 baseline identities differ") + _ensure_dataset_file(bucket, expected.year) + started = time.monotonic() + calculate_simulation_frames(simulation) + local_path = Path(resolve_data_folder()) / identity.filename + if not local_path.exists(): + raise RuntimeError( + f"Stage 12 baseline calculation produced no cache file {identity.filename}" + ) + uploaded, stored = Stage12CacheStore(bucket).upload_file(expected.path, local_path) + return BaselineBuildResult( + **stored.model_dump(), + year=expected.year, + segment_index=expected.segment_index, + simulation_id=expected.simulation_id, + cache_outcome="miss", + uploaded=uploaded, + compute_seconds=round(time.monotonic() - started, 1), + ) + + +def verify_determinism_impl( + bucket: str, plan: PrecomputePlan, expected: BaselinePlanEntry +) -> DeterminismVerdict: + """Compare a stored baseline with a fresh calculation and record a pass.""" + + from policyengine_simulation_executor.simulation_runtime import resolve_data_folder + from policyengine_simulation_executor.stage12_runtime.simulation import ( + calculate_simulation_frames, + ) + + from .runtime import CACHE_HIT + + if expected not in plan.baselines: + raise RuntimeError("Stage 12 determinism probe is not a planned baseline") + store = Stage12CacheStore(bucket) + group = tuple(expected.region_codes) + baseline = _scoped_planned_simulation(expected.year, group) + reform = _scoped_planned_simulation(expected.year, group, reform=True) + _ensure_dataset_file(bucket, expected.year) + local_path = Path(resolve_data_folder()) / expected.filename + if not local_path.exists(): + store.download_file(expected.path, local_path) + cached_calculation = calculate_simulation_frames(baseline) + if cached_calculation.cache_outcome != CACHE_HIT: + raise RuntimeError( + "Stage 12 determinism check did not load the stored baseline " + f"(cache outcome {cached_calculation.cache_outcome})" + ) + cached = cached_calculation.frames + fresh = calculate_simulation_frames(reform).frames + differences = [] + for entity in sorted(set(cached) | set(fresh)): + left = cached.get(entity) + right = fresh.get(entity) + if left is None or right is None: + differences.append(f"{entity}: missing on one side") + continue + for column in sorted(set(left) | set(right)): + if column not in left or column not in right: + differences.append(f"{entity}.{column}: only on one side") + elif str(left[column].dtype) != str(right[column].dtype): + differences.append(f"{entity}.{column}: dtype differs") + elif not left[column].equals(right[column]): + differences.append(f"{entity}.{column}: values differ") + if not differences: + store.write_json( + determinism_record_path(plan.baselines), + determinism_record(plan.baselines), + ) + return DeterminismVerdict( + equal=not differences, + differences=tuple(differences[:50]), + ) + + +def build_manifest(plan: PrecomputePlan, store: Stage12CacheStore) -> CacheManifest: + artifacts = [] + for entry in plan.datasets: + stored = store.describe(entry.path) + artifacts.append( + ManifestArtifact( + type="dataset", + path=entry.path, + filename=entry.filename, + year=entry.year, + identity_digest=entry.digest, + content_sha256=stored.content_sha256, + size_bytes=stored.size_bytes, + ) + ) + for entry in plan.baselines: + stored = store.describe(entry.path) + artifacts.append( + ManifestArtifact( + type="baseline", + path=entry.path, + filename=entry.filename, + year=entry.year, + identity_digest=entry.digest, + content_sha256=stored.content_sha256, + size_bytes=stored.size_bytes, + ) + ) + return CacheManifest( + years=plan.years, + partition_sha256=plan.partition_sha256, + bundle=plan.bundle, + artifacts=tuple(artifacts), + ) + + +def publish_manifest_impl(bucket: str, plan: PrecomputePlan) -> str: + store = Stage12CacheStore(bucket) + # Checked here, not only by the caller, so no run can publish baselines + # that never passed the determinism check. + if not store.exists(determinism_record_path(plan.baselines)): + raise RuntimeError("Stage 12 cache baselines have not passed determinism") + return store.write_manifest(build_manifest(plan, store).canonical_payload()) + + +def run_precompute( + bucket: str, + *, + plan_artifacts: RemoteFunction, + build_dataset: RemoteFunction, + compute_baseline: RemoteFunction, + verify_determinism: RemoteFunction, + publish_manifest: RemoteFunction, + echo: Callable[[str], None] = print, +) -> str: + plan = PrecomputePlan.model_validate(plan_artifacts.remote(bucket)) + work = select_work(plan) + dataset_calls = [ + build_dataset.spawn(bucket, entry.model_dump()) for entry in work.datasets + ] + for call in dataset_calls: + DatasetBuildResult.model_validate(call.get()) + baseline_calls = [ + compute_baseline.spawn(bucket, entry.model_dump()) for entry in work.baselines + ] + for call in baseline_calls: + BaselineBuildResult.model_validate(call.get()) + # A run that computed baselines checks one of them. A run that computed + # none still checks one when no earlier run recorded a pass, so a failed + # check cannot be skipped by running again. + if work.baselines or not plan.determinism_verified: + probe = work.baselines[0] if work.baselines else plan.baselines[0] + verdict = DeterminismVerdict.model_validate( + verify_determinism.remote(bucket, plan.model_dump(), probe.model_dump()) + ) + if not verdict.equal: + raise SystemExit( + "Stage 12 cache determinism validation failed: " + + "; ".join(verdict.differences) + ) + digest = publish_manifest.remote(bucket, plan.model_dump()) + if not isinstance(digest, str) or not digest: + raise RuntimeError("Stage 12 cache manifest publisher returned no digest") + echo(manifest_digest_line(digest)) + return digest diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py new file mode 100644 index 000000000..03b102535 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/runtime.py @@ -0,0 +1,119 @@ +"""Validated Stage 12 cache loading with calculation fallback.""" + +from __future__ import annotations + +import logging +from collections.abc import Mapping +from typing import Any + +from policyengine.core import Simulation +from policyengine.core.scoping_strategy import RegionGroupStrategy +from policyengine_simulation_contract.stage12_bundle import ResolvedStage12Bundle +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + SimulationRole, + Stage12OutputPlan, +) +from pydantic import PrivateAttr + +from policyengine_simulation_executor.stage12_runtime.output_planning import ( + validate_output_frames, +) + +from .keys import Stage12BaselineCacheIdentity, collect_baseline_identity + +logger = logging.getLogger(__name__) + +CACHE_HIT = "hit" +CACHE_INCOMPLETE = "incomplete" +CACHE_MISS = "miss" + + +def qualifying_cache_identity( + simulation: PlannedSimulationExecutionInput, + *, + resolved: ResolvedStage12Bundle, + dataset_is_default: bool, + scoping_strategy: Any, + spm: Mapping[str, Any] | None, +) -> Stage12BaselineCacheIdentity | None: + """Return a safe cache identity, or None for an ineligible simulation.""" + + if simulation.role is not SimulationRole.BASELINE: + return None + if simulation.policy: + return None + if not dataset_is_default: + return None + if not isinstance(scoping_strategy, RegionGroupStrategy): + return None + raw_group = simulation.options.get("region_group") + if not isinstance(raw_group, list) or not raw_group: + return None + if any( + not isinstance(code, str) or not code.startswith("state/") for code in raw_group + ): + return None + region_codes = tuple(code for code in raw_group if isinstance(code, str)) + scope_key = getattr(scoping_strategy, "cache_key", None) + if not isinstance(scope_key, str) or not scope_key: + return None + return collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=region_codes, + scope_key=scope_key, + spm=spm, + ) + + +class Stage12CachedSimulation(Simulation): + """A deterministic baseline that validates any loaded Stage 12 cache.""" + + _stage12_output_plan: Stage12OutputPlan | None = PrivateAttr(default=None) + _computed_this_process: bool = PrivateAttr(default=False) + _stage12_cache_outcome: str | None = PrivateAttr(default=None) + + @property + def stage12_cache_outcome(self) -> str | None: + return self._stage12_cache_outcome + + def configure_stage12_cache(self, output_plan: Stage12OutputPlan) -> None: + self._stage12_output_plan = output_plan + + def run(self) -> None: + self._computed_this_process = True + super().run() + + def _record_outcome(self, outcome: str) -> None: + if self._stage12_cache_outcome is None: + self._stage12_cache_outcome = outcome + + def ensure(self) -> None: + self._computed_this_process = False + super().ensure() + if self._computed_this_process: + self._record_outcome(CACHE_MISS) + return + plan = self._stage12_output_plan + if plan is None: + raise RuntimeError("Stage 12 cached simulation has no output plan") + data = getattr(getattr(self, "output_dataset", None), "data", None) + entity_data = getattr(data, "entity_data", None) + try: + validate_output_frames(entity_data or {}, plan) + except (TypeError, ValueError): + logger.warning( + "Stage 12 cache %s is incomplete; recalculating", + self.id, + ) + self._record_outcome(CACHE_INCOMPLETE) + self.run() + self.save() + from policyengine.core.simulation import _cache + + storage_id = getattr(self, "storage_id", self.id) + _cache._cache.pop(storage_id, None) + _cache.add(storage_id, self) + return + self._record_outcome(CACHE_HIT) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/store.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/store.py new file mode 100644 index 000000000..e7531c090 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_cache/store.py @@ -0,0 +1,164 @@ +"""Dedicated write-once object store for Stage 12 cache files.""" + +from __future__ import annotations + +import json +import os +from collections.abc import Mapping +from hashlib import sha256 +from pathlib import Path +from typing import Any + +from google.api_core.exceptions import NotFound, PreconditionFailed + +from .keys import canonical_digest +from .models import CacheObject + +CACHE_BUCKET_ENV = "STAGE12_CACHE_BUCKET" + + +def _file_sha256(path: Path) -> str: + digest = sha256() + with path.open("rb") as stream: + while chunk := stream.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def resolve_cache_bucket(explicit: str | None = None) -> str: + bucket = explicit or os.environ.get(CACHE_BUCKET_ENV) + if not bucket: + raise RuntimeError( + f"Stage 12 cache bucket is not configured; set {CACHE_BUCKET_ENV}" + ) + return bucket + + +def _default_client(): + from google.cloud import storage + + from policyengine_simulation_executor.simulation_runtime import ( + setup_gcp_credentials, + ) + + with setup_gcp_credentials(): + return storage.Client() + + +class Stage12CacheStore: + def __init__(self, bucket_name: str | None = None, *, client: Any = None): + self.bucket_name = resolve_cache_bucket(bucket_name) + self._client = client + + @property + def client(self): + if self._client is None: + self._client = _default_client() + return self._client + + def _blob(self, path: str): + return self.client.bucket(self.bucket_name).blob(path) + + def exists(self, path: str) -> bool: + return bool(self._blob(path).exists()) + + def upload_file( + self, path: str, local_path: str | Path + ) -> tuple[bool, CacheObject]: + source = Path(local_path) + payload_digest = _file_sha256(source) + size_bytes = source.stat().st_size + blob = self._blob(path) + blob.metadata = { + "sha256": payload_digest, + "size_bytes": str(size_bytes), + } + try: + blob.upload_from_filename(str(source), if_generation_match=0) + uploaded = True + except PreconditionFailed: + uploaded = False + stored = self.describe(path) + if stored.content_sha256 != payload_digest or stored.size_bytes != size_bytes: + raise RuntimeError("immutable Stage 12 cache path contains other data") + return uploaded, stored + + def upload_bytes( + self, + path: str, + payload: bytes, + *, + content_type: str, + ) -> tuple[bool, CacheObject]: + digest = sha256(payload).hexdigest() + blob = self._blob(path) + blob.metadata = {"sha256": digest, "size_bytes": str(len(payload))} + try: + blob.upload_from_string( + payload, + content_type=content_type, + if_generation_match=0, + ) + uploaded = True + except PreconditionFailed: + uploaded = False + stored = self.describe(path) + if stored.content_sha256 != digest or stored.size_bytes != len(payload): + raise RuntimeError("immutable Stage 12 cache path contains other data") + return uploaded, stored + + def describe(self, path: str) -> CacheObject: + blob = self._blob(path) + try: + blob.reload() + except NotFound: + raise FileNotFoundError(path) from None + metadata = blob.metadata or {} + digest = metadata.get("sha256") + raw_size = metadata.get("size_bytes") + if not isinstance(digest, str) or not isinstance(raw_size, str): + payload = blob.download_as_bytes() + digest = sha256(payload).hexdigest() + raw_size = str(len(payload)) + return CacheObject( + path=path, + content_sha256=digest, + size_bytes=int(raw_size), + ) + + def download_file(self, path: str, local_path: str | Path) -> CacheObject: + destination = Path(local_path) + destination.parent.mkdir(parents=True, exist_ok=True) + partial = destination.with_name(destination.name + ".partial") + try: + self._blob(path).download_to_filename(str(partial)) + stored = self.describe(path) + if _file_sha256(partial) != stored.content_sha256: + raise RuntimeError("downloaded Stage 12 cache digest differs") + partial.replace(destination) + return stored + finally: + partial.unlink(missing_ok=True) + + def write_json(self, path: str, payload: Mapping[str, Any]) -> None: + encoded = json.dumps( + payload, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode("utf-8") + self.upload_bytes(path, encoded, content_type="application/json") + + def write_manifest(self, payload: Mapping[str, Any]) -> str: + digest = canonical_digest(payload) + self.write_json(f"manifests/{digest}.json", payload) + return digest + + def read_manifest(self, digest: str) -> Mapping[str, Any] | None: + try: + payload = self._blob(f"manifests/{digest}.json").download_as_bytes() + except NotFound: + return None + value = json.loads(payload) + return value if isinstance(value, Mapping) else None diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py index 49c473f36..13894f1e3 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_qualification.py @@ -2,7 +2,7 @@ from __future__ import annotations -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from hashlib import sha256 from typing import Any, Literal @@ -80,6 +80,46 @@ class Stage12ParityReceipt(BaseModel): AggregateBuilder = Callable[..., dict[str, Any]] +def _run_stage12_simulation_for_qualification( + simulation: PlannedSimulationExecutionInput, + *, + groups: Sequence[Sequence[str]] | None = None, + segment_runner: Callable[[object], object] | None = None, + monolithic_runner: SingleSimulationRunner | None = None, +) -> SimulationCalculation: + """Exercise the segmented calculation and merge without Modal transport.""" + + from policyengine_simulation_executor.stage12_runtime import ( + build_segment_inputs, + calculate_segment, + calculate_simulation_frames, + merge_segment_results, + should_segment_simulation, + stage12_region_groups_for_model, + ) + + run_monolithic = monolithic_runner or calculate_simulation_frames + if not should_segment_simulation(simulation): + return _calculation(run_monolithic(simulation)) + selected_groups = groups + if selected_groups is None: + from policyengine_simulation_executor.simulation_runtime import ( + _country_module, + ) + + country = simulation.geography.country + selected_groups = stage12_region_groups_for_model( + country, + _country_module(country).model, + ) + run_segment = segment_runner or calculate_segment + results = [ + run_segment(segment) + for segment in build_segment_inputs(simulation, groups=selected_groups) + ] + return merge_segment_results(simulation, results, groups=selected_groups) + + def _local_qualification_runtime() -> ObservabilityRuntime: """Create a local runtime that cannot send qualification data remotely.""" @@ -127,6 +167,7 @@ def _validate_matching_input( "policyengine_version": bundle.policyengine_version, "version": bundle.country_package_version, "spm": baseline.options.get("spm"), + "segmented": baseline.options.get("segmented"), } for field, expected_value in expected.items(): actual_value = existing_request.get(field) @@ -243,12 +284,9 @@ def run_existing(request: dict[str, Any]) -> Mapping[str, Any]: return run_simulation_impl(request, runtime=local_runtime) existing_runner = run_existing + default_stage12_runner = single_simulation_runner is None if single_simulation_runner is None: - from policyengine_simulation_executor.stage12_runtime import ( - calculate_simulation_frames, - ) - - single_simulation_runner = calculate_simulation_frames + single_simulation_runner = _run_stage12_simulation_for_qualification if aggregate_builder is None: from policyengine_simulation_executor.stage12_runtime import ( build_aggregate_report, @@ -279,11 +317,18 @@ def run_existing(request: dict[str, Any]) -> Mapping[str, Any]: "report.data_version", ) - with ThreadPoolExecutor(max_workers=2) as executor: - baseline_future = executor.submit(single_simulation_runner, baseline_input) - reform_future = executor.submit(single_simulation_runner, reform_input) - baseline_calculation = _calculation(baseline_future.result()) - reform_calculation = _calculation(reform_future.result()) + if default_stage12_runner: + # The production segments run in isolated Modal processes. Run the + # local qualifier sequentially so country-model process caches are not + # accessed concurrently by Python threads. + baseline_calculation = _calculation(single_simulation_runner(baseline_input)) + reform_calculation = _calculation(single_simulation_runner(reform_input)) + else: + with ThreadPoolExecutor(max_workers=2) as executor: + baseline_future = executor.submit(single_simulation_runner, baseline_input) + reform_future = executor.submit(single_simulation_runner, reform_input) + baseline_calculation = _calculation(baseline_future.result()) + reform_calculation = _calculation(reform_future.result()) compare_simulation_frames( existing_baseline, diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py index dd61bce17..3a876f234 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/__init__.py @@ -2,8 +2,26 @@ from .aggregation import build_aggregate_report, build_spm_result from .coordination import coordinate_report +from .failures import ( + Stage12Cancellation, + Stage12ExecutionError, + Stage12FailureDetail, + Stage12InputError, + redacted_failure_detail, + validate_policy_periods, +) +from .fanout import Stage12SegmentRunner, run_segmented_simulation +from .partition import stage12_region_groups, stage12_region_groups_for_model +from .segment_contracts import Stage12SegmentInput, Stage12SegmentResult +from .segmentation import ( + build_segment_inputs, + calculate_segment, + merge_segment_results, + should_segment_simulation, +) from .simulation import ( SimulationCalculation, + build_stage12_simulation, calculate_simulation_frames, run_single_simulation, simulation_input_sha256, @@ -15,9 +33,26 @@ __all__ = [ "SimulationCalculation", + "Stage12SegmentInput", + "Stage12SegmentResult", + "Stage12SegmentRunner", + "Stage12ExecutionError", + "Stage12FailureDetail", + "Stage12InputError", + "Stage12Cancellation", "build_aggregate_report", + "build_segment_inputs", + "build_stage12_simulation", "calculate_simulation_frames", + "calculate_segment", "coordinate_report", + "redacted_failure_detail", "run_single_simulation", + "run_segmented_simulation", + "merge_segment_results", "simulation_input_sha256", + "stage12_region_groups", + "stage12_region_groups_for_model", + "should_segment_simulation", + "validate_policy_periods", ] diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py index a9dadb23a..b064fccbe 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/aggregation.py @@ -6,6 +6,7 @@ from typing import Any import pandas as pd +from policyengine.core import Simulation from policyengine_simulation_contract.stage12_bundle import CountryId from policyengine_simulation_contract.stage12_execution import ( ReportExecutionInput, @@ -51,6 +52,15 @@ def validate_uk_local_authority_metadata( return None +class PrecomputedSimulation(Simulation): + """A Stage 12 aggregate input whose output dataset is already complete.""" + + def ensure(self) -> None: + """Do not recalculate frames loaded from Stage 12 artifacts.""" + + return + + def validate_aligned_outputs( report: ReportExecutionInput, output_plan: Stage12OutputPlan, @@ -130,9 +140,6 @@ def build_aggregate_report( reform_descriptor: SimulationArtifactDescriptor, uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None, ) -> dict[str, Any]: - from policyengine_simulation_executor.segmented_national_reduce import ( - PrecomputedSimulation, - ) from policyengine_simulation_executor.simulation_output_builder import ( SimulationOutputBuilder, ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py index c61bdd50c..aee0586e7 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/comparison.py @@ -17,11 +17,13 @@ from policyengine_simulation_executor.stage12_artifacts import Stage12ArtifactStore from policyengine_simulation_executor.stage12_result_comparison import compare_results +from policyengine_simulation_executor.stage12_timeouts import ( + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, +) from .dependencies import ChildInvoker, ComparisonStore from .failure_causes import failure_summary -PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS = 900 logger = logging.getLogger(__name__) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py index ae1824909..7c1d8a431 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/coordination.py @@ -2,12 +2,14 @@ from __future__ import annotations +import logging +import time from collections.abc import Callable from datetime import UTC, datetime from functools import partial from hashlib import sha256 from typing import Any -from uuid import UUID, uuid4 +from uuid import uuid4 from policyengine_observability import ObservabilityRuntime from policyengine_simulation_contract.stage12_execution import ( @@ -36,6 +38,9 @@ deserialize_simulation_frames, deserialize_uk_local_authority_metadata, ) +from policyengine_simulation_executor.stage12_timeouts import ( + SIMULATION_WAIT_TIMEOUT_SECONDS, +) from .aggregation import ( build_aggregate_report, @@ -51,21 +56,334 @@ artifact_store, runtime_store, ) -from .failure_causes import ( - StageProgress, - failure_summary, - record_failure_cause, - tracked_stage, +from .failure_causes import StageProgress, record_failure_cause, tracked_stage +from .failure_recording import ( + persist_failure_record, + report_failure, + report_original_failure, +) +from .failures import ( + Stage12ExecutionError, + Stage12FailureDetail, + failure_detail_from_record, ) -from .failure_recording import persist_failure_record, report_original_failure from .output_planning import ( plan_simulation_input, resolve_report_output_plan, validate_output_frames, ) +from .segmentation import should_segment_simulation from .simulation import descriptor_from_record, simulation_input_sha256 -SIMULATION_WAIT_TIMEOUT_SECONDS = 3_000 +SIMULATION_POLL_INITIAL_SECONDS = 0.25 +SIMULATION_POLL_MAX_SECONDS = 1.0 +SIMULATION_POLL_BACKOFF_FACTOR = 2.0 + +logger = logging.getLogger(__name__) + + +def _next_poll_interval(current: float) -> float: + return min( + current * SIMULATION_POLL_BACKOFF_FACTOR, + SIMULATION_POLL_MAX_SECONDS, + ) + + +def _log_child_operation_failure( + level: int, + message: str, + *, + simulation: PlannedSimulationExecutionInput, + error: Exception, +) -> None: + """Log a failed child state update or cancellation with its identifiers. + + ``message`` receives the simulation role, simulation execution ID, + evaluation ID and error type, in that order. The error's own message is + never logged. + """ + + error_type = type(error).__name__ + logger.log( + level, + message, + simulation.role.value, + simulation.simulation_execution_id, + simulation.evaluation_id, + error_type, + extra={ + "evaluation_id": str(simulation.evaluation_id), + "simulation_execution_id": str(simulation.simulation_execution_id), + "simulation_role": simulation.role.value, + "error_type": error_type, + }, + ) + + +def _replace_running_simulation( + *, + store: ComparisonStore, + simulation: PlannedSimulationExecutionInput, + status: ComparisonRunLifecycleStatus, + detail: Stage12FailureDetail, +) -> ComparisonSimulationRecord: + changed_at = datetime.now(UTC) + current = store.get_simulation(simulation.simulation_execution_id) + candidate = current.model_copy( + update={ + "status": status, + "error_code": detail.error_code, + "error_summary": detail.error_summary, + "updated_at": changed_at, + "completed_at": changed_at, + } + ) + result, _ = store.replace_simulation_if_status( + candidate, + expected_status=ComparisonRunLifecycleStatus.RUNNING, + ) + return result + + +SIMULATION_CANCELLED_ERROR_CODE = "simulation_cancelled" + + +def _cancellation(reason: str) -> Stage12FailureDetail: + """Describe a child the coordinator stopped; ``reason`` ends the summary.""" + + return Stage12FailureDetail( + error_code=SIMULATION_CANCELLED_ERROR_CODE, + error_summary=f"Cancelled after {reason}.", + ) + + +def _cancel_pending_simulations( + *, + calls: dict[str, ChildCall], + simulations: dict[str, PlannedSimulationExecutionInput], + roles: set[str], + store: ComparisonStore, + detail: Stage12FailureDetail, +) -> None: + """Record each started, unfinished child as failed, then stop it. + + Only a running record changes, so a child that has already succeeded or + failed keeps its own outcome. An unsegmented child is always cancelled + through Modal, even if its record could not be written. State-update and + cancellation errors are logged rather than raised, so they never replace + the failure that caused the cancellation. + """ + + for role in sorted(roles): + simulation = simulations[role] + try: + _replace_running_simulation( + store=store, + simulation=simulation, + status=ComparisonRunLifecycleStatus.FAILED, + detail=detail, + ) + except Exception as error: # noqa: BLE001 + _log_child_operation_failure( + logging.ERROR, + "Stage 12 could not record cancellation of %s simulation %s " + "in report %s (%s)", + simulation=simulation, + error=error, + ) + cancellation_recorded = False + else: + cancellation_recorded = True + if cancellation_recorded and should_segment_simulation(simulation): + # A segmented worker reads its record, cancels the region-group + # calls it started, and exits. Modal does not cancel spawned calls + # recursively, so cancelling the worker through Modal would leave + # those calls running. That is the fallback only when the record + # could not be written and the worker cannot learn of the + # cancellation. + continue + try: + calls[role].cancel() + except Exception as error: # noqa: BLE001 + _log_child_operation_failure( + logging.WARNING, + "Stage 12 could not cancel %s simulation %s in report %s (%s)", + simulation=simulation, + error=error, + ) + + +def _simulation_failure_attributes( + simulation: PlannedSimulationExecutionInput, +) -> dict[str, object]: + return { + "simulation_role": simulation.role.value, + "evaluation_id": str(simulation.evaluation_id), + "simulation_execution_id": str(simulation.simulation_execution_id), + } + + +def _log_child_state_failure( + *, + simulation: PlannedSimulationExecutionInput, + error: Exception, +) -> None: + _log_child_operation_failure( + logging.ERROR, + "Stage 12 could not record the failure of %s simulation %s in report %s (%s)", + simulation=simulation, + error=error, + ) + + +def _record_child_failure( + *, + store: ComparisonStore, + simulation: PlannedSimulationExecutionInput, + detail: Stage12FailureDetail, +) -> None: + """Record a running child as failed with its own cause. + + A child that has already finished keeps its own outcome. A state error is + logged rather than raised, so it never replaces the child's failure. + """ + + try: + _replace_running_simulation( + store=store, + simulation=simulation, + status=ComparisonRunLifecycleStatus.FAILED, + detail=detail, + ) + except Exception as state_error: # noqa: BLE001 + _log_child_state_failure(simulation=simulation, error=state_error) + + +def _child_failure_detail( + *, + error: Exception, + simulation: PlannedSimulationExecutionInput, + store: ComparisonStore, + runtime: ObservabilityRuntime | None, + progress: StageProgress, +) -> Stage12FailureDetail: + """Report a failed child call and return the failure for its report. + + A failure the child's worker recorded itself takes precedence over the + error its call raised. Reading or writing the child's record can fail, + for example when the database is unavailable. That error is logged, and + the call's own failure is used and still recorded where possible, so a + state error never replaces the child's failure. + """ + + detail = report_failure( + error, + runtime=runtime, + failure_scope="child_wait", + default_code="simulation_invocation_failed", + attributes=_simulation_failure_attributes(simulation), + progress=progress, + ) + try: + child = store.get_simulation(simulation.simulation_execution_id) + except Exception as state_error: # noqa: BLE001 + _log_child_state_failure(simulation=simulation, error=state_error) + else: + if ( + child.status is ComparisonRunLifecycleStatus.FAILED + and child.error_code is not None + ): + return failure_detail_from_record( + error_code=child.error_code, + error_summary=child.error_summary, + ) + _record_child_failure(store=store, simulation=simulation, detail=detail) + return detail + + +def _observe_simulations( + *, + calls: dict[str, ChildCall], + simulations: dict[str, PlannedSimulationExecutionInput], + descriptors: dict[str, SimulationArtifactDescriptor], + store: ComparisonStore, + runtime: ObservabilityRuntime | None, + progress: StageProgress, + wait_stage_name: str, + sleep: Callable[[float], None], + monotonic: Callable[[], float], + deadline: float, + timeout_seconds: float, +) -> None: + pending = set(calls) + poll_interval = SIMULATION_POLL_INITIAL_SECONDS + + while pending: + progress_made = False + for role in sorted(pending): + call = calls[role] + try: + # The probe records the child's role, so a failure names it. + with progress.stage(wait_stage_name, role=role): + result = call.get(timeout=0) + descriptor = SimulationArtifactDescriptor.model_validate(result) + except TimeoutError: + continue + except Exception as error: # noqa: BLE001 + detail = _child_failure_detail( + error=error, + simulation=simulations[role], + store=store, + runtime=runtime, + progress=progress, + ) + _cancel_pending_simulations( + calls=calls, + simulations=simulations, + roles=pending - {role}, + store=store, + detail=_cancellation(f"the {role} simulation failed"), + ) + raise Stage12ExecutionError(detail) from None + descriptors[role] = descriptor + pending.remove(role) + progress_made = True + + if not pending: + return + remaining_seconds = deadline - monotonic() + if remaining_seconds <= 0: + detail = Stage12FailureDetail( + error_code="simulation_wait_timeout", + error_summary=( + "Stage 12 simulations did not finish within " + f"{timeout_seconds:g} seconds of the report coordinator " + "starting." + ), + ) + wait_timeout = Stage12ExecutionError(detail) + report_original_failure( + wait_timeout, + runtime=runtime, + failure_scope="child_wait_timeout", + ) + _cancel_pending_simulations( + calls=calls, + simulations=simulations, + roles=pending, + store=store, + detail=_cancellation( + f"the {timeout_seconds:g}-second wait window closed" + ), + ) + raise wait_timeout + # Never sleep past the deadline, so the window closes on time. + sleep(min(poll_interval, remaining_seconds)) + poll_interval = ( + SIMULATION_POLL_INITIAL_SECONDS + if progress_made + else _next_poll_interval(poll_interval) + ) def _child_record( @@ -194,43 +512,6 @@ def _claim_parent( return claimed, True -def _mark_child_failed( - persistence: ComparisonStore, - simulation_execution_id: UUID, - *, - error_code: str, - error_summary: str, -) -> str | None: - """Mark an unfinished child failed and return the cause its record keeps. - - A worker that failed has already persisted its own redacted cause before - Modal returned the worker's stable exception, so that cause is kept rather - than overwritten. A succeeded child is left unchanged and has no cause. - """ - - latest = persistence.get_simulation(simulation_execution_id) - if latest.status is ComparisonRunLifecycleStatus.SUCCEEDED: - return None - if ( - latest.status is ComparisonRunLifecycleStatus.FAILED - and latest.error_summary is not None - ): - return latest.error_summary - failed_at = datetime.now(UTC) - persistence.replace_simulation( - latest.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": error_code, - "error_summary": error_summary, - "updated_at": failed_at, - "completed_at": failed_at, - } - ) - ) - return error_summary - - def coordinate_report( payload: object, context_payload: object, @@ -245,7 +526,12 @@ def coordinate_report( output_plan_resolver: Callable[[ReportExecutionInput], Stage12OutputPlan] | None = None, runtime: ObservabilityRuntime | None = None, + sleep: Callable[[float], None] = time.sleep, + monotonic: Callable[[], float] = time.monotonic, + wait_timeout_seconds: float = SIMULATION_WAIT_TIMEOUT_SECONDS, ) -> dict[str, Any]: + # Open the shared child-wait window before any setup work. + child_wait_deadline = monotonic() + wait_timeout_seconds report = ReportExecutionInput.model_validate(payload) context = Stage12InvocationContext.model_validate(context_payload) submitted_parent = ComparisonReportRecord.model_validate(parent_payload) @@ -290,12 +576,7 @@ def coordinate_report( } descriptors: dict[str, SimulationArtifactDescriptor] = {} calls: dict[str, ChildCall] = {} - # A child failure is reported where it occurs and then propagates to the - # report handler below, which must not report the same error twice. - reported_child_failure: Exception | None = None - # A failed child records its own cause; the report keeps that cause - # rather than the generic exception Modal returns to the coordinator. - child_failure_summary: str | None = None + simulations: dict[str, PlannedSimulationExecutionInput] = {} try: with stage_scope(Stage.STAGE12_COORDINATOR_PREPARATION): context = context.model_copy( @@ -321,14 +602,19 @@ def coordinate_report( ) else: output_plan = output_plan_resolver(report) - simulations = [] + planned_inputs: list[PlannedSimulationExecutionInput] = [] for simulation in (report.baseline, report.reform): with stage_scope( Stage.STAGE12_CHILD_INPUT_PLANNING, {"simulation_role": simulation.role.value}, ): - simulations.append(plan_simulation_input(simulation, output_plan)) - planned_simulations = tuple(simulations) + planned_inputs.append( + plan_simulation_input(simulation, output_plan) + ) + planned_simulations = tuple(planned_inputs) + simulations = { + simulation.role.value: simulation for simulation in planned_simulations + } with stage_span(Stage.STAGE12_CHILD_STATE_CREATE): children = { simulation.role: persistence.create_or_resolve_simulation( @@ -400,25 +686,29 @@ def coordinate_report( observability_context=child_observability_context, ) except Exception as error: - report_original_failure( + detail = report_failure( error, runtime=runtime, failure_scope="child_dispatch", - attributes={"simulation_role": simulation.role.value}, + default_code="simulation_dispatch_failed", + attributes=_simulation_failure_attributes(simulation), + progress=progress, + ) + _record_child_failure( + store=persistence, + simulation=simulation, + detail=detail, ) - reported_child_failure = error - persist_failure_record( - partial( - _mark_child_failed, - persistence, - simulation.simulation_execution_id, - error_code="simulation_dispatch_failed", - error_summary=failure_summary(error, progress=progress), + _cancel_pending_simulations( + calls=calls, + simulations=simulations, + roles=set(calls), + store=persistence, + detail=_cancellation( + f"the {simulation.role.value} simulation failed" ), - runtime=runtime, - failure_scope="child_dispatch", ) - raise + raise Stage12ExecutionError(detail) from None calls[simulation.role.value] = call persistence.attach_simulation_invocation( simulation.simulation_execution_id, @@ -426,36 +716,22 @@ def coordinate_report( modal_invocation_id=call.object_id, updated_at=datetime.now(UTC), ) - # Every required call has been started before the coordinator waits. - for role, call in calls.items(): - try: - with stage_span(Stage.STAGE12_CHILD_WAIT, {"simulation_role": role}): - descriptors[role] = SimulationArtifactDescriptor.model_validate( - call.get(timeout=SIMULATION_WAIT_TIMEOUT_SECONDS) - ) - except Exception as error: - report_original_failure( - error, - runtime=runtime, - failure_scope="child_wait", - attributes={"simulation_role": role}, - ) - reported_child_failure = error - simulation = next( - item for item in simulations if item.role.value == role - ) - child_failure_summary = persist_failure_record( - partial( - _mark_child_failed, - persistence, - simulation.simulation_execution_id, - error_code="simulation_invocation_failed", - error_summary=failure_summary(error, progress=progress), - ), - runtime=runtime, - failure_scope="child_wait", - ) - raise + # Every required call has been started before the coordinator observes + # baseline and reform together through nonblocking result probes. + with stage_span(Stage.STAGE12_CHILD_WAIT): + _observe_simulations( + calls=calls, + simulations=simulations, + descriptors=descriptors, + store=persistence, + runtime=runtime, + progress=progress, + wait_stage_name=stage_plan.name(Stage.STAGE12_CHILD_WAIT), + sleep=sleep, + monotonic=monotonic, + deadline=child_wait_deadline, + timeout_seconds=wait_timeout_seconds, + ) baseline = descriptors["baseline"] reform = descriptors["reform"] validate_aligned_outputs(report, output_plan, baseline, reform) @@ -525,16 +801,33 @@ def coordinate_report( # Persist a redacted cause for any coordinator, child-call, aggregation, or # artifact failure before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 - if error is not reported_child_failure: - report_original_failure( + if isinstance(error, Stage12ExecutionError): + # Child, dispatch and wait-timeout failures raise + # Stage12ExecutionError only after reporting the original failure + # and stopping their unfinished peers. + detail = error.detail + else: + # Any other failure is reported here, once, before any database + # write. If it occurred after a child started, such as a failed + # placeholder write or invocation attach, every started, + # unfinished child is then stopped so none keeps running for a + # failed report. + detail = report_failure( error, runtime=runtime, failure_scope="report_coordination", + default_code="report_coordination_failed", + attributes={"evaluation_id": str(report.evaluation_id)}, + progress=progress, ) - report_failure_summary = child_failure_summary or failure_summary( - error, progress=progress - ) - record_failure_cause(runtime, error, report_failure_summary) + _cancel_pending_simulations( + calls=calls, + simulations=simulations, + roles=set(calls) - set(descriptors), + store=persistence, + detail=_cancellation("the coordinator failed"), + ) + record_failure_cause(runtime, error, detail.error_summary) failed_at = datetime.now(UTC) persist_failure_record( partial( @@ -543,8 +836,8 @@ def coordinate_report( update={ "status": ComparisonRunLifecycleStatus.FAILED, "aggregation_status": ComparisonRunAggregationStatus.FAILED, - "error_code": "report_coordination_failed", - "error_summary": report_failure_summary, + "error_code": detail.error_code, + "error_summary": detail.error_summary, "updated_at": failed_at, "completed_at": failed_at, } diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py index e5848ebe0..3009789b1 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/dependencies.py @@ -12,6 +12,7 @@ from policyengine_simulation_contract.stage12_execution import ( ComparisonReportPersistenceResult, ComparisonReportRecord, + ComparisonRunLifecycleStatus, ComparisonSimulationPersistenceResult, ComparisonSimulationRecord, ) @@ -53,6 +54,13 @@ def replace_simulation( record: ComparisonSimulationRecord, ) -> ComparisonSimulationRecord: ... + def replace_simulation_if_status( + self, + record: ComparisonSimulationRecord, + *, + expected_status: ComparisonRunLifecycleStatus, + ) -> tuple[ComparisonSimulationRecord, bool]: ... + def attach_simulation_invocation( self, simulation_execution_id: UUID, @@ -68,6 +76,8 @@ class ChildCall(Protocol): def get(self, *, timeout: float | None = None) -> object: ... + def cancel(self) -> None: ... + class ChildInvoker(Protocol): def spawn( diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py index b86c0679c..ac6dc782d 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failure_recording.py @@ -8,9 +8,51 @@ from policyengine_observability import ObservabilityRuntime +from .failure_causes import StageProgress +from .failures import Stage12FailureDetail, redacted_failure_detail + logger = logging.getLogger(__name__) +def report_failure( + error: Exception, + *, + runtime: ObservabilityRuntime | None, + failure_scope: str, + default_code: str, + attributes: Mapping[str, object] | None = None, + segment_index: int | None = None, + progress: StageProgress | None = None, +) -> Stage12FailureDetail: + """Report a failure once and return the safe detail to persist for it. + + The report is emitted before any database write is attempted, as + ``report_original_failure`` requires. An unexpected error's detail has a + new correlation ID, and the report carries the same ID, so the redacted + summary persisted for callers leads back to this one record of the + original error. ``progress`` names the failing stage in that summary. + """ + + detail = redacted_failure_detail( + error, + default_code=default_code, + segment_index=segment_index, + progress=progress, + ) + correlation = ( + {"correlation_id": detail.correlation_id} + if detail.correlation_id is not None + else {} + ) + report_original_failure( + error, + runtime=runtime, + failure_scope=failure_scope, + attributes={**(attributes or {}), **correlation}, + ) + return detail + + def report_original_failure( error: Exception, *, @@ -69,10 +111,14 @@ def _emit_error( attributes: Mapping[str, object], ) -> None: if runtime is None: + failure_label = str(attributes.get("stage12_failure_scope")) + correlation_id = attributes.get("correlation_id") + if correlation_id is not None: + failure_label = f"{failure_label}; correlation_id={correlation_id}" logger.error( "%s (%s)", message, - attributes.get("stage12_failure_scope"), + failure_label, exc_info=error, ) return diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py new file mode 100644 index 000000000..b46c05037 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/failures.py @@ -0,0 +1,158 @@ +"""Safe, structured failure details for Stage 12 execution.""" + +from __future__ import annotations + +from collections.abc import Mapping +from datetime import date +from typing import Any + +from policyengine_simulation_contract.spm import spm_error_detail +from policyengine_simulation_observability.errors import make_correlation_id +from pydantic import BaseModel, ConfigDict + +from .failure_causes import StageProgress, failure_summary + +ERROR_CODE_MAX_LENGTH = 64 +ERROR_SUMMARY_MAX_LENGTH = 512 + + +def _bounded(value: str, maximum: int) -> str: + if len(value) <= maximum: + return value + return value[: maximum - 3] + "..." + + +class Stage12FailureDetail(BaseModel): + """Failure values safe to persist and return from the internal status API.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + error_code: str + error_summary: str + correlation_id: str | None = None + segment_index: int | None = None + + +class Stage12InputError(ValueError): + """An input error whose code and message are safe to return to the caller.""" + + def __init__(self, code: str, message: str) -> None: + self.code = code + self.message = message + super().__init__(message) + + +class Stage12ExecutionError(RuntimeError): + """Carry one safe failure between Stage 12 execution layers.""" + + def __init__(self, detail: Stage12FailureDetail) -> None: + self.detail = detail + super().__init__("Stage 12 execution failed") + + +class Stage12Cancellation(RuntimeError): + """Stop a simulation whose durable record was cancelled by its parent.""" + + +def _safe_detail( + *, + error_code: str, + error_summary: str, + correlation_id: str | None = None, + segment_index: int | None = None, +) -> Stage12FailureDetail: + return Stage12FailureDetail( + error_code=_bounded(error_code, ERROR_CODE_MAX_LENGTH), + error_summary=_bounded(error_summary, ERROR_SUMMARY_MAX_LENGTH), + correlation_id=correlation_id, + segment_index=segment_index, + ) + + +def failure_detail_from_record( + *, + error_code: str, + error_summary: str | None, +) -> Stage12FailureDetail: + """Recreate a propagated failure without replacing its persisted values.""" + + return _safe_detail( + error_code=error_code, + error_summary=error_summary or "Simulation failed", + ) + + +def redacted_failure_detail( + error: BaseException, + *, + default_code: str, + segment_index: int | None = None, + progress: StageProgress | None = None, +) -> Stage12FailureDetail: + """Retain typed public failures and redact unexpected exception messages. + + This only describes the failure; it does not report it. Any other error + is summarized by its redacted cause (see ``failure_summary``), naming the + stage ``progress`` recorded, and a new correlation ID. The caller reports + the error once with that ID so the summary leads back to it. The cause is + shortened, never the ID, to fit the summary limit. + """ + + if isinstance(error, Stage12ExecutionError): + return error.detail + if isinstance(error, Stage12InputError): + return _safe_detail( + error_code=error.code, + error_summary=error.message, + segment_index=segment_index, + ) + spm = spm_error_detail(error) + if spm is not None: + return _safe_detail( + error_code=spm.code, + error_summary=spm.message, + segment_index=segment_index, + ) + + correlation_id = make_correlation_id() + correlation = f" (correlation_id={correlation_id})" + cause = failure_summary(error, progress=progress) + if segment_index is not None: + cause = f"Segment {segment_index}: {cause}" + cause = _bounded(cause, ERROR_SUMMARY_MAX_LENGTH - len(correlation)) + return _safe_detail( + error_code=default_code, + error_summary=f"{cause}{correlation}", + correlation_id=correlation_id, + segment_index=segment_index, + ) + + +def validate_policy_periods(policy: Mapping[str, Any]) -> None: + """Require effective dates in the wire formats emitted by the frontend.""" + + for parameter, value in policy.items(): + if not isinstance(value, Mapping): + continue + for raw_period in value: + period = str(raw_period) + parts = period.split(".") + if len(parts) not in {1, 2}: + _raise_invalid_period(parameter, period) + parsed = [] + for part in parts: + try: + parsed.append(date.fromisoformat(part)) + except ValueError: + _raise_invalid_period(parameter, period) + if len(parsed) == 2 and parsed[1] < parsed[0]: + _raise_invalid_period(parameter, period) + + +def _raise_invalid_period(parameter: str, period: str) -> None: + raise Stage12InputError( + "invalid_policy_period", + f"Policy parameter {parameter!r} has invalid effective period " + f"{period!r}; expected YYYY-MM-DD or " + "YYYY-MM-DD.YYYY-MM-DD.", + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py new file mode 100644 index 000000000..d8026a625 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/fanout.py @@ -0,0 +1,227 @@ +"""Independent Modal orchestration for one Stage 12 segmented simulation.""" + +from __future__ import annotations + +import logging +import time +from collections.abc import Callable, Sequence +from typing import Any + +import modal +from policyengine_observability import ObservabilityRuntime +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, +) +from policyengine_simulation_observability.stages import ( + STAGE12_SIMULATION_STAGES, + Stage, +) + +from .failure_recording import report_failure +from .failures import Stage12Cancellation, Stage12ExecutionError +from .partition import stage12_region_groups_for_model +from .segmentation import build_segment_inputs, merge_segment_results +from .simulation import SimulationCalculation + +logger = logging.getLogger(__name__) + +SEGMENT_FUNCTION_NAME = "run_single_simulation_segment_us" +POLL_INTERVAL_INITIAL_SECONDS = 0.5 +POLL_INTERVAL_MAX_SECONDS = 8.0 +POLL_INTERVAL_BACKOFF_FACTOR = 2.0 + + +def _next_backoff(current: float, *, maximum: float) -> float: + return min(current * POLL_INTERVAL_BACKOFF_FACTOR, maximum) + + +class Stage12SegmentRunner: + """Start, collect, and merge the 20 parts of one logical simulation.""" + + def __init__( + self, + simulation: PlannedSimulationExecutionInput, + *, + app_name: str, + runtime: ObservabilityRuntime, + modal_module: Any = modal, + groups: Sequence[Sequence[str]] | None = None, + sleep: Callable[[float], None] = time.sleep, + poll_interval_seconds: float = POLL_INTERVAL_INITIAL_SECONDS, + poll_interval_max_seconds: float = POLL_INTERVAL_MAX_SECONDS, + cancellation_requested: Callable[[], bool] | None = None, + ) -> None: + self.simulation = simulation + self.runtime = runtime + self.sleep = sleep + self.poll_interval_initial_seconds = poll_interval_seconds + self.poll_interval_max_seconds = poll_interval_max_seconds + self.cancellation_requested = cancellation_requested or (lambda: False) + self.groups = [list(group) for group in groups] if groups is not None else None + self.child_function = modal_module.Function.from_name( + app_name, + SEGMENT_FUNCTION_NAME, + ) + + def _groups(self) -> list[list[str]]: + if self.groups is not None: + return self.groups + from policyengine_simulation_executor.simulation_runtime import ( + _country_module, + ) + + country = self.simulation.geography.country + return stage12_region_groups_for_model( + country, + _country_module(country).model, + ) + + def run(self) -> SimulationCalculation: + groups = self._groups() + segments = build_segment_inputs(self.simulation, groups=groups) + self.runtime.set_context( + stage12_execution="segmented", + stage12_segment_count=len(segments), + ) + handles: list[tuple[int, Any]] = [] + try: + for segment in segments: + call = None + with self.runtime.span( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_DISPATCH), + attributes={"segment_index": segment.segment_index}, + ): + call = self.child_function.spawn( + segment.model_dump(mode="python"), + observability_context=self.runtime.capture_context(), + ) + if call is None: + raise RuntimeError("Stage 12 segment dispatch returned no call") + handles.append((segment.segment_index, call)) + except Exception: + self._cancel_all(handles) + raise + + payloads: list[object] | None = None + with self.runtime.span( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_WAIT) + ): + payloads = self._collect(handles) + if payloads is None: + raise RuntimeError("Stage 12 segment wait returned no payloads") + merged: SimulationCalculation | None = None + with self.runtime.span( + STAGE12_SIMULATION_STAGES.name(Stage.STAGE12_SEGMENT_MERGE) + ): + merged = merge_segment_results( + self.simulation, + payloads, + groups=groups, + ) + if merged is None: + raise RuntimeError("Stage 12 segment merge returned no calculation") + return merged + + def _collect(self, handles: list[tuple[int, Any]]) -> list[object]: + results: list[object | None] = [None] * len(handles) + pending = set(range(len(handles))) + poll_errors: dict[int, int] = {} + current_sleep = self.poll_interval_initial_seconds + poll_count = 0 + + while pending: + self._raise_if_cancelled(handles) + progress_made = False + for index in sorted(pending): + segment_index, call = handles[index] + poll_count += 1 + try: + results[index] = call.get(timeout=0) + except TimeoutError: + poll_errors.pop(index, None) + continue + except Exception as error: + attempts = poll_errors.get(index, 0) + 1 + poll_errors[index] = attempts + if attempts < 2: + logger.warning( + "Stage 12 segment %d poll failed once; retrying (%s)", + segment_index, + type(error).__name__, + ) + continue + detail = report_failure( + error, + runtime=self.runtime, + failure_scope="segment_call", + default_code="segment_execution_failed", + attributes={ + "simulation_role": self.simulation.role.value, + "evaluation_id": str(self.simulation.evaluation_id), + "simulation_execution_id": str( + self.simulation.simulation_execution_id + ), + "segment_index": segment_index, + }, + segment_index=segment_index, + ) + self._cancel_all(handles) + raise Stage12ExecutionError(detail) from None + pending.discard(index) + poll_errors.pop(index, None) + progress_made = True + self.runtime.set_context(stage12_segment_poll_count=poll_count) + if pending and not progress_made: + self._wait_with_cancellation(current_sleep, handles) + current_sleep = _next_backoff( + current_sleep, + maximum=self.poll_interval_max_seconds, + ) + elif progress_made: + current_sleep = self.poll_interval_initial_seconds + + if any(result is None for result in results): + raise RuntimeError("Stage 12 segment collection is incomplete") + return [result for result in results if result is not None] + + def _wait_with_cancellation( + self, + seconds: float, + handles: list[tuple[int, Any]], + ) -> None: + remaining = seconds + while remaining > 0: + self._raise_if_cancelled(handles) + delay = min(1.0, remaining) + self.sleep(delay) + remaining -= delay + self._raise_if_cancelled(handles) + + def _raise_if_cancelled(self, handles: list[tuple[int, Any]]) -> None: + if not self.cancellation_requested(): + return + self._cancel_all(handles) + raise Stage12Cancellation("Stage 12 simulation was cancelled") + + @staticmethod + def _cancel_all(handles: list[tuple[int, Any]]) -> None: + for _, call in handles: + try: + call.cancel() + except Exception: + pass + + +def run_segmented_simulation( + simulation: PlannedSimulationExecutionInput, + *, + app_name: str, + runtime: ObservabilityRuntime, + cancellation_requested: Callable[[], bool] | None = None, +) -> SimulationCalculation: + return Stage12SegmentRunner( + simulation, + app_name=app_name, + runtime=runtime, + cancellation_requested=cancellation_requested, + ).run() diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/partition.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/partition.py new file mode 100644 index 000000000..5e7066b80 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/partition.py @@ -0,0 +1,64 @@ +"""Stage 12-owned US partition for segmented national calculations.""" + +from __future__ import annotations + +from typing import Any + +# This is an independent snapshot of the measured 20-way partition used by +# the v1 runner when Stage 12 segmentation was introduced. Stage 12 must not +# import the v1 partition: either execution path can now evolve without +# changing the other one implicitly. +US_REGION_GROUPS: tuple[tuple[str, ...], ...] = ( + ("state/hi", "state/ia", "state/wi"), + ("state/ak", "state/ga", "state/la"), + ("state/al", "state/co", "state/ri"), + ("state/mn", "state/mo", "state/nh"), + ("state/az", "state/de", "state/me"), + ("state/ny", "state/wv"), + ("state/in", "state/ky", "state/md"), + ("state/mt", "state/sc", "state/wa"), + ("state/id", "state/or", "state/tn"), + ("state/ar", "state/nc", "state/nd"), + ("state/dc", "state/ks", "state/mi"), + ("state/ct", "state/oh", "state/vt"), + ("state/nv", "state/va", "state/wy"), + ("state/pa", "state/ut"), + ("state/sd", "state/tx"), + ("state/ca",), + ("state/il", "state/ms"), + ("state/fl", "state/nm"), + ("state/ne", "state/nj"), + ("state/ma", "state/ok"), +) + + +def stage12_region_groups(country: str) -> list[list[str]] | None: + """Return a mutable copy of the Stage 12 partition when supported.""" + + if country.lower() != "us": + return None + return [list(group) for group in US_REGION_GROUPS] + + +def stage12_region_groups_for_model(country: str, model: Any) -> list[list[str]]: + """Return the partition, retaining newly registered state regions. + + The static partition is measured for the certified dataset. A later model + can register another state-level region before the partition is rebalanced; + assigning it to the final group prevents silently omitting its households. + """ + + groups = stage12_region_groups(country) + if groups is None: + raise ValueError(f"Stage 12 has no national partition for {country!r}") + registry = getattr(model, "region_registry", None) + if registry is None: + return groups + registered = { + region.code + for region in registry.regions + if getattr(region, "region_type", None) == "state" + } + covered = {code for group in groups for code in group} + groups[-1].extend(sorted(registered - covered)) + return groups diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py new file mode 100644 index 000000000..ba4cfb68d --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segment_contracts.py @@ -0,0 +1,77 @@ +"""Strict internal contracts for Stage 12 segment function calls.""" + +from __future__ import annotations + +from hashlib import sha256 +from typing import Annotated, Literal + +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + Sha256Digest, + SimulationRole, +) +from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator + + +class _StrictSegmentModel(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True) + + +class Stage12SegmentInput(_StrictSegmentModel): + """One single-policy region-group calculation.""" + + schema_version: Literal[1] = 1 + simulation: PlannedSimulationExecutionInput + segment_index: Annotated[int, Field(ge=0, lt=20)] + region_codes: Annotated[tuple[str, ...], Field(min_length=1)] + + @field_validator("region_codes") + @classmethod + def require_unique_region_codes(cls, value: tuple[str, ...]) -> tuple[str, ...]: + if len(value) != len(set(value)): + raise ValueError("segment region codes must be unique") + if any(not code.startswith("state/") for code in value): + raise ValueError("segment region codes must identify state regions") + return value + + @model_validator(mode="after") + def require_us_simulation(self) -> Stage12SegmentInput: + if self.simulation.geography.country != "us": + raise ValueError("Stage 12 segments currently support only the US") + if self.simulation.role not in { + SimulationRole.BASELINE, + SimulationRole.REFORM, + }: + raise ValueError("Stage 12 segment role must be baseline or reform") + return self + + +class Stage12SegmentResult(_StrictSegmentModel): + """Compressed, dtype-preserving output from one segment calculation.""" + + schema_version: Literal[1] = 1 + segment_index: Annotated[int, Field(ge=0, lt=20)] + role: SimulationRole + region_codes: Annotated[tuple[str, ...], Field(min_length=1)] + parquet_payload: Annotated[bytes, Field(min_length=1)] + payload_sha256: Sha256Digest + + @field_validator("region_codes") + @classmethod + def require_unique_region_codes(cls, value: tuple[str, ...]) -> tuple[str, ...]: + if len(value) != len(set(value)): + raise ValueError("segment region codes must be unique") + return value + + @field_validator("role") + @classmethod + def require_report_role(cls, value: SimulationRole) -> SimulationRole: + if value not in {SimulationRole.BASELINE, SimulationRole.REFORM}: + raise ValueError("Stage 12 segment role must be baseline or reform") + return value + + @model_validator(mode="after") + def require_payload_digest(self) -> Stage12SegmentResult: + if sha256(self.parquet_payload).hexdigest() != self.payload_sha256: + raise ValueError("segment payload digest differs from its bytes") + return self diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py new file mode 100644 index 000000000..8715d46f2 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/segmentation.py @@ -0,0 +1,192 @@ +"""Independent Stage 12 segment calculation and exact national merge.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from hashlib import sha256 + +import pandas as pd +from policyengine_observability import ObservabilityRuntime +from policyengine_simulation_contract.spm import SPMProvenance +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, +) +from policyengine_simulation_observability.stages import STAGE12_SEGMENT_STAGES + +from policyengine_simulation_executor.stage12_artifacts import ( + deserialize_calculation_provenance, + deserialize_simulation_frames, + serialize_simulation_frames, +) + +from .output_planning import validate_output_frames +from .partition import stage12_region_groups +from .segment_contracts import Stage12SegmentInput, Stage12SegmentResult +from .simulation import SimulationCalculation, calculate_simulation_frames + + +def should_segment_simulation(simulation: PlannedSimulationExecutionInput) -> bool: + """Return whether this logical simulation should use the US partition.""" + + if simulation.geography.country != "us": + return False + if simulation.geography.region.strip().lower() != "us": + return False + if simulation.options.get("segmented") is False: + return False + requirements = simulation.output_plan.requirements + if requirements.include_cliff_impacts: + return False + if requirements.labor_supply_response_active: + return False + return stage12_region_groups("us") is not None + + +def build_segment_inputs( + simulation: PlannedSimulationExecutionInput, + *, + groups: Sequence[Sequence[str]] | None = None, +) -> tuple[Stage12SegmentInput, ...]: + """Build one immutable single-policy input per region group.""" + + selected = groups or stage12_region_groups(simulation.geography.country) + if not selected: + raise ValueError("Stage 12 simulation has no national partition") + if len(selected) != 20: + raise ValueError("Stage 12 national partition must contain 20 groups") + return tuple( + Stage12SegmentInput( + simulation=simulation, + segment_index=index, + region_codes=tuple(group), + ) + for index, group in enumerate(selected) + ) + + +def _scoped_simulation(segment: Stage12SegmentInput) -> PlannedSimulationExecutionInput: + options = { + key: value + for key, value in segment.simulation.options.items() + if key != "segmented" + } + options["region_group"] = list(segment.region_codes) + return segment.simulation.model_copy(update={"options": options}) + + +def calculate_segment( + payload: object, + *, + calculator: Callable[ + [PlannedSimulationExecutionInput], + SimulationCalculation | Mapping[str, pd.DataFrame], + ] = calculate_simulation_frames, + runtime: ObservabilityRuntime | None = None, +) -> Stage12SegmentResult: + """Calculate one segment without persistence or report aggregation.""" + + segment = Stage12SegmentInput.model_validate(payload) + scoped = _scoped_simulation(segment) + if calculator is calculate_simulation_frames: + calculated = calculate_simulation_frames( + scoped, + runtime=runtime, + stage_plan=STAGE12_SEGMENT_STAGES, + ) + else: + calculated = calculator(scoped) + if isinstance(calculated, SimulationCalculation): + frames = calculated.frames + provenance = calculated.calculation_provenance + else: + frames = calculated + provenance = None + validate_output_frames(frames, scoped.output_plan) + parquet_payload, _ = serialize_simulation_frames( + frames, + calculation_provenance=provenance, + ) + return Stage12SegmentResult( + segment_index=segment.segment_index, + role=segment.simulation.role, + region_codes=segment.region_codes, + parquet_payload=parquet_payload, + payload_sha256=sha256(parquet_payload).hexdigest(), + ) + + +def _merge_calculation_provenance( + values: Sequence[SPMProvenance | None], +) -> SPMProvenance | None: + """Keep one shared compact receipt only when every segment agrees.""" + present = [value for value in values if value is not None] + if not present: + return None + if len(present) != len(values): + raise ValueError("segment calculation metadata is incomplete") + receipt = present[0] + if any(value != receipt for value in present[1:]): + raise ValueError("segment calculation provenance does not match") + return receipt + + +def merge_segment_results( + simulation: PlannedSimulationExecutionInput, + payloads: Sequence[object], + *, + groups: Sequence[Sequence[str]] | None = None, +) -> SimulationCalculation: + """Validate and merge all segment payloads into one national result.""" + + expected = build_segment_inputs(simulation, groups=groups) + results = [Stage12SegmentResult.model_validate(payload) for payload in payloads] + if len(results) != len(expected): + raise ValueError( + "Stage 12 segmented calculation returned an incomplete result set" + ) + by_index: dict[int, Stage12SegmentResult] = {} + for result in results: + if result.segment_index in by_index: + raise ValueError( + "Stage 12 segmented calculation returned a duplicate index" + ) + by_index[result.segment_index] = result + + decoded: list[dict[str, pd.DataFrame]] = [] + provenance: list[SPMProvenance | None] = [] + for segment in expected: + result = by_index.get(segment.segment_index) + if result is None: + raise ValueError("Stage 12 segmented calculation omitted a segment") + if result.role is not simulation.role: + raise ValueError("Stage 12 segment result names another simulation role") + if result.region_codes != segment.region_codes: + raise ValueError("Stage 12 segment result names another region group") + decoded.append(deserialize_simulation_frames(result.parquet_payload)) + provenance.append(deserialize_calculation_provenance(result.parquet_payload)) + + expected_entities = set(decoded[0]) + if any(set(frames) != expected_entities for frames in decoded[1:]): + raise ValueError("Stage 12 segment entity schemas do not match") + merged: dict[str, pd.DataFrame] = {} + for entity in sorted(expected_entities): + parts = [frames[entity] for frames in decoded] + first_dtypes = {column: str(dtype) for column, dtype in parts[0].dtypes.items()} + for part in parts[1:]: + if { + column: str(dtype) for column, dtype in part.dtypes.items() + } != first_dtypes: + raise ValueError("Stage 12 segment column dtypes do not match") + frame = pd.concat(parts, ignore_index=True) + identifier = f"{entity}_id" + if identifier not in frame: + raise ValueError(f"Stage 12 segment entity {entity!r} has no identifier") + if frame[identifier].duplicated().any(): + raise ValueError(f"Stage 12 segment entity {entity!r} has duplicate IDs") + merged[entity] = frame + validate_output_frames(merged, simulation.output_plan) + return SimulationCalculation( + frames=merged, + calculation_provenance=_merge_calculation_provenance(provenance), + cache_outcome=None, + ) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py index cd3c9cc30..b661eae03 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_runtime/simulation.py @@ -2,12 +2,15 @@ from __future__ import annotations +import logging from collections.abc import Callable, Mapping +from contextlib import nullcontext from dataclasses import dataclass from datetime import UTC, datetime from functools import partial from hashlib import sha256 from typing import Any +from uuid import UUID import pandas as pd from policyengine_observability import ObservabilityRuntime @@ -26,6 +29,7 @@ from policyengine_simulation_observability.stages import ( STAGE12_SIMULATION_STAGES, Stage, + StagePlan, ) from policyengine_simulation_executor.stage12_artifacts import ( @@ -38,21 +42,85 @@ ) from .dependencies import ComparisonStore, artifact_store, runtime_store -from .failure_causes import ( - StageProgress, - failure_summary, - record_failure_cause, - tracked_stage, +from .failure_causes import StageProgress, record_failure_cause, tracked_stage +from .failure_recording import persist_failure_record, report_failure +from .failures import ( + Stage12Cancellation, + Stage12ExecutionError, + Stage12FailureDetail, + validate_policy_periods, ) -from .failure_recording import persist_failure_record, report_original_failure from .output_planning import apply_output_plan, validate_output_frames +logger = logging.getLogger(__name__) + @dataclass(frozen=True) class SimulationCalculation: frames: Mapping[str, pd.DataFrame] calculation_provenance: SPMProvenance | None = None uk_local_authority_metadata: UKLocalAuthorityMetadata | None = None + cache_outcome: str | None = None + + +def build_stage12_simulation( + params: dict[str, Any], + *, + dataset: Any, + dataset_selection: Any, + country_module: Any, + policy: dict[str, Any] | None, + scoping_strategy: Any, + region_code: str | None, + execution: PlannedSimulationExecutionInput, +): + """Construct a Stage 12 simulation without selecting v1 cache code.""" + + from policyengine.core import Simulation + + from policyengine_simulation_executor.spm import normalize_runtime_spm + + selection = normalize_runtime_spm(params) + from policyengine_simulation_executor.stage12_cache.runtime import ( + Stage12CachedSimulation, + qualifying_cache_identity, + ) + + resolved = load_stage12_bundle() + identity = qualifying_cache_identity( + execution, + resolved=resolved, + dataset_is_default=dataset_selection.is_default, + scoping_strategy=scoping_strategy, + spm=selection, + ) + simulation_type = Stage12CachedSimulation if identity is not None else Simulation + model = simulation_type( + **({"spm": selection} if selection is not None else {}), + **({"id": identity.simulation_id} if identity is not None else {}), + dataset=dataset, + tax_benefit_model_version=country_module.model, + policy=policy, + scoping_strategy=scoping_strategy, + ) + if isinstance(model, Stage12CachedSimulation): + model.configure_stage12_cache(execution.output_plan) + del region_code + return model + + +def simulation_runtime_params( + simulation: PlannedSimulationExecutionInput, +) -> dict[str, Any]: + """Return the parameters a Stage 12 simulation is resolved and built from.""" + + return { + "country": simulation.geography.country, + "scope": "macro", + "time_period": str(simulation.year), + "region": simulation.geography.region, + **simulation.options, + } def simulation_input_sha256(simulation: SimulationExecutionInput) -> str: @@ -116,6 +184,7 @@ def calculate_simulation_frames( simulation: PlannedSimulationExecutionInput, *, runtime: ObservabilityRuntime | None = None, + stage_plan: StagePlan = STAGE12_SIMULATION_STAGES, progress: StageProgress | None = None, ) -> SimulationCalculation: """Run one policy and return the coordinator-planned entity output tables.""" @@ -124,7 +193,7 @@ def calculate_simulation_frames( progress = StageProgress(role=simulation.role.value) stage_span = partial( tracked_stage, - stage_plan=STAGE12_SIMULATION_STAGES, + stage_plan=stage_plan, progress=progress, runtime=runtime, ) @@ -134,15 +203,8 @@ def calculate_simulation_frames( raise ValueError("population artifact differs from bundle dataset provenance") _require_installed_bundle(simulation) country = simulation.geography.country - params: dict[str, Any] = { - "country": country, - "scope": "macro", - "time_period": str(simulation.year), - "region": simulation.geography.region, - **simulation.options, - } + params = simulation_runtime_params(simulation) from policyengine_simulation_executor.simulation_runtime import ( - _build_simulation, _country_module, _load_dataset, _normalise_policy, @@ -186,17 +248,27 @@ def calculate_simulation_frames( with stage_span(Stage.POLICY_NORMALIZATION): policy = _normalise_policy(simulation.policy) with stage_span(Stage.SIMULATION_BUILD): - model = _build_simulation( + model = build_stage12_simulation( params, dataset=dataset, dataset_selection=dataset_selection, + country_module=country_module, policy=policy, scoping_strategy=region.scoping_strategy, region_code=region.code, + execution=simulation, ) with stage_span(Stage.STAGE12_CALCULATION): apply_output_plan(model, simulation.output_plan) - model.ensure() + with ( + stage_span(Stage.STAGE12_CACHE_LOOKUP) + if hasattr(model, "stage12_cache_outcome") + else nullcontext() + ): + model.ensure() + cache_outcome = getattr(model, "stage12_cache_outcome", None) + if runtime is not None and cache_outcome is not None: + runtime.set_context(stage12_cache_outcome=cache_outcome) output_data = getattr(getattr(model, "output_dataset", None), "data", None) entity_data = getattr(output_data, "entity_data", None) if not isinstance(entity_data, Mapping): @@ -218,6 +290,7 @@ def calculate_simulation_frames( frames=frames, calculation_provenance=calculation_provenance, uk_local_authority_metadata=uk_local_authority_metadata, + cache_outcome=cache_outcome, ) @@ -267,6 +340,32 @@ def descriptor_from_record( ) +def _mark_running_simulation_failed( + persistence: ComparisonStore, + simulation_execution_id: UUID, + *, + detail: Stage12FailureDetail, + failed_at: datetime, +) -> None: + """Record a failure only while the simulation is still running.""" + + latest = persistence.get_simulation(simulation_execution_id) + if latest.status is not ComparisonRunLifecycleStatus.RUNNING: + return + persistence.replace_simulation_if_status( + latest.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": detail.error_code, + "error_summary": detail.error_summary, + "updated_at": failed_at, + "completed_at": failed_at, + } + ), + expected_status=ComparisonRunLifecycleStatus.RUNNING, + ) + + def run_single_simulation( payload: object, context_payload: object, @@ -278,6 +377,10 @@ def run_single_simulation( [PlannedSimulationExecutionInput], Mapping[str, pd.DataFrame] | SimulationCalculation, ] = calculate_simulation_frames, + segmented_calculator: Callable[ + [PlannedSimulationExecutionInput, Callable[[], bool]], SimulationCalculation + ] + | None = None, runtime: ObservabilityRuntime | None = None, ) -> dict[str, Any]: simulation = PlannedSimulationExecutionInput.model_validate(payload) @@ -288,6 +391,11 @@ def run_single_simulation( child = persistence.get_simulation(simulation.simulation_execution_id) if child.status is ComparisonRunLifecycleStatus.SUCCEEDED: return descriptor_from_record(child, simulation).model_dump(mode="json") + if child.status not in { + ComparisonRunLifecycleStatus.PENDING, + ComparisonRunLifecycleStatus.RUNNING, + }: + raise RuntimeError("Stage 12 simulation is no longer active") started = datetime.now(UTC) running = child.model_copy( update={ @@ -298,7 +406,25 @@ def run_single_simulation( "error_summary": None, } ) - persistence.replace_simulation(running) + running, claimed = persistence.replace_simulation_if_status( + running, + expected_status=child.status, + ) + if not claimed: + raise RuntimeError("Stage 12 simulation was cancelled") + + def cancellation_requested() -> bool: + try: + latest = persistence.get_simulation(simulation.simulation_execution_id) + except Exception as error: # noqa: BLE001 + logger.warning( + "Stage 12 cancellation check failed for %s (%s)", + simulation.simulation_execution_id, + type(error).__name__, + ) + return False + return latest.status is not ComparisonRunLifecycleStatus.RUNNING + progress = StageProgress(role=simulation.role.value) stage_span = partial( tracked_stage, @@ -307,12 +433,17 @@ def run_single_simulation( runtime=runtime, ) try: + validate_policy_periods(simulation.policy) with stage_span(Stage.STAGE12_INPUT_WRITE): artifact_storage.write_input( prefix=context.artifact_prefix, simulation=simulation, ) - if calculator is calculate_simulation_frames: + from .segmentation import should_segment_simulation + + if segmented_calculator is not None and should_segment_simulation(simulation): + calculated = segmented_calculator(simulation, cancellation_requested) + elif calculator is calculate_simulation_frames: calculated = calculate_simulation_frames( simulation, runtime=runtime, @@ -339,7 +470,7 @@ def run_single_simulation( uk_local_authority_metadata=uk_local_authority_metadata, ) completed = datetime.now(UTC) - persistence.replace_simulation( + _, replaced = persistence.replace_simulation_if_status( running.model_copy( update={ "status": ComparisonRunLifecycleStatus.SUCCEEDED, @@ -352,33 +483,43 @@ def run_single_simulation( "updated_at": completed, "completed_at": completed, } - ) + ), + expected_status=ComparisonRunLifecycleStatus.RUNNING, ) + if not replaced: + raise Stage12Cancellation("Stage 12 simulation was cancelled") return descriptor.model_dump(mode="json") + except Stage12Cancellation: + raise RuntimeError("Stage 12 simulation was cancelled") from None # Persist a redacted cause for any country-package calculation failure # before returning a stable exception to Modal. except Exception as error: # noqa: BLE001 - report_original_failure( - error, - runtime=runtime, - failure_scope="simulation_execution", - attributes={"simulation_role": simulation.role.value}, - ) - execution_failure_summary = failure_summary(error, progress=progress) - record_failure_cause(runtime, error, execution_failure_summary) + if isinstance(error, Stage12ExecutionError): + # A segmented calculation reports its segment failure where it + # occurs, before raising this error. + detail = error.detail + else: + detail = report_failure( + error, + runtime=runtime, + failure_scope="simulation_execution", + default_code="simulation_execution_failed", + attributes={ + "simulation_role": simulation.role.value, + "evaluation_id": str(simulation.evaluation_id), + "simulation_execution_id": str(simulation.simulation_execution_id), + }, + progress=progress, + ) + record_failure_cause(runtime, error, detail.error_summary) failed_at = datetime.now(UTC) persist_failure_record( partial( - persistence.replace_simulation, - running.model_copy( - update={ - "status": ComparisonRunLifecycleStatus.FAILED, - "error_code": "simulation_execution_failed", - "error_summary": execution_failure_summary, - "updated_at": failed_at, - "completed_at": failed_at, - } - ), + _mark_running_simulation_failed, + persistence, + simulation.simulation_execution_id, + detail=detail, + failed_at=failed_at, ), runtime=runtime, failure_scope="simulation_execution", diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_timeouts.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_timeouts.py new file mode 100644 index 000000000..0f7385517 --- /dev/null +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_timeouts.py @@ -0,0 +1,41 @@ +"""Modal timeouts and waits shared by the Stage 12 app and report coordinator. + +This module has no imports, so the Modal app declaration can read it without +loading the Stage 12 runtime. +""" + +# Modal timeouts of the simulation workers that the report coordinator starts. +# A US worker starts 20 region-group calls (3,600 s each), then merges their +# results and writes one artifact. +US_SIMULATION_TIMEOUT_SECONDS = 3_900 +UK_SIMULATION_TIMEOUT_SECONDS = 3_000 + +# One child-wait window opens when `coordinate_report` starts. It covers +# coordinator setup (claiming, output planning including the country model +# load, child-state creation, retained-artifact reads and dispatch) and the +# concurrent wait for both simulations, so setup time shortens the children's +# window instead of the time left after it. The window equals the longest +# simulation worker timeout, so the coordinator only stops waiting for a worker +# early by its own setup time and the time the worker spent queued in Modal. +# The coordinator's last poll never sleeps past the end of the window. +SIMULATION_WAIT_TIMEOUT_SECONDS = max( + US_SIMULATION_TIMEOUT_SECONDS, + UK_SIMULATION_TIMEOUT_SECONDS, +) + +# An automatic (shadow) report then waits up to 15 minutes for the production +# result it compares against. +PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS = 900 + +# Time left after both waits to read and validate both simulation artifacts, +# aggregate, write the aggregate and comparison artifacts, and record the +# outcome or failure. A complete testing-environment US report, including all +# 40 region-group calculations, took 288 s. +REPORT_FINISHING_SECONDS = 600 + +# 3,900 s + 900 s + 600 s = 5,400 s. +COORDINATOR_TIMEOUT_SECONDS = ( + SIMULATION_WAIT_TIMEOUT_SECONDS + + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + + REPORT_FINISHING_SECONDS +) diff --git a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py index 008bfae40..3ef087e41 100644 --- a/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/src/policyengine_simulation_executor/stage12_worker_validation.py @@ -25,8 +25,13 @@ ), ("STAGE12_DATABASE_URL",), ("STAGE12_ARTIFACT_BUCKET",), + ("STAGE12_CACHE_BUCKET",), ) +US_VALIDATION_REFORM = { + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01.2100-12-31": 3_000} +} + def _require_secret_environment(environment: Mapping[str, str]) -> None: for alternatives in REQUIRED_SECRET_ALTERNATIVES: @@ -55,16 +60,40 @@ def _check_dataset_access(path: str, expected_sha256: str) -> None: def _run_non_serving_calculation(country: CountryId) -> None: country_module = import_module(f"policyengine.tax_benefit_models.{country}") - spm = {"geography_kind": "national"} if country == "us" else None + people = [{"age": 40, "employment_income": 50_000}] + extra_variables = ["household_net_income"] + options: dict[str, Any] = {} + if country == "us": + from policyengine_simulation_executor.simulation_runtime import ( + _normalise_policy, + ) + + people = [ + { + "age": 40, + "employment_income": 50_000, + "is_tax_unit_head": True, + }, + {"age": 10, "is_tax_unit_dependent": True}, + ] + extra_variables.append("ctc") + options["spm"] = {"geography_kind": "national"} + options["reform"] = _normalise_policy(US_VALIDATION_REFORM) result = country_module.calculate_household( - people=[{"age": 40, "employment_income": 50_000}], - **({"spm": spm} if spm is not None else {}), + people=people, year=2026, - extra_variables=["household_net_income"], + extra_variables=extra_variables, + **options, ) household = getattr(result, "household", None) if not isinstance(household, Mapping) or "household_net_income" not in household: raise RuntimeError("Stage 12 validation calculation returned no net income") + if country == "us": + tax_unit = getattr(result, "tax_unit", None) + if not isinstance(tax_unit, Mapping) or not isinstance( + tax_unit.get("ctc"), (int, float) + ): + raise RuntimeError("Stage 12 validation reform returned no CTC value") def _check_uk_local_authority_resources() -> None: @@ -95,10 +124,58 @@ def _check_uk_local_authority_dataset(dataset_path: str) -> None: detect_uk_local_authority_metadata_from_hdf(dataset_path) +def _check_installed_cache(expected_manifest_sha256: str) -> dict[str, Any]: + """Validate the baked manifest and require one real baseline cache hit.""" + + from policyengine_simulation_executor.simulation_runtime import resolve_data_folder + from policyengine_simulation_executor.stage12_cache.models import CacheManifest + from policyengine_simulation_executor.stage12_cache.precompute import ( + _scoped_planned_simulation, + ) + from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, + ) + from policyengine_simulation_executor.stage12_runtime.simulation import ( + calculate_simulation_frames, + ) + + data_folder = Path(resolve_data_folder()) + manifest_path = data_folder / ".stage12-cache-manifest.json" + payload = manifest_path.read_bytes() + if sha256(payload).hexdigest() != expected_manifest_sha256: + raise RuntimeError("installed Stage 12 cache manifest digest differs") + manifest = CacheManifest.model_validate_json(payload) + if manifest.years != (2026, 2027, 2025): + raise RuntimeError("installed Stage 12 cache years differ") + if len(manifest.artifacts) != 63: + raise RuntimeError("installed Stage 12 cache artifact count differs") + probe = next( + item + for item in manifest.artifacts + if item.type == "baseline" and item.year == 2026 + ) + probe_path = data_folder / probe.filename + _check_dataset_access(str(probe_path), probe.content_sha256) + calculation = calculate_simulation_frames( + _scoped_planned_simulation(2026, US_REGION_GROUPS[0]) + ) + if calculation.cache_outcome != "hit": + raise RuntimeError("deployed Stage 12 cache probe did not load a cache hit") + return { + "manifest_sha256": expected_manifest_sha256, + "partition_sha256": manifest.partition_sha256, + "years": list(manifest.years), + "segment_count": 20, + "dataset_count": sum(item.type == "dataset" for item in manifest.artifacts), + "baseline_count": sum(item.type == "baseline" for item in manifest.artifacts), + } + + def validate_country_worker( *, country: CountryId, expected_bundle_manifest_sha256: str, + expected_cache_manifest_sha256: str | None = None, environment: Mapping[str, str] | None = None, dataset_path_resolver: Callable[[CountryId], str | None] = ( _resolve_installed_dataset_path @@ -108,6 +185,7 @@ def validate_country_worker( local_authority_resource_check: Callable[[str], None] = ( _check_uk_local_authority_dataset ), + cache_check: Callable[[str], dict[str, Any]] = _check_installed_cache, ) -> dict[str, Any]: """Validate imports, installed bundle, secrets, data, and calculation.""" @@ -145,6 +223,9 @@ def validate_country_worker( if country == "uk": local_authority_resource_check(installed_dataset_path) calculation_check(country) + cache = None + if country == "us" and expected_cache_manifest_sha256: + cache = cache_check(expected_cache_manifest_sha256) return { "validated": True, "country": country, @@ -155,4 +236,5 @@ def validate_country_worker( "dataset_uri": country_bundle.default_dataset_uri, "data_artifact_revision": country_bundle.data_artifact_revision, "bundle_manifest_sha256": resolved.bundle_manifest_sha256, + "cache": cache, } diff --git a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py index 3cbe1ab1d..d7b4053c0 100644 --- a/projects/policyengine-simulation-executor/tests/test_modal_scripts.py +++ b/projects/policyengine-simulation-executor/tests/test_modal_scripts.py @@ -224,6 +224,9 @@ def test_deployment_sources_exclude_legacy_hf_secret_names(self): REPO_ROOT / "projects/policyengine-simulation-executor/src/modal/v2_app.py", REPO_ROOT / "projects/policyengine-simulation-executor/src/modal/precompute_app.py", + REPO_ROOT + / "projects/policyengine-simulation-executor/src/modal" + / "stage12_precompute_app.py", ) sources = "\n".join(path.read_text(encoding="utf-8") for path in paths) @@ -754,6 +757,94 @@ def test_precompute_is_gated_by_the_shared_secret_sync(self): assert "needs: prepare" in executor_job +class TestStage12CachePrecompute: + """Tests for the independent Stage 12 cache deployment script.""" + + script = SCRIPTS_DIR / "stage12-cache-precompute.sh" + + def _run_with_fake_uv(self, tmp_path, output: str, *args): + bin_dir = tmp_path / "bin" + bin_dir.mkdir() + log_path = tmp_path / "uv-calls.log" + github_output = tmp_path / "github-output.txt" + uv_path = bin_dir / "uv" + uv_path.write_text( + '#!/bin/bash\nprintf \'%s\\n\' "$*" >> "$UV_FAKE_LOG"\nprintf \'%s\\n\' "$UV_FAKE_OUTPUT"\n', + encoding="utf-8", + ) + uv_path.chmod(0o755) + env = { + **os.environ, + "PATH": f"{bin_dir}{os.pathsep}{os.environ['PATH']}", + "UV_FAKE_LOG": str(log_path), + "UV_FAKE_OUTPUT": output, + "GITHUB_OUTPUT": str(github_output), + "STAGE12_CACHE_BUCKET": "policyengine-stage12-cache-staging", + } + result = subprocess.run( + ["bash", str(self.script), "staging", *args], + capture_output=True, + text=True, + env=env, + ) + calls = log_path.read_text(encoding="utf-8") if log_path.exists() else "" + emitted = ( + github_output.read_text(encoding="utf-8") if github_output.exists() else "" + ) + return result, calls, emitted + + def test_validates_syntax_and_exports_exact_manifest_digest(self, tmp_path): + syntax = subprocess.run( + ["bash", "-n", str(self.script)], + capture_output=True, + text=True, + ) + assert syntax.returncode == 0, syntax.stderr + digest = "a" * 64 + + result, calls, emitted = self._run_with_fake_uv( + tmp_path, + f"progress\nSTAGE12_CACHE_MANIFEST_DIGEST={digest}", + ) + + assert result.returncode == 0, result.stderr + assert ( + "run modal run --env=staging src/modal/stage12_precompute_app.py" in calls + ) + assert "--force" not in calls + assert emitted == f"manifest_digest={digest}\n" + + def test_rejects_an_invalid_digest(self, tmp_path): + result, calls, emitted = self._run_with_fake_uv( + tmp_path, + "STAGE12_CACHE_MANIFEST_DIGEST=not-a-digest", + ) + + assert result.returncode != 0 + assert calls.strip() == ( + "run modal run --env=staging src/modal/stage12_precompute_app.py" + ) + assert emitted == "" + + def test_v1_force_recompute_never_reaches_the_immutable_cache(self): + """Stored Stage 12 cache files cannot be replaced, so nothing forces + a rebuild of them.""" + reusable_workflow = ( + REPO_ROOT / ".github" / "workflows" / "simulation-deploy.reusable.yml" + ).read_text(encoding="utf-8") + cache_step = next( + line + for line in reusable_workflow.splitlines() + if "stage12-cache-precompute.sh" in line + ) + + assert cache_step.strip() == ( + "run: ../../.github/scripts/stage12-cache-precompute.sh " + '"${{ inputs.modal_environment }}"' + ) + assert "--force" not in self.script.read_text(encoding="utf-8") + + class TestModalRecordDeployment: """Tests for modal-record-deployment.sh""" diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py new file mode 100644 index 000000000..59e4856ef --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_image.py @@ -0,0 +1,107 @@ +"""Tests for the self-contained Stage 12 image cache fetch.""" + +from __future__ import annotations + +import json +from hashlib import sha256 +from pathlib import Path +from types import SimpleNamespace + +import pytest +from src.modal.stage12_cache_image import fetch_stage12_cache + + +class _Blob: + def __init__(self, payload: bytes): + self.payload = payload + + def download_to_filename(self, filename: str) -> None: + Path(filename).write_bytes(self.payload) + + +class _Client: + def __init__(self, payloads: dict[str, bytes]): + self.payloads = payloads + + def bucket(self, _name: str): + return SimpleNamespace(blob=lambda path: _Blob(self.payloads[path])) + + +def _manifest() -> tuple[dict, dict[str, bytes]]: + artifacts = [] + payloads = {} + for kind, count in (("dataset", 3), ("baseline", 60)): + for index in range(count): + filename = f"{kind}-{index}.h5" + path = f"{kind}s/us/{index}/{filename}" + payload = f"{kind}-{index}".encode() + payloads[path] = payload + artifacts.append( + { + "type": kind, + "path": path, + "filename": filename, + "year": (2026, 2027, 2025)[index % 3], + "identity_digest": f"identity-{kind}-{index}", + "content_sha256": sha256(payload).hexdigest(), + "size_bytes": len(payload), + } + ) + return ( + { + "schema": "stage12-mf1", + "country": "us", + "years": [2026, 2027, 2025], + "partition_sha256": "p" * 64, + "bundle": { + "bundle_manifest_sha256": "b" * 64, + "policyengine_version": "5.2.0", + "core_package_version": "3.0.0", + "country_package_version": "1.0.0", + "data_package_version": "2.0.0", + "data_artifact_revision": "revision", + "default_dataset": "populace_us_2024", + }, + "artifacts": artifacts, + }, + payloads, + ) + + +def _digest(manifest: dict) -> str: + return sha256( + json.dumps( + manifest, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + allow_nan=False, + ).encode() + ).hexdigest() + + +def test_fetches_and_verifies_all_cache_files(monkeypatch, tmp_path) -> None: + manifest, payloads = _manifest() + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + monkeypatch.setenv("STAGE12_BUNDLE_MANIFEST_SHA256", "b" * 64) + monkeypatch.setenv("STAGE12_CACHE_MANIFEST_DIGEST", _digest(manifest)) + + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + assert len(list(tmp_path.glob("*.h5"))) == 63 + assert (tmp_path / ".stage12-cache-manifest.json").exists() + + +def test_rejects_foreign_or_corrupt_cache(monkeypatch, tmp_path) -> None: + manifest, payloads = _manifest() + monkeypatch.setenv("POLICYENGINE_DATA_FOLDER", str(tmp_path)) + monkeypatch.setenv("STAGE12_BUNDLE_MANIFEST_SHA256", "different") + monkeypatch.setenv("STAGE12_CACHE_MANIFEST_DIGEST", _digest(manifest)) + with pytest.raises(RuntimeError, match="another installed bundle"): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) + + monkeypatch.setenv("STAGE12_BUNDLE_MANIFEST_SHA256", "b" * 64) + corrupt_path = manifest["artifacts"][0]["path"] + payloads[corrupt_path] = b"corrupt" + with pytest.raises(RuntimeError, match="size differs|digest differs"): + fetch_stage12_cache("cache-bucket", manifest, client=_Client(payloads)) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py new file mode 100644 index 000000000..dec3f0258 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_precompute.py @@ -0,0 +1,496 @@ +"""Unit tests for Stage 12 cache planning and orchestration.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pandas as pd +import pytest +from google.api_core.exceptions import NotFound, PreconditionFailed +from policyengine_simulation_executor import simulation_runtime +from policyengine_simulation_executor.stage12_cache import precompute +from policyengine_simulation_executor.stage12_cache.models import ( + BaselinePlanEntry, + CacheBundleIdentity, + CacheObject, + DatasetPlanEntry, + PrecomputePlan, +) +from policyengine_simulation_executor.stage12_cache.store import Stage12CacheStore +from policyengine_simulation_executor.stage12_runtime import ( + simulation as simulation_module, +) +from policyengine_simulation_executor.stage12_runtime.partition import US_REGION_GROUPS +from policyengine_simulation_executor.stage12_runtime.simulation import ( + SimulationCalculation, +) +from pydantic import ValidationError + + +def _bundle() -> CacheBundleIdentity: + return CacheBundleIdentity( + bundle_manifest_sha256="a" * 64, + policyengine_version="5.2.0", + core_package_version="3.0.0", + country_package_version="1.0.0", + data_package_version="2.0.0", + data_artifact_revision="revision", + default_dataset="populace_us_2024", + ) + + +def _plan(*, exists: bool = False, verified: bool | None = None) -> PrecomputePlan: + datasets = tuple( + DatasetPlanEntry( + year=year, + digest=f"dataset-{year}", + path=f"datasets/us/dataset-{year}/populace_year_{year}.h5", + filename=f"populace_year_{year}.h5", + exists=exists, + ) + for year in precompute.PRECOMPUTE_YEARS + ) + baselines = tuple( + BaselinePlanEntry( + year=year, + segment_index=index, + region_codes=group, + digest=f"baseline-{year}-{index}", + path=f"baselines/us/baseline-{year}-{index}/stage12-bl1-{index}.h5", + filename=f"stage12-bl1-{index}.h5", + simulation_id=f"stage12-bl1-{index}", + output_plan_sha256="b" * 64, + exists=exists, + ) + for year in precompute.PRECOMPUTE_YEARS + for index, group in enumerate(US_REGION_GROUPS) + ) + return PrecomputePlan( + years=precompute.PRECOMPUTE_YEARS, + datasets=datasets, + baselines=baselines, + bundle=_bundle(), + partition_sha256="c" * 64, + # A warm store came from a run whose determinism check passed. + determinism_verified=exists if verified is None else verified, + ) + + +class _FakeBlob: + """A GCS blob over a dict, honouring create-only uploads.""" + + def __init__(self, objects: dict, path: str): + self.objects = objects + self.path = path + self.metadata = None + + def exists(self) -> bool: + return self.path in self.objects + + def reload(self) -> None: + if self.path not in self.objects: + raise NotFound(self.path) + self.metadata = self.objects[self.path]["metadata"] + + def upload_from_string(self, payload, *, content_type, if_generation_match): + assert if_generation_match == 0 + if self.path in self.objects: + raise PreconditionFailed(self.path) + self.objects[self.path] = {"payload": payload, "metadata": self.metadata} + + def download_as_bytes(self) -> bytes: + if self.path not in self.objects: + raise NotFound(self.path) + return self.objects[self.path]["payload"] + + +class _FakeClient: + def __init__(self): + self.objects: dict[str, dict] = {} + + def bucket(self, _name: str): + return SimpleNamespace(blob=lambda path: _FakeBlob(self.objects, path)) + + +def _use_store(monkeypatch, client: _FakeClient) -> None: + monkeypatch.setattr( + precompute, + "Stage12CacheStore", + lambda bucket: Stage12CacheStore(bucket, client=client), + ) + + +def _store_every_planned_file(client: _FakeClient, plan: PrecomputePlan) -> None: + for entry in (*plan.datasets, *plan.baselines): + client.objects[entry.path] = { + "payload": b"h5", + "metadata": {"sha256": "f" * 64, "size_bytes": "2"}, + } + + +def test_plan_shape_matches_v1_year_window_and_twenty_groups() -> None: + plan = _plan() + + assert plan.years == (2026, 2027, 2025) + assert len(plan.datasets) == 3 + assert len(plan.baselines) == 60 + assert [entry.year for entry in plan.datasets] == [2026, 2027, 2025] + assert { + entry.segment_index for entry in plan.baselines if entry.year == 2026 + } == set(range(20)) + + +def test_strict_plan_rejects_unknown_modal_boundary_fields() -> None: + payload = _plan().model_dump() + payload["unexpected"] = True + + with pytest.raises(ValidationError): + PrecomputePlan.model_validate(payload) + + +def test_work_selection_computes_only_missing_files() -> None: + cold = precompute.select_work(_plan()) + warm = precompute.select_work(_plan(exists=True)) + + assert len(cold.datasets) == 3 and len(cold.baselines) == 60 + assert warm.datasets == () and warm.baselines == () + + +class _Call: + def __init__(self, result, events, label): + self.result = result + self.events = events + self.label = label + + def get(self): + self.events.append(f"get:{self.label}") + return self.result + + +class _Function: + def __init__(self, callback, events, label): + self.callback = callback + self.events = events + self.label = label + self.spawn_calls = [] + self.remote_calls = [] + + def spawn(self, *args): + self.spawn_calls.append(args) + self.events.append(f"spawn:{self.label}") + return _Call(self.callback(*args), self.events, self.label) + + def remote(self, *args): + self.remote_calls.append(args) + self.events.append(f"remote:{self.label}") + return self.callback(*args) + + +def test_precompute_finishes_dataset_wave_before_baseline_wave() -> None: + events = [] + plan = _plan() + planner = _Function(lambda _: plan.model_dump(), events, "plan") + dataset = _Function( + lambda _, entry: { + "path": entry["path"], + "content_sha256": "d" * 64, + "size_bytes": 1, + "year": entry["year"], + "uploaded": True, + "build_seconds": 1.0, + }, + events, + "dataset", + ) + baseline = _Function( + lambda _, entry: { + "path": entry["path"], + "content_sha256": "e" * 64, + "size_bytes": 1, + "year": entry["year"], + "segment_index": entry["segment_index"], + "simulation_id": entry["simulation_id"], + "cache_outcome": "miss", + "uploaded": True, + "compute_seconds": 1.0, + }, + events, + "baseline", + ) + verify = _Function(lambda *_: {"equal": True, "differences": []}, events, "verify") + publish = _Function(lambda *_: "manifest-digest", events, "publish") + lines = [] + + digest = precompute.run_precompute( + "cache-bucket", + plan_artifacts=planner, + build_dataset=dataset, + compute_baseline=baseline, + verify_determinism=verify, + publish_manifest=publish, + echo=lines.append, + ) + + first_baseline_spawn = events.index("spawn:baseline") + assert all( + index < first_baseline_spawn + for index, event in enumerate(events) + if event == "get:dataset" + ) + assert len(dataset.spawn_calls) == 3 + assert len(baseline.spawn_calls) == 60 + assert verify.remote_calls == [ + ("cache-bucket", plan.model_dump(), plan.baselines[0].model_dump()) + ] + assert digest == "manifest-digest" + assert lines[-1] == "STAGE12_CACHE_MANIFEST_DIGEST=manifest-digest" + + +def test_warm_verified_precompute_skips_builds_and_determinism_check() -> None: + events = [] + plan = _plan(exists=True) + planner = _Function(lambda _: plan.model_dump(), events, "plan") + unused = _Function(lambda *_: pytest.fail("unexpected build"), events, "build") + verify = _Function( + lambda *_: pytest.fail("unexpected validation"), events, "verify" + ) + publish = _Function(lambda *_: "manifest-digest", events, "publish") + + precompute.run_precompute( + "cache-bucket", + plan_artifacts=planner, + build_dataset=unused, + compute_baseline=unused, + verify_determinism=verify, + publish_manifest=publish, + ) + + assert unused.spawn_calls == [] + assert verify.remote_calls == [] + + +def test_failed_determinism_validation_prevents_manifest_publication() -> None: + events = [] + plan = _plan() + planner = _Function(lambda _: plan.model_dump(), events, "plan") + dataset = _Function( + lambda _, entry: { + "path": entry["path"], + "content_sha256": "d" * 64, + "size_bytes": 1, + "year": entry["year"], + "uploaded": True, + "build_seconds": 1.0, + }, + events, + "dataset", + ) + baseline = _Function( + lambda _, entry: { + "path": entry["path"], + "content_sha256": "e" * 64, + "size_bytes": 1, + "year": entry["year"], + "segment_index": entry["segment_index"], + "simulation_id": entry["simulation_id"], + "cache_outcome": "miss", + "uploaded": True, + "compute_seconds": 1.0, + }, + events, + "baseline", + ) + verify = _Function( + lambda *_: {"equal": False, "differences": ["person.age differs"]}, + events, + "verify", + ) + publish = _Function(lambda *_: "unexpected", events, "publish") + + with pytest.raises(SystemExit, match="person.age differs"): + precompute.run_precompute( + "cache-bucket", + plan_artifacts=planner, + build_dataset=dataset, + compute_baseline=baseline, + verify_determinism=verify, + publish_manifest=publish, + ) + + assert publish.remote_calls == [] + + +def test_manifest_records_content_hashes_and_sizes() -> None: + plan = _plan(exists=True) + store = SimpleNamespace( + describe=lambda path: CacheObject( + path=path, + content_sha256="f" * 64, + size_bytes=123, + ) + ) + + manifest = precompute.build_manifest(plan, store) + + assert len(manifest.artifacts) == 63 + assert all(item.content_sha256 == "f" * 64 for item in manifest.artifacts) + assert all(item.size_bytes == 123 for item in manifest.artifacts) + + +@pytest.mark.parametrize("passes", [True, False]) +def test_a_run_with_no_new_baselines_still_checks_unverified_ones( + passes: bool, +) -> None: + events = [] + plan = _plan(exists=True, verified=False) + planner = _Function(lambda _: plan.model_dump(), events, "plan") + unused = _Function(lambda *_: pytest.fail("unexpected build"), events, "build") + verify = _Function( + lambda *_: { + "equal": passes, + "differences": [] if passes else ["person.age: values differ"], + }, + events, + "verify", + ) + publish = _Function(lambda *_: "manifest-digest", events, "publish") + + def run() -> None: + precompute.run_precompute( + "cache-bucket", + plan_artifacts=planner, + build_dataset=unused, + compute_baseline=unused, + verify_determinism=verify, + publish_manifest=publish, + echo=lambda _line: None, + ) + + if passes: + run() + assert len(publish.remote_calls) == 1 + else: + with pytest.raises(SystemExit, match="person.age"): + run() + assert publish.remote_calls == [] + assert verify.remote_calls == [ + ("cache-bucket", plan.model_dump(), plan.baselines[0].model_dump()) + ] + + +def test_publish_requires_a_recorded_determinism_pass(monkeypatch) -> None: + client = _FakeClient() + _use_store(monkeypatch, client) + plan = _plan(exists=True, verified=False) + _store_every_planned_file(client, plan) + + with pytest.raises(RuntimeError, match="have not passed determinism"): + precompute.publish_manifest_impl("cache-bucket", plan) + assert not any(path.startswith("manifests/") for path in client.objects) + + Stage12CacheStore("cache-bucket", client=client).write_json( + precompute.determinism_record_path(plan.baselines), + precompute.determinism_record(plan.baselines), + ) + digest = precompute.publish_manifest_impl("cache-bucket", plan) + + assert f"manifests/{digest}.json" in client.objects + + +def _frames(age: int) -> dict[str, pd.DataFrame]: + return {"person": pd.DataFrame({"person_id": [1], "age": [age]})} + + +@pytest.mark.parametrize( + ("cached_outcome", "fresh_age", "expected"), + [ + ("hit", 30, "recorded"), + ("hit", 31, "not recorded"), + ("miss", 30, "rejected"), + ], +) +def test_determinism_check_records_only_a_pass_that_loaded_the_stored_file( + monkeypatch, tmp_path, cached_outcome: str, fresh_age: int, expected: str +) -> None: + client = _FakeClient() + _use_store(monkeypatch, client) + plan = _plan(exists=True, verified=False) + probe = plan.baselines[0] + (tmp_path / probe.filename).write_bytes(b"h5") + monkeypatch.setattr(simulation_runtime, "resolve_data_folder", lambda: tmp_path) + monkeypatch.setattr(precompute, "_ensure_dataset_file", lambda *_: None) + monkeypatch.setattr( + precompute, + "_scoped_planned_simulation", + lambda _year, _group, *, reform=False: "reform" if reform else "baseline", + ) + calculations = { + "baseline": SimulationCalculation( + frames=_frames(30), cache_outcome=cached_outcome + ), + "reform": SimulationCalculation(frames=_frames(fresh_age)), + } + monkeypatch.setattr( + simulation_module, + "calculate_simulation_frames", + lambda planned: calculations[planned], + ) + record_path = precompute.determinism_record_path(plan.baselines) + + if expected == "rejected": + with pytest.raises(RuntimeError, match="did not load the stored baseline"): + precompute.verify_determinism_impl("cache-bucket", plan, probe) + else: + verdict = precompute.verify_determinism_impl("cache-bucket", plan, probe) + assert verdict.equal is (expected == "recorded") + assert (record_path in client.objects) is (expected == "recorded") + + +def test_planned_baseline_file_is_the_one_the_worker_saves_and_loads( + monkeypatch, +) -> None: + """Uses the real bundle, US model and .py naming: no stubbed ensure().""" + + from policyengine.core import Dataset + from policyengine.tax_benefit_models.common.model_version import ( + output_dataset_filepath, + ) + from policyengine_simulation_executor.stage12_cache.runtime import ( + Stage12CachedSimulation, + ) + + client = _FakeClient() + _use_store(monkeypatch, client) + plan = precompute.plan_artifacts_impl("cache-bucket") + entry = plan.baselines[0] + planned = precompute._scoped_planned_simulation(entry.year, entry.region_codes) + params = simulation_module.simulation_runtime_params(planned) + country_module = simulation_runtime._country_module("us") + region = simulation_runtime._resolve_region( + country_module=country_module, country="us", params=params + ) + dataset_path = "/opt/policyengine/stage12-data/populace_year_2026.h5" + + worker_model = simulation_module.build_stage12_simulation( + params, + dataset=Dataset( + name="in-memory", + description="Naming only", + year=entry.year, + filepath=dataset_path, + ), + dataset_selection=SimpleNamespace(is_default=True), + country_module=country_module, + policy=None, + scoping_strategy=region.scoping_strategy, + region_code=region.code, + execution=planned, + ) + + assert isinstance(worker_model, Stage12CachedSimulation) + assert worker_model.id == entry.simulation_id + assert output_dataset_filepath(worker_model).name == entry.filename + assert entry.path.endswith(f"/{entry.filename}") + # A US simulation always resolves an SPM selection, which .py names. + assert entry.filename != f"{entry.simulation_id}.h5" + assert plan.determinism_verified is False diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py new file mode 100644 index 000000000..e7c013e00 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_cache_runtime.py @@ -0,0 +1,290 @@ +"""Unit tests for independent Stage 12 cache identities and loading.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pandas as pd +from policyengine.core import Simulation +from policyengine_simulation_contract.stage12_bundle import ( + ResolvedStage12Bundle, + Stage12BundleManifest, + Stage12CountryBundle, + Stage12Dataset, +) +from policyengine_simulation_contract.stage12_execution import SimulationRole +from policyengine_simulation_executor.stage12_cache.keys import ( + canonical_digest, + collect_baseline_identity, + collect_dataset_identity, +) +from policyengine_simulation_executor.stage12_cache.runtime import ( + CACHE_HIT, + CACHE_INCOMPLETE, + CACHE_MISS, + Stage12CachedSimulation, + qualifying_cache_identity, +) +from test_stage12_runtime import _planned_simulation + +# A resolved selection, as normalize_runtime_spm returns for every US request. +SPM_SELECTION = { + "forecast_content_sha256": "e" * 64, + "scenario": "ce_trend", + "geography_kind": "county", + "geography_id": None, + "county_vintage": "2020", + "as_of": None, +} + + +def _resolved() -> ResolvedStage12Bundle: + dataset = Stage12Dataset( + identity="populace_us_2024", + uri="hf://policyengine/data/populace_us_2024.h5@revision", + artifact_revision="revision", + sha256="a" * 64, + repo_type="dataset", + ) + country = Stage12CountryBundle( + country="us", + country_package_name="policyengine-us", + country_package_version="1.0.0", + country_package_requirement="policyengine-us==1.0.0", + data_package_name="policyengine-us-data", + data_package_version="2.0.0", + data_release_version="2.0.0", + data_artifact_revision="revision", + default_dataset=dataset.identity, + default_dataset_uri=dataset.uri, + datasets=(dataset,), + ) + return ResolvedStage12Bundle( + bundle=Stage12BundleManifest( + policyengine_version="5.2.0", + policyengine_requirement="policyengine==5.2.0", + core_package_version="3.0.0", + core_package_requirement="policyengine-core==3.0.0", + countries=(country,), + ), + bundle_manifest_sha256="b" * 64, + ) + + +def _strategy(): + from policyengine.core.scoping_strategy import ( + RegionGroupStrategy, + RowFilterStrategy, + ) + + return RegionGroupStrategy( + members=[RowFilterStrategy(variable_name="state_code", variable_value="CA")] + ) + + +def test_cache_writer_and_reader_share_full_identity() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={"options": {"region_group": ["state/ca"]}} + ) + resolved = _resolved() + strategy = _strategy() + + reader = qualifying_cache_identity( + simulation, + resolved=resolved, + dataset_is_default=True, + scoping_strategy=strategy, + spm=SPM_SELECTION, + ) + writer = collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=("state/ca",), + scope_key=strategy.cache_key, + spm=SPM_SELECTION, + ) + + assert reader == writer + assert reader is not None + assert reader.simulation_id.startswith("stage12-bl1-") + # Named as .py's Simulation.storage_id names the file it saves and loads. + assert reader.filename == ( + f"{reader.simulation_id}-spm-{canonical_digest(SPM_SELECTION)}.h5" + ) + assert reader.store_path.endswith(f"/{reader.filename}") + + +def test_dataset_and_baseline_inputs_rotate_cache_identity() -> None: + resolved = _resolved() + first = collect_dataset_identity(resolved, country="us", year=2026) + later = collect_dataset_identity(resolved, country="us", year=2027) + assert first.digest != later.digest + + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={"options": {"region_group": ["state/ca"]}} + ) + ca = collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=("state/ca",), + scope_key="state/ca", + spm=SPM_SELECTION, + ) + ny = collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=("state/ny",), + scope_key="state/ny", + spm=SPM_SELECTION, + ) + national_spm = collect_baseline_identity( + simulation, + resolved=resolved, + region_codes=("state/ca",), + scope_key="state/ca", + spm={**SPM_SELECTION, "geography_kind": "national"}, + ) + first_entity = simulation.output_plan.entities[0] + changed_plan = collect_baseline_identity( + simulation.model_copy( + update={ + "output_plan": simulation.output_plan.model_copy( + update={ + "entities": ( + first_entity.model_copy( + update={ + "materialized_variables": ( + *first_entity.materialized_variables, + "new_cache_variable", + ) + } + ), + *simulation.output_plan.entities[1:], + ) + } + ) + } + ), + resolved=resolved, + region_codes=("state/ca",), + scope_key="state/ca", + spm=SPM_SELECTION, + ) + assert ca.digest != ny.digest + assert changed_plan.digest != ca.digest + assert national_spm.digest != ca.digest + + +def test_only_current_law_default_dataset_baseline_qualifies() -> None: + baseline = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={"options": {"region_group": ["state/ca"]}} + ) + kwargs = { + "resolved": _resolved(), + "dataset_is_default": True, + "scoping_strategy": _strategy(), + "spm": SPM_SELECTION, + } + assert qualifying_cache_identity(baseline, **kwargs) is not None + assert ( + qualifying_cache_identity( + baseline.model_copy(update={"role": SimulationRole.REFORM}), **kwargs + ) + is None + ) + assert ( + qualifying_cache_identity( + baseline.model_copy(update={"policy": {"reform": {"2026": 1}}}), + **kwargs, + ) + is None + ) + assert ( + qualifying_cache_identity(baseline, **{**kwargs, "dataset_is_default": False}) + is None + ) + + +def _cached_model() -> Stage12CachedSimulation: + model = Stage12CachedSimulation.model_construct( + id="stage12-bl1-test", + output_dataset=None, + dataset=SimpleNamespace(year=2026), + tax_benefit_model_version=None, + policy=None, + scoping_strategy=None, + extra_variables={}, + dynamic=None, + ) + model.configure_stage12_cache( + _planned_simulation(SimulationRole.BASELINE).output_plan + ) + return model + + +def _output(complete: bool = True): + household = { + "household_id": pd.Series([1], dtype="int64"), + "household_net_income": pd.Series([1.0], dtype="float32"), + } + if not complete: + household.pop("household_net_income") + return SimpleNamespace( + data=SimpleNamespace( + entity_data={ + "household": pd.DataFrame(household), + "person": pd.DataFrame( + { + "age": [30], + "household_id": [1], + "person_id": [1], + } + ), + } + ) + ) + + +def test_cached_simulation_reports_hit_and_miss(monkeypatch) -> None: + hit = _cached_model() + monkeypatch.setattr( + Simulation, + "ensure", + lambda self: setattr(self, "output_dataset", _output()), + ) + hit.ensure() + assert hit.stage12_cache_outcome == CACHE_HIT + + miss = _cached_model() + + def ensure_by_running(self): + self.run() + + monkeypatch.setattr(Simulation, "ensure", ensure_by_running) + monkeypatch.setattr( + Simulation, + "run", + lambda self: setattr(self, "output_dataset", _output()), + ) + miss.ensure() + assert miss.stage12_cache_outcome == CACHE_MISS + + +def test_incomplete_cache_recalculates(monkeypatch) -> None: + model = _cached_model() + monkeypatch.setattr( + Simulation, + "ensure", + lambda self: setattr(self, "output_dataset", _output(complete=False)), + ) + monkeypatch.setattr( + Simulation, + "run", + lambda self: setattr(self, "output_dataset", _output()), + ) + monkeypatch.setattr(Simulation, "save", lambda self: None) + + model.ensure() + + assert model.stage12_cache_outcome == CACHE_INCOMPLETE + assert "household_net_income" in model.output_dataset.data.entity_data["household"] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py b/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py index 20f041cca..78fc7945d 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_failure_recording.py @@ -5,10 +5,17 @@ import logging from policyengine_simulation_executor.stage12_runtime import failure_recording +from policyengine_simulation_executor.stage12_runtime.failure_causes import ( + StageProgress, +) from policyengine_simulation_executor.stage12_runtime.failure_recording import ( persist_failure_record, report_original_failure, ) +from policyengine_simulation_executor.stage12_runtime.failures import ( + ERROR_SUMMARY_MAX_LENGTH, + redacted_failure_detail, +) def test_original_failure_is_logged_when_no_observability_runtime_is_available( @@ -66,3 +73,23 @@ def keep_worker_cause() -> str: ) assert kept_cause == worker_cause + + +def test_a_long_unexpected_failure_keeps_its_correlation_id() -> None: + progress = StageProgress(role="baseline") + with progress.stage("stage12_calculation"): + pass + + detail = redacted_failure_detail( + RuntimeError("household data rejected: " + "row " * 400), + default_code="simulation_execution_failed", + progress=progress, + ) + + assert detail.correlation_id is not None + suffix = f" (correlation_id={detail.correlation_id})" + assert len(detail.error_summary) == ERROR_SUMMARY_MAX_LENGTH + assert detail.error_summary.endswith(f"...{suffix}") + assert detail.error_summary.startswith( + "RuntimeError after stage12_calculation (baseline): household data rejected" + ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py index 811d90746..c26558790 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_modal_app.py @@ -58,6 +58,21 @@ def _invoke( ): from policyengine_simulation_executor import stage12_runtime + if function_name == "run_single_simulation_segment_us": + + class Result: + def model_dump(self, *, mode): + assert mode == "python" + return implementation() + + monkeypatch.setattr( + stage12_runtime, "calculate_segment", lambda *_, **__: Result() + ) + return module.run_single_simulation_segment_us( + {"simulation": {"role": "baseline"}, "segment_index": 0}, + observability_context=_observability_context(), + ) + if function_name == "coordinate_report": monkeypatch.setattr( module.modal, @@ -181,6 +196,14 @@ def test_v2_static_runtime_file_smoke_uses_exact_uk_worker_image(monkeypatch) -> } +def test_only_us_worker_image_fetches_stage12_cache(monkeypatch) -> None: + module = _load(monkeypatch) + + assert any(call[0] == "run_function" for call in module.us_worker_image.calls) + assert not any(call[0] == "run_function" for call in module.uk_worker_image.calls) + assert not any(call[0] == "run_function" for call in module.coordinator_image.calls) + + def test_v2_app_declares_validation_workers_and_non_http_coordinator( monkeypatch, ) -> None: @@ -190,18 +213,24 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( assert set(functions) == { "validate_worker_us", "validate_worker_uk", + "run_single_simulation_segment_us", "run_single_simulation_us", "run_single_simulation_uk", "coordinate_report", } assert functions["run_single_simulation_us"]["image"] is module.us_worker_image assert functions["run_single_simulation_uk"]["image"] is module.uk_worker_image - assert functions["run_single_simulation_us"]["timeout"] == 3000 + assert ( + functions["run_single_simulation_segment_us"]["image"] is module.us_worker_image + ) + assert functions["run_single_simulation_segment_us"]["timeout"] == 3600 + assert functions["run_single_simulation_segment_us"]["max_containers"] == 300 + assert functions["run_single_simulation_us"]["timeout"] == 3900 assert functions["run_single_simulation_uk"]["timeout"] == 3000 assert functions["run_single_simulation_us"]["max_containers"] == 10 assert functions["run_single_simulation_uk"]["max_containers"] == 10 assert functions["coordinate_report"]["image"] is module.coordinator_image - assert functions["coordinate_report"]["timeout"] == 4500 + assert functions["coordinate_report"]["timeout"] == 5400 assert functions["coordinate_report"]["max_containers"] == 10 assert "asgi_app" not in vars(module) @@ -211,6 +240,7 @@ def test_v2_app_declares_validation_workers_and_non_http_coordinator( [ "run_single_simulation_us", "run_single_simulation_uk", + "run_single_simulation_segment_us", "coordinate_report", ], ) @@ -237,6 +267,7 @@ def test_stage12_functions_propagate_identifier_and_close_runtime( [ "run_single_simulation_us", "run_single_simulation_uk", + "run_single_simulation_segment_us", "coordinate_report", ], ) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py new file mode 100644 index 000000000..e245af84b --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_precompute_app.py @@ -0,0 +1,70 @@ +"""Declaration tests for the independent Stage 12 precompute app.""" + +from __future__ import annotations + +import importlib +import sys + +from fixtures.fake_modal import install_fake_modal + + +def _load(monkeypatch): + install_fake_modal(monkeypatch) + sys.modules.pop("src.modal.stage12_precompute_app", None) + return importlib.import_module("src.modal.stage12_precompute_app") + + +def test_precompute_app_declares_independent_worker_resources(monkeypatch) -> None: + module = _load(monkeypatch) + functions = dict(module.app.function_calls) + + assert set(functions) == { + "plan_artifacts", + "build_dataset", + "compute_baseline", + "verify_determinism", + "publish_manifest", + } + assert functions["build_dataset"]["cpu"] == 8.0 + assert functions["build_dataset"]["memory"] == 65536 + assert functions["build_dataset"]["timeout"] == 7200 + assert functions["compute_baseline"]["cpu"] == 8.0 + assert functions["compute_baseline"]["memory"] == 32768 + assert functions["compute_baseline"]["timeout"] == 3600 + assert module.app.name == "policyengine-stage12-cache-precompute" + + +def test_precompute_image_installs_only_the_us_bundle(monkeypatch) -> None: + module = _load(monkeypatch) + command = next( + call for call in module.precompute_image.calls if call[0] == "run_commands" + )[1][0] + + assert "--country us" in command + assert "--country uk" not in command + local_source = next( + call + for call in module.precompute_image.calls + if call[0] == "add_local_python_source" + ) + assert "policyengine_simulation_executor" in local_source[1] + assert "policyengine_stage12_persistence" in local_source[1] + + +def test_precompute_app_mounts_only_the_read_only_hugging_face_secret( + monkeypatch, +) -> None: + module = _load(monkeypatch) + read_only_secret = { + "args": ("pe-uk-private-hf-read-token",), + "kwargs": {"required_keys": ["HUGGING_FACE_TOKEN"]}, + } + gcp_secret = {"args": ("stage12-evaluation-gcp-credentials",), "kwargs": {}} + + assert module.hf_secret == read_only_secret + bundle_install = next( + call for call in module.precompute_image.calls if call[0] == "run_commands" + ) + assert bundle_install[2]["secrets"] == [read_only_secret] + for function_name, declaration in module.app.function_calls: + assert declaration["secrets"] == [gcp_secret, read_only_secret], function_name diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py index eae07ae75..820a2e922 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_qualification.py @@ -13,10 +13,17 @@ ) from policyengine_simulation_executor.stage12_parity import ParityMismatch from policyengine_simulation_executor.stage12_qualification import ( + _run_stage12_simulation_for_qualification, qualify_report_parity, ) -from policyengine_simulation_executor.stage12_runtime import SimulationCalculation -from test_stage12_runtime import _report +from policyengine_simulation_executor.stage12_runtime import ( + SimulationCalculation, + calculate_segment, +) +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, +) +from test_stage12_runtime import _planned_simulation, _report def _frames(value: float) -> dict[str, pd.DataFrame]: @@ -60,6 +67,44 @@ def _existing_request() -> dict[str, Any]: } +def test_qualification_runner_exercises_all_segments_and_exact_merge() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + calls = [] + + def run_segment(segment): + calls.append(segment.segment_index) + index = segment.segment_index + frames = { + "household": pd.DataFrame( + { + "household_id": pd.Series([index], dtype="int64"), + "household_net_income": pd.Series([index + 0.5], dtype="float32"), + } + ), + "person": pd.DataFrame( + { + "age": pd.Series([30 + index], dtype="int16"), + "household_id": pd.Series([index], dtype="int64"), + "person_id": pd.Series([100 + index], dtype="int64"), + } + ), + } + return calculate_segment( + segment, + calculator=lambda _: SimulationCalculation(frames=frames), + ) + + result = _run_stage12_simulation_for_qualification( + simulation, + groups=US_REGION_GROUPS, + segment_runner=run_segment, + ) + + assert calls == list(range(20)) + assert result.frames["household"]["household_id"].tolist() == list(range(20)) + assert str(result.frames["person"]["age"].dtype) == "int16" + + def test_qualification_runs_matching_inputs_through_both_complete_paths() -> None: report = _report() baseline = _frames(100.0) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py index 9b5d9b248..679883b0a 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_runtime.py @@ -2,18 +2,24 @@ from __future__ import annotations +import importlib import io import json +import logging +import re +import sys import time from concurrent.futures import ThreadPoolExecutor from contextlib import contextmanager, nullcontext from datetime import UTC, datetime, timedelta from hashlib import sha256 from threading import Lock +from types import SimpleNamespace from uuid import UUID, uuid4 import pandas as pd import pytest +from fixtures.fake_modal import install_fake_modal from fixtures.fake_object_store import InMemoryObjectStore from fixtures.macro_output import CURRENT_SINGLE_YEAR_MACRO_RESULT from policyengine_observability import ( @@ -28,6 +34,7 @@ from policyengine_simulation_contract.macro_output import SingleYearMacroOutput from policyengine_simulation_contract.spm import ( SPMComparisonProvenance, + SPMInputError, SPMRuntimeVersions, SPMSelection, build_spm_provenance, @@ -73,6 +80,7 @@ ) from policyengine_simulation_executor.stage12_runtime import ( SimulationCalculation, + Stage12SegmentRunner, _build_spm_result, build_aggregate_report, database_retry, @@ -91,9 +99,15 @@ from policyengine_simulation_executor.stage12_runtime.aggregation import ( validate_uk_local_authority_metadata, ) +from policyengine_simulation_executor.stage12_runtime.comparison import ( + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS, +) from policyengine_simulation_executor.stage12_runtime.dependencies import ( runtime_store, ) +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, +) from policyengine_simulation_executor.uk_local_authority_metadata import ( detect_uk_local_authority_metadata, ) @@ -380,6 +394,14 @@ def replace_simulation(self, record): self.children[record.simulation_execution_id] = record return record + def replace_simulation_if_status(self, record, *, expected_status): + with self.lock: + current = self.children[record.simulation_execution_id] + if current.status is not expected_status: + return current, False + self.children[record.simulation_execution_id] = record + return record, True + def attach_simulation_invocation( self, simulation_execution_id, @@ -567,7 +589,7 @@ def build_simulation(params, *, dataset, dataset_selection, **kwargs): resolve_dataset, ) monkeypatch.setattr(simulation_runtime, "_load_dataset", load_dataset) - monkeypatch.setattr(simulation_runtime, "_build_simulation", build_simulation) + monkeypatch.setattr(worker, "build_stage12_simulation", build_simulation) result = worker.calculate_simulation_frames( _planned_simulation(SimulationRole.BASELINE) @@ -685,6 +707,55 @@ def test_single_worker_accepts_one_policy_and_persists_one_artifact() -> None: ] +def test_single_worker_routes_eligible_us_calculation_to_segment_runner() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + received = [] + + def segmented_calculator(value, cancellation_requested): + received.append(value) + assert cancellation_requested() is False + return SimulationCalculation(frames=_frames()) + + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + ) + + assert received == [simulation] + + +def test_single_worker_honors_explicit_segmentation_opt_out() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={"options": {"segmented": False}} + ) + store.children[simulation.simulation_execution_id] = _child(simulation) + segmented_calls = [] + + def unexpected_segmented(value, cancellation_requested): + segmented_calls.append(value) + return SimulationCalculation(frames=_frames()) + + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: _frames(), + segmented_calculator=unexpected_segmented, + ) + + assert segmented_calls == [] + + def test_single_worker_retains_detached_calculation_provenance() -> None: store = FakeStore() simulation = _planned_simulation(SimulationRole.BASELINE) @@ -734,7 +805,7 @@ def fail(_simulation): child = store.children[simulation.simulation_execution_id] assert child.status is ComparisonRunLifecycleStatus.FAILED assert child.error_code == "simulation_execution_failed" - assert child.error_summary == ( + assert _failure_cause(child.error_summary) == ( "RuntimeError in stage12_calculation (baseline): " "country package calculation failed" ) @@ -742,7 +813,7 @@ def fail(_simulation): def test_single_worker_persists_and_records_a_redacted_cause() -> None: store = FakeStore() - simulation = _planned_simulation(SimulationRole.REFORM) + simulation = _planned_simulation(SimulationRole.BASELINE) store.children[simulation.simulation_execution_id] = _child(simulation) runtime = RecordingRuntime() @@ -763,8 +834,8 @@ def fail(_simulation): ) child = store.children[simulation.simulation_execution_id] - assert child.error_summary == ( - "RuntimeError in stage12_calculation (reform): " + assert _failure_cause(child.error_summary) == ( + "RuntimeError in stage12_calculation (baseline): " "could not reach postgresql://[REDACTED]@db.example/postgres" ) assert "database-secret" not in child.error_summary @@ -791,11 +862,64 @@ def write_simulation(self, **_kwargs): calculator=lambda _: _frames(), ) - assert store.children[simulation.simulation_execution_id].error_summary == ( + assert _failure_cause( + store.children[simulation.simulation_execution_id].error_summary + ) == ( "OSError in stage12_simulation_artifact_write (baseline): bucket write rejected" ) +def test_single_worker_preserves_typed_failure_values() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + + def fail(_simulation): + raise SPMInputError("SPM_YEAR_UNAVAILABLE", "No SPM values for 2026") + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + ) + + child = store.children[simulation.simulation_execution_id] + assert child.error_code == "SPM_YEAR_UNAVAILABLE" + assert child.error_summary == "No SPM values for 2026" + + +def test_single_worker_rejects_bare_year_policy_period_before_calculation() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE).model_copy( + update={ + "policy": {"gov.irs.credits.ctc.amount.base[0].amount": {"2026": 3_000}} + } + ) + store.children[simulation.simulation_execution_id] = _child(simulation) + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("calculation must not start"), + ) + + child = store.children[simulation.simulation_execution_id] + assert child.error_code == "invalid_policy_period" + assert child.error_summary == ( + "Policy parameter 'gov.irs.credits.ctc.amount.base[0].amount' has invalid " + "effective period '2026'; expected YYYY-MM-DD or " + "YYYY-MM-DD.YYYY-MM-DD." + ) + + def _connection_limit_error() -> OperationalError: # SQLAlchemy wraps a failure to open a connection without a statement. return OperationalError( @@ -805,6 +929,24 @@ def _connection_limit_error() -> OperationalError: ) +CORRELATION_SUFFIX = re.compile(r" \(correlation_id=(?P[0-9a-f]+)\)\Z") + + +def _correlation_id(error_summary: str) -> str: + """Return the correlation ID that ends a redacted failure summary.""" + + match = CORRELATION_SUFFIX.search(error_summary) + assert match is not None, error_summary + return match["correlation_id"] + + +def _failure_cause(error_summary: str) -> str: + """Return the redacted cause before a failure summary's correlation ID.""" + + _correlation_id(error_summary) + return CORRELATION_SUFFIX.sub("", error_summary) + + class TelemetryRecordingRuntime(RecordingRuntime): """Record telemetry in the same ordered list as database writes.""" @@ -848,11 +990,13 @@ def test_single_worker_reports_failure_before_an_unwritable_failure_record() -> ordered_events: list[tuple[str, object]] = [] class UnwritableFailureStore(FakeStore): - def replace_simulation(self, record): + def replace_simulation_if_status(self, record, *, expected_status): if record.status is ComparisonRunLifecycleStatus.FAILED: ordered_events.append(("failure_write", record.error_summary)) raise _connection_limit_error() - return super().replace_simulation(record) + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) store = UnwritableFailureStore() simulation = _planned_simulation(SimulationRole.BASELINE) @@ -880,11 +1024,14 @@ def fail(_simulation): assert attributes == { "stage12_failure_scope": "simulation_execution", "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "correlation_id": _correlation_id(failure_write[1]), } - assert failure_write == ( - "failure_write", + assert failure_write[0] == "failure_write" + assert _failure_cause(failure_write[1]) == ( "ValueError in stage12_calculation (baseline): " - "country column missing from household data", + "country column missing from household data" ) assert persistence_report[0] == "telemetry" assert persistence_report[1][0] == ( @@ -895,6 +1042,44 @@ def fail(_simulation): assert child.status is ComparisonRunLifecycleStatus.RUNNING +def test_late_simulation_success_does_not_overwrite_parent_cancellation() -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + + def segmented_calculator(value, cancellation_requested): + current = store.children[value.simulation_execution_id] + cancelled_at = datetime.now(UTC) + store.replace_simulation( + current.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "simulation_cancelled", + "error_summary": "Cancelled after the reform simulation failed.", + "updated_at": cancelled_at, + "completed_at": cancelled_at, + } + ) + ) + assert cancellation_requested() is True + return SimulationCalculation(frames=_frames()) + + with pytest.raises(RuntimeError, match="Stage 12 simulation was cancelled"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + ) + + child = store.children[simulation.simulation_execution_id] + assert child.status is ComparisonRunLifecycleStatus.FAILED + assert child.error_code == "simulation_cancelled" + + def test_single_worker_rejects_frames_that_do_not_satisfy_the_output_plan() -> None: store = FakeStore() simulation = _planned_simulation(SimulationRole.BASELINE) @@ -1190,17 +1375,24 @@ def __init__(self, object_id, future): def get(self, *, timeout=None): return self.future.result(timeout=timeout) + def cancel(self): + self.future.cancel() + class ImmediateCall: def __init__(self, object_id, result, observed_timeouts): self.object_id = object_id self.result = result self.observed_timeouts = observed_timeouts + self.cancel_count = 0 def get(self, *, timeout=None): self.observed_timeouts.append(timeout) return self.result + def cancel(self): + self.cancel_count += 1 + class ConcurrentInvoker: def __init__( @@ -1669,31 +1861,84 @@ def test_coordinator_never_writes_partial_aggregate_when_a_child_fails() -> None assert "child calculation failed" not in str(error.value) - expected_summary = ( - "RuntimeError in stage12_child_wait (reform): child calculation failed" - ) assert artifacts.aggregate_writes == [] assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == expected_summary + assert _failure_cause(store.parent.error_summary) == ( + "RuntimeError in stage12_child_wait (reform): child calculation failed" + ) failed_child = store.children[REFORM_ID] assert failed_child.status is ComparisonRunLifecycleStatus.FAILED assert failed_child.error_code == "simulation_invocation_failed" - assert failed_child.error_summary == expected_summary + assert failed_child.error_summary == store.parent.error_summary @pytest.mark.parametrize( - "invoker", + ("failed_role", "failed_id", "pending_role", "pending_id"), [ - lambda artifacts: ConcurrentInvoker(artifacts, wrong_plan_digest=True), - lambda artifacts: ConcurrentInvoker(artifacts, missing_column=True), + ("reform", REFORM_ID, "baseline", BASELINE_ID), + ("baseline", BASELINE_ID, "reform", REFORM_ID), ], ) -def test_coordinator_rejects_artifacts_that_do_not_satisfy_the_shared_plan( - invoker, +def test_coordinator_observes_one_role_failure_while_the_peer_is_pending( + failed_role, + failed_id, + pending_role, + pending_id, ) -> None: store = FakeStore() artifacts = FakeArtifacts() + class PendingCall: + object_id = f"call-{pending_role}" + + def __init__(self): + self.timeouts = [] + self.cancel_count = 0 + + def get(self, *, timeout=None): + self.timeouts.append(timeout) + raise TimeoutError + + def cancel(self): + self.cancel_count += 1 + + class FailedCall: + object_id = f"call-{failed_role}" + + def __init__(self): + self.timeouts = [] + + def get(self, *, timeout=None): + self.timeouts.append(timeout) + child = store.children[failed_id] + failed_at = datetime.now(UTC) + store.replace_simulation( + child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + "updated_at": failed_at, + "completed_at": failed_at, + } + ) + ) + raise RuntimeError("remote wrapper") + + def cancel(self): + return None + + baseline_call = PendingCall() + failed_call = FailedCall() + + class ConcurrentFailureInvoker: + def spawn(self, *, simulation, **_): + role = PlannedSimulationExecutionInput.model_validate(simulation).role + return failed_call if role.value == failed_role else baseline_call + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): coordinate_report( _report().model_dump(mode="json"), @@ -1703,20 +1948,40 @@ def test_coordinator_rejects_artifacts_that_do_not_satisfy_the_shared_plan( coordinator_invocation_id="coordinator-1", store=store, artifacts=artifacts, - invoker=invoker(artifacts), + invoker=ConcurrentFailureInvoker(), aggregator=lambda **_: _aggregate_payload({"must": "not run"}), + sleep=lambda _: pytest.fail("failure should be observed before sleeping"), ) + assert failed_call.timeouts == [0] + assert baseline_call.timeouts == ([0] if pending_role == "baseline" else []) + # A segmented peer reads its cancelled record and stops its own + # region-group calls, so it is not cancelled through Modal. + assert baseline_call.cancel_count == 0 + pending = store.children[pending_id] + assert pending.status is ComparisonRunLifecycleStatus.FAILED + assert pending.error_code == "simulation_cancelled" + assert pending.error_summary == ( + f"Cancelled after the {failed_role} simulation failed." + ) + assert store.parent.error_code == "invalid_policy_period" + assert store.parent.error_summary == "The reform period is invalid." assert artifacts.aggregate_writes == [] - assert store.parent.status is ComparisonRunLifecycleStatus.FAILED -def test_coordinator_records_failure_when_child_persistence_fails() -> None: - class FailingChildStore(FakeStore): - def create_or_resolve_simulation(self, record): - raise RuntimeError("database rejected the child record") +def test_dispatch_failure_cancels_an_already_started_peer() -> None: + store = FakeStore() + baseline_call = TimeoutCall("call-baseline") - store = FailingChildStore() + class DispatchFailureInvoker: + def spawn(self, *, simulation, **_): + role = PlannedSimulationExecutionInput.model_validate(simulation).role + if role is SimulationRole.REFORM: + raise ConnectionError("dispatch contains sensitive details") + return baseline_call + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): coordinate_report( @@ -1727,37 +1992,87 @@ def create_or_resolve_simulation(self, record): coordinator_invocation_id="coordinator-1", store=store, artifacts=FakeArtifacts(), - invoker=ConcurrentInvoker(FakeArtifacts()), + invoker=DispatchFailureInvoker(), ) - assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_code == "report_coordination_failed" - assert store.parent.error_summary == ( - "RuntimeError in stage12_child_state_create: database rejected the child record" - ) + baseline = store.children[BASELINE_ID] + reform = store.children[REFORM_ID] + assert baseline_call.cancelled == 0 + assert baseline.status is ComparisonRunLifecycleStatus.FAILED + assert baseline.error_code == "simulation_cancelled" + assert baseline.error_summary == "Cancelled after the reform simulation failed." + assert reform.status is ComparisonRunLifecycleStatus.FAILED + assert reform.error_code == "simulation_dispatch_failed" + assert store.parent.error_code == "simulation_dispatch_failed" -class RefusingDispatchInvoker: - def __init__(self, dispatch_error): - self.dispatch_error = dispatch_error - self.events = [] +def test_coordinator_propagates_an_already_persisted_child_failure() -> None: + store = FakeStore() + artifacts = FakeArtifacts() - def spawn(self, *, simulation, **_): - role = PlannedSimulationExecutionInput.model_validate(simulation).role.value - self.events.append(f"spawn:{role}") - raise self.dispatch_error + class PersistedFailureCall: + object_id = "call-reform" - def restore(self, invocation_id): - raise AssertionError(f"unexpected restored invocation {invocation_id}") + def get(self, *, timeout=None): + child = store.children[REFORM_ID] + failed_at = datetime.now(UTC) + store.replace_simulation( + child.model_copy( + update={ + "status": ComparisonRunLifecycleStatus.FAILED, + "error_code": "invalid_policy_period", + "error_summary": "The reform period is invalid.", + "updated_at": failed_at, + "completed_at": failed_at, + } + ) + ) + raise RuntimeError("remote wrapper must not replace the child failure") + + class PersistedFailureInvoker: + def spawn(self, *, simulation, **_): + parsed = PlannedSimulationExecutionInput.model_validate(simulation) + if parsed.role is SimulationRole.REFORM: + return PersistedFailureCall() + descriptor = artifacts.add_simulation(parsed, _frames(100.0)) + return ImmediateCall( + "call-baseline", descriptor.model_dump(mode="json"), [] + ) + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") -def test_coordinator_records_a_redacted_dispatch_failure_cause() -> None: + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=PersistedFailureInvoker(), + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), + ) + + assert store.children[REFORM_ID].error_code == "invalid_policy_period" + assert store.children[REFORM_ID].error_summary == "The reform period is invalid." + assert store.parent.error_code == "invalid_policy_period" + assert store.parent.error_summary == "The reform period is invalid." + + +@pytest.mark.parametrize( + "invoker", + [ + lambda artifacts: ConcurrentInvoker(artifacts, wrong_plan_digest=True), + lambda artifacts: ConcurrentInvoker(artifacts, missing_column=True), + ], +) +def test_coordinator_rejects_artifacts_that_do_not_satisfy_the_shared_plan( + invoker, +) -> None: store = FakeStore() artifacts = FakeArtifacts() - invoker = RefusingDispatchInvoker( - ConnectionError("Modal rejected the call with token=opaque-token") - ) - runtime = RecordingRuntime() with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): coordinate_report( @@ -1768,25 +2083,38 @@ def test_coordinator_records_a_redacted_dispatch_failure_cause() -> None: coordinator_invocation_id="coordinator-1", store=store, artifacts=artifacts, - invoker=invoker, + invoker=invoker(artifacts), aggregator=lambda **_: _aggregate_payload({"must": "not run"}), - runtime=runtime, ) - expected_summary = ( - "ConnectionError in stage12_child_dispatch (baseline): " - "Modal rejected the call with token=[REDACTED]" - ) - assert invoker.events == ["spawn:baseline"] assert artifacts.aggregate_writes == [] - failed_child = store.children[BASELINE_ID] - assert failed_child.status is ComparisonRunLifecycleStatus.FAILED - assert failed_child.error_code == "simulation_dispatch_failed" - assert failed_child.error_summary == expected_summary + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + + +def test_coordinator_records_failure_when_child_persistence_fails() -> None: + class FailingChildStore(FakeStore): + def create_or_resolve_simulation(self, record): + raise RuntimeError("database rejected the child record") + + store = FailingChildStore() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=ConcurrentInvoker(FakeArtifacts()), + ) + + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED assert store.parent.error_code == "report_coordination_failed" - assert store.parent.error_summary == expected_summary - [recorded_cause] = runtime.recorded_exceptions - assert str(recorded_cause) == expected_summary + assert _failure_cause(store.parent.error_summary) == ( + "RuntimeError in stage12_child_state_create: database rejected the child record" + ) class WorkerRecordedFailureCall: @@ -1816,7 +2144,10 @@ def __init__(self, object_id): self.object_id = object_id def get(self, *, timeout=None): - raise AssertionError("the coordinator waits for baseline first") + raise AssertionError("the coordinator observes baseline first") + + def cancel(self): + return None class WorkerRecordedFailureInvoker: @@ -1843,7 +2174,8 @@ def test_coordinator_keeps_the_cause_recorded_by_a_failed_worker() -> None: store = FakeStore() worker_summary = ( "ValueError in stage12_calculation (baseline): " - "Variable 'country' not found in household data" + "Variable 'country' not found in household data " + "(correlation_id=0123456789abcdef0123456789abcdef)" ) runtime = RecordingRuntime() @@ -1864,7 +2196,7 @@ def test_coordinator_keeps_the_cause_recorded_by_a_failed_worker() -> None: failed_child = store.children[BASELINE_ID] assert failed_child.error_code == "simulation_execution_failed" assert failed_child.error_summary == worker_summary - assert store.parent.error_code == "report_coordination_failed" + assert store.parent.error_code == "simulation_execution_failed" assert store.parent.error_summary == worker_summary [recorded_cause] = runtime.recorded_exceptions assert str(recorded_cause) == worker_summary @@ -1891,7 +2223,7 @@ def write_aggregate(self, *, prefix, payload): aggregator=lambda **_: {"result": {"spm_provenance": None}}, ) - assert store.parent.error_summary == ( + assert _failure_cause(store.parent.error_summary) == ( "ValidationError in stage12_aggregate_artifact_write: " "1 validation error for SPMComparisonProvenance: " "Input should be a valid dictionary or instance " @@ -2127,11 +2459,15 @@ def fail_aggregation(**_): assert error.value.__cause__ is None original_report, failure_write, persistence_report = ordered_events - assert original_report[1][1] == {"stage12_failure_scope": "report_coordination"} + assert original_report[1][1] == { + "stage12_failure_scope": "report_coordination", + "evaluation_id": str(EVALUATION_ID), + "correlation_id": _correlation_id(failure_write[1]), + } assert original_report[1][2] is aggregation_error - assert failure_write == ( - "failure_write", - "ValueError in stage12_aggregation: aggregate failed validation", + assert failure_write[0] == "failure_write" + assert _failure_cause(failure_write[1]) == ( + "ValueError in stage12_aggregation: aggregate failed validation" ) assert isinstance(persistence_report[1][2], OperationalError) assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING @@ -2170,26 +2506,47 @@ def get_simulation(self, simulation_execution_id): if event[0] == "telemetry" and event[1][0] == "Stage 12 execution failed" ] assert [report[1] for report in failure_reports] == [ - {"stage12_failure_scope": "child_wait", "simulation_role": "baseline"} + { + "stage12_failure_scope": "child_wait", + "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "correlation_id": _correlation_id(store.parent.error_summary), + } ] assert ordered_events[0][0] == "telemetry" assert ordered_events[1] == ("child_read", str(BASELINE_ID)) # The report record is still written after the child record could not be. assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == ( + assert store.parent.error_code == "simulation_invocation_failed" + assert _failure_cause(store.parent.error_summary) == ( "RuntimeError in stage12_child_wait (baseline): child calculation failed" ) -def test_coordinator_reports_a_child_dispatch_failure_once() -> None: - ordered_events: list[tuple[str, object]] = [] - dispatch_error = ConnectionError("Modal refused the dispatch") - store = FakeStore() - invoker = RefusingDispatchInvoker(dispatch_error) +class RefusingDispatchInvoker: + def __init__(self, dispatch_error): + self.dispatch_error = dispatch_error + self.events = [] - with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): - coordinate_report( - _report().model_dump(mode="json"), + def spawn(self, *, simulation, **_): + role = PlannedSimulationExecutionInput.model_validate(simulation).role.value + self.events.append(f"spawn:{role}") + raise self.dispatch_error + + def restore(self, invocation_id): + raise AssertionError(f"unexpected restored invocation {invocation_id}") + + +def test_coordinator_reports_a_child_dispatch_failure_once() -> None: + ordered_events: list[tuple[str, object]] = [] + dispatch_error = ConnectionError("Modal refused the dispatch") + store = FakeStore() + invoker = RefusingDispatchInvoker(dispatch_error) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + _report().model_dump(mode="json"), _context().model_dump(mode="json"), _parent().model_dump(mode="json"), application_name=_context().modal_application, @@ -2206,23 +2563,65 @@ def test_coordinator_reports_a_child_dispatch_failure_once() -> None: for event in ordered_events if event[0] == "telemetry" and event[1][0] == "Stage 12 execution failed" ] + baseline_child = store.children[BASELINE_ID] assert [(attributes, error) for _message, attributes, error in failure_reports] == [ ( - {"stage12_failure_scope": "child_dispatch", "simulation_role": "baseline"}, + { + "stage12_failure_scope": "child_dispatch", + "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "correlation_id": _correlation_id(baseline_child.error_summary), + }, dispatch_error, ) ] assert invoker.events == ["spawn:baseline"] - baseline_child = store.children[BASELINE_ID] assert baseline_child.status is ComparisonRunLifecycleStatus.FAILED assert baseline_child.error_code == "simulation_dispatch_failed" - expected_summary = ( + assert _failure_cause(baseline_child.error_summary) == ( "ConnectionError in stage12_child_dispatch (baseline): " "Modal refused the dispatch" ) - assert baseline_child.error_summary == expected_summary assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == expected_summary + assert store.parent.error_summary == baseline_child.error_summary + + +def test_coordinator_records_a_redacted_dispatch_failure_cause() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = RefusingDispatchInvoker( + ConnectionError("Modal rejected the call with token=opaque-token") + ) + runtime = RecordingRuntime() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=invoker, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), + runtime=runtime, + ) + + assert invoker.events == ["spawn:baseline"] + assert artifacts.aggregate_writes == [] + failed_child = store.children[BASELINE_ID] + assert failed_child.status is ComparisonRunLifecycleStatus.FAILED + assert failed_child.error_code == "simulation_dispatch_failed" + assert _failure_cause(failed_child.error_summary) == ( + "ConnectionError in stage12_child_dispatch (baseline): " + "Modal rejected the call with token=[REDACTED]" + ) + assert store.parent.error_code == "simulation_dispatch_failed" + assert store.parent.error_summary == failed_child.error_summary + [recorded_cause] = runtime.recorded_exceptions + assert str(recorded_cause) == failed_child.error_summary def test_coordinator_rejects_incompatible_row_identity() -> None: @@ -2313,9 +2712,13 @@ def test_coordinator_resumes_a_matching_running_child_invocation() -> None: class TimeoutCall: def __init__(self, object_id): self.object_id = object_id + self.cancelled = 0 def get(self, *, timeout=None): - raise TimeoutError("child result was not ready") + raise TimeoutError("contains sensitive timeout details") + + def cancel(self): + self.cancelled += 1 class TimeoutInvoker: @@ -2349,17 +2752,1113 @@ def test_coordinator_records_bounded_timeout_without_an_aggregate() -> None: artifacts=artifacts, invoker=invoker, aggregator=lambda **_: _aggregate_payload({"must": "not run"}), + wait_timeout_seconds=0, ) - assert "child result was not ready" not in str(error.value) + assert "sensitive" not in str(error.value) assert invoker.events == ["spawn:baseline", "spawn:reform"] assert artifacts.aggregate_writes == [] - expected_summary = ( - "TimeoutError in stage12_child_wait (baseline): child result was not ready" + assert store.parent.status is ComparisonRunLifecycleStatus.FAILED + assert store.parent.error_code == "simulation_wait_timeout" + assert store.parent.error_summary == ( + "Stage 12 simulations did not finish within 0 seconds of the report " + "coordinator starting." + ) + assert all( + child.status is ComparisonRunLifecycleStatus.FAILED + for child in store.children.values() + ) + assert all( + child.error_code == "simulation_cancelled" for child in store.children.values() + ) + + +def test_closed_wait_window_is_reported_before_its_failure_records() -> None: + ordered_events: list[tuple[str, object]] = [] + + class UnwritableFailureStore(FakeStore): + def replace_simulation_if_status(self, record, *, expected_status): + if record.status is ComparisonRunLifecycleStatus.FAILED: + ordered_events.append(("failure_write", record.error_code)) + raise _connection_limit_error() + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + def replace_report(self, record): + if record.status is ComparisonRunLifecycleStatus.FAILED: + ordered_events.append(("failure_write", record.error_code)) + raise _connection_limit_error() + return super().replace_report(record) + + store = UnwritableFailureStore() + invoker = TimeoutInvoker() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=FakeArtifacts(), + invoker=invoker, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), + runtime=TelemetryRecordingRuntime(ordered_events), + wait_timeout_seconds=0, + ) + + original_report, *failure_writes, persistence_report = ordered_events + event_kind, (message, attributes, reported_error) = original_report + assert event_kind == "telemetry" + assert message == "Stage 12 execution failed" + assert attributes == {"stage12_failure_scope": "child_wait_timeout"} + assert reported_error.detail.error_code == "simulation_wait_timeout" + # Both cancellation records and the report record were attempted after + # the closed window had already been reported. + assert failure_writes == [ + ("failure_write", "simulation_cancelled"), + ("failure_write", "simulation_cancelled"), + ("failure_write", "simulation_wait_timeout"), + ] + assert persistence_report[1][0] == ( + "Stage 12 failure record could not be persisted" + ) + assert store.parent.status is ComparisonRunLifecycleStatus.RUNNING + + +class StartupWriteError(Exception): + """A coordinator persistence failure after a child may have started.""" + + +class ScriptedChildCall: + """A started child that keeps running unless the test scripts an outcome.""" + + def __init__( + self, + object_id, + *, + clock=None, + result=None, + wait_error=None, + cancel_error=None, + during_wait=None, + ): + self.object_id = object_id + self.clock = clock + self.result = result + self.wait_error = wait_error + self.cancel_error = cancel_error + self.during_wait = during_wait + self.probed_at = [] + self.cancel_count = 0 + + def get(self, *, timeout=None): + self.probed_at.append( + self.clock.monotonic() if self.clock is not None else None + ) + if self.during_wait is not None: + self.during_wait() + if self.wait_error is not None: + raise self.wait_error + if self.result is not None: + return self.result + raise TimeoutError("still running") + + def cancel(self): + self.cancel_count += 1 + if self.cancel_error is not None: + raise self.cancel_error + + +class ScriptedChildInvoker: + def __init__( + self, + *, + clock=None, + results=None, + wait_errors=None, + cancel_errors=None, + during_wait=None, + ): + self.clock = clock + self.results = results or {} + self.wait_errors = wait_errors or {} + self.cancel_errors = cancel_errors or {} + self.during_wait = during_wait or {} + self.events = [] + self.calls = {} + + def spawn(self, *, simulation, **_): + role = PlannedSimulationExecutionInput.model_validate(simulation).role.value + self.events.append(f"spawn:{role}") + return self._call(role) + + def restore(self, invocation_id): + role = invocation_id.removeprefix("call-") + self.events.append(f"restore:{role}") + return self._call(role) + + def _call(self, role): + call = ScriptedChildCall( + f"call-{role}", + clock=self.clock, + result=self.results.get(role), + wait_error=self.wait_errors.get(role), + cancel_error=self.cancel_errors.get(role), + during_wait=self.during_wait.get(role), + ) + self.calls[role] = call + return call + + +class PlaceholderWriteFailingStore(FakeStore): + def __init__(self, failing_simulation_id): + super().__init__() + self.failing_simulation_id = failing_simulation_id + + def replace_simulation(self, record): + # The coordinator's only plain replacement is the dispatch placeholder. + if record.simulation_execution_id == self.failing_simulation_id: + raise StartupWriteError("placeholder write contains sensitive details") + return super().replace_simulation(record) + + +class AttachFailingStore(FakeStore): + def __init__(self, failing_simulation_id): + super().__init__() + self.failing_simulation_id = failing_simulation_id + + def attach_simulation_invocation(self, simulation_execution_id, **values): + if simulation_execution_id == self.failing_simulation_id: + raise StartupWriteError("attach contains sensitive details") + return super().attach_simulation_invocation(simulation_execution_id, **values) + + +class CancellationStateFailingStore(FakeStore): + def replace_simulation_if_status(self, record, *, expected_status): + if record.error_code == "simulation_cancelled": + raise ConnectionError("database contains sensitive details") + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + +def _simulation_id(role: str) -> UUID: + return BASELINE_ID if role == "baseline" else REFORM_ID + + +def _report_with_segmentation(*, segmented: bool) -> ReportExecutionInput: + if segmented: + return _report() + + def unsegmented(role: SimulationRole) -> SimulationExecutionInput: + return _simulation(role).model_copy(update={"options": {"segmented": False}}) + + return ReportExecutionInput( + evaluation_id=EVALUATION_ID, + baseline=unsegmented(SimulationRole.BASELINE), + reform=unsegmented(SimulationRole.REFORM), + requested_aggregates=tuple(ReportAggregate), ) + + +def _coordinate_with(report, store, invoker, **coordinator_options): + coordinator_options.setdefault("artifacts", FakeArtifacts()) + return coordinate_report( + report.model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + invoker=invoker, + aggregator=lambda **_: _aggregate_payload({"must": "not run"}), + **coordinator_options, + ) + + +def _assert_cancelled(child, *, after: str) -> None: + assert child.status is ComparisonRunLifecycleStatus.FAILED + assert child.error_code == "simulation_cancelled" + assert child.error_summary == f"Cancelled after {after}." + assert child.completed_at is not None + + +def _assert_original_coordinator_failure(error, store) -> None: + assert isinstance(error.value.__context__, StartupWriteError) assert store.parent.status is ComparisonRunLifecycleStatus.FAILED - assert store.parent.error_summary == expected_summary - timed_out_child = store.children[BASELINE_ID] - assert timed_out_child.status is ComparisonRunLifecycleStatus.FAILED - assert timed_out_child.error_summary == expected_summary + assert store.parent.error_code == "report_coordination_failed" + original_error = error.value.__context__ + assert _failure_cause(store.parent.error_summary).startswith( + f"{type(original_error).__name__} " + ) + assert _failure_cause(store.parent.error_summary).endswith(f": {original_error}") + + +def _coordination_log(caplog, level): + (record,) = [ + record + for record in caplog.records + if record.name == coordination_module.__name__ and record.levelno == level + ] + return record + + +def _assert_log_identifies_child(record, *, role: str, error_type: str) -> None: + simulation_execution_id = str(_simulation_id(role)) + assert record.evaluation_id == str(EVALUATION_ID) + assert record.simulation_execution_id == simulation_execution_id + assert record.simulation_role == role + assert record.error_type == error_type + message = record.getMessage() + assert str(EVALUATION_ID) in message + assert simulation_execution_id in message + assert "sensitive" not in message + assert record.exc_info is None + + +def test_baseline_failure_cancels_and_records_the_reform_child() -> None: + store = FakeStore() + artifacts = FakeArtifacts() + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline worker failed")} + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + artifacts=artifacts, + ) + + reform_call = invoker.calls["reform"] + assert reform_call.cancel_count == 1 + assert reform_call.probed_at == [] + assert artifacts.aggregate_writes == [] + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert store.parent.error_code == "simulation_invocation_failed" + + +def test_reform_failure_leaves_the_finished_baseline_alone() -> None: + report = _report_with_segmentation(segmented=False) + store = FakeStore() + artifacts = FakeArtifacts() + baseline_descriptor = artifacts.add_simulation(report.baseline) + invoker = ScriptedChildInvoker( + results={"baseline": baseline_descriptor.model_dump(mode="json")}, + wait_errors={"reform": LookupError("reform worker failed")}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(report, store, invoker, artifacts=artifacts) + + assert invoker.calls["baseline"].cancel_count == 0 + assert store.children[BASELINE_ID].status is ComparisonRunLifecycleStatus.RUNNING + assert store.children[BASELINE_ID].error_code is None + assert store.children[REFORM_ID].error_code == "simulation_invocation_failed" + + +def test_baseline_failure_cancels_a_reform_call_restored_on_retry() -> None: + report = _report_with_segmentation(segmented=False) + store = FakeStore() + # An earlier coordinator attempt started the reform child before stopping. + store.children[REFORM_ID] = _child(report.reform).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": "call-reform", + "started_at": NOW, + } + ) + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline worker failed")} + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(report, store, invoker) + + assert invoker.events == ["spawn:baseline", "restore:reform"] + assert invoker.calls["reform"].cancel_count == 1 + _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + + +@pytest.mark.parametrize( + ("worker_status", "worker_error_code"), + [ + (ComparisonRunLifecycleStatus.SUCCEEDED, None), + (ComparisonRunLifecycleStatus.FAILED, "simulation_execution_failed"), + ], +) +def test_cancellation_leaves_a_sibling_that_already_finished_unchanged( + worker_status, + worker_error_code, +) -> None: + store = FakeStore() + worker_records = [] + + def reform_worker_records_its_outcome(): + # The reform worker finishes while the coordinator probes baseline. + finished_at = NOW + timedelta(minutes=5) + worker_records.append( + store.replace_simulation( + store.get_simulation(REFORM_ID).model_copy( + update={ + "status": worker_status, + "error_code": worker_error_code, + "updated_at": finished_at, + "completed_at": finished_at, + } + ) + ) + ) + + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline worker failed")}, + during_wait={"baseline": reform_worker_records_its_outcome}, + ) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + # Cancelling a finished Modal call is harmless; its record is not touched. + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[REFORM_ID] == worker_records[0] + assert store.children[BASELINE_ID].error_code == "simulation_invocation_failed" + + +def test_peer_cancellation_failure_does_not_replace_the_child_failure( + caplog, +) -> None: + store = FakeStore() + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline failure is sensitive")}, + cancel_errors={ + "reform": ConnectionError("cancel request contains sensitive details") + }, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + assert invoker.calls["reform"].cancel_count == 1 + _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert store.parent.error_code == "simulation_invocation_failed" + _assert_log_identifies_child( + _coordination_log(caplog, logging.WARNING), + role="reform", + error_type="ConnectionError", + ) + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_failed_cancellation_state_write_still_cancels_through_modal( + caplog, + segmented, +) -> None: + store = CancellationStateFailingStore() + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline failure is sensitive")}, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with(_report_with_segmentation(segmented=segmented), store, invoker) + + # A segmented worker cannot read a cancellation that was never written, so + # it is cancelled through Modal as well. + assert invoker.calls["reform"].cancel_count == 1 + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.RUNNING + assert store.parent.error_code == "simulation_invocation_failed" + _assert_log_identifies_child( + _coordination_log(caplog, logging.ERROR), + role="reform", + error_type="ConnectionError", + ) + + +def test_wait_timeout_cancels_and_records_every_unfinished_child() -> None: + store = FakeStore() + invoker = ScriptedChildInvoker() + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + wait_timeout_seconds=0, + ) + + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + for role in ("baseline", "reform"): + _assert_cancelled( + store.children[_simulation_id(role)], + after="the 0-second wait window closed", + ) + assert store.parent.error_code == "simulation_wait_timeout" + + +def test_timeout_cancellation_failure_log_identifies_the_report_and_child( + caplog, +) -> None: + store = FakeStore() + invoker = ScriptedChildInvoker( + cancel_errors={ + "reform": ConnectionError("cancel request contains sensitive details") + }, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + wait_timeout_seconds=0, + ) + + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + assert store.parent.error_code == "simulation_wait_timeout" + _assert_log_identifies_child( + _coordination_log(caplog, logging.WARNING), + role="reform", + error_type="ConnectionError", + ) + + +@pytest.mark.parametrize("segmented", [True, False]) +def test_failed_placeholder_write_stops_the_started_baseline(segmented) -> None: + store = PlaceholderWriteFailingStore(REFORM_ID) + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(_report_with_segmentation(segmented=segmented), store, invoker) + + _assert_original_coordinator_failure(error, store) + assert invoker.events == ["spawn:baseline"] + _assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") + # A segmented worker stops itself, and its region-group calls, once it + # reads the cancelled record; Modal does not cancel spawned calls + # recursively. + assert invoker.calls["baseline"].cancel_count == (0 if segmented else 1) + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.PENDING + + +@pytest.mark.parametrize( + ("failing_role", "started_roles"), + [ + ("baseline", ["baseline"]), + ("reform", ["baseline", "reform"]), + ], +) +def test_failed_invocation_attach_stops_every_started_child( + failing_role, + started_roles, +) -> None: + store = AttachFailingStore(_simulation_id(failing_role)) + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + _assert_original_coordinator_failure(error, store) + assert invoker.events == [f"spawn:{role}" for role in started_roles] + # The child whose invocation could not be attached is stopped as well. + for role in started_roles: + _assert_cancelled( + store.children[_simulation_id(role)], after="the coordinator failed" + ) + assert invoker.calls[role].cancel_count == 1 + if failing_role == "baseline": + assert store.children[REFORM_ID].status is ComparisonRunLifecycleStatus.PENDING + + +def test_startup_failure_stops_a_call_restored_from_an_earlier_attempt() -> None: + report = _report_with_segmentation(segmented=False) + store = PlaceholderWriteFailingStore(REFORM_ID) + store.children[BASELINE_ID] = _child(report.baseline).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.RUNNING, + "modal_invocation_id": "call-baseline", + "started_at": NOW, + } + ) + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(report, store, invoker) + + _assert_original_coordinator_failure(error, store) + assert invoker.events == ["restore:baseline"] + _assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") + assert invoker.calls["baseline"].cancel_count == 1 + + +def test_startup_failure_leaves_a_finished_child_unchanged() -> None: + report = _report_with_segmentation(segmented=False) + store = AttachFailingStore(REFORM_ID) + artifacts = FakeArtifacts() + baseline_descriptor = artifacts.add_simulation(report.baseline) + finished_baseline = _child(report.baseline).model_copy( + update={ + "status": ComparisonRunLifecycleStatus.SUCCEEDED, + "output_uri": baseline_descriptor.artifact.uri, + "output_sha256": baseline_descriptor.artifact.content_sha256, + "output_schema_version": 1, + "row_identity_columns": baseline_descriptor.row_identity.identifier_columns, + "row_count": baseline_descriptor.row_identity.row_count, + "row_identity_sha256": baseline_descriptor.row_identity.identity_sha256, + "completed_at": NOW, + } + ) + store.children[BASELINE_ID] = finished_baseline + invoker = ScriptedChildInvoker() + + with pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error: + _coordinate_with(report, store, invoker, artifacts=artifacts) + + _assert_original_coordinator_failure(error, store) + assert invoker.events == ["spawn:reform"] + assert store.children[BASELINE_ID] == finished_baseline + _assert_cancelled(store.children[REFORM_ID], after="the coordinator failed") + assert invoker.calls["reform"].cancel_count == 1 + + +def test_startup_cancellation_error_does_not_replace_the_original_failure( + caplog, +) -> None: + store = PlaceholderWriteFailingStore(REFORM_ID) + invoker = ScriptedChildInvoker( + cancel_errors={ + "baseline": ConnectionError("cancel request contains sensitive details") + } + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error, + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + _assert_original_coordinator_failure(error, store) + assert invoker.calls["baseline"].cancel_count == 1 + _assert_cancelled(store.children[BASELINE_ID], after="the coordinator failed") + _assert_log_identifies_child( + _coordination_log(caplog, logging.WARNING), + role="baseline", + error_type="ConnectionError", + ) + + +def test_startup_cancellation_state_error_does_not_replace_the_original_failure( + caplog, +) -> None: + class FailingStore(PlaceholderWriteFailingStore, CancellationStateFailingStore): + pass + + store = FailingStore(REFORM_ID) + invoker = ScriptedChildInvoker() + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises( + RuntimeError, match="Stage 12 report coordination failed" + ) as error, + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + _assert_original_coordinator_failure(error, store) + assert invoker.calls["baseline"].cancel_count == 1 + _assert_log_identifies_child( + _coordination_log(caplog, logging.ERROR), + role="baseline", + error_type="ConnectionError", + ) + + +class ChildStateError(Exception): + """A database failure while reading or writing a child's own record.""" + + +class OwnChildStateFailingStore(FakeStore): + """Fail one read, or one failed-status write, of one child's record once. + + Cancellation writes are left alone, so only the failed child's own state + update can fail. + """ + + def __init__(self, *, failing_simulation_id, failing_operation): + super().__init__() + self.failing_simulation_id = failing_simulation_id + self.failing_operation = failing_operation + self.remaining_failures = 1 + + def _fail_once(self, operation, simulation_execution_id): + if ( + operation == self.failing_operation + and simulation_execution_id == self.failing_simulation_id + and self.remaining_failures > 0 + ): + self.remaining_failures -= 1 + raise ChildStateError("database contains sensitive details") + + def get_simulation(self, simulation_execution_id): + self._fail_once("read", simulation_execution_id) + return super().get_simulation(simulation_execution_id) + + def replace_simulation_if_status(self, record, *, expected_status): + if record.error_code != "simulation_cancelled": + self._fail_once("write", record.simulation_execution_id) + return super().replace_simulation_if_status( + record, expected_status=expected_status + ) + + +@pytest.mark.parametrize( + ("failing_operation", "baseline_status", "baseline_error_code"), + [ + # The failure is still recorded: the read before the conditional + # write succeeds. + ( + "read", + ComparisonRunLifecycleStatus.FAILED, + "simulation_invocation_failed", + ), + ("write", ComparisonRunLifecycleStatus.RUNNING, None), + ], +) +def test_failed_child_state_error_does_not_replace_the_child_failure( + caplog, + failing_operation, + baseline_status, + baseline_error_code, +) -> None: + store = OwnChildStateFailingStore( + failing_simulation_id=BASELINE_ID, + failing_operation=failing_operation, + ) + invoker = ScriptedChildInvoker( + wait_errors={"baseline": LookupError("baseline failure is sensitive")}, + ) + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + assert store.parent.error_code == "simulation_invocation_failed" + _correlation_id(store.parent.error_summary) + baseline = store.children[BASELINE_ID] + assert baseline.status is baseline_status + assert baseline.error_code == baseline_error_code + if baseline_error_code is not None: + assert baseline.error_summary == store.parent.error_summary + # The failed child is neither cancelled nor recorded as cancelled. + assert invoker.calls["baseline"].cancel_count == 0 + _assert_cancelled(store.children[REFORM_ID], after="the baseline simulation failed") + assert invoker.calls["reform"].cancel_count == 1 + _assert_log_identifies_child( + _coordination_log(caplog, logging.ERROR), + role="baseline", + error_type="ChildStateError", + ) + + +class ReformDispatchFailureInvoker(ScriptedChildInvoker): + def spawn(self, *, simulation, **values): + role = PlannedSimulationExecutionInput.model_validate(simulation).role + if role is SimulationRole.REFORM: + self.events.append("spawn:reform") + raise ConnectionError("dispatch contains sensitive details") + return super().spawn(simulation=simulation, **values) + + +@pytest.mark.parametrize("failing_operation", ["read", "write"]) +def test_dispatch_failure_state_error_does_not_replace_the_child_failure( + caplog, + failing_operation, +) -> None: + store = OwnChildStateFailingStore( + failing_simulation_id=REFORM_ID, + failing_operation=failing_operation, + ) + invoker = ReformDispatchFailureInvoker() + + with ( + caplog.at_level(logging.WARNING, logger=coordination_module.__name__), + pytest.raises(RuntimeError, match="Stage 12 report coordination failed"), + ): + _coordinate_with(_report_with_segmentation(segmented=False), store, invoker) + + assert invoker.events == ["spawn:baseline", "spawn:reform"] + assert store.parent.error_code == "simulation_dispatch_failed" + _correlation_id(store.parent.error_summary) + # The reform record could not be updated, so it keeps its dispatch + # placeholder rather than a cancellation. + reform = store.children[REFORM_ID] + assert reform.status is ComparisonRunLifecycleStatus.RUNNING + assert reform.error_code is None + _assert_cancelled(store.children[BASELINE_ID], after="the reform simulation failed") + assert invoker.calls["baseline"].cancel_count == 1 + _assert_log_identifies_child( + _coordination_log(caplog, logging.ERROR), + role="reform", + error_type="ChildStateError", + ) + + +class ExceptionRecordingRuntime(RecordingRuntime): + """Record every telemetry call that can report a failure.""" + + def __init__(self) -> None: + super().__init__() + self.reports: list[tuple[str, object, dict, object]] = [] + + def log(self, message, *, severity="INFO", attributes=None, error=None): + self.reports.append(("log", message, dict(attributes or {}), error)) + + def record_exception(self, error, *, handled, status_code=None): + self.reports.append(("exception", None, {}, error)) + + def event(self, name, *, attributes=None, **_options): + self.reports.append(("event", name, dict(attributes or {}), None)) + + def set_context(self, **_attributes): + return None + + +def _fail_simulation_calculation(runtime): + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + calculation_error = ValueError("calculation contains sensitive details") + + def fail(_simulation): + raise calculation_error + + with pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=fail, + runtime=runtime, + ) + return calculation_error, store.children[BASELINE_ID].error_summary + + +def _fail_child_call(runtime): + store = FakeStore() + child_error = LookupError("baseline failure contains sensitive details") + invoker = ScriptedChildInvoker(wait_errors={"baseline": child_error}) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + runtime=runtime, + ) + return child_error, store.parent.error_summary + + +def _fail_child_dispatch(runtime): + store = FakeStore() + dispatch_error = ConnectionError("dispatch contains sensitive details") + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report(), + store, + RefusingDispatchInvoker(dispatch_error), + runtime=runtime, + ) + return dispatch_error, store.parent.error_summary + + +def _fail_report_aggregation(runtime): + store = FakeStore() + artifacts = FakeArtifacts() + aggregation_error = ValueError("aggregation contains sensitive details") + + def fail_aggregation(**_values): + raise aggregation_error + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + coordinate_report( + _report().model_dump(mode="json"), + _context().model_dump(mode="json"), + _parent().model_dump(mode="json"), + application_name=_context().modal_application, + coordinator_invocation_id="coordinator-1", + store=store, + artifacts=artifacts, + invoker=ConcurrentInvoker(artifacts), + aggregator=fail_aggregation, + runtime=runtime, + ) + return aggregation_error, store.parent.error_summary + + +@pytest.mark.parametrize( + "fail", + [ + _fail_simulation_calculation, + _fail_child_call, + _fail_child_dispatch, + _fail_report_aggregation, + ], +) +@pytest.mark.parametrize("with_runtime", [True, False], ids=["runtime", "no-runtime"]) +def test_an_unexpected_failure_is_reported_exactly_once( + caplog, + fail, + with_runtime, +) -> None: + runtime = ExceptionRecordingRuntime() if with_runtime else None + + with caplog.at_level(logging.DEBUG): + original_error, error_summary = fail(runtime) + + correlation_id = _correlation_id(error_summary) + logged_failures = [record for record in caplog.records if record.exc_info] + if runtime is not None: + [(kind, message, attributes, reported_error)] = [ + report for report in runtime.reports if report[0] == "log" + ] + assert (kind, message) == ("log", "Stage 12 execution failed") + assert reported_error is original_error + assert attributes["correlation_id"] == correlation_id + # The redacted cause, never the original error, is also attached to + # the failing operation, once. + [recorded_cause] = [ + report[3] for report in runtime.reports if report[0] == "exception" + ] + assert str(recorded_cause) == error_summary + assert sorted(report[0] for report in runtime.reports) == ["exception", "log"] + assert logged_failures == [] + else: + [logged_failure] = logged_failures + assert logged_failure.exc_info[1] is original_error + assert f"correlation_id={correlation_id}" in logged_failure.getMessage() + + +def test_a_segment_failure_is_reported_exactly_once(caplog) -> None: + store = FakeStore() + simulation = _planned_simulation(SimulationRole.BASELINE) + store.children[simulation.simulation_execution_id] = _child(simulation) + runtime = ExceptionRecordingRuntime() + segment_error = RuntimeError("segment contains sensitive details") + + class SegmentFunction: + @staticmethod + def spawn(payload, *, observability_context): + segment_index = payload["segment_index"] + # The first segment fails on every probe; the others keep running. + return ScriptedChildCall( + f"segment-{segment_index}", + wait_error=segment_error if segment_index == 0 else None, + ) + + class SegmentFunctionLookup: + @staticmethod + def from_name(_app_name, _function_name): + return SegmentFunction + + def segmented_calculator(value, cancellation_requested): + return Stage12SegmentRunner( + value, + app_name="stage12-app", + runtime=runtime, + modal_module=SimpleNamespace(Function=SegmentFunctionLookup), + groups=US_REGION_GROUPS, + sleep=lambda _seconds: None, + cancellation_requested=cancellation_requested, + ).run() + + with ( + caplog.at_level(logging.DEBUG), + pytest.raises(RuntimeError, match="Stage 12 simulation execution failed"), + ): + run_single_simulation( + simulation.model_dump(mode="json"), + _context().model_dump(mode="json"), + required_country="us", + store=store, + artifacts=FakeArtifacts(), + calculator=lambda _: pytest.fail("monolithic calculator was selected"), + segmented_calculator=segmented_calculator, + runtime=runtime, + ) + + child = store.children[BASELINE_ID] + assert child.error_code == "segment_execution_failed" + [(kind, message, attributes, reported_error)] = [ + report for report in runtime.reports if report[0] == "log" + ] + [recorded_cause] = [ + report[3] for report in runtime.reports if report[0] == "exception" + ] + assert str(recorded_cause) == child.error_summary + assert (kind, message) == ("log", "Stage 12 execution failed") + assert reported_error is segment_error + assert attributes == { + "stage12_failure_scope": "segment_call", + "simulation_role": "baseline", + "evaluation_id": str(EVALUATION_ID), + "simulation_execution_id": str(BASELINE_ID), + "segment_index": 0, + "correlation_id": _correlation_id( + child.error_summary.removeprefix("Segment 0: ") + ), + } + assert [record for record in caplog.records if record.exc_info] == [] + + +class SteppedClock: + def __init__(self, start_seconds=10_000.0): + self.current_seconds = start_seconds + + def monotonic(self): + return self.current_seconds + + def advance(self, elapsed_seconds): + self.current_seconds += elapsed_seconds + + +def _declared_modal_timeouts(monkeypatch): + install_fake_modal(monkeypatch) + sys.modules.pop("src.modal.v2_app", None) + modal_app = importlib.import_module("src.modal.v2_app") + return { + name: options.get("timeout") for name, options in modal_app.app.function_calls + } + + +# After the child window closes and the production-result wait ends, the +# coordinator still has to read both simulation artifacts, aggregate, write its +# artifacts, and record the outcome before Modal stops it. +MINIMUM_FINISHING_SECONDS = 600 + + +def test_declared_modal_timeouts_fit_the_shared_wait_window(monkeypatch) -> None: + declared_timeouts = _declared_modal_timeouts(monkeypatch) + window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + child_worker_timeouts = { + name: timeout + for name, timeout in declared_timeouts.items() + if name.startswith("run_single_simulation") + } + + assert set(child_worker_timeouts) == { + "run_single_simulation_us", + "run_single_simulation_uk", + "run_single_simulation_segment_us", + } + # The window covers every worker's own Modal timeout. + assert all(window_seconds >= timeout for timeout in child_worker_timeouts.values()) + assert ( + declared_timeouts["coordinate_report"] + >= window_seconds + + PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + + MINIMUM_FINISHING_SECONDS + ) + + +@pytest.mark.parametrize("setup_seconds", [0.0, 400.0, 2_000.0]) +def test_child_wait_window_opens_when_the_coordinator_starts( + monkeypatch, + setup_seconds, +) -> None: + declared_timeouts = _declared_modal_timeouts(monkeypatch) + window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + clock = SteppedClock() + coordinator_started_at = clock.monotonic() + + def slow_output_planning(_report): + # Loading the country model is coordinator setup, not child time. + clock.advance(setup_seconds) + return _output_plan() + + store = FakeStore() + invoker = ScriptedChildInvoker(clock=clock) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report(), + store, + invoker, + output_plan_resolver=slow_output_planning, + monotonic=clock.monotonic, + sleep=clock.advance, + ) + + probed_at = invoker.calls["baseline"].probed_at + child_waits_began_after = probed_at[0] - coordinator_started_at + window_closed_after = probed_at[-1] - coordinator_started_at + assert child_waits_began_after == setup_seconds + # Setup shortens the children's window rather than the time after it. + assert window_closed_after == pytest.approx(window_seconds) + assert probed_at[-1] - probed_at[0] == pytest.approx(window_seconds - setup_seconds) + assert ( + declared_timeouts["coordinate_report"] + - window_closed_after + - PRODUCTION_RESULT_WAIT_TIMEOUT_SECONDS + >= MINIMUM_FINISHING_SECONDS + ) + assert store.parent.error_code == "simulation_wait_timeout" + for role in ("baseline", "reform"): + _assert_cancelled( + store.children[_simulation_id(role)], + after=f"the {window_seconds:g}-second wait window closed", + ) + + +def test_exhausted_wait_window_gives_the_children_no_extra_time() -> None: + window_seconds = coordination_module.SIMULATION_WAIT_TIMEOUT_SECONDS + clock = SteppedClock() + + def setup_that_uses_the_whole_window(_report): + clock.advance(window_seconds + 60.0) + return _output_plan() + + store = FakeStore() + invoker = ScriptedChildInvoker(clock=clock) + + with pytest.raises(RuntimeError, match="Stage 12 report coordination failed"): + _coordinate_with( + _report_with_segmentation(segmented=False), + store, + invoker, + output_plan_resolver=setup_that_uses_the_whole_window, + monotonic=clock.monotonic, + sleep=lambda _: pytest.fail("an exhausted window must not wait"), + ) + + assert [len(call.probed_at) for call in invoker.calls.values()] == [1, 1] + assert [call.cancel_count for call in invoker.calls.values()] == [1, 1] + assert store.parent.error_code == "simulation_wait_timeout" diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py b/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py new file mode 100644 index 000000000..28d15f6c2 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_calculation.py @@ -0,0 +1,249 @@ +"""Unit tests for Stage 12 segment calculation and merging.""" + +from __future__ import annotations + +from hashlib import sha256 + +import pandas as pd +import pytest +from policyengine_simulation_contract.spm import ( + SPMProvenance, + SPMRuntimeVersions, + SPMSelection, + build_spm_provenance, +) +from policyengine_simulation_contract.stage12_execution import ( + PlannedSimulationExecutionInput, + SimulationRole, +) +from policyengine_simulation_executor.stage12_artifacts import ( + serialize_simulation_frames, +) +from policyengine_simulation_executor.stage12_runtime import ( + SimulationCalculation, + Stage12SegmentResult, + build_segment_inputs, + calculate_segment, + merge_segment_results, + should_segment_simulation, +) +from test_stage12_runtime import _planned_simulation + + +def _frames(index: int) -> dict[str, pd.DataFrame]: + return { + "household": pd.DataFrame( + { + "household_id": pd.Series([index], dtype="int64"), + "household_net_income": pd.Series([index + 0.5], dtype="float32"), + } + ), + "person": pd.DataFrame( + { + "age": pd.Series([30 + index], dtype="int16"), + "household_id": pd.Series([index], dtype="int64"), + "person_id": pd.Series([100 + index], dtype="int64"), + } + ), + } + + +def _results( + simulation: PlannedSimulationExecutionInput, provenance: SPMProvenance | None = None +) -> list[Stage12SegmentResult]: + results = [] + for segment in build_segment_inputs(simulation): + payload, _ = serialize_simulation_frames( + _frames(segment.segment_index), calculation_provenance=provenance + ) + results.append( + Stage12SegmentResult( + segment_index=segment.segment_index, + role=simulation.role, + region_codes=segment.region_codes, + parquet_payload=payload, + payload_sha256=sha256(payload).hexdigest(), + ) + ) + return results + + +def _spm_provenance() -> SPMProvenance: + """Build an explicit test fixture, not installed runtime version metadata.""" + selection = SPMSelection( + forecast_content_sha256="f" * 64, + scenario="official", + geography_kind="national", + geography_id=None, + county_vintage="2020", + as_of=None, + ) + return build_spm_provenance( + forecast_id="test-forecast", + forecast_sha256="f" * 64, + selection=selection, + years=("2026",), + runtime_versions=SPMRuntimeVersions( + policyengine="test-fixture", + policyengine_core="test-fixture", + policyengine_us="test-fixture", + spm_calculator="test-fixture", + ), + ) + + +def test_merge_preserves_one_identical_compact_spm_receipt() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + receipt = _spm_provenance() + + calculated = merge_segment_results(simulation, _results(simulation, receipt)) + + assert calculated.calculation_provenance == receipt + assert isinstance(calculated.calculation_provenance, SPMProvenance) + + +@pytest.mark.parametrize("difference", ["missing", "scenario", "versions", "years"]) +def test_merge_rejects_missing_or_incompatible_segment_spm_receipts( + difference: str, +) -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + receipt = _spm_provenance() + results = _results(simulation, receipt) + alternative = { + "missing": None, + "scenario": receipt.model_copy(update={"scenario": "alternative"}), + "versions": receipt.model_copy( + update={ + "runtime_versions": receipt.runtime_versions.model_copy( + update={"policyengine": "different-test-fixture"} + ) + } + ), + "years": receipt.model_copy(update={"years": ("2027",)}), + }[difference] + payload, _ = serialize_simulation_frames( + _frames(0), calculation_provenance=alternative + ) + results[0] = results[0].model_copy( + update={ + "parquet_payload": payload, + "payload_sha256": sha256(payload).hexdigest(), + } + ) + + with pytest.raises(ValueError, match="segment calculation"): + merge_segment_results(simulation, results) + + +def test_eligibility_matches_stage12_supported_national_shape() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + assert should_segment_simulation(simulation) is True + assert ( + should_segment_simulation( + simulation.model_copy(update={"options": {"segmented": False}}) + ) + is False + ) + assert ( + should_segment_simulation( + simulation.model_copy( + update={ + "geography": simulation.geography.model_copy( + update={"region": "state/ca"} + ) + } + ) + ) + is False + ) + assert ( + should_segment_simulation( + simulation.model_copy( + update={ + "output_plan": simulation.output_plan.model_copy( + update={ + "requirements": simulation.output_plan.requirements.model_copy( + update={"include_cliff_impacts": True} + ) + } + ) + } + ) + ) + is False + ) + + +def test_builds_twenty_single_policy_segment_inputs() -> None: + simulation = _planned_simulation(SimulationRole.REFORM) + + segments = build_segment_inputs(simulation) + + assert len(segments) == 20 + assert {segment.segment_index for segment in segments} == set(range(20)) + assert all(segment.simulation.role is SimulationRole.REFORM for segment in segments) + assert all(segment.simulation.policy == simulation.policy for segment in segments) + + +def test_segment_calculation_scopes_one_group_and_preserves_dtypes() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + segment = build_segment_inputs(simulation)[0] + received = [] + + def calculator(scoped): + received.append(scoped) + return SimulationCalculation(frames=_frames(0)) + + result = calculate_segment(segment, calculator=calculator) + merged = merge_segment_results( + simulation, + [ + calculate_segment( + item, + calculator=lambda scoped, index=item.segment_index: ( + SimulationCalculation(frames=_frames(index)) + ), + ) + for item in build_segment_inputs(simulation) + ], + ) + + assert received[0].options["region_group"] == list(segment.region_codes) + assert "segmented" not in received[0].options + assert result.role is SimulationRole.BASELINE + assert str(merged.frames["household"]["household_net_income"].dtype) == "float32" + assert str(merged.frames["person"]["age"].dtype) == "int16" + assert merged.frames["household"]["household_id"].tolist() == list(range(20)) + + +def test_merge_rejects_missing_duplicate_and_cross_role_results() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + results = _results(simulation) + + with pytest.raises(ValueError, match="incomplete"): + merge_segment_results(simulation, results[:-1]) + with pytest.raises(ValueError, match="duplicate index"): + merge_segment_results(simulation, [*results[:-1], results[0]]) + with pytest.raises(ValueError, match="another simulation role"): + merge_segment_results( + simulation, + [ + results[0].model_copy(update={"role": SimulationRole.REFORM}), + *results[1:], + ], + ) + + +def test_merge_rejects_duplicate_entity_ids_across_groups() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + results = _results(simulation) + duplicate_payload, _ = serialize_simulation_frames(_frames(0)) + results[1] = results[1].model_copy( + update={ + "parquet_payload": duplicate_payload, + "payload_sha256": sha256(duplicate_payload).hexdigest(), + } + ) + + with pytest.raises(ValueError, match="duplicate IDs"): + merge_segment_results(simulation, results) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py new file mode 100644 index 000000000..278c4919f --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segment_fanout.py @@ -0,0 +1,215 @@ +"""Unit tests for independent Stage 12 Modal segment orchestration.""" + +from __future__ import annotations + +from contextlib import contextmanager +from types import SimpleNamespace +from typing import cast + +import pandas as pd +import pytest +from policyengine_observability import ObservabilityRuntime +from policyengine_simulation_contract.stage12_execution import SimulationRole +from policyengine_simulation_executor.stage12_runtime import ( + SimulationCalculation, + Stage12Cancellation, + Stage12ExecutionError, + Stage12SegmentRunner, + build_segment_inputs, + calculate_segment, +) +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, +) +from test_stage12_runtime import _planned_simulation + + +def _frames(index: int) -> dict[str, pd.DataFrame]: + return { + "household": pd.DataFrame( + { + "household_id": pd.Series([index], dtype="int64"), + "household_net_income": pd.Series([index + 0.5], dtype="float32"), + } + ), + "person": pd.DataFrame( + { + "age": pd.Series([30 + index], dtype="int16"), + "household_id": pd.Series([index], dtype="int64"), + "person_id": pd.Series([100 + index], dtype="int64"), + } + ), + } + + +class FakeRuntime: + def __init__(self) -> None: + self.context: dict[str, object] = {} + self.spans: list[tuple[str, dict | None]] = [] + + def set_context(self, **attributes) -> None: + self.context.update(attributes) + + def capture_context(self) -> dict[str, str]: + return {"observability_id": "00000000-0000-4000-8000-000000000001"} + + @contextmanager + def span(self, name, *, attributes=None): + self.spans.append((name, attributes)) + yield + + +class FakeCall: + def __init__(self, outcomes: list[object]) -> None: + self.outcomes = outcomes + self.get_timeouts: list[float] = [] + self.cancelled = 0 + + def get(self, *, timeout: float): + self.get_timeouts.append(timeout) + outcome = self.outcomes.pop(0) + if isinstance(outcome, BaseException): + raise outcome + return outcome + + def cancel(self) -> None: + self.cancelled += 1 + + +class FakeFunction: + def __init__(self, results: list[dict], errors: dict[int, list[Exception]]): + self.results = results + self.errors = errors + self.spawns: list[tuple[dict, dict | None]] = [] + self.calls: list[FakeCall] = [] + + def spawn(self, payload: dict, *, observability_context: dict | None): + index = payload["segment_index"] + self.spawns.append((payload, observability_context)) + call = FakeCall([*self.errors.get(index, []), self.results[index]]) + self.calls.append(call) + return call + + +def _results(simulation) -> list[dict]: + return [ + calculate_segment( + segment, + calculator=lambda scoped, index=segment.segment_index: ( + SimulationCalculation(frames=_frames(index)) + ), + ).model_dump(mode="python") + for segment in build_segment_inputs(simulation) + ] + + +def _runner(*, errors: dict[int, list[Exception]] | None = None): + simulation = _planned_simulation(SimulationRole.BASELINE) + function = FakeFunction(_results(simulation), errors or {}) + + class FunctionLookup: + @staticmethod + def from_name(app_name, function_name): + assert app_name == "stage12-app" + assert function_name == "run_single_simulation_segment_us" + return function + + runtime = FakeRuntime() + sleeps: list[float] = [] + runner = Stage12SegmentRunner( + simulation, + app_name="stage12-app", + runtime=cast(ObservabilityRuntime, runtime), + modal_module=SimpleNamespace(Function=FunctionLookup), + groups=US_REGION_GROUPS, + sleep=sleeps.append, + poll_interval_seconds=0.5, + poll_interval_max_seconds=2.0, + ) + return runner, function, runtime, sleeps + + +def test_starts_twenty_segments_and_merges_them_in_partition_order() -> None: + runner, function, runtime, sleeps = _runner() + + result = runner.run() + + assert len(function.spawns) == 20 + assert [payload["segment_index"] for payload, _ in function.spawns] == list( + range(20) + ) + assert all(context is not None for _, context in function.spawns) + assert all(call.get_timeouts == [0] for call in function.calls) + assert result.frames["household"]["household_id"].tolist() == list(range(20)) + assert runtime.context["stage12_segment_count"] == 20 + assert sleeps == [] + + +def test_retries_one_poll_error_without_restarting_the_segment() -> None: + runner, function, _, _ = _runner(errors={0: [ConnectionError("transient")]}) + + runner.run() + + assert len(function.spawns) == 20 + assert function.calls[0].get_timeouts == [0, 0] + assert not any(call.cancelled for call in function.calls) + + +def test_not_ready_polls_use_bounded_backoff() -> None: + runner, function, _, sleeps = _runner( + errors={index: [TimeoutError()] for index in range(20)} + ) + + runner.run() + + assert sleeps == [0.5] + assert all(call.get_timeouts == [0, 0] for call in function.calls) + + +def test_second_poll_error_cancels_all_started_segments() -> None: + runner, function, _, _ = _runner( + errors={0: [ConnectionError("first"), RuntimeError("second")]} + ) + + with pytest.raises(Stage12ExecutionError) as error: + runner.run() + + assert len(function.calls) == 20 + assert all(call.cancelled == 1 for call in function.calls) + assert error.value.detail.error_code == "segment_execution_failed" + assert error.value.detail.segment_index == 0 + assert error.value.detail.error_summary.startswith( + "Segment 0: RuntimeError: second (correlation_id=" + ) + + +def test_dispatch_error_cancels_segments_that_already_started() -> None: + runner, function, _, _ = _runner() + spawn = function.spawn + + def fail_during_dispatch(payload, *, observability_context): + if payload["segment_index"] == 5: + raise ConnectionError("dispatch failed") + return spawn(payload, observability_context=observability_context) + + function.spawn = fail_during_dispatch + + with pytest.raises(ConnectionError, match="dispatch failed"): + runner.run() + + assert len(function.calls) == 5 + assert all(call.cancelled == 1 for call in function.calls) + + +def test_persisted_parent_cancellation_cancels_all_segments() -> None: + runner, function, _, sleeps = _runner( + errors={index: [TimeoutError()] for index in range(20)} + ) + checks = iter([False, False, True]) + runner.cancellation_requested = lambda: next(checks, True) + + with pytest.raises(Stage12Cancellation): + runner.run() + + assert sleeps == [0.5] + assert all(call.cancelled == 1 for call in function.calls) diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py b/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py new file mode 100644 index 000000000..c2e23da81 --- /dev/null +++ b/projects/policyengine-simulation-executor/tests/test_stage12_segmentation_foundation.py @@ -0,0 +1,112 @@ +"""Foundational tests for Stage 12-owned segmented execution.""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import SimpleNamespace + +import pytest +from policyengine_simulation_contract.stage12_execution import SimulationRole +from policyengine_simulation_executor.stage12_runtime.partition import ( + US_REGION_GROUPS, + stage12_region_groups, + stage12_region_groups_for_model, +) +from policyengine_simulation_executor.stage12_runtime.segment_contracts import ( + Stage12SegmentInput, +) +from pydantic import ValidationError +from test_stage12_runtime import _planned_simulation + +PROJECT = Path(__file__).resolve().parents[1] +PACKAGE = PROJECT / "src" / "policyengine_simulation_executor" +FORBIDDEN_IMPORTS = { + "policyengine_simulation_executor.artifact_keys", + "policyengine_simulation_executor.baseline_artifacts", + "policyengine_simulation_executor.national_partition", + "policyengine_simulation_executor.precompute", + "policyengine_simulation_executor.precompute_models", + "policyengine_simulation_executor.segmented_national_reduce", + "src.modal._image_setup", + "src.modal.precompute_app", + "src.modal.segmented_national", +} + + +def test_us_partition_has_twenty_unique_groups_and_complete_state_coverage() -> None: + groups = stage12_region_groups("us") + + assert groups is not None + assert len(groups) == 20 + assert all(groups) + flattened = [code for group in groups for code in group] + assert len(flattened) == len(set(flattened)) == 51 + assert set(flattened) == {code for group in US_REGION_GROUPS for code in group} + assert stage12_region_groups("uk") is None + + +def test_partition_returns_independent_mutable_copies() -> None: + first = stage12_region_groups("us") + second = stage12_region_groups("us") + assert first is not None and second is not None + + first[-1].append("state/test") + + assert "state/test" not in second[-1] + + +def test_new_model_regions_are_added_to_the_final_group_in_canonical_order() -> None: + model = SimpleNamespace( + region_registry=SimpleNamespace( + regions=[ + SimpleNamespace(code="state/zz", region_type="state"), + SimpleNamespace(code="state/aa", region_type="state"), + SimpleNamespace(code="district/1", region_type="district"), + ] + ) + ) + + groups = stage12_region_groups_for_model("us", model) + + assert groups[-1][-2:] == ["state/aa", "state/zz"] + + +def test_segment_input_rejects_duplicate_or_non_state_regions() -> None: + simulation = _planned_simulation(SimulationRole.BASELINE) + + with pytest.raises(ValidationError, match="must be unique"): + Stage12SegmentInput( + simulation=simulation, + segment_index=0, + region_codes=("state/ca", "state/ca"), + ) + with pytest.raises(ValidationError, match="state regions"): + Stage12SegmentInput( + simulation=simulation, + segment_index=0, + region_codes=("district/1",), + ) + + +def test_stage12_feature_modules_do_not_import_v1_feature_implementations() -> None: + feature_paths = list((PACKAGE / "stage12_runtime").glob("*.py")) + cache_package = PACKAGE / "stage12_cache" + if cache_package.exists(): + feature_paths.extend(cache_package.glob("*.py")) + + violations: list[str] = [] + for path in feature_paths: + tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + for node in ast.walk(tree): + if isinstance(node, ast.Import): + imported = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + imported = [node.module or ""] + else: + continue + for module in imported: + if module in FORBIDDEN_IMPORTS: + violations.append(f"{path.name}: {module}") + + assert violations == [] diff --git a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py index 057790992..fb86ead35 100644 --- a/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py +++ b/projects/policyengine-simulation-executor/tests/test_stage12_worker_validation.py @@ -19,6 +19,7 @@ def _environment() -> dict[str, str]: "GOOGLE_APPLICATION_CREDENTIALS_JSON": "{}", "STAGE12_DATABASE_URL": "postgresql://runtime:secret@db.example/postgres", "STAGE12_ARTIFACT_BUCKET": "policyengine-stage12-staging", + "STAGE12_CACHE_BUCKET": "policyengine-stage12-cache-staging", } @@ -34,6 +35,7 @@ def test_validation_checks_dataset_and_non_serving_calculation() -> None: dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", dataset_check=lambda path, digest: datasets.append((path, digest)), calculation_check=countries.append, + cache_check=lambda digest: {"manifest_sha256": digest}, ) assert result["validated"] is True @@ -48,6 +50,7 @@ def test_validation_checks_dataset_and_non_serving_calculation() -> None: ) assert datasets == [("/installed/populace_us_2024.h5", expected_dataset.sha256)] assert countries == ["us"] + assert result["cache"] is None def test_non_serving_us_calculation_supports_national_spm() -> None: @@ -62,7 +65,10 @@ def test_non_serving_us_calculation_selects_national_without_county( def calculate_household(**kwargs): calls.append(kwargs) - return SimpleNamespace(household={"household_net_income": 1}) + return SimpleNamespace( + household={"household_net_income": 1}, + tax_unit={"ctc": 1}, + ) monkeypatch.setattr( "policyengine_simulation_executor.stage12_worker_validation.import_module", @@ -75,6 +81,35 @@ def calculate_household(**kwargs): assert "household" not in calls[0] +def test_us_validation_preserves_the_frontend_reform_interval(monkeypatch) -> None: + from policyengine_simulation_executor import stage12_worker_validation + + received = {} + + def calculate_household(**kwargs): + received.update(kwargs) + return SimpleNamespace( + household={"household_net_income": 50_000}, + tax_unit={"ctc": 3_000}, + ) + + monkeypatch.setattr( + stage12_worker_validation, + "import_module", + lambda _: SimpleNamespace(calculate_household=calculate_household), + ) + + _run_non_serving_calculation("us") + + assert received["reform"] == { + # The current wrapper compiles the complete interval; do not discard + # its end date while passing the front-end reform to calculate_household. + "gov.irs.credits.ctc.amount.base[0].amount": {"2026-01-01.2100-12-31": 3_000} + } + assert received["people"][1]["is_tax_unit_dependent"] is True + assert "ctc" in received["extra_variables"] + + def test_uk_validation_loads_packaged_local_authority_resources() -> None: resolved = load_stage12_bundle() checks: list[str] = [] @@ -107,6 +142,33 @@ def test_validation_rejects_digest_mismatch_before_dataset_access() -> None: assert accessed == [] +def test_us_validation_reports_the_verified_cache_release() -> None: + resolved = load_stage12_bundle() + seen: list[str] = [] + cache = { + "manifest_sha256": "c" * 64, + "partition_sha256": "d" * 64, + "years": [2026, 2027, 2025], + "segment_count": 20, + "dataset_count": 3, + "baseline_count": 60, + } + + result = validate_country_worker( + country="us", + expected_bundle_manifest_sha256=resolved.bundle_manifest_sha256, + expected_cache_manifest_sha256="c" * 64, + environment=_environment(), + dataset_path_resolver=lambda _: "/installed/populace_us_2024.h5", + dataset_check=lambda *_: None, + calculation_check=lambda _: None, + cache_check=lambda digest: seen.append(digest) or cache, + ) + + assert seen == ["c" * 64] + assert result["cache"] == cache + + @pytest.mark.parametrize( "environment", [ diff --git a/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py b/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py index d6f37a66a..c5b555c9d 100644 --- a/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py +++ b/projects/policyengine-simulation-executor/tests/test_validate_v2_worker_release.py @@ -28,13 +28,28 @@ def fake_spawn_validation(*, application_name, function_name, environment): "dataset_uri": bundle.default_dataset_uri, "data_artifact_revision": bundle.data_artifact_revision, "bundle_manifest_sha256": resolved.bundle_manifest_sha256, + "cache": ( + { + "manifest_sha256": "c" * 64, + "partition_sha256": "d" * 64, + "years": [2026, 2027, 2025], + "segment_count": 20, + "dataset_count": 3, + "baseline_count": 60, + } + if country == "us" + else None + ), }, invocation_id, ) monkeypatch.setattr(validation, "_spawn_validation", fake_spawn_validation) - result = validation.validate_release(environment="staging") + result = validation.validate_release( + environment="staging", + expected_cache_manifest_sha256="c" * 64, + ) assert result.validated is True assert result.country_validation_invocation_ids == { @@ -42,6 +57,8 @@ def fake_spawn_validation(*, application_name, function_name, environment): "uk": "call-uk", } assert invocation_ids == ["call-us", "call-uk"] + assert result.cache is not None + assert result.cache.manifest_sha256 == "c" * 64 def test_country_result_mismatch_stops_publication(monkeypatch) -> None: @@ -59,4 +76,7 @@ def test_country_result_mismatch_stops_publication(monkeypatch) -> None: ) with pytest.raises(RuntimeError, match="differ from deployment"): - validation.validate_release(environment="staging") + validation.validate_release( + environment="staging", + expected_cache_manifest_sha256="c" * 64, + ) diff --git a/rules/registry.toml b/rules/registry.toml index ce8b51016..8895bd17c 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -197,7 +197,7 @@ invocation = "uv run pytest tests/" [[rule]] id = "sim.stage12.deployment-scope" category = "stage12" -statement = "Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources; the Stage 12 worker can request a Google OIDC token." +statement = "Main-branch deploys roll out Stage 12 in both environments; a Stage 12-only deploy runs candidate tests and the requested promotion but never redeploys existing Modal resources or runs the existing stack's complete integration suite; the Stage 12 worker can request a Google OIDC token." rationale = "docs/operations" scope = [".github/workflows/*.yml"] tier = "mechanical" @@ -210,6 +210,7 @@ nodes = [ "tests/test_deployment_assets.py::test_main_deployment_automatically_deploys_stage12_in_both_environments", "tests/test_deployment_assets.py::test_stage12_only_deployment_runs_candidate_tests_and_requested_promotion", "tests/test_deployment_assets.py::test_stage12_only_deployment_cannot_redeploy_existing_modal_resources", + "tests/test_deployment_assets.py::test_stage12_only_deployment_skips_existing_stack_integration_suite", "tests/test_deployment_assets.py::test_stage12_worker_deployment_can_request_a_google_oidc_token", "tests/test_deployment_assets.py::test_stage12_only_promotion_runs_whatever_the_skipped_stack_jobs", ]