Coverage for qreward / utils / scheduler / async_runner.py: 100.00%

1"""Async runner implementation for asynchronous task execution.""" 

2 

3import asyncio (empty)

4import logging (empty)

5from typing import Any, Callable, List, Optional, Sequence (empty)

6 

7from .base import BaseRunner, _CancelledErrorGroups (empty)

8from .config import ScheduleConfig (empty)

9from .context import ExecutionContext (empty)

10 

11logger = logging.getLogger(__name__) (empty)

12 

13 

14class AsyncRunner(BaseRunner): (empty)

15 """Async runner implementation.""" 

16 

17 def create_task( (empty)

18 self, func: Callable, *args, **kwargs 

19 ) -> asyncio.Task: 

20 """Create an asyncio task.""" 

21 return asyncio.create_task(func(*args, **kwargs)) (empty)

22 

23 def get_task_result( (empty)

24 self, task: asyncio.Task 

25 ) -> Any: 

26 """Get result from a completed asyncio task.""" 

27 return task.result() (empty)

28 

29 def get_task_exception( (empty)

30 self, task: asyncio.Task 

31 ) -> Optional[BaseException]: 

32 """Get exception from a completed asyncio task.""" 

33 return task.exception() (empty)

34 

35 def is_task_cancelled( (empty)

36 self, task: asyncio.Task 

37 ) -> bool: 

38 """Check if an asyncio task was cancelled.""" 

39 return task.cancelled() (empty)

40 

41 async def _cancel_tasks( (empty)

42 self, 

43 pending: Sequence[asyncio.Task], (empty)

44 done: List[asyncio.Task], (empty)

45 retry_interval: float, (empty)

46 ) -> None: # pragma: no cover (empty)

47 """Cancel pending tasks and consume done tasks. 

48 

49 Args: 

50 pending: Pending tasks to cancel. 

51 done: Done tasks to consume. 

52 retry_interval: Time to wait for cancellation. 

53 """ 

54 while done: (empty)

55 try: (empty)

56 done.pop() (empty)

57 except _CancelledErrorGroups: 

58 pass 

59 

60 if pending: (empty)

61 for task in pending: (empty)

62 if not task.done(): (empty)

63 task.cancel() (empty)

64 try: (empty)

65 await asyncio.wait_for( (empty)

66 asyncio.gather(*pending, return_exceptions=True), (empty)

67 timeout=retry_interval if retry_interval > 0 else None, (empty)

68 ) 

69 except _CancelledErrorGroups: (empty)

70 pass (empty)

71 

72 async def execute_impl( (empty)

73 self, 

74 context: ExecutionContext, 

75 config: ScheduleConfig, 

76 *args: Any, 

77 **kwargs: Any, 

78 ) -> Any: 

79 """Async execution implementation. 

80 

81 Args: 

82 context: Execution context. 

83 config: Schedule configuration. 

84 *args: Function arguments. 

85 **kwargs: Function keyword arguments. 

86 

87 Returns: 

88 Function result. 

89 """ 

90 run_tasks: List[asyncio.Task] = [] (empty)

91 

92 try: (empty)

93 while not context.finish and ( (empty)

94 context.cur_times <= config.retry_times or len(run_tasks) > 0 

95 ): 

96 # 1. Submit new tasks if allowed 

97 if context.can_submit_task(len(run_tasks)): (empty)

98 limiter_timeout = context.get_limiter_timeout( (empty)

99 len(run_tasks) 

100 ) 

101 

102 if not context.limiter or await ( (empty)

103 context.limiter.async_allow( 

104 limiter_timeout if limiter_timeout > 0 else None 

105 ) 

106 ): 

107 if context.is_hedge_submit( (empty)

108 len(run_tasks) 

109 ): # pragma: no cover 

110 context.record_hedge() (empty)

111 elif ( (empty)

112 context.result_exception is not None 

113 ): # pragma: no cover 

114 context.record_exception(context.result_exception) (empty)

115 context.result_exception = None (empty)

116 

117 run_tasks.append( (empty)

118 asyncio.create_task(context.func(*args, **kwargs)) 

119 ) 

120 context.mark_task_submitted() (empty)

121 

122 # 2. Wait for tasks 

123 done, pending = set(), set() (empty)

124 if run_tasks: # pragma: no branch (empty)

125 cur_timeout = context.compute_timeout(len(run_tasks)) (empty)

126 if cur_timeout > 0: (empty)

127 done, pending = await asyncio.wait( (empty)

128 run_tasks, 

129 timeout=cur_timeout, 

130 return_when=asyncio.FIRST_COMPLETED, 

131 ) 

132 else: # pragma: no cover 

133 done, pending = await asyncio.wait( (empty)

134 run_tasks, return_when=asyncio.FIRST_COMPLETED 

135 ) 

136 

137 # 3. Process results 

138 can_add_speed_up = context.can_increase_speed() (empty)

139 cancel_wait = self._compute_cancel_wait(config, context) (empty)

140 

141 for finished in list(done): (empty)

142 if finished.cancelled(): # pragma: no cover (empty)

143 run_tasks.remove(finished) (empty)

144 continue (empty)

145 

146 exception = finished.exception() (empty)

147 if exception is None: (empty)

148 # Success 

149 await self._cancel_tasks( (empty)

150 list(pending), list(done), cancel_wait 

151 ) 

152 context.result = finished.result() (empty)

153 context.result_exception = None (empty)

154 context.finish = True (empty)

155 run_tasks.remove(finished) (empty)

156 return context.result (empty)

157 

158 # Handle exception 

159 run_tasks.remove(finished) (empty)

160 retryable, increased = self._handle_exception( (empty)

161 context, config, finished, can_add_speed_up 

162 ) 

163 if increased: (empty)

164 can_add_speed_up = False (empty)

165 

166 if retryable: (empty)

167 await asyncio.sleep(cancel_wait) (empty)

168 break (empty)

169 else: 

170 await self._cancel_tasks( (empty)

171 list(pending), list(done), cancel_wait 

172 ) 

173 context.finish = True (empty)

174 break (empty)

175 

176 # 4. Check timeout 

177 if config.timeout > 0 and context.elapsed > config.timeout: (empty)

178 context.result_exception = asyncio.TimeoutError( (empty)

179 f"execute more than {config.timeout} seconds" 

180 ) 

181 await self._cancel_tasks( (empty)

182 list(pending) + run_tasks, [], cancel_wait 

183 ) 

184 context.finish = True (empty)

185 

186 # 5. Return result 

187 return self._return_result(context, config, args, kwargs) (empty)

188 

189 finally: 

190 context.running_task_pool.add(-1) (empty)

191 if config.debug: (empty)

192 self._log_finish(context) (empty)