Python: Fix sub-workflow checkpoint restore to preserve sub-workflow state#7097
Python: Fix sub-workflow checkpoint restore to preserve sub-workflow state#7097TaoChenOSU wants to merge 2 commits into
Conversation
Add Runner.capture_checkpoint_object/restore_from_checkpoint_object (quiescent-only nested checkpoint) and embed a sub_workflow_checkpoint in WorkflowExecutor.on_checkpoint_save/on_checkpoint_restore so a resumed parent restores each sub-workflow's mid-progress state instead of only replaying pending request-info events. Keeps a backward-compat fallback when sub_workflow_checkpoint is absent.
There was a problem hiding this comment.
Pull request overview
This PR fixes a correctness gap in Python workflow checkpointing for nested workflows: when a parent workflow is resumed, the sub-workflow now resumes with its own internal runner/executor state restored (not just the parent WorkflowExecutor’s pending-request bookkeeping).
Changes:
- Add
Runner.capture_checkpoint_object()/Runner.restore_from_checkpoint_object()to support in-memory checkpoint capture/restore for embedded (sub-workflow) state. - Update
WorkflowExecutorcheckpoint save/restore to embed and restore the child workflow’s checkpoint, with a backward-compat fallback when the embedded checkpoint is absent. - Add regression tests covering checkpoint object roundtrip, quiescent-only capture, graph-signature mismatch rejection, and mid-progress sub-workflow resume behavior.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| python/packages/core/agent_framework/_workflows/_runner.py | Adds in-memory checkpoint capture/restore APIs used for nested workflow checkpoint embedding. |
| python/packages/core/agent_framework/_workflows/_workflow_executor.py | Embeds the sub-workflow checkpoint on save and restores it on resume, keeping a fallback for older checkpoints. |
| python/packages/core/tests/workflow/test_runner.py | Adds unit tests for the new runner checkpoint object APIs and invariants. |
| python/packages/core/tests/workflow/test_sub_workflow.py | Adds regression coverage ensuring sub-workflow mid-progress state survives parent checkpoint restore. |
| sub_workflow_checkpoint = state.get("sub_workflow_checkpoint") | ||
| if sub_workflow_checkpoint is not None: | ||
| sub_workflow_checkpoint = decode_checkpoint_value(sub_workflow_checkpoint) | ||
| await self.workflow._runner.restore_from_checkpoint_object(sub_workflow_checkpoint) # pyright: ignore[reportPrivateUsage] | ||
| return |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Python Test Coverage Report •
Python Unit Test Overview
|
|||||||||||||||||||||||||||||||||||
Add RunnerContext.create_checkpoint_object alongside create_checkpoint (create_checkpoint now delegates to it and persists), so Runner.capture_checkpoint_object builds the snapshot via the context instead of a one-off get_messages peek primitive. In-flight messages are captured non-destructively (per-source lists copied). The checkpoint-less capturing contexts (azurefunctions, durabletask) raise NotImplementedError to match create_checkpoint.
| self._workflow_name, | ||
| self._graph_signature_hash, | ||
| self._state, | ||
| None, |
There was a problem hiding this comment.
What happens to the child checkpoint chain when the wrapped workflow also has checkpoint storage? This creates a fresh checkpoint ID but never persists it; after parent restore, _mark_resumed() makes the child's next persisted checkpoint point to that missing ID. I reproduced this with separate parent and child InMemoryCheckpointStorage instances: the first child checkpoint after resume had a previous_checkpoint_id absent from child storage, so ancestry traversal is broken.
|
|
||
| await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id="START")) | ||
|
|
||
| with pytest.raises(WorkflowCheckpointException, match="in-flight executor messages"): |
There was a problem hiding this comment.
Could we align this test with the intended behavior? The code now saves queued messages in the checkpoint, but this test still expects that to raise an error, so it fails locally and in CI. If saving queued messages is intentional, the test and PR description should be updated; otherwise, the code should continue rejecting them.
Motivation & Context
When a workflow that contains sub-workflows (via
WorkflowExecutor) is checkpointed and later resumed, each sub-workflow's own mid-progress state was lost. On restore, the parent only replayed the sub-workflow's pending request-info events; any executor state or in-flight progress inside the sub-workflow that had advanced before the checkpoint was not restored. A resumed parent therefore re-ran sub-workflows from an effectively empty state instead of continuing where they left off, producing incorrect results for nested/hierarchical workflows that checkpoint mid-run.This change makes sub-workflow checkpoint/restore preserve the full nested state, so resuming a parent workflow faithfully continues each sub-workflow.
Description & Review Guide
What are the major changes?
Runnergains two internal methods:capture_checkpoint_object()— captures a sub-workflow's state as aWorkflowCheckpointobject. It is quiescent-only: it raisesWorkflowCheckpointExceptionif the runner still has in-flight messages, so only settled state is embedded.restore_from_checkpoint_object(checkpoint)— validates the graph signature, clears and re-imports shared/executor state, applies the checkpoint, and marks the runner as resumed.WorkflowExecutor.on_checkpoint_save()now embeds the nested checkpoint under asub_workflow_checkpointkey, andon_checkpoint_restore()decodes it and restores the sub-workflow viarestore_from_checkpoint_object().sub_workflow_checkpointis absent (checkpoints written by older versions), restore falls back to the previous behavior of replaying pending request-info events.test_runner.py(capture/restore roundtrip, quiescent-only guard, graph-signature-mismatch rejection) andtest_sub_workflow.py(mid-progress sub-workflow resume preserves state).What is the impact of these changes?
Runnermethods are internal. Existing checkpoints remain loadable via the fallback path.What do you want reviewers to focus on?
capture_checkpoint_object()and the graph-signature validation inrestore_from_checkpoint_object().WorkflowExecutor.on_checkpoint_restore()for checkpoints lackingsub_workflow_checkpoint.Related Issue
Part of the multi-PR workflow engine refactor series (follows #6695 and #6776). No standalone tracking issue.
Contribution Checklist
breaking changelabel (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.