Coverage for core / src / sensorkit / core / task.py: 92%
133 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 uuid
6from collections.abc import Generator
7from datetime import datetime, timedelta
8from typing import TYPE_CHECKING, Annotated, Any, ClassVar, Literal, override
10from pydantic import BaseModel, BeforeValidator, Field, PrivateAttr
12from sensorkit.common.keyword import KeywordDict, declare_keyword, validated_items
13from sensorkit.common.model import ModelRegistry, RegistryBaseModel
15if TYPE_CHECKING:
16 from sensorkit.core.controller import InternalControllerState, TaskExecutionResult
19class Task(RegistryBaseModel):
20 """Base task definition.
22 This is the user-extensible part of the task system: subclasses add domain-specific fields and
23 register automatically via the `task_type` discriminator. Identity and execution-envelope data
24 (`task_id`, `controller_id`, `context`, `expiry_time`) live separately on `TaskExecution`,
25 which the controller mints when a task is submitted for execution.
26 """
28 task_type: Literal[None] = None
30 # The execution envelope is associated by the controller immediately before a task handler is
31 # invoked. It is excluded from serialization and only meaningful server-side within a handler.
32 _execution: TaskExecution | None = PrivateAttr(default=None)
34 # Task model registry.
35 registry: ClassVar[ModelRegistry[Task]] = ModelRegistry(discriminator="task_type")
37 @classmethod
38 def model_registry(cls):
39 return cls.registry
41 @property
42 def execution(self) -> TaskExecution:
43 """The `TaskExecution` envelope associated with this task.
45 The controller associates the envelope just before invoking a task handler, so it is
46 always available from within a handler.
48 Returns:
49 The associated execution envelope.
51 Raises:
52 RuntimeError: If accessed when no execution has been associated (e.g. outside a task
53 handler).
54 """
55 if self._execution is None:
56 raise RuntimeError("task has no execution context (accessed outside a task handler?)")
58 return self._execution
60 def associate_execution(self, execution: TaskExecution) -> None:
61 """Associate an execution envelope with this task.
63 The controller calls this immediately before invoking a task handler so the handler can
64 reach the envelope via [`execution`][sensorkit.core.task.Task.execution].
66 Args:
67 execution: The execution envelope to associate with this task.
69 Raises:
70 RuntimeError: If an execution is already associated with this task.
71 """
72 if self._execution is not None:
73 raise RuntimeError("task already has an associated execution")
75 self._execution = execution
77 def __eq__(self, other: object) -> bool:
78 # The execution back-link is transient runtime state (excluded from serialization), so it
79 # must not affect equality: two tasks with the same semantic content are equal whether or
80 # not either is currently associated with an execution.
81 if not isinstance(other, Task):
82 return NotImplemented
84 return (
85 type(self) is type(other)
86 and self.__dict__ == other.__dict__
87 and self.__pydantic_extra__ == other.__pydantic_extra__
88 )
90 __hash__ = None
92 def target_state(self) -> InternalControllerState | None:
93 """Return the state-transition target if this is a lifecycle task.
95 Returns:
96 The target controller state, or `None` for non-lifecycle tasks.
97 """
98 # FIXME: This mapping should be defined by other means in controller.py.
99 return None
101 def default_expiry(self) -> datetime | timedelta:
102 """Return the default execution expiry for this task.
104 Used by the tasking loop when the task factory supplies no explicit `expiry_time`. A
105 `timedelta` is interpreted relative to dispatch time; a `datetime` is absolute. Subclasses
106 may override to express a domain-specific deadline.
108 Returns:
109 The default expiry as an absolute time or a duration from dispatch.
110 """
111 return timedelta(seconds=300)
113 def submit(
114 self,
115 *,
116 context: KeywordDict | None = None,
117 expiry_time: datetime | None = None,
118 ) -> TaskSubmission:
119 """Bundle this task with execution parameters into a `TaskSubmission`.
121 A convenience for task factories: `yield task.submit(expiry_time=...)` reads more fluently
122 than constructing a `TaskSubmission` by hand. Yielding a bare task is equivalent to
123 `task.submit()` with no parameters.
125 Args:
126 context: Optional keyword context to attach to the execution.
127 expiry_time: Optional time after which the execution should be considered expired.
129 Returns:
130 A `TaskSubmission` wrapping this task and the given parameters.
131 """
132 return TaskSubmission(task=self, context=context, expiry_time=expiry_time)
135class InitTask(Task):
136 """Init Task"""
138 task_type: Literal["init"] = "init"
140 @override
141 def target_state(self):
142 from sensorkit.core.controller import InternalControllerState
144 return InternalControllerState.OPERATE
147class StandbyTask(Task):
148 """Standby Task"""
150 task_type: Literal["standby"] = "standby"
152 @override
153 def target_state(self):
154 from sensorkit.core.controller import InternalControllerState
156 return InternalControllerState.STANDBY
159class ShutdownTask(Task):
160 """Shutdown Task"""
162 task_type: Literal["shutdown"] = "shutdown"
164 @override
165 def target_state(self):
166 from sensorkit.core.controller import InternalControllerState
168 return InternalControllerState.SHUTDOWN
171class CalibrateTask(Task):
172 """Calibrate Task"""
174 task_type: Literal["calibrate"] = "calibrate"
177class RecoverTask(Task):
178 """Recover Task"""
180 task_type: Literal["recover"] = "recover"
183class CollectTask(Task):
184 """Collect Task"""
186 task_type: Literal["collect"] = "collect"
189@declare_keyword
190class TaskInfo(BaseModel):
191 """Keyword describing a task that has been executed on a controller."""
192 task: Task
193 task_id: uuid.UUID
194 controller_id: str
196 def get_fits_cards(self):
197 yield "SKTASK", (self.task.task_type, "SensorKit task type")
198 yield "SKTASKID", (str(self.task_id), "SensorKit task ID")
199 yield "SKCTRL", (self.controller_id, "SensorKit controller name")
202class TaskExecution(BaseModel):
203 """Execution envelope wrapping a `Task`.
205 Carries server-minted identity (`task_id`, `controller_id`) and client-supplied execution
206 parameters (`context`, `expiry_time`). This type is not user-extensible; the extensible
207 semantic definition is the embedded `task`.
208 """
210 task: Task
211 task_id: uuid.UUID
212 controller_id: str
213 context: Annotated[
214 KeywordDict,
215 BeforeValidator(lambda v: KeywordDict() if v is None else v),
216 Field(default_factory=KeywordDict),
217 ]
218 expiry_time: datetime | None = None
220 # Live result future, bound only in the context that owns the in-flight call (the client
221 # tasking loop today). Transient runtime state: excluded from serialization and equality, so a
222 # deserialized execution carries no result future and is not awaitable.
223 _result: asyncio.Future[TaskExecutionResult] | None = PrivateAttr(default=None)
225 def bind_result(self, result: asyncio.Future[TaskExecutionResult]) -> None:
226 """Bind the live result future for this execution.
228 The dispatching context binds the in-flight call's future so that holders of this execution
229 can await its completion or inspect the result.
231 Args:
232 result: The future resolving to this execution's final result.
234 Raises:
235 RuntimeError: If a result future is already bound.
236 """
237 if self._result is not None:
238 raise RuntimeError("execution already has a bound result")
240 self._result = result
242 def _require_result(self) -> asyncio.Future[TaskExecutionResult]:
243 if self._result is None:
244 raise RuntimeError("execution is not awaitable in this context")
246 return self._result
248 def __await__(self) -> Generator[Any, None, TaskExecutionResult]:
249 """Await the final result of this execution."""
250 return self._require_result().__await__()
252 def done(self) -> bool:
253 """Return True if the execution has completed (result or error)."""
254 return self._require_result().done()
256 def result(self) -> TaskExecutionResult:
257 """Return the final result, raising if the execution has not yet completed."""
258 return self._require_result().result()
261class TaskSubmission(BaseModel):
262 """A Task bundled with client-supplied execution parameters.
264 A task factory may return this in place of a bare `Task` to attach per-instance execution
265 parameters (`context`, `expiry_time`) that the controller records on the minted
266 `TaskExecution`. Returning a bare `Task` is equivalent to a `TaskSubmission` with no parameters;
267 `Task.submit` is the convenient way to build one.
268 """
270 task: Task
271 context: KeywordDict | None = None
272 expiry_time: datetime | None = None
275type RawKeywords = dict[str, Any]
276"""A sparse, unvalidated mapping of keyword key to raw payload, as declared in config."""
279def _deep_merge(base: Any, override: Any) -> Any:
280 """Recursively merge `override` onto `base`, with `override` winning on conflict.
282 Dicts are merged key-by-key (so a keyword payload's fields combine across layers); any
283 other value, including a list, is replaced wholesale. A `None` override keeps `base`,
284 which lets a layer contribute a key another layer omits.
285 """
286 if override is None:
287 return base
289 if isinstance(base, dict) and isinstance(override, dict):
290 merged = dict(base)
292 for key, value in override.items():
293 merged[key] = _deep_merge(base.get(key), value)
295 return merged
297 return override
300class TaskContextOverlay(BaseModel, extra="allow"):
301 """Raw, sparse, layered task context as declared in config, prior to validation.
303 The `all` field provides keywords that apply to every task type. Additional fields
304 (`init`, `standby`, `shutdown`, or any custom task type name) provide type-specific
305 overrides. Extra fields (via `extra="allow"`) support custom task types.
307 Values are kept as raw payloads and merged field-wise across layers; validation into
308 keyword instances is deferred to `build`, so overlapping keywords combine rather than
309 clobber and every payload is validated exactly once.
310 """
312 all: RawKeywords = Field(default_factory=dict)
313 init: RawKeywords = Field(default_factory=dict)
314 standby: RawKeywords = Field(default_factory=dict)
315 shutdown: RawKeywords = Field(default_factory=dict)
316 __pydantic_extra__: dict[str, RawKeywords] = Field(init=False)
318 def _task_type_items(self):
319 for field, keywords in self:
320 if field != "all":
321 yield field, keywords
323 def propagate(self, into: TaskContextOverlay):
324 """Merge this (parent) overlay into `into` (child), with the child taking precedence.
326 A keyword present in both layers is merged field-wise, so a child override fills in
327 rather than replaces the parent's fields. A parent's type-specific keyword is skipped
328 when the child already carries that keyword in its `all`, since the child's `all`
329 already supplies it for every task type.
330 """
331 for task_type, keywords in self._task_type_items():
332 target = getattr(into, task_type, None)
334 if target is None:
335 target = {}
336 setattr(into, task_type, target)
338 for key, payload in keywords.items():
339 if key in into.all:
340 continue
342 target[key] = _deep_merge(payload, target.get(key))
344 into.all = _deep_merge(self.all, into.all)
346 def build(self) -> TaskContextMap:
347 """Flatten the `all` layer into each task type and validate every payload as a keyword."""
348 return TaskContextMap(
349 defaults=KeywordDict(validated_items(self.all)),
350 by_type={
351 task_type: KeywordDict(validated_items(_deep_merge(self.all, keywords)))
352 for task_type, keywords in self._task_type_items()
353 },
354 )
357class TaskContextMap(BaseModel):
358 """Flattened, validated task contexts keyed by task type."""
360 defaults: KeywordDict = Field(default_factory=KeywordDict)
361 by_type: dict[str, KeywordDict] = Field(default_factory=dict)
363 def get(self, task_type: str) -> KeywordDict:
364 """Return the effective context for `task_type`, falling back to `defaults`."""
365 return self.by_type.get(task_type, self.defaults).copy()