Coverage for core / src / sensorkit / core / client.py: 91%

1# SPDX-License-Identifier: Apache-2.0 

2import asyncio 

3import contextlib 

4from dataclasses import dataclass 

5from datetime import datetime 

6from typing import Any, Self, overload, override 

7 

8from loguru import logger 

9from pydantic import BaseModel 

10 

11from sensorkit.backend.base import Backend, BackendError, BackendImpl, Entity, ServiceInfo 

12from sensorkit.backend.lease import LeaseGroup 

13from sensorkit.core.controller import ControllerClient 

14from sensorkit.core.device import DeviceClient 

15from sensorkit.core.entity import DeviceDetails, EntityClient, EntityInfo 

16from sensorkit.core.impl.controller import ControllerImpl 

17from sensorkit.core.impl.device import DeviceImpl 

18from sensorkit.core.impl.entity import EntityImpl 

19from sensorkit.core.impl.program import ProgramImpl 

20from sensorkit.core.program import ProgramClient 

21from sensorkit.core.trait import Trait 

22 

23 

24class SensorKit: 

25 """The main SensorKit API entrypoint.""" 

26 

27 def __init__(self, backend: Backend | BackendImpl): 

28 if isinstance(backend, BackendImpl): 

29 backend = Backend(impl=backend) 

30 

31 self.backend: Backend = backend 

32 self._clients: dict[Entity, EntityClient] = {} 

33 self._service_registered = False 

34 

35 async def register_service(self, name: str, version: str): 

36 """Register this SensorKit instance as a named service and return a ServiceContext.""" 

37 if self._service_registered: 

38 raise RuntimeError("Service already registered") 

39 

40 self._service_registered = True 

41 return await ServiceContext.register( 

42 self, 

43 ServiceInfo(name=name, version=version), 

44 ) 

45 

46 async def list_services(self): 

47 """Return a mapping of entity to ServiceStatus for all services visible in the KV store.""" 

48 records: dict[Entity, ServiceRecord] = {} 

49 online: set[Entity] = set() 

50 

51 for entry in await self.backend.key_value().get_all(deep=True): 

52 match entry.key.prop: 

53 case "ServiceRecord": 

54 records[entry.key.entity()] = ServiceRecord.model_validate_json(entry.value) 

55 case "EntityLease": 

56 online.add(entry.key.entity()) 

57 

58 return { 

59 entity: ServiceStatus(service=record, online=entity in online) 

60 for entity, record in records.items() 

61 } 

62 

63 async def list_entities(self) -> dict[Entity, EntityInfo]: 

64 """Return a mapping of entity to lease-value string for all currently online entities.""" 

65 return { 

66 entry.key.entity(): EntityInfo.model_validate_json(entry.value) 

67 for entry in await self.backend.key_value().get_all(deep=True) 

68 if entry.key.prop == "EntityInfo" 

69 } 

70 

71 async def find_devices(self, *, match_trait: Trait) -> list[DeviceClient]: 

72 """Return DeviceClients for all online devices matching the given trait.""" 

73 results: list[DeviceClient] = [] 

74 

75 for entry in await self.backend.key_value().get_all(deep=True): 

76 if entry.key.prop != "EntityInfo": 

77 continue 

78 

79 info = EntityInfo.model_validate_json(entry.value) 

80 if not isinstance(info.details, DeviceDetails): 

81 continue 

82 

83 if match_trait.match(info.details): 

84 results.append(self.device(entry.key.entity())) 

85 

86 return results 

87 

88 @overload 

89 def entity(self, entity: Entity) -> EntityClient: ... 

90 

91 @overload 

92 def entity(self, *path: str) -> EntityClient: ... 

93 

94 def entity(self, *args: Any): 

95 """Return (or create) an EntityClient for the given entity path or Entity object.""" 

96 entity = args[0] if args and isinstance(args[0], Entity) else Entity(args) 

97 

98 if entity in self._clients: 

99 return self._clients[entity] 

100 

101 return self._clients.setdefault(entity, EntityClient(self, entity)) 

102 

103 @overload 

104 def device(self, entity: Entity) -> DeviceClient: ... 

105 

106 @overload 

107 def device(self, *path: str) -> DeviceClient: ... 

108 

109 def device(self, *args: Any): 

110 """Return (or create) a DeviceClient for the given entity path or Entity object.""" 

