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

1"""Main schedule decorator implementation.""" 

2 

3import atexit (empty)

4import concurrent.futures (empty)

5import functools (empty)

6import inspect (empty)

7import threading (empty)

8from typing import ( (empty)

9 Any, Callable, List, Optional, Tuple, Type, Union, 

10) 

11 

12from .adaptive_limiter import AdaptiveRateLimiter (empty)

13from .async_runner import AsyncRunner (empty)

14from .circuit_breaker import CircuitBreaker (empty)

15from .config import ScheduleConfig, _sentinel_none (empty)

16from .context import ExecutionContext (empty)

17from .sync_runner import SyncRunner (empty)

18from .limiter import LimiterPool (empty)

19from .metrics import ScheduleMetrics (empty)

20from .pools import RunningTaskPool (empty)

21from .telemetry import TelemetryExporter (empty)

22 

23# Global registry for ThreadPoolExecutors created by sync schedule decorators. 

24# Protected by a lock to ensure thread safety during registration. 

25_executor_registry_lock = threading.Lock() (empty)

26_executor_registry: List[concurrent.futures.ThreadPoolExecutor] = [] (empty)

27_atexit_registered = False (empty)

28 

29 

30def _shutdown_executors() -> None: (empty)

31 """Shutdown all registered ThreadPoolExecutors. 

32 

33 Called automatically at process exit via atexit to prevent thread leaks. 

34 Uses wait=False to avoid blocking process shutdown. 

35 """ 

36 with _executor_registry_lock: (empty)

37 for executor in _executor_registry: (empty)

38 executor.shutdown(wait=False) (empty)

39 _executor_registry.clear() (empty)

40 

41 

42def _register_executor( (empty)

43 executor: concurrent.futures.ThreadPoolExecutor 

44) -> None: 

45 """Register a ThreadPoolExecutor for cleanup at process exit. 

46 

47 Args: 

48 executor: The ThreadPoolExecutor to register. 

49 """ 

50 global _atexit_registered 

51 with _executor_registry_lock: (empty)

52 _executor_registry.append(executor) (empty)

53 if not _atexit_registered: (empty)

54 atexit.register(_shutdown_executors) (empty)

55 _atexit_registered = True (empty)

56 

57 

58def schedule( (empty)

59 timeout: float = 0, 

60 hedged_request_time: float = 0, 

61 hedged_request_proportion: float = 0.05, 

62 hedged_request_max_times: int = 2, 

63 speed_up_max_multiply: int = 5, 

64 retry_times: int = 0, 

65 retry_interval: float = 1, 

66 limit_size: int = 0, 

67 limit_window: float = 1.0, 

68 key_func: object = _sentinel_none, 

69 exception_types: Union[ 

70 Type[BaseException], 

71 Tuple[Type[BaseException], ...], 

72 ] = BaseException, 

73 default_result: object = _sentinel_none, 

74 debug: bool = False, 

75 metrics_callback: Optional[Callable[[ScheduleMetrics], None]] = None, 

76 circuit_breaker_threshold: int = 0, 

77 circuit_breaker_recovery: float = 30.0, 

78 adaptive_limit: bool = False, 

79 adaptive_limit_min: int = 10, 

80 adaptive_limit_max: int = 500, 

81 adaptive_error_threshold: float = 0.3, 

82 adaptive_latency_threshold: float = 5.0, 

83 adaptive_window_seconds: float = 10.0, 

84 priority: int = 5, 

85 telemetry_exporter: Optional[TelemetryExporter] = None, 

86): 

