Coverage for core / src / sensorkit / core / impl / controller.py: 85%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import contextlib 

6from _contextvars import ContextVar 

7from datetime import UTC, datetime 

8from typing import Callable, ClassVar, Collection, Mapping, override 

9 

10import uuid_utils.compat as uuid 

11from loguru import logger 

12 

13from sensorkit.backend.base import Entity 

14from sensorkit.backend.request import CallContext 

15from sensorkit.core.controller import ( 

16 AbortRequestMessage, 

17 AbortResponseMessage, 

18 ControllerDevice, 

19 ControllerEnableState, 

20 ControllerEnableStateRequest, 

21 ControllerInterface, 

22 ControllerOperatingState, 

23 ControllerState, 

24 ExecuteRequestMessage, 

25 ExecuteResponseMessage, 

26 InternalControllerState, 

27 TaskExecutionResult, 

28 TaskExecutionState, 

29 TaskFinishInfo, 

30 TaskHandlerCallback, 

31 abort_task_request, 

32 execute_task_request, 

33 set_enable_state_request, 

34) 

35from sensorkit.core.device import DeviceClient 

36from sensorkit.core.entity import ControllerDetails, EntityInfo 

37from sensorkit.core.impl.entity import EntityImpl 

38from sensorkit.core.task import Task, TaskExecution, TaskInfo 

39from sensorkit.data.context import Context, ContextSubscription 

40 

41 

42class ControllerDeviceMap(Mapping[str, ControllerDevice]): 

43 """Ordered mapping of device names to their ControllerDevice instances, created on demand.""" 

44 

45 def __init__(self): 

46 self._dict: dict[str, ControllerDevice] = {} 

47 

48 def get_or_add(self, name: str, *, get_client_func: Callable[[str], DeviceClient]): 

49 """Return the ControllerDevice for the given name, creating it if it does not yet exist.""" 

50 if name not in self._dict: 

51 client = get_client_func(name) 

52 self._dict[name] = ControllerDevice(client, ContextSubscription(client)) 

53 

54 return self._dict[name] 

55 

56 def __getitem__(self, key, /): 

57 return self._dict[key] 

58 

59 def __len__(self): 

60 return len(self._dict) 

61 

62 def __iter__(self): 

63 return iter(self._dict) 

64 

65 

66class ControllerImpl(EntityImpl, ControllerInterface): 

67 """Helper for implementing server-side functionality of a Controller.""" 

68 

69 current: ClassVar[ContextVar[ControllerImpl | None]] = ContextVar("current_controller", default=None) 

70 

71 def __init__(self, **kwargs): 

72 super().__init__(**kwargs) 

73 

74 self._enable_hooks: list[Callable[[], None]] = [] 

75 self._disable_hooks: list[Callable[[], None]] = [] 

76 self._task_handlers: dict[type[Task], TaskHandlerCallback] = {} 

77 self._task_asyncio: asyncio.Task | None = None 

78 self._devices = ControllerDeviceMap() 

79 

80 @override 

81 def task_running(self): 

82 return bool(self._task_asyncio) 

83 

84 @override 

85 def on_enable(self, func: Callable[[], None]): 

86 self._enable_hooks.append(func) 

87 return func 

88 

89 @override 

90 def on_disable(self, func: Callable[[], None]): 

91 self._disable_hooks.append(func) 

92 return func 

93 

94 @override 

95 async def init_impl(self): 

96 self._state = await ControllerState.recover_or_init( 

97 self, 

98 enable_state=ControllerEnableState(enabled=True), 

99 operating_state=ControllerOperatingState(current=InternalControllerState.UNKNOWN), 

100 execution_state=TaskExecutionState(execution=None), 

101 ) 

102 

103 @override 

104 async def attach_impl(self): 

105 if self._state.enable_state.enabled: 

106 await self._call_with_context(self._enable_hooks) 

107 

108 await self.handle_request(set_enable_state_request, self._set_enable_state) 

109 await self.handle_request(abort_task_request, self._abort_request) 

110 await self.handle_request(execute_task_request, self._execute_request) 

111 

112 await self.start_device_subscriptions() 

113 

114 @override 

115 async def detach_impl(self): 

116 await self.stop_device_subscriptions() 

117 

118 @override 

119 def use_device(self, name: str, *, subscribe: list[type] | None = None) -> DeviceClient: 

