Coverage for core / src / sensorkit / core / program.py: 94%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import collections 

6import contextlib 

7import functools 

8from abc import abstractmethod 

9from datetime import datetime 

10from typing import TYPE_CHECKING, Any, Callable, Literal 

11 

12from intervaltree import Interval 

13from loguru import logger 

14from pydantic import BaseModel, Field, model_validator 

15 

16from sensorkit.backend.event import Event 

17from sensorkit.backend.request import ExtendedResponse, Request 

18from sensorkit.common.aio import AsyncObserver 

19from sensorkit.common.keyword import declare_keyword 

20from sensorkit.core.entity import EntityClient, EntityInterface, EntityRef 

21from sensorkit.core.executor import TaskFactoryFunc 

22from sensorkit.core.task import TaskContextMap, TaskExecution 

23 

24if TYPE_CHECKING: 

25 from sensorkit.core.client import SensorKit 

26 

27 

28class OfferInterval(Interval): 

29 """An Interval subclass with typed begin/end datetimes used to represent a program offer window.""" 

30 

31 begin: datetime 

32 end: datetime 

33 data: Any = None 

34 

35 

36@declare_keyword 

37class ProgramOffering(BaseModel): 

38 """Keyword publishing the set of offer windows during which this program can task a controller.""" 

39 

40 offer_windows: list[OfferInterval] 

41 

42 

43class ProgramEnableState(Event): 

44 """Event indicating the Program enable state has changed.""" 

45 enabled: bool = False 

46 controller: str | None = None 

47 

48 

49class ProgramActiveState(Event): 

50 """Event indicating the Program active state has changed.""" 

51 active: bool 

52 origin: Literal["request", "init", "error"] 

53 stopping: bool = False 

54 contexts: TaskContextMap = Field(default_factory=TaskContextMap) 

55 

56 

57class ProgramTaskingState(Event): 

58 """Event indicating the Task a Program is currently executing, if any. 

59 

60 Carries the full `TaskExecution` envelope (not just an id) so observers see the same task shape 

61 the controller publishes on its `TaskExecutionState`, including the controller-minted 

62 `task_id`. `None` while the program is between tasks or idle. 

63 """ 

64 executing_task: TaskExecution | None = None 

65 

66 @model_validator(mode="before") 

67 @classmethod 

68 def _migrate_legacy_executing_task(cls, data: Any) -> Any: 

69 """Upgrade state persisted before the executing task carried a full `TaskExecution`. 

70 

71 The pre-split `ProgramTaskingStatus` recorded only `executing_task_id` — a bare task id. 

72 A bare id cannot be rehydrated into a `TaskExecution` envelope, so a legacy id collapses to 

73 `executing_task=None` (equivalent to "idle"); observers of the old shape only ever saw the 

74 id, which the program re-publishes the moment it next dispatches a task. 

75 """ 

76 if not isinstance(data, dict) or "executing_task" in data: 

77 return data 

78 

79 if "executing_task_id" in data: 

80 data = {k: v for k, v in data.items() if k != "executing_task_id"} 

81 data["executing_task"] = None 

82 

83 return data 

84 

85 

86class ProgramState(BaseModel): 

87 """Internal state snapshot of a Program.""" 

88 enable_state: ProgramEnableState 

89 active_state: ProgramActiveState 

90 tasking_state: ProgramTaskingState 

91 

92 @model_validator(mode="before") 

93 @classmethod 

94 def _migrate_legacy_tasking_state(cls, data: Any) -> Any: 

95 """Upgrade KV snapshots persisted before `tasking_status` was renamed to `tasking_state`. 

96 

97 Versions predating the Task/TaskExecution split stored the program's executing task under a 

98 `tasking_status` key (a `ProgramTaskingStatus` carrying `executing_task_id`); this 

99 version expects `tasking_state`. Renaming the key here lets `ProgramTaskingState`'s own 

100 `before` validator finish upgrading the inner shape, so a program starting against a NATS 

101 broker holding older state no longer raises a `ValidationError`. A snapshot missing both 

102 keys entirely is treated as idle. 

103 """ 

104 if not isinstance(data, dict) or "tasking_state" in data: 

105 return data 

106 

107 data = dict(data) 

108 data["tasking_state"] = data.pop("tasking_status", {}) 

109 return data 

110 

111 

112class ProgramEnableStateRequest(BaseModel): 

113 """Request that a Program enable or disable task sourcing. 

114 

115 The `controller` field specifies the name of the target Controller to task. It must be 

116 specified before activating the Program, and some Program implementations may require that it 

117 be specified whenever `enable` is True. 

118 """ 

119 enable: bool 

120 controller: str | None 

121 

122 

123set_enable_state_request = Request.define( 

124 "set_enable_state", 

125 payload=ProgramEnableStateRequest, 

126) 