87 """Decorator for scheduling function execution with retry, hedging, 

88 rate limiting, etc. 

89 

90 Args: 

91 timeout: Timeout in seconds, 0 means no timeout 

92 hedged_request_time: Time to trigger hedged requests 

93 hedged_request_proportion: Max proportion for hedged requests 

94 hedged_request_max_times: Max number of hedged requests 

95 speed_up_max_multiply: Max speed up multiplier 

96 retry_times: Max retry times (not including first call) 

97 retry_interval: Retry interval in seconds 

98 limit_size: Rate limit size 

99 limit_window: Rate limit window in seconds 

100 key_func: Function to generate pool key 

101 exception_types: Exceptions to catch and retry 

102 default_result: Default result on failure 

103 debug: Enable debug logging 

104 metrics_callback: Optional callback invoked after each execution 

105 with a ScheduleMetrics instance containing execution statistics 

106 circuit_breaker_threshold: Failure threshold to trip the circuit 

107 breaker. 0 disables the circuit breaker. 

108 circuit_breaker_recovery: Seconds before the breaker transitions 

109 from OPEN to HALF_OPEN for probe requests. 

110 adaptive_limit: Enable adaptive rate limiting. 

111 adaptive_limit_min: Minimum rate limit floor. 

112 adaptive_limit_max: Maximum rate limit ceiling. 

113 adaptive_error_threshold: Error rate threshold to trigger slowdown. 

114 adaptive_latency_threshold: Avg latency threshold to trigger slowdown. 

115 adaptive_window_seconds: Sliding window duration for statistics. 

116 priority: Task priority (0-9, lower = higher). Default 5 (NORMAL). 

117 telemetry_exporter: Optional TelemetryExporter for OTel metrics export. 

118 

119 Returns: 

120 Decorated function 

121 """ 

122 config = ScheduleConfig( (empty)

123 timeout=timeout, 

124 hedged_request_time=hedged_request_time, 

125 hedged_request_proportion=hedged_request_proportion, 

126 hedged_request_max_times=hedged_request_max_times, 

127 speed_up_max_multiply=speed_up_max_multiply, 

128 retry_times=retry_times, 

129 retry_interval=retry_interval, 

130 limit_size=limit_size, 

131 limit_window=limit_window, 

132 key_func=key_func, 

133 exception_types=exception_types, 

134 default_result=default_result, 

135 debug=debug, 

136 adaptive_limit=adaptive_limit, 

137 adaptive_limit_min=adaptive_limit_min, 

138 adaptive_limit_max=adaptive_limit_max, 

139 adaptive_error_threshold=adaptive_error_threshold, 

140 adaptive_latency_threshold=adaptive_latency_threshold, 

141 adaptive_window_seconds=adaptive_window_seconds, 

142 priority=priority, 

143 ) 

144 

145 # Create adaptive rate limiter if enabled 

146 adaptive_limiter: Optional[AdaptiveRateLimiter] = None (empty)

147 if adaptive_limit and limit_size > 0: (empty)

148 adaptive_limiter = AdaptiveRateLimiter( (empty)

149 initial_limit=limit_size, 

150 limit_min=adaptive_limit_min, 

151 limit_max=adaptive_limit_max, 

152 error_threshold=adaptive_error_threshold, 

153 latency_threshold=adaptive_latency_threshold, 

154 window_seconds=adaptive_window_seconds, 

155 ) 

156 

157 # Create circuit breaker if threshold > 0 

158 breaker: Optional[CircuitBreaker] = None (empty)

159 if circuit_breaker_threshold > 0: (empty)

160 breaker = CircuitBreaker( (empty)

161 failure_threshold=circuit_breaker_threshold, 

162 recovery_timeout=circuit_breaker_recovery, 

163 ) 

164 

165 def decorator(func: Callable[..., Any]) -> Callable[..., Any]: (empty)

166 if inspect.iscoroutinefunction(func): (empty)

167 # Async function 

168 runner = AsyncRunner() (empty)

169 

170 @functools.wraps(func) (empty)

171 async def async_wrapper(*args, **kwargs): (empty)

172 # Generate key 

173 key = func.__qualname__ (empty)

174 if key_func is not _sentinel_none and callable(key_func): (empty)

175 key = f"{func.__qualname__}.{key_func(*args, **kwargs)}" (empty)

176 

177 # Get pools 

178 if timeout <= 0: (empty)

179 running_task = RunningTaskPool.get_pool(key) (empty)

180 else: 

181 running_task = RunningTaskPool.get_pool( (empty)

182 key, window_interval=timeout 

183 ) 

