diff --git a/python/packages/main/agent_framework/workflow/__init__.py b/python/packages/main/agent_framework/workflow/__init__.py index 14669fecab..bbfd05c2a9 100644 --- a/python/packages/main/agent_framework/workflow/__init__.py +++ b/python/packages/main/agent_framework/workflow/__init__.py @@ -28,6 +28,10 @@ _IMPORTS = [ "RequestInfoMessage", "WorkflowRunResult", "Workflow", + "FileCheckpointStorage", + "InMemoryCheckpointStorage", + "CheckpointStorage", + "WorkflowCheckpoint", ] diff --git a/python/packages/main/agent_framework/workflow/__init__.pyi b/python/packages/main/agent_framework/workflow/__init__.pyi index 30ea6fde9c..6506d4a936 100644 --- a/python/packages/main/agent_framework/workflow/__init__.pyi +++ b/python/packages/main/agent_framework/workflow/__init__.pyi @@ -6,15 +6,19 @@ from agent_framework_workflow import ( AgentExecutorResponse, AgentRunEvent, AgentRunStreamingEvent, + CheckpointStorage, Executor, ExecutorCompletedEvent, ExecutorEvent, ExecutorInvokeEvent, + FileCheckpointStorage, + InMemoryCheckpointStorage, RequestInfoEvent, RequestInfoExecutor, RequestInfoMessage, Workflow, WorkflowBuilder, + WorkflowCheckpoint, WorkflowCompletedEvent, WorkflowContext, WorkflowEvent, @@ -30,15 +34,19 @@ __all__ = [ "AgentExecutorResponse", "AgentRunEvent", "AgentRunStreamingEvent", + "CheckpointStorage", "Executor", "ExecutorCompletedEvent", "ExecutorEvent", "ExecutorInvokeEvent", + "FileCheckpointStorage", + "InMemoryCheckpointStorage", "RequestInfoEvent", "RequestInfoExecutor", "RequestInfoMessage", "Workflow", "WorkflowBuilder", + "WorkflowCheckpoint", "WorkflowCompletedEvent", "WorkflowContext", "WorkflowEvent", diff --git a/python/packages/workflow/agent_framework_workflow/__init__.py b/python/packages/workflow/agent_framework_workflow/__init__.py index ee9404e07b..1d59918525 100644 --- a/python/packages/workflow/agent_framework_workflow/__init__.py +++ b/python/packages/workflow/agent_framework_workflow/__init__.py @@ -2,6 +2,15 @@ import importlib.metadata +from ._checkpoint import ( + CheckpointStorage, + FileCheckpointStorage, + InMemoryCheckpointStorage, + WorkflowCheckpoint, +) +from ._const import ( + DEFAULT_MAX_ITERATIONS, +) from ._events import ( AgentRunEvent, AgentRunStreamingEvent, @@ -22,6 +31,11 @@ from ._executor import ( RequestInfoMessage, handler, ) +from ._runner_context import ( + InProcRunnerContext, + Message, + RunnerContext, +) from ._validation import ( EdgeDuplicationError, GraphConnectivityError, @@ -40,26 +54,34 @@ except importlib.metadata.PackageNotFoundError: __all__ = [ + "DEFAULT_MAX_ITERATIONS", "AgentExecutor", "AgentExecutorRequest", "AgentExecutorResponse", "AgentRunEvent", "AgentRunStreamingEvent", + "CheckpointStorage", "EdgeDuplicationError", "Executor", "ExecutorCompletedEvent", "ExecutorEvent", "ExecutorInvokeEvent", + "FileCheckpointStorage", "GraphConnectivityError", + "InMemoryCheckpointStorage", + "InProcRunnerContext", + "Message", "RequestInfoEvent", "RequestInfoEvent", "RequestInfoExecutor", "RequestInfoExecutor", "RequestInfoMessage", + "RunnerContext", "TypeCompatibilityError", "ValidationTypeEnum", "Workflow", "WorkflowBuilder", + "WorkflowCheckpoint", "WorkflowCompletedEvent", "WorkflowContext", "WorkflowEvent", diff --git a/python/packages/workflow/agent_framework_workflow/_checkpoint.py b/python/packages/workflow/agent_framework_workflow/_checkpoint.py new file mode 100644 index 0000000000..5ef3ae3b06 --- /dev/null +++ b/python/packages/workflow/agent_framework_workflow/_checkpoint.py @@ -0,0 +1,199 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import json +import logging +import os +import uuid +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Protocol + +from ._const import DEFAULT_MAX_ITERATIONS + +logger = logging.getLogger(__name__) + + +@dataclass(slots=True) +class WorkflowCheckpoint: + """Represents a complete checkpoint of workflow state.""" + + checkpoint_id: str = field(default_factory=lambda: str(uuid.uuid4())) + workflow_id: str = "" + timestamp: str = field(default_factory=lambda: datetime.now(timezone.utc).isoformat()) + + # Core workflow state + messages: dict[str, list[dict[str, Any]]] = field(default_factory=dict) # type: ignore[misc] + shared_state: dict[str, Any] = field(default_factory=dict) # type: ignore[misc] + executor_states: dict[str, dict[str, Any]] = field(default_factory=dict) # type: ignore[misc] + + # Runtime state + iteration_count: int = 0 + max_iterations: int = DEFAULT_MAX_ITERATIONS + + # Metadata + metadata: dict[str, Any] = field(default_factory=dict) # type: ignore[misc] + version: str = "1.0" + + def to_dict(self) -> dict[str, Any]: + return asdict(self) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "WorkflowCheckpoint": + return cls(**data) + + +class CheckpointStorage(Protocol): + """Protocol for checkpoint storage backends.""" + + async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str: + """Save a checkpoint and return its ID.""" + ... + + async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None: + """Load a checkpoint by ID.""" + ... + + async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]: + """List checkpoint IDs. If workflow_id is provided, filter by that workflow.""" + ... + + async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]: + """List checkpoint objects. If workflow_id is provided, filter by that workflow.""" + ... + + async def delete_checkpoint(self, checkpoint_id: str) -> bool: + """Delete a checkpoint by ID.""" + ... + + +class InMemoryCheckpointStorage: + """In-memory checkpoint storage for testing and development.""" + + def __init__(self): + """Initialize the memory storage.""" + self._checkpoints: dict[str, WorkflowCheckpoint] = {} + + async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str: + """Save a checkpoint and return its ID.""" + self._checkpoints[checkpoint.checkpoint_id] = checkpoint + logger.debug(f"Saved checkpoint {checkpoint.checkpoint_id} to memory") + return checkpoint.checkpoint_id + + async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None: + """Load a checkpoint by ID.""" + checkpoint = self._checkpoints.get(checkpoint_id) + if checkpoint: + logger.debug(f"Loaded checkpoint {checkpoint_id} from memory") + return checkpoint + + async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]: + """List checkpoint IDs. If workflow_id is provided, filter by that workflow.""" + if workflow_id is None: + return list(self._checkpoints.keys()) + return [cp.checkpoint_id for cp in self._checkpoints.values() if cp.workflow_id == workflow_id] + + async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]: + """List checkpoint objects. If workflow_id is provided, filter by that workflow.""" + if workflow_id is None: + return list(self._checkpoints.values()) + return [cp for cp in self._checkpoints.values() if cp.workflow_id == workflow_id] + + async def delete_checkpoint(self, checkpoint_id: str) -> bool: + """Delete a checkpoint by ID.""" + if checkpoint_id in self._checkpoints: + del self._checkpoints[checkpoint_id] + logger.debug(f"Deleted checkpoint {checkpoint_id} from memory") + return True + return False + + +class FileCheckpointStorage: + """File-based checkpoint storage for persistence.""" + + def __init__(self, storage_path: str | Path): + """Initialize the file storage.""" + self.storage_path = Path(storage_path) + self.storage_path.mkdir(parents=True, exist_ok=True) + logger.info(f"Initialized file checkpoint storage at {self.storage_path}") + + async def save_checkpoint(self, checkpoint: WorkflowCheckpoint) -> str: + """Save a checkpoint and return its ID.""" + file_path = self.storage_path / f"{checkpoint.checkpoint_id}.json" + checkpoint_dict = asdict(checkpoint) + + def _write_atomic() -> None: + tmp_path = file_path.with_suffix(".json.tmp") + with open(tmp_path, "w") as f: + json.dump(checkpoint_dict, f, indent=2, ensure_ascii=False) + os.replace(tmp_path, file_path) + + await asyncio.to_thread(_write_atomic) + + logger.info(f"Saved checkpoint {checkpoint.checkpoint_id} to {file_path}") + return checkpoint.checkpoint_id + + async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None: + """Load a checkpoint by ID.""" + file_path = self.storage_path / f"{checkpoint_id}.json" + + if not file_path.exists(): + return None + + def _read() -> dict[str, Any]: + with open(file_path) as f: + return json.load(f) + + checkpoint_dict = await asyncio.to_thread(_read) + + checkpoint = WorkflowCheckpoint(**checkpoint_dict) + logger.info(f"Loaded checkpoint {checkpoint_id} from {file_path}") + return checkpoint + + async def list_checkpoint_ids(self, workflow_id: str | None = None) -> list[str]: + """List checkpoint IDs. If workflow_id is provided, filter by that workflow.""" + + def _list_ids() -> list[str]: + checkpoint_ids: list[str] = [] + for file_path in self.storage_path.glob("*.json"): + try: + with open(file_path) as f: + data = json.load(f) + if workflow_id is None or data.get("workflow_id") == workflow_id: + checkpoint_ids.append(data.get("checkpoint_id", file_path.stem)) + except Exception as e: + logger.warning(f"Failed to read checkpoint file {file_path}: {e}") + return checkpoint_ids + + return await asyncio.to_thread(_list_ids) + + async def list_checkpoints(self, workflow_id: str | None = None) -> list[WorkflowCheckpoint]: + """List checkpoint objects. If workflow_id is provided, filter by that workflow.""" + + def _list_checkpoints() -> list[WorkflowCheckpoint]: + checkpoints: list[WorkflowCheckpoint] = [] + for file_path in self.storage_path.glob("*.json"): + try: + with open(file_path) as f: + data = json.load(f) + if workflow_id is None or data.get("workflow_id") == workflow_id: + checkpoints.append(WorkflowCheckpoint.from_dict(data)) + except Exception as e: + logger.warning(f"Failed to read checkpoint file {file_path}: {e}") + return checkpoints + + return await asyncio.to_thread(_list_checkpoints) + + async def delete_checkpoint(self, checkpoint_id: str) -> bool: + """Delete a checkpoint by ID.""" + file_path = self.storage_path / f"{checkpoint_id}.json" + + def _delete() -> bool: + if file_path.exists(): + file_path.unlink() + logger.info(f"Deleted checkpoint {checkpoint_id} from {file_path}") + return True + return False + + return await asyncio.to_thread(_delete) diff --git a/python/packages/workflow/agent_framework_workflow/_const.py b/python/packages/workflow/agent_framework_workflow/_const.py new file mode 100644 index 0000000000..b426692725 --- /dev/null +++ b/python/packages/workflow/agent_framework_workflow/_const.py @@ -0,0 +1,3 @@ +# Copyright (c) Microsoft. All rights reserved. + +DEFAULT_MAX_ITERATIONS = 100 # Default maximum iterations for workflow execution. diff --git a/python/packages/workflow/agent_framework_workflow/_runner.py b/python/packages/workflow/agent_framework_workflow/_runner.py index dd53d8c4da..1688a2f5fa 100644 --- a/python/packages/workflow/agent_framework_workflow/_runner.py +++ b/python/packages/workflow/agent_framework_workflow/_runner.py @@ -4,16 +4,16 @@ import asyncio import logging from collections import defaultdict from collections.abc import AsyncIterable +from typing import Any from ._edge import Edge from ._events import WorkflowEvent +from ._executor import Executor from ._runner_context import Message, RunnerContext from ._shared_state import SharedState logger = logging.getLogger(__name__) -DEFAULT_MAX_ITERATIONS = 100 - class Runner: """A class to run a workflow in Pregel supersteps.""" @@ -23,8 +23,9 @@ class Runner: edges: list[Edge], shared_state: SharedState, ctx: RunnerContext, - max_iterations: int = DEFAULT_MAX_ITERATIONS, - ) -> None: + max_iterations: int = 100, + workflow_id: str | None = None, + ): """Initialize the runner with edges, shared state, and context. Args: @@ -32,66 +33,107 @@ class Runner: shared_state: The shared state for the workflow. ctx: The runner context for the workflow. max_iterations: The maximum number of iterations to run. + workflow_id: The workflow ID for checkpointing. """ self._edge_map = self._parse_edges(edges) self._ctx = ctx self._iteration = 0 self._max_iterations = max_iterations self._shared_state = shared_state - self._is_running = False + self._workflow_id = workflow_id + self._running = False + self._resumed_from_checkpoint = False # Track whether we resumed + + # Set workflow ID in context if provided + if workflow_id: + self._ctx.set_workflow_id(workflow_id) @property def context(self) -> RunnerContext: """Get the workflow context.""" return self._ctx + def mark_resumed(self, iteration: int | None = None, max_iterations: int | None = None) -> None: + """Mark the runner as having resumed from a checkpoint. + + Optionally set the current iteration and max iterations. + """ + self._resumed_from_checkpoint = True + if iteration is not None: + self._iteration = iteration + if max_iterations is not None: + self._max_iterations = max_iterations + async def run_until_convergence(self) -> AsyncIterable[WorkflowEvent]: """Run the workflow until no more messages are sent.""" + if self._running: + raise RuntimeError("Runner is already running.") + + self._running = True try: - if self._is_running: - raise RuntimeError("Runner is already running.") - self._is_running = True + # Process any events from initial execution before checkpointing + if await self._ctx.has_events(): + logger.info("Processing events from initial execution") + events = await self._ctx.drain_events() + for event in events: + logger.info(f"Yielding initial event: {event}") + yield event + + # Create first checkpoint if there are messages from initial execution + if await self._ctx.has_messages() and self._ctx.has_checkpointing(): + if not self._resumed_from_checkpoint: + logger.info("Creating checkpoint after initial execution") + await self._create_checkpoint_if_enabled("after_initial_execution") + else: + logger.info("Skipping 'after_initial_execution' checkpoint because we resumed from a checkpoint") + + # Initialize context with starting iteration state + await self._update_context_with_shared_state() + while self._iteration < self._max_iterations: + logger.info(f"Starting superstep {self._iteration + 1}") await self._run_iteration() self._iteration += 1 + # Update context with current iteration state immediately + await self._update_context_with_shared_state() + + logger.info(f"Completed superstep {self._iteration}") + + # Process events first before any checkpointing if await self._ctx.has_events(): + logger.info("Processing events before checkpointing") events = await self._ctx.drain_events() for event in events: + logger.debug(f"Yielding event: {event}") yield event + # Create checkpoint after each superstep iteration + await self._create_checkpoint_if_enabled(f"superstep_{self._iteration}") + if not await self._ctx.has_messages(): break - else: + + if self._iteration >= self._max_iterations and await self._ctx.has_messages(): raise RuntimeError(f"Runner did not converge after {self._max_iterations} iterations.") - finally: - self._is_running = False + + logger.info(f"Workflow completed after {self._iteration} supersteps") self._iteration = 0 + self._resumed_from_checkpoint = False # Reset resume flag for next run + finally: + self._running = False async def _run_iteration(self): - """Run a superstep of the workflow execution.""" - async def _deliver_messages(source_executor_id: str, messages: list[Message]) -> None: - """Deliver messages to the executors. - - Outer loop to concurrently deliver messages from all sources to their targets. - """ - async def _deliver_messages_inner( edge: Edge, messages: list[Message], ) -> None: - """Deliver messages to a specific target executor. - - Inner loop to deliver messages to a specific target executor. - """ for message in messages: if message.target_id is not None and message.target_id != edge.target_id: continue - if not edge.can_handle(message.data): continue - await edge.send_message(message, self._shared_state, self._ctx) associated_edges = self._edge_map.get(source_executor_id, []) @@ -105,6 +147,127 @@ class Runner: ] await asyncio.gather(*tasks) + async def _create_checkpoint_if_enabled(self, checkpoint_type: str) -> str | None: + """Create a checkpoint if checkpointing is enabled and attach a label and metadata.""" + if not self._ctx.has_checkpointing(): + return None + + try: + # Auto-snapshot executor states + await self._auto_snapshot_executor_states() + await self._update_context_with_shared_state() + checkpoint_category = "initial" if checkpoint_type == "after_initial_execution" else "superstep" + metadata = { + "superstep": self._iteration, + "checkpoint_type": checkpoint_category, + } + checkpoint_id = await self._ctx.create_checkpoint(metadata=metadata) + logger.info(f"Created {checkpoint_type} checkpoint: {checkpoint_id}") + return checkpoint_id + except Exception as e: + logger.warning(f"Failed to create {checkpoint_type} checkpoint: {e}") + return None + + async def _auto_snapshot_executor_states(self) -> None: + """Populate executor state by calling snapshot hooks on executors if available. + + Convention: + - If an executor defines an async or sync method `snapshot_state(self) -> dict`, use it. + - Else if it has a plain attribute `state` that is a dict, use that. + Only JSON-serializable dicts should be provided by executors. + """ + executors: dict[str, Executor] = {} + for edge_list in self._edge_map.values(): + for edge in edge_list: + executors[edge.source.id] = edge.source + executors[edge.target.id] = edge.target + for exec_id, executor in executors.items(): + state_dict: dict[str, Any] | None = None + snapshot = getattr(executor, "snapshot_state", None) + try: + if callable(snapshot): + maybe = snapshot() + if asyncio.iscoroutine(maybe): # type: ignore[arg-type] + maybe = await maybe # type: ignore[assignment] + if isinstance(maybe, dict): + state_dict = maybe # type: ignore[assignment] + else: + state_attr = getattr(executor, "state", None) + if isinstance(state_attr, dict): + state_dict = state_attr # type: ignore[assignment] + except Exception as ex: # pragma: no cover + logger.debug(f"Executor {exec_id} snapshot_state failed: {ex}") + if state_dict is not None: + try: + await self._ctx.set_state(exec_id, state_dict) + except Exception as ex: # pragma: no cover + logger.debug(f"Failed to persist state for executor {exec_id}: {ex}") + + async def _update_context_with_shared_state(self) -> None: + if not self._ctx.has_checkpointing(): + return + + try: + current_state = await self._ctx.get_checkpoint_state() + + shared_state_data = {} + async with self._shared_state.hold(): + if hasattr(self._shared_state, "_state"): + shared_state_data = dict(self._shared_state._state) # type: ignore[attr-defined] + + current_state["shared_state"] = shared_state_data + current_state["iteration_count"] = self._iteration + current_state["max_iterations"] = self._max_iterations + + await self._ctx.set_checkpoint_state(current_state) + except Exception as e: + logger.warning(f"Failed to update context with shared state: {e}") + + async def restore_from_checkpoint(self, checkpoint_id: str) -> bool: + """Restore workflow state from a checkpoint. + + Args: + checkpoint_id: The ID of the checkpoint to restore from + + Returns: + True if restoration was successful, False otherwise + """ + if not self._ctx.has_checkpointing(): + logger.warning("Context does not support checkpointing") + return False + + try: + success = await self._ctx.restore_from_checkpoint(checkpoint_id) + if not success: + return False + + await self._restore_shared_state_from_context() + self.mark_resumed() # mark resumed; iteration/max already restored from context + logger.info(f"Successfully restored workflow from checkpoint: {checkpoint_id}") + return True + except Exception as e: + logger.error(f"Failed to restore from checkpoint {checkpoint_id}: {e}") + return False + + async def _restore_shared_state_from_context(self) -> None: + if not self._ctx.has_checkpointing(): + return + + try: + restored_state = await self._ctx.get_checkpoint_state() + + shared_state_data = restored_state.get("shared_state", {}) + if shared_state_data and hasattr(self._shared_state, "_state"): + async with self._shared_state.hold(): + self._shared_state._state.clear() # type: ignore[attr-defined] + self._shared_state._state.update(shared_state_data) # type: ignore[attr-defined] + + self._iteration = restored_state.get("iteration_count", 0) + self._max_iterations = restored_state.get("max_iterations", self._max_iterations) + + except Exception as e: + logger.warning(f"Failed to restore shared state from context: {e}") + def _parse_edges(self, edges: list[Edge]) -> dict[str, list[Edge]]: """Parse the edges of the workflow into a more convenient format. diff --git a/python/packages/workflow/agent_framework_workflow/_runner_context.py b/python/packages/workflow/agent_framework_workflow/_runner_context.py index bcb1e1ba34..e6f1fab27a 100644 --- a/python/packages/workflow/agent_framework_workflow/_runner_context.py +++ b/python/packages/workflow/agent_framework_workflow/_runner_context.py @@ -1,11 +1,15 @@ # Copyright (c) Microsoft. All rights reserved. import logging +import uuid from collections import defaultdict from dataclasses import dataclass -from typing import Any, Protocol, TypeVar, runtime_checkable +from typing import Any, Protocol, TypedDict, TypeVar, runtime_checkable +from ._checkpoint import CheckpointStorage, WorkflowCheckpoint +from ._const import DEFAULT_MAX_ITERATIONS from ._events import WorkflowEvent +from ._shared_state import SharedState logger = logging.getLogger(__name__) @@ -21,9 +25,21 @@ class Message: target_id: str | None = None +class CheckpointState(TypedDict): + messages: dict[str, list[dict[str, Any]]] + shared_state: dict[str, Any] + executor_states: dict[str, dict[str, Any]] + iteration_count: int + max_iterations: int + + @runtime_checkable class RunnerContext(Protocol): - """Protocol for the execution context used by the runner.""" + """Protocol for the execution context used by the runner. + + A single context that supports messaging, events, and optional checkpointing. + If checkpoint storage is not configured, checkpoint methods may raise. + """ async def send_message(self, message: Message) -> None: """Send a message from the executor to the context. @@ -73,43 +89,208 @@ class RunnerContext(Protocol): """ ... + async def set_state(self, executor_id: str, state: dict[str, Any]) -> None: + """Set the state for a specific executor. -class InProcRunnerContext(RunnerContext): - """In-process execution context for local execution of workflows.""" + Args: + executor_id: The ID of the executor whose state is being set. + state: The state data to be set for the executor. + """ + ... - def __init__(self): - """Initialize the in-process execution context.""" + async def get_state(self, executor_id: str) -> dict[str, Any] | None: + """Get the state for a specific executor. + + Args: + executor_id: The ID of the executor whose state is being retrieved. + + Returns: + The state data for the executor, or None if not found. + """ + ... + + # Checkpointing capability + def has_checkpointing(self) -> bool: + """Check if the context supports checkpointing. + + Returns: + True if checkpointing is supported, False otherwise. + """ + ... + + # Checkpointing APIs (optional, enabled by storage) + def set_workflow_id(self, workflow_id: str) -> None: + """Set the workflow ID for the context.""" + ... + + def reset_for_new_run(self, workflow_shared_state: SharedState | None = None) -> None: + """Reset the context for a new workflow run.""" + ... + + async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str: + """Create a checkpoint of the current workflow state. + + Args: + metadata: Optional metadata to associate with the checkpoint. + """ + ... + + async def restore_from_checkpoint(self, checkpoint_id: str) -> bool: + """Restore the context from a checkpoint. + + Args: + checkpoint_id: The ID of the checkpoint to restore from. + + Returns: + True if the restoration was successful, False otherwise. + """ + ... + + async def get_checkpoint_state(self) -> CheckpointState: + """Get the current state of the context suitable for checkpointing.""" + ... + + async def set_checkpoint_state(self, state: CheckpointState) -> None: + """Set the state of the context from a checkpoint. + + Args: + state: The state data to set for the context. + """ + ... + + +class InProcRunnerContext: + """In-process execution context for local execution and optional checkpointing.""" + + def __init__(self, checkpoint_storage: CheckpointStorage | None = None): + """Initialize the in-process execution context. + + Args: + checkpoint_storage: Optional storage to enable checkpointing. + """ self._messages: defaultdict[str, list[Message]] = defaultdict(list) self._events: list[WorkflowEvent] = [] + # Checkpointing configuration/state + self._checkpoint_storage = checkpoint_storage + self._workflow_id: str | None = None + self._shared_state: dict[str, Any] = {} + self._executor_states: dict[str, dict[str, Any]] = {} + self._iteration_count: int = 0 + self._max_iterations: int = 100 + async def send_message(self, message: Message) -> None: - """Send a message from the executor to the context.""" self._messages[message.source_id].append(message) async def drain_messages(self) -> dict[str, list[Message]]: - """Drain all messages from the context.""" messages = dict(self._messages) self._messages.clear() return messages async def has_messages(self) -> bool: - """Check if there are any messages in the context.""" return bool(self._messages) async def add_event(self, event: WorkflowEvent) -> None: - """Add an event to the execution context. - - Args: - event: The event to be added. - """ self._events.append(event) async def drain_events(self) -> list[WorkflowEvent]: - """Drain all events from the context.""" events = self._events.copy() self._events.clear() return events async def has_events(self) -> bool: - """Check if there are any events in the context.""" return bool(self._events) + + async def set_state(self, executor_id: str, state: dict[str, Any]) -> None: + self._executor_states[executor_id] = state + + async def get_state(self, executor_id: str) -> dict[str, Any] | None: + return self._executor_states.get(executor_id) + + def has_checkpointing(self) -> bool: + return self._checkpoint_storage is not None + + def set_workflow_id(self, workflow_id: str) -> None: + self._workflow_id = workflow_id + + def reset_for_new_run(self, workflow_shared_state: "SharedState | None" = None) -> None: + self._messages.clear() + self._events.clear() + self._shared_state.clear() + self._executor_states.clear() + self._iteration_count = 0 + if workflow_shared_state is not None and hasattr(workflow_shared_state, "_state"): + workflow_shared_state._state.clear() # type: ignore[attr-defined] + + async def create_checkpoint(self, metadata: dict[str, Any] | None = None) -> str: + if not self._checkpoint_storage: + raise ValueError("Checkpoint storage not configured") + + wf_id = self._workflow_id or str(uuid.uuid4()) + self._workflow_id = wf_id + state = await self.get_checkpoint_state() + + checkpoint = WorkflowCheckpoint( + workflow_id=wf_id, + messages=state["messages"], + shared_state=state.get("shared_state", {}), + executor_states=state.get("executor_states", {}), + iteration_count=state.get("iteration_count", 0), + max_iterations=state.get("max_iterations", DEFAULT_MAX_ITERATIONS), + metadata=metadata or {}, + ) + checkpoint_id = await self._checkpoint_storage.save_checkpoint(checkpoint) + logger.info(f"Created checkpoint {checkpoint_id} for workflow {wf_id}'") + return checkpoint_id + + async def restore_from_checkpoint(self, checkpoint_id: str) -> bool: + if not self._checkpoint_storage: + raise ValueError("Checkpoint storage not configured") + + checkpoint = await self._checkpoint_storage.load_checkpoint(checkpoint_id) + if not checkpoint: + logger.error(f"Checkpoint {checkpoint_id} not found") + return False + + state: CheckpointState = { + "messages": checkpoint.messages, + "shared_state": checkpoint.shared_state, + "executor_states": checkpoint.executor_states, + "iteration_count": checkpoint.iteration_count, + "max_iterations": checkpoint.max_iterations, + } + await self.set_checkpoint_state(state) + self._workflow_id = checkpoint.workflow_id + logger.info(f"Restored state from checkpoint {checkpoint_id}'") + return True + + async def get_checkpoint_state(self) -> CheckpointState: + serializable_messages: dict[str, list[dict[str, Any]]] = {} + for source_id, message_list in self._messages.items(): + serializable_messages[source_id] = [ + {"data": msg.data, "source_id": msg.source_id, "target_id": msg.target_id} for msg in message_list + ] + return { + "messages": serializable_messages, + "shared_state": self._shared_state, + "executor_states": self._executor_states, + "iteration_count": self._iteration_count, + "max_iterations": self._max_iterations, + } + + async def set_checkpoint_state(self, state: CheckpointState) -> None: + self._messages.clear() + messages_data = state.get("messages", {}) + for source_id, message_list in messages_data.items(): + self._messages[source_id] = [ + Message( + data=msg.get("data"), + source_id=msg.get("source_id", ""), + target_id=msg.get("target_id"), + ) + for msg in message_list + ] + self._shared_state = state.get("shared_state", {}) + self._executor_states = state.get("executor_states", {}) + self._iteration_count = state.get("iteration_count", 0) + self._max_iterations = state.get("max_iterations", 100) diff --git a/python/packages/workflow/agent_framework_workflow/_workflow.py b/python/packages/workflow/agent_framework_workflow/_workflow.py index 84cc8178b5..bc3ae00f19 100644 --- a/python/packages/workflow/agent_framework_workflow/_workflow.py +++ b/python/packages/workflow/agent_framework_workflow/_workflow.py @@ -1,15 +1,19 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import logging import sys +import uuid from collections.abc import AsyncIterable, Callable, Sequence from typing import Any +from ._checkpoint import CheckpointStorage +from ._const import DEFAULT_MAX_ITERATIONS from ._edge import Edge from ._events import RequestInfoEvent, WorkflowCompletedEvent, WorkflowEvent from ._executor import Executor, RequestInfoExecutor -from ._runner import DEFAULT_MAX_ITERATIONS, Runner -from ._runner_context import InProcRunnerContext, RunnerContext +from ._runner import Runner +from ._runner_context import CheckpointState, InProcRunnerContext, RunnerContext from ._shared_state import SharedState from ._validation import validate_workflow_graph from ._workflow_context import WorkflowContext @@ -20,6 +24,9 @@ else: from typing_extensions import Self # pragma: no cover +logger = logging.getLogger(__name__) + + class WorkflowRunResult(list[WorkflowEvent]): """A list of events generated during the workflow execution in non-streaming mode.""" @@ -48,6 +55,9 @@ class WorkflowRunResult(list[WorkflowEvent]): return [event for event in self if isinstance(event, RequestInfoEvent)] +# region Workflow + + class Workflow: """A class representing a workflow that can be executed. @@ -77,7 +87,11 @@ class Workflow: } self._shared_state = SharedState() - self._runner = Runner(self._edges, self._shared_state, runner_context, max_iterations=max_iterations) + + workflow_id = str(uuid.uuid4()) + self._runner = Runner( + self._edges, self._shared_state, runner_context, max_iterations=max_iterations, workflow_id=workflow_id + ) @property def edges(self) -> list[Edge]: @@ -101,7 +115,7 @@ class Workflow: return list(self._executors.values()) async def run_streaming(self, message: Any) -> AsyncIterable[WorkflowEvent]: - """Send a message to the starting executor of the workflow and stream the events generated by the workflow. + """Run the workflow with a starting message and stream events. Args: message: The message to be sent to the starting executor. @@ -109,6 +123,9 @@ class Workflow: Yields: WorkflowEvent: The events generated during the workflow execution. """ + # Reset context for a new run if supported + self._runner.context.reset_for_new_run(self._shared_state) + executor = self._start_executor if isinstance(executor, str): executor = self._get_executor_by_id(executor) @@ -117,15 +134,71 @@ class Workflow: message, WorkflowContext( executor.id, - [ - # Using the workflow class name as the source executor ID when - # delivering the first message to the starting executor - self.__class__.__name__ - ], + [self.__class__.__name__], self._shared_state, self._runner.context, ), ) + + async for event in self._runner.run_until_convergence(): + yield event + + async def run_streaming_from_checkpoint( + self, + checkpoint_id: str, + checkpoint_storage: CheckpointStorage | None = None, + responses: dict[str, Any] | None = None, + ) -> AsyncIterable[WorkflowEvent]: + """Resume workflow execution from a checkpoint and stream events. + + Args: + checkpoint_id: The ID of the checkpoint to restore from. + checkpoint_storage: Optional checkpoint storage to use for restoration. + If not provided, the workflow must have been built with checkpointing enabled. + responses: Optional dictionary of responses to inject into the workflow + after restoration. Keys are request IDs, values are response data. + + Yields: + WorkflowEvent: Events generated during workflow execution. + + Raises: + ValueError: If neither checkpoint_storage is provided nor checkpointing is enabled. + RuntimeError: If checkpoint restoration fails. + """ + has_checkpointing = self._runner.context.has_checkpointing() + + if not has_checkpointing and not checkpoint_storage: + raise ValueError( + "Cannot restore from checkpoint: either provide checkpoint_storage parameter " + "or build workflow with WorkflowBuilder.with_checkpointing(checkpoint_storage)." + ) + + if has_checkpointing: + # restore via Runner so shared state and iteration are synchronized + restored = await self._runner.restore_from_checkpoint(checkpoint_id) + else: + if checkpoint_storage is None: + raise ValueError("checkpoint_storage cannot be None.") + restored = await self._restore_from_external_checkpoint(checkpoint_id, checkpoint_storage) + + if not restored: + raise RuntimeError(f"Failed to restore from checkpoint: {checkpoint_id}") + + if responses: + request_info_executor = self._get_executor_by_id(RequestInfoExecutor.EXECUTOR_ID) + if isinstance(request_info_executor, RequestInfoExecutor): + for request_id, response_data in responses.items(): + await request_info_executor.handle_response( + response_data, + request_id, + WorkflowContext( + request_info_executor.id, + [self.__class__.__name__], + self._shared_state, + self._runner.context, + ), + ) + async for event in self._runner.run_until_convergence(): yield event @@ -150,11 +223,7 @@ class Workflow: request_id, WorkflowContext( request_info_executor.id, - [ - # Using the workflow class name as the source executor ID when - # delivering the first message to the starting executor - self.__class__.__name__ - ], + [self.__class__.__name__], self._shared_state, self._runner.context, ), @@ -177,6 +246,33 @@ class Workflow: events = [event async for event in self.run_streaming(message)] return WorkflowRunResult(events) + async def run_from_checkpoint( + self, + checkpoint_id: str, + checkpoint_storage: CheckpointStorage | None = None, + responses: dict[str, Any] | None = None, + ) -> WorkflowRunResult: + """Resume workflow execution from a checkpoint. + + Args: + checkpoint_id: The ID of the checkpoint to restore from. + checkpoint_storage: Optional checkpoint storage to use for restoration. + If not provided, the workflow must have been built with checkpointing enabled. + responses: Optional dictionary of responses to inject into the workflow + after restoration. Keys are request IDs, values are response data. + + Returns: + A WorkflowRunResult instance containing a list of events generated during the workflow execution. + + Raises: + ValueError: If neither checkpoint_storage is provided nor checkpointing is enabled. + RuntimeError: If checkpoint restoration fails. + """ + events = [ + event async for event in self.run_streaming_from_checkpoint(checkpoint_id, checkpoint_storage, responses) + ] + return WorkflowRunResult(events) + async def send_responses(self, responses: dict[str, Any]) -> WorkflowRunResult: """Send responses back to the workflow. @@ -202,6 +298,104 @@ class Workflow: raise ValueError(f"Executor with ID {executor_id} not found.") return self._executors[executor_id] + async def _restore_from_external_checkpoint( + self, checkpoint_id: str, checkpoint_storage: CheckpointStorage + ) -> bool: + """Restore workflow state from an external checkpoint storage. + + This method implements the state transfer pattern: load checkpoint data + from external storage and transfer it to the current workflow context. + + Args: + checkpoint_id: The ID of the checkpoint to restore from. + checkpoint_storage: The checkpoint storage to load from. + + Returns: + True if restoration was successful, False otherwise. + """ + try: + checkpoint = await checkpoint_storage.load_checkpoint(checkpoint_id) + if not checkpoint: + return False + + temp_context = InProcRunnerContext(checkpoint_storage) + state: CheckpointState = { + "messages": checkpoint.messages, + "shared_state": checkpoint.shared_state, + "executor_states": checkpoint.executor_states, + "iteration_count": checkpoint.iteration_count, + "max_iterations": checkpoint.max_iterations, + } + + await temp_context.set_checkpoint_state(state) + restored_state = await temp_context.get_checkpoint_state() + await self._transfer_state_to_context(restored_state) + + # Also set runner iteration/max so superstep numbering continues + self._runner.mark_resumed(iteration=checkpoint.iteration_count, max_iterations=checkpoint.max_iterations) + + return True + + except Exception as e: + import logging + + logger = logging.getLogger(__name__) + logger.error(f"Failed to restore from external checkpoint {checkpoint_id}: {e}") + return False + + async def _transfer_state_to_context(self, restored_state: CheckpointState) -> None: + """Transfer restored checkpoint state into the current workflow runtime. + + This transfers: + - messages -> into the current RunnerContext so delivery can continue + - executor_states -> into the current RunnerContext so ctx.get_state() works after resume + - shared_state -> into the Workflow's SharedState so executors can read values set before the checkpoint + """ + # Best-effort restoration + # Restore shared state so downstream executors can read values (e.g., original_input) + try: + shared_state_data = restored_state.get("shared_state", {}) + if shared_state_data and hasattr(self._shared_state, "_state"): + async with self._shared_state.hold(): + self._shared_state._state.clear() # type: ignore[attr-defined] + self._shared_state._state.update(shared_state_data) # type: ignore[attr-defined] + except Exception as exc: # pragma: no cover + logger.debug("Failed to restore shared_state during external restore: %s", exc) + + # Restore executor states into the context so ctx.get_state() calls after resume succeed + try: + executor_states = restored_state.get("executor_states", {}) + for exec_id, state in executor_states.items(): + try: + await self._runner.context.set_state(exec_id, state) + except Exception as exc: # pragma: no cover - ignore per-executor failures + logger.debug("Failed to restore executor state for %s during external restore: %s", exec_id, exc) + except Exception as exc: # pragma: no cover + logger.debug("Failed to iterate executor_states during external restore: %s", exc) + + # Transfer pending messages into the context for delivery in the next superstep + messages_data = restored_state["messages"] + for _, message_list in messages_data.items(): + for msg_data in message_list: + source_any = msg_data.get("source_id", "") + source_id: str = source_any if isinstance(source_any, str) else str(source_any) + if not source_id: + source_id = "" + target_raw = msg_data.get("target_id") + target_id: str | None = ( + target_raw if target_raw is None or isinstance(target_raw, str) else str(target_raw) + ) + + # Build and send Message via runner context + from ._runner_context import Message as _Msg + + await self._runner.context.send_message( + _Msg(data=msg_data.get("data"), source_id=source_id, target_id=target_id) + ) + + +# region WorkflowBuilder + class WorkflowBuilder: """A builder class for constructing workflows. @@ -209,11 +403,12 @@ class WorkflowBuilder: This class provides methods to add edges and set the starting executor for the workflow. """ - def __init__(self): + def __init__(self, max_iterations: int = DEFAULT_MAX_ITERATIONS): """Initialize the WorkflowBuilder with an empty list of edges and no starting executor.""" self._edges: list[Edge] = [] self._start_executor: Executor | str | None = None - self._max_iterations: int = DEFAULT_MAX_ITERATIONS + self._checkpoint_storage: CheckpointStorage | None = None + self._max_iterations: int = max_iterations def add_edge( self, @@ -329,6 +524,15 @@ class WorkflowBuilder: self._max_iterations = max_iterations return self + def with_checkpointing(self, checkpoint_storage: CheckpointStorage) -> "Self": + """Enable checkpointing with the specified storage. + + Args: + checkpoint_storage: The checkpoint storage to use. + """ + self._checkpoint_storage = checkpoint_storage + return self + def build(self) -> Workflow: """Build and return the constructed workflow. @@ -347,4 +551,9 @@ class WorkflowBuilder: validate_workflow_graph(self._edges, self._start_executor) - return Workflow(self._edges, self._start_executor, InProcRunnerContext(), self._max_iterations) + context = InProcRunnerContext(self._checkpoint_storage) + + return Workflow(self._edges, self._start_executor, context, self._max_iterations) + + +# endregion diff --git a/python/packages/workflow/agent_framework_workflow/_workflow_context.py b/python/packages/workflow/agent_framework_workflow/_workflow_context.py index ca24748543..7b5a070888 100644 --- a/python/packages/workflow/agent_framework_workflow/_workflow_context.py +++ b/python/packages/workflow/agent_framework_workflow/_workflow_context.py @@ -89,3 +89,18 @@ class WorkflowContext: def shared_state(self) -> SharedState: """Get the shared state.""" return self._shared_state + + async def set_state(self, state: dict[str, Any]) -> None: + """Persist this executor's state into the checkpointable context. + + Executors call this with a JSON-serializable dict capturing the minimal + state needed to resume. It replaces any previously stored state. + """ + if hasattr(self._runner_context, "set_state"): + await self._runner_context.set_state(self._executor_id, state) # type: ignore[arg-type] + + async def get_state(self) -> dict[str, Any] | None: + """Retrieve previously persisted state for this executor, if any.""" + if hasattr(self._runner_context, "get_state"): + return await self._runner_context.get_state(self._executor_id) # type: ignore[return-value] + return None diff --git a/python/packages/workflow/tests/test_checkpoint.py b/python/packages/workflow/tests/test_checkpoint.py new file mode 100644 index 0000000000..4ae5304672 --- /dev/null +++ b/python/packages/workflow/tests/test_checkpoint.py @@ -0,0 +1,334 @@ +# Copyright (c) Microsoft. All rights reserved. + +import json +import tempfile +from datetime import datetime, timezone +from pathlib import Path + +from agent_framework_workflow._checkpoint import ( + FileCheckpointStorage, + InMemoryCheckpointStorage, + WorkflowCheckpoint, +) + + +def test_workflow_checkpoint_default_values(): + checkpoint = WorkflowCheckpoint() + + assert checkpoint.checkpoint_id != "" + assert checkpoint.workflow_id == "" + assert checkpoint.timestamp != "" + assert checkpoint.messages == {} + assert checkpoint.shared_state == {} + assert checkpoint.executor_states == {} + assert checkpoint.iteration_count == 0 + assert checkpoint.max_iterations == 100 + assert checkpoint.metadata == {} + assert checkpoint.version == "1.0" + + +def test_workflow_checkpoint_custom_values(): + custom_timestamp = datetime.now(timezone.utc).isoformat() + checkpoint = WorkflowCheckpoint( + checkpoint_id="test-checkpoint-123", + workflow_id="test-workflow-456", + timestamp=custom_timestamp, + messages={"executor1": [{"data": "test"}]}, + shared_state={"key": "value"}, + executor_states={"executor1": {"state": "active"}}, + iteration_count=5, + max_iterations=50, + metadata={"test": True}, + version="2.0", + ) + + assert checkpoint.checkpoint_id == "test-checkpoint-123" + assert checkpoint.workflow_id == "test-workflow-456" + assert checkpoint.timestamp == custom_timestamp + assert checkpoint.messages == {"executor1": [{"data": "test"}]} + assert checkpoint.shared_state == {"key": "value"} + assert checkpoint.executor_states == {"executor1": {"state": "active"}} + assert checkpoint.iteration_count == 5 + assert checkpoint.max_iterations == 50 + assert checkpoint.metadata == {"test": True} + assert checkpoint.version == "2.0" + + +async def test_memory_checkpoint_storage_save_and_load(): + storage = InMemoryCheckpointStorage() + checkpoint = WorkflowCheckpoint(workflow_id="test-workflow", messages={"executor1": [{"data": "hello"}]}) + + # Save checkpoint + saved_id = await storage.save_checkpoint(checkpoint) + assert saved_id == checkpoint.checkpoint_id + + # Load checkpoint + loaded_checkpoint = await storage.load_checkpoint(checkpoint.checkpoint_id) + assert loaded_checkpoint is not None + assert loaded_checkpoint.checkpoint_id == checkpoint.checkpoint_id + assert loaded_checkpoint.workflow_id == checkpoint.workflow_id + assert loaded_checkpoint.messages == checkpoint.messages + + +async def test_memory_checkpoint_storage_load_nonexistent(): + storage = InMemoryCheckpointStorage() + + result = await storage.load_checkpoint("nonexistent-id") + assert result is None + + +async def test_memory_checkpoint_storage_list_checkpoints(): + storage = InMemoryCheckpointStorage() + + # Create checkpoints for different workflows + checkpoint1 = WorkflowCheckpoint(workflow_id="workflow-1") + checkpoint2 = WorkflowCheckpoint(workflow_id="workflow-1") + checkpoint3 = WorkflowCheckpoint(workflow_id="workflow-2") + + await storage.save_checkpoint(checkpoint1) + await storage.save_checkpoint(checkpoint2) + await storage.save_checkpoint(checkpoint3) + + # Test list_checkpoint_ids for workflow-1 + workflow1_checkpoint_ids = await storage.list_checkpoint_ids("workflow-1") + assert len(workflow1_checkpoint_ids) == 2 + assert checkpoint1.checkpoint_id in workflow1_checkpoint_ids + assert checkpoint2.checkpoint_id in workflow1_checkpoint_ids + + # Test list_checkpoints for workflow-1 (returns objects) + workflow1_checkpoints = await storage.list_checkpoints("workflow-1") + assert len(workflow1_checkpoints) == 2 + assert all(isinstance(cp, WorkflowCheckpoint) for cp in workflow1_checkpoints) + assert {cp.checkpoint_id for cp in workflow1_checkpoints} == {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id} + + # Test list_checkpoint_ids for workflow-2 + workflow2_checkpoint_ids = await storage.list_checkpoint_ids("workflow-2") + assert len(workflow2_checkpoint_ids) == 1 + assert checkpoint3.checkpoint_id in workflow2_checkpoint_ids + + # Test list_checkpoints for workflow-2 (returns objects) + workflow2_checkpoints = await storage.list_checkpoints("workflow-2") + assert len(workflow2_checkpoints) == 1 + assert workflow2_checkpoints[0].checkpoint_id == checkpoint3.checkpoint_id + + # Test list_checkpoint_ids for non-existent workflow + empty_checkpoint_ids = await storage.list_checkpoint_ids("nonexistent-workflow") + assert len(empty_checkpoint_ids) == 0 + + # Test list_checkpoints for non-existent workflow + empty_checkpoints = await storage.list_checkpoints("nonexistent-workflow") + assert len(empty_checkpoints) == 0 + + # Test list_checkpoint_ids without workflow filter (all checkpoints) + all_checkpoint_ids = await storage.list_checkpoint_ids() + assert len(all_checkpoint_ids) == 3 + expected_ids = {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id, checkpoint3.checkpoint_id} + assert expected_ids.issubset(set(all_checkpoint_ids)) + + # Test list_checkpoints without workflow filter (all checkpoints) + all_checkpoints = await storage.list_checkpoints() + assert len(all_checkpoints) == 3 + assert all(isinstance(cp, WorkflowCheckpoint) for cp in all_checkpoints) + + +async def test_memory_checkpoint_storage_delete(): + storage = InMemoryCheckpointStorage() + checkpoint = WorkflowCheckpoint(workflow_id="test-workflow") + + # Save checkpoint + await storage.save_checkpoint(checkpoint) + assert await storage.load_checkpoint(checkpoint.checkpoint_id) is not None + + # Delete checkpoint + result = await storage.delete_checkpoint(checkpoint.checkpoint_id) + assert result is True + + # Verify deletion + assert await storage.load_checkpoint(checkpoint.checkpoint_id) is None + + # Try to delete again + result = await storage.delete_checkpoint(checkpoint.checkpoint_id) + assert result is False + + +async def test_file_checkpoint_storage_save_and_load(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + checkpoint = WorkflowCheckpoint( + workflow_id="test-workflow", + messages={"executor1": [{"data": "hello", "source_id": "test", "target_id": None}]}, + shared_state={"key": "value"}, + ) + + # Save checkpoint + saved_id = await storage.save_checkpoint(checkpoint) + assert saved_id == checkpoint.checkpoint_id + + # Verify file was created + file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json" + assert file_path.exists() + + # Load checkpoint + loaded_checkpoint = await storage.load_checkpoint(checkpoint.checkpoint_id) + assert loaded_checkpoint is not None + assert loaded_checkpoint.checkpoint_id == checkpoint.checkpoint_id + assert loaded_checkpoint.workflow_id == checkpoint.workflow_id + assert loaded_checkpoint.messages == checkpoint.messages + assert loaded_checkpoint.shared_state == checkpoint.shared_state + + +async def test_file_checkpoint_storage_load_nonexistent(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + result = await storage.load_checkpoint("nonexistent-id") + assert result is None + + +async def test_file_checkpoint_storage_list_checkpoints(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + # Create checkpoints for different workflows + checkpoint1 = WorkflowCheckpoint(workflow_id="workflow-1") + checkpoint2 = WorkflowCheckpoint(workflow_id="workflow-1") + checkpoint3 = WorkflowCheckpoint(workflow_id="workflow-2") + + await storage.save_checkpoint(checkpoint1) + await storage.save_checkpoint(checkpoint2) + await storage.save_checkpoint(checkpoint3) + + # Test list_checkpoint_ids for workflow-1 + workflow1_checkpoint_ids = await storage.list_checkpoint_ids("workflow-1") + assert len(workflow1_checkpoint_ids) == 2 + assert checkpoint1.checkpoint_id in workflow1_checkpoint_ids + assert checkpoint2.checkpoint_id in workflow1_checkpoint_ids + + # Test list_checkpoints for workflow-1 (returns objects) + workflow1_checkpoints = await storage.list_checkpoints("workflow-1") + assert len(workflow1_checkpoints) == 2 + assert all(isinstance(cp, WorkflowCheckpoint) for cp in workflow1_checkpoints) + checkpoint_ids = {cp.checkpoint_id for cp in workflow1_checkpoints} + assert checkpoint_ids == {checkpoint1.checkpoint_id, checkpoint2.checkpoint_id} + + # Test list_checkpoint_ids for workflow-2 + workflow2_checkpoint_ids = await storage.list_checkpoint_ids("workflow-2") + assert len(workflow2_checkpoint_ids) == 1 + assert checkpoint3.checkpoint_id in workflow2_checkpoint_ids + + # Test list_checkpoints for workflow-2 (returns objects) + workflow2_checkpoints = await storage.list_checkpoints("workflow-2") + assert len(workflow2_checkpoints) == 1 + assert workflow2_checkpoints[0].checkpoint_id == checkpoint3.checkpoint_id + + # Test list all checkpoints + all_checkpoint_ids = await storage.list_checkpoint_ids() + assert len(all_checkpoint_ids) == 3 + + all_checkpoints = await storage.list_checkpoints() + assert len(all_checkpoints) == 3 + assert all(isinstance(cp, WorkflowCheckpoint) for cp in all_checkpoints) + + +async def test_file_checkpoint_storage_delete(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + checkpoint = WorkflowCheckpoint(workflow_id="test-workflow") + + # Save checkpoint + await storage.save_checkpoint(checkpoint) + file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json" + assert file_path.exists() + + # Delete checkpoint + result = await storage.delete_checkpoint(checkpoint.checkpoint_id) + assert result is True + assert not file_path.exists() + + # Try to delete again + result = await storage.delete_checkpoint(checkpoint.checkpoint_id) + assert result is False + + +async def test_file_checkpoint_storage_directory_creation(): + with tempfile.TemporaryDirectory() as temp_dir: + nested_path = Path(temp_dir) / "nested" / "checkpoint" / "storage" + storage = FileCheckpointStorage(nested_path) + + # Directory should be created + assert nested_path.exists() + assert nested_path.is_dir() + + # Should be able to save checkpoints + checkpoint = WorkflowCheckpoint(workflow_id="test") + await storage.save_checkpoint(checkpoint) + + file_path = nested_path / f"{checkpoint.checkpoint_id}.json" + assert file_path.exists() + + +async def test_file_checkpoint_storage_corrupted_file(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + # Create a corrupted JSON file + corrupted_file = Path(temp_dir) / "corrupted.json" + with open(corrupted_file, "w") as f: # noqa: ASYNC230 + f.write("{ invalid json }") + + # list_checkpoints should handle the corrupted file gracefully + checkpoints = await storage.list_checkpoints("any-workflow") + assert checkpoints == [] + + +async def test_file_checkpoint_storage_json_serialization(): + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + # Create checkpoint with complex nested data + checkpoint = WorkflowCheckpoint( + workflow_id="complex-workflow", + messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]}, + shared_state={"list": [1, 2, 3], "dict": {"a": "b", "c": {"d": "e"}}, "bool": True, "null": None}, + executor_states={"executor1": {"state": "active", "config": {"timeout": 30, "retries": 3}}}, + ) + + # Save and load + await storage.save_checkpoint(checkpoint) + loaded = await storage.load_checkpoint(checkpoint.checkpoint_id) + + assert loaded is not None + assert loaded.messages == checkpoint.messages + assert loaded.shared_state == checkpoint.shared_state + assert loaded.executor_states == checkpoint.executor_states + + # Verify the JSON file is properly formatted + file_path = Path(temp_dir) / f"{checkpoint.checkpoint_id}.json" + with open(file_path) as f: # noqa: ASYNC230 + data = json.load(f) + + assert data["messages"]["executor1"][0]["data"]["nested"]["value"] == 42 + assert data["shared_state"]["list"] == [1, 2, 3] + assert data["shared_state"]["bool"] is True + assert data["shared_state"]["null"] is None + + +def test_checkpoint_storage_protocol_compliance(): + # This test ensures both implementations have all required methods + memory_storage = InMemoryCheckpointStorage() + + with tempfile.TemporaryDirectory() as temp_dir: + file_storage = FileCheckpointStorage(temp_dir) + + for storage in [memory_storage, file_storage]: + # Test that all protocol methods exist and are callable + assert hasattr(storage, "save_checkpoint") + assert callable(storage.save_checkpoint) + assert hasattr(storage, "load_checkpoint") + assert callable(storage.load_checkpoint) + assert hasattr(storage, "list_checkpoint_ids") + assert callable(storage.list_checkpoint_ids) + assert hasattr(storage, "list_checkpoints") + assert callable(storage.list_checkpoints) + assert hasattr(storage, "delete_checkpoint") + assert callable(storage.delete_checkpoint) diff --git a/python/packages/workflow/tests/test_workflow.py b/python/packages/workflow/tests/test_workflow.py index b5be8dc2ca..1aa7a0446b 100644 --- a/python/packages/workflow/tests/test_workflow.py +++ b/python/packages/workflow/tests/test_workflow.py @@ -1,10 +1,12 @@ # Copyright (c) Microsoft. All rights reserved. +import tempfile from dataclasses import dataclass import pytest from agent_framework.workflow import ( Executor, + FileCheckpointStorage, RequestInfoEvent, RequestInfoExecutor, RequestInfoMessage, @@ -15,6 +17,8 @@ from agent_framework.workflow import ( handler, ) +from agent_framework_workflow import Message + @dataclass class MockMessage: @@ -275,3 +279,278 @@ async def test_fan_in(): completed_event = events.get_completed_event() assert completed_event is not None and completed_event.data == 4 + + +@pytest.fixture +def simple_executor() -> Executor: + class SimpleExecutor(Executor): + @handler + async def handle_message(self, message: Message, context: WorkflowContext) -> None: + pass + + return SimpleExecutor("test_executor") + + +async def test_workflow_with_checkpointing_enabled(simple_executor: Executor): + """Test that a workflow can be built with checkpointing enabled.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + # Build workflow with checkpointing - should not raise any errors + workflow = ( + WorkflowBuilder() + .add_edge(simple_executor, simple_executor) # Self-loop to satisfy graph requirements + .set_start_executor(simple_executor) + .with_checkpointing(storage) + .build() + ) + + # Verify workflow was created and can run + test_message = Message(data="test message", source_id="test", target_id=None) + result = await workflow.run(test_message) + assert result is not None + + +async def test_workflow_checkpointing_not_enabled_for_external_restore(simple_executor: Executor): + """Test that external checkpoint restoration fails when workflow doesn't support checkpointing.""" + # Build workflow WITHOUT checkpointing + workflow = ( + WorkflowBuilder() + .add_edge(simple_executor, simple_executor) # Self-loop to satisfy graph requirements + .set_start_executor(simple_executor) + .build() + ) + + # Attempt to restore from checkpoint without providing external storage should fail + try: + [event async for event in workflow.run_streaming_from_checkpoint("fake-checkpoint-id")] + raise AssertionError("Expected ValueError to be raised") + except ValueError as e: + assert "Cannot restore from checkpoint" in str(e) + assert "either provide checkpoint_storage parameter" in str(e) + + +async def test_workflow_run_stream_from_checkpoint_no_checkpointing_enabled(simple_executor: Executor): + # Build workflow WITHOUT checkpointing + workflow = ( + WorkflowBuilder() + .add_edge(simple_executor, simple_executor) # Self-loop to satisfy graph requirements + .set_start_executor(simple_executor) + .build() + ) + + # Attempt to run from checkpoint should fail + try: + async for _ in workflow.run_streaming_from_checkpoint("fake_checkpoint_id"): + pass + raise AssertionError("Expected ValueError to be raised") + except ValueError as e: + assert "Cannot restore from checkpoint" in str(e) + assert "either provide checkpoint_storage parameter" in str(e) + + +async def test_workflow_run_stream_from_checkpoint_invalid_checkpoint(simple_executor: Executor): + """Test that attempting to restore from a non-existent checkpoint fails appropriately.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + # Build workflow with checkpointing + workflow = ( + WorkflowBuilder() + .add_edge(simple_executor, simple_executor) # Self-loop to satisfy graph requirements + .set_start_executor(simple_executor) + .with_checkpointing(storage) + .build() + ) + + # Attempt to run from non-existent checkpoint should fail + try: + async for _ in workflow.run_streaming_from_checkpoint("nonexistent_checkpoint_id"): + pass + raise AssertionError("Expected RuntimeError to be raised") + except RuntimeError as e: + assert "Failed to restore from checkpoint" in str(e) + + +async def test_workflow_run_stream_from_checkpoint_with_external_storage(simple_executor: Executor): + """Test that external checkpoint storage can be provided for restoration.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + # Create a test checkpoint manually in storage + from agent_framework_workflow._checkpoint import WorkflowCheckpoint + + test_checkpoint = WorkflowCheckpoint( + workflow_id="test-workflow", + messages={}, + shared_state={}, + executor_states={}, + iteration_count=0, + max_iterations=100, + ) + checkpoint_id = await storage.save_checkpoint(test_checkpoint) + + # Create a workflow WITHOUT checkpointing + workflow_without_checkpointing = ( + WorkflowBuilder().add_edge(simple_executor, simple_executor).set_start_executor(simple_executor).build() + ) + + # Resume from checkpoint using external storage parameter + try: + events: list[WorkflowEvent] = [] + async for event in workflow_without_checkpointing.run_streaming_from_checkpoint( + checkpoint_id, checkpoint_storage=storage + ): + events.append(event) + if len(events) >= 2: # Limit to avoid infinite loops + break + except Exception: + # Expected since we have minimal setup, but method should accept the parameters + pass + + +async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Executor): + """Test the non-streaming run_from_checkpoint method.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + # Create a test checkpoint manually in storage + from agent_framework_workflow._checkpoint import WorkflowCheckpoint + + test_checkpoint = WorkflowCheckpoint( + workflow_id="test-workflow", + messages={}, + shared_state={}, + executor_states={}, + iteration_count=0, + max_iterations=100, + ) + checkpoint_id = await storage.save_checkpoint(test_checkpoint) + + # Build workflow with checkpointing + workflow = ( + WorkflowBuilder() + .add_edge(simple_executor, simple_executor) + .set_start_executor(simple_executor) + .with_checkpointing(storage) + .build() + ) + + # Test non-streaming run_from_checkpoint method + result = await workflow.run_from_checkpoint(checkpoint_id) + assert isinstance(result, list) # Should return WorkflowRunResult which extends list + assert hasattr(result, "get_completed_event") # Should have WorkflowRunResult methods + + +async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executor: Executor): + """Test that run_streaming_from_checkpoint accepts responses parameter.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + # Create a test checkpoint manually in storage + from agent_framework_workflow._checkpoint import WorkflowCheckpoint + + test_checkpoint = WorkflowCheckpoint( + workflow_id="test-workflow", + messages={}, + shared_state={}, + executor_states={}, + iteration_count=0, + max_iterations=100, + ) + checkpoint_id = await storage.save_checkpoint(test_checkpoint) + + # Build workflow with checkpointing + workflow = ( + WorkflowBuilder() + .add_edge(simple_executor, simple_executor) + .set_start_executor(simple_executor) + .with_checkpointing(storage) + .build() + ) + + # Test that run_stream_from_checkpoint accepts responses parameter + responses = {"request_123": {"data": "test_response"}} + + try: + events: list[WorkflowEvent] = [] + async for event in workflow.run_streaming_from_checkpoint(checkpoint_id, responses=responses): + events.append(event) + if len(events) >= 2: # Limit to avoid infinite loops + break + except Exception: + # Expected since we have minimal setup, but method should accept the parameters + pass + + +@dataclass +class StateTrackingMessage: + """A message that tracks state for testing context reset behavior.""" + + data: str + run_id: str + + +class StateTrackingExecutor(Executor): + """An executor that tracks state in shared state to test context reset behavior.""" + + @handler(output_types=[]) + async def handle_message(self, message: StateTrackingMessage, ctx: WorkflowContext) -> None: + """Handle the message and track it in shared state.""" + # Get existing messages from shared state + try: + existing_messages = await ctx.get_shared_state("processed_messages") + except KeyError: + existing_messages = [] + + # Record this message + message_record = f"{message.run_id}:{message.data}" + existing_messages.append(message_record) # type: ignore + + # Update shared state + await ctx.set_shared_state("processed_messages", existing_messages) + + # Complete workflow with current shared state + await ctx.add_event(WorkflowCompletedEvent(data=existing_messages.copy())) # type: ignore + + +async def test_workflow_multiple_runs_no_state_collision(): + """Test that running the same workflow instance multiple times doesn't have state collision.""" + with tempfile.TemporaryDirectory() as temp_dir: + storage = FileCheckpointStorage(temp_dir) + + # Create executor that tracks state in shared state + state_executor = StateTrackingExecutor("state_executor") + + # Build workflow with checkpointing + workflow = ( + WorkflowBuilder() + .add_edge(state_executor, state_executor) # Self-loop to satisfy graph requirements + .set_start_executor(state_executor) + .with_checkpointing(storage) + .build() + ) + + # Run 1: Should only see messages from run 1 + result1 = await workflow.run(StateTrackingMessage(data="message1", run_id="run1")) + completed1 = result1.get_completed_event() + assert completed1 is not None + assert completed1.data == ["run1:message1"] + + # Run 2: Should only see messages from run 2, not run 1 + result2 = await workflow.run(StateTrackingMessage(data="message2", run_id="run2")) + completed2 = result2.get_completed_event() + assert completed2 is not None + assert completed2.data == ["run2:message2"] # Should NOT contain run1 data + + # Run 3: Should only see messages from run 3 + result3 = await workflow.run(StateTrackingMessage(data="message3", run_id="run3")) + completed3 = result3.get_completed_event() + assert completed3 is not None + assert completed3.data == ["run3:message3"] # Should NOT contain run1 or run2 data + + # Verify that each run only processed its own message + # This confirms that the checkpointable context properly resets between runs + assert completed1.data != completed2.data + assert completed2.data != completed3.data + assert completed1.data != completed3.data diff --git a/python/samples/getting_started/workflow/step_02_simple_workflow_condition.py b/python/samples/getting_started/workflow/step_02_simple_workflow_condition.py index 1d8625c02c..c6f4f21f60 100644 --- a/python/samples/getting_started/workflow/step_02_simple_workflow_condition.py +++ b/python/samples/getting_started/workflow/step_02_simple_workflow_condition.py @@ -3,13 +3,7 @@ import asyncio from dataclasses import dataclass -from agent_framework.workflow import ( - Executor, - WorkflowBuilder, - WorkflowCompletedEvent, - WorkflowContext, - handler, -) +from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler """ The following sample demonstrates a basic workflow with two executors diff --git a/python/samples/getting_started/workflow/step_06_map_reduce.py b/python/samples/getting_started/workflow/step_06_map_reduce.py index 8879bce982..2e8c5ca995 100644 --- a/python/samples/getting_started/workflow/step_06_map_reduce.py +++ b/python/samples/getting_started/workflow/step_06_map_reduce.py @@ -3,24 +3,11 @@ import ast import asyncio import os -import sys from collections import defaultdict from dataclasses import dataclass import aiofiles -from agent_framework.workflow import ( - Executor, - WorkflowBuilder, - WorkflowCompletedEvent, - WorkflowContext, - handler, -) - -if sys.version_info >= (3, 12): - pass # pragma: no cover -else: - pass # pragma: no cover - +from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler """ The following sample demonstrates a basic map reduce workflow that @@ -119,7 +106,8 @@ class Map(Executor): data: An instance of SplitCompleted signaling the map step can be started. ctx: The execution context containing the shared state and other information. """ - # Retrieve the data to be processed from the shared state.# Define a key for the shared state to store the data to be processed + # Retrieve the data to be processed from the shared state. + # Define a key for the shared state to store the data to be processed data_to_be_processed: list[str] = await ctx.get_shared_state(SHARED_STATE_DATA_KEY) chunk_start, chunk_end = await ctx.get_shared_state(self.id) diff --git a/python/samples/getting_started/workflow/step_07_checkpoint_resume.py b/python/samples/getting_started/workflow/step_07_checkpoint_resume.py new file mode 100644 index 0000000000..15a6493368 --- /dev/null +++ b/python/samples/getting_started/workflow/step_07_checkpoint_resume.py @@ -0,0 +1,215 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import os +from pathlib import Path + +from agent_framework.workflow import ( + Executor, + FileCheckpointStorage, + WorkflowBuilder, + WorkflowCompletedEvent, + WorkflowContext, + handler, +) + +""" +Demonstrates workflow checkpointing, shared state, and resumption at superstep boundaries. + +Flow: +1) UpperCaseExecutor: "hello world" -> "HELLO WORLD" (writes shared_state: original_input, upper_output) +2) ReverseTextExecutor: "HELLO WORLD" -> "DLROW OLLEH" +3) LowerCaseExecutor: "DLROW OLLEH" -> "dlrow olleh" (reads shared_state, emits WorkflowCompletedEvent) + +Initial run checkpoints: +- after_initial_execution: messages from upper_case_executor +- superstep_1: messages from reverse_text_executor +- superstep_2: no messages (final events only) + +Resume: +- Resume from the checkpoint containing "DLROW OLLEH" (superstep_1); only LowerCaseExecutor runs. +- Iteration continues from the checkpoint; one checkpoint is created after the resumed superstep. +""" + +# Define the temporary directory for storing checkpoints +DIR = os.path.dirname(__file__) +TEMP_DIR = os.path.join(DIR, "tmp", "checkpoints") +os.makedirs(TEMP_DIR, exist_ok=True) + + +class UpperCaseExecutor(Executor): + @handler(output_types=[str]) + async def to_upper_case(self, text: str, ctx: WorkflowContext) -> None: + result = text.upper() + print(f"UpperCaseExecutor: '{text}' -> '{result}'") + # Persist executor state into checkpointable context + prev = await ctx.get_state() or {} + count = int(prev.get("count", 0)) + 1 + await ctx.set_state({ + "count": count, + "last_input": text, + "last_output": result, + }) + # Write to shared_state so downstream executors (and checkpoints) can see it + await ctx.set_shared_state("original_input", text) + await ctx.set_shared_state("upper_output", result) + await ctx.send_message(result) + + +class LowerCaseExecutor(Executor): + @handler(output_types=[str]) + async def to_lower_case(self, text: str, ctx: WorkflowContext) -> None: + result = text.lower() + print(f"LowerCaseExecutor: '{text}' -> '{result}'") + # Read from shared_state written by UpperCaseExecutor + orig = await ctx.get_shared_state("original_input") + upper = await ctx.get_shared_state("upper_output") + print(f"LowerCaseExecutor (shared_state): original_input='{orig}', upper_output='{upper}'") + # Persist executor state into checkpointable context + prev = await ctx.get_state() or {} + count = int(prev.get("count", 0)) + 1 + await ctx.set_state({ + "count": count, + "last_input": text, + "last_output": result, + "final": True, + }) + await ctx.add_event(WorkflowCompletedEvent(result)) + + +class ReverseTextExecutor(Executor): + def __init__(self, id: str): + """Initialize the executor with an ID.""" + super().__init__(id=id) + + @handler(output_types=[str]) + async def reverse_text(self, text: str, ctx: WorkflowContext) -> None: + result = text[::-1] + print(f"ReverseTextExecutor: '{text}' -> '{result}'") + # Persist executor state into checkpointable context + prev = await ctx.get_state() or {} + count = int(prev.get("count", 0)) + 1 + await ctx.set_state({ + "count": count, + "last_input": text, + "last_output": result, + }) + await ctx.send_message(result) + + +async def find_checkpoint_with_message( + checkpoint_storage: FileCheckpointStorage, workflow_id: str, needle: str +) -> str | None: + """Find the checkpoint that contains a message data value exactly equal to 'needle'.""" + checkpoints = await checkpoint_storage.list_checkpoints(workflow_id=workflow_id) + # Sort by timestamp ascending so earlier checkpoints appear first + checkpoints.sort(key=lambda cp: cp.timestamp) + for checkpoint in checkpoints: + for executor_messages in checkpoint.messages.values(): + for message in executor_messages: + if message.get("data") == needle: + return checkpoint.checkpoint_id + return None + + +async def main(): + # Clear existing checkpoints in this sample directory + checkpoint_dir = Path(TEMP_DIR) + for file in checkpoint_dir.glob("*.json"): + file.unlink() + + upper_case_executor = UpperCaseExecutor(id="upper_case_executor") + reverse_text_executor = ReverseTextExecutor(id="reverse_text_executor") + lower_case_executor = LowerCaseExecutor(id="lower_case_executor") + + checkpoint_storage = FileCheckpointStorage(storage_path=TEMP_DIR) + + workflow = ( + WorkflowBuilder(max_iterations=5) + .add_edge(upper_case_executor, reverse_text_executor) + .add_edge(reverse_text_executor, lower_case_executor) + .set_start_executor(upper_case_executor) + .with_checkpointing(checkpoint_storage=checkpoint_storage) + .build() + ) + + print("Running workflow with initial message...") + async for event in workflow.run_streaming(message="hello world"): + print(f"Event: {event}") + + # Inspect checkpoints + all_checkpoints = await checkpoint_storage.list_checkpoints() + if not all_checkpoints: + print("No checkpoints found!") + return + + # All checkpoints from this run share one workflow_id + workflow_id = all_checkpoints[0].workflow_id + + # Dump a quick summary including shared_state keys of interest + print("\nCheckpoint summary:") + for cp in sorted(all_checkpoints, key=lambda c: c.timestamp): + msg_count = sum(len(v) for v in cp.messages.values()) + state_keys = sorted(list(cp.executor_states.keys())) if hasattr(cp, "executor_states") else [] + orig = cp.shared_state.get("original_input") if hasattr(cp, "shared_state") else None + upper = cp.shared_state.get("upper_output") if hasattr(cp, "shared_state") else None + print( + f"- {cp.checkpoint_id} | " + f"iter={cp.iteration_count} | messages={msg_count} | states={state_keys} | " + f"shared_state: original_input='{orig}', upper_output='{upper}'" + ) + + # Find the checkpoint with DLROW OLLEH + # This will have us resume at the third (last) executor (node) + checkpoint_id = await find_checkpoint_with_message(checkpoint_storage, workflow_id, "DLROW OLLEH") + if not checkpoint_id: + print("Could not find checkpoint with 'DLROW OLLEH'!") + return + + # The previous workflow can also be used. + # Showing that the workflow can run from a previous checkpoint, + # when checkpointing is not enabled for the particular instance. + new_workflow = ( + WorkflowBuilder(max_iterations=5) + .add_edge(upper_case_executor, reverse_text_executor) + .add_edge(reverse_text_executor, lower_case_executor) + .set_start_executor(upper_case_executor) + .build() + ) + + print(f"\nResuming from checkpoint: {checkpoint_id}") + async for event in new_workflow.run_streaming_from_checkpoint(checkpoint_id, checkpoint_storage=checkpoint_storage): + print(f"Resumed Event: {event}") + + """ + Sample Output: + + Running workflow with initial message... + UpperCaseExecutor: 'hello world' -> 'HELLO WORLD' + Event: ExecutorInvokeEvent(executor_id=upper_case_executor) + Event: ExecutorCompletedEvent(executor_id=upper_case_executor) + ReverseTextExecutor: 'HELLO WORLD' -> 'DLROW OLLEH' + Event: ExecutorInvokeEvent(executor_id=reverse_text_executor) + Event: ExecutorCompletedEvent(executor_id=reverse_text_executor) + LowerCaseExecutor: 'DLROW OLLEH' -> 'dlrow olleh' + LowerCaseExecutor (shared_state): original_input='hello world', upper_output='HELLO WORLD' + Event: ExecutorInvokeEvent(executor_id=lower_case_executor) + Event: WorkflowCompletedEvent(data=dlrow olleh) + Event: ExecutorCompletedEvent(executor_id=lower_case_executor) + + Checkpoint summary: + - dfc63e72-8e8d-454f-9b6d-0d740b9062e6 | label='after_initial_execution' | iter=0 | messages=1 | states=['upper_case_executor'] | shared_state: original_input='hello world', upper_output='HELLO WORLD' + - a78c345a-e5d9-45ba-82c0-cb725452d91b | label='superstep_1' | iter=1 | messages=1 | states=['reverse_text_executor', 'upper_case_executor'] | shared_state: original_input='hello world', upper_output='HELLO WORLD' + - 637c1dbd-a525-4404-9583-da03980537a2 | label='superstep_2' | iter=2 | messages=0 | states=['lower_case_executor', 'reverse_text_executor', 'upper_case_executor'] | shared_state: original_input='hello world', upper_output='HELLO WORLD' + + Resuming from checkpoint: a78c345a-e5d9-45ba-82c0-cb725452d91b + LowerCaseExecutor: 'DLROW OLLEH' -> 'dlrow olleh' + LowerCaseExecutor (shared_state): original_input='hello world', upper_output='HELLO WORLD' + Resumed Event: ExecutorInvokeEvent(executor_id=lower_case_executor) + Resumed Event: WorkflowCompletedEvent(data=dlrow olleh) + Resumed Event: ExecutorCompletedEvent(executor_id=lower_case_executor) + """ # noqa: E501 + + +if __name__ == "__main__": + asyncio.run(main())