mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
[BREAKING] Python: Refactor SharedState to State with sync methods and superstep caching (#3667)
* Refactor SharedState to State with sync methods and superstep caching * Fixes * Address PR feedback * Remove dead links * Fix lab test import
This commit is contained in:
committed by
GitHub
Unverified
parent
4e25917644
commit
10afb86213
+12
-12
@@ -10,7 +10,7 @@ import aiofiles
|
||||
from agent_framework import (
|
||||
Executor, # Base class for custom workflow steps
|
||||
WorkflowBuilder, # Fluent builder for executors and edges
|
||||
WorkflowContext, # Per run context with shared state and messaging
|
||||
WorkflowContext, # Per run context with workflow state and messaging
|
||||
WorkflowOutputEvent, # Event emitted when workflow yields output
|
||||
WorkflowViz, # Utility to visualize a workflow graph
|
||||
handler, # Decorator to expose an Executor method as a step
|
||||
@@ -26,7 +26,7 @@ It also demonstrates WorkflowViz for graph visualization.
|
||||
|
||||
Purpose:
|
||||
Show how to:
|
||||
- Partition input once and coordinate parallel mappers with shared state.
|
||||
- Partition input once and coordinate parallel mappers with workflow state.
|
||||
- Implement map, shuffle, and reduce executors that pass file paths instead of large payloads.
|
||||
- Use fan out and fan in edges to express parallelism and joins.
|
||||
- Persist intermediate results to disk to bound memory usage for large inputs.
|
||||
@@ -49,8 +49,8 @@ TEMP_DIR = os.path.join(DIR, "tmp")
|
||||
# Ensure the temporary directory exists
|
||||
os.makedirs(TEMP_DIR, exist_ok=True)
|
||||
|
||||
# Define a key for the shared state to store the data to be processed
|
||||
SHARED_STATE_DATA_KEY = "data_to_be_processed"
|
||||
# Define a key for the workflow state to store the data to be processed
|
||||
STATE_DATA_KEY = "data_to_be_processed"
|
||||
|
||||
|
||||
class SplitCompleted:
|
||||
@@ -69,17 +69,17 @@ class Split(Executor):
|
||||
|
||||
@handler
|
||||
async def split(self, data: str, ctx: WorkflowContext[SplitCompleted]) -> None:
|
||||
"""Tokenize input and assign contiguous index ranges to each mapper via shared state.
|
||||
"""Tokenize input and assign contiguous index ranges to each mapper via workflow state.
|
||||
|
||||
Args:
|
||||
data: The raw text to process.
|
||||
ctx: Workflow context to persist shared state and send messages.
|
||||
ctx: Workflow context to persist state and send messages.
|
||||
"""
|
||||
# Process data into a list of words and remove empty lines or words.
|
||||
word_list = self._preprocess(data)
|
||||
|
||||
# Store tokenized words once so all mappers can read by index.
|
||||
await ctx.set_shared_state(SHARED_STATE_DATA_KEY, word_list)
|
||||
ctx.set_state(STATE_DATA_KEY, word_list)
|
||||
|
||||
# Divide indices into contiguous slices for each mapper.
|
||||
map_executor_count = len(self._map_executor_ids)
|
||||
@@ -90,8 +90,8 @@ class Split(Executor):
|
||||
start_index = i * chunk_size
|
||||
end_index = start_index + chunk_size if i < map_executor_count - 1 else len(word_list)
|
||||
|
||||
# The mapper reads its slice from shared state keyed by its own executor id.
|
||||
await ctx.set_shared_state(self._map_executor_ids[i], (start_index, end_index))
|
||||
# The mapper reads its slice from workflow state keyed by its own executor id.
|
||||
ctx.set_state(self._map_executor_ids[i], (start_index, end_index))
|
||||
await ctx.send_message(SplitCompleted(), self._map_executor_ids[i])
|
||||
|
||||
tasks = [asyncio.create_task(_process_chunk(i)) for i in range(map_executor_count)]
|
||||
@@ -119,11 +119,11 @@ class Map(Executor):
|
||||
|
||||
Args:
|
||||
_: SplitCompleted marker indicating maps can begin.
|
||||
ctx: Workflow context for shared state access and messaging.
|
||||
ctx: Workflow context for workflow state access and messaging.
|
||||
"""
|
||||
# Retrieve tokens and our assigned slice.
|
||||
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)
|
||||
data_to_be_processed: list[str] = ctx.get_state(STATE_DATA_KEY)
|
||||
chunk_start, chunk_end = ctx.get_state(self.id)
|
||||
|
||||
results = [(item, 1) for item in data_to_be_processed[chunk_start:chunk_end]]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user