120 """Register a controlled device and return a client for it. 

121 

122 Registers a device that this controller interacts with. Optionally subscribes to 

123 specific keyword types published by the device. Subscribed keywords are automatically 

124 cached and made available in contexts via `update_context()`. 

125 

126 Note: 

127 Subscriptions remain inactive until `start_device_subscriptions()` is called. 

128 When using the declarative API, this occurs automatically after initialization. 

129 

130 Args: 

131 name: Entity identifier for the device (e.g., "mount", "camera"). 

132 subscribe: Optional list of keyword model types (e.g., Pydantic models) to 

133 subscribe to. Each subscribed keyword will be monitored, and its latest 

134 value cached for inclusion in contexts. 

135 

136 Returns: 

137 A client instance for interacting with the registered device. 

138 Access the subscription object via `get_device(name).subscription` to 

139 retrieve cached values directly. 

140 

141 Raises: 

142 KeyError: If attempting to access a device via `get_device(name)` that 

143 hasn't been registered with `use_device()`. 

144 

145 Examples: 

146 Basic device registration without subscriptions: 

147 

148 >>> camera = self.use_device("camera") 

149 

150 Register device and subscribe to specific keyword types: 

151 

152 >>> mount = self.use_device("mount", subscribe=[AltAzPointing, RADecPointing]) 

153 

154 Access cached subscription values from the ControlledDevice instance: 

155 

156 >>> pointing = self.get_device("mount").get(AltAzPointing) 

157 """ 

158 device = self._devices.get_or_add( 

159 name, 

160 get_client_func=lambda _: self.sensorkit().device(Entity((name,))), 

161 ) 

162 

163 if subscribe: 

164 for kw_type in subscribe: 

165 device.subscription.add(kw_type) 

166 

167 return device.client 

168 

169 @override 

170 def get_device(self, name: str) -> ControllerDevice: 

171 return self._devices[name] 

172 

173 @override 

174 def all_devices(self) -> Collection[ControllerDevice]: 

175 return self._devices.values() 

176 

177 @override 

178 async def start_device_subscriptions(self): 

179 """Start all keyword subscriptions registered via `use_device`. 

180 

181 This is called automatically by the declarative API after init 

182 callbacks have run. It may also be called manually if the 

183 controller is not using the declarative API. 

184 """ 

185 for device in self._devices.values(): 

186 await device.subscription.start() 

187 

188 @override 

189 async def stop_device_subscriptions(self): 

190 """Stop all keyword subscriptions registered via `use_device`. 

191 

192 Called automatically by the declarative API before deinit 

193 callbacks run. 

194 """ 

195 for device in self._devices.values(): 

196 await device.subscription.stop() 

197 

198 @override 

199 async def update_context( 

200 self, 

201 *args, 

202 **kwargs, 

203 ) -> Context: 

204 """Update the context of the currently executing task. 

205 

206 Merges the base context with fresh snapshots from all device keyword 

207 subscriptions and additional provided values. The updated context is 

208 persisted to the execution state and returned. 

209 

210 This method allows task handlers to refresh their context mid-execution, 

211 incorporating the latest device state. 

212 

213 Args: 

214 **kwargs: Additional literal key-value pairs to include in the context. 

215 

216 Returns: 

217 The newly updated `Context` containing (in precedence order, 

218 highest-last): current task context, fresh device snapshots, 

219 and *kwargs*. 

220 

221 Raises: 

222 RuntimeError: If called when no task is currently executing. 

223 

224 Examples: 

225 Update context with latest device data during task execution: 

226 

227 >>> async def my_task_handler(self, task): 

228 ... # ... some work ... 

229 ... ctx = await self.update_context() 

230 ... pointing = ctx.get(AltAzPointing) 

231 

232 Add custom values while updating: 

233 

234 >>> ctx = await self.update_context(iteration=5, timestamp=time.time()) 

235 """ 

236 current = self._state.execution_state 

237 

238 if not current.executing or current.execution is None: 

239 raise RuntimeError("no task executing") 

240 

241 ctx = Context(current.context) 

242 

243 for device in self._devices.values(): 

244 device.subscription.snapshot(into=ctx) 

245 

246 ctx.set(*args) 

247 

248 for key, value in kwargs.items(): 

249 ctx.set_value(key, value) 

250 

