Share one Stage 12 child wait deadline and cancel unfinished children by anth-volk · Pull Request #793 · PolicyEngine/policyengine-sim-api

anth-volk · 2026-10-13T19:23:04Z

Fixes #792

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

This is an interim change before #712. It keeps the sequential baseline-then-reform wait, but bounds it and cleans up after a failure.

Shared wait deadline:
- One 3,000 s window (SIMULATION_WAIT_TIMEOUT_SECONDS, the same as the workers' Modal timeout) opens when coordinate_report starts. It covers coordinator setup (claiming, output planning including the country model load, retained-artifact reads and dispatch) and both child waits. Each call.get receives only the remaining time, never less than zero.
- Setup time therefore shortens the children's waits, not the time after them. However long setup takes within the window, 1,500 s of the coordinator's 4,500 s Modal timeout remain once the window closes: 900 s for the production-result wait, and 600 s to aggregate, persist and record a failure. This matches the budget already described on coordinate_report in v2_app.py.
- Previously a slow baseline could give reform a second full 3,000 s.
Cancel unfinished workers:
- When a wait fails or times out, _cancel_unfinished_children cancels every unfinished child, best effort. This includes a timed-out child that is still running, and a call restored from an earlier coordinator attempt.
- It records a cancelled sibling that is still RUNNING as FAILED with error_code="simulation_cancelled". A sibling that has already recorded SUCCEEDED or FAILED is left unchanged. Modal's InputCancellation is a BaseException, so the worker can't record its own outcome.
- Cancellation and state-update errors never replace the original failure. Their warnings (stage12_child_cancel_failed, stage12_cancelled_child_state_update_failed) carry evaluation_id, the child's simulation_execution_id, simulation_role and error_type, but no error message.
ChildCall gains cancel(). Modal's FunctionCall.cancel already exists.
Behaviour change: a reform worker still queued when the shared window is used up now fails, instead of running past the coordinator's budget. Coordinator setup time now also comes out of the children's window.
No timeout or Modal configuration changes. No environment variables change.

Interaction with other PRs

Segment and cache Stage 12 US simulations #712: it replaces this wait loop with concurrent polling. When it rebases:
- take Segment and cache Stage 12 US simulations #712's version of the loop and delete _cancel_unfinished_children;
- open Segment and cache Stage 12 US simulations #712's polling deadline when coordinate_report starts, not when polling starts, to keep the guaranteed 1,500 s;
- Segment and cache Stage 12 US simulations #712 adds the same cancel() line to dependencies.py;
- both PRs append tests and cancel() doubles to tests/test_stage12_runtime.py.
Release Stage 12 database sessions and keep original failures visible #789 and Record a redacted Stage 12 failure cause instead of only the exception class #791: they edit the failure handlers in the same file. This PR only changes the wait loop, opens its deadline at coordinator entry, and adds the cancellation helper.

Testing

New tests in tests/test_stage12_runtime.py:
- The children share one wait window, and an exhausted window gives the next child no extra time.
- With 400 s of coordinator setup, the baseline wait receives 2,600 s, and the declared Modal timeouts leave at least 900 s plus 600 s after the window. This test reads the timeouts from the Modal app declaration.
- A baseline failure cancels and records the reform child, including a reform call restored from an earlier coordinator attempt.
- A baseline timeout cancels both unfinished children.
- A reform failure leaves a finished baseline alone.
- Cancellation leaves a sibling that already recorded SUCCEEDED or FAILED unchanged.
- Cancellation errors and sibling state errors don't mask the child failure, and their warnings identify the report and child without the error message.
Suites, on Python 3.13 with uv run --frozen:
- All 117 Stage 12 executor tests pass.
- The full executor suite passes: 744 passed, 23 skipped, 3 deselected.
Regression checks: I reverted each behaviour in turn. Each time, the matching new test failed:
- anchoring the deadline at dispatch fails the setup-time test;
- removing the RUNNING check fails both finished-sibling tests;
- dropping the identifiers fails both warning tests.
The original cancellation tests still fail against the code before this PR only at their clock fixture, because they need the new monotonic-clock hook.
Formatting and lint:
- ruff format --check and git diff --check pass.
- ruff check with the locked ruff (0.15.11) is clean on the changed files.
- Ruff 0.16.8 reports four findings in tests/test_stage12_runtime.py, and the same four are already on main.
Not run: Pyright, which isn't installed in the test environment.

🤖 Generated with Claude Code

The coordinator waited up to 3,000 s for baseline and then up to a fresh 3,000 s for reform, so queued workers could push it past its own 4,500 s Modal timeout before it recorded the failure. Both waits now draw on one 3,000 s window, leaving 1,500 s for the 900 s production-result wait and 600 s of aggregation and failure persistence. When a child wait fails or times out, the coordinator now cancels every unfinished child, best effort, and records a cancelled sibling as failed, because Modal's InputCancellation is a BaseException the worker cannot record itself. Cancellation and state-update errors are logged and never replace the original child failure. This keeps the sequential baseline-then-reform wait; concurrent observation remains part of #712. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Anchor the shared child-wait deadline when coordinate_report starts, so claiming, output planning, retained-artifact reads and dispatch shorten the children's waits rather than the 1,500 s left for the production-result wait, aggregation, persistence and failure recording. Add evaluation_id and the child's simulation_execution_id to the cancellation warnings, and cover finished siblings, restored calls and the ImmediateCall cancel double. The coordinator timeout test now runs the coordinator with simulated setup time against the declared Modal timeouts, so it lives beside the other coordinator tests. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

anth-volk · 2026-10-14T15:14:22Z

Closing this in favour of #712, which now covers #792 and lists Fixes #792. #712 opens the shared child wait window at coordinator entry, cancels every unfinished worker when one fails or times out (including failures during coordinator startup), and logs cancellations with the evaluation and simulation execution IDs. It observes baseline and reform together, which #792 always left to #712.

…eouts The coordinator stopped waiting 3,000 s after it started, although a US simulation worker may run for 3,900 s and each of its region-group calls for 3,600 s, so US simulations could never use the time Modal allowed them. The coordinator's 5,100 s timeout was not derived from either. Following the rule from #793, a new dependency-free stage12_timeouts module defines the worker timeouts, sets the shared child-wait window to the longest of them (3,900 s), and derives the coordinator's Modal timeout as that window plus the 900 s production-result wait plus 600 s to read results, aggregate, write artifacts and record the outcome: 5,400 s. A complete testing-environment US report took 288 s, so 600 s of finishing time is ample. The Modal app, coordinator and comparison step all read these constants. The last poll no longer sleeps past the deadline, so the window closes on time instead of up to 1 s late and the finishing time is not eroded. Tests read the declared Modal timeouts and check that the window covers every simulation worker's timeout and that the coordinator's timeout leaves the production wait plus at least 600 s. They also check that setup time shortens the window and that an exhausted window gives the children no extra time. Only the Modal timeout values change: no image, secret or environment variable attachment changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Adopt #793's cancellation practice in the concurrent coordinator: - Every simulation the coordinator stops, after a peer failure, a wait timeout or a coordinator failure, is recorded as failed with simulation_cancelled, and its error_summary gives the reason, for example "Cancelled after the baseline simulation failed.". The cancelled_after_peer_failure and cancelled_after_coordinator_failure codes and the incomplete status for cancellations are gone; the wait timeout stays the report's error code. - Only a running record changes, so a simulation that already succeeded or failed keeps its outcome. Workers already treat any status other than running as a cancellation, so a segmented worker still sees the failed record and cancels its own region-group calls. - An unsegmented simulation is always cancelled through Modal, even when its record could not be written. A segmented simulation is cancelled through Modal only when its record could not be written, because Modal does not cancel spawned calls recursively. Cancellation and state-update errors are still logged without replacing the original failure. Carry over #793's tests for a cancelled and recorded reform, a finished baseline left alone, a restored call cancelled on retry, finished siblings left unchanged, and both unfinished children cancelled on timeout. Add tests for the Modal fallback after a failed state write, for segmented and unsegmented simulations. ImmediateCall now counts cancellations, and the entry and persistence fixtures use the new cancellation record. Add a changelog line, as #793 did. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>