Coverage for haystack/utils/async_utils.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import asyncio 

6from typing import Any, TypeVar 

7 

8from haystack import logging 

9 

10logger = logging.getLogger(__name__) 

11 

12T = TypeVar("T") 

13 

14 

15async def _gather_tasks_with_cancel(tasks: list[asyncio.Task[T]]) -> list[T]: 

16 """ 

17 Wait for all tasks, cancelling and draining unfinished siblings if one fails. 

18 

19 :param tasks: Tasks to wait for. 

20 :returns: 

21 The task results in input order. 

22 """ 

23 try: 

24 return await asyncio.gather(*tasks) 

25 except Exception: 

26 for task in tasks: 

27 task.cancel() 

28 await asyncio.gather(*tasks, return_exceptions=True) 

29 raise 

30 

31 

32async def _execute_component_async(component_instance: Any, **kwargs: Any) -> dict[str, Any]: 

33 """ 

34 Run a component asynchronously, preferring its `run_async` method when implemented. 

35 

36 If the component does not implement `run_async`, its synchronous `run` method is executed in a thread 

37 to avoid blocking the event loop. 

38 

39 :param component_instance: The component to run. Any object exposing a `run` method and optionally a 

40 `run_async` coroutine method. 

41 :param kwargs: Keyword arguments passed to the component's `run_async` or `run` method. 

42 :returns: 

43 The component's output dictionary. 

44 """ 

45 run_async = getattr(component_instance, "run_async", None) 

46 if callable(run_async): 

47 return await run_async(**kwargs) 

48 

49 logger.debug( 

50 "{component_type} does not implement 'run_async'. Running the synchronous 'run' method in a thread " 

51 "to avoid blocking the event loop.", 

52 component_type=type(component_instance).__name__, 

53 ) 

54 return await asyncio.to_thread(component_instance.run, **kwargs)