Coverage for core / src / sensorkit / core / impl / program.py: 85%
190 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-02 00:03 +0000
1# SPDX-License-Identifier: Apache-2.0
2from __future__ import annotations
4import asyncio
5import contextlib
6from _contextvars import ContextVar
7from datetime import datetime
8from typing import Any, Callable, ClassVar, override
10from intervaltree import IntervalTree
11from loguru import logger
13from sensorkit.backend.base import Entity, KVError
14from sensorkit.backend.event import Event
15from sensorkit.backend.request import CallContext
16from sensorkit.core.entity import EntityInfo
17from sensorkit.core.executor import TaskFactoryFunc, TaskingLoop
18from sensorkit.core.impl.entity import EntityImpl
19from sensorkit.core.program import (
20 OfferInterval,
21 ProgramActiveState,
22 ProgramActiveStateRequest,
23 ProgramEnableState,
24 ProgramEnableStateRequest,
25 ProgramInterface,
26 ProgramOffering,
27 ProgramState,
28 ProgramTaskingState,
29 set_active_state_request,
30 set_enable_state_request,
31)
32from sensorkit.core.task import TaskContextMap, TaskExecution
35class ProgramOffers:
36 """Manages an IntervalTree of offer windows and signals updates to waiters."""
38 def __init__(self):
39 self._tree = IntervalTree()
40 self._updated = asyncio.Event()
42 def wait(self):
43 """Return an awaitable that resolves when the offer windows have been updated."""
44 return self._updated.wait()
46 def poll(self):
47 """Return True and clear the update flag if a new update is pending, else False."""
48 if self._updated.is_set():
49 self._updated.clear()
50 return True
52 return False
54 def get_offer_windows(self) -> list[OfferInterval]:
55 """Return a sorted list of merged offer windows from the current interval tree."""
56 # Create a tree that *references* the offer tree data and merge overlapping intervals. This
57 # avoids an extra copy, and we know this is safe because `merge_overlaps` builds a new set
58 # of intervals and does not modify `all_intervals`.
59 assert_no_change = self._tree.all_intervals.copy()
60 ref = IntervalTree(self._tree.all_intervals)
61 ref.merge_overlaps()
62 assert self._tree.all_intervals == assert_no_change
63 return sorted(ref)
65 def add(self, start: datetime, end: datetime, obj: Any = None):
66 """Add an offer window to the interval tree and signal waiters."""
67 self._tree.addi(
68 begin=start,
69 end=end,
70 data=obj,
71 )
72 self._updated.set()
74 def remove(self, start: datetime, end: datetime, obj: Any = None):
75 """Remove an offer window from the interval tree, logging a warning if it does not exist."""
76 try:
77 self._tree.removei(
78 begin=start,
79 end=end,
80 data=obj,
81 )
82 self._updated.set()
83 except ValueError:
84 logger.warning(f"Nonexistent offer window removed: {start} -> {end} {obj=}")
86 def clear(self):
87 if not self._tree.is_empty():
88 self._tree.clear()
89 self._updated.set()
92class ProgramImpl(EntityImpl, ProgramInterface):
93 """Helper for implementing server-side functionality of a Program."""
95 current: ClassVar[ContextVar[ProgramImpl | None]] = ContextVar("current_program", default=None)
97 def __init__(self, **kwargs):
98 super().__init__(**kwargs)
100 # Store the user-supplied intervals that are to make up the offered operating windows.
101 self._offers = ProgramOffers()
103 self._enable_hooks: list[Callable[[], None]] = []
104 self._disable_hooks: list[Callable[[], None]] = []
105 self._state = ProgramState(
106 enable_state=ProgramEnableState(enabled=False),
107 active_state=ProgramActiveState(active=False, origin="init"),
108 tasking_state=ProgramTaskingState(),
109 )
110 self._state_lock = asyncio.Lock()
111 self._task_factory: TaskFactoryFunc | None = None
112 self._task_loop: TaskingLoop | None = None
114 @override
115 async def init_impl(self):
116 try:
117 # Restore state.
118 async with self._state_lock:
119 self._state = await self.kv_get_model(ProgramState)
120 logger.debug(f"restored program state: {self._state}")
121 except KVError:
122 logger.debug("initializing program state")
123 await self.kv_put_model(self._state)
125 if self._state.active_state.active:
126 # Our last known state was tasking. We aren't going to automatically start tasking,
127 # but we can emit an event indicating the previous tasking state stopped, giving
128 # observers an opportunity to react themselves.
129 logger.debug("ending lingering tasking active state")
130 await self._update_state(
131 "active_state",
132 ProgramActiveState(active=False, origin="init"),
133 )
135 @override
136 async def attach_impl(self):
137 if self._state.enable_state.enabled:
138 await self._call_with_context(self._enable_hooks)
140 await self.handle_request(set_enable_state_request, self._set_enable_state)
141 await self.handle_request(set_active_state_request, self._set_active_state)
143 @override
144 def on_enable(self, func: Callable[[], None]):
145 self._enable_hooks.append(func)
146 return func
148 @override
149 def on_disable(self, func: Callable[[], None]):
150 self._disable_hooks.append(func)
151 return func
153 @override
154 def task_factory(self, func: TaskFactoryFunc):
155 self._task_factory = func
156 return func
158 async def _start_loop(self, contexts: TaskContextMap):
159 if self._task_loop is not None:
160 return
162 logger.debug(f"starting tasking loop with {contexts=}")
164 async def _on_task_change(execution: TaskExecution | None):
165 # Publish the program's current tasking state as the loop starts/finishes each task.
166 await self._update_state(
167 "tasking_state", ProgramTaskingState(executing_task=execution)
168 )
170 # Create the tasking loop.
171 self._task_loop = TaskingLoop(
172 controller=self.sensorkit().controller(
173 Entity.at(self._state.enable_state.controller)
174 ),
175 factory_func=self._task_factory,
176 contexts=contexts,
177 task_group=self.task_group,
178 on_task_change=_on_task_change,
179 )
181 # Start a background task to make sure the end states are properly handled whether the
182 # loop ends by request or by error. We take pains to do this before starting the loop
183 # itself to eliminate data races and ensure event ordering.
184 queue = asyncio.Queue()
186 async def _finalize_loop():
187 aio_task = await queue.get()
189 try:
190 await aio_task
191 finally:
192 # Clear any lingering tasking state: the loop may have exited mid-task (error or
193 # abort), where clearing is intentionally deferred to teardown here.
194 if self._state.tasking_state.executing_task is not None:
195 await self._update_state("tasking_state", ProgramTaskingState())
197 await self._update_state(
198 "active_state",
199 ProgramActiveState(active=False, origin="request")
200 if self._task_loop.stop_requested
201 else ProgramActiveState(active=False, origin="error"),
202 )
203 logger.info(f"Tasking loop exiting by {self._state.active_state.origin}")
204 self._task_loop = None
205 self._task_finalizer = None
207 self._task_finalizer = self.task_group.create_task(_finalize_loop())
209 # Finally, we can start the tasking loop and update our state.
210 with self.enter_context():
211 logger.info("Starting tasking loop")
212 aio_task = self._task_loop.start()
214 await self._update_state(
215 "active_state",
216 ProgramActiveState(active=True, origin="request"),
217 )
219 # Tell the finalizer about the loop task.
220 queue.put_nowait(aio_task)
222 async def _stop_loop(self, *, timeout=None):
223 if self._task_loop is None:
224 return
226 await self._update_state(
227 "active_state",
228 ProgramActiveState(active=True, origin=self._state.active_state.origin, stopping=True),
229 )
231 # Due to the await above, we have to check again whether the loop is still running, as it
232 # may have been concurrently stopped.
233 if self._task_loop is None:
234 return
236 try:
237 await self._task_loop.stop(timeout=timeout)
238 except (Exception, asyncio.CancelledError) as e:
239 logger.warning(f"Error while stopping tasking loop: {e}")
241 # Make sure the end state has been updated by the finalizer.
242 if self._task_finalizer is not None:
243 with contextlib.suppress(asyncio.CancelledError):
244 await self._task_finalizer
246 async def _set_enable_state(self, request: ProgramEnableStateRequest):
247 if not request.enable and not self._state.enable_state.enabled:
248 return
250 enablement_changed = request.enable ^ self._state.enable_state.enabled
251 controller_changed = request.controller != self._state.enable_state.controller
253 if not enablement_changed and not controller_changed:
254 return
256 logger.info(
257 f"{'Enabling' if request.enable else 'Disabling'} "
258 f"{self.entity} for target {request.controller}"
259 )
261 await self._update_state(
262 "enable_state",
263 ProgramEnableState(enabled=request.enable, controller=request.controller),
264 )
266 if request.enable:
267 # Call the enable hook.
268 await self._call_with_context(self._enable_hooks)
269 else:
270 # Disabling the program implies setting the active state low. We consider this an abort
271 # case; if the caller wants a graceful stop, they can explicitly stop tasking first.
272 await self._stop_loop(timeout=0)
274 # Call the disable hook.
275 await self._call_with_context(self._disable_hooks)
277 async def _set_active_state(
278 self,
279 request: ProgramActiveStateRequest,
280 call: CallContext[None, None],
281 ):
282 with self.enter_context():
283 logger.info(f"Requested to {request.action} the tasking loop")
285 if request.action == "start":
286 # Cannot activate tasking if we aren't enabled with a target Controller configured.
287 if not self._state.enable_state.enabled or not self._state.enable_state.controller:
288 call.reject(response=None)
289 return
291 if (
292 self._state.active_state.active
293 and request.contexts != self._state.active_state.contexts
294 ):
295 # Don't support restart-with-different-context in a single request.
296 call.reject(response=None)
297 return
299 call.accept(response=None)
301 if request.active_state() == self._state.active_state.active:
302 # We are already in the requested state, so claim success.
303 await call.succeed(result=None)
304 return
306 try:
307 match request.action:
308 case "start":
309 # Start a new tasking loop.
310 with self.enter_context():
311 await self._start_loop(request.contexts)
312 case "stop":
313 # Stop the tasking loop gracefully.
314 # TODO: add backstop timeout reflecting maximum task time
315 await call.progress_from_task(
316 self.task_group.create_task(self._stop_loop()),
317 cadence=6,
318 ttl=10,
319 )
320 case "abort":
321 # Abort the tasking loop immediately.
322 with self.enter_context():
323 await self._stop_loop(timeout=0)
324 except Exception as e:
325 with self.enter_context():
326 logger.exception("Error setting active state")
328 await call.fail(f"{type(e).__name__} setting active state ({e})")
329 else:
330 await call.succeed(result=None)
332 async def _update_state(self, key: str, event: Event):
333 async with self._state_lock:
334 setattr(self._state, key, event)
335 await self.emit_event(event)
336 await self.kv_put_model(self._state)
338 @override
339 def get_offers(self):
340 return self._offers.get_offer_windows()
342 @override
343 async def publish_offers(self):
344 """Publish the current set of offers if they have changed."""
345 if self._offers.poll():
346 await self.publish(
347 ProgramOffering(offer_windows=self._offers.get_offer_windows())
348 )
350 @override
351 def add_offer(self, start: datetime, end: datetime, obj: Any = None):
352 self._offers.add(start, end, obj)
354 @override
355 def remove_offer(self, start: datetime, end: datetime, obj: Any = None):
356 self._offers.remove(start, end, obj)
358 @override
359 def clear_offers(self):
360 self._offers.clear()
362 @override
363 def entity_info(self) -> EntityInfo:
364 return EntityInfo(entity_type="program", details=None)