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

1"""Sync runner module for synchronous task execution.""" 

2 

3import concurrent.futures (empty)

4import logging (empty)

5import time (empty)

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

7 

8from .base import BaseRunner, _CancelledErrorGroups (empty)

9from .config import ScheduleConfig (empty)

10from .context import ExecutionContext (empty)

11 

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

13 

14 

15class SyncRunner(BaseRunner): (empty)

16 """Synchronous runner for executing scheduled tasks. 

17 

18 Uses ThreadPoolExecutor to execute tasks synchronously, 

19 supporting retry, hedge, and rate limiting features. 

20 """ 

21 

22 def __init__(self): (empty)

23 super().__init__() (empty)

24 self._executor: Optional[ (empty)

25 concurrent.futures.ThreadPoolExecutor 

26 ] = None 

27 

28 def set_executor( (empty)

29 self, executor: concurrent.futures.ThreadPoolExecutor 

30 ): 

31 """Set the thread pool executor. 

32 

33 Args: 

34 executor: Thread pool executor for task execution. 

35 """ 

36 self._executor = executor (empty)

37 

38 def create_task( (empty)

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

40 ) -> concurrent.futures.Future: 

41 """Create and submit a task to the executor.""" 

42 if self._executor is None: (empty)

43 raise RuntimeError("Executor not set") (empty)

44 return self._executor.submit(func, *args, **kwargs) (empty)

45 

46 def get_task_result( (empty)

47 self, task: concurrent.futures.Future 

48 ) -> Any: 

49 """Get the result of a completed task.""" 

50 return task.result() (empty)

51 

52 def get_task_exception( (empty)

53 self, task: concurrent.futures.Future 

54 ) -> Optional[BaseException]: 

55 """Get the exception from a failed task, or None.""" 

56 return task.exception() (empty)

57 

58 def is_task_cancelled( (empty)

59 self, task: concurrent.futures.Future 

60 ) -> bool: 

61 """Check if a task was cancelled.""" 

62 return task.cancelled() (empty)

63 

64 def _cancel_tasks( (empty)

65 self, 

66 pending: Sequence[concurrent.futures.Future], (empty)

67 done: List, (empty)

68 retry_interval: float, (empty)

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

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

71 while done: (empty)

72 try: (empty)

73 done.pop() (empty)

74 except _CancelledErrorGroups: 

75 pass 

76 

77 for task in pending: (empty)

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

79 task.cancel() (empty)

80 

81 if pending: (empty)

82 try: (empty)

83 concurrent.futures.wait( (empty)

84 pending, (empty)

85 timeout=retry_interval, (empty)

86 return_when=concurrent.futures.ALL_COMPLETED, (empty)

87 ) 

88 except _CancelledErrorGroups: 

89 pass 

90 

91 def execute_impl( (empty)

92 self, 

93 context: ExecutionContext, 

94 config: ScheduleConfig, 

95 *args: Any, 

96 **kwargs: Any, 

97 ) -> Any: 

98 """Sync execution implementation.""" 

99 run_tasks: List[concurrent.futures.Future] = [] (empty)

100 

101 try: (empty)

102 while not context.finish and ( (empty)

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

104 ): 

105 # 1. Submit new tasks if allowed 

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

107 limiter_timeout = context.get_limiter_timeout( (empty)

108 len(run_tasks) 

109 ) 

110 

111 if not context.limiter or context.limiter.allow( (empty)

112 limiter_timeout if limiter_timeout > 0 else None 

113 ): 

114 if context.is_hedge_submit( (empty)

115 len(run_tasks) 

116 ): # pragma: no cover 

117 context.record_hedge() (empty)

118 elif ( (empty)

119 context.result_exception is not None 

120 ): # pragma: no cover 

121 context.record_exception(context.result_exception) (empty)

122 context.result_exception = None (empty)

123 

124 run_tasks.append( (empty)

125 self._executor.submit( 

126 context.func, *args, **kwargs 

127 ) 

128 ) 

129 context.mark_task_submitted() (empty)

130 

131 # 2. Wait for tasks 

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

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

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

135 if cur_timeout > 0: (empty)

136 done, pending = concurrent.futures.wait( (empty)

137 run_tasks, 

138 timeout=cur_timeout, 

139 return_when=concurrent.futures.FIRST_COMPLETED, 

140 ) 

141 else: # pragma: no cover 

142 done, pending = concurrent.futures.wait( (empty)

143 run_tasks, 

144 return_when=concurrent.futures.FIRST_COMPLETED, 

145 ) 

146 

147 # 3. Process results 

148 can_add_speed_up = context.can_increase_speed() (empty)

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

150 

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

152 run_tasks.remove(finished) (empty)

153 

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

155 continue 

156 

157 exception = finished.exception() (empty)

158 if exception is None: (empty)

159 # Success 

160 self._cancel_tasks( (empty)

161 list(pending), list(done), cancel_wait 

162 ) 

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

164 context.result_exception = None (empty)

165 context.finish = True (empty)

166 break (empty)

167 

168 # Handle exception 

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

170 context, config, finished, can_add_speed_up 

171 ) 

172 if increased: (empty)

173 can_add_speed_up = False (empty)

174 

175 if retryable: (empty)

176 time.sleep(cancel_wait) (empty)

177 break (empty)

178 else: 

179 self._cancel_tasks( (empty)

180 list(pending), [], cancel_wait 

181 ) 

182 context.finish = True (empty)

183 break (empty)

184 

185 # 4. Check timeout 

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

187 context.result_exception = TimeoutError( (empty)

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

189 ) 

190 self._cancel_tasks( (empty)

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

192 ) 

193 context.finish = True (empty)

194 

195 # 5. Return result 

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

197 

198 finally: 

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

200 if config.debug: (empty)

201 self._log_finish(context) (empty)