Coverage for core / src / sensorkit / webapi / serve.py: 94%

1# SPDX-License-Identifier: Apache-2.0 

2import asyncio 

3import collections 

4import contextlib 

5import pathlib 

6from abc import ABC, abstractmethod 

7from collections.abc import AsyncIterator 

8from datetime import UTC, datetime 

9from typing import Any, Literal, NamedTuple, override 

10 

11from loguru import logger 

12from pydantic import BaseModel 

13 

14import sensorkit.api as sk 

15from sensorkit.common.aio import AsyncObserver 

16from sensorkit.common.filewatch import FileEventKind, watch_dir 

17from sensorkit.common.keyword import KeywordDict 

18from sensorkit.data.fits import FITSHeader 

19 

20DEFAULT_CONTROLLER_ID_FIELD = "SKCTRL" 

21 

22# A file is announced as soon as it is created, so the first sighting can precede 

23# its contents: the writer is still streaming, and a Docker bind mount raises 

24# `created` when the file appears rather than when its data lands. Reading then 

25# yields a truncated frame — and a header-only FITS opens *without* error, so the 

26# damage can be a wrong data_size rather than an exception. Re-read until the whole 

27# frame is there; a frame already complete passes first time, so the initial scan 

28# of an existing directory pays nothing. 

29_SETTLE_POLL_S = 0.25 

30_SETTLE_ATTEMPTS = 20 

31 

32type ServeDataConfig = ServeLocalFITSConfig 

33 

34 

35@sk.declare_keyword 

36class ProductInfo(BaseModel): 

37 controller_id: str 

38 product_id: str 

39 register_time: datetime 

40 data_size: int 

41 

42 

43class _CacheEntry(NamedTuple): 

44 info: ProductInfo 

45 path: pathlib.Path 

46 metadata: KeywordDict 

47 

48 

49class ServeHandler(ABC): 

50 """Abstract base for data product serving backends.""" 

51 

52 @abstractmethod 

53 async def get_listing(self) -> list[ProductInfo]: 

54 """Return all known data product records. 

55 

56 This method waits until the initial listing is complete before returning. 

57 """ 

58 

59 @abstractmethod 

60 def watch_listing(self) -> AsyncIterator[ProductInfo]: 

61 """Yield all data product records as they become known.""" 

62 

63 @abstractmethod 

64 def has_product(self, controller_id: str, product_id: str) -> bool: 

65 """Return whether the given product is currently known for the controller. 

66 

67 Reflects only what is known right now; unlike `get_listing`, it does not wait 

68 for the initial listing to complete. 

69 """ 

70 

71 @abstractmethod 

72 def get_metadata(self, controller_id: str, product_id: str) -> KeywordDict: 

73 """Return the metadata dict for the given controller and product. 

74 

75 The `ProductInfo` keyword must always be present in the returned metadata, 

76 whether it was embedded in the persisted metadata or injected by the handler. 

77 """ 

78 

79 @abstractmethod 

80 async def get_data(self, controller_id: str, product_id: str) -> bytes: 

81 """Return the raw bytes for the given controller and product.""" 

82 

83 @abstractmethod 

84 def start_monitor(self, *, task_group: asyncio.TaskGroup): 

85 """Start the background monitoring task.""" 

86 

87 @abstractmethod 

88 async def stop_monitor(self): 

89 """Stop the background monitoring task.""" 

90 

91 

92class ServeLocalFITSConfig(BaseModel): 

93 """Configure serving of FITS files.""" 

94 

95 kind: Literal["local_fits"] = "local_fits" 

96 root_directory: str 

97 controller_id: Literal["from_path", "from_metadata"] = "from_path" 

98 controller_id_field: str | None = None 

99 

100 def create_handler(self): 

101 return ServeLocalFITSHandler(self) 

102 

103 

104class ServeLocalFITSHandler(ServeHandler): 

105 """ServeHandler that watches a local directory for FITS files.""" 

106 

107 def __init__(self, config: ServeLocalFITSConfig): 

108 self.config = config 

109 self._task: asyncio.Task | None = None 

110 self._observer: AsyncObserver[ProductInfo] = AsyncObserver() 

111 self._cache = collections.defaultdict(dict[str, _CacheEntry]) 

112 self._listing_ready = asyncio.Event() 

113 

114 @override 

115 async def get_listing(self): 

116 await self._listing_ready.wait() 

117 return [entry.info for products in self._cache.values() for entry in products.values()] 

118 

119 @override 

120 async def watch_listing(self): 

121 # Subscribe before reading the listing so that updates in between are not missed. 

122 with self._observer.subscription() as queue: 

123 for info in await self.get_listing(): 

124 yield info 

125 

126 while True: 

127 yield await queue.get() 

128 

129 @override 

130 def has_product(self, controller_id: str, product_id: str): 

