Coverage for mlos_bench/mlos_bench/event_loop_context.py: 92%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5"""EventLoopContext class definition.""" 

6 

7import asyncio 

8import logging 

9import sys 

10from asyncio import AbstractEventLoop 

11from concurrent.futures import Future 

12from threading import Lock as ThreadLock 

13from threading import Thread 

14from typing import Any, Coroutine, Optional, TypeVar 

15 

16if sys.version_info >= (3, 10): 

17 from typing import TypeAlias 

18else: 

19 from typing_extensions import TypeAlias 

20 

21CoroReturnType = TypeVar("CoroReturnType") # pylint: disable=invalid-name 

22if sys.version_info >= (3, 9): 

23 FutureReturnType: TypeAlias = Future[CoroReturnType] 

24else: 

25 FutureReturnType: TypeAlias = Future 

26 

27_LOG = logging.getLogger(__name__) 

28 

29 

30class EventLoopContext: 

31 """ 

32 EventLoopContext encapsulates a background thread for asyncio event loop processing 

33 as an aid for context managers. 

34 

35 There is generally only expected to be one of these, either as a base class instance 

36 if it's specific to that functionality or for the full mlos_bench process to support 

37 parallel trial runners, for instance. 

38 

39 It's enter() and exit() routines are expected to be called from the caller's context 

40 manager routines (e.g., __enter__ and __exit__). 

41 """ 

42 

43 def __init__(self) -> None: 

44 self._event_loop: Optional[AbstractEventLoop] = None 

45 self._event_loop_thread: Optional[Thread] = None 

46 self._event_loop_thread_lock = ThreadLock() 

47 self._event_loop_thread_refcnt: int = 0 

48 

49 def _run_event_loop(self) -> None: 

50 """Runs the asyncio event loop in a background thread.""" 

51 assert self._event_loop is not None 

52 asyncio.set_event_loop(self._event_loop) 

53 self._event_loop.run_forever() 

54 

55 def enter(self) -> None: 

56 """Manages starting the background thread for event loop processing.""" 

57 # Start the background thread if it's not already running. 

58 with self._event_loop_thread_lock: 

59 if not self._event_loop_thread: 

60 assert self._event_loop_thread_refcnt == 0 

61 if self._event_loop is None: 

62 if sys.platform == "win32": 

63 asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) 

64 self._event_loop = asyncio.new_event_loop() 

65 assert not self._event_loop.is_running() 

66 self._event_loop_thread = Thread(target=self._run_event_loop, daemon=True) 

67 self._event_loop_thread.start() 

68 self._event_loop_thread_refcnt += 1 

69 

70 def exit(self) -> None: 

71 """Manages cleaning up the background thread for event loop processing.""" 

72 with self._event_loop_thread_lock: 

73 self._event_loop_thread_refcnt -= 1 

74 assert self._event_loop_thread_refcnt >= 0 

75 if self._event_loop_thread_refcnt == 0: 

76 assert self._event_loop is not None 

77 self._event_loop.call_soon_threadsafe(self._event_loop.stop) 

78 _LOG.info("Waiting for event loop thread to stop...") 

79 assert self._event_loop_thread is not None 

80 self._event_loop_thread.join(timeout=3) 

81 if self._event_loop_thread.is_alive(): 

82 raise RuntimeError("Failed to stop event loop thread.") 

83 self._event_loop_thread = None 

84 

85 def run_coroutine(self, coro: Coroutine[Any, Any, CoroReturnType]) -> FutureReturnType: 

86 """ 

87 Runs the given coroutine in the background event loop thread and returns a 

88 Future that can be used to wait for the result. 

89 

90 Parameters 

91 ---------- 

92 coro : Coroutine[Any, Any, CoroReturnType] 

93 The coroutine to run. 

94 

95 Returns 

96 ------- 

97 Future[CoroReturnType] 

98 A future that will be completed when the coroutine completes. 

99 """ 

100 assert self._event_loop_thread_refcnt > 0 

101 assert self._event_loop is not None 

102 return asyncio.run_coroutine_threadsafe(coro, self._event_loop)