From 4748fe7c38a8eb8f4c723309f146071a4397a97e Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 02:02:37 +0400 Subject: [PATCH 1/2] Add the simulation job store and background result collection Add simulation.jobs and simulation.job_results (revision 0002) and the JobStore that owns them. One jobs row records each calculation: submitter, an input hash covering the job kind, policyengine.py bundle version and request, the Modal application, function and call ID, status, timings, collection checks and any error. Every transition is a conditional update on the expected status, and no lease is taken: submission inserts with ON CONFLICT DO NOTHING on a unique index over the input hash that excludes failed jobs; only the instance that moves a job from pending to dispatching spawns its Modal call; a result or failure is written only while the job is running, so a second collector changes nothing. PostgreSQL tests force the interleaving (the second writer blocks on the first's row lock, then finds nothing to change) and race eight threads per job. In the Simulation Entrypoint, SimulationJobService submits jobs, collects results from Cloud Tasks requests (checking the Modal call without waiting, validating the result against the job kind's model, and re-enqueuing with exponential backoff while it runs), and sweeps from Cloud Scheduler requests (re-enqueuing running jobs not checked for five minutes and failing dispatches stalled for ten). The internal collection and sweep routes accept only Google OIDC tokens of the invoker service account. A status route reads a job by UUID or Modal fc- call ID from the database only. No job kind is registered yet, no route submits jobs, and none of the routes is in the OpenAPI document. The deploy step passes RECONCILIATION_QUEUE and RECONCILIATION_INVOKER_SERVICE_ACCOUNT from GitHub environment variables, the stable service URL as RECONCILIATION_TARGET_URL, and the Modal environment. Co-Authored-By: Claude Opus 5.5 --- .../deploy-cloud-run-simulation-entry.sh | 16 + .../workflows/simulation-deploy.reusable.yml | 3 + AGENTS.md | 51 +- docs/generated/rules.md | 24 +- .../README.md | 32 + .../jobs.py | 561 +++++++++++++++ .../schema.py | 2 +- .../tests/conftest.py | 9 + .../tests/job_store_support.py | 58 ++ .../tests/test_job_store.py | 256 +++++++ .../tests/test_job_store_concurrency.py | 268 ++++++++ migrations/versions/0002_simulation_jobs.py | 116 ++++ .../DEPLOYMENT.md | 26 + .../policyengine-simulation-entry/README.md | 5 + .../pyproject.toml | 2 + .../src/policyengine_simulation_entry/app.py | 215 +++++- .../policyengine_simulation_entry/config.py | 75 ++ .../google_oidc.py | 69 ++ .../policyengine_simulation_entry/schemas.py | 62 ++ .../simulation_jobs.py | 643 ++++++++++++++++++ .../tests/simulation_job_support.py | 295 ++++++++ .../tests/test_config.py | 77 +++ .../tests/test_deployment_assets.py | 9 + .../tests/test_google_oidc.py | 94 +++ .../tests/test_simulation_job_adapters.py | 239 +++++++ .../tests/test_simulation_job_routes.py | 263 +++++++ .../tests/test_simulation_jobs.py | 404 +++++++++++ .../tests/test_simulation_jobs_deployment.py | 106 +++ .../policyengine-simulation-entry/uv.lock | 7 + rules/registry.toml | 66 ++ 30 files changed, 4004 insertions(+), 49 deletions(-) create mode 100644 libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/jobs.py create mode 100644 libs/policyengine-simulation-persistence/tests/job_store_support.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_job_store.py create mode 100644 libs/policyengine-simulation-persistence/tests/test_job_store_concurrency.py create mode 100644 migrations/versions/0002_simulation_jobs.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/google_oidc.py create mode 100644 projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_jobs.py create mode 100644 projects/policyengine-simulation-entry/tests/simulation_job_support.py create mode 100644 projects/policyengine-simulation-entry/tests/test_google_oidc.py create mode 100644 projects/policyengine-simulation-entry/tests/test_simulation_job_adapters.py create mode 100644 projects/policyengine-simulation-entry/tests/test_simulation_job_routes.py create mode 100644 projects/policyengine-simulation-entry/tests/test_simulation_jobs.py create mode 100644 projects/policyengine-simulation-entry/tests/test_simulation_jobs_deployment.py diff --git a/.github/scripts/deploy-cloud-run-simulation-entry.sh b/.github/scripts/deploy-cloud-run-simulation-entry.sh index d0c042930..5dee7c33d 100755 --- a/.github/scripts/deploy-cloud-run-simulation-entry.sh +++ b/.github/scripts/deploy-cloud-run-simulation-entry.sh @@ -78,6 +78,9 @@ required_variables=( OLD_GATEWAY_AUTH_CLIENT_ID_VALUE OLD_GATEWAY_AUTH_CLIENT_SECRET_SECRET_NAME SIMULATION_DATABASE_URL_SECRET_NAME + RECONCILIATION_QUEUE_VALUE + RECONCILIATION_INVOKER_SERVICE_ACCOUNT_VALUE + ENTRYPOINT_STABLE_URL OBSERVABILITY_SERVICE_NAMESPACE OBSERVABILITY_TRACE_PROJECT_ID OTEL_EXPORTER_OTLP_ENDPOINT @@ -110,6 +113,15 @@ esac || fail "ENTRYPOINT_MAX_INSTANCES must be a non-negative integer" ((10#${ENTRYPOINT_MAX_INSTANCES} >= 10#${ENTRYPOINT_MIN_INSTANCES})) \ || fail "ENTRYPOINT_MAX_INSTANCES must be at least ENTRYPOINT_MIN_INSTANCES" +# Background collection of simulation job results: Cloud Tasks calls the +# stable service URL (never the candidate tag) with an OIDC token of the +# invoker service account, which Cloud Scheduler also uses for the sweep. +[[ "${RECONCILIATION_QUEUE_VALUE}" =~ ^projects/[a-z][a-z0-9-]{4,28}[a-z0-9]/locations/[a-z0-9-]+/queues/[A-Za-z0-9-]{1,100}$ ]] \ + || fail "RECONCILIATION_QUEUE_VALUE must be a full Cloud Tasks queue name" +[[ "${RECONCILIATION_INVOKER_SERVICE_ACCOUNT_VALUE}" =~ ^[a-z][a-z0-9-]{4,28}[a-z0-9]@[a-z][a-z0-9-]{4,28}[a-z0-9]\.iam\.gserviceaccount\.com$ ]] \ + || fail "RECONCILIATION_INVOKER_SERVICE_ACCOUNT_VALUE must be a service account email" +[[ "${ENTRYPOINT_STABLE_URL}" =~ ^https://[a-z0-9.-]+$ ]] \ + || fail "ENTRYPOINT_STABLE_URL must be the service's HTTPS URL without a path" # An unset GitHub variable arrives as an empty value. Default it to log-only # permission checks; switching to 1 is a separate operator decision. export SIMULATION_ENTRYPOINT_ENFORCE_PERMISSIONS_VALUE="${SIMULATION_ENTRYPOINT_ENFORCE_PERMISSIONS_VALUE:-0}" @@ -173,6 +185,10 @@ jq -n ' OTEL_METRICS_EXPORTER: "otlp", OTEL_TRACES_SAMPLER_ARG: "1.0", POLICYENGINE_OTEL_GOOGLE_AUDIENCE: env.POLICYENGINE_OTEL_GOOGLE_AUDIENCE, + RECONCILIATION_QUEUE: env.RECONCILIATION_QUEUE_VALUE, + RECONCILIATION_INVOKER_SERVICE_ACCOUNT: env.RECONCILIATION_INVOKER_SERVICE_ACCOUNT_VALUE, + RECONCILIATION_TARGET_URL: env.ENTRYPOINT_STABLE_URL, + SIMULATION_MODAL_ENVIRONMENT: env.MODAL_ENVIRONMENT, STAGE12_ENABLED: ( if env.DEPLOY_STAGE12_V2 == "true" then env.STAGE12_ENABLED_VALUE else "0" end ) diff --git a/.github/workflows/simulation-deploy.reusable.yml b/.github/workflows/simulation-deploy.reusable.yml index a48931da4..7acf5b852 100644 --- a/.github/workflows/simulation-deploy.reusable.yml +++ b/.github/workflows/simulation-deploy.reusable.yml @@ -265,6 +265,9 @@ jobs: OLD_GATEWAY_AUTH_CLIENT_ID_VALUE: ${{ secrets.OLD_GATEWAY_AUTH_CLIENT_ID }} OLD_GATEWAY_AUTH_CLIENT_SECRET_SECRET_NAME: ${{ vars.OLD_GATEWAY_AUTH_CLIENT_SECRET_SECRET_NAME }} SIMULATION_DATABASE_URL_SECRET_NAME: ${{ vars.SIMULATION_DATABASE_URL_SECRET_NAME }} + RECONCILIATION_QUEUE_VALUE: ${{ vars.RECONCILIATION_QUEUE }} + RECONCILIATION_INVOKER_SERVICE_ACCOUNT_VALUE: ${{ vars.RECONCILIATION_INVOKER_SERVICE_ACCOUNT }} + ENTRYPOINT_STABLE_URL: ${{ steps.previous.outputs.stable_url }} STAGE12_ENABLED_VALUE: ${{ vars.STAGE12_ENABLED }} STAGE12_ARTIFACT_BUCKET_VALUE: ${{ vars.STAGE12_ARTIFACT_BUCKET }} STAGE12_DATABASE_URL_SECRET_NAME: ${{ vars.STAGE12_DATABASE_URL_SECRET_NAME }} diff --git a/AGENTS.md b/AGENTS.md index 337bd8967..9901a42dc 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -97,45 +97,20 @@ 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 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-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`) -- `sim.stage12.no-schema-ddl`: Simulation runtimes and the Stage 12 persistence library contain no schema DDL (CREATE, ALTER, DROP or TRUNCATE of tables, types, schemas or databases, or create_all); policyengine-api owns the schema. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_stage12_schema_authority.py`) -- `sim.stage12.persistence-typed-dml-only`: The Stage 12 persistence library and runtime use typed SQLAlchemy Core DML only: no raw SQL text, driver-level SQL or schema operations. (blocking: pytest `libs/policyengine-stage12-persistence/tests/test_source_boundaries.py`) -- `sim.stage12.production-lock-has-db-runtime`: The production lock file includes the direct Stage 12 database runtime dependencies. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) +- Checked: `sim.stage12.deployment-scope`, `sim.stage12.deployment-uses-api-schema`, `sim.stage12.no-production-user-associations`, `sim.stage12.no-schema-ddl`, `sim.stage12.persistence-typed-dml-only`, `sim.stage12.production-lock-has-db-runtime` ### Simulation persistence -- `sim.persistence.executor-independent`: The Modal executor and gateway never declare, lock or import policyengine-simulation-persistence; only the Simulation Entrypoint and the migration job install it. (blocking: pytest `libs/policyengine-simulation-persistence/tests/test_ownership_boundaries.py`) -- `sim.persistence.migrate-before-entrypoint`: Each deployment applies the simulation chain in its own job, the only one that receives SIMULATION_MIGRATION_DATABASE_URL, before the entrypoint candidate deploys; the entrypoint receives SIMULATION_DATABASE_URL as a pinned Secret Manager version; and the candidate smoke requires the database at the image's revision. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_simulation_database_deployment.py`) -- `sim.persistence.migration-cycle`: On a fresh PostgreSQL, the simulation chain upgrades to its head, passes alembic check, downgrades to base and upgrades again, run by a non-superuser migrator, and the runtime role can read the revision but cannot change the schema. (blocking: pytest `libs/policyengine-simulation-persistence/tests/test_migration_cycle.py`) -- `sim.persistence.no-modal-database-credentials`: No Modal secret, Modal deployment or Modal application receives SIMULATION_DATABASE_URL or SIMULATION_MIGRATION_DATABASE_URL; the Stage 12 comparison's own database secret is the only database credential Modal holds. (blocking: pytest `libs/policyengine-simulation-persistence/tests/test_ownership_boundaries.py`) -- `sim.persistence.one-migration-chain`: The simulation schema changes only through the one Alembic chain at migrations/ (configured by migrations/alembic.ini, versioned in simulation.alembic_version); no other Alembic configuration, environment or revision directory exists, and runtime source code contains no schema DDL. (blocking: pytest `libs/policyengine-simulation-persistence/tests/test_ownership_boundaries.py`) +- Checked: `sim.persistence.conditional-job-writes`, `sim.persistence.executor-independent`, `sim.persistence.internal-routes-google-oidc`, `sim.persistence.migrate-before-entrypoint`, `sim.persistence.migration-cycle`, `sim.persistence.no-modal-database-credentials`, `sim.persistence.one-migration-chain`, `sim.persistence.status-reads-never-advance-work` ### Deployment pipeline -- `sim.deploy.container-context-minimal`: Container builds exclude local environments and unrelated projects, and the entrypoint container builds for the requested target platform. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.docker-images-build`: The simulation entry and executor Docker images build. (blocking: CI job `docker-build`) -- `sim.deploy.gcloud-without-terraform`: Deployment uses the gcloud-based workflow, not Terraform. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) +- Checked: `sim.deploy.container-context-minimal`, `sim.deploy.docker-images-build`, `sim.deploy.gcloud-without-terraform`, `sim.deploy.modal-image-smoke`, `sim.deploy.numerical-sentinel-gate`, `sim.deploy.permission-enforcement-wired`, `sim.deploy.promotion-order-explicit`, `sim.deploy.release-identity-smoke`, `sim.deploy.revision-tags-valid`, `sim.deploy.routing-publication-journaled`, `sim.deploy.shell-scripts-valid`, `sim.deploy.smoke-tests-strict`, `sim.deploy.traffic-changes-exact`, `sim.deploy.workflow-shell-safety` - `sim.deploy.image-smoke-covers-every-image-lib`: The image smoke path filter includes every library mounted into a Modal image, so a change to any of them runs the smoke. (doc-only, not checked: follow it yourself) -- `sim.deploy.modal-image-smoke`: App entrypoints import inside the real Modal images before merge, whenever a library baked into an image changes. (blocking: CI job `image-smoke`) -- `sim.deploy.numerical-sentinel-gate`: An existing-stack deployment computes the numerical release sentinel from the exact bundle it installs, and the candidate and stable services must reproduce it; a failing sentinel check fails its step. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.permission-enforcement-wired`: Simulation Entrypoint permission enforcement defaults to log-only, rejects any value other than 0 or 1, and reaches the candidate revision and every deployed authentication run. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.promotion-order-explicit`: Full-stack promotion order is explicit, and the reusable deployment rejects unexpectedly skipped phases. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.release-identity-smoke`: Deploy gates pass the whole release identity to the entrypoint smoke, which requires every expectation, refuses malformed ones, and rejects a release whose routes, latest pointers, bundle, UK model or UK dataset differ from that identity. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.revision-tags-valid`: Cloud Run revision tags are valid on a workflow's first run and fit the existing service names. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.routing-publication-journaled`: Each routing publication is named by its run, attempt and environment and journaled before it is applied; an always-running finalizer resolves it from the journal with the release-gate outcome, a release reports success only after that finalizer completes, and routing jobs never tolerate a failed command. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.shell-scripts-valid`: Deployment shell scripts have valid shell syntax. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.smoke-tests-strict`: Candidate and stable smoke tests reject a revision mismatch without retrying, and fail if the previous revision does not converge. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.traffic-changes-exact`: Traffic changes target an exact, validated revision, are reversible, and refuse to run over an intervening promotion. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.deploy.workflow-shell-safety`: Cloud Run deployment never interpolates GitHub expressions into shell, escapes environment values, and pins secret versions. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) ### API contracts and integration -- `sim.contracts.clients-publish-deployed-commit`: Client publication checks out the commit that was actually deployed. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.contracts.full-integration-beta-only`: The complete integration suite runs against beta deployments only, never production. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_deployment_assets.py`) -- `sim.contracts.gateway-openapi-golden`: The gateway's OpenAPI spec matches the checked-in golden file; regenerate it deliberately when the API changes. (blocking: pytest `projects/policyengine-simulation-gateway/tests/test_openapi_golden.py`) -- `sim.contracts.local-integration`: Generated-client integration tests pass against the locally built services. (blocking: CI job `test-integration`) +- Checked: `sim.contracts.clients-publish-deployed-commit`, `sim.contracts.full-integration-beta-only`, `sim.contracts.gateway-openapi-golden`, `sim.contracts.local-integration` - `sim.contracts.regenerate-clients`: When public API schemas change, regenerate the clients and run the generated-client integration tests. (doc-only, not checked: follow it yourself) ### Observability @@ -147,32 +122,20 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. ### Code quality and tests -- `sim.quality.ci-cd-scripts-ruff-exemptions`: CI/CD scripts under .github may print their output (T201); every other ruff rule applies to them. (blocking: ruff per-file ignore of T201 for `.github/**`) -- `sim.quality.command-line-scripts-ruff-exemptions`: Command-line scripts under a scripts directory may print their output (T201), run other programs (S603, S607) and use today's date (DTZ011); every other ruff rule applies to them. (blocking: ruff per-file ignore of DTZ011, S603, S607, T201 for `**/scripts/**`) -- `sim.quality.deployed-tests-ruff-exemptions`: The deployed-service test suite (projects/*/authenticated_tests, pytest run against a live candidate) gets the same ruff exemptions as unit tests (sim.quality.tests-ruff-exemptions); every other ruff rule applies to it. (blocking: ruff per-file ignore of ANN, D1, DTZ001, DTZ007, INP001, PLC0415, PLR2004, S101, S102, S105, S106, S108, S301, S603, S607, SLF001, TID253 for `**/authenticated_tests/**`) +- Checked: `sim.quality.ci-cd-scripts-ruff-exemptions`, `sim.quality.command-line-scripts-ruff-exemptions`, `sim.quality.deployed-tests-ruff-exemptions`, `sim.quality.ruff-format-src`, `sim.quality.ruff-lint`, `sim.quality.runtime-evaluated-models-resolve`, `sim.quality.src-programs-ruff-exemptions`, `sim.quality.tests-pass`, `sim.quality.tests-ruff-exemptions` - `sim.quality.every-project-tested`: Every project and library with unit tests runs them in the PR test matrix. (doc-only, not checked: follow it yourself) -- `sim.quality.ruff-format-src`: Source code under projects/*/src and libs/*/src is formatted with ruff format. (blocking: ruff format) - `sim.quality.ruff-format-tests`: Test code is formatted with ruff format too; CI currently checks only src directories. (doc-only, not checked: follow it yourself) -- `sim.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in ruff.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (A, ANN, BLE, D, DTZ, E, E4, E7, E9, EM, ERA, EXE, F, FAST, FIX, FLY, FURB, G, I, INP, LOG, N, PIE, PLC, PYI, RET, RSE, S, SLF, T20, TC, TD, TID253, TRY, UP, W)) -- `sim.quality.runtime-evaluated-models-resolve`: Every Pydantic model is runtime-evaluated for ruff: one of its bases, read through the imports of the model's own module, is listed in ruff.toml's runtime-evaluated-base-classes or is a runtime-evaluated class in that same module; and every Pydantic model resolves its annotations once its package is imported. (blocking: pytest `projects/policyengine-simulation-entry/tests/test_runtime_evaluated_annotations.py`, `projects/policyengine-simulation-executor/tests/test_runtime_evaluated_annotations.py`) -- `sim.quality.src-programs-ruff-exemptions`: Programs under src (the Modal deployment utilities, and the OpenAPI and Stage 12 infrastructure commands) may print their output (T201), open URLs from their arguments (S310) and, for the Modal utilities run with python -m, live in a directory that isn't a package (INP001); every other ruff rule applies to them. (blocking: ruff per-file ignore of INP001, S310, T201 for `projects/*/src/{modal/utils/*,policyengine_simulation_entry/generate_openapi,policyengine_simulation_entry/stage12_infrastructure,policyengine_simulation_gateway/generate_openapi}.py`) -- `sim.quality.tests-pass`: Each project's unit tests pass in its own locked environment. (blocking: CI job `test`) -- `sim.quality.tests-ruff-exemptions`: Test files may use pytest's bare assert, need no docstrings, annotations or named constants for magic numbers, and may reach into private members, live outside packages and import inside a test (S101, D1, ANN, PLR2004, SLF001, INP001, PLC0415). They may also run repository scripts (S603, S607) and use fake credentials (S105, S106), temporary paths (S108), naive datetimes (DTZ001, DTZ007), pickle (S301) and exec (S102), and may import the lazily loaded packages at module level (TID253). Every other ruff rule applies to tests. (blocking: ruff per-file ignore of ANN, D1, DTZ001, DTZ007, INP001, PLC0415, PLR2004, S101, S102, S105, S106, S108, S301, S603, S607, SLF001, TID253 for `**/tests/**`) ### CI workflows and scripts -- `sim.ci.multiline-run-is-script`: A GitHub Actions run: block holds one line of code and no chained commands (&&, ||, ;); anything more lives in a tested script under .github/scripts/, called from a single-line step. (blocking: quality guard `ci-scripts`) -- `sim.ci.shell-scripts-pass-shellcheck`: Every shell script under .github/scripts/ passes shellcheck. (blocking: quality guard `ci-scripts`) -- `sim.ci.workflow-scripts-tested`: Every shell script a workflow calls is run by at least one test. (blocking: quality guard `ci-scripts`) +- Checked: `sim.ci.multiline-run-is-script`, `sim.ci.shell-scripts-pass-shellcheck`, `sim.ci.workflow-scripts-tested` ### Process and repository hygiene -- `sim.agents.adapters-point-to-agents-md`: CLAUDE.md and .github/copilot-instructions.md only point to AGENTS.md; skills and rules are listed in AGENTS.md alone. (blocking: quality guard `rules-registry`) +- Checked: `sim.agents.adapters-point-to-agents-md`, `sim.process.rules-registry`, `sim.process.stacked-prs-full-ci-at-bottom` - `sim.process.agents-never-bypass`: Agents never bypass branch rules, even with admin credentials: no gh pr merge --admin and no bypass option; a blocked merge is reported to a human. (doc-only, not checked: follow it yourself) - `sim.process.pr-tracking-issue`: Every PR has a tracking issue in this repository, and the first line of its description is Fixes #ISSUE_NUMBER. (doc-only, not checked: follow it yourself) - `sim.process.rule-machinery-code-owned`: Changes to the rule machinery (rules/, scripts/rules/, scripts/run_quality_guards.py and CODEOWNERS) need approval from a core developer other than the author. (doc-only, not checked: follow it yourself) -- `sim.process.rules-registry`: Every rule is registered in rules/registry.toml, is enforced where it claims to be, and its rendered views are current. (blocking: quality guard `rules-registry`) - `sim.process.same-repo-draft-prs`: PRs are opened as drafts from branches in PolicyEngine/policyengine-sim-api, never from forks. (doc-only, not checked: follow it yourself) -- `sim.process.stacked-prs-full-ci-at-bottom`: In a stack of pull requests, only the lowest unmerged one runs the full PR workflow and the Modal image smoke. Every pr.yml job is needed by the CI complete gate, which fails on the pull requests above it and otherwise passes only when every job passed, so a stack merges bottom-up and nothing merges untested. (blocking: pytest `projects/policyengine-simulation-executor/tests/test_ci_complete_gate.py`) diff --git a/docs/generated/rules.md b/docs/generated/rules.md index f1edaf91a..2a49381da 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -8,11 +8,11 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf | Tier | enforced | ratcheted | advisory | doc-only | Total | | --- | ---: | ---: | ---: | ---: | ---: | -| mechanical | 43 | 1 | 0 | 14 | 58 | +| mechanical | 46 | 1 | 0 | 14 | 61 | | metric | 0 | 0 | 0 | 0 | 0 | | judgment | 0 | 0 | 0 | 0 | 0 | -**Blocking:** 44 of 58 rules. +**Blocking:** 47 of 61 rules. ## Gaps @@ -54,8 +54,11 @@ Every rule that governs `policyengine-sim-api` is listed here with how it is enf - `sim.deploy.smoke-tests-strict` - `sim.deploy.traffic-changes-exact` - `sim.deploy.workflow-shell-safety` +- `sim.persistence.conditional-job-writes` +- `sim.persistence.internal-routes-google-oidc` - `sim.persistence.migrate-before-entrypoint` - `sim.persistence.migration-cycle` +- `sim.persistence.status-reads-never-advance-work` - `sim.process.rules-registry` - `sim.stage12.deployment-scope` - `sim.stage12.deployment-uses-api-schema` @@ -125,32 +128,47 @@ flowchart LR | Rule | Statement | Tier | Status | Blocking | Enforced by | CI | Proof | | --- | --- | --- | --- | --- | --- | --- | --- | +| `sim.persistence.conditional-job-writes` | Every simulation job transition and result write is one conditional update on the job's expected status (and submission an insert that does nothing on a conflicting input hash), so concurrent collectors, dispatchers or submitters change a job at most once; no lease is taken. | mechanical | enforced | yes | pytest `libs/policyengine-simulation-persistence/tests/test_job_store_concurrency.py` | `pr.yml:simulation-database` | **unproven** | | `sim.persistence.executor-independent` | The Modal executor and gateway never declare, lock or import policyengine-simulation-persistence; only the Simulation Entrypoint and the migration job install it. | mechanical | enforced | yes | pytest `libs/policyengine-simulation-persistence/tests/test_ownership_boundaries.py` | `pr.yml:test` | yes | +| `sim.persistence.internal-routes-google-oidc` | The simulation job collection and sweep routes accept only Google-signed OIDC tokens whose audience is the service URL and whose email is the configured invoker service account. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_google_oidc.py`, `projects/policyengine-simulation-entry/tests/test_simulation_job_routes.py` | `pr.yml:test` | **unproven** | | `sim.persistence.migrate-before-entrypoint` | Each deployment applies the simulation chain in its own job, the only one that receives SIMULATION_MIGRATION_DATABASE_URL, before the entrypoint candidate deploys; the entrypoint receives SIMULATION_DATABASE_URL as a pinned Secret Manager version; and the candidate smoke requires the database at the image's revision. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_simulation_database_deployment.py` | `pr.yml:test` | **unproven** | | `sim.persistence.migration-cycle` | On a fresh PostgreSQL, the simulation chain upgrades to its head, passes alembic check, downgrades to base and upgrades again, run by a non-superuser migrator, and the runtime role can read the revision but cannot change the schema. | mechanical | enforced | yes | pytest `libs/policyengine-simulation-persistence/tests/test_migration_cycle.py` | `pr.yml:simulation-database` | **unproven** | | `sim.persistence.no-modal-database-credentials` | No Modal secret, Modal deployment or Modal application receives SIMULATION_DATABASE_URL or SIMULATION_MIGRATION_DATABASE_URL; the Stage 12 comparison's own database secret is the only database credential Modal holds. | mechanical | enforced | yes | pytest `libs/policyengine-simulation-persistence/tests/test_ownership_boundaries.py` | `pr.yml:test` | yes | | `sim.persistence.one-migration-chain` | The simulation schema changes only through the one Alembic chain at migrations/ (configured by migrations/alembic.ini, versioned in simulation.alembic_version); no other Alembic configuration, environment or revision directory exists, and runtime source code contains no schema DDL. | mechanical | enforced | yes | pytest `libs/policyengine-simulation-persistence/tests/test_ownership_boundaries.py` | `pr.yml:test` | yes | +| `sim.persistence.status-reads-never-advance-work` | Simulation job status reads use only the Simulation Entrypoint's database: they never inspect, spawn or advance a Modal call; results are collected by Cloud Tasks and the Cloud Scheduler sweep. | mechanical | enforced | yes | pytest `projects/policyengine-simulation-entry/tests/test_simulation_job_routes.py`, `projects/policyengine-simulation-entry/tests/test_simulation_jobs.py` | `pr.yml:test` | **unproven** | ```mermaid flowchart LR + r_sim_persistence_conditional_job_writes("sim.persistence.conditional-job-writes (enforced)") + e_pytest_libs_policyengine_simulation_persistence_tests_test_job_store_concurrency_py["pytest libs/policyengine-simulation-persistence/tests/test_job_store_concurrency.py"] + j_CI_pr_yml_simulation_database["CI pr.yml:simulation-database"] r_sim_persistence_executor_independent("sim.persistence.executor-independent (enforced)") e_pytest_libs_policyengine_simulation_persistence_tests_test_ownership_boundaries_py["pytest libs/policyengine-simulation-persistence/tests/test_ownership_boundaries.py"] j_CI_pr_yml_test["CI pr.yml:test"] + r_sim_persistence_internal_routes_google_oidc("sim.persistence.internal-routes-google-oidc (enforced)") + e_pytest_projects_policyengine_simulation_entry_tests_test_google_oidc_py__projects_policyengine_simulation_entry_tests_test_simulation_job_routes_py["pytest projects/policyengine-simulation-entry/tests/test_google_oidc.py, projects/policyengine-simulation-entry/tests/test_simulation_job_routes.py"] r_sim_persistence_migrate_before_entrypoint("sim.persistence.migrate-before-entrypoint (enforced)") e_pytest_projects_policyengine_simulation_entry_tests_test_simulation_database_deployment_py["pytest projects/policyengine-simulation-entry/tests/test_simulation_database_deployment.py"] r_sim_persistence_migration_cycle("sim.persistence.migration-cycle (enforced)") e_pytest_libs_policyengine_simulation_persistence_tests_test_migration_cycle_py["pytest libs/policyengine-simulation-persistence/tests/test_migration_cycle.py"] - j_CI_pr_yml_simulation_database["CI pr.yml:simulation-database"] r_sim_persistence_no_modal_database_credentials("sim.persistence.no-modal-database-credentials (enforced)") r_sim_persistence_one_migration_chain("sim.persistence.one-migration-chain (enforced)") + r_sim_persistence_status_reads_never_advance_work("sim.persistence.status-reads-never-advance-work (enforced)") + e_pytest_projects_policyengine_simulation_entry_tests_test_simulation_job_routes_py__projects_policyengine_simulation_entry_tests_test_simulation_jobs_py["pytest projects/policyengine-simulation-entry/tests/test_simulation_job_routes.py, projects/policyengine-simulation-entry/tests/test_simulation_jobs.py"] + r_sim_persistence_conditional_job_writes --> e_pytest_libs_policyengine_simulation_persistence_tests_test_job_store_concurrency_py + e_pytest_libs_policyengine_simulation_persistence_tests_test_job_store_concurrency_py --> j_CI_pr_yml_simulation_database r_sim_persistence_executor_independent --> e_pytest_libs_policyengine_simulation_persistence_tests_test_ownership_boundaries_py e_pytest_libs_policyengine_simulation_persistence_tests_test_ownership_boundaries_py --> j_CI_pr_yml_test + r_sim_persistence_internal_routes_google_oidc --> e_pytest_projects_policyengine_simulation_entry_tests_test_google_oidc_py__projects_policyengine_simulation_entry_tests_test_simulation_job_routes_py + e_pytest_projects_policyengine_simulation_entry_tests_test_google_oidc_py__projects_policyengine_simulation_entry_tests_test_simulation_job_routes_py --> j_CI_pr_yml_test r_sim_persistence_migrate_before_entrypoint --> e_pytest_projects_policyengine_simulation_entry_tests_test_simulation_database_deployment_py e_pytest_projects_policyengine_simulation_entry_tests_test_simulation_database_deployment_py --> j_CI_pr_yml_test r_sim_persistence_migration_cycle --> e_pytest_libs_policyengine_simulation_persistence_tests_test_migration_cycle_py e_pytest_libs_policyengine_simulation_persistence_tests_test_migration_cycle_py --> j_CI_pr_yml_simulation_database r_sim_persistence_no_modal_database_credentials --> e_pytest_libs_policyengine_simulation_persistence_tests_test_ownership_boundaries_py r_sim_persistence_one_migration_chain --> e_pytest_libs_policyengine_simulation_persistence_tests_test_ownership_boundaries_py + r_sim_persistence_status_reads_never_advance_work --> e_pytest_projects_policyengine_simulation_entry_tests_test_simulation_job_routes_py__projects_policyengine_simulation_entry_tests_test_simulation_jobs_py + e_pytest_projects_policyengine_simulation_entry_tests_test_simulation_job_routes_py__projects_policyengine_simulation_entry_tests_test_simulation_jobs_py --> j_CI_pr_yml_test ``` ## Deployment pipeline diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index c9e927b8f..78daac6d0 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -30,6 +30,38 @@ The temporary Stage 12 comparison tables are not part of this schema: they are owned and migrated by the Public API, and `libs/policyengine-stage12-persistence` writes them with DML only. +## The job store + +`simulation.jobs` has one row per calculation the Simulation Entrypoint sends +to Modal: the submitter, a SHA-256 of the job kind, policyengine.py bundle +version and request (`input_sha256`), the Modal application, function and +call ID, the status, timings, the number of collection checks and any error. +`simulation.job_results` holds the validated result of each succeeded job that +is not a report. + +A job moves `pending` -> `dispatching` -> `running` -> `succeeded` or +`failed`. Every transition in `jobs.py` is one conditional +`UPDATE ... WHERE status = ` (`sim.persistence.conditional-job-writes`): + +- **Submission** inserts with `ON CONFLICT DO NOTHING` on a unique index over + `input_sha256` that excludes failed jobs, so identical requests share one + job and a failed job never blocks a new attempt. +- **Dispatch** happens only after the conditional `pending` -> `dispatching` + update succeeds, so one Modal call starts per job. +- **Collection** writes a result (in the same transaction as the + `running` -> `succeeded` update) or a failure only while the job is + `running`. When two collectors race, PostgreSQL makes the second wait for the + first's row lock, re-check the condition and change nothing. No lease is + taken. `tests/test_job_store_concurrency.py` forces that interleaving and + races many threads. + +The Simulation Entrypoint collects results in the background: Cloud Tasks +calls its collection route, which checks the Modal call without waiting and +either stores the outcome or schedules the next check with backoff; a Cloud +Scheduler sweep re-enqueues running jobs nobody has checked for five minutes +and fails jobs whose dispatch stalled for ten. Status reads use only these +tables (`sim.persistence.status-reads-never-advance-work`). + ## Database roles | Role | Used by | Privileges | diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/jobs.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/jobs.py new file mode 100644 index 000000000..0712393f3 --- /dev/null +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/jobs.py @@ -0,0 +1,561 @@ +"""The simulation job store: one ``simulation.jobs`` row per calculation. + +The Simulation Entrypoint writes every row; Modal never does. A job moves +``pending`` -> ``dispatching`` -> ``running`` -> ``succeeded`` or ``failed``. +Every transition is one conditional ``UPDATE ... WHERE status = ``, +so when two entrypoint instances act on the same job at once, PostgreSQL makes +the second wait for the first's row lock and then find nothing to change. No +lease or lock table is needed: + +- submission is idempotent: identical inputs share one unfinished or + succeeded job (``INSERT ... ON CONFLICT DO NOTHING`` on the input hash); +- only the instance that moves a job from ``pending`` to ``dispatching`` + starts its Modal call; +- only one collector moves a ``running`` job to ``succeeded`` (writing its + result in the same transaction) or ``failed``. +""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, fields +from datetime import UTC, datetime +from enum import StrEnum +from typing import TYPE_CHECKING, Any +from uuid import UUID, uuid4 + +from sqlalchemy import ( + CheckConstraint, + Column, + DateTime, + ForeignKey, + Index, + Integer, + String, + Table, + Text, + Uuid, + func, + insert, + or_, + select, + text, + update, +) +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy.dialects.postgresql import insert as postgres_insert + +from policyengine_simulation_persistence.schema import simulation_metadata + +if TYPE_CHECKING: + from collections.abc import Callable + + from sqlalchemy import Connection, Engine, RowMapping + +MAX_ERROR_MESSAGE_LENGTH = 2000 +DISPATCH_ABANDONED = "dispatch_abandoned" + + +class JobStatus(StrEnum): + """Where a job is in its lifecycle.""" + + PENDING = "pending" + DISPATCHING = "dispatching" + RUNNING = "running" + SUCCEEDED = "succeeded" + FAILED = "failed" + + +UNFINISHED_STATUSES = (JobStatus.PENDING, JobStatus.DISPATCHING, JobStatus.RUNNING) +_STATUS_VALUES = ", ".join(f"'{status.value}'" for status in JobStatus) +_SHA256_PATTERN = "^[0-9a-f]{64}$" + +jobs = Table( + "jobs", + simulation_metadata, + Column("id", Uuid(as_uuid=True), primary_key=True), + Column("kind", Text, nullable=False), + Column("status", Text, nullable=False), + Column("submitted_by", Text, nullable=False), + Column("input_sha256", String(64), nullable=False), + Column("policyengine_version", Text, nullable=False), + Column("modal_application", Text, nullable=False), + Column("modal_function", Text, nullable=False), + Column("modal_call_id", Text, nullable=True, unique=True), + Column("observability_id", Text, nullable=True), + Column("error_code", Text, nullable=True), + Column("error_message", Text, nullable=True), + Column("collection_attempts", Integer, nullable=False, server_default=text("0")), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("updated_at", DateTime(timezone=True), nullable=False), + Column("dispatched_at", DateTime(timezone=True), nullable=True), + Column("last_checked_at", DateTime(timezone=True), nullable=True), + Column("completed_at", DateTime(timezone=True), nullable=True), + CheckConstraint(f"status IN ({_STATUS_VALUES})", name="status"), + CheckConstraint(f"input_sha256 ~ '{_SHA256_PATTERN}'", name="input_sha256"), + CheckConstraint("collection_attempts >= 0", name="collection_attempts"), +) +# Identical inputs share one job unless it failed; a failed job never blocks a +# new submission of the same input. +Index( + "uq_jobs_input_sha256_unfailed", + jobs.c.input_sha256, + unique=True, + postgresql_where=jobs.c.status != JobStatus.FAILED.value, +) +# The scheduled sweep reads only unfinished jobs. +Index( + "ix_jobs_unfinished_status", + jobs.c.status, + postgresql_where=jobs.c.status.in_( + [status.value for status in UNFINISHED_STATUSES] + ), +) + +job_results = Table( + "job_results", + simulation_metadata, + Column( + "job_id", + Uuid(as_uuid=True), + ForeignKey(jobs.c.id), + primary_key=True, + ), + Column("result", JSONB, nullable=False), + Column("result_schema_version", Integer, nullable=False), + Column("result_sha256", String(64), nullable=False), + Column("created_at", DateTime(timezone=True), nullable=False), + CheckConstraint(f"result_sha256 ~ '{_SHA256_PATTERN}'", name="result_sha256"), + CheckConstraint("result_schema_version >= 1", name="result_schema_version"), +) + + +def _canonical_json(value: object) -> bytes: + return json.dumps( + value, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ).encode("utf-8") + + +def input_sha256(*, kind: str, policyengine_version: str, request: object) -> str: + """Return the idempotency hash of one calculation's inputs. + + It covers the job kind, the policyengine.py bundle version and the request + as canonical JSON, so the same request under another bundle is another job. + """ + return hashlib.sha256( + _canonical_json( + { + "kind": kind, + "policyengine_version": policyengine_version, + "request": request, + } + ) + ).hexdigest() + + +def result_sha256(result: object) -> str: + """Return the SHA-256 of a result as canonical JSON.""" + return hashlib.sha256(_canonical_json(result)).hexdigest() + + +def utc_now() -> datetime: + """Return the current time in UTC.""" + return datetime.now(UTC) + + +@dataclass(frozen=True) +class NewJob: + """What a submission records before anything is dispatched.""" + + kind: str + submitted_by: str + input_sha256: str + policyengine_version: str + modal_application: str + modal_function: str + observability_id: str | None = None + + +@dataclass(frozen=True) +class JobRecord: + """One ``simulation.jobs`` row.""" + + id: UUID + kind: str + status: JobStatus + submitted_by: str + input_sha256: str + policyengine_version: str + modal_application: str + modal_function: str + modal_call_id: str | None + observability_id: str | None + error_code: str | None + error_message: str | None + collection_attempts: int + created_at: datetime + updated_at: datetime + dispatched_at: datetime | None + last_checked_at: datetime | None + completed_at: datetime | None + + @property + def finished(self) -> bool: + """Whether the job has reached ``succeeded`` or ``failed``.""" + return self.status not in UNFINISHED_STATUSES + + @classmethod + def from_row(cls, row: RowMapping) -> JobRecord: + """Build the record from a selected ``jobs`` row.""" + values = {field.name: row[field.name] for field in fields(cls)} + values["status"] = JobStatus(values["status"]) + return cls(**values) + + +@dataclass(frozen=True) +class JobResultRecord: + """One ``simulation.job_results`` row.""" + + job_id: UUID + result: dict[str, Any] + result_schema_version: int + result_sha256: str + created_at: datetime + + +@dataclass(frozen=True) +class Submission: + """The job a submission resolved to, and whether this call created it.""" + + job: JobRecord + created: bool + + +def _bounded(message: str) -> str: + return message[:MAX_ERROR_MESSAGE_LENGTH] + + +class JobStore: + """Typed SQLAlchemy Core reads and conditional writes for simulation jobs.""" + + def __init__( + self, + engine: Engine, + *, + clock: Callable[[], datetime] = utc_now, + ) -> None: + """Use ``engine`` (the runtime role) and ``clock`` for every timestamp.""" + self.engine = engine + self._clock = clock + + # --- Submission and dispatch --------------------------------------------- + + def submit(self, new_job: NewJob) -> Submission: + """Create a pending job, or return the existing one for the same input. + + A failed job does not count: the same input then creates a new job. + """ + now = self._clock() + values = { + "id": uuid4(), + "status": JobStatus.PENDING.value, + "created_at": now, + "updated_at": now, + **{field.name: getattr(new_job, field.name) for field in fields(new_job)}, + } + statement = ( + postgres_insert(jobs) + .values(**values) + .on_conflict_do_nothing( + index_elements=[jobs.c.input_sha256], + index_where=jobs.c.status != JobStatus.FAILED.value, + ) + .returning(*jobs.c) + ) + existing_statement = select(jobs).where( + jobs.c.input_sha256 == new_job.input_sha256, + jobs.c.status != JobStatus.FAILED.value, + ) + # The conflicting job can fail between the insert and the read; the + # next attempt then creates a new job. + for _attempt in range(3): + with self.engine.begin() as connection: + created = connection.execute(statement).mappings().first() + if created is not None: + return Submission(job=JobRecord.from_row(created), created=True) + existing = connection.execute(existing_statement).mappings().first() + if existing is not None: + return Submission(job=JobRecord.from_row(existing), created=False) + msg = "job submission did not settle on one job" + raise RuntimeError(msg) + + def _claim_dispatch(self, connection: Connection, job_id: UUID) -> bool: + claimed = connection.execute( + update(jobs) + .where(jobs.c.id == job_id, jobs.c.status == JobStatus.PENDING.value) + .values(status=JobStatus.DISPATCHING.value, updated_at=self._clock()) + .returning(jobs.c.id) + ).first() + return claimed is not None + + def claim_dispatch(self, job_id: UUID) -> bool: + """Move a pending job to dispatching; only the winner starts its call.""" + with self.engine.begin() as connection: + return self._claim_dispatch(connection, job_id) + + def record_dispatched(self, job_id: UUID, *, modal_call_id: str) -> bool: + """Record the acknowledged Modal call and move the job to running.""" + now = self._clock() + with self.engine.begin() as connection: + return ( + connection.execute( + update(jobs) + .where( + jobs.c.id == job_id, + jobs.c.status == JobStatus.DISPATCHING.value, + ) + .values( + status=JobStatus.RUNNING.value, + modal_call_id=modal_call_id, + dispatched_at=now, + updated_at=now, + ) + .returning(jobs.c.id) + ).first() + is not None + ) + + def record_dispatch_failure( + self, + job_id: UUID, + *, + error_code: str, + error_message: str, + ) -> bool: + """Fail a job whose Modal call was refused or never acknowledged.""" + with self.engine.begin() as connection: + return self._finish_failed( + connection, + job_id, + expected=JobStatus.DISPATCHING, + error_code=error_code, + error_message=error_message, + ) + + # --- Reads ----------------------------------------------------------------- + + def get(self, job_id: UUID) -> JobRecord | None: + """Return the job with this ID, or None.""" + with self.engine.connect() as connection: + row = ( + connection.execute(select(jobs).where(jobs.c.id == job_id)) + .mappings() + .first() + ) + return None if row is None else JobRecord.from_row(row) + + def get_by_modal_call_id(self, modal_call_id: str) -> JobRecord | None: + """Return the job whose Modal call has this ID, or None.""" + with self.engine.connect() as connection: + row = ( + connection.execute( + select(jobs).where(jobs.c.modal_call_id == modal_call_id) + ) + .mappings() + .first() + ) + return None if row is None else JobRecord.from_row(row) + + def get_result(self, job_id: UUID) -> JobResultRecord | None: + """Return the stored result of a succeeded job, or None.""" + with self.engine.connect() as connection: + row = ( + connection.execute( + select(job_results).where(job_results.c.job_id == job_id) + ) + .mappings() + .first() + ) + return None if row is None else JobResultRecord(**row) + + # --- Collection ------------------------------------------------------------ + + def record_check(self, job_id: UUID) -> int | None: + """Count one check of a still-running call; None if it is not running.""" + now = self._clock() + with self.engine.begin() as connection: + return connection.execute( + update(jobs) + .where(jobs.c.id == job_id, jobs.c.status == JobStatus.RUNNING.value) + .values( + collection_attempts=jobs.c.collection_attempts + 1, + last_checked_at=now, + updated_at=now, + ) + .returning(jobs.c.collection_attempts) + ).scalar() + + def _complete( + self, + connection: Connection, + job_id: UUID, + *, + result: dict[str, Any], + result_schema_version: int, + ) -> bool: + now = self._clock() + finished = connection.execute( + update(jobs) + .where(jobs.c.id == job_id, jobs.c.status == JobStatus.RUNNING.value) + .values( + status=JobStatus.SUCCEEDED.value, + completed_at=now, + last_checked_at=now, + updated_at=now, + ) + .returning(jobs.c.id) + ).first() + if finished is None: + return False + connection.execute( + insert(job_results).values( + job_id=job_id, + result=result, + result_schema_version=result_schema_version, + result_sha256=result_sha256(result), + created_at=now, + ) + ) + return True + + def complete( + self, + job_id: UUID, + *, + result: dict[str, Any], + result_schema_version: int, + ) -> bool: + """Store a running job's validated result and mark it succeeded. + + Return False, writing nothing, when the job is no longer running + because another collector already finished it. + """ + with self.engine.begin() as connection: + return self._complete( + connection, + job_id, + result=result, + result_schema_version=result_schema_version, + ) + + def _finish_failed( + self, + connection: Connection, + job_id: UUID, + *, + expected: JobStatus, + error_code: str, + error_message: str, + ) -> bool: + now = self._clock() + return ( + connection.execute( + update(jobs) + .where(jobs.c.id == job_id, jobs.c.status == expected.value) + .values( + status=JobStatus.FAILED.value, + error_code=error_code, + error_message=_bounded(error_message), + completed_at=now, + updated_at=now, + ) + .returning(jobs.c.id) + ).first() + is not None + ) + + def _fail( + self, + connection: Connection, + job_id: UUID, + *, + error_code: str, + error_message: str, + ) -> bool: + return self._finish_failed( + connection, + job_id, + expected=JobStatus.RUNNING, + error_code=error_code, + error_message=error_message, + ) + + def fail(self, job_id: UUID, *, error_code: str, error_message: str) -> bool: + """Mark a running job failed; False if another collector finished it.""" + with self.engine.begin() as connection: + return self._fail( + connection, + job_id, + error_code=error_code, + error_message=error_message, + ) + + # --- The scheduled sweep --------------------------------------------------- + + def jobs_to_collect(self, *, checked_before: datetime, limit: int) -> list[UUID]: + """Return running jobs nobody has checked since ``checked_before``.""" + last_activity = func.coalesce( + jobs.c.last_checked_at, + jobs.c.dispatched_at, + jobs.c.updated_at, + ) + with self.engine.connect() as connection: + return list( + connection.execute( + select(jobs.c.id) + .where( + jobs.c.status == JobStatus.RUNNING.value, + last_activity < checked_before, + ) + .order_by(last_activity) + .limit(limit) + ).scalars() + ) + + def abandon_stalled_dispatches(self, *, updated_before: datetime) -> list[UUID]: + """Fail jobs stuck before their Modal call was recorded. + + A job stays ``pending`` or ``dispatching`` only if the instance + submitting it stopped before recording the call. Failing it lets the + same input be submitted again. + """ + now = self._clock() + with self.engine.begin() as connection: + return list( + connection.execute( + update(jobs) + .where( + or_( + jobs.c.status == JobStatus.PENDING.value, + jobs.c.status == JobStatus.DISPATCHING.value, + ), + jobs.c.updated_at < updated_before, + ) + .values( + status=JobStatus.FAILED.value, + error_code=DISPATCH_ABANDONED, + error_message=( + "The submitting instance stopped before the Modal " + "call was recorded." + ), + completed_at=now, + updated_at=now, + ) + .returning(jobs.c.id) + ).scalars() + ) diff --git a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py index 8cf367b96..0a63de00b 100644 --- a/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py +++ b/libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/schema.py @@ -17,7 +17,7 @@ # The newest revision in migrations/versions. Each migration PR updates it, and # tests/test_schema.py fails when it differs from the chain's head. The # readiness check reports it next to the database's revision. -HEAD_REVISION = "0001" +HEAD_REVISION = "0002" NAMING_CONVENTION = { "ix": "ix_%(table_name)s_%(column_0_N_name)s", diff --git a/libs/policyengine-simulation-persistence/tests/conftest.py b/libs/policyengine-simulation-persistence/tests/conftest.py index c31819dba..ca07fe629 100644 --- a/libs/policyengine-simulation-persistence/tests/conftest.py +++ b/libs/policyengine-simulation-persistence/tests/conftest.py @@ -15,10 +15,12 @@ from uuid import uuid4 import pytest +from alembic import command from migration_support import ( MIGRATION_DATABASE_URL_VARIABLE, MIGRATOR_ROLE, TEST_ROLE_PASSWORD, + alembic_config, ) from policyengine_simulation_persistence.schema import RUNTIME_ROLE from sqlalchemy import create_engine, text @@ -112,3 +114,10 @@ def disposable_database( text(f'DROP DATABASE IF EXISTS "{database_name}" WITH (FORCE)') ) server_engine.dispose() + + +@pytest.fixture +def migrated_database(disposable_database: DisposableDatabase) -> DisposableDatabase: + """Return a disposable database with the whole chain applied.""" + command.upgrade(alembic_config(), "head") + return disposable_database diff --git a/libs/policyengine-simulation-persistence/tests/job_store_support.py b/libs/policyengine-simulation-persistence/tests/job_store_support.py new file mode 100644 index 000000000..70d1943c6 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/job_store_support.py @@ -0,0 +1,58 @@ +"""Builders shared by the job store tests.""" + +from __future__ import annotations + +from datetime import UTC, datetime, timedelta + +from policyengine_simulation_persistence import create_simulation_engine +from policyengine_simulation_persistence.jobs import JobStore, NewJob, input_sha256 + +START = datetime(2026, 10, 20, 12, tzinfo=UTC) + + +class SteppingClock: + """A clock the test moves forward explicitly.""" + + def __init__(self, start: datetime = START) -> None: + self.now = start + + def __call__(self) -> datetime: + return self.now + + def advance(self, **elapsed: float) -> None: + self.now += timedelta(**elapsed) + + +def new_job(request: dict | None = None, **overrides) -> NewJob: + """Return a household calculation job for ``request``.""" + values = { + "kind": "household_calculation", + "submitted_by": "api-v1-client@clients", + "input_sha256": input_sha256( + kind="household_calculation", + policyengine_version="6.7.5", + request=request or {"household": {"people": {"you": {"age": 40}}}}, + ), + "policyengine_version": "6.7.5", + "modal_application": "policyengine-simulation-v2-py6-7-5", + "modal_function": "calculate_household", + "observability_id": "00000000-0000-4000-8000-000000000001", + } + values.update(overrides) + return NewJob(**values) + + +def runtime_store(database, clock=None) -> JobStore: + """Return a store that connects as the runtime role, as the entrypoint does.""" + return JobStore( + create_simulation_engine(database.runtime_url), + clock=clock or SteppingClock(), + ) + + +def running_job(store: JobStore, request: dict | None = None): + """Submit a job and take it through dispatch, as the entrypoint would.""" + job = store.submit(new_job(request)).job + assert store.claim_dispatch(job.id) + assert store.record_dispatched(job.id, modal_call_id=f"fc-{job.id.hex}") + return store.get(job.id) diff --git a/libs/policyengine-simulation-persistence/tests/test_job_store.py b/libs/policyengine-simulation-persistence/tests/test_job_store.py new file mode 100644 index 000000000..b4c19b483 --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_job_store.py @@ -0,0 +1,256 @@ +"""The simulation.jobs store: idempotent submission and conditional transitions.""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from job_store_support import ( + START, + SteppingClock, + new_job, + running_job, + runtime_store, +) +from policyengine_simulation_persistence.jobs import ( + JobStatus, + input_sha256, + result_sha256, +) + +RESULT = {"net_income": 51234.5, "programs": {"snap": 0}} + + +# --- Hashing, without a database ------------------------------------------- + + +def test_the_input_hash_ignores_key_order_but_covers_kind_and_bundle() -> None: + first = input_sha256( + kind="household_calculation", + policyengine_version="6.7.5", + request={"b": 1, "a": {"y": 2, "x": 3}}, + ) + + assert first == input_sha256( + kind="household_calculation", + policyengine_version="6.7.5", + request={"a": {"x": 3, "y": 2}, "b": 1}, + ) + assert first != input_sha256( + kind="household_calculation", + policyengine_version="6.7.6", + request={"b": 1, "a": {"y": 2, "x": 3}}, + ) + assert first != input_sha256( + kind="economy_report", + policyengine_version="6.7.5", + request={"b": 1, "a": {"y": 2, "x": 3}}, + ) + assert len(first) == 64 + + +def test_the_result_hash_is_canonical() -> None: + assert result_sha256({"b": 1, "a": 2}) == result_sha256({"a": 2, "b": 1}) + + +# --- Submission and dispatch ------------------------------------------------ + + +@pytest.mark.postgres +def test_identical_submissions_share_one_job(migrated_database) -> None: + store = runtime_store(migrated_database) + + first = store.submit(new_job()) + second = store.submit(new_job(submitted_by="another-client@clients")) + + assert first.created is True + assert second.created is False + assert second.job.id == first.job.id + assert second.job.submitted_by == "api-v1-client@clients" + assert first.job.status is JobStatus.PENDING + assert first.job.created_at == START + + +@pytest.mark.postgres +def test_a_failed_job_lets_the_same_input_be_submitted_again(migrated_database) -> None: + store = runtime_store(migrated_database) + failed = running_job(store) + assert store.fail(failed.id, error_code="worker_error", error_message="boom") + + retried = store.submit(new_job()) + + assert retried.created is True + assert retried.job.id != failed.id + assert store.get(failed.id).status is JobStatus.FAILED + + +@pytest.mark.postgres +def test_dispatch_is_claimed_once_and_records_the_modal_call(migrated_database) -> None: + clock = SteppingClock() + store = runtime_store(migrated_database, clock) + job = store.submit(new_job()).job + + assert store.claim_dispatch(job.id) is True + assert store.claim_dispatch(job.id) is False + clock.advance(seconds=2) + assert store.record_dispatched(job.id, modal_call_id="fc-01ABC") is True + assert store.record_dispatched(job.id, modal_call_id="fc-other") is False + + dispatched = store.get(job.id) + assert dispatched.status is JobStatus.RUNNING + assert dispatched.modal_call_id == "fc-01ABC" + assert dispatched.dispatched_at == clock.now + assert store.get_by_modal_call_id("fc-01ABC").id == job.id + assert store.get_by_modal_call_id("fc-unknown") is None + + +@pytest.mark.postgres +def test_a_rejected_dispatch_fails_the_job(migrated_database) -> None: + store = runtime_store(migrated_database) + job = store.submit(new_job()).job + assert store.claim_dispatch(job.id) + + assert store.record_dispatch_failure( + job.id, error_code="dispatch_failed", error_message="Modal refused" + ) + + failed = store.get(job.id) + assert failed.status is JobStatus.FAILED + assert failed.error_code == "dispatch_failed" + assert failed.completed_at is not None + + +@pytest.mark.postgres +def test_an_unknown_job_reads_as_none(migrated_database) -> None: + store = runtime_store(migrated_database) + + assert store.get(uuid4()) is None + assert store.get_result(uuid4()) is None + assert store.claim_dispatch(uuid4()) is False + + +# --- Collection --------------------------------------------------------------- + + +@pytest.mark.postgres +def test_a_result_is_written_once_with_its_hash(migrated_database) -> None: + clock = SteppingClock() + store = runtime_store(migrated_database, clock) + job = running_job(store) + clock.advance(minutes=3) + + assert store.complete(job.id, result=RESULT, result_schema_version=1) is True + assert ( + store.complete(job.id, result={"other": True}, result_schema_version=1) is False + ) + + finished = store.get(job.id) + stored = store.get_result(job.id) + assert finished.status is JobStatus.SUCCEEDED + assert finished.completed_at == clock.now + assert stored.result == RESULT + assert stored.result_schema_version == 1 + assert stored.result_sha256 == result_sha256(RESULT) + + +@pytest.mark.postgres +def test_only_a_running_job_can_finish(migrated_database) -> None: + store = runtime_store(migrated_database) + pending = store.submit(new_job()).job + + assert store.complete(pending.id, result=RESULT, result_schema_version=1) is False + assert store.fail(pending.id, error_code="x", error_message="y") is False + assert store.get(pending.id).status is JobStatus.PENDING + assert store.get_result(pending.id) is None + + +@pytest.mark.postgres +def test_a_failure_cannot_overwrite_a_result(migrated_database) -> None: + store = runtime_store(migrated_database) + job = running_job(store) + assert store.complete(job.id, result=RESULT, result_schema_version=1) + + assert store.fail(job.id, error_code="worker_error", error_message="late") is False + + assert store.get(job.id).status is JobStatus.SUCCEEDED + assert store.get(job.id).error_code is None + + +@pytest.mark.postgres +def test_long_error_messages_are_truncated(migrated_database) -> None: + store = runtime_store(migrated_database) + job = running_job(store) + + assert store.fail(job.id, error_code="worker_error", error_message="x" * 10_000) + + assert len(store.get(job.id).error_message) == 2000 + + +@pytest.mark.postgres +def test_each_check_of_a_running_job_is_counted(migrated_database) -> None: + clock = SteppingClock() + store = runtime_store(migrated_database, clock) + job = running_job(store) + + clock.advance(seconds=5) + assert store.record_check(job.id) == 1 + assert store.record_check(job.id) == 2 + assert store.get(job.id).last_checked_at == clock.now + assert store.complete(job.id, result=RESULT, result_schema_version=1) + assert store.record_check(job.id) is None + + +# --- The sweep ------------------------------------------------------------------ + + +@pytest.mark.postgres +def test_the_sweep_selects_running_jobs_not_checked_recently(migrated_database) -> None: + clock = SteppingClock() + store = runtime_store(migrated_database, clock) + never_checked = running_job(store, {"case": 1}) + checked_long_ago = running_job(store, {"case": 2}) + store.record_check(checked_long_ago.id) + clock.advance(minutes=10) + checked_recently = running_job(store, {"case": 3}) + store.record_check(checked_recently.id) + finished = running_job(store, {"case": 4}) + store.complete(finished.id, result=RESULT, result_schema_version=1) + pending = store.submit(new_job({"case": 5})).job + + selected = store.jobs_to_collect( + checked_before=clock.now - _minutes(5), + limit=10, + ) + + assert set(selected) == {never_checked.id, checked_long_ago.id} + assert pending.id not in selected + assert len(store.jobs_to_collect(checked_before=clock.now, limit=1)) == 1 + + +@pytest.mark.postgres +def test_the_sweep_fails_dispatches_that_stalled(migrated_database) -> None: + clock = SteppingClock() + store = runtime_store(migrated_database, clock) + stalled_pending = store.submit(new_job({"case": "pending"})).job + stalled_dispatch = store.submit(new_job({"case": "dispatching"})).job + store.claim_dispatch(stalled_dispatch.id) + running = running_job(store, {"case": "running"}) + clock.advance(minutes=15) + recent = store.submit(new_job({"case": "recent"})).job + + abandoned = store.abandon_stalled_dispatches( + updated_before=clock.now - _minutes(10) + ) + + assert set(abandoned) == {stalled_pending.id, stalled_dispatch.id} + for job_id in abandoned: + assert store.get(job_id).status is JobStatus.FAILED + assert store.get(job_id).error_code == "dispatch_abandoned" + assert store.get(recent.id).status is JobStatus.PENDING + assert store.get(running.id).status is JobStatus.RUNNING + + +def _minutes(count): + from datetime import timedelta + + return timedelta(minutes=count) diff --git a/libs/policyengine-simulation-persistence/tests/test_job_store_concurrency.py b/libs/policyengine-simulation-persistence/tests/test_job_store_concurrency.py new file mode 100644 index 000000000..2eaa6d00d --- /dev/null +++ b/libs/policyengine-simulation-persistence/tests/test_job_store_concurrency.py @@ -0,0 +1,268 @@ +"""Two or more collectors, dispatchers or submitters acting on one job at once. + +Result collection takes no lease: every write is a conditional update, so a +second writer changes nothing. These tests run against PostgreSQL. The first +group forces the exact interleaving (the second writer blocks on the first +writer's row lock, then re-checks its condition once the first commits); the +second group races many threads through the public store methods. +""" + +from __future__ import annotations + +import threading +import time +from concurrent.futures import ThreadPoolExecutor + +import pytest +from job_store_support import new_job, running_job, runtime_store +from policyengine_simulation_persistence import normalized_database_url +from policyengine_simulation_persistence.jobs import JobStatus, job_results, jobs +from sqlalchemy import create_engine, func, select, text + +pytestmark = pytest.mark.postgres + +WRITERS = 8 +BLOCK_TIMEOUT_SECONDS = 10 + + +def _engine(database_url: str, application_name: str): + return create_engine( + normalized_database_url(database_url), + connect_args={"application_name": application_name}, + ) + + +def _wait_until_waiting_on_a_lock(observer, application_name: str) -> None: + deadline = time.monotonic() + BLOCK_TIMEOUT_SECONDS + query = text( + "SELECT count(*) FROM pg_stat_activity " + "WHERE application_name = :name AND wait_event_type = 'Lock'" + ) + while time.monotonic() < deadline: + with observer.connect() as connection: + if connection.execute(query, {"name": application_name}).scalar(): + return + time.sleep(0.02) + pytest.fail(f"{application_name} never waited on the first writer's row lock") + + +class _Interleaving: + """Hold writer A's transaction open while writer B runs into its lock.""" + + def __init__(self, database) -> None: + self.first = _engine(database.runtime_url, "first-writer") + self.second = _engine(database.runtime_url, "second-writer") + self.observer = _engine(database.superuser_url, "observer") + + def run(self, first_write, second_write, *, first_commits: bool = True): + with ThreadPoolExecutor(max_workers=1) as executor: + with self.first.connect() as first_connection: + first_transaction = first_connection.begin() + first_outcome = first_write(first_connection) + + def second() -> bool: + with self.second.begin() as second_connection: + return second_write(second_connection) + + second_future = executor.submit(second) + _wait_until_waiting_on_a_lock(self.observer, "second-writer") + if first_commits: + first_transaction.commit() + else: + first_transaction.rollback() + return first_outcome, second_future.result(timeout=BLOCK_TIMEOUT_SECONDS) + + def dispose(self) -> None: + for engine in (self.first, self.second, self.observer): + engine.dispose() + + +@pytest.fixture +def interleaving(migrated_database): + controller = _Interleaving(migrated_database) + yield controller + controller.dispose() + + +def _stored_results(store, job_id): + with store.engine.connect() as connection: + return connection.execute( + select(job_results.c.result).where(job_results.c.job_id == job_id) + ).all() + + +# --- Forced interleavings ------------------------------------------------------ + + +def test_a_second_collector_waits_for_the_first_then_changes_nothing( + migrated_database, interleaving +) -> None: + store = runtime_store(migrated_database) + job = running_job(store) + + first, second = interleaving.run( + lambda connection: store._complete( + connection, job.id, result={"collector": "first"}, result_schema_version=1 + ), + lambda connection: store._complete( + connection, job.id, result={"collector": "second"}, result_schema_version=1 + ), + ) + + assert (first, second) == (True, False) + assert store.get(job.id).status is JobStatus.SUCCEEDED + assert _stored_results(store, job.id) == [({"collector": "first"},)] + + +def test_a_failure_racing_a_result_cannot_overwrite_it( + migrated_database, interleaving +) -> None: + store = runtime_store(migrated_database) + job = running_job(store) + + first, second = interleaving.run( + lambda connection: store._complete( + connection, job.id, result={"collector": "first"}, result_schema_version=1 + ), + lambda connection: store._fail( + connection, job.id, error_code="worker_error", error_message="late" + ), + ) + + assert (first, second) == (True, False) + finished = store.get(job.id) + assert finished.status is JobStatus.SUCCEEDED + assert finished.error_code is None + + +def test_a_result_racing_a_failure_is_discarded( + migrated_database, interleaving +) -> None: + store = runtime_store(migrated_database) + job = running_job(store) + + first, second = interleaving.run( + lambda connection: store._fail( + connection, job.id, error_code="result_expired", error_message="expired" + ), + lambda connection: store._complete( + connection, job.id, result={"collector": "late"}, result_schema_version=1 + ), + ) + + assert (first, second) == (True, False) + assert store.get(job.id).status is JobStatus.FAILED + assert _stored_results(store, job.id) == [] + + +def test_if_the_first_collector_rolls_back_the_second_writes_the_result( + migrated_database, interleaving +) -> None: + store = runtime_store(migrated_database) + job = running_job(store) + + first, second = interleaving.run( + lambda connection: store._complete( + connection, job.id, result={"collector": "first"}, result_schema_version=1 + ), + lambda connection: store._complete( + connection, job.id, result={"collector": "second"}, result_schema_version=1 + ), + first_commits=False, + ) + + assert (first, second) == (True, True) + assert _stored_results(store, job.id) == [({"collector": "second"},)] + + +def test_a_second_dispatcher_waits_for_the_first_then_claims_nothing( + migrated_database, interleaving +) -> None: + store = runtime_store(migrated_database) + job = store.submit(new_job()).job + + first, second = interleaving.run( + lambda connection: store._claim_dispatch(connection, job.id), + lambda connection: store._claim_dispatch(connection, job.id), + ) + + assert (first, second) == (True, False) + assert store.get(job.id).status is JobStatus.DISPATCHING + + +# --- Many threads through the public methods ------------------------------------ + + +def _race(action, count: int = WRITERS): + barrier = threading.Barrier(count) + + def contender(index: int): + barrier.wait() + return action(index) + + with ThreadPoolExecutor(max_workers=count) as executor: + return list(executor.map(contender, range(count))) + + +def test_concurrent_collectors_write_exactly_one_result(migrated_database) -> None: + store = runtime_store(migrated_database) + for case in range(5): + job = running_job(store, {"case": case}) + + outcomes = _race( + lambda index, job_id=job.id: store.complete( + job_id, result={"collector": index}, result_schema_version=1 + ) + ) + + assert outcomes.count(True) == 1 + winner = outcomes.index(True) + assert _stored_results(store, job.id) == [({"collector": winner},)] + assert store.get(job.id).status is JobStatus.SUCCEEDED + + +def test_concurrent_results_and_failures_leave_one_consistent_outcome( + migrated_database, +) -> None: + store = runtime_store(migrated_database) + for case in range(5): + job = running_job(store, {"case": case}) + + def collect(index, job_id=job.id): + if index % 2: + return store.fail(job_id, error_code="worker_error", error_message="x") + return store.complete( + job_id, result={"collector": index}, result_schema_version=1 + ) + + outcomes = _race(collect) + + assert outcomes.count(True) == 1 + winner = outcomes.index(True) + finished = store.get(job.id) + if winner % 2: + assert finished.status is JobStatus.FAILED + assert _stored_results(store, job.id) == [] + else: + assert finished.status is JobStatus.SUCCEEDED + assert _stored_results(store, job.id) == [({"collector": winner},)] + + +def test_concurrent_dispatch_claims_start_one_modal_call(migrated_database) -> None: + store = runtime_store(migrated_database) + job = store.submit(new_job()).job + + outcomes = _race(lambda _index: store.claim_dispatch(job.id)) + + assert outcomes.count(True) == 1 + + +def test_concurrent_identical_submissions_create_one_job(migrated_database) -> None: + store = runtime_store(migrated_database) + + submissions = _race(lambda _index: store.submit(new_job({"shared": True}))) + + assert [submission.created for submission in submissions].count(True) == 1 + assert len({submission.job.id for submission in submissions}) == 1 + with store.engine.connect() as connection: + assert connection.execute(select(func.count()).select_from(jobs)).scalar() == 1 diff --git a/migrations/versions/0002_simulation_jobs.py b/migrations/versions/0002_simulation_jobs.py new file mode 100644 index 000000000..8e4f66931 --- /dev/null +++ b/migrations/versions/0002_simulation_jobs.py @@ -0,0 +1,116 @@ +"""Add the simulation job store: simulation.jobs and simulation.job_results. + +One ``jobs`` row records each calculation the Simulation Entrypoint submits to +Modal: who submitted it, its input hash and policyengine.py bundle version, the +Modal call, its status, timings and any error. ``job_results`` holds the +validated result of each succeeded job that is not a report. The runtime role +reads and writes both through the default privileges granted in 0001. + +Revision ID: 0002 +Revises: 0001 +Create Date: 2026-10-20 +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +revision = "0002" +down_revision = "0001" +branch_labels = None +depends_on = None + +SCHEMA = "simulation" + + +def upgrade() -> None: + """Create the jobs and job_results tables.""" + op.create_table( + "jobs", + sa.Column("id", postgresql.UUID(as_uuid=True), nullable=False), + sa.Column("kind", sa.Text(), nullable=False), + sa.Column("status", sa.Text(), nullable=False), + sa.Column("submitted_by", sa.Text(), nullable=False), + sa.Column("input_sha256", sa.String(length=64), nullable=False), + sa.Column("policyengine_version", sa.Text(), nullable=False), + sa.Column("modal_application", sa.Text(), nullable=False), + sa.Column("modal_function", sa.Text(), nullable=False), + sa.Column("modal_call_id", sa.Text(), nullable=True), + sa.Column("observability_id", sa.Text(), nullable=True), + sa.Column("error_code", sa.Text(), nullable=True), + sa.Column("error_message", sa.Text(), nullable=True), + sa.Column( + "collection_attempts", + sa.Integer(), + server_default=sa.text("0"), + nullable=False, + ), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("updated_at", sa.DateTime(timezone=True), nullable=False), + sa.Column("dispatched_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("last_checked_at", sa.DateTime(timezone=True), nullable=True), + sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True), + sa.CheckConstraint( + "status IN ('pending', 'dispatching', 'running', 'succeeded', 'failed')", + name="ck_jobs_status", + ), + sa.CheckConstraint( + "input_sha256 ~ '^[0-9a-f]{64}$'", + name="ck_jobs_input_sha256", + ), + sa.CheckConstraint( + "collection_attempts >= 0", + name="ck_jobs_collection_attempts", + ), + sa.PrimaryKeyConstraint("id", name="pk_jobs"), + sa.UniqueConstraint("modal_call_id", name="uq_jobs_modal_call_id"), + schema=SCHEMA, + ) + op.create_index( + "uq_jobs_input_sha256_unfailed", + "jobs", + ["input_sha256"], + unique=True, + schema=SCHEMA, + postgresql_where=sa.text("status != 'failed'"), + ) + op.create_index( + "ix_jobs_unfinished_status", + "jobs", + ["status"], + schema=SCHEMA, + postgresql_where=sa.text("status IN ('pending', 'dispatching', 'running')"), + ) + op.create_table( + "job_results", + sa.Column("job_id", postgresql.UUID(as_uuid=True), nullable=False), + sa.Column("result", postgresql.JSONB(), nullable=False), + sa.Column("result_schema_version", sa.Integer(), nullable=False), + sa.Column("result_sha256", sa.String(length=64), nullable=False), + sa.Column("created_at", sa.DateTime(timezone=True), nullable=False), + sa.CheckConstraint( + "result_sha256 ~ '^[0-9a-f]{64}$'", + name="ck_job_results_result_sha256", + ), + sa.CheckConstraint( + "result_schema_version >= 1", + name="ck_job_results_result_schema_version", + ), + sa.ForeignKeyConstraint( + ["job_id"], + [f"{SCHEMA}.jobs.id"], + name="fk_job_results_job_id_jobs", + ), + sa.PrimaryKeyConstraint("job_id", name="pk_job_results"), + schema=SCHEMA, + ) + + +def downgrade() -> None: + """Drop the job_results and jobs tables.""" + op.drop_table("job_results", schema=SCHEMA) + op.drop_index("ix_jobs_unfinished_status", table_name="jobs", schema=SCHEMA) + op.drop_index("uq_jobs_input_sha256_unfailed", table_name="jobs", schema=SCHEMA) + op.drop_table("jobs", schema=SCHEMA) diff --git a/projects/policyengine-simulation-entry/DEPLOYMENT.md b/projects/policyengine-simulation-entry/DEPLOYMENT.md index d22fab8ca..4eef82b72 100644 --- a/projects/policyengine-simulation-entry/DEPLOYMENT.md +++ b/projects/policyengine-simulation-entry/DEPLOYMENT.md @@ -46,6 +46,32 @@ rollback checks do not, because an older image may run against a newer schema. The ownership rules and roles are described in `libs/policyengine-simulation-persistence/README.md`. +## Simulation job collection + +The job store's results are collected in the background by Cloud Tasks and a +Cloud Scheduler sweep, never by the clients that poll. The deploy step passes +the entrypoint: + +- `RECONCILIATION_QUEUE`, the full Cloud Tasks queue name, from the GitHub + environment variable of the same name; +- `RECONCILIATION_INVOKER_SERVICE_ACCOUNT`, the service account whose OIDC + token Cloud Tasks attaches to each collection task, from the GitHub + environment variable of the same name; +- `RECONCILIATION_TARGET_URL`, the stable service URL that the deploy job reads + from Cloud Run before deploying, so tasks never target a candidate revision; +- `SIMULATION_MODAL_ENVIRONMENT`, the deployment's Modal environment. + +`POST /internal/simulation-jobs/{job_id}/collect` (called by Cloud Tasks) and +`POST /internal/simulation-jobs/sweep` (called by Cloud Scheduler) accept only +a Google-signed OIDC token whose audience is `RECONCILIATION_TARGET_URL` and +whose email is the invoker service account. The Cloud Scheduler job must +therefore use that service account and audience. The entrypoint's runtime +service account needs permission to create tasks in the queue and to act as +the invoker service account; the invoker needs permission to invoke the +service. `GET /simulation-jobs/{job_reference}` reads a job by its UUID or +Modal `fc-` call ID from the database only. None of these routes is in the +OpenAPI document. + ## Candidate qualification For each eligible revision, the workflow first deploys `beta` diff --git a/projects/policyengine-simulation-entry/README.md b/projects/policyengine-simulation-entry/README.md index 3a3f72815..dbe877e93 100644 --- a/projects/policyengine-simulation-entry/README.md +++ b/projects/policyengine-simulation-entry/README.md @@ -39,6 +39,11 @@ Auth0 permission: | `GET /budget-window-jobs/{batch_job_id}` | `simulation:read` | | `POST /internal/stage12/reports` | `stage12:operate` | | `GET /internal/stage12/reports/{evaluation_id}` | `stage12:operate` | +| `GET /simulation-jobs/{job_reference}` | `simulation:read` | + +`POST /internal/simulation-jobs/{job_id}/collect` and +`POST /internal/simulation-jobs/sweep` do not use Auth0: they accept only the +Google OIDC token Cloud Tasks and Cloud Scheduler send (see `DEPLOYMENT.md`). `GET /health`, `GET /ready`, `GET /versions`, `GET /versions/{kind}` and `POST /ping` stay public. diff --git a/projects/policyengine-simulation-entry/pyproject.toml b/projects/policyengine-simulation-entry/pyproject.toml index e4793eb6a..56df83a3c 100644 --- a/projects/policyengine-simulation-entry/pyproject.toml +++ b/projects/policyengine-simulation-entry/pyproject.toml @@ -14,6 +14,8 @@ requires-python = ">=3.13,<3.14" dependencies = [ "cryptography>=41.0.0", "fastapi>=0.115.0,<1", + # Default credentials for creating Cloud Tasks collection tasks. + "google-auth[requests]>=2.30,<3", "httpx>=0.28.0,<1", "modal>=1.4,<2", "policyengine-observability[fastapi,google,httpx,otlp-grpc]>=3.0.2,<4", diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py index 95e3039f2..b8929e26e 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py @@ -10,8 +10,9 @@ from typing import TYPE_CHECKING, Any, Literal, Protocol from uuid import UUID -from fastapi import Depends, FastAPI, Request +from fastapi import Depends, FastAPI, HTTPException, Request from fastapi.responses import JSONResponse, Response +from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer from policyengine_observability import REQUEST_ID_HEADER from policyengine_simulation_contract.gateway_models import ( BudgetWindowBatchRequest, @@ -42,6 +43,7 @@ init_simulation_observability, ) from policyengine_simulation_persistence import HEAD_REVISION +from policyengine_simulation_persistence.jobs import JobStore from pydantic import BaseModel, TypeAdapter, ValidationError from starlette.datastructures import MutableHeaders @@ -58,16 +60,32 @@ SimulationBackend, ) from policyengine_simulation_entry.config import Settings +from policyengine_simulation_entry.google_oidc import ( + GoogleOidcVerifier, + InternalCallerRejectedError, +) from policyengine_simulation_entry.schemas import ( BackendTelemetryAttributes, BackendTelemetryPayload, RequestIdentifiers, SimulationDatabaseReadiness, SimulationEntrypointReadinessResponse, + SimulationJobCollectionResponse, + SimulationJobError, + SimulationJobStatusResponse, + SimulationJobSweepResponse, TemporaryStage12ReportResponse, TemporaryStage12SubmissionResponse, ) from policyengine_simulation_entry.simulation_database import SimulationDatabase +from policyengine_simulation_entry.simulation_jobs import ( + REGISTERED_JOB_KINDS, + CloudTasksCollectionQueue, + JobView, + ModalCallDispatcher, + ModalCallInspector, + SimulationJobService, +) from policyengine_simulation_entry.stage12_backend import ( Stage12ComparisonBackend, TemporaryStage12DispatchFailed, @@ -87,6 +105,19 @@ type ResponseIdentifier = Literal["job_id", "batch_job_id"] STAGE12_MODAL_SUBMISSION_TIMEOUT_SECONDS = 5.0 SIMULATION_DATABASE_READINESS_TIMEOUT_SECONDS = 3.0 +SIMULATION_JOB_RETRY_AFTER_SECONDS = "5" +_internal_bearer = HTTPBearer(auto_error=False) + + +class InternalCallerVerifier(Protocol): + """Verify the token on a Cloud Tasks or Cloud Scheduler request. + + The real implementation is ``GoogleOidcVerifier``; tests inject fakes. + """ + + def verify(self, token: str) -> None: + """Return when the token is valid; raise InternalCallerRejectedError.""" + ... class SimulationDatabaseConnection(Protocol): @@ -192,6 +223,8 @@ def create_app( auth_dependency: AuthenticationDependency | None = None, comparison_backend: ComparisonBackend | None = None, simulation_database: SimulationDatabaseConnection | None = None, + simulation_jobs: SimulationJobService | None = None, + internal_caller_verifier: InternalCallerVerifier | None = None, ) -> FastAPI: """Build the app with injectable auth/backend seams for hermetic tests.""" runtime_settings = settings or Settings.from_env() @@ -200,6 +233,37 @@ def create_app( runtime_database = SimulationDatabase.from_url( runtime_settings.simulation_database_url ) + runtime_jobs = simulation_jobs + runtime_collection_queue: CloudTasksCollectionQueue | None = None + if ( + runtime_jobs is None + and runtime_settings.simulation_jobs_configured + and isinstance(runtime_database, SimulationDatabase) + ): + runtime_collection_queue = CloudTasksCollectionQueue( + queue=runtime_settings.reconciliation_queue, + target_url=runtime_settings.reconciliation_target_url, + invoker_service_account=( + runtime_settings.reconciliation_invoker_service_account + ), + ) + runtime_jobs = SimulationJobService( + store=JobStore(runtime_database.engine), + dispatcher=ModalCallDispatcher( + runtime_settings.simulation_modal_environment + ), + inspector=ModalCallInspector(), + queue=runtime_collection_queue, + job_kinds=REGISTERED_JOB_KINDS, + ) + runtime_verifier = internal_caller_verifier + if runtime_verifier is None and runtime_settings.simulation_jobs_configured: + runtime_verifier = GoogleOidcVerifier( + audience=runtime_settings.reconciliation_target_url, + service_account_email=( + runtime_settings.reconciliation_invoker_service_account + ), + ) @asynccontextmanager async def lifespan(_: FastAPI) -> AsyncIterator[None]: @@ -209,6 +273,8 @@ async def lifespan(_: FastAPI) -> AsyncIterator[None]: yield finally: await runtime_backend.close() + if runtime_collection_queue is not None: + await runtime_collection_queue.close() if runtime_database is not None: runtime_database.dispose() runtime.shutdown() @@ -594,6 +660,153 @@ async def get_temporary_stage12_report( headers=response_headers or None, ) + def simulation_jobs_unavailable() -> JSONResponse: + return JSONResponse( + status_code=503, + content={"detail": "Simulation jobs are unavailable."}, + headers={"Retry-After": "10"}, + ) + + async def internal_caller( + credentials: HTTPAuthorizationCredentials | None = Depends(_internal_bearer), + ) -> None: + """Admit only Cloud Tasks and Cloud Scheduler, by their Google token.""" + verifier = runtime_verifier + if verifier is None: + raise HTTPException(status_code=503, headers={"Retry-After": "10"}) + if credentials is None: + raise HTTPException(status_code=403) + try: + await asyncio.to_thread(verifier.verify, credentials.credentials) + except InternalCallerRejectedError as error: + logger.info( + "simulation_job_internal_caller_rejected", + extra={"reason": str(error)}, + ) + raise HTTPException(status_code=403) from error + + def job_status_payload(view: JobView) -> SimulationJobStatusResponse: + job = view.job + return SimulationJobStatusResponse( + job_id=job.id, + kind=job.kind, + status=job.status.value, + policyengine_version=job.policyengine_version, + modal_call_id=job.modal_call_id, + created_at=job.created_at, + dispatched_at=job.dispatched_at, + completed_at=job.completed_at, + error=( + SimulationJobError(code=job.error_code, message=job.error_message) + if job.error_code is not None + else None + ), + result=None if view.result is None else view.result.result, + result_sha256=None if view.result is None else view.result.result_sha256, + ) + + # Simulation job store (API v2 migration, Step 1). These routes are not in + # the OpenAPI document and no production caller uses them yet. Reading a + # job uses only the database and never checks or advances its Modal call. + @app.get( + "/simulation-jobs/{job_reference}", + include_in_schema=False, + dependencies=[Depends(simulation_reader)], + ) + async def get_simulation_job(job_reference: str, request: Request) -> Response: + service = runtime_jobs + if service is None: + return simulation_jobs_unavailable() + try: + view = await service.find(job_reference) + # A failed database read becomes a 503; the caller polls again. + except Exception as error: # noqa: BLE001 # any read failure becomes a 503 + logger.error( # noqa: TRY400 # the error type is enough; the traceback could expose the job reference's data + "simulation_job_status_failed", + extra={ + "request_id": request.state.request_id, + "error_type": type(error).__name__, + }, + ) + return simulation_jobs_unavailable() + if view is None: + return JSONResponse( + status_code=404, + content={"detail": "Simulation job was not found."}, + ) + headers: dict[str, str] = {} + job_observability_id = bind_observability_id(request, view.job.observability_id) + if job_observability_id is not None: + headers[OBSERVABILITY_ID_HEADER] = job_observability_id + if not view.job.finished: + headers["Retry-After"] = SIMULATION_JOB_RETRY_AFTER_SECONDS + return JSONResponse( + status_code=200 if view.job.finished else 202, + content=job_status_payload(view).model_dump(mode="json"), + headers=headers or None, + ) + + @app.post( + "/internal/simulation-jobs/{job_id}/collect", + include_in_schema=False, + dependencies=[Depends(internal_caller)], + ) + async def collect_simulation_job(job_id: UUID, request: Request) -> Response: + service = runtime_jobs + if service is None: + return simulation_jobs_unavailable() + try: + report = await service.collect(job_id) + # Modal, Cloud Tasks or the database failed; a 503 makes Cloud Tasks + # retry this task with its queue's backoff. + except Exception as error: # noqa: BLE001 # every transient failure becomes a 503 + logger.warning( + "simulation_job_collection_failed", + extra={ + "request_id": request.state.request_id, + "job_id": str(job_id), + "error_type": type(error).__name__, + }, + ) + return simulation_jobs_unavailable() + payload = SimulationJobCollectionResponse( + outcome=report.outcome.value, + job_status=None if report.job_status is None else report.job_status.value, + next_check_seconds=report.next_check_seconds, + ) + return JSONResponse(content=payload.model_dump(mode="json")) + + @app.post( + "/internal/simulation-jobs/sweep", + include_in_schema=False, + dependencies=[Depends(internal_caller)], + ) + async def sweep_simulation_jobs(request: Request) -> Response: + service = runtime_jobs + if service is None: + return simulation_jobs_unavailable() + try: + report = await service.sweep() + except Exception as error: # noqa: BLE001 # a failed sweep becomes a 503; the next run repeats it + logger.warning( + "simulation_job_sweep_failed", + extra={ + "request_id": request.state.request_id, + "error_type": type(error).__name__, + }, + ) + return simulation_jobs_unavailable() + payload = SimulationJobSweepResponse( + abandoned=report.abandoned, + enqueued=report.enqueued, + enqueue_failures=report.enqueue_failures, + ) + # Report refused tasks as a failed run, so Cloud Scheduler shows it. + return JSONResponse( + status_code=503 if report.enqueue_failures else 200, + content=payload.model_dump(mode="json"), + ) + @app.post( "/simulate/economy/comparison", summary="Submit Simulation", diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py index a88531302..c88a43865 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/config.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +import re from dataclasses import dataclass from urllib.parse import urlparse @@ -14,6 +15,12 @@ PRODUCTION_ENVIRONMENTS = frozenset({"main", "prod", "production"}) MODAL_GATEWAY_HOST_SUFFIX = ".modal.run" +CLOUD_TASKS_QUEUE_PATTERN = re.compile( + r"projects/[a-z][a-z0-9-]{4,28}[a-z0-9]/locations/[a-z0-9-]+/queues/[A-Za-z0-9-]{1,100}" +) +SERVICE_ACCOUNT_PATTERN = re.compile( + r"[a-z][a-z0-9-]{4,28}[a-z0-9]@[a-z][a-z0-9-]{4,28}[a-z0-9]\.iam\.gserviceaccount\.com" +) class ConfigurationError(RuntimeError): @@ -94,6 +101,14 @@ class Settings: # outside deployed environments; nothing on the old-gateway or Stage 12 # paths reads it. simulation_database_url: str = "" + # Background collection of simulation job results. Cloud Tasks calls the + # collection route of RECONCILIATION_TARGET_URL with a Google OIDC token of + # RECONCILIATION_INVOKER_SERVICE_ACCOUNT; Cloud Scheduler calls the sweep + # route with the same identity. + reconciliation_queue: str = "" + reconciliation_invoker_service_account: str = "" + reconciliation_target_url: str = "" + simulation_modal_environment: str = "" @classmethod def from_env(cls) -> Settings: @@ -145,6 +160,14 @@ def from_env(cls) -> Settings: stage12_artifact_bucket=os.getenv("STAGE12_ARTIFACT_BUCKET", ""), stage12_retention_days=int(os.getenv("STAGE12_RETENTION_DAYS", "30")), simulation_database_url=os.getenv("SIMULATION_DATABASE_URL", ""), + reconciliation_queue=os.getenv("RECONCILIATION_QUEUE", ""), + reconciliation_invoker_service_account=os.getenv( + "RECONCILIATION_INVOKER_SERVICE_ACCOUNT", "" + ), + reconciliation_target_url=os.getenv("RECONCILIATION_TARGET_URL", "").rstrip( + "/" + ), + simulation_modal_environment=os.getenv("SIMULATION_MODAL_ENVIRONMENT", ""), ) @property @@ -163,6 +186,57 @@ def stage12_resources_configured(self) -> bool: ) ) + @property + def _simulation_job_settings(self) -> dict[str, str]: + return { + "SIMULATION_DATABASE_URL": self.simulation_database_url, + "RECONCILIATION_QUEUE": self.reconciliation_queue, + "RECONCILIATION_INVOKER_SERVICE_ACCOUNT": ( + self.reconciliation_invoker_service_account + ), + "RECONCILIATION_TARGET_URL": self.reconciliation_target_url, + "SIMULATION_MODAL_ENVIRONMENT": self.simulation_modal_environment, + } + + @property + def simulation_jobs_configured(self) -> bool: + """Whether the job store, its queue and its Modal environment are all set.""" + return all(self._simulation_job_settings.values()) + + def _validate_simulation_jobs(self) -> None: + reconciliation = { + name: value + for name, value in self._simulation_job_settings.items() + if name != "SIMULATION_DATABASE_URL" + } + if not any(reconciliation.values()): + return + missing = [ + name for name, value in self._simulation_job_settings.items() if not value + ] + if missing: + msg = f"Simulation jobs require: {', '.join(missing)}." + raise ConfigurationError(msg) + if not CLOUD_TASKS_QUEUE_PATTERN.fullmatch(self.reconciliation_queue): + msg = ( + "RECONCILIATION_QUEUE must be a full Cloud Tasks queue name: " + "projects//locations//queues/." + ) + raise ConfigurationError(msg) + if not SERVICE_ACCOUNT_PATTERN.fullmatch( + self.reconciliation_invoker_service_account + ): + msg = ( + "RECONCILIATION_INVOKER_SERVICE_ACCOUNT must be a Google service " + "account email." + ) + raise ConfigurationError(msg) + _validate_https_url( + "RECONCILIATION_TARGET_URL", + self.reconciliation_target_url, + allow_path=False, + ) + def validate(self) -> None: """Raise ConfigurationError unless the service can start safely. @@ -262,3 +336,4 @@ def validate(self) -> None: ): msg = "SIMULATION_DATABASE_URL must use the canonical postgresql:// scheme." raise ConfigurationError(msg) + self._validate_simulation_jobs() diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/google_oidc.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/google_oidc.py new file mode 100644 index 000000000..9c6ddc54b --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/google_oidc.py @@ -0,0 +1,69 @@ +"""Verify the Google OIDC tokens Cloud Tasks and Cloud Scheduler send. + +The internal simulation job routes are reachable on the public service URL, +so each request must carry a Google-signed ID token whose audience is the +service URL and whose email is the configured invoker service account. +""" + +from __future__ import annotations + +from typing import Protocol + +import jwt + +GOOGLE_CERTIFICATES_URL = "https://www.googleapis.com/oauth2/v3/certs" +GOOGLE_ISSUERS = ("https://accounts.google.com", "accounts.google.com") +REQUIRED_CLAIMS = ["exp", "iat", "iss", "aud", "email"] + + +class InternalCallerRejectedError(Exception): + """The token is missing, invalid, or from another identity.""" + + +class SigningKeySource(Protocol): + """Where verification finds the key that signed a token.""" + + def get_signing_key_from_jwt(self, token: str) -> jwt.PyJWK: + """Return the key named by the token's header.""" + ... + + +class GoogleOidcVerifier: + """Accept only Google ID tokens of one service account for one audience.""" + + def __init__( + self, + *, + audience: str, + service_account_email: str, + jwks_client: SigningKeySource | None = None, + ) -> None: + """Bind the expected audience and account; keys come from Google.""" + self.audience = audience + self.service_account_email = service_account_email + self.jwks_client = jwks_client or jwt.PyJWKClient(GOOGLE_CERTIFICATES_URL) + + def verify(self, token: str) -> None: + """Return when ``token`` is valid; otherwise raise the rejection. + + The rejection's message names only the failure type, never the token. + Fetching Google's keys may block, so callers run this off the event loop. + """ + try: + signing_key = self.jwks_client.get_signing_key_from_jwt(token).key + claims = jwt.decode( + token, + signing_key, + algorithms=["RS256"], + audience=self.audience, + issuer=GOOGLE_ISSUERS, + options={"require": REQUIRED_CLAIMS}, + ) + except jwt.PyJWTError as error: + raise InternalCallerRejectedError(type(error).__name__) from error + if ( + claims.get("email") != self.service_account_email + or claims.get("email_verified") is not True + ): + msg = "unexpected_service_account" + raise InternalCallerRejectedError(msg) diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/schemas.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/schemas.py index 74aa0f74e..d471e6d2b 100644 --- a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/schemas.py +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/schemas.py @@ -2,9 +2,11 @@ from __future__ import annotations +from datetime import datetime from typing import Literal, TypedDict from uuid import UUID +from policyengine_simulation_contract.json_types import JsonObject from policyengine_simulation_contract.stage12_execution import ( ComparisonReportRecord, ComparisonRunLifecycleStatus, @@ -140,3 +142,63 @@ class SimulationEntrypointReadinessResponse(BaseModel): status: Literal["ready", "not_ready"] simulation_database: SimulationDatabaseReadiness | None = None + + +class SimulationJobError(BaseModel): + """Why a simulation job failed; never worker data.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + code: str + message: str | None + + +class SimulationJobStatusResponse(BaseModel): + """One simulation job as the database records it. + + ``result`` is present once the job succeeded. Reading a job never checks + or advances its Modal call. + """ + + model_config = ConfigDict(extra="forbid", frozen=True) + + job_id: UUID + kind: str + status: Literal["pending", "dispatching", "running", "succeeded", "failed"] + policyengine_version: str + modal_call_id: str | None + created_at: datetime + dispatched_at: datetime | None + completed_at: datetime | None + error: SimulationJobError | None + result: JsonObject | None + result_sha256: str | None + + +class SimulationJobCollectionResponse(BaseModel): + """What one Cloud Tasks collection request did.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + outcome: Literal[ + "unknown_job", + "not_running", + "still_running", + "succeeded", + "failed", + "already_finished", + ] + job_status: ( + Literal["pending", "dispatching", "running", "succeeded", "failed"] | None + ) + next_check_seconds: float | None + + +class SimulationJobSweepResponse(BaseModel): + """What one Cloud Scheduler sweep did.""" + + model_config = ConfigDict(extra="forbid", frozen=True) + + abandoned: int + enqueued: int + enqueue_failures: int diff --git a/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_jobs.py b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_jobs.py new file mode 100644 index 000000000..152caeb9d --- /dev/null +++ b/projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_jobs.py @@ -0,0 +1,643 @@ +"""Simulation jobs: idempotent submission and background result collection. + +The Simulation Entrypoint records one ``simulation.jobs`` row per +calculation, starts its Modal call once, and collects the result in the +background: + +1. A submission inserts the job (identical inputs share one job), moves it + from ``pending`` to ``dispatching`` with a conditional update, spawns the + Modal call, records the call ID and schedules a Cloud Tasks collection task. +2. Each collection task checks the Modal call without waiting. A running call + schedules the next task with exponential backoff; a finished call's result + is validated and written with a conditional update, so a second collector + changes nothing. No lease is taken. +3. A Cloud Scheduler sweep enqueues a task for running jobs that nobody has + checked recently, and fails jobs whose dispatch stalled. + +Status reads use only the database. Modal never receives database +credentials and never writes these tables. +""" + +from __future__ import annotations + +import asyncio +import logging +import re +import threading +from dataclasses import dataclass, field +from datetime import timedelta +from enum import StrEnum +from typing import TYPE_CHECKING, Any, Protocol +from uuid import UUID + +import google.auth +import google.auth.exceptions +import google.auth.transport.requests +import httpx +import modal +from policyengine_simulation_persistence.jobs import ( + JobRecord, + JobResultRecord, + JobStatus, + NewJob, + Submission, + input_sha256, + utc_now, +) +from pydantic import BaseModel, ValidationError + +if TYPE_CHECKING: + from collections.abc import Awaitable, Callable, Mapping + from datetime import datetime + +logger = logging.getLogger(__name__) + +FIRST_COLLECTION_DELAY_SECONDS = 2.0 +COLLECTION_BACKOFF_BASE_SECONDS = 2 +COLLECTION_BACKOFF_MAX_SECONDS = 60 +# Longer than the longest backoff, so a job whose task chain is alive is +# never swept. +COLLECTION_STALE_AFTER = timedelta(minutes=5) +DISPATCH_TIMEOUT = timedelta(minutes=10) +SWEEP_BATCH_SIZE = 200 +SWEEP_CONCURRENCY = 10 +MODAL_CALL_ID_PATTERN = re.compile(r"fc-[0-9A-Za-z]+") +CLOUD_TASKS_API = "https://cloudtasks.googleapis.com/v2" +CLOUD_PLATFORM_SCOPE = "https://www.googleapis.com/auth/cloud-platform" +CLOUD_TASKS_TIMEOUT_SECONDS = 5.0 + + +def collection_delay_seconds(attempts: int) -> int: + """Return the wait before the next check after ``attempts`` checks.""" + return min( + COLLECTION_BACKOFF_BASE_SECONDS * 2 ** max(attempts - 1, 0), + COLLECTION_BACKOFF_MAX_SECONDS, + ) + + +@dataclass(frozen=True) +class JobKind: + """A kind of calculation the job store can collect. + + Step 2 registers the household calculation and report kinds; until then no + kind is registered and nothing submits jobs. + """ + + name: str + result_model: type[BaseModel] + result_schema_version: int + + +# The job kinds production collects, by name. +REGISTERED_JOB_KINDS: dict[str, JobKind] = {} + + +@dataclass(frozen=True) +class ModalTarget: + """The deployed Modal function a job calls.""" + + application: str + function: str + + +@dataclass(frozen=True) +class CallRunning: + """The Modal call has not produced its output yet.""" + + +@dataclass(frozen=True) +class CallSucceeded: + """The Modal call returned ``output``.""" + + output: object + + +@dataclass(frozen=True) +class CallFailed: + """The Modal call failed; the message carries no worker data.""" + + error_code: str + error_message: str + + +type CallOutcome = CallRunning | CallSucceeded | CallFailed + + +class CallInspectionError(RuntimeError): + """Modal could not be asked about a call; try again later.""" + + +class CollectionQueueError(RuntimeError): + """Cloud Tasks did not accept a collection task.""" + + +class UnknownJobKindError(ValueError): + """The job kind is not registered.""" + + +class CallDispatcher(Protocol): + """Start a Modal call without waiting for it.""" + + async def spawn( + self, + target: ModalTarget, + request_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Return the acknowledged Modal call ID.""" + ... + + +class CallInspector(Protocol): + """Check a Modal call without waiting for it.""" + + async def inspect(self, modal_call_id: str) -> CallOutcome: + """Return the call's outcome; raise CallInspectionError if unknown.""" + ... + + +class CollectionQueue(Protocol): + """Schedule a collection task for a job.""" + + async def enqueue(self, job_id: UUID, *, delay_seconds: float) -> None: + """Schedule one collection; raise CollectionQueueError if refused.""" + ... + + +class JobStoreOperations(Protocol): + """The ``JobStore`` operations the service uses.""" + + def submit(self, new_job: NewJob) -> Submission: + """Create the job, or return the existing one for the same input.""" + ... + + def claim_dispatch(self, job_id: UUID) -> bool: + """Move a pending job to dispatching.""" + ... + + def record_dispatched(self, job_id: UUID, *, modal_call_id: str) -> bool: + """Move a dispatching job to running.""" + ... + + def record_dispatch_failure( + self, job_id: UUID, *, error_code: str, error_message: str + ) -> bool: + """Fail a dispatching job.""" + ... + + def get(self, job_id: UUID) -> JobRecord | None: + """Return a job by ID.""" + ... + + def get_by_modal_call_id(self, modal_call_id: str) -> JobRecord | None: + """Return a job by Modal call ID.""" + ... + + def get_result(self, job_id: UUID) -> JobResultRecord | None: + """Return a succeeded job's result.""" + ... + + def record_check(self, job_id: UUID) -> int | None: + """Count one check of a running job.""" + ... + + def complete( + self, job_id: UUID, *, result: dict[str, Any], result_schema_version: int + ) -> bool: + """Store a running job's result.""" + ... + + def fail(self, job_id: UUID, *, error_code: str, error_message: str) -> bool: + """Fail a running job.""" + ... + + def jobs_to_collect(self, *, checked_before: datetime, limit: int) -> list[UUID]: + """Return running jobs not checked recently.""" + ... + + def abandon_stalled_dispatches(self, *, updated_before: datetime) -> list[UUID]: + """Fail jobs whose dispatch stalled.""" + ... + + +class CollectionOutcome(StrEnum): + """What one collection task did.""" + + UNKNOWN_JOB = "unknown_job" + NOT_RUNNING = "not_running" + STILL_RUNNING = "still_running" + SUCCEEDED = "succeeded" + FAILED = "failed" + ALREADY_FINISHED = "already_finished" + + +@dataclass(frozen=True) +class CollectionReport: + """The result of one collection task.""" + + outcome: CollectionOutcome + job_status: JobStatus | None + next_check_seconds: float | None = None + + +@dataclass(frozen=True) +class SweepReport: + """What one sweep did.""" + + abandoned: int + enqueued: int + enqueue_failures: int + + +@dataclass(frozen=True) +class JobView: + """A job and, once it succeeded, its stored result.""" + + job: JobRecord + result: JobResultRecord | None + + +@dataclass +class SimulationJobService: + """Submit simulation jobs and collect their results in the background.""" + + store: JobStoreOperations + dispatcher: CallDispatcher + inspector: CallInspector + queue: CollectionQueue + job_kinds: Mapping[str, JobKind] = field(default_factory=dict) + clock: Callable[[], datetime] = utc_now + + # --- Submission ------------------------------------------------------------ + + async def submit( # noqa: PLR0913 # every argument is part of the job's identity or provenance + self, + *, + kind: str, + request_payload: dict[str, Any], + policyengine_version: str, + submitted_by: str, + target: ModalTarget, + observability_id: str | None = None, + observability_context: dict[str, Any] | None = None, + ) -> JobRecord: + """Record the job and start its Modal call unless one already exists. + + Return the job as it stands afterwards. A refused Modal call fails the + job instead of raising; a refused collection task is left to the sweep. + """ + if kind not in self.job_kinds: + raise UnknownJobKindError(kind) + submission = await asyncio.to_thread( + self.store.submit, + NewJob( + kind=kind, + submitted_by=submitted_by, + input_sha256=input_sha256( + kind=kind, + policyengine_version=policyengine_version, + request=request_payload, + ), + policyengine_version=policyengine_version, + modal_application=target.application, + modal_function=target.function, + observability_id=observability_id, + ), + ) + job = submission.job + if not submission.created or not await asyncio.to_thread( + self.store.claim_dispatch, job.id + ): + return await self._current(job) + try: + modal_call_id = await self.dispatcher.spawn( + target, + request_payload, + observability_context=observability_context, + ) + # Any failure to start the call fails the job; the caller may resubmit. + except Exception as error: # noqa: BLE001 # recorded on the job + logger.warning( + "simulation_job_dispatch_failed", + extra={"job_id": str(job.id), "error_type": type(error).__name__}, + ) + await asyncio.to_thread( + self.store.record_dispatch_failure, + job.id, + error_code="dispatch_failed", + error_message=( + f"Modal did not accept the call ({type(error).__name__})." + ), + ) + return await self._current(job) + await asyncio.to_thread( + self.store.record_dispatched, job.id, modal_call_id=modal_call_id + ) + try: + await self.queue.enqueue( + job.id, delay_seconds=FIRST_COLLECTION_DELAY_SECONDS + ) + except CollectionQueueError as error: + logger.warning( + "simulation_job_collection_not_scheduled", + extra={"job_id": str(job.id), "error_type": type(error).__name__}, + ) + return await self._current(job) + + async def _current(self, job: JobRecord) -> JobRecord: + return await asyncio.to_thread(self.store.get, job.id) or job + + # --- Collection -------------------------------------------------------------- + + async def collect(self, job_id: UUID) -> CollectionReport: + """Check one job's Modal call and record what it found. + + Raise CallInspectionError or CollectionQueueError for a transient + failure, so Cloud Tasks retries this task. + """ + job = await asyncio.to_thread(self.store.get, job_id) + if job is None: + return CollectionReport(CollectionOutcome.UNKNOWN_JOB, None) + if job.status is not JobStatus.RUNNING or job.modal_call_id is None: + return CollectionReport(CollectionOutcome.NOT_RUNNING, job.status) + outcome = await self.inspector.inspect(job.modal_call_id) + if isinstance(outcome, CallRunning): + attempts = await asyncio.to_thread(self.store.record_check, job.id) + if attempts is None: + return await self._finished_elsewhere(job.id) + delay = collection_delay_seconds(attempts) + await self.queue.enqueue(job.id, delay_seconds=delay) + return CollectionReport( + CollectionOutcome.STILL_RUNNING, JobStatus.RUNNING, delay + ) + if isinstance(outcome, CallFailed): + return await self._fail(job.id, outcome.error_code, outcome.error_message) + return await self._complete(job, outcome.output) + + async def _complete(self, job: JobRecord, output: object) -> CollectionReport: + kind = self.job_kinds.get(job.kind) + if kind is None: + return await self._fail( + job.id, + "unknown_job_kind", + f"No collector is registered for job kind {job.kind!r}.", + ) + try: + result = kind.result_model.model_validate(output).model_dump(mode="json") + except ValidationError as error: + return await self._fail( + job.id, + "invalid_result", + f"The result failed validation with {error.error_count()} errors.", + ) + if not isinstance(result, dict): + return await self._fail( + job.id, "invalid_result", "The result is not a JSON object." + ) + stored = await asyncio.to_thread( + self.store.complete, + job.id, + result=result, + result_schema_version=kind.result_schema_version, + ) + if not stored: + return await self._finished_elsewhere(job.id) + return CollectionReport(CollectionOutcome.SUCCEEDED, JobStatus.SUCCEEDED) + + async def _fail( + self, job_id: UUID, error_code: str, error_message: str + ) -> CollectionReport: + failed = await asyncio.to_thread( + self.store.fail, + job_id, + error_code=error_code, + error_message=error_message, + ) + if not failed: + return await self._finished_elsewhere(job_id) + return CollectionReport(CollectionOutcome.FAILED, JobStatus.FAILED) + + async def _finished_elsewhere(self, job_id: UUID) -> CollectionReport: + job = await asyncio.to_thread(self.store.get, job_id) + return CollectionReport( + CollectionOutcome.ALREADY_FINISHED, + None if job is None else job.status, + ) + + # --- The sweep ------------------------------------------------------------------ + + async def sweep(self) -> SweepReport: + """Fail stalled dispatches and enqueue running jobs not checked recently.""" + now = self.clock() + abandoned = await asyncio.to_thread( + lambda: self.store.abandon_stalled_dispatches( + updated_before=now - DISPATCH_TIMEOUT + ) + ) + job_ids = await asyncio.to_thread( + lambda: self.store.jobs_to_collect( + checked_before=now - COLLECTION_STALE_AFTER, + limit=SWEEP_BATCH_SIZE, + ) + ) + limit = asyncio.Semaphore(SWEEP_CONCURRENCY) + + async def enqueue(job_id: UUID) -> bool: + async with limit: + try: + await self.queue.enqueue(job_id, delay_seconds=0) + except CollectionQueueError as error: + logger.warning( + "simulation_job_sweep_enqueue_failed", + extra={ + "job_id": str(job_id), + "error_type": type(error).__name__, + }, + ) + return False + return True + + enqueued = await asyncio.gather(*(enqueue(job_id) for job_id in job_ids)) + return SweepReport( + abandoned=len(abandoned), + enqueued=enqueued.count(True), + enqueue_failures=enqueued.count(False), + ) + + # --- Status reads ---------------------------------------------------------------- + + async def find(self, reference: str) -> JobView | None: + """Return the job a UUID or Modal ``fc-`` call ID names, from the database.""" + try: + job_id = UUID(reference) + except ValueError: + if not MODAL_CALL_ID_PATTERN.fullmatch(reference): + return None + job = await asyncio.to_thread(self.store.get_by_modal_call_id, reference) + else: + job = await asyncio.to_thread(self.store.get, job_id) + if job is None: + return None + result = None + if job.status is JobStatus.SUCCEEDED: + result = await asyncio.to_thread(self.store.get_result, job.id) + return JobView(job=job, result=result) + + +# --- Modal ------------------------------------------------------------------------ + +# Modal client and service problems: the call's state is unknown, so retry. +TRANSIENT_MODAL_ERRORS = ( + modal.exception.ConnectionError, + modal.exception.AuthError, + modal.exception.InternalError, + modal.exception.ServiceError, + modal.exception.ResourceExhaustedError, +) + + +class ModalCallDispatcher: + """Spawn deployed Modal functions in one Modal environment.""" + + def __init__(self, environment: str) -> None: + """Look functions up in ``environment``.""" + self._environment = environment + + async def spawn( + self, + target: ModalTarget, + request_payload: dict[str, Any], + *, + observability_context: dict[str, Any] | None, + ) -> str: + """Spawn ``target`` and return Modal's call ID once Modal acknowledges it.""" + function = modal.Function.from_name( + target.application, + target.function, + environment_name=self._environment, + ) + call = await function.spawn.aio( + request_payload, + observability_context=observability_context, + ) + modal_call_id = getattr(call, "object_id", None) + if not isinstance(modal_call_id, str) or not modal_call_id: + msg = "Modal acknowledged the spawn with no call ID" + raise RuntimeError(msg) + return modal_call_id + + +class ModalCallInspector: + """Read a Modal call's output without waiting for it.""" + + async def inspect(self, modal_call_id: str) -> CallOutcome: # noqa: PLR0911 # one return per Modal outcome + """Return the call's outcome, or raise CallInspectionError.""" + try: + call = modal.FunctionCall.from_id(modal_call_id) + output = await call.get.aio(timeout=0) + except modal.exception.OutputExpiredError: + return CallFailed("result_expired", "Modal no longer holds the output.") + except modal.exception.FunctionTimeoutError: + return CallFailed("worker_timeout", "The Modal call timed out.") + except modal.exception.NotFoundError: + return CallFailed("call_not_found", "Modal does not know the call.") + except TRANSIENT_MODAL_ERRORS as error: + raise CallInspectionError(type(error).__name__) from error + except TimeoutError: + return CallRunning() + # The call re-raises whatever the worker raised. Record only its type: + # the message may contain calculation inputs. + except Exception as error: # noqa: BLE001 # every worker failure fails the job + return CallFailed( + "worker_error", + f"The Modal call raised {type(error).__name__}.", + ) + return CallSucceeded(output) + + +# --- Cloud Tasks ------------------------------------------------------------------- + + +class GoogleAccessTokens: + """OAuth access tokens from the runtime's default Google credentials.""" + + def __init__(self) -> None: + """Load credentials on first use.""" + self._credentials: Any = None + self._lock = threading.Lock() + + def _token(self) -> str: + with self._lock: + if self._credentials is None: + self._credentials, _ = google.auth.default( + scopes=[CLOUD_PLATFORM_SCOPE] + ) + if not self._credentials.valid: + self._credentials.refresh(google.auth.transport.requests.Request()) + return self._credentials.token + + async def __call__(self) -> str: + """Return a valid access token, refreshing it off the event loop.""" + return await asyncio.to_thread(self._token) + + +class CloudTasksCollectionQueue: + """Create Cloud Tasks HTTP tasks that call the collection route.""" + + def __init__( # noqa: PLR0913 # the queue, target, identity and three test seams + self, + *, + queue: str, + target_url: str, + invoker_service_account: str, + access_token: Callable[[], Awaitable[str]] | None = None, + http_client: httpx.AsyncClient | None = None, + clock: Callable[[], datetime] = utc_now, + ) -> None: + """Send tasks to ``queue`` that call ``target_url`` as the invoker.""" + self.queue = queue + self.target_url = target_url.rstrip("/") + self.invoker_service_account = invoker_service_account + self._access_token = access_token or GoogleAccessTokens() + self._http_client = http_client or httpx.AsyncClient( + timeout=CLOUD_TASKS_TIMEOUT_SECONDS + ) + self._clock = clock + + def collection_url(self, job_id: UUID) -> str: + """Return the collection route URL for ``job_id``.""" + return f"{self.target_url}/internal/simulation-jobs/{job_id}/collect" + + async def enqueue(self, job_id: UUID, *, delay_seconds: float) -> None: + """Create one task that calls the collection route after the delay.""" + task: dict[str, Any] = { + "httpRequest": { + "httpMethod": "POST", + "url": self.collection_url(job_id), + "oidcToken": { + "serviceAccountEmail": self.invoker_service_account, + "audience": self.target_url, + }, + } + } + if delay_seconds > 0: + scheduled = self._clock() + timedelta(seconds=delay_seconds) + task["scheduleTime"] = scheduled.isoformat().replace("+00:00", "Z") + try: + token = await self._access_token() + response = await self._http_client.post( + f"{CLOUD_TASKS_API}/{self.queue}/tasks", + json={"task": task}, + headers={"Authorization": f"Bearer {token}"}, + ) + except (httpx.HTTPError, google.auth.exceptions.GoogleAuthError) as error: + msg = f"Cloud Tasks could not be reached ({type(error).__name__})" + raise CollectionQueueError(msg) from error + if response.status_code >= httpx.codes.MULTIPLE_CHOICES: + msg = f"Cloud Tasks refused the task with HTTP {response.status_code}" + raise CollectionQueueError(msg) + + async def close(self) -> None: + """Close the HTTP client.""" + await self._http_client.aclose() diff --git a/projects/policyengine-simulation-entry/tests/simulation_job_support.py b/projects/policyengine-simulation-entry/tests/simulation_job_support.py new file mode 100644 index 000000000..c10536abd --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/simulation_job_support.py @@ -0,0 +1,295 @@ +"""Fakes for the simulation job service: store, Modal calls and Cloud Tasks. + +``InMemoryJobStore`` keeps the real store's contract: every transition is +conditional on the current status and applied under one lock, so a second +writer finds nothing to change. The real conditional SQL is tested against +PostgreSQL in libs/policyengine-simulation-persistence. +""" + +from __future__ import annotations + +import threading +from dataclasses import replace +from datetime import UTC, datetime, timedelta +from uuid import UUID, uuid4 + +from policyengine_simulation_entry.simulation_jobs import ( + CallRunning, + CollectionQueueError, + JobKind, + ModalTarget, +) +from policyengine_simulation_persistence.jobs import ( + DISPATCH_ABANDONED, + UNFINISHED_STATUSES, + JobRecord, + JobResultRecord, + JobStatus, + Submission, + result_sha256, +) +from pydantic import BaseModel, ConfigDict + +START = datetime(2026, 10, 20, 12, tzinfo=UTC) +TARGET = ModalTarget( + application="policyengine-simulation-v2-py6-7-5", + function="calculate_household", +) + + +class HouseholdResult(BaseModel): + """A stand-in result model for a registered job kind.""" + + model_config = ConfigDict(extra="forbid") + + net_income: float + + +HOUSEHOLD_KIND = JobKind( + name="household_calculation", + result_model=HouseholdResult, + result_schema_version=1, +) + + +class SteppingClock: + """A clock the test moves forward explicitly.""" + + def __init__(self) -> None: + self.now = START + + def __call__(self) -> datetime: + return self.now + + def advance(self, **elapsed: float) -> None: + self.now += timedelta(**elapsed) + + +class InMemoryJobStore: + """The JobStore contract in memory, with conditional transitions.""" + + def __init__(self, clock: SteppingClock) -> None: + self.clock = clock + self.jobs: dict[UUID, JobRecord] = {} + self.results: dict[UUID, JobResultRecord] = {} + self.lock = threading.Lock() + self.calls: list[str] = [] + + def _transition(self, job_id, expected, **changes) -> JobRecord | None: + job = self.jobs.get(job_id) + if job is None or job.status not in expected: + return None + updated = replace(job, updated_at=self.clock(), **changes) + self.jobs[job_id] = updated + return updated + + def submit(self, new_job): + self.calls.append("submit") + with self.lock: + for job in self.jobs.values(): + if ( + job.input_sha256 == new_job.input_sha256 + and job.status is not JobStatus.FAILED + ): + return Submission(job=job, created=False) + now = self.clock() + job = JobRecord( + id=uuid4(), + kind=new_job.kind, + status=JobStatus.PENDING, + submitted_by=new_job.submitted_by, + input_sha256=new_job.input_sha256, + policyengine_version=new_job.policyengine_version, + modal_application=new_job.modal_application, + modal_function=new_job.modal_function, + modal_call_id=None, + observability_id=new_job.observability_id, + error_code=None, + error_message=None, + collection_attempts=0, + created_at=now, + updated_at=now, + dispatched_at=None, + last_checked_at=None, + completed_at=None, + ) + self.jobs[job.id] = job + return Submission(job=job, created=True) + + def claim_dispatch(self, job_id): + self.calls.append("claim_dispatch") + with self.lock: + return ( + self._transition( + job_id, {JobStatus.PENDING}, status=JobStatus.DISPATCHING + ) + is not None + ) + + def record_dispatched(self, job_id, *, modal_call_id): + self.calls.append("record_dispatched") + with self.lock: + return ( + self._transition( + job_id, + {JobStatus.DISPATCHING}, + status=JobStatus.RUNNING, + modal_call_id=modal_call_id, + dispatched_at=self.clock(), + ) + is not None + ) + + def record_dispatch_failure(self, job_id, *, error_code, error_message): + self.calls.append("record_dispatch_failure") + with self.lock: + return ( + self._transition( + job_id, + {JobStatus.DISPATCHING}, + status=JobStatus.FAILED, + error_code=error_code, + error_message=error_message, + completed_at=self.clock(), + ) + is not None + ) + + def get(self, job_id): + self.calls.append("get") + return self.jobs.get(job_id) + + def get_by_modal_call_id(self, modal_call_id): + self.calls.append("get_by_modal_call_id") + return next( + (job for job in self.jobs.values() if job.modal_call_id == modal_call_id), + None, + ) + + def get_result(self, job_id): + self.calls.append("get_result") + return self.results.get(job_id) + + def record_check(self, job_id): + self.calls.append("record_check") + with self.lock: + job = self.jobs.get(job_id) + if job is None or job.status is not JobStatus.RUNNING: + return None + updated = self._transition( + job_id, + {JobStatus.RUNNING}, + collection_attempts=job.collection_attempts + 1, + last_checked_at=self.clock(), + ) + return updated.collection_attempts + + def complete(self, job_id, *, result, result_schema_version): + self.calls.append("complete") + with self.lock: + finished = self._transition( + job_id, + {JobStatus.RUNNING}, + status=JobStatus.SUCCEEDED, + completed_at=self.clock(), + ) + if finished is None: + return False + self.results[job_id] = JobResultRecord( + job_id=job_id, + result=result, + result_schema_version=result_schema_version, + result_sha256=result_sha256(result), + created_at=self.clock(), + ) + return True + + def fail(self, job_id, *, error_code, error_message): + self.calls.append("fail") + with self.lock: + return ( + self._transition( + job_id, + {JobStatus.RUNNING}, + status=JobStatus.FAILED, + error_code=error_code, + error_message=error_message, + completed_at=self.clock(), + ) + is not None + ) + + def jobs_to_collect(self, *, checked_before, limit): + self.calls.append("jobs_to_collect") + candidates = [ + job + for job in self.jobs.values() + if job.status is JobStatus.RUNNING + and (job.last_checked_at or job.dispatched_at or job.updated_at) + < checked_before + ] + return [job.id for job in candidates[:limit]] + + def abandon_stalled_dispatches(self, *, updated_before): + self.calls.append("abandon_stalled_dispatches") + with self.lock: + abandoned = [] + for job in list(self.jobs.values()): + if ( + job.status in UNFINISHED_STATUSES + and job.status is not JobStatus.RUNNING + and job.updated_at < updated_before + ): + self._transition( + job.id, + {job.status}, + status=JobStatus.FAILED, + error_code=DISPATCH_ABANDONED, + error_message="abandoned", + completed_at=self.clock(), + ) + abandoned.append(job.id) + return abandoned + + +class FakeDispatcher: + """Acknowledge each spawn with a new call ID, or fail like Modal.""" + + def __init__(self, failure: Exception | None = None) -> None: + self.failure = failure + self.spawned: list[tuple[ModalTarget, dict, dict | None]] = [] + + async def spawn(self, target, request_payload, *, observability_context): + self.spawned.append((target, request_payload, observability_context)) + if self.failure is not None: + raise self.failure + return f"fc-{len(self.spawned):026d}" + + +class FakeInspector: + """Return a scripted outcome for each Modal call, or raise.""" + + def __init__(self, outcome=None, failure: Exception | None = None) -> None: + self.outcome = outcome if outcome is not None else CallRunning() + self.failure = failure + self.inspected: list[str] = [] + + async def inspect(self, modal_call_id): + self.inspected.append(modal_call_id) + if self.failure is not None: + raise self.failure + return self.outcome + + +class FakeQueue: + """Record each collection task, or refuse like Cloud Tasks.""" + + def __init__(self, *, refuse: bool = False) -> None: + self.refuse = refuse + self.enqueued: list[tuple[UUID, float]] = [] + + async def enqueue(self, job_id, *, delay_seconds): + if self.refuse: + msg = "Cloud Tasks refused the task with HTTP 403" + raise CollectionQueueError(msg) + self.enqueued.append((job_id, delay_seconds)) diff --git a/projects/policyengine-simulation-entry/tests/test_config.py b/projects/policyengine-simulation-entry/tests/test_config.py index 3f28a1d80..5f1656887 100644 --- a/projects/policyengine-simulation-entry/tests/test_config.py +++ b/projects/policyengine-simulation-entry/tests/test_config.py @@ -201,3 +201,80 @@ def test_the_simulation_database_url_must_be_canonical(database_url): make_settings(simulation_database_url=database_url).validate() assert "secret" not in str(error.value) + + +RECONCILIATION = { + "simulation_database_url": "postgresql://simulation_runtime:secret@db/postgres", + "reconciliation_queue": ( + "projects/policyengine-staging/locations/us-central1/queues/simulation-jobs" + ), + "reconciliation_invoker_service_account": ( + "simulation-jobs-invoker@policyengine-staging.iam.gserviceaccount.com" + ), + "reconciliation_target_url": ( + "https://policyengine-simulation-entry-staging-abc-uc.a.run.app" + ), + "simulation_modal_environment": "staging", +} + + +def test_reconciliation_settings_are_read_from_the_environment(monkeypatch): + from policyengine_simulation_entry.config import Settings + + monkeypatch.setenv( + "SIMULATION_DATABASE_URL", RECONCILIATION["simulation_database_url"] + ) + monkeypatch.setenv("RECONCILIATION_QUEUE", RECONCILIATION["reconciliation_queue"]) + monkeypatch.setenv( + "RECONCILIATION_INVOKER_SERVICE_ACCOUNT", + RECONCILIATION["reconciliation_invoker_service_account"], + ) + monkeypatch.setenv( + "RECONCILIATION_TARGET_URL", RECONCILIATION["reconciliation_target_url"] + ) + monkeypatch.setenv("SIMULATION_MODAL_ENVIRONMENT", "staging") + + settings = Settings.from_env() + + assert {name: getattr(settings, name) for name in RECONCILIATION} == RECONCILIATION + assert settings.simulation_jobs_configured is True + make_settings(**RECONCILIATION).validate() + + +def test_simulation_jobs_are_off_without_reconciliation_settings(): + settings = make_settings( + simulation_database_url=RECONCILIATION["simulation_database_url"] + ) + + settings.validate() + assert settings.simulation_jobs_configured is False + + +@pytest.mark.parametrize( + "missing", + [ + "reconciliation_queue", + "reconciliation_invoker_service_account", + "reconciliation_target_url", + "simulation_modal_environment", + "simulation_database_url", + ], +) +def test_reconciliation_settings_are_all_or_nothing(missing): + with pytest.raises(ConfigurationError, match="Simulation jobs require"): + make_settings(**{**RECONCILIATION, missing: ""}).validate() + + +@pytest.mark.parametrize( + ("setting", "value"), + [ + ("reconciliation_queue", "simulation-jobs"), + ("reconciliation_queue", "projects/p/locations/us-central1/queues/q/tasks"), + ("reconciliation_invoker_service_account", "invoker@example.com"), + ("reconciliation_target_url", "http://entry.a.run.app"), + ("reconciliation_target_url", "https://entry.a.run.app/internal"), + ], +) +def test_reconciliation_settings_must_be_well_formed(setting, value): + with pytest.raises(ConfigurationError, match=setting.upper()): + make_settings(**{**RECONCILIATION, setting: value}).validate() diff --git a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py index 259c732f8..9328f422a 100644 --- a/projects/policyengine-simulation-entry/tests/test_deployment_assets.py +++ b/projects/policyengine-simulation-entry/tests/test_deployment_assets.py @@ -492,6 +492,15 @@ def _cloud_run_deploy_environment(tmp_path: Path) -> dict[str, str]: "SIMULATION_DATABASE_URL_SECRET_NAME": ( "simulation-staging-runtime-database-url" ), + "RECONCILIATION_QUEUE_VALUE": ( + "projects/policyengine-staging/locations/us-central1/queues/simulation-jobs" + ), + "RECONCILIATION_INVOKER_SERVICE_ACCOUNT_VALUE": ( + "simulation-jobs-invoker@policyengine-staging.iam.gserviceaccount.com" + ), + "ENTRYPOINT_STABLE_URL": ( + "https://policyengine-simulation-entry-staging-abc-uc.a.run.app" + ), } diff --git a/projects/policyengine-simulation-entry/tests/test_google_oidc.py b/projects/policyengine-simulation-entry/tests/test_google_oidc.py new file mode 100644 index 000000000..4298cdcda --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_google_oidc.py @@ -0,0 +1,94 @@ +"""Google OIDC verification for the Cloud Tasks and Cloud Scheduler routes.""" + +from __future__ import annotations + +import time +from types import SimpleNamespace + +import jwt +import pytest +from cryptography.hazmat.primitives.asymmetric import rsa +from policyengine_simulation_entry.google_oidc import ( + GoogleOidcVerifier, + InternalCallerRejectedError, +) + +AUDIENCE = "https://policyengine-simulation-entry-staging-abc-uc.a.run.app" +INVOKER = "simulation-jobs-invoker@policyengine-staging.iam.gserviceaccount.com" +SIGNING_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) +OTHER_KEY = rsa.generate_private_key(public_exponent=65537, key_size=2048) + + +class FakeGoogleKeys: + """Return the test signing key, as PyJWKClient returns Google's.""" + + def get_signing_key_from_jwt(self, _token): + return SimpleNamespace(key=SIGNING_KEY.public_key()) + + +def _token(key=SIGNING_KEY, algorithm="RS256", **overrides): + now = int(time.time()) + claims = { + "iss": "https://accounts.google.com", + "aud": AUDIENCE, + "email": INVOKER, + "email_verified": True, + "iat": now, + "exp": now + 300, + "sub": "1234567890", + } + claims.update(overrides) + claims = {name: value for name, value in claims.items() if value is not None} + signing_key = key if algorithm == "RS256" else "shared-secret-" + "x" * 32 + return jwt.encode(claims, signing_key, algorithm=algorithm) + + +def _verifier(): + return GoogleOidcVerifier( + audience=AUDIENCE, + service_account_email=INVOKER, + jwks_client=FakeGoogleKeys(), + ) + + +@pytest.mark.parametrize( + "issuer", ["https://accounts.google.com", "accounts.google.com"] +) +def test_a_token_from_the_invoker_account_is_accepted(issuer): + _verifier().verify(_token(iss=issuer)) + + +@pytest.mark.parametrize( + "token", + [ + _token(aud="https://another-service.a.run.app"), + _token(iss="https://evil.example"), + _token(email="someone-else@policyengine-staging.iam.gserviceaccount.com"), + _token(email_verified=False), + _token(email=None), + _token(exp=int(time.time()) - 60), + _token(key=OTHER_KEY), + _token(algorithm="HS256"), + "not-a-token", + ], + ids=[ + "audience", + "issuer", + "account", + "unverified-email", + "no-email", + "expired", + "wrong-key", + "symmetric", + "malformed", + ], +) +def test_every_other_token_is_rejected(token): + with pytest.raises(InternalCallerRejectedError): + _verifier().verify(token) + + +def test_the_verifier_reads_google_signing_keys_by_default(): + verifier = GoogleOidcVerifier(audience=AUDIENCE, service_account_email=INVOKER) + + assert verifier.jwks_client.uri == "https://www.googleapis.com/oauth2/v3/certs" diff --git a/projects/policyengine-simulation-entry/tests/test_simulation_job_adapters.py b/projects/policyengine-simulation-entry/tests/test_simulation_job_adapters.py new file mode 100644 index 000000000..7e28fc555 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_simulation_job_adapters.py @@ -0,0 +1,239 @@ +"""The Modal and Cloud Tasks adapters behind the simulation job service.""" + +from __future__ import annotations + +import asyncio +import json +from datetime import UTC, datetime +from types import SimpleNamespace +from uuid import UUID + +import httpx +import modal +import pytest +from policyengine_simulation_entry import simulation_jobs as jobs_module +from policyengine_simulation_entry.simulation_jobs import ( + CallFailed, + CallInspectionError, + CallRunning, + CallSucceeded, + CloudTasksCollectionQueue, + CollectionQueueError, + ModalCallDispatcher, + ModalCallInspector, + ModalTarget, +) + +JOB_ID = UUID("00000000-0000-4000-8000-0000000000aa") +QUEUE = "projects/policyengine-staging/locations/us-central1/queues/simulation-jobs" +TARGET_URL = "https://policyengine-simulation-entry-staging-abc-uc.a.run.app" +INVOKER = "simulation-jobs-invoker@policyengine-staging.iam.gserviceaccount.com" + + +# --- Modal call inspection -------------------------------------------------------- + + +class FakeFunctionCall: + def __init__(self, behaviour): + self.behaviour = behaviour + self.timeouts = [] + self.get = SimpleNamespace(aio=self._get) + + async def _get(self, timeout=None): + self.timeouts.append(timeout) + if isinstance(self.behaviour, BaseException): + raise self.behaviour + return self.behaviour + + +def _inspect(monkeypatch, behaviour): + calls = {} + + def from_id(call_id): + calls[call_id] = FakeFunctionCall(behaviour) + return calls[call_id] + + monkeypatch.setattr(modal.FunctionCall, "from_id", from_id) + outcome = asyncio.run(ModalCallInspector().inspect("fc-01ABC")) + assert calls["fc-01ABC"].timeouts == [0] + return outcome + + +def test_a_call_without_output_is_still_running(monkeypatch): + assert _inspect(monkeypatch, TimeoutError()) == CallRunning() + + +def test_a_finished_call_returns_its_output(monkeypatch): + assert _inspect(monkeypatch, {"net_income": 1.0}) == CallSucceeded( + {"net_income": 1.0} + ) + + +@pytest.mark.parametrize( + ("failure", "error_code"), + [ + (modal.exception.OutputExpiredError(), "result_expired"), + (modal.exception.FunctionTimeoutError(), "worker_timeout"), + (modal.exception.NotFoundError(), "call_not_found"), + (KeyError("household secret"), "worker_error"), + ], +) +def test_a_failed_call_fails_with_a_code_and_no_worker_data( + monkeypatch, failure, error_code +): + outcome = _inspect(monkeypatch, failure) + + assert isinstance(outcome, CallFailed) + assert outcome.error_code == error_code + assert "secret" not in outcome.error_message + + +@pytest.mark.parametrize( + "failure", + [modal.exception.ConnectionError(), modal.exception.AuthError()], +) +def test_modal_being_unreachable_is_transient(monkeypatch, failure): + with pytest.raises(CallInspectionError): + _inspect(monkeypatch, failure) + + +# --- Modal dispatch ---------------------------------------------------------------- + + +def test_the_dispatcher_spawns_the_target_function_with_observability_context( + monkeypatch, +): + looked_up = [] + spawned = [] + + async def spawn(payload, *, observability_context): + spawned.append((payload, observability_context)) + return SimpleNamespace(object_id="fc-01NEW") + + def from_name(application, function, *, environment_name): + looked_up.append((application, function, environment_name)) + return SimpleNamespace(spawn=SimpleNamespace(aio=spawn)) + + monkeypatch.setattr(modal.Function, "from_name", from_name) + target = ModalTarget(application="app-v2", function="calculate_household") + + call_id = asyncio.run( + ModalCallDispatcher("staging").spawn( + target, {"household": {}}, observability_context={"traceparent": "x"} + ) + ) + + assert call_id == "fc-01NEW" + assert looked_up == [("app-v2", "calculate_household", "staging")] + assert spawned == [({"household": {}}, {"traceparent": "x"})] + + +def test_the_dispatcher_requires_an_acknowledged_call_id(monkeypatch): + async def spawn(_payload, *, observability_context): + return SimpleNamespace(object_id="") + + monkeypatch.setattr( + modal.Function, + "from_name", + lambda *_args, **_options: SimpleNamespace(spawn=SimpleNamespace(aio=spawn)), + ) + + with pytest.raises(RuntimeError, match="no call ID"): + asyncio.run( + ModalCallDispatcher("staging").spawn( + ModalTarget("app", "function"), {}, observability_context=None + ) + ) + + +# --- Cloud Tasks --------------------------------------------------------------------- + + +def _queue(handler, clock=None): + async def access_token(): + return "access-token" + + return CloudTasksCollectionQueue( + queue=QUEUE, + target_url=TARGET_URL, + invoker_service_account=INVOKER, + access_token=access_token, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(handler)), + clock=clock or (lambda: datetime(2026, 10, 20, 12, tzinfo=UTC)), + ) + + +def test_a_collection_task_targets_the_stable_service_with_an_oidc_token(): + requests = [] + + def handler(request): + requests.append(request) + return httpx.Response(200, json={"name": f"{QUEUE}/tasks/1"}) + + asyncio.run(_queue(handler).enqueue(JOB_ID, delay_seconds=8)) + + request = requests[0] + assert str(request.url) == f"https://cloudtasks.googleapis.com/v2/{QUEUE}/tasks" + assert request.headers["Authorization"] == "Bearer access-token" + assert json.loads(request.content) == { + "task": { + "scheduleTime": "2026-10-20T12:00:08Z", + "httpRequest": { + "httpMethod": "POST", + "url": f"{TARGET_URL}/internal/simulation-jobs/{JOB_ID}/collect", + "oidcToken": { + "serviceAccountEmail": INVOKER, + "audience": TARGET_URL, + }, + }, + } + } + + +def test_an_immediate_task_has_no_schedule_time(): + bodies = [] + + def handler(request): + bodies.append(json.loads(request.content)) + return httpx.Response(200, json={}) + + asyncio.run(_queue(handler).enqueue(JOB_ID, delay_seconds=0)) + + assert "scheduleTime" not in bodies[0]["task"] + + +def test_a_refused_or_unreachable_queue_raises_a_queue_error(): + refused = _queue(lambda _request: httpx.Response(403, json={"error": {}})) + + def unreachable(request): + message = "no route" + raise httpx.ConnectError(message, request=request) + + with pytest.raises(CollectionQueueError, match="403"): + asyncio.run(refused.enqueue(JOB_ID, delay_seconds=0)) + with pytest.raises(CollectionQueueError): + asyncio.run(_queue(unreachable).enqueue(JOB_ID, delay_seconds=0)) + + +def test_default_credentials_supply_the_access_token(monkeypatch): + refreshed = [] + + class FakeCredentials: + valid = False + token = None + + def refresh(self, _request): + refreshed.append(True) + self.valid = True + self.token = "fresh-token" + + monkeypatch.setattr( + jobs_module.google.auth, + "default", + lambda scopes: (FakeCredentials(), "policyengine-staging"), + ) + tokens = jobs_module.GoogleAccessTokens() + + assert asyncio.run(tokens()) == "fresh-token" + assert asyncio.run(tokens()) == "fresh-token" + assert refreshed == [True] diff --git a/projects/policyengine-simulation-entry/tests/test_simulation_job_routes.py b/projects/policyengine-simulation-entry/tests/test_simulation_job_routes.py new file mode 100644 index 000000000..151aa2ae2 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_simulation_job_routes.py @@ -0,0 +1,263 @@ +"""HTTP routes of the simulation job store: status reads, collection and sweep. + +None of them is in the OpenAPI document, and no production caller uses them +yet. The internal routes accept only Google OIDC tokens of the configured +invoker service account (Cloud Tasks and Cloud Scheduler). +""" + +from __future__ import annotations + +import asyncio + +import pytest +from conftest import FakeBackend, make_settings +from fastapi.testclient import TestClient +from policyengine_simulation_entry.app import create_app +from policyengine_simulation_entry.google_oidc import InternalCallerRejectedError +from policyengine_simulation_entry.schemas import CallerIdentity +from policyengine_simulation_entry.simulation_jobs import ( + CallInspectionError, + CallRunning, + CallSucceeded, + SimulationJobService, +) +from policyengine_simulation_observability.identifiers import OBSERVABILITY_ID_HEADER +from simulation_job_support import ( + HOUSEHOLD_KIND, + TARGET, + FakeDispatcher, + FakeInspector, + FakeQueue, + InMemoryJobStore, + SteppingClock, +) + +VALID_TOKEN = "google-signed-token" +OBSERVABILITY_ID = "00000000-0000-4000-8000-000000000001" + + +class FakeVerifier: + """Accept only VALID_TOKEN, as the Google OIDC verifier would.""" + + def __init__(self): + self.tokens = [] + + def verify(self, token): + self.tokens.append(token) + if token != VALID_TOKEN: + msg = "InvalidAudienceError" + raise InternalCallerRejectedError(msg) + + +def _service(inspector=None, queue=None): + clock = SteppingClock() + store = InMemoryJobStore(clock) + service = SimulationJobService( + store=store, + dispatcher=FakeDispatcher(), + inspector=inspector or FakeInspector(CallRunning()), + queue=queue or FakeQueue(), + job_kinds={HOUSEHOLD_KIND.name: HOUSEHOLD_KIND}, + clock=clock, + ) + return service, store, clock + + +def _running_job(service): + return asyncio.run( + service.submit( + kind=HOUSEHOLD_KIND.name, + request_payload={"household": {}}, + policyengine_version="6.7.5", + submitted_by="api-v1@clients", + target=TARGET, + observability_id=OBSERVABILITY_ID, + ) + ) + + +def _client(service=None, verifier=None, **app_options): + return TestClient( + create_app( + settings=make_settings(), + backend=FakeBackend(), + auth_dependency=app_options.pop("auth_dependency", lambda: None), + simulation_jobs=service, + internal_caller_verifier=verifier, + **app_options, + ) + ) + + +def _bearer(token=VALID_TOKEN): + return {"Authorization": f"Bearer {token}"} + + +# --- Status reads ----------------------------------------------------------------- + + +def test_status_reads_return_the_stored_result_from_the_database(): + service, _, _ = _service( + inspector=FakeInspector(CallSucceeded({"net_income": 2.5})) + ) + job = _running_job(service) + asyncio.run(service.collect(job.id)) + service.inspector = FakeInspector(failure=AssertionError("polling called Modal")) + + with _client(service) as client: + by_id = client.get(f"/simulation-jobs/{job.id}") + by_call = client.get(f"/simulation-jobs/{job.modal_call_id}") + + assert by_id.status_code == 200 + assert by_id.json() == by_call.json() + body = by_id.json() + assert body["job_id"] == str(job.id) + assert body["status"] == "succeeded" + assert body["result"] == {"net_income": 2.5} + assert body["result_sha256"] + assert body["error"] is None + assert by_id.headers[OBSERVABILITY_ID_HEADER] == OBSERVABILITY_ID + + +def test_an_unfinished_job_reads_as_accepted_with_a_retry_hint(): + service, _, _ = _service() + job = _running_job(service) + + with _client(service) as client: + result = client.get(f"/simulation-jobs/{job.id}") + + assert result.status_code == 202 + assert result.headers["Retry-After"] == "5" + assert result.json()["status"] == "running" + assert result.json()["result"] is None + + +def test_an_unknown_job_reads_as_not_found(): + service, _, _ = _service() + + with _client(service) as client: + assert client.get("/simulation-jobs/fc-unknown").status_code == 404 + assert client.get("/simulation-jobs/not-an-id").status_code == 404 + + +def test_status_reads_only_log_a_missing_permission_by_default(): + service, _, _ = _service() + job = _running_job(service) + caller = CallerIdentity.model_validate( + {"sub": "user", "iss": "https://issuer.example/", "aud": "x", "exp": 1} + ) + + with _client(service, auth_dependency=lambda: caller) as client: + allowed_while_logging = client.get(f"/simulation-jobs/{job.id}") + + assert allowed_while_logging.status_code == 202 + + +def test_status_reads_reject_callers_without_the_permission_when_enforced(): + service, _, _ = _service() + job = _running_job(service) + caller = CallerIdentity.model_validate( + {"sub": "user", "iss": "https://issuer.example/", "aud": "x", "exp": 1} + ) + app = create_app( + settings=make_settings(enforce_permissions=True), + backend=FakeBackend(), + auth_dependency=lambda: caller, + simulation_jobs=service, + ) + + with TestClient(app) as client: + assert client.get(f"/simulation-jobs/{job.id}").status_code == 403 + + +def test_job_routes_are_unavailable_without_the_job_store(): + with _client() as client: + assert client.get("/simulation-jobs/fc-anything").status_code == 503 + assert ( + client.post( + "/internal/simulation-jobs/sweep", headers=_bearer() + ).status_code + == 503 + ) + + +# --- Collection and the sweep -------------------------------------------------------- + + +@pytest.mark.parametrize( + "headers", + [{}, _bearer("forged"), {"Authorization": "Basic abc"}], + ids=["missing", "rejected", "not-bearer"], +) +def test_internal_routes_reject_callers_without_a_valid_google_token(headers): + service, _, _ = _service() + job = _running_job(service) + + with _client(service, FakeVerifier()) as client: + collect = client.post( + f"/internal/simulation-jobs/{job.id}/collect", headers=headers + ) + sweep = client.post("/internal/simulation-jobs/sweep", headers=headers) + + assert collect.status_code == 403 + assert sweep.status_code == 403 + + +def test_the_collection_route_reports_what_collection_did(): + service, store, _ = _service( + inspector=FakeInspector(CallSucceeded({"net_income": 3.0})) + ) + job = _running_job(service) + + with _client(service, FakeVerifier()) as client: + result = client.post( + f"/internal/simulation-jobs/{job.id}/collect", headers=_bearer() + ) + + assert result.status_code == 200 + assert result.json() == { + "outcome": "succeeded", + "job_status": "succeeded", + "next_check_seconds": None, + } + assert store.get_result(job.id).result == {"net_income": 3.0} + + +def test_the_collection_route_asks_cloud_tasks_to_retry_transient_failures(): + service, _, _ = _service( + inspector=FakeInspector(failure=CallInspectionError("Modal unavailable")) + ) + job = _running_job(service) + + with _client(service, FakeVerifier()) as client: + result = client.post( + f"/internal/simulation-jobs/{job.id}/collect", headers=_bearer() + ) + + assert result.status_code == 503 + assert result.headers["Retry-After"] + + +def test_the_sweep_route_reports_its_counts_and_fails_when_tasks_were_refused(): + service, _, clock = _service() + _running_job(service) + clock.advance(minutes=6) + + with _client(service, FakeVerifier()) as client: + swept = client.post("/internal/simulation-jobs/sweep", headers=_bearer()) + service.queue = FakeQueue(refuse=True) + clock.advance(minutes=6) + refused = client.post("/internal/simulation-jobs/sweep", headers=_bearer()) + + assert swept.status_code == 200 + assert swept.json() == {"abandoned": 0, "enqueued": 1, "enqueue_failures": 0} + assert refused.status_code == 503 + assert refused.json()["enqueue_failures"] == 1 + + +def test_job_routes_stay_out_of_the_openapi_document(): + paths = create_app(settings=make_settings(), backend=FakeBackend()).openapi()[ + "paths" + ] + + assert not [path for path in paths if "simulation-jobs" in path] diff --git a/projects/policyengine-simulation-entry/tests/test_simulation_jobs.py b/projects/policyengine-simulation-entry/tests/test_simulation_jobs.py new file mode 100644 index 000000000..f09768ba1 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_simulation_jobs.py @@ -0,0 +1,404 @@ +"""The simulation job service: submission, background collection and the sweep.""" + +from __future__ import annotations + +import asyncio + +import pytest +from policyengine_simulation_entry.simulation_jobs import ( + COLLECTION_STALE_AFTER, + DISPATCH_TIMEOUT, + FIRST_COLLECTION_DELAY_SECONDS, + CallFailed, + CallInspectionError, + CallRunning, + CallSucceeded, + CollectionOutcome, + CollectionQueueError, + SimulationJobService, + UnknownJobKindError, + collection_delay_seconds, +) +from policyengine_simulation_persistence.jobs import JobStatus, input_sha256 +from simulation_job_support import ( + HOUSEHOLD_KIND, + TARGET, + FakeDispatcher, + FakeInspector, + FakeQueue, + InMemoryJobStore, + SteppingClock, +) + +REQUEST = {"household": {"people": {"you": {"age": 40}}}} +OBSERVABILITY_CONTEXT = { + "traceparent": "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" +} + + +def _service( + *, + dispatcher=None, + inspector=None, + queue=None, + clock=None, +): + clock = clock or SteppingClock() + store = InMemoryJobStore(clock) + service = SimulationJobService( + store=store, + dispatcher=dispatcher or FakeDispatcher(), + inspector=inspector or FakeInspector(), + queue=queue or FakeQueue(), + job_kinds={HOUSEHOLD_KIND.name: HOUSEHOLD_KIND}, + clock=clock, + ) + return service, store + + +async def _submit(service, request=REQUEST, submitted_by="api-v1@clients"): + return await service.submit( + kind=HOUSEHOLD_KIND.name, + request_payload=request, + policyengine_version="6.7.5", + submitted_by=submitted_by, + target=TARGET, + observability_id="00000000-0000-4000-8000-000000000001", + observability_context=OBSERVABILITY_CONTEXT, + ) + + +# --- Submission ---------------------------------------------------------------- + + +def test_a_submission_records_the_job_dispatches_it_and_schedules_collection(): + dispatcher = FakeDispatcher() + queue = FakeQueue() + service, store = _service(dispatcher=dispatcher, queue=queue) + + job = asyncio.run(_submit(service)) + + assert job.status is JobStatus.RUNNING + assert job.modal_call_id == "fc-" + "1".zfill(26) + assert job.submitted_by == "api-v1@clients" + assert job.input_sha256 == input_sha256( + kind=HOUSEHOLD_KIND.name, policyengine_version="6.7.5", request=REQUEST + ) + assert (job.modal_application, job.modal_function) == ( + TARGET.application, + TARGET.function, + ) + assert dispatcher.spawned == [(TARGET, REQUEST, OBSERVABILITY_CONTEXT)] + assert queue.enqueued == [(job.id, FIRST_COLLECTION_DELAY_SECONDS)] + assert store.calls[:3] == ["submit", "claim_dispatch", "record_dispatched"] + + +def test_an_identical_submission_reuses_the_job_without_a_second_modal_call(): + dispatcher = FakeDispatcher() + service, _ = _service(dispatcher=dispatcher) + + first = asyncio.run(_submit(service)) + second = asyncio.run(_submit(service, submitted_by="another@clients")) + + assert second.id == first.id + assert len(dispatcher.spawned) == 1 + + +def test_concurrent_identical_submissions_start_one_modal_call(): + dispatcher = FakeDispatcher() + service, _ = _service(dispatcher=dispatcher) + + async def submit_together(): + return await asyncio.gather(*(_submit(service) for _ in range(5))) + + submitted = asyncio.run(submit_together()) + + assert len({job.id for job in submitted}) == 1 + assert len(dispatcher.spawned) == 1 + + +def test_a_refused_modal_call_fails_the_job_without_raising(): + dispatcher = FakeDispatcher(failure=ConnectionError("password=secret")) + queue = FakeQueue() + service, _ = _service(dispatcher=dispatcher, queue=queue) + + job = asyncio.run(_submit(service)) + + assert job.status is JobStatus.FAILED + assert job.error_code == "dispatch_failed" + assert "secret" not in job.error_message + assert queue.enqueued == [] + + +def test_a_refused_collection_task_leaves_the_job_for_the_sweep(): + service, _ = _service(queue=FakeQueue(refuse=True)) + + job = asyncio.run(_submit(service)) + + assert job.status is JobStatus.RUNNING + + +def test_an_unregistered_kind_is_rejected_before_anything_is_written(): + service, store = _service() + + with pytest.raises(UnknownJobKindError): + asyncio.run( + service.submit( + kind="unregistered", + request_payload=REQUEST, + policyengine_version="6.7.5", + submitted_by="api-v1@clients", + target=TARGET, + ) + ) + + assert store.jobs == {} + + +# --- Collection ---------------------------------------------------------------- + + +def _running_job(service): + return asyncio.run(_submit(service)) + + +def test_a_running_call_is_counted_and_checked_again_with_backoff(): + inspector = FakeInspector(CallRunning()) + queue = FakeQueue() + service, store = _service(inspector=inspector, queue=queue) + job = _running_job(service) + queue.enqueued.clear() + + first = asyncio.run(service.collect(job.id)) + second = asyncio.run(service.collect(job.id)) + + assert first.outcome is CollectionOutcome.STILL_RUNNING + assert (first.next_check_seconds, second.next_check_seconds) == (2, 4) + assert queue.enqueued == [(job.id, 2), (job.id, 4)] + assert store.get(job.id).collection_attempts == 2 + assert inspector.inspected == [job.modal_call_id, job.modal_call_id] + + +def test_the_backoff_doubles_up_to_a_minute(): + assert [collection_delay_seconds(attempt) for attempt in range(1, 9)] == [ + 2, + 4, + 8, + 16, + 32, + 60, + 60, + 60, + ] + + +def test_a_finished_call_stores_its_validated_result(): + inspector = FakeInspector(CallSucceeded({"net_income": 51234.5})) + service, store = _service(inspector=inspector) + job = _running_job(service) + + report = asyncio.run(service.collect(job.id)) + + assert report.outcome is CollectionOutcome.SUCCEEDED + assert store.get(job.id).status is JobStatus.SUCCEEDED + assert store.get_result(job.id).result == {"net_income": 51234.5} + assert store.get_result(job.id).result_schema_version == 1 + + +def test_a_result_that_fails_validation_fails_the_job(): + inspector = FakeInspector(CallSucceeded({"net_income": "not a number"})) + service, store = _service(inspector=inspector) + job = _running_job(service) + + report = asyncio.run(service.collect(job.id)) + + assert report.outcome is CollectionOutcome.FAILED + assert store.get(job.id).error_code == "invalid_result" + assert "not a number" not in store.get(job.id).error_message + assert store.get_result(job.id) is None + + +def test_a_failed_call_fails_the_job(): + inspector = FakeInspector(CallFailed("worker_error", "The call raised KeyError.")) + service, store = _service(inspector=inspector) + job = _running_job(service) + + report = asyncio.run(service.collect(job.id)) + + assert report.outcome is CollectionOutcome.FAILED + assert store.get(job.id).error_code == "worker_error" + + +def test_two_collectors_store_one_result_and_the_second_changes_nothing(): + service, store = _service() + job = _running_job(service) + + class BothCollectorsSeeTheCallFinish: + """Hold each collector until both have read the job as running.""" + + def __init__(self): + self.barrier = asyncio.Barrier(2) + + async def inspect(self, _modal_call_id): + await self.barrier.wait() + return CallSucceeded({"net_income": 1.0}) + + async def collect_together(): + service.inspector = BothCollectorsSeeTheCallFinish() + return await asyncio.gather(service.collect(job.id), service.collect(job.id)) + + reports = asyncio.run(collect_together()) + + assert sorted(report.outcome.value for report in reports) == sorted( + [CollectionOutcome.SUCCEEDED.value, CollectionOutcome.ALREADY_FINISHED.value] + ) + assert store.calls.count("complete") == 2 + assert list(store.results) == [job.id] + assert store.get_result(job.id).result == {"net_income": 1.0} + + +def test_a_finished_job_is_never_inspected_again(): + inspector = FakeInspector(CallSucceeded({"net_income": 1.0})) + service, _ = _service(inspector=inspector) + job = _running_job(service) + asyncio.run(service.collect(job.id)) + inspector.inspected.clear() + + report = asyncio.run(service.collect(job.id)) + + assert report.outcome is CollectionOutcome.NOT_RUNNING + assert report.job_status is JobStatus.SUCCEEDED + assert inspector.inspected == [] + + +def test_an_unknown_job_is_dropped(): + from uuid import uuid4 + + service, _ = _service() + + report = asyncio.run(service.collect(uuid4())) + + assert report.outcome is CollectionOutcome.UNKNOWN_JOB + + +def test_a_job_of_an_unregistered_kind_fails_on_collection(): + inspector = FakeInspector(CallSucceeded({"net_income": 1.0})) + service, store = _service(inspector=inspector) + job = _running_job(service) + service.job_kinds = {} + + report = asyncio.run(service.collect(job.id)) + + assert report.outcome is CollectionOutcome.FAILED + assert store.get(job.id).error_code == "unknown_job_kind" + + +@pytest.mark.parametrize( + "failure", + [CallInspectionError("Modal unavailable"), CollectionQueueError("refused")], +) +def test_transient_failures_propagate_so_cloud_tasks_retries(failure): + if isinstance(failure, CallInspectionError): + service, store = _service(inspector=FakeInspector(failure=failure)) + else: + service, store = _service(queue=FakeQueue()) + job = _running_job(service) + if isinstance(failure, CollectionQueueError): + service.queue = FakeQueue(refuse=True) + + with pytest.raises(type(failure)): + asyncio.run(service.collect(job.id)) + + assert store.get(job.id).status is JobStatus.RUNNING + + +# --- The sweep ------------------------------------------------------------------- + + +def test_the_sweep_enqueues_stale_running_jobs_and_fails_stalled_dispatches(): + clock = SteppingClock() + queue = FakeQueue() + service, store = _service(clock=clock, queue=queue) + stale = _running_job(service) + stalled = asyncio.run( + asyncio.to_thread(store.submit, _new_job_like(stale, {"case": "stalled"})) + ).job + clock.advance(seconds=DISPATCH_TIMEOUT.total_seconds() + 1) + fresh = asyncio.run(_submit(service, request={"case": "fresh"})) + queue.enqueued.clear() + + report = asyncio.run(service.sweep()) + + assert report.abandoned == 1 + assert store.get(stalled.id).status is JobStatus.FAILED + assert report.enqueued == 1 + assert queue.enqueued == [(stale.id, 0)] + assert fresh.id not in [job_id for job_id, _ in queue.enqueued] + assert COLLECTION_STALE_AFTER < DISPATCH_TIMEOUT + + +def test_the_sweep_reports_refused_tasks(): + clock = SteppingClock() + service, _ = _service(clock=clock) + _running_job(service) + clock.advance(seconds=COLLECTION_STALE_AFTER.total_seconds() + 1) + service.queue = FakeQueue(refuse=True) + + report = asyncio.run(service.sweep()) + + assert (report.enqueued, report.enqueue_failures) == (0, 1) + + +def _new_job_like(job, request): + from policyengine_simulation_persistence.jobs import NewJob + + return NewJob( + kind=job.kind, + submitted_by=job.submitted_by, + input_sha256=input_sha256( + kind=job.kind, policyengine_version="6.7.5", request=request + ), + policyengine_version="6.7.5", + modal_application=job.modal_application, + modal_function=job.modal_function, + ) + + +# --- Status reads ---------------------------------------------------------------- + + +def test_status_reads_accept_a_job_id_or_a_modal_call_id_and_never_call_modal(): + inspector = FakeInspector(CallSucceeded({"net_income": 1.0})) + dispatcher = FakeDispatcher() + service, _ = _service(inspector=inspector, dispatcher=dispatcher) + job = _running_job(service) + asyncio.run(service.collect(job.id)) + inspector.inspected.clear() + + by_id = asyncio.run(service.find(str(job.id))) + by_call = asyncio.run(service.find(job.modal_call_id)) + + assert by_id.job.id == by_call.job.id == job.id + assert by_id.result.result == {"net_income": 1.0} + assert inspector.inspected == [] + assert len(dispatcher.spawned) == 1 + + +@pytest.mark.parametrize( + "reference", + ["not-a-job", "fc-", "fc-has spaces", "00000000-0000-4000-8000-000000000099"], +) +def test_unknown_or_malformed_references_find_nothing(reference): + service, _ = _service() + + assert asyncio.run(service.find(reference)) is None + + +def test_an_unfinished_job_reads_without_a_result(): + service, _ = _service() + job = _running_job(service) + + view = asyncio.run(service.find(str(job.id))) + + assert view.job.status is JobStatus.RUNNING + assert view.result is None diff --git a/projects/policyengine-simulation-entry/tests/test_simulation_jobs_deployment.py b/projects/policyengine-simulation-entry/tests/test_simulation_jobs_deployment.py new file mode 100644 index 000000000..3ca5e1fc5 --- /dev/null +++ b/projects/policyengine-simulation-entry/tests/test_simulation_jobs_deployment.py @@ -0,0 +1,106 @@ +"""Deployment wiring for background collection of simulation job results. + +The queue and the invoker service account come from GitHub environment +variables; the collection target is the stable service URL the deploy job +reads before deploying; the Modal environment is the deployment's own. +""" + +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest +from test_deployment_assets import ( + CLOUD_RUN_DEPLOY_SCRIPT, + REUSABLE_DEPLOY_WORKFLOW, + _cloud_run_deploy_environment, + _workflow_jobs, +) + +QUEUE = "projects/policyengine-staging/locations/us-central1/queues/simulation-jobs" +INVOKER = "simulation-jobs-invoker@policyengine-staging.iam.gserviceaccount.com" +STABLE_URL = "https://policyengine-simulation-entry-staging-abc-uc.a.run.app" + + +def _run_deploy(tmp_path, **overrides): + environment = _cloud_run_deploy_environment(tmp_path) | overrides + result = subprocess.run( + [CLOUD_RUN_DEPLOY_SCRIPT], + check=False, + capture_output=True, + text=True, + env=environment, + ) + return result, environment + + +def test_the_entrypoint_receives_the_reconciliation_settings(tmp_path): + result, environment = _run_deploy(tmp_path) + + assert result.returncode == 0, result.stderr + runtime_environment = json.loads( + Path(environment["GCLOUD_ENV_FILE_CAPTURE"]).read_text(encoding="utf-8") + ) + assert runtime_environment["RECONCILIATION_QUEUE"] == QUEUE + assert runtime_environment["RECONCILIATION_INVOKER_SERVICE_ACCOUNT"] == INVOKER + assert runtime_environment["RECONCILIATION_TARGET_URL"] == STABLE_URL + assert runtime_environment["SIMULATION_MODAL_ENVIRONMENT"] == "staging" + + +@pytest.mark.parametrize( + "variable", + [ + "RECONCILIATION_QUEUE_VALUE", + "RECONCILIATION_INVOKER_SERVICE_ACCOUNT_VALUE", + "ENTRYPOINT_STABLE_URL", + ], +) +def test_the_deploy_requires_every_reconciliation_setting(tmp_path, variable): + result, _ = _run_deploy(tmp_path, **{variable: ""}) + + assert result.returncode != 0 + assert variable in result.stderr + + +@pytest.mark.parametrize( + ("variable", "value"), + [ + ("RECONCILIATION_QUEUE_VALUE", "simulation-jobs"), + ("RECONCILIATION_QUEUE_VALUE", f"{QUEUE}/tasks"), + ("RECONCILIATION_INVOKER_SERVICE_ACCOUNT_VALUE", "invoker@example.com"), + ("ENTRYPOINT_STABLE_URL", "http://entry.a.run.app"), + ("ENTRYPOINT_STABLE_URL", f"{STABLE_URL}/path"), + ], +) +def test_the_deploy_rejects_malformed_reconciliation_settings( + tmp_path, variable, value +): + result, _ = _run_deploy(tmp_path, **{variable: value}) + + assert result.returncode != 0 + assert variable in result.stderr + + +def test_the_deploy_step_passes_the_settings_and_the_stable_url(): + deploy_entrypoint = _workflow_jobs(REUSABLE_DEPLOY_WORKFLOW)["deploy_entrypoint"] + steps = [step.get("name") for step in deploy_entrypoint["steps"]] + deploy_step = next( + step + for step in deploy_entrypoint["steps"] + if step.get("name") == "Deploy Cloud Run entrypoint candidate" + ) + + assert deploy_step["env"]["RECONCILIATION_QUEUE_VALUE"] == ( + "${{ vars.RECONCILIATION_QUEUE }}" + ) + assert deploy_step["env"]["RECONCILIATION_INVOKER_SERVICE_ACCOUNT_VALUE"] == ( + "${{ vars.RECONCILIATION_INVOKER_SERVICE_ACCOUNT }}" + ) + assert deploy_step["env"]["ENTRYPOINT_STABLE_URL"] == ( + "${{ steps.previous.outputs.stable_url }}" + ) + assert steps.index("Capture stable service state before deployment") < ( + steps.index("Deploy Cloud Run entrypoint candidate") + ) diff --git a/projects/policyengine-simulation-entry/uv.lock b/projects/policyengine-simulation-entry/uv.lock index a5971d22f..89243bb07 100644 --- a/projects/policyengine-simulation-entry/uv.lock +++ b/projects/policyengine-simulation-entry/uv.lock @@ -347,6 +347,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/59/13/477d90d09591b3938b45c4e11f4d8a51291682112cb5efcac961e815d562/google_auth-2.58.0-py3-none-any.whl", hash = "sha256:8a9c4645bb4c8e91668fb1934b95ae6a8687084232753639220ba9bf04a1610d", size = 262404, upload-time = "2026-09-09T20:49:33.951Z" }, ] +[package.optional-dependencies] +requests = [ + { name = "requests" }, +] + [[package]] name = "google-cloud-appengine-logging" version = "1.10.0" @@ -896,6 +901,7 @@ source = { editable = "." } dependencies = [ { name = "cryptography" }, { name = "fastapi" }, + { name = "google-auth", extra = ["requests"] }, { name = "httpx" }, { name = "modal" }, { name = "policyengine-observability", extra = ["fastapi", "google", "httpx", "otlp-grpc"] }, @@ -923,6 +929,7 @@ test = [ requires-dist = [ { name = "cryptography", specifier = ">=41.0.0" }, { name = "fastapi", specifier = ">=0.115.0,<1" }, + { name = "google-auth", extras = ["requests"], specifier = ">=2.30,<3" }, { name = "httpx", specifier = ">=0.28.0,<1" }, { name = "modal", specifier = ">=1.4,<2" }, { name = "openapi-python-client", marker = "extra == 'build'", specifier = ">=0.21.6" }, diff --git a/rules/registry.toml b/rules/registry.toml index b64ca0287..8cf17bcef 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -218,6 +218,72 @@ workflow = ".github/workflows/pr.yml" job = "test" invocation = "uv run pytest tests/" +[[rule]] +id = "sim.persistence.conditional-job-writes" +category = "persistence" +statement = "Every simulation job transition and result write is one conditional update on the job's expected status (and submission an insert that does nothing on a conflicting input hash), so concurrent collectors, dispatchers or submitters change a job at most once; no lease is taken." +rationale = "libs/policyengine-simulation-persistence/README.md#the-job-store" +scope = ["libs/policyengine-simulation-persistence/src/policyengine_simulation_persistence/jobs.py"] +tier = "mechanical" +status = "enforced" +autofix = false +[rule.enforcer] +kind = "pytest" +project = "libs/policyengine-simulation-persistence" +nodes = [ + "tests/test_job_store_concurrency.py::test_a_second_collector_waits_for_the_first_then_changes_nothing", + "tests/test_job_store_concurrency.py::test_a_failure_racing_a_result_cannot_overwrite_it", + "tests/test_job_store_concurrency.py::test_a_second_dispatcher_waits_for_the_first_then_claims_nothing", + "tests/test_job_store_concurrency.py::test_concurrent_collectors_write_exactly_one_result", + "tests/test_job_store_concurrency.py::test_concurrent_identical_submissions_create_one_job", +] +[[rule.ci]] +workflow = ".github/workflows/pr.yml" +job = "simulation-database" +invocation = "uv run pytest tests/ -v -m postgres" + +[[rule]] +id = "sim.persistence.status-reads-never-advance-work" +category = "persistence" +statement = "Simulation job status reads use only the Simulation Entrypoint's database: they never inspect, spawn or advance a Modal call; results are collected by Cloud Tasks and the Cloud Scheduler sweep." +rationale = "libs/policyengine-simulation-persistence/README.md#the-job-store" +scope = ["projects/policyengine-simulation-entry/src/policyengine_simulation_entry/simulation_jobs.py", "projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py"] +tier = "mechanical" +status = "enforced" +autofix = false +[rule.enforcer] +kind = "pytest" +project = "projects/policyengine-simulation-entry" +nodes = [ + "tests/test_simulation_jobs.py::test_status_reads_accept_a_job_id_or_a_modal_call_id_and_never_call_modal", + "tests/test_simulation_job_routes.py::test_status_reads_return_the_stored_result_from_the_database", +] +[[rule.ci]] +workflow = ".github/workflows/pr.yml" +job = "test" +invocation = "uv run pytest tests/" + +[[rule]] +id = "sim.persistence.internal-routes-google-oidc" +category = "persistence" +statement = "The simulation job collection and sweep routes accept only Google-signed OIDC tokens whose audience is the service URL and whose email is the configured invoker service account." +rationale = "projects/policyengine-simulation-entry/DEPLOYMENT.md" +scope = ["projects/policyengine-simulation-entry/src/policyengine_simulation_entry/google_oidc.py", "projects/policyengine-simulation-entry/src/policyengine_simulation_entry/app.py"] +tier = "mechanical" +status = "enforced" +autofix = false +[rule.enforcer] +kind = "pytest" +project = "projects/policyengine-simulation-entry" +nodes = [ + "tests/test_google_oidc.py::test_every_other_token_is_rejected", + "tests/test_simulation_job_routes.py::test_internal_routes_reject_callers_without_a_valid_google_token", +] +[[rule.ci]] +workflow = ".github/workflows/pr.yml" +job = "test" +invocation = "uv run pytest tests/" + # --- Stage 12 persistence -------------------------------------------------- [[rule]] From 09f22547d092f9dc7e5f53b1a8635799893e1ba5 Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:13:12 +0400 Subject: [PATCH 2/2] Import the job tables in the migration environment migrations/env.py compared the database against simulation_metadata without importing the module that defines simulation.jobs and simulation.job_results, so in a standalone Alembic process the metadata held no tables. "alembic check" then reported both tables as removed, and "alembic revision --autogenerate" would have written a revision that drops them. The CI migration-cycle test passed only because other test modules imported the job store during collection. The environment now imports the job store module. A new test runs "alembic check" in a separate process after upgrading to the head; it failed with the removed-table differences before this change. The test is added to rule sim.persistence.migration-cycle, whose scope now includes the environment, and the README's revision steps say that new table modules must be imported there. Co-Authored-By: Claude Opus 5.5 --- .../README.md | 7 +++-- .../tests/test_migration_cycle.py | 27 ++++++++++++++++++- migrations/env.py | 4 +++ rules/registry.toml | 3 ++- 4 files changed, 37 insertions(+), 4 deletions(-) diff --git a/libs/policyengine-simulation-persistence/README.md b/libs/policyengine-simulation-persistence/README.md index 78daac6d0..7df36c97e 100644 --- a/libs/policyengine-simulation-persistence/README.md +++ b/libs/policyengine-simulation-persistence/README.md @@ -90,11 +90,14 @@ reads them. ## Adding a revision -1. Change the tables in `src/policyengine_simulation_persistence/`. +1. Change the tables in `src/policyengine_simulation_persistence/`. A new + module that defines tables must also be imported by `migrations/env.py`, + because the schema comparison sees only tables whose module was imported. 2. Add a revision under `migrations/versions/` whose `down_revision` is the current head, and set `HEAD_REVISION` in `schema.py` to its revision ID. 3. Run the tests against a disposable PostgreSQL server (below). They run the - chain's upgrade, `alembic check`, downgrade and upgrade cycle. + chain's upgrade, `alembic check`, downgrade and upgrade cycle, and run + `alembic check` again in a separate process, as the command line does. Only one migration PR in this chain may be open at a time; a second one rebases onto the new head after the first merges. diff --git a/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py b/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py index 9a3f9931b..bad6db0c7 100644 --- a/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py +++ b/libs/policyengine-simulation-persistence/tests/test_migration_cycle.py @@ -2,9 +2,12 @@ from __future__ import annotations +import subprocess +import sys + import pytest from alembic import command -from migration_support import MIGRATOR_ROLE, alembic_config +from migration_support import ALEMBIC_CONFIG_PATH, MIGRATOR_ROLE, alembic_config from policyengine_simulation_persistence import ( HEAD_REVISION, SIMULATION_SCHEMA, @@ -41,6 +44,28 @@ def test_upgrade_check_downgrade_upgrade_cycle(disposable_database) -> None: command.check(config) +def test_alembic_check_run_on_its_own_finds_no_differences( + disposable_database, +) -> None: + """Run ``alembic check`` in a new process, as the command line does. + + Inside the test session other test modules import the table definitions, + so an in-process check passes even when the migration environment does not + import them. A new process sees only what the environment imports. + """ + command.upgrade(alembic_config(), "head") + + check = subprocess.run( + [sys.executable, "-m", "alembic", "-c", str(ALEMBIC_CONFIG_PATH), "check"], + capture_output=True, + text=True, + check=False, + timeout=120, + ) + + assert check.returncode == 0, check.stdout + check.stderr + + def test_the_version_table_lives_only_in_the_simulation_schema( disposable_database, ) -> None: diff --git a/migrations/env.py b/migrations/env.py index a8b00eab5..72576d470 100644 --- a/migrations/env.py +++ b/migrations/env.py @@ -6,6 +6,10 @@ from logging.config import fileConfig from typing import TYPE_CHECKING +# Importing a table module registers its tables in simulation_metadata, which +# the schema comparison (alembic check and autogenerate) reads. Import every +# module that defines tables, or the comparison reports them as removed. +import policyengine_simulation_persistence.jobs # noqa: F401 # registers the job tables from alembic import context from policyengine_simulation_persistence import ( ALEMBIC_VERSION_TABLE, diff --git a/rules/registry.toml b/rules/registry.toml index 8cf17bcef..5675bddfb 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -132,7 +132,7 @@ id = "sim.persistence.migration-cycle" category = "persistence" statement = "On a fresh PostgreSQL, the simulation chain upgrades to its head, passes alembic check, downgrades to base and upgrades again, run by a non-superuser migrator, and the runtime role can read the revision but cannot change the schema." rationale = "libs/policyengine-simulation-persistence/README.md#tests" -scope = ["migrations/versions/*.py", "libs/policyengine-simulation-persistence/src/**/*.py"] +scope = ["migrations/env.py", "migrations/versions/*.py", "libs/policyengine-simulation-persistence/src/**/*.py"] tier = "mechanical" status = "enforced" autofix = false @@ -141,6 +141,7 @@ kind = "pytest" project = "libs/policyengine-simulation-persistence" nodes = [ "tests/test_migration_cycle.py::test_upgrade_check_downgrade_upgrade_cycle", + "tests/test_migration_cycle.py::test_alembic_check_run_on_its_own_finds_no_differences", "tests/test_migration_cycle.py::test_the_version_table_lives_only_in_the_simulation_schema", "tests/test_migration_cycle.py::test_the_runtime_role_reads_the_revision_but_cannot_change_the_schema", ]