Repository: microsoft/promptflow
File: src/promptflow-core/promptflow/executor/_async_nodes_scheduler.py (Python)

Code before:
        # Then the event loop will wait for all tasks to be completed before raising the cancellation error.
        # See reference: https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.Executor
        try:
            outputs = await self._execute_with_thread_pool(executor, nodes, inputs, context, timeout_seconds)
        except asyncio.CancelledError:
            await self.cancel()
            raise
        executor.shutdown()
        return outputs

    async def _execute_with_thread_pool(
        self,
        executor: ThreadPoolExecutor,
        nodes: List[Node],
        inputs: Dict[str, Any],
        context: FlowExecutionContext,
        timeout_seconds: Optional[int] = None,
    ) -> Tuple[dict, dict]:
        flow_logger.info(f"Start to run {len(nodes)} nodes with the current event loop.")
        dag_manager = DAGManager(nodes, inputs)
        timeout_task = asyncio.create_task(asyncio.sleep(timeout_seconds)) if timeout_seconds else None
        task2nodes = self._execute_nodes(dag_manager, context, executor)
        while not dag_manager.completed():
            task2nodes = await self._wait_and_complete_nodes(
                task2nodes, dag_manager, context, timeout_task, timeout_seconds
            )
            submitted_tasks2nodes = self._execute_nodes(dag_manager, context, executor)
            task2nodes.update(submitted_tasks2nodes)
        # Set the event to notify the monitor thread to exit
        # Ref: https://docs.python.org/3/library/threading.html#event-objects
        self._dag_manager_completed_event.set()
        for node in dag_manager.bypassed_nodes:
            dag_manager.completed_nodes_outputs[node] = None
        return dag_manager.completed_nodes_outputs, dag_manager.bypassed_nodes

    async def _wait_and_complete_nodes(
        self,
        task2nodes: Dict[Task, Node],
        dag_manager: DAGManager,
        context: FlowExecutionContext,
        timeout_task: Optional[Task] = None,
        line_timeout_sec: Optional[int] = None,
    ) -> Dict[Task, Node]:
        if not task2nodes:
            raise NoNodeExecutedError("No nodes are ready for execution, but the flow is not completed.")
        tasks = [task for task in task2nodes]
        for task in tasks:
            self._task_start_time[task] = time.time()
        if timeout_task is not None:
            tasks.append(timeout_task)
        done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)

Review comment (suggestion):
How about adding another method called `_cancel_nodes()` and then we modify the logic here to something like:

```python
while not dag_manager.completed():
    if execution_time >= timeout_seconds:
        await self._cancel_nodes()
        break

    # Do the exiting thing here, like:
    task2nodes = await self._wait_and_complete_nodes(task2nodes, dag_manager)
```

Code after:
        # Then the event loop will wait for all tasks to be completed before raising the cancellation error.
        # See reference: https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.Executor
        try:
            outputs = await self._execute_with_thread_pool(executor, nodes, inputs, context, timeout_seconds)
        except asyncio.CancelledError:
            await self.cancel()
            raise
        executor.shutdown()
        return outputs

    async def _execute_with_thread_pool(
        self,
        executor: ThreadPoolExecutor,
        nodes: List[Node],
        inputs: Dict[str, Any],
        context: FlowExecutionContext,
        timeout_seconds: Optional[int] = None,
    ) -> Tuple[dict, dict]:
        flow_logger.info(f"Start to run {len(nodes)} nodes with the current event loop.")
        start_time = time.time()
        dag_manager = DAGManager(nodes, inputs)
        task2nodes = self._execute_nodes(dag_manager, context, executor)
        while not dag_manager.completed():
            remaining_timeout = None if timeout_seconds is None else timeout_seconds - (time.time() - start_time)
            task = asyncio.create_task(self._wait_and_complete_nodes(task2nodes, dag_manager))
            try:
                task2nodes = await asyncio.wait_for(task, remaining_timeout)
            except asyncio.TimeoutError:
                flow_logger.warning(f"Line execution timeout after {timeout_seconds} seconds.")
                await self.cancel_tasks(task2nodes.keys())
                raise LineExecutionTimeoutError(context._line_number, timeout_seconds)
            submitted_tasks2nodes = self._execute_nodes(dag_manager, context, executor)
            task2nodes.update(submitted_tasks2nodes)
        # Set the event to notify the monitor thread to exit
        # Ref: https://docs.python.org/3/library/threading.html#event-objects
        self._dag_manager_completed_event.set()
        for node in dag_manager.bypassed_nodes:
            dag_manager.completed_nodes_outputs[node] = None
        return dag_manager.completed_nodes_outputs, dag_manager.bypassed_nodes

    async def cancel_tasks(self, tasks: List[asyncio.Task]):
        for task in tasks:
            if not task.done():
                task.cancel()
        cancel_timeout = 1
        # Wait at most 1 second for the tasks to cleanup
        await asyncio.wait(tasks, return_when=asyncio.ALL_COMPLETED, timeout=cancel_timeout)

    async def _wait_and_complete_nodes(self, task2nodes: Dict[Task, Node], dag_manager: DAGManager) -> Dict[Task, Node]:
        if not task2nodes:
            raise NoNodeExecutedError("No nodes are ready for execution, but the flow is not completed.")