diff --git a/src/promptflow-core/promptflow/executor/_async_nodes_scheduler.py b/src/promptflow-core/promptflow/executor/_async_nodes_scheduler.py index 2ccb50d68c2c..e618e3048ba0 100644 --- a/src/promptflow-core/promptflow/executor/_async_nodes_scheduler.py +++ b/src/promptflow-core/promptflow/executor/_async_nodes_scheduler.py @@ -132,6 +132,8 @@ async def _wait_and_complete_nodes( for task in tasks: if not task.done(): task.cancel() + # Wait for the cancelled tasks to be scheduled to cleanup + await asyncio.sleep(0) err_msg = f"Line execution timeout after {line_timeout_sec} seconds." context.cancel_node_runs(err_msg) raise LineExecutionTimeoutError(context._line_number, line_timeout_sec)