251 await self._state.update( 

252 self, 

253 TaskExecutionState( 

254 executing=current.executing, 

255 aborting=current.aborting, 

256 execution=current.execution, 

257 context=ctx, 

258 ), 

259 ) 

260 

261 return ctx 

262 

263 async def _set_enable_state(self, request: ControllerEnableStateRequest): 

264 if request.enable == self._state.enable_state.enabled: 

265 return 

266 

267 await self._state.update(self, ControllerEnableState(enabled=request.enable)) 

268 

269 if request.enable: 

270 await self._call_with_context(self._enable_hooks) 

271 else: 

272 #TODO: Abort. Need a nice way of calling a local "loopback" request. 

273 # await self.call_loopback(abort_task_request, AbortRequestMessage(task_id=None)) 

274 

275 await self._call_with_context(self._disable_hooks) 

276 

277 async def _abort_request( 

278 self, 

279 requested: AbortRequestMessage, 

280 call: CallContext[AbortResponseMessage, None], 

281 ): 

282 # Cannot abort if no task is running. 

283 if not self.task_running(): 

284 call.reject(response=AbortResponseMessage(aborting=False, task_id=None)) 

285 return 

286 

287 aio_task = self._task_asyncio 

288 task_id = self._state.execution_state.execution.task_id 

289 

290 # If given, ensure the task_id matches what is running. 

291 if requested.task_id and task_id != requested.task_id: 

292 call.reject(response=AbortResponseMessage(aborting=False, task_id=None)) 

293 return 

294 

295 # Send an accept response and update status indicating we are about to abort. 

296 logger.debug(f"aborting running task {task_id}") 

297 call.accept(response=AbortResponseMessage(aborting=True, task_id=task_id)) 

298 await self._state.update( 

299 self, 

300 TaskExecutionState( 

301 executing=True, 

302 aborting=True, 

303 execution=self._state.execution_state.execution 

304 ), 

305 ) 

306 

307 # Abort the asyncio.Task running the SensorKit Task. 

308 aio_task.cancel("Aborted by remote request") 

309 

310 # Wait for the task to die. 

311 with contextlib.suppress(asyncio.CancelledError, Exception): 

312 await call.progress_from_task(aio_task, cadence=5, ttl=10) 

313 

314 if aio_task.done(): 

315 # Success. 

316 await call.succeed(result=None) 

317 else: 

318 await call.fail() 

319 

320 async def _execute_request( 

321 self, 

322 msg: ExecuteRequestMessage, 

323 call: CallContext[ExecuteResponseMessage, TaskExecutionResult], 

324 ): 

325 task = msg.task 

326 

327 # Mint the execution envelope: the controller owns identity (task_id, controller_id) and 

328 # records the client-supplied execution parameters (context, end_time). 

329 execution = TaskExecution( 

330 task=task, 

331 task_id=uuid.uuid7(), 

332 controller_id=str(self.entity), 

333 context=msg.context, 

334 expiry_time=msg.expiry_time, 

335 ) 

336 

337 # Inject information about this task execution into the task context. 

338 execution.context.set( 

339 TaskInfo( 

340 task=task, 

341 task_id=execution.task_id, 

342 controller_id=execution.controller_id, 

343 ) 

344 ) 

345 

346 # Send the initial response. 

347 response = ExecuteResponseMessage(execution=execution) 

348 

349 if not self._state.enable_state.enabled: 

350 logger.warning(f"Rejecting {task.task_type} task: Controller is disabled") 

351 call.reject(response=response) 

352 return 

353 

354 if not msg.interrupt and self.task_running(): 

355 logger.warning(f"Rejecting {task.task_type} task: Task already in progress") 

356 call.reject(response=response) 

357 return 

358 

359 if type(task) not in self._task_handlers: 

360 logger.warning(f"Rejecting {task.task_type} task: No handler registered") 

361 call.reject(response=response) 

362 return 

363 

364 call.accept(response=response) 

365 aio_task: asyncio.Task | None = None 

366 

367 try: 

368 # Interrupt the current task, if there is one. 

369 if self._task_asyncio: 

370 logger.info(f"Interrupting {self._state.execution_state.execution.task.task_type} task") 

371 self._task_asyncio.cancel("Interrupted") 

372 

373 with contextlib.suppress(asyncio.CancelledError): 

374 await call.progress_from_task( 

375 self._task_asyncio, 

376 cadence=6.0, 

377 ttl=10.0, 

378 ) 

