[anth-volk]

Fixes #895

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

The v2 Modal report coordinator can now hand a report's result straight back to its caller, instead of only recording it in the Stage 12 tables and the Stage 12 artifact bucket. A new dispatch-context field, result_delivery, selects the mode. Leaving it out keeps today's Stage 12 behavior, and the Simulation Entrypoint does not send it, so every Stage 12 comparison run behaves as before. Nothing calls the new mode yet: the Simulation Entrypoint's v2 economy path, a later change, will.

This replaces the closed #836 and #837 as one PR, keeping what still fits, with the fixes from their review and from three review rounds on this branch.

The contract (policyengine_simulation_contract)

• Stage12InvocationContext.result_delivery: stage12_database (the default) or return_to_caller. A return_to_caller context cannot name a production job to compare with, because that comparison is recorded on the Stage 12 report row.
• ReportExecutionResult (new report_execution module): the report's totals inline as a SingleYearMacroOutput, their SHA-256 (report_aggregate_sha256, over compact key-sorted JSON, rejecting NaN and infinity), the output plan, the baseline and reform artifact descriptors, and provenance (bundle, worker version, Modal environment, application and simulation function, version manifest digest, coordinator call ID). It rejects artifacts that do not form the report, a digest that does not match, and requested aggregates that differ from the ones the output plan was built for, in content or order, so the two cannot disagree undetected.
• ReportExecutionFailure and ReportExecutionError: the redacted error code and summary a Stage 12 row would record for the same failure, under a fixed message. The error pickles its failure as plain JSON values and rebuilds the model when it is unpickled, so the two sides' pydantic versions need not agree.

Does Modal raise ReportExecutionError in the caller as the same type? Yes.

• A Modal function container pickles the exception its function raised with modal._serialization.pickle_exception, and FunctionCall.get in the caller unpickles and raises it in modal._utils.function_utils._process_result. test_report_execution_modal_transport runs both of Modal's own functions on a ReportExecutionError and gets back the same type, failure and message. Modal's error hints (exc_with_hints) return the same object.
• test_stage12_modal_app checks that the Modal coordinate_report function raises the returned report's error unchanged, as the same object, so Modal pickles exactly that error.
• Across Modal client versions: the v2 worker image locks Modal 1.4.2 and the Simulation Entrypoint locks 1.5.5. A one-off check pickled the error with the executor environment's Modal 1.4.2 and raised it with the entrypoint environment's Modal 1.5.5: same type, same failure, same message.
• Limits: no test runs a real Modal container. The caller must be able to import policyengine_simulation_contract.report_execution; the Simulation Entrypoint depends on the contract library.
• A successful result crosses Modal as the result's JSON-mode dictionary, as the Stage 12 coordinator's descriptor does, so the container's and the caller's pydantic versions never need to agree on a pickled model. The caller validates it with ReportExecutionResult.model_validate.

The returned mode (stage12_runtime.coordinate_report_for_caller, and the Modal coordinate_report with result_delivery = return_to_caller)

• Reads and writes no Stage 12 rows and never builds the database store. It claims nothing, so a repeated dispatch runs the report again; the caller deduplicates.
• Takes no parent record (passing one raises ValueError) and checks the application, worker version and bundle digest as the Stage 12 path does.
• Writes the simulation artifacts under the Stage 12 layout, stage-12-runs/<environment>/<YYYY>/<MM>/<evaluation_id>, built from the context; a prefix the caller sends is replaced.
• Plans the output, starts both workers with the same delivery, waits for both within the shared window, and reads and checks their artifacts. These steps are the code the Stage 12 coordinator runs, moved unchanged into a shared base class, ReportSimulationCoordination, whose subclasses decide what a failed or stopped worker records.
• Aggregates, applies the checks a Stage 12 aggregate artifact write applies, requires the aggregate to describe this report (report ID, bundle, requested aggregates), and returns a ReportExecutionResult. No aggregate artifact is written and no production comparison runs.
• A failure is reported to telemetry once, recorded as the operation's cause, and raised as ReportExecutionError with the redacted code and summary a Stage 12 row would record. A worker's own ReportExecutionError takes precedence over its call error, as a Stage 12 worker's recorded failure does. Failures before any worker starts (invalid payloads, the dispatch checks, an unset artifact bucket) raise ValueError or ValidationError, as the Stage 12 coordinator does at the same points before it writes any row.
• A worker that must stop (its peer failed, the wait window closed, or the coordinator failed after starting it) is cancelled through Modal; a cancellation error is logged and never replaces the failure. If Modal cancels the coordinator's own input, it cancels its unfinished workers the same way and re-raises the cancellation without reporting a failure.
• The coordinator's and both workers' telemetry operations carry result_delivery=return_to_caller, so they can be told apart from Stage 12 runs. The region-group segment function receives no dispatch context and is unchanged; its spans join the worker's trace.

