agent/agent-framework
1
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/agent_framework_orchestrations/_group_chat.py
+2 -3
View File
@@ -32,7 +32,6 @@ from agent_framework import Agent, AgentThread, Message, SupportsAgentRun
from agent_framework._workflows._agent_executor import AgentExecutor, AgentExecutorRequest, AgentExecutorResponse
from agent_framework._workflows._agent_utils import resolve_agent_id
from agent_framework._workflows._checkpoint import CheckpointStorage
from agent_framework._workflows._conversation_state import decode_chat_messages, encode_chat_messages
from agent_framework._workflows._executor import Executor
from agent_framework._workflows._workflow import Workflow
from agent_framework._workflows._workflow_builder import WorkflowBuilder
@@ -476,7 +475,7 @@ class AgentBasedGroupChatOrchestrator(BaseGroupChatOrchestrator):
async def on_checkpoint_save(self) -> dict[str, Any]:
"""Capture current orchestrator state for checkpointing."""
state = await super().on_checkpoint_save()
state["cache"] = encode_chat_messages(self._cache)
state["cache"] = self._cache
serialized_thread = await self._thread.serialize()
state["thread"] = serialized_thread
@@ -486,7 +485,7 @@ class AgentBasedGroupChatOrchestrator(BaseGroupChatOrchestrator):
async def on_checkpoint_restore(self, state: dict[str, Any]) -> None:
"""Restore executor state from checkpoint."""
await super().on_checkpoint_restore(state)
self._cache = decode_chat_messages(state.get("cache", []))
self._cache = state.get("cache", [])
serialized_thread = state.get("thread")
if serialized_thread:
self._thread = await self._agent.deserialize_thread(serialized_thread)
python/packages/orchestrations/agent_framework_orchestrations/_orchestration_state.py
+6 -8
View File
@@ -59,15 +59,14 @@ class OrchestrationState:
Returns:
Dict with encoded conversation and metadata for persistence
"""
from agent_framework._workflows._conversation_state import encode_chat_messages
result: dict[str, Any] = {
"conversation": encode_chat_messages(self.conversation),
"conversation": self.conversation,
"round_index": self.round_index,
"orchestrator_name": self.orchestrator_name,
"metadata": dict(self.metadata),
}
if self.task is not None:
result["task"] = encode_chat_messages([self.task])[0]
result["task"] = self.task
return result
@classmethod
@@ -80,16 +79,15 @@ class OrchestrationState:
Returns:
Restored OrchestrationState instance
"""
from agent_framework._workflows._conversation_state import decode_chat_messages
task = None
if "task" in data:
decoded_tasks = decode_chat_messages([data["task"]])
decoded_tasks = [data["task"]]
task = decoded_tasks[0] if decoded_tasks else None
return cls(
conversation=decode_chat_messages(data.get("conversation", [])),
conversation=data.get("conversation", []),
round_index=data.get("round_index", 0),
orchestrator_name=data.get("orchestrator_name", ""),
metadata=dict(data.get("metadata", {})),
task=task,
)
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"