Coverage for core / src / sensorkit / core / executor.py: 84%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import contextlib 

6from collections.abc import AsyncGenerator, Coroutine 

7from datetime import UTC, datetime, timedelta 

8from typing import Any, Callable 

9 

10from loguru import logger 

11 

12from sensorkit.common.keyword import KeywordDict 

13from sensorkit.core.controller import ControllerClient 

14from sensorkit.core.task import Task, TaskContextMap, TaskExecution, TaskSubmission 

15 

16type TaskFactoryResult = Task | TaskSubmission | None 

17type TaskFactoryFunc = Callable[ 

18 [], 

19 Coroutine[Any, Any, TaskFactoryResult] 

20 | AsyncGenerator[TaskFactoryResult, TaskExecution] 

21 | TaskFactoryResult, 

22] 

23type TaskChangeCallback = Callable[[TaskExecution | None], Coroutine[Any, Any, None]] 

24 

25 

26class NoTaskAvailable(Exception): 

27 """There is no task available for execution.""" 

28 

29 

30async def _invoke_factory( 

31 func: TaskFactoryFunc, 

32) -> tuple[AsyncGenerator[TaskFactoryResult, TaskExecution] | None, TaskSubmission]: 

33 """Invoke a task factory and normalize its output. 

34 

35 The factory may be a plain function or coroutine returning a `Task`/`TaskSubmission`/`None`, or 

36 an async generator that yields one of those and is later resumed with the minted 

37 `TaskExecution` (see `_resume_factory`). 

38 

39 Returns: 

40 A tuple of the generator to resume after dispatch (or `None` for non-generator factories) 

41 and the normalized `TaskSubmission`. 

42 

43 Raises: 

44 NoTaskAvailable: If the factory offers no task this cycle. 

45 """ 

46 obj = func() 

47 

48 match obj: 

49 case AsyncGenerator(): 

50 agen = obj 

51 try: 

52 result = await agen.asend(None) 

53 except StopAsyncIteration as e: 

54 # A generator that never yields is a no-task scenario rather than an error. 

55 raise NoTaskAvailable from e 

56 case Coroutine(): 

57 agen = None 

58 result = await obj 

59 case Task() | TaskSubmission() | None: 

60 agen = None 

61 result = obj 

62 case _: 

63 raise RuntimeError("task factory returned invalid type") 

64 

65 if result is None: 

66 # The factory had no task for us. Close any generator so its finally-blocks run. 

67 if agen is not None: 

68 await agen.aclose() 

69 

70 raise NoTaskAvailable 

71 

72 # Normalize a bare task into a submission with no execution parameters. 

73 submission = result if isinstance(result, TaskSubmission) else TaskSubmission(task=result) 

74 

75 return agen, submission 

76 

77 

78async def _resume_factory( 

79 agen: AsyncGenerator[TaskFactoryResult, TaskExecution], execution: TaskExecution 

80) -> None: 

81 """Resume a factory generator once, handing it the minted execution. 

82 

83 The generator may simply end (ignoring the result), or `await` the execution itself to consume 

84 the result before ending. Either way it is expected to stop after a single resumption. 

85 """ 

86 try: 

87 await agen.asend(execution) 

88 except StopAsyncIteration: 

89 # Expected: the factory ran its post-dispatch logic (if any) and ended. 

90 return 

91 else: 

92 logger.warning( 

93 "Task factory generator yielded more than once! This indicates a programming error." 

94 ) 

95 await agen.aclose() 

96 

97 

98class TaskingLoop: 

99 """Executes a loop that sources tasks from a factory and runs them on a target Controller.""" 

100 

101 def __init__( 

102 self, 

103 controller: ControllerClient, 

104 factory_func: TaskFactoryFunc, 

105 contexts: TaskContextMap, 

106 task_group: asyncio.TaskGroup, 

107 *, 

108 on_task_change: TaskChangeCallback | None = None, 

109 ): 

110 self.controller = controller 

111 self.task_factory: TaskFactoryFunc = factory_func 

112 self.contexts = contexts 

113 self._task_group = task_group 

114 # Invoked with the live `TaskExecution` when a task begins running and with `None` when it 

115 # completes, so an owner (the program) can publish its current tasking state. Error/abort 

116 # exits are left to the owner's loop-teardown path rather than signalled here. 

117 self._on_task_change = on_task_change 

118 self._aio_task: asyncio.Task | None = None 

119 self._stop_requested = asyncio.Event() 

120 

121 def start(self): 

122 """Activate the tasking loop.""" 

123 if self._aio_task: 

124 raise RuntimeError("Tasking loop was already started") 

125 

126 self._aio_task = self._task_group.create_task(self._tasking_loop()) 

127 return self._aio_task 

128 

129 async def stop(self, *, timeout=None): 

130 """Deactivate the tasking loop, invalidating the object.""" 

131 try: 

132 # Wait for a graceful stop. 

133 self._stop_requested.set() 

134 

135 async with asyncio.timeout(timeout): 

136 await self._aio_task 

137 except TimeoutError: 

138 # The loop task was cancelled and will call abort on the Controller. 

139 with contextlib.suppress(asyncio.CancelledError): 

140 await self._aio_task 

141 

142 @property 

143 def stop_requested(self): 