184 limiter = LimiterPool.get_pool(key, limit_size, limit_window) (empty)

185 

186 # Check circuit breaker before execution 

187 if breaker is not None and not breaker.allow_request(): (empty)

188 raise RuntimeError( (empty)

189 f"Circuit breaker is open for {func.__qualname__}" 

190 ) 

191 

192 # Create context 

193 context = ExecutionContext( (empty)

194 func, config, key, running_task, limiter 

195 ) 

196 running_task.add(1) (empty)

197 

198 # Execute 

199 result = await runner.execute_impl( (empty)

200 context, config, *args, **kwargs 

201 ) 

202 

203 # Record to adaptive limiter 

204 if adaptive_limiter is not None: (empty)

205 metrics = context.build_metrics() (empty)

206 adaptive_limiter.record( (empty)

207 latency_seconds=metrics.total_latency_ms / 1000.0, 

208 success=context.result_exception is None, 

209 ) 

210 

211 # Update circuit breaker based on execution outcome 

212 if breaker is not None: (empty)

213 if context.result_exception is not None: (empty)

214 breaker.record_failure() (empty)

215 else: 

216 breaker.record_success() (empty)

217 

218 # Invoke metrics callback 

219 built_metrics = context.build_metrics() (empty)

220 if metrics_callback is not None: (empty)

221 metrics_callback(built_metrics) (empty)

222 if telemetry_exporter is not None: (empty)

223 built_metrics.export_to_otel(telemetry_exporter) (empty)

224 

225 return result (empty)

226 

227 return async_wrapper (empty)

228 

229 else: 

230 # Sync function 

231 runner = SyncRunner() (empty)

232 executor = concurrent.futures.ThreadPoolExecutor() (empty)

233 _register_executor(executor) (empty)

234 

235 @functools.wraps(func) (empty)

236 def sync_wrapper(*args, **kwargs): (empty)

237 # Generate key 

238 key = func.__qualname__ (empty)

239 if key_func is not _sentinel_none and callable(key_func): (empty)

240 key = f"{func.__qualname__}.{key_func(*args, **kwargs)}" (empty)

241 

242 # Get pools 

243 if timeout <= 0: (empty)

244 running_task = RunningTaskPool.get_pool(key) (empty)

245 else: 

246 running_task = RunningTaskPool.get_pool( (empty)

247 key, window_interval=timeout 

248 ) 

249 limiter = LimiterPool.get_pool(key, limit_size, limit_window) (empty)

250 

251 # Check circuit breaker before execution 

252 if breaker is not None and not breaker.allow_request(): (empty)

253 raise RuntimeError( (empty)

254 f"Circuit breaker is open for {func.__qualname__}" 

255 ) 

256 

257 # Create context and execute with shared executor 

258 runner.set_executor(executor) (empty)

259 context = ExecutionContext( (empty)

260 func, config, key, running_task, limiter 

261 ) 

262 running_task.add(1) (empty)

263 

264 result = runner.execute_impl( (empty)

265 context, config, *args, **kwargs 

266 ) 

267 

268 # Record to adaptive limiter 

269 if adaptive_limiter is not None: (empty)

270 metrics = context.build_metrics() (empty)

271 adaptive_limiter.record( (empty)

272 latency_seconds=metrics.total_latency_ms / 1000.0, 

273 success=context.result_exception is None, 

274 ) 

275 

276 # Update circuit breaker based on execution outcome 

277 if breaker is not None: (empty)

278 if context.result_exception is not None: (empty)

279 breaker.record_failure() (empty)

280 else: 

281 breaker.record_success() (empty)

282 

283 # Invoke metrics callback 

284 built_metrics = context.build_metrics() (empty)

285 if metrics_callback is not None: (empty)

286 metrics_callback(built_metrics) (empty)

287 if telemetry_exporter is not None: (empty)

288 built_metrics.export_to_otel(telemetry_exporter) (empty)

289 

290 return result (empty)

291 

292 return sync_wrapper (empty)

293 

294 return decorator (empty)