From 3235f9d98ca14d15ad6042c9ca17d7738f9ea557 Mon Sep 17 00:00:00 2001 From: Heyi Tang Date: Sat, 11 May 2024 06:51:14 +0000 Subject: [PATCH] Add sleep 0 for scheling cleanup tasks. --- .../promptflow/executor/_async_nodes_scheduler.py | 2 ++ 1 file changed, 2 insertions(+) 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)