Coverage for core / src / sensorkit / core / controller.py: 95%
156 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
4from abc import abstractmethod
5from collections.abc import Coroutine
6from dataclasses import dataclass
7from datetime import UTC, datetime
8from enum import StrEnum, auto
9from typing import TYPE_CHECKING, Any, Callable, ClassVar, Collection, Mapping, Self
11import uuid_utils.compat as uuid
12from pydantic import BaseModel, Field, model_validator
14from sensorkit.backend.event import Event
15from sensorkit.backend.request import Call, ExtendedResponse, Request
16from sensorkit.common.keyword import KeywordDict
17from sensorkit.core.device import DeviceClient
18from sensorkit.core.entity import EntityClient, EntityInterface, EntityRef
19from sensorkit.core.state import EventSourcedState
20from sensorkit.core.task import Task, TaskExecution
21from sensorkit.data.context import Context, ContextSubscription
23if TYPE_CHECKING:
24 from sensorkit.core.client import SensorKit
27class InternalControllerState(StrEnum):
28 """High-level Controller states."""
29 OPERATE = auto()
30 STANDBY = auto()
31 SHUTDOWN = auto()
32 ERROR = auto()
33 UNKNOWN = auto()
36class ControllerEnableState(Event):
37 """Event indicating the Controller enable state has changed."""
38 enabled: bool
41class TaskFinishInfo(BaseModel):
42 """Metadata recorded when a task completes, including whether it was aborted or raised an error."""
44 timestamp: datetime = Field(default_factory=lambda: datetime.now(UTC))
45 aborted: bool = False
46 error: bool | str = False
49class TaskExecutionState(Event):
50 """Event representing the current task execution status of a Controller."""
52 executing: bool = False
53 aborting: bool = False
54 finished: TaskFinishInfo | None = None
55 execution: TaskExecution | None
56 context: dict | None = None
58 @model_validator(mode="before")
59 @classmethod
60 def _migrate_legacy_task(cls, data: Any) -> Any:
61 """Upgrade events persisted by versions predating the Task/TaskExecution split.
63 Two legacy shapes are upgraded here:
65 * Versions before the split stored the executing task on `task` as a flat
66 `ControllerTask` — identity (`task_id`, `controller_id`), `context` and the
67 `end_time` deadline embedded alongside the discriminated task fields. This version
68 expects a `TaskExecution` envelope wrapping the semantic `task`.
69 * Versions after the split but before `task` was renamed to `execution` stored the
70 envelope under the `task` key.
72 Both are carried onto `execution`. Rewriting the legacy shape here covers both reads:
73 the KV `ControllerState` snapshot (whose `execution_state` is validated as this model)
74 and event-stream replay (where each event is validated through the registry). Without it,
75 state carried by a NATS broker from an older version would raise a `ValidationError`.
76 """
77 if not isinstance(data, dict):
78 return data
80 # The pre-rename wire format carried the executing task (flat or enveloped) under `task`.
81 # Lift it onto `execution` so the rest of the migration — and the model — see one key.
82 if "execution" not in data and "task" in data:
83 data = dict(data)
84 data["execution"] = data.pop("task")
86 legacy = data.get("execution")
88 # A legacy task is a flat `ControllerTask`: it carries the discriminator at the top level
89 # and has no nested `task` envelope. A current `TaskExecution` always nests `task`.
90 if not isinstance(legacy, dict) or "task_type" not in legacy or "task" in legacy:
91 return data
93 # Lift the envelope fields off the task; the remaining keys (`task_type` plus domain
94 # fields) become the wrapped semantic task. The legacy `end_time` was the execution
95 # deadline, so it maps to the envelope's `expiry_time`; task types that still declare
96 # `end_time` as a domain field (e.g. `standard_collect`) keep their own copy because it
97 # is left in the wrapped task.
98 envelope = {
99 "task": {
100 k: v for k, v in legacy.items() if k not in ("task_id", "controller_id", "context")
101 },
102 "task_id": legacy.get("task_id"),
103 "controller_id": legacy.get("controller_id"),
104 "context": legacy.get("context"),
105 "expiry_time": legacy.get("end_time"),
106 }
108 return {**data, "execution": envelope}
110 @model_validator(mode="after")
111 def _validate(self):
112 if self.executing and self.execution is None:
113 raise ValueError("inconsistent state fields: executing with no task")
115 if self.aborting and not self.executing:
116 raise ValueError("inconsistent state fields: aborting but not executing")
118 if self.finished is not None and (self.executing or self.aborting):
119 raise ValueError("inconsistent state fields: finished but executing or aborting")
121 return self
124class ControllerOperatingState(Event):
125 """Event indicating the Controller operating state has changed."""
126 current: InternalControllerState
127 previous: InternalControllerState | None = None
128 target: InternalControllerState | None = None
130 NO_CHANGE: ClassVar[object] = object()
132 def derive(
133 self,
134 *,
135 current: InternalControllerState = NO_CHANGE,
136 target: InternalControllerState | None = NO_CHANGE,
137 ) -> Self:
138 """Return a new ControllerOperatingState based on this one, updating only the specified fields."""
139 if current == self.current:
140 # This ensures `previous` isn't unnecessarily rotated out.
141 current = self.NO_CHANGE
143 return ControllerOperatingState(
144 current=self.current if current == self.NO_CHANGE else current,
145 previous=self.previous if current == self.NO_CHANGE else self.current,
146 target=self.target if target == self.NO_CHANGE else target,
147 )
150class ControllerState(EventSourcedState):
151 """Controller state."""
152 enable_state: ControllerEnableState
153 operating_state: ControllerOperatingState
154 execution_state: TaskExecutionState
157class ControllerEnableStateRequest(BaseModel):
158 """Request that a Controller enable or disable its task handlers."""
159 enable: bool
162set_enable_state_request = Request.define(
163 "set_enable_state",
164 payload=ControllerEnableStateRequest,
165)
166"""Set the enable state of a Controller."""
169class ExecuteRequestMessage(BaseModel):
170 """A task execution request.
172 The client supplies the semantic `task` plus optional execution parameters. The controller
173 mints the identity (`task_id`, `controller_id`) and returns the resulting `TaskExecution` in
174 the response.
175 """
176 task: Task
177 context: KeywordDict | None = None
178 expiry_time: datetime | None = None
179 interrupt: bool = False
182class ExecuteResponseMessage(ExtendedResponse):
183 """A response to a task execution request.
185 Carries the minted `TaskExecution` envelope so the client learns the assigned `task_id` (e.g.
186 for a subsequent abort) without having to supply it.
187 """
188 execution: TaskExecution | None = None
191class AbortRequestMessage(BaseModel):
192 """A Task abort request."""
193 task_id: uuid.UUID | None
194 """If given, verifies the running task ID matches before aborting."""
197class AbortResponseMessage(ExtendedResponse):
198 """Response to a Task abort request."""
199 aborting: bool
200 """True if the abort is underway, False if the abort was rejected."""
202 task_id: uuid.UUID | None
203 """If the abort was accepted, the task ID of the task being aborted."""
206class TaskExecutionResult(BaseModel):
207 """Result of a successful Task execution."""
208 task_id: uuid.UUID
209 start_time: datetime
210 end_time: datetime
213execute_task_request = Request.define(
214 name="execute_task",
215 payload=ExecuteRequestMessage,
216 response=ExecuteResponseMessage,
217 result=TaskExecutionResult,
218)
220abort_task_request = Request.define(
221 name="abort_task",
222 payload=AbortRequestMessage,
223 response=AbortResponseMessage,
224)
227class ControllerClient(EntityClient):
228 """Object that exposes client-side functionality of a Controller."""
230 async def enable(self):
231 """Request that the Program enable task sourcing for the target Controller."""
232 return await self.call(
233 set_enable_state_request,
234 ControllerEnableStateRequest(enable=True)
235 )
237 async def disable(self):
238 """Request that the Program disable task sourcing."""
239 return await self.call(
240 set_enable_state_request,
241 ControllerEnableStateRequest(enable=False)
242 )
244 def execute_task(
245 self,
246 task: Task,
247 *,
248 context: KeywordDict | None = None,
249 expiry_time: datetime | None = None,
250 interrupt=False,
251 ) -> Call[ExecuteResponseMessage, TaskExecutionResult]:
252 """Send a task execution request to the controller and return a `Call` tracking completion.
254 The controller assigns the `task_id` and `controller_id`. The optional execution
255 parameters are recorded on the resulting `TaskExecution`.
257 Args:
258 task: The semantic task to execute.
259 context: Optional keyword context to attach to the execution.
260 expiry_time: Optional time after which the execution should be considered expired.
261 interrupt: Whether to interrupt any task currently in progress.
263 Returns:
264 A `Call` that yields the final `TaskExecutionResult` when awaited.
265 """
266 return self.call(
267 execute_task_request,
268 ExecuteRequestMessage(
269 task=task,
270 context=context,
271 expiry_time=expiry_time,
272 interrupt=interrupt,
273 ),
274 )
276 async def start_task(
277 self,
278 task: Task,
279 *,
280 context: KeywordDict | None = None,
281 expiry_time: datetime | None = None,
282 interrupt=False,
283 ) -> TaskExecution:
284 """Submit a task and return its execution envelope once the controller acknowledges it.
286 Unlike `execute_task`, this awaits only the controller's initial response, then returns the
287 minted `TaskExecution`. The controller assigns the identity (`task_id`, `controller_id`);
288 the returned execution is bound to the in-flight call, so the caller can read the assigned
289 identity immediately and `await` the execution itself for the final `TaskExecutionResult`.
291 Args:
292 task: The semantic task to execute.
293 context: Optional keyword context to attach to the execution.
294 expiry_time: Optional time after which the execution should be considered expired.
295 interrupt: Whether to interrupt any task currently in progress.
297 Returns:
298 The minted execution envelope, awaitable for the final result.
299 """
300 call = self.execute_task(
301 task, context=context, expiry_time=expiry_time, interrupt=interrupt
302 )
303 response = await call.invoke()
304 execution = response.execution
305 assert execution is not None, "controller acknowledged task without an execution envelope"
307 # The in-flight call's future resolves to the result in this (client) context.
308 execution.bind_result(call.get_future())
310 return execution
312 def abort_task(self, task_id: uuid.UUID | None = None) -> Call[AbortResponseMessage, None]:
313 """Send an abort request to the controller for the optionally specified task ID."""
314 return self.call(abort_task_request, AbortRequestMessage(task_id=task_id))
316 async def wait_for_task(self, task_id: uuid.UUID | None = None):
317 """Wait until the controller reports that no task (or the specified task) is executing."""
318 # FIXME: This is unreliable at present due to backend limitations.
319 async for event in ControllerState.event_stream(self, TaskExecutionState):
320 if task_id is not None:
321 if event.execution is None or event.execution.task_id != task_id:
322 return None
324 if not event.executing:
325 return event
327 raise RuntimeError("stream ended unexpectedly")
330class ControllerRef(EntityRef[ControllerClient]):
331 """A serializable reference to a controller client."""
333 def _get_client(self, kit: SensorKit) -> ControllerClient:
334 return kit.controller(self.name)
337@dataclass
338class ControllerDevice(Mapping[Any, Any]):
339 """A device attached to a controller, providing cached keyword access via its ContextSubscription."""
341 client: DeviceClient
342 subscription: ContextSubscription
344 def __getitem__(self, key, /):
345 return self.subscription.cache[key]
347 def __len__(self):
348 return len(self.subscription.cache)
350 def __iter__(self):
351 return iter(self.subscription.cache)
354type TaskHandlerCallback[T: Task] = Callable[[T], Coroutine[Any, Any, None]]
357class ControllerInterface(EntityInterface):
358 """Interface describing an implementation of a controller."""
360 @abstractmethod
361 def on_enable(self, func: Callable[[], None]):
362 """Register a callback to invoke when the controller is enabled."""
363 ...
365 @abstractmethod
366 def on_disable(self, func: Callable[[], None]):
367 """Register a callback to invoke when the controller is disabled."""
368 ...
370 @abstractmethod
371 def use_device(self, name: str, *, subscribe: list[type] | None = None):
372 """Declare a device dependency, optionally subscribing to the listed keyword types."""
373 ...
375 @abstractmethod
376 def get_device(self, name: str) -> ControllerDevice:
377 """Return the attached ControllerDevice for the given name."""
378 ...
380 @abstractmethod
381 def all_devices(self) -> Collection[ControllerDevice]:
382 """Return all attached ControllerDevice objects."""
383 ...
385 @abstractmethod
386 async def start_device_subscriptions(self):
387 """Start keyword subscriptions for all declared devices."""
388 ...
390 @abstractmethod
391 async def stop_device_subscriptions(self):
392 """Stop keyword subscriptions for all declared devices."""
393 ...
395 @abstractmethod
396 async def update_context(
397 self,
398 *args,
399 **kwargs,
400 ) -> Context:
401 """Update the context of the currently executing task."""
402 ...
404 @abstractmethod
405 def task_handler(self, task_type: type[Task]) -> Callable[..., TaskHandlerCallback]:
406 """Register a handler for the given task type and return a decorator."""
407 ...
409 @abstractmethod
410 def task_running(self) -> bool:
411 """Return True if a task is currently executing on this controller."""
412 ...
414 @abstractmethod
415 async def set_internal_state(self, state: InternalControllerState):
416 """Transition the controller to the given internal operating state."""
417 ...