Release Stage 12 database sessions and keep original failures visible

[anth-volk]

Fixes #788

Same-Repository Draft PR Check

• This PR is a draft.
• This PR is from a branch in PolicyEngine/policyengine-sim-api, not a personal fork.

Summary

Stage 12 processes held pooled database sessions while they calculated or waited, so concurrent runs exhausted the session pooler's 15-client limit. From 6 to 13 October 2026, this caused 106 of the 153 production Stage 12 failures.

• Release sessions after each operation: create_stage12_engine now uses NullPool, so each physical connection closes when its operation finishes, including after a rollback.
• Retry only the step that opens a connection:
  • Each Stage12PersistenceStore operation now opens its connection in one place, through a ConnectionOpener (persistence library). The opener receives only the connect step. The operation's statements, and the COMMIT or ROLLBACK that ends it, run after the opener returns, so they are never retried.
  • The retry is decided by where the failure happened, not by the error. SQLAlchemy raises a failed COMMIT, ROLLBACK or BEGIN the same way as a refused connection: an OperationalError with no statement. A COMMIT can fail after the database has applied it. An earlier version of this PR retried on that error, so a coordinator whose claim committed but whose reply was lost would re-run its claim, find its own running row, and return a deduplicated success without dispatching any simulation.
  • The coordinator and worker build their store with open_connection_with_retry (stage12_runtime/database_retry.py). It backs off from 1 s, doubling up to 15 s with jitter, within a 60 s budget for each connection. Each retry is logged at WARNING with the store operation's name.
  • This covers every coordinator and worker store operation, not only the claims: child-state creation, the dispatch placeholder, attaching the Modal invocation, the result comparison and the failure records. With NullPool each of these opens its own pooler session, so each could otherwise fail a report on a single refused connection. On main they reused the coordinator's pooled connection.
  • Failed statements, failed commits and validation errors still fail immediately.
• Keep the original failure visible in telemetry:
  • The coordinator and worker failure handlers now report the original error to telemetry first. They then attempt the failure-record write (stage12_runtime/failure_recording.py).
  • If that write fails, its error is logged and no longer replaces the error the handler raises.
  • The handlers still raise their stable RuntimeError with no cause (from None), as on main. The original error is therefore visible only in telemetry, not in the exception the Modal call returns.
  • A child failure is reported once.
• No change to error_summary contents. Recording a redacted cause is a separate change.

Deployment

No environment variables, secrets or Modal function settings change. STAGE12_DATABASE_URL is read exactly as before. This PR doesn't change any workflow, Dockerfile or Modal configuration.

Two services use the changed persistence library. Both deploy through simulation-deploy.reusable.yml, which runs from simulation-deploy.yml on every push to main (staging, then production) and from the manual stage12-deploy.yml.

• Stage 12 coordinator and workers (versioned v2 Modal app, deploy_stage12_v2 job): they get NullPool and the connection retry described above.
• Cloud Run Simulation Entrypoint (deploy_entrypoint job):
  • stage12_backend builds one store at startup and reads it on each Stage 12 status request: get_report, then list_simulations.
  • On main, its pool kept up to 5 sessions open per instance between requests (up to 10 under load). Under NullPool, each status poll opens and closes 2 new pooler sessions, adding a connection handshake to each read, and the entrypoint holds no idle sessions between polls.
  • The entrypoint keeps the store's default opener, which does not retry. A refused connection fails that poll immediately, as on main, rather than holding the request for up to 60 s.
  • stage12_infrastructure (the configure_stage12_v2 job's validation) opens a single connection per run, so NullPool doesn't change it in practice.

Not included

• Marking stale running rows as failed. This needs a scheduled function; see the follow-ups in the issue.
• Switching to the transaction-mode pooler.

Merge notes

This PR edits the same failure handlers in coordination.py and simulation.py as #791 (record a redacted failure cause) and #793 (share the child wait deadline). The handlers here pass error_summary through a shared helper, so the redacted cause slots into that argument when rebasing. I suggest merging this PR first.

Testing

• Persistence library (libs/policyengine-stage12-persistence/tests/test_store.py):
  • Regression test for the connection lifecycle: it runs the runtime engine configuration against a real SQLite driver, and checks that the physical connection closes after a successful and a failed transaction and that the next operation gets a new connection. Both cases fail with the previous QueuePool configuration.
  • Store operations open their connection through the opener, with their own operation name, before any statement. A refused connection reaches the opener, and the operation completes after the opener retries it.
  • A COMMIT failure produced by real SQLAlchemy (it has no statement) is raised after the opener returns and is not retried.
  • A store with the default opener, as the entrypoint uses, does not retry a refused connection.
• tests/test_stage12_database_retry.py:
  • Delays grow and stay bounded, and jitter stays between half and all of the ceiling.
  • Retries stop when the budget is spent.
  • Errors other than a refused connection are not retried.
  • A telemetry failure doesn't stop a retry.
  • Without an observability runtime, each retry is logged at WARNING.
• tests/test_stage12_runtime.py:
  • The runtime store retries a refused connection when creating child state, and logs the retry with that operation's name.
  • The runtime store does not retry a failed COMMIT.
  • The coordinator does not deduplicate against its own claim after a COMMIT reply is lost: it raises and dispatches nothing. This test fails on the previous version of this PR, which returned a deduplicated success.
  • The coordinator and the worker build their store with their own observability runtime.
  • Both report the original failure before a failure record they can't write.
  • A child failure is reported once, both when waiting for it fails and when dispatching it fails.
• tests/test_stage12_failure_recording.py: without an observability runtime, the original failure and an unwritable failure record are each logged at ERROR.
• Suites, on Python 3.13 with uv run --frozen:
  • libs/policyengine-stage12-persistence: 17 passed.
  • projects/policyengine-simulation-executor: 750 passed, 23 skipped.
  • projects/policyengine-simulation-entry: 158 passed.
  • libs/policyengine-simulation-contract was not run; it is unchanged.
• Formatting, lint and types:
  • With the locked ruff (0.15.11), ruff format --check and ruff check pass on the changed files.
  • pyright reports no errors for stage12_runtime and the persistence library.
• Not done: a production load test.

🤖 Generated with Claude Code