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

1"""Execution context for managing state during a single schedule 

2invocation.""" 

3 

4import time (empty)

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

6 

7from .config import ScheduleConfig (empty)

8from .limiter import LimiterPool (empty)

9from .metrics import ScheduleMetrics (empty)

10from .pools import RunningTaskPool (empty)

11 

12# Exponent base and step for hedged request threshold calculation: 

13# threshold = pool_size + multiply ** ( 

14# HEDGE_EXPONENT_BASE + times * HEDGE_EXPONENT_STEP 

15# ) 

16HEDGE_EXPONENT_BASE = 0.5 (empty)

17HEDGE_EXPONENT_STEP = 0.5 (empty)

18 

19# Minimum timeout (seconds) returned by get_limiter_timeout to avoid zero-waits 

20MIN_LIMITER_TIMEOUT = 0.001 (empty)

21 

22 

23class ExecutionContext: (empty)

24 """Manages execution state for a single schedule invocation.""" 

25 

26 def __init__( (empty)

27 self, 

28 func: Callable, 

29 config: ScheduleConfig, 

30 key: str, 

31 running_task_pool: RunningTaskPool, 

32 limiter: Optional[LimiterPool], 

33 ): 

34 """Initialize execution context. 

35 

36 Args: 

37 func: Function being decorated 

38 config: Schedule configuration 

39 key: Pool key 

40 running_task_pool: Running task pool instance 

41 limiter: Rate limiter instance or None 

42 """ 

43 self.func = func (empty)

44 self.config = config (empty)

45 self.key = key (empty)

46 self.running_task_pool = running_task_pool (empty)

47 self.limiter = limiter (empty)

48 self.priority: int = config.priority (empty)

49 

50 # Execution state 

51 self.cur_times = 0 (empty)

52 self.cur_speed_up_multiply = 0 (empty)

53 self.start_time = time.perf_counter() (empty)

54 self.last_submit_time = self.start_time (empty)

55 self.cur_hedged_request_times = 1 (empty)

56 

57 # Results 

58 self.result: Any = None (empty)

59 self.result_exception: Optional[BaseException] = None (empty)

60 self.result_exception_list: List[str] = [] (empty)

61 

62 # Control 

63 self.finish = False (empty)

64 

65 def can_submit_task(self, run_tasks_count: int) -> bool: (empty)

66 """Check if a new task can be submitted. 

67 

68 Args: 

69 run_tasks_count: Current number of running tasks 

70 

71 Returns: 

72 True if can submit 

73 """ 

74 if self.cur_times > self.config.retry_times: (empty)

75 return False (empty)

76 

77 # Always allow if no running tasks 

78 if run_tasks_count <= 0: (empty)

79 return True (empty)

80 

81 # Allow if under speed up multiplier 

82 if run_tasks_count < self.cur_speed_up_multiply: (empty)

83 if self.running_task_pool.can_submit( (empty)

84 run_tasks_count + 1 

85 ): 

86 return True (empty)

87 

88 # Allow hedged request 

89 if self._should_hedge(run_tasks_count): (empty)

90 return True (empty)

91 

92 return False (empty)

93 

94 def _should_hedge(self, run_tasks_count: int) -> bool: (empty)

95 """Check if hedged request should be submitted. 

96 

97 Args: 

98 run_tasks_count: Current number of running tasks 

99 

100 Returns: 

101 True if should hedge 

102 """ 

103 if self.config.hedged_request_time <= 0: (empty)

104 return False (empty)

105 

106 if self.cur_hedged_request_times > ( (empty)

107 self.config.hedged_request_max_times 

108 ): 

109 return False (empty)

110 

111 time_since_last = ( (empty)

112 time.perf_counter() - self.last_submit_time 

113 ) 

114 if not ( (empty)

115 self.config.hedged_request_time < time_since_last 

116 or self.cur_hedged_request_times > 1 

117 ): 

118 return False (empty)

119 

120 threshold = run_tasks_count + ( (empty)

121 self.config.hedged_request_multiply ** ( 

122 HEDGE_EXPONENT_BASE 

123 + self.cur_hedged_request_times * HEDGE_EXPONENT_STEP 

124 ) 

125 ) 

126 return self.running_task_pool.can_submit(threshold) (empty)

127 

128 def is_hedge_submit(self, run_tasks_count: int) -> bool: (empty)

129 """Check if current submission is a hedge. 

130 

131 Args: 

132 run_tasks_count: Current number of running tasks 

133 

134 Returns: 

135 True if this is a hedge submission 

136 """ 

137 if self.config.hedged_request_time <= 0: (empty)

138 return False (empty)

139 

140 if self.cur_speed_up_multiply > run_tasks_count: (empty)

141 return False (empty)

142 

143 time_since_last = ( (empty)

144 time.perf_counter() - self.last_submit_time 

145 ) 

146 if not ( (empty)

147 self.config.hedged_request_time < time_since_last 

148 or self.cur_hedged_request_times > 1 

149 ): 

150 return False (empty)

151 

152 if self.cur_hedged_request_times > ( (empty)

153 self.config.hedged_request_max_times 

154 ): 

155 return False (empty)

156 

157 threshold = run_tasks_count + ( (empty)

158 self.config.hedged_request_multiply ** ( 

159 HEDGE_EXPONENT_BASE 

160 + self.cur_hedged_request_times * HEDGE_EXPONENT_STEP 

161 ) 

162 ) 