127"""Control the enable state of a Program.""" 

128 

129 

130class ProgramActiveStateRequest(BaseModel): 

131 """Request that a Program activate or deactivate tasking of the target Controller.""" 

132 action: Literal["start", "stop", "abort"] 

133 contexts: TaskContextMap = Field(default_factory=TaskContextMap) 

134 

135 def active_state(self): 

136 """Return True if the requested action is 'start', False otherwise.""" 

137 return True if self.action == "start" else False 

138 

139 

140class ProgramActiveStateResult(BaseModel): 

141 """Result of a request to change the Program active state.""" 

142 success: bool 

143 

144 

145set_active_state_request = Request.define( 

146 "set_active_state", 

147 payload=ProgramActiveStateRequest, 

148 response=ExtendedResponse, 

149) 

150"""Control the active state of a Program.""" 

151 

152 

153class ProgramClient(EntityClient): 

154 """Object that exposes client-side functionality of a Program.""" 

155 

156 async def enable(self, target_controller: str): 

157 """Request that the Program enable task sourcing for the target Controller.""" 

158 return await self.call( 

159 set_enable_state_request, 

160 ProgramEnableStateRequest( 

161 enable=True, 

162 controller=target_controller, 

163 ) 

164 ) 

165 

166 async def disable(self): 

167 """Request that the Program disable task sourcing.""" 

168 return await self.call( 

169 set_enable_state_request, 

170 ProgramEnableStateRequest( 

171 enable=False, 

172 controller=None, 

173 ) 

174 ) 

175 

176 async def start_tasking(self, contexts: TaskContextMap | None = None): 

177 """Request that the Program begin tasking the target Controller.""" 

178 await self.call( 

179 set_active_state_request, 

180 ProgramActiveStateRequest(action="start", contexts=contexts or TaskContextMap()), 

181 ) 

182 

183 async def stop_tasking(self): 

184 """Request that the Program stop tasking and wait for any in flight task to complete.""" 

185 await self.call( 

186 set_active_state_request, 

187 ProgramActiveStateRequest(action="stop"), 

188 ) 

189 

190 async def abort_tasking(self): 

191 """Request that the Program stop tasking and abort any in flight task immediately.""" 

192 await self.call( 

193 set_active_state_request, 

194 ProgramActiveStateRequest(action="abort"), 

195 ) 

196 

197 async def wait_until_tasking_stops(self): 

198 """ 

199 Wait until the Program tasking loop stops. 

200 

201 Returns the corresponding ProgramActiveState event, if any. 

202 """ 

203 # TODO: This needs to detect if the hosting service goes down and raise if it doesn't 

204 # come back and publish new state within a grace period. 

205 stream = await self.tasking_change_events() 

206 state = await self.kv_get_model(ProgramState) 

207 

208 if state.active_state.active: 

209 async for event in stream: 

210 if not event.active: 

211 return event 

212 

213 return None 

214 

215 def tasking_change_events(self): 

216 """Return an async generator that yields ProgramActiveState events as tasking starts or stops.""" 

217 return self.monitor_event(ProgramActiveState) 

218 

219 @functools.cache 

220 def monitor_enable_state(self): 

221 """Return (creating if needed) a ProgramStateMonitor tracking this program's enable state.""" 

222 monitor = ProgramStateMonitor(self) 

223 monitor.start() 

224 return monitor 

225 

226 

227class ProgramRef(EntityRef[ProgramClient]): 

228 """A serializable reference to a program client.""" 

229 

230 def _get_client(self, kit: SensorKit) -> ProgramClient: 

231 return kit.program(self.name) 

232 

233 

234class ControllerOffers: 

235 """Monitor published offers for all Programs associated with a Controller.""" 

236 

237 def __init__(self, controller: str, program_discovery: ProgramDiscovery): 

238 self.controller = controller 

239 self.program_discovery = program_discovery 

240 self._monitors: dict[str, asyncio.Task] = {} 

241 self._offers: dict[str, list[OfferInterval]] = {} 

242 

243 async def start(self, *, task_group: Any = asyncio): 

244 """Start the program-offer discovery loop, waiting until the first update is received.""" 

245 self._task_group = task_group 

246 

247 # Begin discovering controller offers, making sure to wait for the first set of updates. 

248 ready = asyncio.Event() 

249 self._discover_task = task_group.create_task(self._discover_programs(ready)) 

250 await ready.wait() 

251 

252 async def _discover_programs(self, initial_update: asyncio.Event): 

253 prev = set() 

254 

255 async for programs in self.program_discovery.controller_programs(self.controller): 

256 events = [] 

257 

258 for remove in prev.difference(programs): 

259 self._monitors[remove].cancel() 

260 del self._monitors[remove] 

