Python: carry fan-in edge buffers through checkpoint restore - #7948
Python: carry fan-in edge buffers through checkpoint restore#7948Yashvant Mahadev Hange (YashvantHange) wants to merge 5 commits into
Conversation
A fan-in buffers messages until every source has produced one. That buffer was neither checkpointed nor reset on restore, so buffered messages were lost when resuming a rebuilt workflow, and a buffer left behind by a run that failed mid-superstep survived the restore and duplicated deliveries. Edge runner delivery state now travels with the checkpoint under the reserved shared-state key _edge_state. Both restore paths reset every edge runner before reapplying what the checkpoint held, so a checkpoint written before this change resets the buffer instead of leaking it.
GitHub Actions
Active
GitHub Actions
Active
GitHub Actions
Active
agent-framework-automation
Bot
added
the
python
GitHub Actions
Active
There was a problem hiding this comment.
Pull request overview
Adds checkpoint persistence and restoration for partially filled fan-in edge buffers.
Changes:
- Adds topology-keyed edge-runner checkpoint state.
- Restores or clears fan-in buffers during checkpoint recovery.
- Adds runner, workflow, and sub-workflow regression tests.
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
python/packages/core/agent_framework/_workflows/_const.py |
Defines the reserved edge-state key. |
python/packages/core/agent_framework/_workflows/_checkpoint.py |
Documents edge state in checkpoints. |
python/packages/core/agent_framework/_workflows/_edge_runner.py |
Snapshots and restores fan-in buffers. |
python/packages/core/agent_framework/_workflows/_runner.py |
Integrates edge state into checkpoint lifecycle. |
python/packages/core/tests/workflow/test_runner.py |
Tests buffer restoration and stale-state clearing. |
python/packages/core/tests/workflow/test_checkpoint.py |
Tests workflow restoration from a partial fan-in. |
python/packages/core/tests/workflow/test_sub_workflow.py |
Tests fan-in restoration within sub-workflows. |
💡 Add a code-review agent skill for context-aware, tailored reviews. Learn more in the docs.
Executor ids are only required to be non-empty, so joining source and target ids with "->" and "," let two different fan-in groups produce the same key: sources ['a', 'b->t,c'] and ['a->t,b', 'c'] into the same target collide. One group's buffer snapshot would then overwrite the other's, and the restore would hand both runners the same messages. JSON-encode the sorted pairs instead.
GitHub Actions
Active
@moonbox3 on PR microsoft#7948: _run_iteration's asyncio.gather calls did not cancel their other tasks when one raised, so a still-in-flight fan-in delivery could resume after restore_checkpoint/restore_from_checkpoint had already reset the FanInEdgeRunner's buffer on the same runner instance, appending a message from the failed, never-checkpointed superstep into the freshly restored buffer. Confirmed the race is real before fixing it: reproduced it with a test that fails a sibling source's delivery while another source's fan-in delivery is parked mid-flight, restores from a checkpoint taken before either message was sent, then releases the parked delivery. Without the fix it appends the stale message into the just-restored buffer, and the fan-in later fires using that stale value instead of the fresh resumed delivery. Added _gather_cancelling_siblings_on_error to replace both gather() call sites in _run_iteration, so a failing delivery cancels and awaits every other one before the exception propagates.
GitHub Actions
Active
@moonbox3 follow-up on PR microsoft#7948: the previous fix wrapped _run_iteration's two gather() sites (across sources, and across edge runners for one source), but FanOutEdgeRunner.send_message has its own separate asyncio.gather() across one message's fan-out targets. When a fan-out has a single edge runner for its source - the common case - that inner gather is the only level where one target's failure has a sibling to race against, so the outer fix never reached it. Confirmed the race is real before fixing it: reproduced a fan-out target still executing when its sibling fails, calling ctx.send_message after restore_checkpoint has already cleared RunnerContext._messages, landing its stale output in the queue the restore reset for the resumed run. The test fails on the pre-fix code and passes after. Moved the cancel-and-await helper (renamed gather_cancelling_siblings_on_error, no longer private) into _edge_runner.py, since _runner.py already imports from there and the reverse would be circular, and used it at the fan-out gather site too.
GitHub Actions
Active
Retroactive check on our own claim, per mistake microsoft#12: this PR's review reply said the other gather() sites in _workflow.py/_workflow_executor.py "hold no cross-invocation buffered state a restore can race with." Tested rather than re-argued it. _workflow.py:1034 (send_request_info_response) survives: it pops a still- pending event before writing anything, and a restore clears pending events too, so an orphaned write fails closed (raises) instead of landing. Reproduced directly against InProcRunnerContext to confirm. _workflow_executor.py's four gather() sites do not survive. ctx.send_message and ctx.yield_output/add_event write into RunnerContext state (_messages, the event queue) with no equivalent precondition check. Reproduced the same shape as the fan-in/fan-out bugs: a sub-workflow output still being forwarded when a sibling output's send fails calls ctx.send_message after a checkpoint restore has already cleared the parent's message queue, landing stale output in it. Fixed all four sites with the same gather_cancelling_siblings_on_error helper already used in _run_iteration and FanOutEdgeRunner - moved to _edge_runner.py in the prior commit specifically so _workflow_executor.py could import it without a cycle.
GitHub Actions
Active
|
Evan Mattson (@moonbox3) — flagging this myself rather than waiting for you to find it: after fixing the fan-out gather you reported, I went back and re-examined my own reply from earlier in this thread — the one that excluded "a couple of other
Net: the "these other sites are fine" line I gave you earlier in this thread wasn't reliable — one of the two claims in it was wrong. Wanted you to have the accurate picture rather than assume silence meant it held up. |
Motivation & Context
A fan-in holds messages in
FanInEdgeRunner._bufferuntil every source has produced one. That buffer is neithercheckpointed nor reset when a workflow is restored, and both halves of that cause a bug:
buffer is the only place it exists. A fan-in whose sources arrive in different supersteps holds messages across a
superstep boundary, so a checkpoint taken there omits them. Resuming on a rebuilt workflow leaves a fan-in that is
permanently short and never fires.
buffer. Restoring an earlier checkpoint re-executes those sources, so the same message is delivered twice, and the
duplicate can complete the fan-in before the restored superstep has produced all of its messages.
Description & Review Guide
Edge runner delivery state now travels with the checkpoint, under the reserved shared-state key
_edge_state,alongside the existing
_executor_state.What are the major changes?
EdgeRunnergainsstate_key,snapshot_state(), andrestore_state(). The base implementations hold nostate, so every runner other than the fan-in is unaffected.
FanInEdgeRunnerimplements them over its buffer.RunnerImplcaptures edge state in_prepare_checkpoint_state(the single path behind bothcreate_checkpoint_if_enabledandbuild_checkpoint) and, in bothrestore_from_checkpointandrestore_checkpoint, resets every runner before reapplying what the checkpoint held. The reset runs even forrunners the checkpoint has no entry for, which is what discards a buffer left by an interrupted run.
fan-in holds only one branch, and the same case inside a sub-workflow resumed through its parent's checkpoint.
What is the impact of these changes?
state_keyis derived from the group's topology rather thanEdgeGroup.id, which defaults to a random UUID andso differs between two instances of the same workflow definition - restoring onto a rebuilt instance is the
normal case. Builder validation rejects two edges sharing a
source -> targetpair anywhere in the workflow, sono two groups can produce the same key.
_edge_stateentry: every runner is reset and nothing isreapplied, which is the old behavior minus the stale buffer.
State, which deep-copies on write and again on export, where in-flightmessages in
WorkflowCheckpoint.messagesdo not. This matches how executor state is already handled and costs acopy per checkpointed superstep for a fan-in that is mid-aggregation.
run()again on a workflow instance whose fan-in still holds messages is out of scope here. That is nota restore, and a pause-and-continue on the same instance (
run(responses=...)after an idle-with-pending-requests stop) legitimately depends on the buffer surviving, so telling the two apart is a separate change.
Related Issue
Fixes #7371
Contribution Checklist