Coverage for core / src / sensorkit / api / entrypoint.py: 85%

1# SPDX-License-Identifier: Apache-2.0 

2import asyncio 

3import concurrent.futures 

4import contextlib 

5import os 

6import signal 

7import time 

8from collections.abc import Callable, Coroutine, Mapping 

9from random import random 

10from typing import Any 

11 

12from loguru import logger 

13 

14from sensorkit.api.declarative import Service 

15from sensorkit.common.aio import cleanup_future, scoped_waiter 

16from sensorkit.common.importutil import obj_from_spec 

17 

18type ServiceEntrypointFunc = Callable[[Service], Coroutine[Any, Any, None]] 

19 

20FORCE_QUIT_INTERRUPTS = 3 

21 

22 

23class ServiceEntrypoint: 

24 """Proxy for a service entrypoint function.""" 

25 

26 @classmethod 

27 def from_spec(cls, spec: str, load_file: bool = False): 

28 """Finds a service entrypoint in a module or python file.""" 

29 return obj_from_spec(spec=spec, base=ServiceEntrypoint, load_file=load_file) 

30 

31 def __init__(self, func: ServiceEntrypointFunc, version: str): 

32 self._func = func 

33 self._version = version 

34 

35 async def run(self, name: str): 

36 """Start the service with *name* and return `(service, task)` once it is running.""" 

37 # Invoke the user entrypoint function passing in a new Service instance. 

38 service = Service(name, self._version) 

39 task = asyncio.create_task(self._func(service)) 

40 

41 try: 

42 # Wait for the service to start up or for the task to fail. 

43 await asyncio.wait( 

44 [service.running, task], 

45 return_when=asyncio.FIRST_COMPLETED, 

46 ) 

47 

48 if not service.running.done(): 

49 if e := task.exception(): 

50 logger.debug(f"Service '{name}' failed to start: {type(e).__name__}: {e}") 

51 raise e 

52 else: 

53 raise RuntimeError(f"Entrypoint '{self._func.__name__}' did not start the service!") 

54 

55 # If the service raised, propagate the exception. 

56 await service.running 

57 

58 return service, task 

59 except asyncio.CancelledError: 

60 task.cancel() 

61 

62 try: 

63 await task 

64 except asyncio.CancelledError: 

65 pass 

66 except Exception as e: 

67 logger.opt(exception=e).debug("exception during service shutdown") 

68 

69 for fut in (service.running, service.shutdown): 

70 cleanup_future(fut) 

71 

72 raise 

73 

74 def __call__(self, *args, **kwargs): 

75 raise RuntimeError("Service entrypoint cannot be called directly") 

76 

77 

78def service_entrypoint(*, version: str): 

79 """Returns a decorator for defining a service entrypoint.""" 

80 def decorator(func: ServiceEntrypointFunc): 

81 return ServiceEntrypoint(func, version=version) 

82 

83 return decorator 

84 

85 

86class ShutdownSignal(Exception): 

87 """Raised to signal an orderly service shutdown across threads.""" 

88 

89 

90async def run_services( 

91 entrypoints: Mapping[str, ServiceEntrypoint], 

92 trap_signals: bool = True, 

93 **kwargs, 

94): 

95 """Run a set of service entrypoints.""" 

96 interrupt_count = 0 

97 

98 # We have to use a `concurrent.futures.Future` here so it can be used by different event loops 

99 # running on different threads. 

100 shutdown = concurrent.futures.Future() 

101 

102 def signal_shutdown(sig=None, _frame=None): 

103 nonlocal interrupt_count 

104 

105 match sig: 

106 case signal.SIGTERM: 

107 logger.info("Service process received shutdown signal") 

108 case signal.SIGINT: 

109 interrupt_count += 1 

110 remaining = FORCE_QUIT_INTERRUPTS - interrupt_count 

111 

112 if interrupt_count == 1: 

113 logger.warning("Service process received interrupt signal") 

114 elif remaining > 0: 

115 times = "one more time" if remaining == 1 else f"{remaining} more times" 

116 logger.warning(f"Interrupt {times} to force exit") 

117 else: 

118 logger.error("Forced exit") 

119 os._exit(130) 

120 

121 if not shutdown.done(): 

122 shutdown.set_exception(ShutdownSignal()) 

123 

124 if trap_signals: 

125 signal.signal(signal.SIGINT, signal_shutdown) 

126 signal.signal(signal.SIGTERM, signal_shutdown) 

127 

128 # Run all services, each in its own thread with its own event loop per invocation. 

129 try: 

130 await asyncio.gather( 

131 asyncio.shield(asyncio.wrap_future(shutdown)), 

132 *( 

133 asyncio.to_thread(_service_loop, name, entrypoint, shutdown, **kwargs) 

134 for name, entrypoint in entrypoints.items() 

135 ) 

136 ) 

137 except Exception: 

138 # Shut down other services. 

139 signal_shutdown() 