163 return self.running_task_pool.can_submit(threshold) (empty)

164 

165 def compute_timeout(self, run_tasks_count: int) -> float: (empty)

166 """Compute timeout for current wait. 

167 

168 Args: 

169 run_tasks_count: Current number of running tasks 

170 

171 Returns: 

172 Timeout in seconds (0 means no timeout) 

173 """ 

174 cur_timeout = 0.0 (empty)

175 

176 # Base timeout 

177 if self.config.timeout > 0: (empty)

178 cur_timeout = ( (empty)

179 self.start_time 

180 + self.config.timeout 

181 - time.perf_counter() 

182 ) 

183 

184 # Hedge request timeout 

185 if ( (empty)

186 self.cur_hedged_request_times 

187 <= self.config.hedged_request_max_times 

188 and self.config.hedged_request_time > 0 

189 ): 

190 hedge_timeout = ( (empty)

191 self.start_time 

192 + self.config.hedged_request_time 

193 - time.perf_counter() 

194 ) 

195 if cur_timeout == 0 or hedge_timeout < cur_timeout: (empty)

196 cur_timeout = hedge_timeout (empty)

197 

198 # Speed up timeout (retry_interval) 

199 if run_tasks_count < self.cur_speed_up_multiply: (empty)

200 if self.cur_times < self.config.retry_times: (empty)

201 if ( (empty)

202 cur_timeout > self.config.retry_interval 

203 or cur_timeout == 0 

204 ): 

205 cur_timeout = self.config.retry_interval (empty)

206 

207 if cur_timeout < 0: (empty)

208 cur_timeout = self.config.retry_interval (empty)

209 

210 return cur_timeout (empty)

211 

212 def get_limiter_timeout(self, run_tasks_count: int) -> float: (empty)

213 """Get timeout for rate limiter. 

214 

215 Args: 

216 run_tasks_count: Current number of running tasks 

217 

218 Returns: 

219 Timeout in seconds (0 means no timeout) 

220 """ 

221 if run_tasks_count == 0: (empty)

222 if self.config.timeout > 0: (empty)

223 return max( (empty)

224 MIN_LIMITER_TIMEOUT, 

225 self.config.timeout - self.elapsed 

226 ) 

227 return 0 (empty)

228 

229 cur_timeout = self.config.retry_interval (empty)

230 if self.config.timeout > 0: (empty)

231 remaining = self.config.timeout - self.elapsed (empty)

232 if remaining < self.config.retry_interval: (empty)

233 cur_timeout = max( (empty)

234 MIN_LIMITER_TIMEOUT, remaining 

235 ) 

236 return cur_timeout (empty)

237 

238 @property (empty)

239 def elapsed(self) -> float: (empty)

240 """Get elapsed time since start.""" 

241 return time.perf_counter() - self.start_time (empty)

242 

243 @property (empty)

244 def is_timeout(self) -> bool: (empty)

245 """Check if execution has timed out.""" 

246 if self.config.timeout <= 0: (empty)

247 return False (empty)

248 return self.elapsed > self.config.timeout (empty)

249 

250 def record_exception(self, exception: BaseException): (empty)

251 """Record exception info.""" 

252 self.result_exception_list.append( (empty)

253 f"{type(exception).__name__} {str(exception)}" 

254 ) 

255 

256 def record_hedge(self): (empty)

257 """Record hedge request.""" 

258 self.result_exception_list.append(f"hedged_request: {self.elapsed}") (empty)

259 self.cur_hedged_request_times += 1 (empty)

260 

261 def mark_task_submitted(self): (empty)

262 """Mark that a task was submitted.""" 

263 self.cur_times += 1 (empty)

264 self.last_submit_time = time.perf_counter() (empty)

265 

266 def can_increase_speed(self) -> bool: (empty)

267 """Check if speed can be increased.""" 

268 return self.cur_speed_up_multiply < self.config.speed_up_max_multiply (empty)

269 

270 def increase_speed(self): (empty)

271 """Increase speed multiplier.""" 

272 self.cur_speed_up_multiply += 1 (empty)

273 

274 def reset_speed(self): (empty)

275 """Reset speed multiplier (on overload).""" 

276 self.cur_speed_up_multiply = 0 (empty)

277 

278 def build_metrics(self) -> ScheduleMetrics: (empty)

279 """Build execution metrics from the current context state. 

280 

281 Returns: 

282 ScheduleMetrics with aggregated execution statistics. 

283 """ 

284 total_latency_ms = self.elapsed * 1000.0 (empty)

285 retry_count = max(0, self.cur_times - 1) (empty)

286 failure_count = len(self.result_exception_list) (empty)

287 success_count = 1 if self.result_exception is None else 0 (empty)

288 total_calls = self.cur_times (empty)

289 avg_latency_ms = ( (empty)

290 total_latency_ms / total_calls if total_calls > 0 else 0.0 

291 ) 

292 return ScheduleMetrics( (empty)

293 total_calls=total_calls, 

294 success_count=success_count, 

295 failure_count=failure_count, 

296 retry_count=retry_count, 

297 total_latency_ms=total_latency_ms, 

298 avg_latency_ms=avg_latency_ms, 

299 )