Coverage for core / src / sensorkit / auto / lifecycle.py: 75%
365 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
6import uuid
7from abc import ABC, abstractmethod
8from dataclasses import dataclass, field
9from enum import StrEnum, auto
10from typing import Any, Literal, final, overload, override
12from loguru import logger
14from sensorkit.common.aio import AsyncValueLatch
15from sensorkit.common.keyword import KeywordDict
16from sensorkit.core.controller import ControllerClient, ControllerState, InternalControllerState
17from sensorkit.core.program import ProgramClient
18from sensorkit.core.task import (
19 InitTask,
20 RecoverTask,
21 ShutdownTask,
22 StandbyTask,
23 Task,
24 TaskContextMap,
25)
28class LifecycleStep(StrEnum):
29 """Steps in the Controller lifecycle logic."""
30 INIT = auto()
31 TASKING = auto()
32 STANDBY = auto()
33 SHUTDOWN = auto()
34 RECOVER = auto()
35 WAIT = auto()
38@dataclass(slots=True, frozen=True)
39class DemandState:
40 """The demanded InternalState and associated context."""
41 state: InternalControllerState | None
42 contexts: TaskContextMap | None = None
43 program: ProgramClient | None = None
45 # Other fields must have `compare=False` to ensure correct demand state difference checks.
46 interrupt: bool = field(default=False, compare=False)
49# TODO: Move to sensorkit.core.controller
50class ControllerStateMonitor:
51 """Monitors a Controller's error state."""
53 def __init__(self, controller: ControllerClient):
54 self.controller = controller
56 async def start(self):
57 """Begin monitoring the controller state and wait until the initial state is known."""
58 ready = asyncio.get_running_loop().create_future()
59 self._task = asyncio.create_task(self._monitor(ready))
60 await ready
62 async def _monitor(self, initial_update: asyncio.Future):
63 try:
64 # TODO: For now just query the state once, since we're only using it for initial state.
65 state = await self.controller.kv_get_model(ControllerState)
66 self.current_state = state.operating_state.current
67 initial_update.set_result(True)
68 except asyncio.CancelledError:
69 initial_update.cancel()
70 raise
71 except BaseException as e:
72 initial_update.set_exception(e)
73 raise
75 def in_error(self):
76 """Return True if the controller is currently in an error state (stub)."""
77 # TODO
78 return False
80 async def wait(self):
81 """Wait until a passive error is detected (stub, waits forever)."""
82 # TODO
83 await asyncio.sleep(float("inf"))
86class DemandProcError(Exception):
87 """Raised by a DemandProc to signal a lifecycle error with a categorized origin."""
89 type ErrorKind = Literal["program", "controller", "internal"]
91 def __init__(self, *args, kind: ErrorKind = "internal", **kwargs):
92 super().__init__(*args, **kwargs)
93 self.kind = kind
96class DemandProc(ABC):
97 """Wraps an asyncio.Task and implements a component of the Controller lifecycle logic."""
99 def __init__(
100 self,
101 lifecycle: ControllerLifecycle,
102 demand: DemandState,
103 ):
104 self.lifecycle = lifecycle
105 self.demand = demand
106 self._interrupt_count = 0
107 self._interrupt_level = 0
108 self._cleanup_tasks: set[uuid.UUID] = set()
109 self._cleanup_futures: set[asyncio.Future] = set()
111 @final
112 def start(self):
113 """Create and schedule the underlying asyncio task for this demand procedure."""
114 self._aio_task = asyncio.create_task(self._demand_proc())
116 async def _demand_proc(self):
117 logger.debug(f"{self.demand.state} procedure starting")
118 prereq_coro = asyncio.sleep(0)
119 logic_coro = self.state_logic()
121 while True:
122 try:
123 # Await prerequisite actions. Initially this is nothing, but after an interrupt
124 # this is operations to effect the stop or abort.
125 await prereq_coro
127 # Run the user logic. Initially this is the main state driving logic, but after
128 # an interrupt this is stop or abort logic.
129 await logic_coro
130 except asyncio.CancelledError:
131 if next_stage := self._demand_proc_cancelled():
132 # Managed cancellation. Loop to the next stage.
133 prereq_coro, logic_coro = next_stage
134 else:
135 logger.debug(f"{self.demand.state} procedure exiting due to hard cancel")
136 raise
137 except BaseException as e:
138 logger.debug(f"{self.demand.state} procedure exiting due to error ({e})")
139 raise
140 else:
141 logger.warning(f"{self.demand.state} procedure exiting without raising")
142 break
144 def _demand_proc_cancelled(self) -> tuple[asyncio.Future, asyncio.Future] | None:
145 if self._aio_task.cancelling() > self._interrupt_count:
146 # This means the task was externally cancelled coincident with a stop or abort!
147 # Fall through to a hard cancel.
148 self._interrupt_level = 3
150 match self._interrupt_level:
151 case 1:
152 logger.debug(f"{self.demand.state} is stopping gracefully")
154 while self._aio_task.cancelling() > 0:
155 self._aio_task.uncancel()
157 # Gracefully stop awaitables and SK tasks.
158 prereq_coro = asyncio.wait_for(
159 asyncio.gather(
160 *(future for future in self._cleanup_futures if not future.done()),
161 *(
162 self.lifecycle.controller.wait_for_task(task_id)
163 for task_id in self._cleanup_tasks
164 ),
165 ),
166 150,
167 )
169 # Run the user stop logic at the next iteration.
170 logic_coro = self.stop_logic()
171 case 2:
172 logger.debug(f"{self.demand.state} procedure is aborting immediately")
174 while self._aio_task.cancelling() > 0:
175 self._aio_task.uncancel()
177 # Forcefully stop awaitables and SK tasks.
178 for fut in self._cleanup_futures:
179 fut.cancel("Demand procedure abort")
181 prereq_coro = asyncio.wait_for(
182 asyncio.gather(
183 *(
184 self.lifecycle.controller.abort_task(task_id)
185 for task_id in self._cleanup_tasks
186 )
187 ),
188 15,
189 )
191 # Run the user abort logic at the next iteration.
192 logic_coro = self.abort_logic()
193 case _:
194 return None
196 self._interrupt_count = 0
197 self._interrupt_level = 0
198 self._cleanup_tasks.clear()
199 self._cleanup_futures.clear()
201 return prereq_coro, logic_coro
203 async def shield_with_cleanup[T](self, fut: asyncio.Future[T]):
204 """Awaits the input Future while shielding it from cancellation and handling stop/abort."""
205 self._cleanup_futures.add(fut)
206 result = await asyncio.shield(fut)
207 self._cleanup_futures.discard(fut)
208 return result
210 def run_with_cleanup(self, task: Task, *, context: KeywordDict | None = None):
211 """Dispatch *task* (interrupting any current task) and await its completion under cleanup.
213 The execution future is registered for graceful/forced cleanup *synchronously*, before any
214 cancellable await, so the dispatch and completion happen inside the shielded region: a
215 demand change that cancels the procedure cannot slip in between dispatch and registration.
216 The controller-minted task id is registered for forced abort as soon as it is known.
218 Args:
219 task: The lifecycle task to execute.
220 context: Optional keyword context to attach to the execution.
222 Returns:
223 An awaitable yielding the final task execution result.
224 """
225 async def _execute():
226 execution = await self.lifecycle.controller.start_task(
227 task, context=context, interrupt=True
228 )
229 self._cleanup_tasks.add(execution.task_id)
231 try:
232 return await execution
233 finally:
234 self._cleanup_tasks.discard(execution.task_id)
236 return self.shield_with_cleanup(asyncio.ensure_future(_execute()))
238 @abstractmethod
239 async def state_logic(self):
240 """Run logic for this demand procedure."""
242 async def stop_logic(self):
243 """Graceful stop logic for this demand procedure."""
245 async def abort_logic(self):
246 """Immediate abort logic for this demand procedure."""
248 @final
249 async def interrupt_no_raise(self, msg: Any | None = None, *, abort: bool = False):
250 """Interrupt the current demand procedure and return any exceptions that occurred."""
251 try:
252 if abort:
253 await self.abort()
254 else:
255 await self.stop()
256 except Exception as e:
257 return e
259 return None
261 @final
262 async def stop(self, msg: Any | None = None):
263 """End the task gracefully and wait for completion."""
264 if self._interrupt_level > 1:
265 return
267 self._interrupt_level = 1
268 self._interrupt_count += 1
269 self._aio_task.cancel(msg)
271 # Wait until the task ends.
272 if not self._aio_task.done():
273 with contextlib.suppress(asyncio.CancelledError):
274 await self._aio_task
276 @final
277 async def abort(self, msg: Any | None = None):
278 """End the task as soon as possible and wait for completion."""
279 if self._interrupt_level > 2:
280 return
282 self._interrupt_level = 2
283 self._interrupt_count += 1
284 self._aio_task.cancel(msg)
286 # Wait until the task ends.
287 if not self._aio_task.done():
288 with contextlib.suppress(asyncio.CancelledError):
289 await self._aio_task
291 @final
292 def cancel(self, msg: Any | None = None):
293 """Hard-cancel the demand procedure without waiting for cleanup."""
294 self._interrupt_level = 3
295 self._aio_task.cancel(msg)
297 @final
298 def done(self):
299 """Return True if the underlying asyncio task has finished."""
300 return self._aio_task.done()
302 @final
303 def error(self):
304 """Return the exception raised by the task, or None if still running or successful."""
305 return self._aio_task.exception() if self._aio_task.done() else None
307 @final
308 def future(self) -> asyncio.Future[None]:
309 """Return the underlying asyncio.Task for use as an awaitable or in `asyncio.wait`."""
310 return self._aio_task
313class ControllerLifecycle:
314 """Implements full lifecycle orchestration of a Controller."""
316 def __init__(self):
317 self.step = LifecycleStep.WAIT
318 # TODO: Belief state needs to be a triple like DemandState in order to properly no-op
319 # init and standby transitions.
320 self.belief_state: InternalControllerState | None = None
321 self._demand = AsyncValueLatch(DemandState(state=None))
322 self._can_operate = asyncio.Event()
324 # Define the mapping of demand states to demand procedure implementations.
325 self._demand_procs = {
326 InternalControllerState.OPERATE: self.OperateProc,
327 InternalControllerState.STANDBY: self.StandbyProc,
328 InternalControllerState.SHUTDOWN: self.ShutdownProc,
329 }
331 def start(
332 self,
333 controller: ControllerClient,
334 *,
335 task_group=asyncio,
336 ):
337 """Bind to *controller* and launch the lifecycle loop in *task_group*."""
338 self.controller = controller
340 # Create a monitor to observe the actual Controller state.
341 self._monitor = ControllerStateMonitor(self.controller)
343 # Start the main task, which will block until enable is called.
344 self._main_task = task_group.create_task(self._lifecycle())
346 async def stop(self):
347 """Cancel the lifecycle loop and wait for it to finish."""
348 self._main_task.cancel()
350 with contextlib.suppress(asyncio.CancelledError):
351 await self._main_task
353 def enable(self):
354 """Allow the lifecycle to issue demand commands to the controller."""
355 self._can_operate.set()
357 def disable(self):
358 """Prevent any further demand commands and interrupt the current demand procedure."""
359 self._can_operate.clear()
361 # Interrupt the current demand procedure.
362 self._demand.update(DemandState(state=None, interrupt=True))
364 @property
365 def enabled(self):
366 """True if the lifecycle is currently allowed to command the controller."""
367 return self._can_operate.is_set()
369 @property
370 def demand_state(self):
371 """The currently applied DemandState."""
372 return self._demand.value
374 @overload
375 def set_demand_state(self, state: DemandState): ...
377 @overload
378 def set_demand_state(
379 self,
380 demand: Literal[InternalControllerState.OPERATE],
381 *,
382 contexts: TaskContextMap = ...,
383 program: ProgramClient,
384 interrupt: bool = ...,
385 ): ...
387 @overload
388 def set_demand_state(
389 self,
390 demand: Literal[InternalControllerState.STANDBY, InternalControllerState.SHUTDOWN],
391 *,
392 contexts: TaskContextMap = ...,
393 interrupt: bool = ...,
394 ): ...
396 def set_demand_state(self, demand: InternalControllerState, **kwargs):
397 """Set the demanded controller state, returning True if the value changed."""
398 if demand not in self._demand_procs:
399 raise ValueError(f"Cannot demand controller state {demand}")
401 return self._demand.update(
402 DemandState(demand, **kwargs),
403 only_if_different=True
404 )
406 async def _lifecycle(self):
407 # Start the controller monitor and get an initial belief state.
408 await self._monitor.start()
409 self.belief_state = self._monitor.current_state
410 logger.debug(f"got initial state from controller: {self.belief_state}")
412 # Loop forever, pausing when the Controller is marked inoperable. Note that cancellation
413 # here is OK because we are guaranteed to not have a demand procedure running.
414 while await self._can_operate.wait():
415 # Get the Controller demand state. If there was a pending change to the demand state,
416 # it is applied here.
417 demand = self._demand.apply()
419 if demand.state is not None:
420 # Delegate to a demand procedure that corresponds to the demand state.
421 demand_proc_type = self._demand_procs[demand.state]
422 demand_proc = demand_proc_type(self, demand)
423 demand_proc.start()
424 else:
425 # If there is no demanded Controller state, we just fall through to wait for a
426 # demand state change or a passive error to occur.
427 demand_proc = None
429 try:
430 # Wait for our procedure to complete or an error to occur.
431 error_kind = await self._lifecycle_reactor(demand_proc)
432 except asyncio.CancelledError:
433 logger.debug(f"{self.controller.entity} lifecycle loop cancelled")
435 if demand_proc:
436 # We were externally cancelled with our demand procedure potentially still
437 # running. Await successful cancellation of the task and re-raise.
438 with contextlib.suppress(Exception):
439 await demand_proc.abort()
441 raise
443 # Sleep to avoid thrashing, but only when there was an error.
444 if error_kind is not None:
445 await asyncio.sleep(5.0)
447 async def _lifecycle_reactor(self, demand_proc: DemandProc | None):
448 # Build a list of futures that capture when:
449 # a) the demand procedure, if there is one, ends
450 # - it should have raised; this is considered an active error
451 # - if it did not raise, this is a bug
452 # b) a demand state change is detected
453 # c) the Controller explicitly reports an error
454 futures: list[asyncio.Future[None]] = [
455 asyncio.create_task(self._demand.wait_until_pending()),
456 asyncio.create_task(self._monitor.wait()),
457 *([demand_proc.future()] if demand_proc else []),
458 ]
460 try:
461 # Wait for any of the above conditions to occur.
462 await asyncio.wait(futures, return_when=asyncio.FIRST_COMPLETED)
464 # Determine whether we are in an error state.
465 active_error: BaseException | None = None
467 if demand_proc:
468 if demand_proc.done():
469 # The demand procedure has ended. This can only happen through cancellation
470 # or by an error occurring, so by contract it must have raised.
471 active_error = demand_proc.error()
472 elif self._demand.pending_change():
473 # The demand procedure is still executing and our demand has changed. Stop
474 # it according to the interrupt policy of the incoming demand.
475 active_error = await demand_proc.interrupt_no_raise(
476 abort=self._demand.pending_value.interrupt
477 )
478 else:
479 # The demand procedure is still executing and a passive error was detected.
480 # In this case we always send a hard abort.
481 assert self._monitor.in_error()
482 active_error = await demand_proc.interrupt_no_raise(abort=True)
484 error_kind = (
485 "controller"
486 if self._monitor.in_error()
487 else active_error.kind
488 if isinstance(active_error, DemandProcError)
489 else "internal"
490 if active_error
491 else None
492 )
494 match error_kind:
495 case "controller":
496 # Run the recovery loop if the controller is in an error state.
497 logger.error(f"Controller lifecycle error: {active_error or '<passive>'}")
499 try:
500 await self._error_recovery()
501 except Exception:
502 # If we were unable to recover from the error state, the controller is
503 # inoperable!
504 # FIXME: This will be replaced by blacklisting/backoff, and throwing in
505 # the towel will be left as a higher-level decision.
506 logger.critical("Controller is inoperable!")
507 self._can_operate.clear()
508 case "program":
509 logger.opt(exception=active_error).error(f"Program lifecycle error: {active_error}")
510 # TODO: Handle retry counting and blacklisting via ProgramStateManager.
512 return error_kind
513 finally:
514 # Clean up tasks.
515 for future in futures:
516 future.cancel()
518 async def _error_recovery(self):
519 self.belief_state = InternalControllerState.ERROR
520 retries_until_shutdown = 6
521 time_between_retries = 20.0
522 time_between_cycles = 120.0
523 recover_timeout = 120.0
524 shutdown_timeout = 240.0
526 while True:
527 self.step = LifecycleStep.RECOVER
529 # Attempt recovery several times.
530 for attempt in range(1, retries_until_shutdown + 1):
531 logger.info("Attempting recovery...")
533 try:
534 async with asyncio.timeout(recover_timeout):
535 await self.controller.execute_task(
536 RecoverTask(),
537 interrupt=True,
538 )
540 logger.info("Recovery succeeded.")
541 return
542 except Exception as e:
543 logger.error(f"Recovery attempt failed with {type(e).__name__}: {e}")
545 if attempt < retries_until_shutdown:
546 await asyncio.sleep(time_between_retries)
548 # We can't seem to recover, so we are getting desperate. Just try to shut down.
549 logger.warning("Failed to recover. Attempting shutdown.")
550 self.step = LifecycleStep.SHUTDOWN
552 try:
553 async with asyncio.timeout(shutdown_timeout):
554 await self.controller.execute_task(
555 ShutdownTask(),
556 interrupt=True,
557 )
558 except Exception as e:
559 logger.error(f"Emergency shutdown failed with {type(e).__name__}: {e}")
560 else:
561 # The shutdown task succeeded, in spite of the initial error state and recovery
562 # troubles. Let's play it safe and raise an error, which will cause the
563 # lifecycle logic to mark this Controller as inoperable pending external
564 # intervention.
565 logger.info("Emergency shutdown succeeded.")
566 self.belief_state = InternalControllerState.SHUTDOWN
567 raise RuntimeError("Emergency shutdown")
569 logger.warning(
570 "Emergency shutdown failed. "
571 f"Resuming recovery attempts in {time_between_cycles} sec."
572 )
573 self.step = LifecycleStep.WAIT
574 await asyncio.sleep(time_between_cycles)
576 class OperateProc(DemandProc):
577 """Drives the Controller to the OPERATE state."""
579 @override
580 async def state_logic(self):
581 controller = str(self.lifecycle.controller.entity)
583 # Begin INIT step.
584 self.lifecycle.step = LifecycleStep.INIT
586 logger.info(f"Preparing to operate ({controller})")
587 task = InitTask()
588 context = self.demand.contexts.get("init") if self.demand.contexts else None
590 self.lifecycle.belief_state = InternalControllerState.UNKNOWN # FIXME: clear to avoid shutdown no-op
591 await self.run_with_cleanup(task, context=context)
593 # Init was successful.
594 self.lifecycle.belief_state = InternalControllerState.OPERATE
596 # If there is no program specified, just wait.
597 if self.demand.program is None:
598 self.lifecycle.step = LifecycleStep.WAIT
599 logger.info(f"Waiting for a program for tasking ({controller})")
600 await asyncio.sleep(float("inf"))
602 while True:
603 # Begin TASKING step.
604 logger.info(f"Activating {self.demand.program.entity} program ({controller})")
605 self.lifecycle.step = LifecycleStep.TASKING
607 try:
608 async with asyncio.timeout(5.0):
609 await self.demand.program.start_tasking(self.demand.contexts)
610 except Exception as e:
611 raise DemandProcError("Failed to start tasking", kind="program") from e
613 logger.info(f"Controller is operating ({controller})")
615 # Wait until tasking stops out of band of our own stop logic. This will typically
616 # happen due to an error, but it is also possible that there was a third-party
617 # request for the Program to stop normally or the Program service restarted.
618 stop_event = await self.demand.program.wait_until_tasking_stops()
620 if stop_event is None:
621 raise DemandProcError("Failed to start tasking", kind="program")
623 match stop_event.origin:
624 case "error":
625 logger.error(f"Error reported ({controller})")
626 raise DemandProcError("Program reported error", kind="controller")
627 case "init":
628 # It looks like the Program service was restarted.
629 # TODO: Make sure this was a fast enough restart. Currently, there isn't
630 # enough information published to determine this.
631 logger.warning(f"Active program restarted ({controller})")
633 # Pause briefly and try again.
634 self.lifecycle.step = LifecycleStep.WAIT
635 await asyncio.sleep(3)
636 case "request":
637 logger.info(f"Tasking stopped by request ({controller})")
638 raise DemandProcError("Program was stopped", kind="program")
640 @override
641 async def stop_logic(self):
642 if self.demand.program:
643 try:
644 async with asyncio.timeout(150):
645 await self.demand.program.stop_tasking()
646 except TimeoutError as e:
647 # The procedure ends with this phase, so the abort rung has to be taken here.
648 # The reactor is blocked awaiting us and cannot escalate on our behalf.
649 logger.error("Program stop timed out, aborting tasking")
651 with contextlib.suppress(Exception):
652 await self.abort_logic()
654 raise DemandProcError("Program stop timed out", kind="program") from e
655 except BaseException as e:
656 logger.error(f"Program stop logic failed with {type(e).__name__}: {e}")
657 raise
659 @override
660 async def abort_logic(self):
661 if self.demand.program:
662 try:
663 async with asyncio.timeout(5):
664 await self.demand.program.abort_tasking()
665 except TimeoutError as e:
666 raise DemandProcError("Program abort timed out", kind="program") from e
668 class StandbyProc(DemandProc):
669 """Drives the Controller to the STANDBY state."""
671 @override
672 async def state_logic(self):
673 controller = str(self.lifecycle.controller.entity)
675 # Begin STANDBY step.
676 self.lifecycle.step = LifecycleStep.STANDBY
678 logger.info(f"Entering standby mode ({controller})")
679 task = StandbyTask()
680 context = self.demand.contexts.get("standby") if self.demand.contexts else None
682 self.lifecycle.belief_state = InternalControllerState.UNKNOWN # FIXME: clear to avoid shutdown no-op
683 await self.run_with_cleanup(task, context=context)
685 # Standby was successful.
686 self.lifecycle.belief_state = InternalControllerState.STANDBY
687 self.lifecycle.step = LifecycleStep.WAIT
689 logger.info(f"Controller is idle in standby ({controller})")
690 await asyncio.sleep(float("inf"))
692 class ShutdownProc(DemandProc):
693 """Drives the Controller to the SHUTDOWN state."""
695 @override
696 async def state_logic(self):
697 controller = str(self.lifecycle.controller.entity)
699 # Begin SHUTDOWN step.
700 self.lifecycle.step = LifecycleStep.SHUTDOWN
702 # FIXME: Need a passive monitor to check whether belief state equals actual (reported
703 # by Controller) state, and consider it a local error (trigger recovery logic)
704 # if not. That will force a re-command and hopefully re-synchronization of the
705 # belief state. Then, the no-op below should be safe.
706 if self.lifecycle.belief_state != InternalControllerState.SHUTDOWN:
707 logger.info(f"Commencing shutdown ({controller})")
708 task = ShutdownTask()
709 context = (
710 self.demand.contexts.get("shutdown") if self.demand.contexts else None
711 )
713 await self.run_with_cleanup(task, context=context)
715 # Shutdown was successful.
716 self.lifecycle.belief_state = InternalControllerState.SHUTDOWN
717 self.lifecycle.step = LifecycleStep.WAIT
719 logger.info(f"Controller is shut down ({controller})")
720 await asyncio.sleep(float("inf"))