140 raise 

141 

142 

143def _restart_loop( 

144 wait_fn: Callable[[float], bool], 

145 *, 

146 max_restarts: int | None = None, 

147 startup_failure_threshold: float = 30.0, 

148 restart_backoff_min: float = 5.0, 

149 restart_backoff_max: float = 300.0, 

150 restart_backoff_random: float = 5.0, 

151 restart_backoff_factor: float = 1.8, 

152): 

153 """Restart loop generator. Yields once per iteration; caller runs the service body. 

154 

155 wait_fn(seconds) must return True to continue or False for graceful shutdown, and 

156 raise to signal an abnormal stop. 

157 """ 

158 restarts = 0 

159 backoff = 0.0 

160 

161 while max_restarts is None or restarts <= max_restarts: 

162 if not wait_fn(backoff): 

163 return 

164 

165 start = time.monotonic() 

166 yield 

167 elapsed = time.monotonic() - start 

168 

169 if elapsed < startup_failure_threshold: 

170 restarts += 1 

171 backoff = backoff * restart_backoff_factor + random() * restart_backoff_random 

172 backoff = min(max(backoff, restart_backoff_min), restart_backoff_max) 

173 else: 

174 restarts = 1 

175 backoff = 0.0 

176 

177 

178def _service_loop( 

179 name: str, 

180 entrypoint: ServiceEntrypoint, 

181 shutdown_signal: concurrent.futures.Future, 

182 startup_timeout: float = 300.0, 

183 shutdown_timeout: float = 5.0, 

184 **kwargs, 

185): 

186 last_err: Exception | None = None 

187 

188 def _wait_for_shutdown(timeout: float) -> bool: 

189 if timeout > 0.0 and not shutdown_signal.done(): 

190 logger.info(f"Waiting {timeout:1.0f} sec before restarting {name} ...") 

191 try: 

192 shutdown_signal.result(timeout=timeout) 

193 return False 

194 except concurrent.futures.TimeoutError: 

195 return True 

196 

197 try: 

198 for _ in _restart_loop(_wait_for_shutdown, **kwargs): 

199 logger.info(f"Starting {name}") 

200 last_err = None 

201 

202 try: 

203 asyncio.run( 

204 _service_proc( 

205 name, entrypoint, shutdown_signal, startup_timeout, shutdown_timeout 

206 ) 

207 ) 

208 except Exception as e: 

209 logger.error(f"Service {name} exited due to {type(e).__name__}") 

210 last_err = e 

211 except ShutdownSignal as e: 

212 last_err = e 

213 except Exception as e: 

214 logger.error(f"Terminating {name} service loop due to {type(e).__name__}") 

215 last_err = e 

216 finally: 

217 logger.debug(f"service loop for {name} exiting (last_err is {type(last_err).__name__}))") 

218 

219 # Make sure all other services exit, too. 

220 if not shutdown_signal.done(): 

221 shutdown_signal.set_result(None) 

222 

223 if last_err: 

224 raise last_err 

225 

226 

227async def _service_proc( 

228 name: str, 

229 entrypoint: ServiceEntrypoint, 

230 shutdown_signal: concurrent.futures.Future, 

231 startup_timeout: float, 

232 shutdown_timeout: float, 

233): 

234 service: Service | None = None 

235 service_task: asyncio.Task | None = None 

236 

237 # Bind the cross-thread shutdown signal to this event loop. 

238 shutdown = asyncio.wrap_future(shutdown_signal) 

239 

240 try: 

241 async with scoped_waiter(asyncio.shield(shutdown)) as shutdown_wait: 

242 startup = asyncio.create_task(entrypoint.run(name)) 

243 

244 await asyncio.wait( 

245 [startup, shutdown_wait], 

246 return_when=asyncio.FIRST_COMPLETED, 

247 timeout=startup_timeout, 

248 ) 

249 

250 if shutdown_wait.done(): 

251 startup.cancel() 

252 

253 with contextlib.suppress(asyncio.CancelledError): 

254 await startup 

255 

256 return 

257 

258 if not startup.done(): 

259 raise TimeoutError(f"Service failed to start within {startup_timeout:.0f} seconds") 

260 

261 service, service_task = startup.result() 

262 

263 await asyncio.wait( 

264 [service_task, shutdown_wait], 

265 return_when=asyncio.FIRST_COMPLETED, 

266 ) 

267 

268 if not shutdown_wait.done(): 

269 await service_task 

270 finally: 

271 if service and not service.shutdown.done(): 

272 logger.info(f"Shutting down {name}") 

273 

274 # Try an orderly shutdown first. 

275 with contextlib.suppress(Exception): 

276 async with asyncio.timeout(shutdown_timeout): 

277 await service.stop() 

278 

279 # Kill the user entrypoint task. 

280 if service_task and not service_task.done(): 

281 service_task.cancel() 

282 

283 with contextlib.suppress(asyncio.CancelledError): 

284 await service_task