agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

[BREAKING] Python: Checkpoint refactor: encode/decode, checkpoint format, etc (#3744)

Browse Source 
* WIP: Checkpoint refactor: encode/decode, checkpoint format, etc

* WIP: Remove workflow ID in checkpoints

* Refactor checkpointing

* Add get_latest tests

* Increase test coverage

* Fix formatting

* Fix unit tests

* Fix samples

* fix unit tests

* fix pipeline

* Copilot comments

* Fix tests

* Fix more tests

* Address comments part 1

* Address comments part 2

* Comments
This commit is contained in:
Tao Chen
2026-02-11 12:57:15 -08:00
committed by GitHub
Unverified
parent a2a672b687
commit 7db6c4ab4e
43 changed files with 3335 additions and 2075 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/orchestrations/tests/test_concurrent.py
+10 -14
View File
@@ -224,13 +224,10 @@ async def test_concurrent_checkpoint_resume_round_trip() -> None:
assert baseline_output is not None
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert checkpoints
checkpoints.sort(key=lambda cp: cp.timestamp)
resume_checkpoint = next(
(cp for cp in checkpoints if (cp.metadata or {}).get("checkpoint_type") == "superstep"),
checkpoints[-1],
)
resume_checkpoint = checkpoints[1]
resumed_participants = (
_FakeAgentExec("agentA", "Alpha"),
@@ -270,14 +267,13 @@ async def test_concurrent_checkpoint_runtime_only() -> None:
assert baseline_output is not None
checkpoints = await storage.list_checkpoints()
assert checkpoints
checkpoints.sort(key=lambda cp: cp.timestamp)
resume_checkpoint = next(
(cp for cp in checkpoints if (cp.metadata or {}).get("checkpoint_type") == "superstep"),
checkpoints[-1],
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert len(checkpoints) >= 2, (
"Expected at least 2 checkpoints. The first one is after the start executor, "
"and the second one is after the first round of agent executions."
)
checkpoints.sort(key=lambda cp: cp.timestamp)
resume_checkpoint = checkpoints[1]
resumed_agents = [_FakeAgentExec(id="agent1", reply_text="A1"), _FakeAgentExec(id="agent2", reply_text="A2")]
wf_resume = ConcurrentBuilder(participants=resumed_agents).build()
@@ -320,8 +316,8 @@ async def test_concurrent_checkpoint_runtime_overrides_buildtime() -> None:
assert baseline_output is not None
buildtime_checkpoints = await buildtime_storage.list_checkpoints()
runtime_checkpoints = await runtime_storage.list_checkpoints()
buildtime_checkpoints = await buildtime_storage.list_checkpoints(workflow_name=wf.name)
runtime_checkpoints = await runtime_storage.list_checkpoints(workflow_name=wf.name)
assert len(runtime_checkpoints) > 0, "Runtime storage should have checkpoints"
assert len(buildtime_checkpoints) == 0, "Build-time storage should have no checkpoints when overridden"
python/packages/orchestrations/tests/test_group_chat.py
+3 -3
View File
@@ -620,7 +620,7 @@ async def test_group_chat_checkpoint_runtime_only() -> None:
assert baseline_output is not None
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert len(checkpoints) > 0, "Runtime-only checkpointing should have created checkpoints"
@@ -656,8 +656,8 @@ async def test_group_chat_checkpoint_runtime_overrides_buildtime() -> None:
assert baseline_output is not None
buildtime_checkpoints = await buildtime_storage.list_checkpoints()
runtime_checkpoints = await runtime_storage.list_checkpoints()
buildtime_checkpoints = await buildtime_storage.list_checkpoints(workflow_name=wf.name)
runtime_checkpoints = await runtime_storage.list_checkpoints(workflow_name=wf.name)
assert len(runtime_checkpoints) > 0, "Runtime storage should have checkpoints"
assert len(buildtime_checkpoints) == 0, "Build-time storage should have no checkpoints when overridden"
python/packages/orchestrations/tests/test_magentic.py
+16 -14
View File
@@ -362,7 +362,7 @@ async def test_magentic_checkpoint_resume_round_trip():
assert req_event is not None
assert isinstance(req_event.data, MagenticPlanReviewRequest)
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert checkpoints
checkpoints.sort(key=lambda cp: cp.timestamp)
resume_checkpoint = checkpoints[-1]
@@ -605,8 +605,9 @@ async def test_agent_executor_invoke_with_assistants_client_messages():
async def _collect_checkpoints(
storage: InMemoryCheckpointStorage,
workflow_name: str,
) -> list[WorkflowCheckpoint]:
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=workflow_name)
assert checkpoints
checkpoints.sort(key=lambda cp: cp.timestamp)
return checkpoints
@@ -619,12 +620,13 @@ async def test_magentic_checkpoint_resume_inner_loop_superstep():
participants=[StubThreadAgent()], checkpoint_storage=storage, manager=InvokeOnceManager()
).build()
async for event in workflow.run("inner-loop task", stream=True):
if event.type == "output":
break
async for _ in workflow.run("inner-loop task", stream=True):
continue
checkpoints = await _collect_checkpoints(storage)
inner_loop_checkpoint = next(cp for cp in checkpoints if cp.metadata.get("superstep") == 1) # type: ignore[reportUnknownMemberType]
checkpoints = await _collect_checkpoints(storage, workflow.name)
# The first checkpoint is after the manager has run.
# The second checkpoint is after the participant has run.
inner_loop_checkpoint = checkpoints[1]
resumed = MagenticBuilder(
participants=[StubThreadAgent()], checkpoint_storage=storage, manager=InvokeOnceManager()
@@ -651,7 +653,7 @@ async def test_magentic_checkpoint_resume_from_saved_state():
if event.type == "output":
break
checkpoints = await _collect_checkpoints(storage)
checkpoints = await _collect_checkpoints(storage, workflow.name)
# Verify we can resume from the last saved checkpoint
resumed_state = checkpoints[-1] # Use the last checkpoint
@@ -688,7 +690,7 @@ async def test_magentic_checkpoint_resume_rejects_participant_renames():
assert req_event is not None
assert isinstance(req_event.data, MagenticPlanReviewRequest)
checkpoints = await _collect_checkpoints(storage)
checkpoints = await _collect_checkpoints(storage, workflow.name)
target_checkpoint = checkpoints[-1]
renamed_workflow = MagenticBuilder(
@@ -772,7 +774,7 @@ async def test_magentic_checkpoint_runtime_only() -> None:
assert baseline_output is not None
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert len(checkpoints) > 0, "Runtime-only checkpointing should have created checkpoints"
@@ -806,8 +808,8 @@ async def test_magentic_checkpoint_runtime_overrides_buildtime() -> None:
assert baseline_output is not None
buildtime_checkpoints = await buildtime_storage.list_checkpoints()
runtime_checkpoints = await runtime_storage.list_checkpoints()
buildtime_checkpoints = await buildtime_storage.list_checkpoints(workflow_name=wf.name)
runtime_checkpoints = await runtime_storage.list_checkpoints(workflow_name=wf.name)
assert len(runtime_checkpoints) > 0, "Runtime storage should have checkpoints"
assert len(buildtime_checkpoints) == 0, "Build-time storage should have no checkpoints when overridden"
@@ -856,13 +858,13 @@ async def test_magentic_checkpoint_restore_no_duplicate_history():
break
# Get checkpoint
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert len(checkpoints) > 0, "Should have created checkpoints"
latest_checkpoint = checkpoints[-1]
# Load checkpoint and verify no duplicates in state
checkpoint_data = await storage.load_checkpoint(latest_checkpoint.checkpoint_id)
checkpoint_data = await storage.load(latest_checkpoint.checkpoint_id)
assert checkpoint_data is not None
# Check the magentic_context in the checkpoint
python/packages/orchestrations/tests/test_sequential.py
+6 -14
View File
@@ -146,14 +146,10 @@ async def test_sequential_checkpoint_resume_round_trip() -> None:
assert baseline_output is not None
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert checkpoints
checkpoints.sort(key=lambda cp: cp.timestamp)
resume_checkpoint = next(
(cp for cp in checkpoints if (cp.metadata or {}).get("checkpoint_type") == "superstep"),
checkpoints[-1],
)
resume_checkpoint = checkpoints[0]
resumed_agents = (_EchoAgent(id="agent1", name="A1"), _EchoAgent(id="agent2", name="A2"))
wf_resume = SequentialBuilder(participants=list(resumed_agents), checkpoint_storage=storage).build()
@@ -189,14 +185,10 @@ async def test_sequential_checkpoint_runtime_only() -> None:
assert baseline_output is not None
checkpoints = await storage.list_checkpoints()
checkpoints = await storage.list_checkpoints(workflow_name=wf.name)
assert checkpoints
checkpoints.sort(key=lambda cp: cp.timestamp)
resume_checkpoint = next(
(cp for cp in checkpoints if (cp.metadata or {}).get("checkpoint_type") == "superstep"),
checkpoints[-1],
)
resume_checkpoint = checkpoints[0]
resumed_agents = (_EchoAgent(id="agent1", name="A1"), _EchoAgent(id="agent2", name="A2"))
wf_resume = SequentialBuilder(participants=list(resumed_agents)).build()
@@ -240,8 +232,8 @@ async def test_sequential_checkpoint_runtime_overrides_buildtime() -> None:
assert baseline_output is not None
buildtime_checkpoints = await buildtime_storage.list_checkpoints()
runtime_checkpoints = await runtime_storage.list_checkpoints()
buildtime_checkpoints = await buildtime_storage.list_checkpoints(workflow_name=wf.name)
runtime_checkpoints = await runtime_storage.list_checkpoints(workflow_name=wf.name)
assert len(runtime_checkpoints) > 0, "Runtime storage should have checkpoints"
assert len(buildtime_checkpoints) == 0, "Build-time storage should have no checkpoints when overridden"