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

Python: workflow instance should not be invoked concurrently (#948)

Browse Source 
* workflow instance should not be invoked concurrently

* address comments
This commit is contained in:
Eric Zhu
2025-09-26 13:11:05 -07:00
committed by GitHub
Unverified
parent 0d0231b512
commit a9608465d9
2 changed files with 334 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/agent_framework/_workflow/_workflow.py
+235 -96
View File
@@ -227,12 +227,25 @@ class Workflow(AFBaseModel):
workflow_id=id,
)
# Flag to prevent concurrent workflow executions
self._is_running = False
# Capture a canonical fingerprint of the workflow graph so checkpoints
# can assert they are resumed with an equivalent topology.
self._graph_signature = self._compute_graph_signature()
self._graph_signature_hash = self._hash_graph_signature(self._graph_signature)
self._runner.graph_signature_hash = self._graph_signature_hash
def _ensure_not_running(self) -> None:
"""Ensure the workflow is not already running."""
if self._is_running:
raise RuntimeError("Workflow is already running. Concurrent executions are not allowed.")
self._is_running = True
def _reset_running_flag(self) -> None:
"""Reset the running flag."""
self._is_running = False
def model_dump(self, **kwargs: Any) -> dict[str, Any]:
"""Custom serialization that properly handles WorkflowExecutor nested workflows."""
data = super().model_dump(**kwargs)
@@ -370,20 +383,26 @@ class Workflow(AFBaseModel):
Yields:
WorkflowEvent: The events generated during the workflow execution.
"""
self._ensure_not_running()
try:
async def initial_execution() -> None:
executor = self.get_start_executor()
await executor.execute(
message,
[self.__class__.__name__], # source_executor_ids
self._shared_state, # shared_state
self._runner.context, # runner_context
trace_contexts=None, # No parent trace context for workflow start
source_span_ids=None, # No source span for workflow start
)
async def initial_execution() -> None:
executor = self.get_start_executor()
await executor.execute(
message,
[self.__class__.__name__], # source_executor_ids
self._shared_state, # shared_state
self._runner.context, # runner_context
trace_contexts=None, # No parent trace context for workflow start
source_span_ids=None, # No source span for workflow start
)
async for event in self._run_workflow_with_tracing(initial_executor_fn=initial_execution, reset_context=True):
yield event
async for event in self._run_workflow_with_tracing(
initial_executor_fn=initial_execution, reset_context=True
):
yield event
finally:
self._reset_running_flag()
async def run_stream_from_checkpoint(
self,
@@ -407,60 +426,64 @@ class Workflow(AFBaseModel):
ValueError: If neither checkpoint_storage is provided nor checkpointing is enabled.
RuntimeError: If checkpoint restoration fails.
"""
self._ensure_not_running()
try:
async def checkpoint_restoration() -> None:
has_checkpointing = self._runner.context.has_checkpointing()
async def checkpoint_restoration() -> None:
has_checkpointing = self._runner.context.has_checkpointing()
if not has_checkpointing and checkpoint_storage is None:
raise ValueError(
"Cannot restore from checkpoint: either provide checkpoint_storage parameter "
"or build workflow with WorkflowBuilder.with_checkpointing(checkpoint_storage)."
)
if not has_checkpointing and checkpoint_storage is None:
raise ValueError(
"Cannot restore from checkpoint: either provide checkpoint_storage parameter "
"or build workflow with WorkflowBuilder.with_checkpointing(checkpoint_storage)."
)
restored = await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage)
restored = await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage)
if not restored:
raise RuntimeError(f"Failed to restore from checkpoint: {checkpoint_id}")
if not restored:
raise RuntimeError(f"Failed to restore from checkpoint: {checkpoint_id}")
# Process any pending messages from the checkpoint first
# This ensures that RequestInfoExecutor state is properly populated
# before we try to handle responses
if await self._runner.context.has_messages():
# Run one iteration to process pending messages
# This will populate RequestInfoExecutor._request_events properly
await self._runner._run_iteration() # type: ignore
# Process any pending messages from the checkpoint first
# This ensures that RequestInfoExecutor state is properly populated
# before we try to handle responses
if await self._runner.context.has_messages():
# Run one iteration to process pending messages
# This will populate RequestInfoExecutor._request_events properly
await self._runner._run_iteration() # type: ignore
if responses:
request_info_executor = self._find_request_info_executor()
if request_info_executor:
for request_id, response_data in responses.items():
ctx: WorkflowContext[Any] = WorkflowContext(
request_info_executor.id,
[self.__class__.__name__],
self._shared_state,
self._runner.context,
trace_contexts=None, # No parent trace context for new workflow span
source_span_ids=None, # No source span for response handling
)
if not await request_info_executor.has_pending_request(request_id, ctx):
logger.debug(
f"Skipping pre-supplied response for request {request_id}; no pending request found "
f"after checkpoint restoration."
if responses:
request_info_executor = self._find_request_info_executor()
if request_info_executor:
for request_id, response_data in responses.items():
ctx: WorkflowContext[Any] = WorkflowContext(
request_info_executor.id,
[self.__class__.__name__],
self._shared_state,
self._runner.context,
trace_contexts=None, # No parent trace context for new workflow span
source_span_ids=None, # No source span for response handling
)
continue
await request_info_executor.handle_response(
response_data,
request_id,
ctx,
)
if not await request_info_executor.has_pending_request(request_id, ctx):
logger.debug(
f"Skipping pre-supplied response for request {request_id}; "
f"no pending request found after checkpoint restoration."
)
continue
async for event in self._run_workflow_with_tracing(
initial_executor_fn=checkpoint_restoration,
reset_context=False, # Don't reset context when resuming from checkpoint
):
yield event
await request_info_executor.handle_response(
response_data,
request_id,
ctx,
)
async for event in self._run_workflow_with_tracing(
initial_executor_fn=checkpoint_restoration,
reset_context=False, # Don't reset context when resuming from checkpoint
):
yield event
finally:
self._reset_running_flag()
async def send_responses_streaming(self, responses: dict[str, Any]) -> AsyncIterable[WorkflowEvent]:
"""Send responses back to the workflow and stream the events generated by the workflow.
@@ -472,36 +495,40 @@ class Workflow(AFBaseModel):
Yields:
WorkflowEvent: The events generated during the workflow execution after sending the responses.
"""
self._ensure_not_running()
try:
async def send_responses() -> None:
request_info_executor = self._find_request_info_executor()
if not request_info_executor:
raise ValueError("No RequestInfoExecutor found in workflow.")
async def send_responses() -> None:
request_info_executor = self._find_request_info_executor()
if not request_info_executor:
raise ValueError("No RequestInfoExecutor found in workflow.")
async def _handle_response(response: Any, request_id: str) -> None:
"""Handle the response from the RequestInfoExecutor."""
await request_info_executor.handle_response(
response,
request_id,
WorkflowContext(
request_info_executor.id,
[self.__class__.__name__],
self._shared_state,
self._runner.context,
trace_contexts=None, # No parent trace context for new workflow span
source_span_ids=None, # No source span for response handling
),
)
async def _handle_response(response: Any, request_id: str) -> None:
"""Handle the response from the RequestInfoExecutor."""
await request_info_executor.handle_response(
response,
request_id,
WorkflowContext(
request_info_executor.id,
[self.__class__.__name__],
self._shared_state,
self._runner.context,
trace_contexts=None, # No parent trace context for new workflow span
source_span_ids=None, # No source span for response handling
),
)
await asyncio.gather(*[
_handle_response(response, request_id) for request_id, response in responses.items()
])
await asyncio.gather(*[
_handle_response(response, request_id) for request_id, response in responses.items()
])
async for event in self._run_workflow_with_tracing(
initial_executor_fn=send_responses,
reset_context=False, # Don't reset context when sending responses
):
yield event
async for event in self._run_workflow_with_tracing(
initial_executor_fn=send_responses,
reset_context=False, # Don't reset context when sending responses
):
yield event
finally:
self._reset_running_flag()
async def run(self, message: Any, *, include_status_events: bool = False) -> WorkflowRunResult:
"""Run the workflow with the given message.
@@ -513,11 +540,31 @@ class Workflow(AFBaseModel):
Returns:
A WorkflowRunResult instance containing a list of events generated during the workflow execution.
"""
from agent_framework import AgentRunResponse, AgentRunResponseUpdate
self._ensure_not_running()
try:
from agent_framework import AgentRunResponse, AgentRunResponseUpdate
from ._events import AgentRunEvent, AgentRunUpdateEvent # Local import to avoid cycles
from ._events import AgentRunEvent, AgentRunUpdateEvent # Local import to avoid cycles
raw_events = [event async for event in self.run_stream(message)]
async def initial_execution() -> None:
executor = self.get_start_executor()
await executor.execute(
message,
[self.__class__.__name__], # source_executor_ids
self._shared_state, # shared_state
self._runner.context, # runner_context
trace_contexts=None, # No parent trace context for workflow start
source_span_ids=None, # No source span for workflow start
)
raw_events = [
event
async for event in self._run_workflow_with_tracing(
initial_executor_fn=initial_execution, reset_context=True
)
]
finally:
self._reset_running_flag()
# Coalesce streaming update events into a single AgentRunEvent per executor sequence.
coalesced: list[WorkflowEvent] = [] # type: ignore[name-defined]
@@ -589,12 +636,69 @@ class Workflow(AFBaseModel):
ValueError: If neither checkpoint_storage is provided nor checkpointing is enabled.
RuntimeError: If checkpoint restoration fails.
"""
events = [
event async for event in self.run_stream_from_checkpoint(checkpoint_id, checkpoint_storage, responses)
]
status_events = [e for e in events if isinstance(e, WorkflowStatusEvent)]
filtered_events = [e for e in events if not isinstance(e, (WorkflowStatusEvent, WorkflowStartedEvent))]
return WorkflowRunResult(filtered_events, status_events)
self._ensure_not_running()
try:
async def checkpoint_restoration() -> None:
has_checkpointing = self._runner.context.has_checkpointing()
if not has_checkpointing and checkpoint_storage is None:
raise ValueError(
"Cannot restore from checkpoint: either provide checkpoint_storage parameter "
"or build workflow with WorkflowBuilder.with_checkpointing(checkpoint_storage)."
)
restored = await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage)
if not restored:
raise RuntimeError(f"Failed to restore from checkpoint: {checkpoint_id}")
# Process any pending messages from the checkpoint first
# This ensures that RequestInfoExecutor state is properly populated
# before we try to handle responses
if await self._runner.context.has_messages():
# Run one iteration to process pending messages
# This will populate RequestInfoExecutor._request_events properly
await self._runner._run_iteration() # type: ignore
if responses:
request_info_executor = self._find_request_info_executor()
if request_info_executor:
for request_id, response_data in responses.items():
ctx: WorkflowContext[Any] = WorkflowContext(
request_info_executor.id,
[self.__class__.__name__],
self._shared_state,
self._runner.context,
trace_contexts=None, # No parent trace context for new workflow span
source_span_ids=None, # No source span for response handling
)
if not await request_info_executor.has_pending_request(request_id, ctx):
logger.debug(
f"Skipping pre-supplied response for request {request_id}; "
f"no pending request found after checkpoint restoration."
)
continue
await request_info_executor.handle_response(
response_data,
request_id,
ctx,
)
events = [
event
async for event in self._run_workflow_with_tracing(
initial_executor_fn=checkpoint_restoration,
reset_context=False, # Don't reset context when resuming from checkpoint
)
]
status_events = [e for e in events if isinstance(e, WorkflowStatusEvent)]
filtered_events = [e for e in events if not isinstance(e, (WorkflowStatusEvent, WorkflowStartedEvent))]
return WorkflowRunResult(filtered_events, status_events)
finally:
self._reset_running_flag()
async def send_responses(self, responses: dict[str, Any]) -> WorkflowRunResult:
"""Send responses back to the workflow.
@@ -605,10 +709,45 @@ class Workflow(AFBaseModel):
Returns:
A WorkflowRunResult instance containing a list of events generated during the workflow execution.
"""
events = [event async for event in self.send_responses_streaming(responses)]
status_events = [e for e in events if isinstance(e, WorkflowStatusEvent)]
filtered_events = [e for e in events if not isinstance(e, (WorkflowStatusEvent, WorkflowStartedEvent))]
return WorkflowRunResult(filtered_events, status_events)
self._ensure_not_running()
try:
async def send_responses_internal() -> None:
request_info_executor = self._find_request_info_executor()
if not request_info_executor:
raise ValueError("No RequestInfoExecutor found in workflow.")
async def _handle_response(response: Any, request_id: str) -> None:
"""Handle the response from the RequestInfoExecutor."""
await request_info_executor.handle_response(
response,
request_id,
WorkflowContext(
request_info_executor.id,
[self.__class__.__name__],
self._shared_state,
self._runner.context,
trace_contexts=None, # No parent trace context for new workflow span
source_span_ids=None, # No source span for response handling
),
)
await asyncio.gather(*[
_handle_response(response, request_id) for request_id, response in responses.items()
])
events = [
event
async for event in self._run_workflow_with_tracing(
initial_executor_fn=send_responses_internal,
reset_context=False, # Don't reset context when sending responses
)
]
status_events = [e for e in events if isinstance(e, WorkflowStatusEvent)]
filtered_events = [e for e in events if not isinstance(e, (WorkflowStatusEvent, WorkflowStartedEvent))]
return WorkflowRunResult(filtered_events, status_events)
finally:
self._reset_running_flag()
def _get_executor_by_id(self, executor_id: str) -> Executor:
"""Get an executor by its ID.
python/packages/main/tests/workflow/test_workflow.py
+99
View File
@@ -1,5 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import tempfile
from dataclasses import dataclass
from typing import Any
@@ -688,3 +689,101 @@ async def test_workflow_with_simple_cycle_and_exit_condition():
# Should have multiple events due to cycling
assert len(events) >= 4, f"Expected at least 4 events due to cycling, got {len(events)}"
async def test_workflow_concurrent_execution_prevention():
"""Test that concurrent workflow executions are prevented."""
# Create a simple workflow that takes some time to execute
executor = IncrementExecutor(id="slow_executor", limit=3, increment=1)
workflow = WorkflowBuilder().set_start_executor(executor).build()
# Create a task that will run the workflow
async def run_workflow():
return await workflow.run(NumberMessage(data=0))
# Start the first workflow execution
task1 = asyncio.create_task(run_workflow())
# Give it a moment to start
await asyncio.sleep(0.01)
# Try to start a second concurrent execution - this should fail
with pytest.raises(RuntimeError, match="Workflow is already running. Concurrent executions are not allowed."):
await workflow.run(NumberMessage(data=0))
# Wait for the first task to complete
result = await task1
assert result.get_final_state() == WorkflowRunState.IDLE
# After the first execution completes, we should be able to run again
result2 = await workflow.run(NumberMessage(data=0))
assert result2.get_final_state() == WorkflowRunState.IDLE
async def test_workflow_concurrent_execution_prevention_streaming():
"""Test that concurrent workflow streaming executions are prevented."""
# Create a simple workflow
executor = IncrementExecutor(id="slow_executor", limit=3, increment=1)
workflow = WorkflowBuilder().set_start_executor(executor).build()
# Create an async generator that will consume the stream slowly
async def consume_stream_slowly():
result = []
async for event in workflow.run_stream(NumberMessage(data=0)):
result.append(event)
await asyncio.sleep(0.01) # Slow consumption
return result
# Start the first streaming execution
task1 = asyncio.create_task(consume_stream_slowly())
# Give it a moment to start
await asyncio.sleep(0.02)
# Try to start a second concurrent execution - this should fail
with pytest.raises(RuntimeError, match="Workflow is already running. Concurrent executions are not allowed."):
await workflow.run(NumberMessage(data=0))
# Wait for the first task to complete
result = await task1
assert len(result) > 0 # Should have received some events
# After the first execution completes, we should be able to run again
result2 = await workflow.run(NumberMessage(data=0))
assert result2.get_final_state() == WorkflowRunState.IDLE
async def test_workflow_concurrent_execution_prevention_mixed_methods():
"""Test that concurrent executions are prevented across different execution methods."""
# Create a simple workflow
executor = IncrementExecutor(id="slow_executor", limit=3, increment=1)
workflow = WorkflowBuilder().set_start_executor(executor).build()
# Start a streaming execution
async def consume_stream():
result = []
async for event in workflow.run_stream(NumberMessage(data=0)):
result.append(event)
await asyncio.sleep(0.01)
return result
task1 = asyncio.create_task(consume_stream())
await asyncio.sleep(0.02) # Let it start
# Try different execution methods - all should fail
with pytest.raises(RuntimeError, match="Workflow is already running. Concurrent executions are not allowed."):
await workflow.run(NumberMessage(data=0))
with pytest.raises(RuntimeError, match="Workflow is already running. Concurrent executions are not allowed."):
async for _ in workflow.run_stream(NumberMessage(data=0)):
break
with pytest.raises(RuntimeError, match="Workflow is already running. Concurrent executions are not allowed."):
await workflow.send_responses({"test": "data"})
# Wait for the original task to complete
await task1
# Now all methods should work again
result = await workflow.run(NumberMessage(data=0))
assert result.get_final_state() == WorkflowRunState.IDLE