From 6da6c2a44cf1f8b9225487a711c3e6094ca5355f Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Tue, 5 Aug 2025 09:43:37 -0700 Subject: [PATCH] fix: using async lock for accessing shared object in parallel executions and update tests for testing various type of functions 1. given we are running parallel functions in one event loop (one thread) , we should use async lock instead of thread lock 2. test three kind of functions: a. sync function b. async function that doesn't yield c. async function that yield PiperOrigin-RevId: 791255012 --- src/google/adk/flows/llm_flows/functions.py | 14 +- .../flows/llm_flows/test_functions_simple.py | 199 ++++++++++++++---- 2 files changed, 170 insertions(+), 43 deletions(-) diff --git a/src/google/adk/flows/llm_flows/functions.py b/src/google/adk/flows/llm_flows/functions.py index 05ab2e65..86f7e30a 100644 --- a/src/google/adk/flows/llm_flows/functions.py +++ b/src/google/adk/flows/llm_flows/functions.py @@ -320,8 +320,8 @@ async def handle_function_calls_live( if not function_calls: return None - # Create thread-safe lock for active_streaming_tools modifications - streaming_lock = threading.Lock() + # Create async lock for active_streaming_tools modifications + streaming_lock = asyncio.Lock() # Create tasks for parallel execution tasks = [ @@ -368,7 +368,7 @@ async def _execute_single_function_call_live( function_call: types.FunctionCall, tools_dict: dict[str, BaseTool], agent: LlmAgent, - streaming_lock: threading.Lock, + streaming_lock: asyncio.Lock, ) -> Optional[Event]: """Execute a single function call for live mode with thread safety.""" tool, tool_context = _get_tool_and_context( @@ -448,7 +448,7 @@ async def _process_function_live_helper( function_call, function_args, invocation_context, - streaming_lock: threading.Lock, + streaming_lock: asyncio.Lock, ): function_response = None # Check if this is a stop_streaming function call @@ -458,7 +458,7 @@ async def _process_function_live_helper( ): function_name = function_args['function_name'] # Thread-safe access to active_streaming_tools - with streaming_lock: + async with streaming_lock: active_tasks = invocation_context.active_streaming_tools if ( active_tasks @@ -491,7 +491,7 @@ async def _process_function_live_helper( } if not function_response: # Clean up the reference under lock - with streaming_lock: + async with streaming_lock: if ( invocation_context.active_streaming_tools and function_name in invocation_context.active_streaming_tools @@ -533,7 +533,7 @@ async def _process_function_live_helper( ) # Register streaming tool using original logic - with streaming_lock: + async with streaming_lock: if invocation_context.active_streaming_tools is None: invocation_context.active_streaming_tools = {} diff --git a/tests/unittests/flows/llm_flows/test_functions_simple.py b/tests/unittests/flows/llm_flows/test_functions_simple.py index dbaf3c8c..166800cc 100644 --- a/tests/unittests/flows/llm_flows/test_functions_simple.py +++ b/tests/unittests/flows/llm_flows/test_functions_simple.py @@ -843,30 +843,31 @@ async def test_parallel_state_modifications_thread_safety(): @pytest.mark.asyncio -async def test_parallel_mixed_sync_async_functions(): - """Test parallel execution with mix of sync and async functions.""" - execution_log = [] +async def test_sync_function_blocks_async_functions(): + """Test that sync functions block async functions from running concurrently.""" + execution_order = [] - def sync_function(value: int) -> dict: - execution_log.append(f'sync_start_{value}') - # Simulate some work - import time + def blocking_sync_function() -> dict: + execution_order.append('sync_A') + # Simulate CPU-intensive work that blocks the event loop + result = 0 + for i in range(1000000): # This blocks the event loop + result += i + execution_order.append('sync_B') + return {'result': 'sync_done'} - time.sleep(0.05) # 50ms - execution_log.append(f'sync_end_{value}') - return {'result': f'sync_{value}'} + async def yielding_async_function() -> dict: + execution_order.append('async_C') + await asyncio.sleep( + 0.001 + ) # This should yield, but can't if event loop is blocked + execution_order.append('async_D') + return {'result': 'async_done'} - async def async_function(value: int) -> dict: - execution_log.append(f'async_start_{value}') - await asyncio.sleep(0.05) # 50ms - execution_log.append(f'async_end_{value}') - return {'result': f'async_{value}'} - - # Create function calls + # Create function calls - these should run "in parallel" function_calls = [ - types.Part.from_function_call(name='sync_function', args={'value': 1}), - types.Part.from_function_call(name='async_function', args={'value': 2}), - types.Part.from_function_call(name='sync_function', args={'value': 3}), + types.Part.from_function_call(name='blocking_sync_function', args={}), + types.Part.from_function_call(name='yielding_async_function', args={}), ] responses: list[types.Content] = [function_calls, 'response1'] @@ -875,24 +876,150 @@ async def test_parallel_mixed_sync_async_functions(): agent = Agent( name='test_agent', model=mock_model, - tools=[sync_function, async_function], + tools=[blocking_sync_function, yielding_async_function], ) runner = testing_utils.TestInMemoryRunner(agent) - - import time - - start_time = time.time() events = await runner.run_async_with_new_session('test') - total_time = time.time() - start_time - # Should complete in less than 120ms (parallel) rather than 150ms (sequential) - # Allow for overhead from task creation and synchronization - assert total_time < 0.12, f'Execution took {total_time}s, expected < 0.12s' + # With blocking sync function, execution should be sequential: A, B, C, D + # The sync function blocks, preventing the async function from yielding properly + assert execution_order == ['sync_A', 'sync_B', 'async_C', 'async_D'] - # Verify all functions were called - assert 'sync_start_1' in execution_log - assert 'sync_end_1' in execution_log - assert 'async_start_2' in execution_log - assert 'async_end_2' in execution_log - assert 'sync_start_3' in execution_log - assert 'sync_end_3' in execution_log + +@pytest.mark.asyncio +async def test_async_function_without_yield_blocks_others(): + """Test that async functions without yield statements block other functions.""" + execution_order = [] + + async def non_yielding_async_function() -> dict: + execution_order.append('non_yield_A') + # CPU-intensive work without any await statements - blocks like sync function + result = 0 + for i in range(1000000): # No await here, so this blocks the event loop + result += i + execution_order.append('non_yield_B') + return {'result': 'non_yielding_done'} + + async def yielding_async_function() -> dict: + execution_order.append('yield_C') + await asyncio.sleep( + 0.001 + ) # This should yield, but can't if event loop is blocked + execution_order.append('yield_D') + return {'result': 'yielding_done'} + + # Create function calls + function_calls = [ + types.Part.from_function_call( + name='non_yielding_async_function', args={} + ), + types.Part.from_function_call(name='yielding_async_function', args={}), + ] + + responses: list[types.Content] = [function_calls, 'response1'] + mock_model = testing_utils.MockModel.create(responses=responses) + + agent = Agent( + name='test_agent', + model=mock_model, + tools=[non_yielding_async_function, yielding_async_function], + ) + runner = testing_utils.TestInMemoryRunner(agent) + events = await runner.run_async_with_new_session('test') + + # Non-yielding async function blocks, so execution is sequential: A, B, C, D + assert execution_order == ['non_yield_A', 'non_yield_B', 'yield_C', 'yield_D'] + + +@pytest.mark.asyncio +async def test_yielding_async_functions_run_concurrently(): + """Test that async functions with proper yields run concurrently.""" + execution_order = [] + + async def yielding_async_function_1() -> dict: + execution_order.append('func1_A') + await asyncio.sleep(0.001) # Yield control + execution_order.append('func1_B') + return {'result': 'func1_done'} + + async def yielding_async_function_2() -> dict: + execution_order.append('func2_C') + await asyncio.sleep(0.001) # Yield control + execution_order.append('func2_D') + return {'result': 'func2_done'} + + # Create function calls + function_calls = [ + types.Part.from_function_call(name='yielding_async_function_1', args={}), + types.Part.from_function_call(name='yielding_async_function_2', args={}), + ] + + responses: list[types.Content] = [function_calls, 'response1'] + mock_model = testing_utils.MockModel.create(responses=responses) + + agent = Agent( + name='test_agent', + model=mock_model, + tools=[yielding_async_function_1, yielding_async_function_2], + ) + runner = testing_utils.TestInMemoryRunner(agent) + events = await runner.run_async_with_new_session('test') + + # With proper yielding, execution should interleave: A, C, B, D + # Both functions start, yield, then complete + assert execution_order == ['func1_A', 'func2_C', 'func1_B', 'func2_D'] + + +@pytest.mark.asyncio +async def test_mixed_function_types_execution_order(): + """Test execution order with all three types of functions.""" + execution_order = [] + + def sync_function() -> dict: + execution_order.append('sync_A') + # Small amount of blocking work + result = sum(range(100000)) + execution_order.append('sync_B') + return {'result': 'sync_done'} + + async def non_yielding_async() -> dict: + execution_order.append('non_yield_C') + # CPU work without yield + result = sum(range(100000)) + execution_order.append('non_yield_D') + return {'result': 'non_yield_done'} + + async def yielding_async() -> dict: + execution_order.append('yield_E') + await asyncio.sleep(0.001) # Proper yield + execution_order.append('yield_F') + return {'result': 'yield_done'} + + # Create function calls + function_calls = [ + types.Part.from_function_call(name='sync_function', args={}), + types.Part.from_function_call(name='non_yielding_async', args={}), + types.Part.from_function_call(name='yielding_async', args={}), + ] + + responses: list[types.Content] = [function_calls, 'response1'] + mock_model = testing_utils.MockModel.create(responses=responses) + + agent = Agent( + name='test_agent', + model=mock_model, + tools=[sync_function, non_yielding_async, yielding_async], + ) + runner = testing_utils.TestInMemoryRunner(agent) + events = await runner.run_async_with_new_session('test') + + # All blocking functions run sequentially, then the yielding one + # Expected order: sync_A, sync_B, non_yield_C, non_yield_D, yield_E, yield_F + assert execution_order == [ + 'sync_A', + 'sync_B', + 'non_yield_C', + 'non_yield_D', + 'yield_E', + 'yield_F', + ]