Workers (run_single_simulation, Modal run_single_simulation_us and _uk)

With return_to_caller, a worker checks the policy periods, writes its input JSON and Parquet artifact as a Stage 12 worker does, and returns the artifact descriptor, touching no rows. It stores both under keys naming their SHA-256 (inputs/<role>-<sha256>.json and simulations/<role>-<sha256>.parquet), so a restarted or repeated call never collides with an earlier attempt's objects (see "Pre-emption" below); the comparison path keeps its fixed keys. A failure raises ReportExecutionError, including a segment's failure and typed input failures such as invalid_policy_period.

Segment calls of a cancelled segmented US worker

Modal stops a cancelled synchronous input by raising InputCancellation (a BaseException) in it, and does not cancel the calls that input spawned. The segment runner caught only Exception, so up to 20 region-group segment calls kept running after their worker was cancelled through Modal. It now keeps each call as soon as it starts, inside its dispatch span, and cancels every started segment call when its dispatch or wait stops on a BaseException that is not an Exception, then re-raises it; failures that are Exceptions keep their existing handling, and each call is cancelled once. This is how a returned segmented run stops its segments. On the Stage 12 path, the coordinator cancels a segmented worker through Modal only when it could not write the cancellation to the worker's record; that fallback now stops the segment calls too. No Stage 12 record or result changes.

The Stage 12 path is unchanged

• The Stage 12 coordinator leaves result_delivery out of the context it sends its workers, so their payload has exactly the fields it had before (tested).
• The Modal coordinate_report function's parent argument now defaults to None; the Simulation Entrypoint still passes all three arguments, and a Stage 12 report without a parent record is still rejected before anything is claimed or started (tested). Stage 12 operation attributes are unchanged (tested).
• The Stage 12 coordinator refuses a return_to_caller context before claiming anything.
• The shared dispatch step now keeps a started worker's call before its dispatch span closes, so a stop in that moment still finds it. On the Stage 12 path this only matters if closing the span fails: the worker that had just started is then stopped instead of being left running. Stage 12 records are otherwise unchanged.

Pre-emption: decision A, and what was verified

Decision A: when Modal pre-empts or cancels a call, that call cancels its child calls. Modal then restarts the pre-empted call on another machine, and it starts over. Pre-emption may make a result slower but must never make it fail.