131 return product_id in cache if (cache := self._cache.get(controller_id)) else False 

132 

133 @override 

134 def get_metadata(self, controller_id: str, product_id: str): 

135 return self._cache[controller_id][product_id].metadata 

136 

137 @override 

138 async def get_data(self, controller_id: str, product_id: str): 

139 path = self._cache[controller_id][product_id].path 

140 raw_bytes = await asyncio.to_thread(path.read_bytes) 

141 return raw_bytes 

142 

143 @override 

144 def start_monitor(self, *, task_group: asyncio.TaskGroup): 

145 self._task = task_group.create_task(self._monitor()) 

146 

147 @override 

148 async def stop_monitor(self): 

149 if self._task is not None: 

150 self._task.cancel() 

151 

152 with contextlib.suppress(asyncio.CancelledError): 

153 await self._task 

154 

155 self._task = None 

156 

157 async def _monitor(self): 

158 root = pathlib.Path(self.config.root_directory).resolve() 

159 

160 while not await asyncio.to_thread(root.exists): 

161 logger.debug(f"waiting for fits server directory {root} to exist...") 

162 await asyncio.sleep(30.0) 

163 

164 async with watch_dir( 

165 root, 

166 recursive=True, 

167 existing=True, 

168 existing_done=self._listing_ready, 

169 kinds=(FileEventKind.CREATED, FileEventKind.MOVED, FileEventKind.EXISTING), 

170 ) as events: 

171 async for event in events: 

172 if event.is_directory or event.path.suffix != ".fits": 

173 continue 

174 await self._found_file(event.path) 

175 

176 async def _read_whole_frame(self, path: pathlib.Path): 

177 """Read *path* once all of it is on disk, or None if it never gets there.""" 

178 from astropy.io import fits 

179 

180 def read_file(): 

181 # The header declares how many bytes the frame occupies, so a 

182 # half-written one is recognisable rather than merely suspected. 

183 # (astropy usually raises on the short read before we compare.) 

184 with fits.open(path) as hdul: 

185 stat = path.stat() 

186 if stat.st_size < sum(hdu.filebytes() for hdu in hdul): 

187 return None 

188 

189 # Metadata comes from the HDU that carries the image. A 

190 # tile-compressed frame leaves an empty stub in the primary HDU 

191 # and holds the image and its cards in an extension, so the 

192 # shape test is what distinguishes the two. Neither is_image nor 

193 # shape touches pixel data, so nothing is decompressed here. 

194 for hdu in hdul: 

195 if hdu.is_image and hdu.shape: 

196 return stat, hdu.header 

197 

198 return stat, hdul[0].header 

199 

200 for _ in range(_SETTLE_ATTEMPTS): 

201 try: 

202 found = await asyncio.to_thread(read_file) 

203 except Exception: 

204 found = None # empty or unparsable: nothing written yet, or garbage 

205 if found is not None: 

206 return found 

207 await asyncio.sleep(_SETTLE_POLL_S) 

208 

209 return None 

210 

211 async def _found_file(self, path: pathlib.Path): 

212 try: 

213 found = await self._read_whole_frame(path) 

214 if found is None: 

215 logger.warning(f"{path} is not a complete FITS frame, skipping") 

216 return 

217 

218 stat, header = found 

219 controller_id: Any = None 

220 

221 # Metadata takes precedence over path-based controller ID resolution. 

222 match self.config.controller_id: 

223 case "from_metadata": 

224 controller_id = header.get( 

225 self.config.controller_id_field or DEFAULT_CONTROLLER_ID_FIELD 

226 ) 

227 case "from_path": 

228 root = pathlib.Path(self.config.root_directory).resolve() 

229 relative_path = path.relative_to(root) 

230 controller_id = ( 

231 relative_path.parts[0] if len(relative_path.parts) > 1 else root.parts[-1] 

232 ) 

233 

234 if controller_id is None or not isinstance(controller_id, str): 

235 logger.warning(f"cannot determine controller for {path}, skipping") 

236 return 

237 

238 # Create ProductInfo and inject into the metadata. The product ID is the name of the 

239 # file. This makes name clashes with other handlers unlikely, even if the stem of the 

240 # filename is the same as another product ID (possibly the same one). 

241 info = ProductInfo( 

242 controller_id=controller_id, 

243 product_id=path.name, 

244 register_time=datetime.fromtimestamp(getattr(stat, "st_birthtime", stat.st_mtime), UTC), 

245 data_size=stat.st_size, 

246 ) 

247 metadata = KeywordDict(info, FITSHeader.from_astropy_header(header)) 

248 

249 # Cache and notify observers. 

250 self._cache[controller_id][info.product_id] = _CacheEntry(info, path, metadata) 

251 self._observer.notify(info) 

252 except Exception: 

253 logger.exception(f"error processing {path}")