Python: Fix sub-workflow checkpoint restore to preserve sub-workflow state#7097
Open
TaoChenOSU wants to merge 6 commits into
Open
Python: Fix sub-workflow checkpoint restore to preserve sub-workflow state#7097TaoChenOSU wants to merge 6 commits into
TaoChenOSU wants to merge 6 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.
Contributor
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. |
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
Contributor
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.
moonbox3
reviewed
Jul 14, 2026
The sub-workflow is a single shared instance, so per-execution ExecutionContext/request routing never provided real isolation. Delegate request/response tracking to the sub-workflow itself: can_handle accepts targeted propagated responses, _handle_response validates against the sub-workflow's pending requests and forwards responses immediately, and on_checkpoint_save embeds only the sub-workflow checkpoint (on_checkpoint_restore keeps a legacy reader for older checkpoints). Also emit the fresh-message/checkpoint-while-pending warning from FunctionalWorkflow.run to match Workflow.run.
The storage backend already materializes the full checkpoint on load (FileCheckpointStorage decodes recursively; InMemoryCheckpointStorage deep-copies), so the embedded sub_workflow_checkpoint (and legacy execution_contexts) arrive already decoded - like every other executor's on_checkpoint_restore state. Remove the no-op decode_checkpoint_value calls and the now-unused import.
| self._previous_checkpoint_id = checkpoint.checkpoint_id | ||
| return checkpoint | ||
|
|
||
| async def restore_from_checkpoint_object(self, checkpoint: WorkflowCheckpoint) -> None: |
Member
There was a problem hiding this comment.
In the description this is called a internal method, do should be "_restore...", we treat the whole set of packages as internal so even if this flags from within the orchestration package it should still be underscored
| """ | ||
| ... | ||
|
|
||
| async def create_checkpoint_object( |
Member
There was a problem hiding this comment.
The description has "capture" instead of create, and also should be underscored
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.