144 """Return True if a stop has been requested for this tasking loop.""" 

145 return self._stop_requested.is_set() 

146 

147 def _resolve_context(self, submission: TaskSubmission) -> KeywordDict | None: 

148 """Merge the factory-supplied context with the per-task-type context bundle.""" 

149 context = KeywordDict(submission.context or {}) 

150 

151 if type_context := self.contexts.get(submission.task.task_type): 

152 context.update(type_context) 

153 

154 return context or None 

155 

156 def _resolve_expiry(self, submission: TaskSubmission) -> datetime | None: 

157 """Resolve the absolute expiry time for a submission. 

158 

159 Uses the factory-supplied `expiry_time` when present, otherwise the task's 

160 `default_expiry`. A `timedelta` default is anchored to the current time. 

161 """ 

162 if submission.expiry_time is not None: 

163 return submission.expiry_time 

164 

165 default = submission.task.default_expiry() 

166 

167 if isinstance(default, timedelta): 

168 return datetime.now(UTC) + default 

169 

170 return default 

171 

172 async def _tasking_loop(self): 

173 while not self._stop_requested.is_set(): 

174 # Source the next task from the factory. 

175 try: 

176 agen, submission = await _invoke_factory(self.task_factory) 

177 except NoTaskAvailable: 

178 with contextlib.suppress(TimeoutError): 

179 await asyncio.wait_for(self._stop_requested.wait(), timeout=5.0) 

180 

181 continue 

182 

183 # Dispatch and run it. A fatal error ends the loop. 

184 if not await self._dispatch(agen, submission): 

185 return 

186 

187 async def _dispatch( 

188 self, 

189 agen: AsyncGenerator[TaskFactoryResult, TaskExecution] | None, 

190 submission: TaskSubmission, 

191 ) -> bool: 

192 """Dispatch a single task and run it to completion. 

193 

194 Returns: 

195 True if the task completed (success, or a failure the factory handled without 

196 re-raising); False if a fatal error occurred and the loop should stop. 

197 """ 

198 task = submission.task 

199 context = self._resolve_context(submission) 

200 expiry_time = self._resolve_expiry(submission) 

201 ttl = (expiry_time - datetime.now(UTC)).total_seconds() if expiry_time else None 

202 

203 logger.info( 

204 f"Dispatching {task.task_type} ({type(task).__name__}) to {self.controller.entity}" 

205 ) 

206 

207 # Instruct the controller to begin execution. The controller mints the identity and binds 

208 # the in-flight call as the execution's result future. 

209 try: 

210 execution = await self.controller.start_task( 

211 task, context=context, expiry_time=expiry_time 

212 ) 

213 except (Exception, asyncio.CancelledError) as e: 

214 # Dispatch failed before the factory was resumed. Deliver the error into the still- 

215 # suspended generator so its except/finally blocks run, then abandon this task. 

216 logger.error(f"Failed to dispatch {task.task_type} to {self.controller.entity}: {e!r}") 

217 

218 if agen is not None: 

219 with contextlib.suppress(BaseException): 

220 await agen.athrow(e) 

221 

222 return False 

223 

224 # Client-side association: the program can reach the execution (and its result future) via 

225 # `task.execution`. 

226 task.associate_execution(execution) 

227 await self._notify_task_change(execution) 

228 

229 try: 

230 async with asyncio.timeout(ttl): 

231 # Resume the factory with the running execution (task_id now known). It may await 

232 # the execution to consume the result and run post-task logic. 

233 if agen is not None: 

234 await _resume_factory(agen, execution) 

235 

236 # Guarantee completion before the next iteration, regardless of whether the factory 

237 # awaited the execution itself. If it did, the call is already settled and this 

238 # returns (or re-raises) immediately; if it consumed an error without re-raising, 

239 # the execution is done and we leave that decision to the factory. 

240 if not execution.done(): 

241 await execution 

242 except (Exception, asyncio.CancelledError) as e: 

243 # On error/abort the loop terminates; clearing the tasking state is left to the owner's 

244 # teardown so we don't emit while unwinding (possibly under cancellation). 

245 await self._abort(execution, e) 

246 return False 

247 

248 await self._notify_task_change(None) 

249 logger.info(f"Finished {task.task_type} ({type(task).__name__})") 

250 

251 return True 

252 

253 async def _notify_task_change(self, execution: TaskExecution | None): 

254 """Signal the owner that the currently running task changed, if a callback is registered.""" 

255 if self._on_task_change is not None: 

256 await self._on_task_change(execution) 

257 

258 async def _abort(self, execution: TaskExecution, e: BaseException): 

259 """Log and abort an in-flight task following a timeout, interrupt, or unhandled error.""" 

260 msg = f"Aborting {execution.task.task_type} on {self.controller.entity}" 

261 

262 match e: 

263 case asyncio.CancelledError(): 

264 logger.error(f"{msg} due to interrupt") 

265 case _: 

266 logger.error(f"{msg} due to error") 

267 logger.opt(exception=e).debug( 

268 f"{type(e).__name__} during execution of {execution.task_id}" 

269 ) 

270 

271 with contextlib.suppress(BaseException): 

272 # Kick off an abort here, but don't wait for it to complete. 

273 await self.controller.abort_task(execution.task_id).invoke()