111 entity = args[0] if args and isinstance(args[0], Entity) else Entity(args) 

112 

113 if entity in self._clients: 

114 return self._clients[entity] 

115 

116 return self._clients.setdefault(entity, DeviceClient(self, entity)) 

117 

118 @overload 

119 def controller(self, entity: Entity) -> ControllerClient: ... 

120 

121 @overload 

122 def controller(self, *path: str) -> ControllerClient: ... 

123 

124 def controller(self, *args: Any): 

125 """Return (or create) a ControllerClient for the given entity path or Entity object.""" 

126 entity = args[0] if args and isinstance(args[0], Entity) else Entity(args) 

127 

128 if entity in self._clients: 

129 return self._clients[entity] 

130 

131 return self._clients.setdefault(entity, ControllerClient(self, entity)) 

132 

133 @overload 

134 def program(self, entity: Entity) -> ProgramClient: ... 

135 

136 @overload 

137 def program(self, *path: str) -> ProgramClient: ... 

138 

139 def program(self, *args: Any): 

140 """Return (or create) a ProgramClient for the given entity path or Entity object.""" 

141 entity = args[0] if args and isinstance(args[0], Entity) else Entity(args) 

142 

143 if entity in self._clients: 

144 return self._clients[entity] 

145 

146 return self._clients.setdefault(entity, ProgramClient(self, entity)) 

147 

148 

149class ServiceRecord(BaseModel): 

150 """Information about a service entity.""" 

151 info: ServiceInfo 

152 last_registered: datetime 

153 

154 

155@dataclass 

156class ServiceStatus: 

157 """Status of a service.""" 

158 service: ServiceRecord 

159 online: bool 

160 

161 

162class ServiceContext(EntityImpl): 

163 """An object that exposes SensorKit service API.""" 

164 

165 ENTITY_LEASE_TTL = 60.0 

166 SERVICE_LEASE_TTL = 10.0 

167 

168 @classmethod 

169 async def register(cls, kit: SensorKit, info: ServiceInfo) -> Self: 

170 """Create a ServiceContext and register itself as an entity.""" 

171 instance = cls(kit, info) 

172 return await instance.register_impl(instance, lease_ttl=cls.SERVICE_LEASE_TTL) 

173 

174 def __init__(self, sensorkit: SensorKit, info: ServiceInfo): 

175 super().__init__( 

176 sensorkit=sensorkit, entity=Entity.at(info.name), task_group=asyncio.TaskGroup() 

177 ) 

178 

179 self.info = info 

180 self._lease_group = LeaseGroup() 

181 self._impls: list[EntityImpl] = [] 

182 self._shutdown = asyncio.get_running_loop().create_future() 

183 self._shutdown_called = False 

184 

185 async def register_impl[T: EntityImpl]( 

186 self, 

187 impl: T, 

188 *, 

189 acquire_lease: bool = True, 

190 lease_ttl: float = ENTITY_LEASE_TTL, 

191 ) -> T: 

192 """Register an entity implementation on the backend. 

193 

194 If `acquire_lease` is False, no lease is acquired. This should only be done in special 

195 scenarios where the entity is already leased by another implementation in the same process. 

196 """ 

197 if acquire_lease: 

198 await self._lease_group.acquire( 

199 impl._kv, 

200 key="EntityLease", 

201 ttl=lease_ttl, 

202 record=self.info, 

203 ) 

204 

205 # Run baseline entity initialization. 

206 await impl.init_impl() 

207 

208 # Track the impl before attaching so that an impl whose attach raises partway is still 

209 # detached on shutdown. 

210 self._impls.append(impl) 

211 await impl.attach() 

212 

213 return impl 

214 

215 @override 

216 async def init_impl(self): 

217 """Start the main service task.""" 

218 # Publish information about this service. 

219 await self.kv_put_model( 

220 ServiceRecord(info=self.info, last_registered=datetime.now()) 

221 ) 

222 

223 # Notify the backend implementation that we are a service. 

224 with contextlib.suppress(Exception): 

225 await self.backend.register_service(self.info) 

226 

227 # Start the service task. 

228 ready = asyncio.Event() 

229 self._task = asyncio.create_task( 

230 self._service_task(ready), 

231 name=f"service-{self.info.name}", 

232 ) 

233 await ready.wait() 

234 

235 async def _service_task(self, ready: asyncio.Event): 