261 

262 for add in programs.difference(prev): 

263 assert add not in self._monitors 

264 events.append(asyncio.Event()) 

265 self._monitors[add] = self._task_group.create_task( 

266 self._monitor_program(add, events[-1]) 

267 ) 

268 

269 await asyncio.gather(*(event.wait() for event in events), return_exceptions=True) 

270 initial_update.set() 

271 prev = programs 

272 

273 async def _monitor_program(self, program: str, initial_update: asyncio.Event): 

274 logger.debug(f"starting offer monitor for program {program}") 

275 self._offers[program] = [] 

276 

277 try: 

278 # Suppress propagation of cancellation to the parent task group. 

279 with contextlib.suppress(asyncio.CancelledError): 

280 # Get a stream of offers for the target Program. 

281 client = self.program_discovery.client.program(program) 

282 stream = await client.monitor(ProgramOffering) 

283 

284 initial_update.set() 

285 

286 async for _, offering in stream: 

287 self._offers[program] = sorted(offering.offer_windows) 

288 finally: 

289 initial_update.set() 

290 del self._offers[program] 

291 logger.debug(f"cancelled offer monitor for program {program}") 

292 

293 def get_offers(self): 

294 """Return a snapshot of the current offer windows keyed by program name. 

295 

296 Each program's offer windows are sorted in ascending order. 

297 """ 

298 return self._offers.copy() 

299 

300 

301class ProgramDiscovery: 

302 """Discover Programs and their associated Controllers.""" 

303 

304 def __init__(self): 

305 self._enabled_programs: AsyncObserver[frozenset[str]] = AsyncObserver(frozenset()) 

306 self._known_programs: AsyncObserver[frozenset[str]] = AsyncObserver(frozenset()) 

307 self._controllers: dict[str, AsyncObserver[frozenset[str]]] = collections.defaultdict( 

308 lambda: AsyncObserver(frozenset()) 

309 ) 

310 self._monitor_tasks: dict[str, asyncio.Task] = {} 

311 self._discover_task: asyncio.Task | None = None 

312 

313 async def start(self, client: SensorKit): 

314 """Begin monitoring the backend for program enable-state changes, waiting for the first update.""" 

315 self.client = client 

316 

317 # Begin discovering enabled programs, making sure to wait for the first set of updates. 

318 ready = asyncio.Event() 

319 self._discover_task = asyncio.create_task(self._discover_programs(ready)) 

320 await ready.wait() 

321 

322 async def stop(self): 

323 """Cancel discovery and every per-program enable-state monitor it started.""" 

324 if self._discover_task is None: 

325 return 

326 

327 self._discover_task.cancel() 

328 

329 for task in self._monitor_tasks.values(): 

330 task.cancel() 

331 

332 await asyncio.gather( 

333 self._discover_task, *self._monitor_tasks.values(), return_exceptions=True 

334 ) 

335 

336 self._monitor_tasks.clear() 

337 self._discover_task = None 

338 

339 async def _discover_programs(self, initial_update: asyncio.Event): 

340 tasks = self._monitor_tasks 

341 kv = self.client.backend.key_value() 

342 

343 # Programs present in the initial snapshot. We're "ready" once each of them 

344 # has been discovered and reported its initial enable state — or immediately 

345 # if there are none. 

346 pending = { 

347 str(entry.key.entity()) 

348 for entry in await kv.get_all(deep=True) 

349 if entry.key.prop == "ProgramState" 

350 } 

351 if not pending: 

352 initial_update.set() 

353 

354 # Monitor all registered Programs. 

355 # FIXME: Need an alternate way to do this without firehose. Cannot assume the backend 

356 # allows prefix wildcards to capture '*.ProgramState' (and NATS, in fact, does not). 

357 # This most likely means we need to define a dedicated KV prefix for system info, 

358 # where we can define keys to explicitly indicate existence, like EntityLease. 

359 monitor = await kv.monitor_all(deep=True) 

360 

361 async for entry in monitor: 

362 if entry.key.prop != "ProgramState": 

363 continue 

364 

365 program = str(entry.key.entity()) 

366 client = self.client.program(program) 

367 events = [] 

368 

369 if entry.deleted(): 

370 # Remove the monitor task associated with the offline program, if any. 

371 if program in tasks: 

372 task = tasks.pop(program) 

373 task.cancel() 

374 self._known_programs.notify(frozenset(tasks.keys())) 

375 elif program not in tasks: 

376 # Start a new enable state monitor for the newly discovered program. 

377 events.append(asyncio.Event()) 

378 tasks[program] = asyncio.create_task(self._monitor_task(client, events[-1])) 

379 self._known_programs.notify(frozenset(tasks.keys())) 

380 

381 # Wait for the initial enablement status of each discovered program. 