379 

380 # Execute the task. 

381 start_time = datetime.now(UTC) 

382 

383 with self.enter_context(): 

384 logger.info(f"Executing {task.task_type} task") 

385 aio_task = asyncio.create_task(self._execute_task(execution)) 

386 

387 self._task_asyncio = aio_task 

388 

389 await call.progress_from_task( 

390 self._task_asyncio, 

391 cadence=6.0, 

392 ttl=10.0, 

393 ) 

394 

395 end_time = datetime.now(UTC) 

396 except asyncio.CancelledError: 

397 with self.enter_context(): 

398 logger.warning(f"Execution of {task.task_type} task cancelled") 

399 

400 try: 

401 await call.fail() 

402 except Exception as e: 

403 logger.warning(f"Error logging cancelled task event ({type(e).__name__})") 

404 

405 raise 

406 except Exception as e: 

407 with self.enter_context(): 

408 logger.error(f"Error executing {task.task_type} task") 

409 logger.opt(exception=e).debug(f"{task.task_type} ({execution.task_id}) failed") 

410 

411 try: 

412 await call.fail() 

413 except Exception: 

414 logger.warning(f"Error logging failed task event ({type(e).__name__})") 

415 else: 

416 with self.enter_context(): 

417 logger.info(f"Finished {task.task_type} task") 

418 

419 await call.succeed( 

420 result=TaskExecutionResult( 

421 task_id=execution.task_id, 

422 start_time=start_time, 

423 end_time=end_time, 

424 ) 

425 ) 

426 finally: 

427 # Clear the task object only if it wasn't overwritten in the meantime. 

428 if self._task_asyncio is aio_task: 

429 self._task_asyncio = None 

430 

431 if not aio_task.done(): 

432 aio_task.cancel() 

433 

434 with contextlib.suppress(asyncio.CancelledError, Exception): 

435 await aio_task 

436 

437 async def _execute_task(self, execution: TaskExecution): 

438 # Associate the envelope with the task so the handler can reach it via `task.execution`. 

439 task = execution.task 

440 task.associate_execution(execution) 

441 logger.debug(f"execution begun {execution=}") 

442 finish_info = TaskFinishInfo(aborted=True) 

443 

444 # Emit the execution start event, seeding the executing state's context from the 

445 # incoming execution context. This becomes the base layer that `update_context` 

446 # merges device snapshots onto and that is ultimately sent to downstream cameras. 

447 await self._state.update( 

448 self, 

449 TaskExecutionState(executing=True, execution=execution, context=execution.context), 

450 self._state.operating_state.derive(target=task.target_state()), 

451 ) 

452 

453 try: 

454 # Find the user-supplied task handler callback for this task type. 

455 handler = self._task_handlers.get(type(task)) 

456 assert handler 

457 

458 # Run the task handler. 

459 await handler(task) 

460 except Exception as e: 

461 finish_info = TaskFinishInfo(error=str(e)) 

462 raise 

463 else: 

464 finish_info = TaskFinishInfo() 

465 finally: 

466 try: 

467 await self._state.update( 

468 self, 

469 TaskExecutionState(finished=finish_info, execution=execution), 

470 self._state.operating_state.derive( 

471 current=ts if (ts := task.target_state()) is not None else ControllerOperatingState.NO_CHANGE, 

472 target=None, 

473 ), 

474 ) 

475 except Exception as e: 

476 logger.warning(f"Error logging final task execution state ({type(e).__name__})") 

477 

478 @override 

479 def task_handler[T: Task](self, task_model: type[T]): 

480 def decorator(func: TaskHandlerCallback[T]): 

481 if task_model in self._task_handlers: 

482 raise RuntimeError("multiple handlers for task type") 

483 

484 self._task_handlers[task_model] = func 

485 return func 

486 

487 return decorator 

488 

489 @override 

490 def entity_info(self) -> EntityInfo: 

491 return EntityInfo( 

492 entity_type="controller", 

493 details=ControllerDetails( 

494 supported_tasks=[t.__name__ for t in self._task_handlers.keys()], 

495 controlled_devices=list(self._devices.keys()), 

496 ), 

497 ) 

498 

499 @override 

500 async def set_internal_state(self, state: InternalControllerState): 

501 await self._state.update(self, self._state.operating_state.derive(current=state))