236 ran_detach = False 

237 logger.debug(f"Service {self.info.name} startup") 

238 

239 try: 

240 async with self.task_group: 

241 # Signal the outer coroutine that the TaskGroup is active. 

242 ready.set() 

243 

244 # Start lease maintenance. 

245 await self._lease_group.refresh_loop() 

246 

247 # Normal exit. Run detach while the task group is still active and the backend 

248 # is still live. 

249 ran_detach = True 

250 await self._run_detach() 

251 

252 # Force all tasks to shut down. 

253 raise asyncio.CancelledError("Service shutdown") 

254 except asyncio.CancelledError: 

255 if self._shutdown_called: 

256 logger.debug(f"Service {self.info.name} normal shutdown") 

257 self._shutdown.set_result(True) 

258 else: 

259 logger.error(f"Service {self.info.name} abnormal shutdown (cancelled)") 

260 self._shutdown.cancel() 

261 

262 raise 

263 except BaseExceptionGroup as eg: 

264 logger.error(f"Service {self.info.name} abnormal shutdown") 

265 errors = [e for e in eg.exceptions if not isinstance(e, asyncio.CancelledError)] 

266 

267 for i, e in enumerate(errors, 1): 

268 log = logger.opt(exception=e) if not isinstance(e, GeneratorExit) else logger 

269 log.debug(f"{e.__class__.__name__} (error {i} of {len(errors)})") 

270 

271 if errors: 

272 self._shutdown.set_exception(BaseExceptionGroup("ServiceContext errors", errors)) 

273 else: 

274 self._shutdown.cancel() 

275 finally: 

276 # Shut down request listeners. 

277 with contextlib.suppress(Exception): 

278 async with asyncio.timeout(5.0): 

279 await self.backend.shutdown_request_listeners() 

280 

281 # Ensure all impls are detached on every shutdown path. 

282 if not ran_detach: 

283 await self._run_detach() 

284 

285 async def _run_detach(self): 

286 impls = list(self._impls) 

287 

288 # Detach all impls in parallel. 

289 results = await asyncio.gather( 

290 *(impl.detach() for impl in impls), 

291 return_exceptions=True, 

292 ) 

293 

294 for impl, result in zip(impls, results, strict=True): 

295 match result: 

296 case Exception(): 

297 logger.warning(f"Error during {impl.entity} detach ({type(result).__name__})") 

298 case asyncio.CancelledError(): 

299 pass 

300 case BaseException(): 

301 # Propagate nuclear exceptions (SystemExit / KeyboardInterrupt). 

302 raise result 

303 

304 async def shutdown(self): 

305 """Shut down this service context, making it no longer usable.""" 

306 self._shutdown_called = True 

307 

308 # Expire our leases to trigger shutdown. 

309 try: 

310 await self._lease_group.expire() 

311 except* BackendError: 

312 # If the backend is unavailable, the task group is in the process of tearing down 

313 # anyway, so we continue to await the shutdown future. 

314 pass 

315 

316 # The caller has requested an intentional shutdown, so we won't bother propagating an error 

317 # if the service context happened to simultaneously blow up. 

318 with contextlib.suppress(asyncio.CancelledError, Exception): 

319 await self.join() 

320 

321 async def join(self): 

322 """Wait until the service shuts down. 

323 

324 If the service exited abnormally, the exception that caused it to exit will be raised. 

325 """ 

326 with contextlib.suppress(asyncio.CancelledError): 

327 # The task can only raise CancelledError (and nuclear BaseExceptions) due to exception 

328 # handling in _service_task above. 

329 await self._task 

330 

331 await self._shutdown 

332 

333 async def register_entity(self, entity: str): 

334 """Register a generic entity implementation on the backend.""" 

335 return await self.register_impl(EntityImpl.for_service_context(self, entity)) 

336 

337 async def register_device(self, entity: str): 

338 """Register a Device implementation on the backend.""" 

339 return await self.register_impl(DeviceImpl.for_service_context(self, entity)) 

340 

341 async def register_controller(self, entity: str): 

342 """Register a Controller implementation on the backend.""" 

343 return await self.register_impl(ControllerImpl.for_service_context(self, entity)) 

344 

345 async def register_program(self, entity: str): 

346 """Register a Program implementation on the backend.""" 

347 return await self.register_impl(ProgramImpl.for_service_context(self, entity))