Coverage for core / src / sensorkit / core / impl / entity.py: 97%
121 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 typing import TYPE_CHECKING, Any, Callable, ClassVar, Iterable, Literal, Self, final, override
9from loguru import logger
10from pydantic import BaseModel, TypeAdapter
12from sensorkit.backend.base import Entity, KVError, SpecialProperty
13from sensorkit.backend.event import Event
14from sensorkit.backend.request import ExtendedHandlerFunc, HandlerFunc, Request
15from sensorkit.common.keyword import Keyword, dump_keyword_json, get_keyword_info
16from sensorkit.core.entity import EntityBase, EntityInfo, EntityInterface
17from sensorkit.data.graph import DataGraph
19if TYPE_CHECKING:
20 from sensorkit.core.client import SensorKit, ServiceContext
23class EntityImpl(EntityBase, EntityInterface):
24 """Object that exposes implementation-side functionality of an entity."""
26 current: ClassVar[ContextVar[EntityImpl | None]] = ContextVar("current_entity", default=None)
27 """The instance in the current execution context, if any."""
29 @classmethod
30 def for_service_context(cls, context: ServiceContext, entity: str) -> Self:
31 """Return the entity in the current execution context."""
32 return cls(
33 sensorkit=context.sensorkit(),
34 entity=Entity.at(entity),
35 task_group=context.task_group,
36 )
38 def __init__(self, sensorkit: SensorKit, entity: Entity, *, task_group: asyncio.TaskGroup):
39 super().__init__(sensorkit, entity)
40 self._data_graph: DataGraph | None = None
41 self._task_group = task_group
42 self._attach_hooks: list[Callable] = []
43 self._detach_hooks: list[Callable] = []
45 async def init_impl(self):
46 """Implements initialization of the entity.
48 This method is called after the backend is up and available, but before `attach` is run.
49 Implementations will typically retrieve configuration, restore persisted state, and
50 perform any other initialization to prepare resources required by user `on_attach` hooks.
51 """
52 pass
54 def on_attach(self, func: Callable):
55 """Register a callback to run when this entity is being attached."""
56 self._attach_hooks.append(func)
57 return func
59 @final
60 async def attach(self):
61 """Attach the entity implementation to the running service context.
63 Runs all attach hooks and publishes entity info. An exception raised in any attach hook
64 is propagated as fatal.
65 """
66 await self._call_with_context(self._attach_hooks, error_policy="raise")
68 with self.enter_context():
69 await self.attach_impl()
71 await self.publish_entity_info()
73 async def attach_impl(self):
74 """Implements "attach". Runs after the user attach hooks."""
75 pass
77 def on_detach(self, func: Callable):
78 """Register a callback to run when this entity is being detached."""
79 self._detach_hooks.append(func)
80 return func
82 @final
83 async def detach(self):
84 """Detach the entity implementation from the running service context.
86 Runs all detach hooks. An exception raised in the internal detach hook (`detach_impl`) is
87 propagated as fatal, while exceptions from user detach hooks are logged.
88 """
89 with self.enter_context():
90 await self.detach_impl()
92 for result in await self._call_with_context(self._detach_hooks, error_policy="ignore"):
93 match result:
94 case asyncio.CancelledError():
95 logger.warning(f"Cancelled during {self.entity} detach")
96 case Exception():
97 logger.warning(f"Error during {self.entity} detach ({type(result).__name__})")
98 logger.opt(exception=result).debug("detach hook raised")
99 case BaseException():
100 raise result
102 async def detach_impl(self):
103 """Implements "detach". Runs before the user detach hooks."""
104 pass
106 @contextlib.contextmanager
107 def enter_context(self):
108 """Store a reference to this entity in the current execution context."""
109 var = type(self).current
110 token = var.set(self)
112 # Also set the base EntityImpl.current so sk.entity() resolves for all entity types.
113 base_var = EntityImpl.current
114 base_token = base_var.set(self) if var is not base_var else None
116 try:
117 yield
118 finally:
119 var.reset(token)
121 if base_token is not None:
122 base_var.reset(base_token)
124 @override
125 @property
126 def task_group(self):
127 return self._task_group
129 @override
130 async def emit_event(self, event: Event):
131 """Serialize and publish the event to the entity's event stream subject."""
132 await self._stream.publish(
133 SpecialProperty.EVENTS,
134 event.model_dump_json().encode(),
135 )
137 @override
138 async def publish(self, model: Keyword):
139 """Serialize and publish a keyword model to the entity's data stream."""
140 if info := get_keyword_info(model):
141 await self._stream.publish(
142 info.key,
143 dump_keyword_json(model),
144 )
145 else:
146 await self._stream.publish(
147 model.__class__.__name__,
148 model.model_dump_json().encode()
149 if isinstance(model, BaseModel)
150 else TypeAdapter(model).dump_json(model),
151 )
153 @override
154 async def handle_request[P: BaseModel | None, R: BaseModel | None, V: BaseModel | None](
155 self,
156 request: Request[P, R, V],
157 func: HandlerFunc[P, R] | ExtendedHandlerFunc[P, R, V],
158 ):
159 """Register a handler for the given Request definition on the backend."""
160 await self._request.handle_request(
161 request.name,
162 request.create_handler(func, self._stream),
163 )
165 @override
166 async def data_graph(self):
167 """Return an AppSource if there is a DataGraph is associated with this entity."""
168 if not self._data_graph:
169 with contextlib.suppress(KVError):
170 dg = await self.kv_get_model(DataGraph)
172 # Make sure a concurrent call does not result in two graphs.
173 if not self._data_graph:
174 self._data_graph = dg
175 self._data_graph.start(task_group=self.task_group)
176 logger.info(f"Started DataGraph with {len(self._data_graph.nodes)} ops")
178 return self._data_graph
180 def entity_info(self) -> EntityInfo:
181 """Build this entity's info record from its current state."""
182 return EntityInfo(entity_type="generic", details=None)
184 @override
185 async def publish_entity_info(self) -> EntityInfo:
186 """Publish this entity's EntityInfo entry to the KV store."""
187 info = self.entity_info()
188 await self.kv_put_model(info)
189 return info
191 async def _call_with_context(
192 self,
193 funcs: Iterable[Callable],
194 args: tuple[Any, ...] = (),
195 kwargs: dict[str, Any] | None = None,
196 *,
197 error_policy: Literal["raise", "log", "ignore"] = "log",
198 ) -> list[Any]:
199 """Invoke hooks concurrently within this entity's execution context."""
200 call_kwargs = kwargs if kwargs is not None else {}
202 async def invoke(func: Callable):
203 aw = func(*args, **call_kwargs)
204 return await aw if asyncio.iscoroutine(aw) else aw
206 with self.enter_context():
207 tasks = [asyncio.ensure_future(invoke(func)) for func in funcs]
209 if error_policy == "raise":
210 # Fail fast: propagate the first hook exception immediately, cancelling any
211 # still-running siblings so none are left orphaned on the loop.
212 try:
213 return await asyncio.gather(*tasks)
214 finally:
215 for task in tasks:
216 task.cancel()
217 await asyncio.gather(*tasks, return_exceptions=True)
219 results = await asyncio.gather(*tasks, return_exceptions=True)
221 if error_policy == "log":
222 for result in results:
223 match result:
224 case Exception():
225 logger.opt(exception=result).debug("error in callback")
226 case BaseException():
227 logger.debug(f"error in callback ({type(result).__name__})")
229 return results