[ShiYuan Xiong (hehe764)]

Motivation & Context

A functional workflow with checkpoint storage can lose the per-step checkpoint of a step that already completed when a parallel sibling step fails, so a resumed run re-runs the completed step and its side effects. Whether the checkpoint is lost depends on the storage implementation's await timing, not on user code. Fixes #9615 (analysis and approach proposed in the linked issue comment).

Description & Review Guide

• What are the major changes?

  • The per-step checkpoint save now runs as its own task, with both the normal and the drain await wrapped in asyncio.shield, so a cancellation of the step task during the framework's failure cleanup after a sibling raises cannot reach the save through the await.
  • RunContext tracks in-flight per-step saves in a dedicated set; _cancel_workflow_tasks expands protection through the task-parent map to include storage-created descendants, and drains these tasks rather than cancelling them.
  • Regression test test_completed_parallel_step_is_checkpointed_when_sibling_fails covers inline asynchronous writes, asyncio.create_task writes, and asyncio.gather writes; all three resume without repeating the completed step.

• What is the impact of these changes?

  • Failure cleanup no longer interrupts an in-flight save or its storage-created child tasks. When storage completes successfully, the persisted result can be replayed without repeating the completed step. Storage failures, direct cancellation, and process termination are outside this guarantee.
  • Run teardown waits for in-flight storage saves. An adapter that never finishes can delay teardown indefinitely; storage timeouts remain the adapter's responsibility.

• What do you want reviewers to focus on?

  • The interaction between the loop task factory (which registers the save task) and _cancel_workflow_tasks: the sweep protects _checkpoint_save_tasks and their descendants from cancellation but still awaits them in its drain gather.
  • Whether waiting for in-flight saves in the failure sweep is acceptable, versus an alternative where storages are individually responsible for surviving cancellation.

Related Issue

Fixes #9615

AI Assistance

• This is an AI-assisted contribution. I reviewed, understood, and verified all submitted content and accept responsibility for it.

AI assistance details: AI assistance was used for the root-cause analysis, implementation, tests, and this PR description. All changes were reviewed, executed, and verified locally by the contributor (reproduction script, regression test red/green against unmodified main, full tests/workflow/ suite on the merged head (14445d9e): 1282 passed, 6 skipped, 1 xfailed; ruff and pyright passed). The descendant-task regression failed for task/gather storage before this follow-up fix. The complete core suite was not rerun for the follow-up.

Contribution Checklist

• The code builds clean without any errors or warnings
• All unit tests pass, and I have added new tests where possible
• The PR follows the Contribution Guidelines
• This PR links to an agreed issue with no competing open PR, or the Related Issue section documents a trivial-change or repository-automation exception.
• This is not a breaking change.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[ShiYuan Xiong (hehe764)]

@microsoft-github-policy-service agree