How it works here. A synchronous Modal function receives pre-emption as KeyboardInterrupt (as Modal's own modal.exception.simulate_preemption documents) and a cancellation as InputCancellation; both are BaseExceptions that are not Exceptions. The returned coordinator and the segment runner handle any such exception by cancelling what they started and re-raising it. The handlers do not import Modal's class, because the project's src/modal package shadows the installed Modal package in unit tests.

Every object a returned report writes, and whether a restart can write it again:

• The returned coordinator writes no object: no aggregate artifact and no comparison artifact.
• The region-group segment calls write no object; they return their results to the worker.
• Each worker writes its planned input JSON and its Parquet output. Both used fixed keys, and the bucket accepts a second write to a key only when the bytes are identical. Identical frames do encode to identical Parquet bytes (checked across processes with different hash seeds), but a restart on another machine can calculate different last bits, and a restart on a redeployed copy of the application can plan the input differently. Either way the restarted worker failed with "immutable Stage 12 artifact path contains other data", and so did the workers a restarted coordinator starts. Both are now stored under keys naming their SHA-256: an identical object reuses the earlier one, and a different one is stored beside it. A cancelled earlier attempt still writing at the same time is also safe: on one key exactly one create succeeds and the other finds identical bytes.
• A report writes no cache object and no cache manifest. The workers read the cache baked into their image at deploy time; only the deploy-time cache precompute writes to the cache bucket.

Tests: a restarted worker whose output differs by one bit, or whose planned input differs, stores it; one with the same output reuses its artifact; and a returned report whose coordinator is interrupted after the baseline stored its output finishes when restarted with a calculation that differs by one bit. Without the change, the differing-output, differing-input and restarted-report tests fail with the collision.

Not changed: fixed time limits across attempts. A restart starts the pre-empted call's work again, but these limits are fixed and are not extended by a restart:

• The coordinator stops waiting for its workers 3,900 s after its own attempt started (SIMULATION_WAIT_TIMEOUT_SECONDS, equal to the US worker's Modal timeout), then fails the report with simulation_wait_timeout.
• A US worker has 3,900 s (its Modal timeout) to collect its 20 segment calls, which have 3,600 s each; a segment call restarted late can make the worker time out.
• The coordinator's own Modal timeout is 5,400 s.

A complete US report took 288 s in testing, so a single pre-emption leads to a failure only when a run is already close to these limits, about an hour in, or after repeated pre-emptions. Removing that last risk means budgeting the waits and Modal timeouts for restarts, for example dropping the returned coordinator's own wait window and raising the coordinator's and the US worker's Modal timeouts. Those Modal functions are shared with the comparison path, and longer timeouts let idle waiting containers run longer. By decision, these limits stay as they are and will be revisited with plan item 2.01.

Other notes. The planned resumption after pre-emption, which reattaches to a re-run's children, is replaced by decision A for the returned mode. Nothing calls the returned mode yet: the Simulation Entrypoint's v2 economy path will call it, validate the returned dictionary with ReportExecutionResult.model_validate, and handle ReportExecutionError.

Environment variables

This PR adds, removes or rewires no environment variable, workflow input, GitHub secret, Secret Manager secret or Modal secret. Variables the changed runtime code reads:

• STAGE12_ARTIFACT_BUCKET (both modes; artifact_store() in the coordinator and workers). Provider: the GitHub environment variable STAGE12_ARTIFACT_BUCKET, present in beta and prod. Wiring: simulation-deploy.yml and stage12-deploy.yml call simulation-deploy.reusable.yml with deploy_stage12_v2: true and secrets: inherit; its deploy_stage12_v2 job runs in the beta or prod environment, and its step "Synchronize the separate Stage 12 runtime secret" passes the variable to .github/scripts/stage12-sync-modal-secrets.sh, which writes it into the Modal secret stage12-evaluation-runtime in the staging or main Modal environment. v2_app.py attaches that secret to the coordinator, both country workers and the segment worker (worker_secrets). The variable's value was not inspected.
• GOOGLE_APPLICATION_CREDENTIALS_JSON (both modes; read by setup_gcp_credentials when the artifact store builds its Cloud Storage client, so the returned mode needs it for every artifact read and write). Provider: the Secret Manager secret named by the GitHub environment variable STAGE12_GCP_CREDENTIALS_SECRET_NAME, in the project named by SIMULATION_ENTRYPOINT_GCP_PROJECT_ID (both present in beta and prod). Wiring: the same step authenticates to Google Cloud with the environment secrets SIMULATION_ENTRYPOINT_GCP_WIF_PROVIDER and SIMULATION_ENTRYPOINT_GCP_DEPLOY_SERVICE_ACCOUNT, reads the secret with gcloud, and writes it into the Modal secret stage12-evaluation-gcp-credentials, which worker_secrets attaches to the same functions. Not verified: the secret's content.
• STAGE12_DATABASE_URL (Stage 12 mode only; the returned mode never reads it). Provider: the Secret Manager secret named by STAGE12_DATABASE_URL_SECRET_NAME (present in beta and prod), through the same step into stage12-evaluation-runtime. Unchanged.
• HUGGING_FACE_TOKEN (both modes; the workers' dataset loading, and the failure redaction, which reads it only to remove its value from summaries). Provider: the organization secret PE_UK_PRIVATE_HF_READ_TOKEN, written by the deploy workflow into the Modal secret v2_app.py attaches as hf_secret. Unchanged.
• MODAL_ENVIRONMENT (the Modal coordinate_report function, which this PR changes, reads it only to label its telemetry, falling back to local). Provider: Modal, which documents setting it in every function container to the name of the Modal environment the container runs in. No workflow, script or secret in this repository provides it. Not verified in the deployed containers. Unchanged.

The deploy job also uses MODAL_TOKEN_ID and MODAL_TOKEN_SECRET (repository and organization secrets, through secrets: inherit); the changed runtime code does not read them, because a coordinator starts its workers with Modal's own container credentials.

What a merge deploys, and what it waits for

Merging to main runs the complete staging and then production deployment. The redeployed v2 Modal application gains the return mode, but nothing calls it. The Simulation Entrypoint sends no result_delivery, so every Stage 12 report keeps its behavior. No image, function resource setting, timeout, secret attachment or workflow changes.

Before it may merge, this PR waits only for CI (CI complete) and a maintainer's review. It depends on no other PR and no operator step, and it adds no database migration.

Testing

Each behavior change has tests that fail without it, checked commit by commit against the parent commit's code. The tests that cover existing behavior (the Modal coordinator passing the error through unchanged, the aggregate mismatch cases, a Stage 12 report still needing its parent record) passed on their first run; the aggregate mismatch cases fail when the coordinator's comparisons are removed.

• Contract (libs/policyengine-simulation-contract): test_report_result_delivery.py (5 tests: the default and return_to_caller delivery, no production comparison for a returned report, an unknown delivery rejected); test_report_execution.py (24 tests: JSON round trip, totals as SingleYearMacroOutput, digest checks and stability, each artifact consistency check, requested aggregates against the output plan in content and order, timezone, unique aggregates, unknown fields, failure bounds, pickling, and a pickle holding only JSON values); test_report_execution_modal_transport.py (1 test, Modal's exception transport). 352 passed.
• Executor (projects/policyengine-simulation-executor): test_stage12_returned_report.py (37 tests: returned workers without the database, their failures and segmented runs; the returned coordinator's typed result, stages, repeated dispatch, dispatch checks, each failure and cancellation path, Modal cancelling the coordinator during the wait or while a dispatch span closes, aggregate and artifact checks, restarts after pre-emption, and a comparison-path worker keeping its fixed keys); test_stage12_artifacts.py (2 new tests: keys named by digest for inputs and outputs); test_stage12_segment_fanout.py (4 new tests: Modal cancelling a segmented worker during the wait, during dispatch, while a dispatch span closes, and before the wait); test_stage12_modal_app.py (8 new tests: routing by delivery, refusing a parent record, unchanged Stage 12 attributes, the error passed through unchanged, and both workers' operation attributes in each mode); test_stage12_runtime.py (1 new test: Stage 12 workers receive the context without the delivery field). Full suite: 1,917 passed, 22 skipped.
• Simulation Entrypoint 668 passed, gateway 268 passed, Stage 12 persistence 40 passed (the contract library is shared).
• ruff format --check projects/*/src libs/*/src and ruff check . with the pinned Ruff 0.15.22 pass, and python -m scripts.rules pyright --sync passes for the contract library (no errors) and the executor (within its recorded baseline), at every commit. scripts/run_quality_guards.py passes.
• Generated clients are unchanged: no Simulation Entrypoint or gateway API schema changed.
• Not run locally: the Modal image smoke and the integration tests, which need Modal and the CI services. CI passed every check on 2026-10-22 before the pre-emption commits; runs on this branch are now cancelled automatically, so the latest commits rely on the local runs above.

🤖 Generated with Claude Code