382 await asyncio.gather(*(event.wait() for event in events), return_exceptions=True) 

383 

384 # Signal ready once every program from the initial snapshot has been seen. 

385 pending.discard(program) 

386 if not pending: 

387 initial_update.set() 

388 

389 async def _monitor_task(self, client: ProgramClient, initial_update: asyncio.Event): 

390 monitor = client.monitor_enable_state() 

391 program = str(client.entity) 

392 

393 try: 

394 async for state in monitor.observe(): 

395 observers = [self._enabled_programs] 

396 logger.debug(f"discovery for {program} observed {state=}") 

397 

398 if state: 

399 observers.append(self._controllers[state.controller]) 

400 

401 for observer in observers: 

402 if state and state.enabled: 

403 new_set = observer.value.union({program}) 

404 else: 

405 new_set = observer.value.difference({program}) 

406 

407 observer.notify(new_set) 

408 

409 initial_update.set() 

410 finally: 

411 monitor.stop() 

412 

413 def enabled_programs(self): 

414 """Return an async generator yielding the current set of enabled program names on each change.""" 

415 return self._enabled_programs.consume(initial_value=True) 

416 

417 def known_programs(self): 

418 """Return an async generator yielding the current set of all known program names on each change.""" 

419 return self._known_programs.consume(initial_value=True) 

420 

421 def controller_programs(self, controller: str): 

422 """Return an async generator yielding the set of programs targeting the given controller on each change.""" 

423 return self._controllers[controller].consume(initial_value=True) 

424 

425 

426class ProgramStateMonitor: 

427 """Monitor a Program's effective enable-state, taking into account the entity liveness.""" 

428 

429 def __init__(self, client: ProgramClient): 

430 self.client = client 

431 self.online = False 

432 self.enable_state: ProgramEnableState | None = None 

433 self._observer: AsyncObserver[ProgramEnableState | None] = AsyncObserver(None) 

434 

435 def start(self): 

436 """Start background tasks that monitor program liveness and enable-state.""" 

437 self._tasks = ( 

438 asyncio.create_task(self._monitor_liveness()), 

439 asyncio.create_task(self._monitor_state()), 

440 ) 

441 

442 def stop(self): 

443 """Cancel the background monitoring tasks.""" 

444 for t in self._tasks: 

445 t.cancel() 

446 

447 def observe(self): 

448 """Return an async generator yielding the current ProgramEnableState (or None if offline) on each change.""" 

449 return self._observer.consume(initial_value=True) 

450 

451 def _notify(self): 

452 if not self.online: 

453 self._observer.notify(None) 

454 else: 

455 self._observer.notify(self.enable_state) 

456 

457 async def _monitor_liveness(self): 

458 async for is_online in self.client.observe_online_state(): 

459 self.online = is_online 

460 self._notify() 

461 

462 async def _monitor_state(self): 

463 program = str(self.client.entity) 

464 

465 try: 

466 stream = await self.client.monitor_event(ProgramEnableState) 

467 initial_state = await self.client.kv_get_model(ProgramState) 

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

469 logger.debug(f"Failed to monitor program {program}: {e}") 

470 raise 

471 

472 logger.debug(f"monitoring discovered program {program}") 

473 self.enable_state = initial_state.enable_state 

474 self._notify() 

475 

476 async for event in stream: 

477 if event.timestamp() > self.enable_state.timestamp(): 

478 self.enable_state = event 

479 self._notify() 

480 

481 

482class ProgramInterface(EntityInterface): 

483 """Interface describing an implementation of a program.""" 

484 

485 @abstractmethod 

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

487 """Register a callback to invoke when the program is enabled.""" 

488 ... 

489 

490 @abstractmethod 

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

492 """Register a callback to invoke when the program is disabled.""" 

493 ... 

494 

495 @abstractmethod 

496 def task_factory(self, func: TaskFactoryFunc) -> TaskFactoryFunc: 

497 """Register the task factory function that the program calls to generate tasks.""" 

498 ... 

499 

500 @abstractmethod 

501 def get_offers(self) -> list[Any]: 

502 """Return the current list of offer windows.""" 

503 ... 

504 

505 @abstractmethod 

506 async def publish_offers(self): 

507 """Publish the current offer windows to the entity's keyword stream.""" 

508 ... 

509 

510 @abstractmethod 

511 def add_offer(self, start: datetime, end: datetime, obj: Any = None): 

512 """Add an offer window covering the given time range.""" 

513 ... 

514 

515 @abstractmethod 

516 def remove_offer(self, start: datetime, end: datetime, obj: Any = None): 

517 """Remove the offer window covering the given time range.""" 

518 ... 

519 

520 @abstractmethod 

521 def clear_offers(self): 

522 """Remove all offer windows.""" 

523 ...