Coverage for src/bioimageio/core/_model_adapter.py: 88%

1import gc 

2import warnings 

3from abc import ABC, abstractmethod 

4from queue import LifoQueue 

5from typing import ( 

6 Any, 

7 Dict, 

8 Generic, 

9 Iterable, 

10 List, 

11 Optional, 

12 Sequence, 

13 Tuple, 

14 Union, 

15) 

16 

17from exceptiongroup import ExceptionGroup 

18from loguru import logger 

19from numpy.typing import NDArray 

20from typing_extensions import TypeVar 

21 

22from bioimageio.spec import ValidationSummary 

23from bioimageio.spec.model import AnyModelDescr, v0_4 

24 

25from ._restore_batch_multi_index import restore_batch_multi_index 

26from ._sample_serializer import SampleSerializer, SerializedSampleBlockType 

27from .common import PerMember 

28from .digest_spec import get_axes_infos, get_member_ids 

29from .sample import Sample 

30from .tensor import Tensor 

31 

32 

33class ModelAdapter(ABC): 

34 """ 

35 Represents model *without* any preprocessing or postprocessing. 

36 

37 ``` 

38 from bioimageio.core import load_description 

39 

40 model = load_description(...) 

41 

42 # option 1: 

43 adapter = create_model_adapter(model) 

44 adapter.forward(...) 

45 adapter.unload() 

46 

47 # option 2: 

48 with create_model_adapter(model) as adapter: 

49 adapter.forward(...) 

50 ``` 

51 """ 

52 

53 def __init__( 

54 self, model_description: AnyModelDescr, devices: Optional[Sequence[str]] 

55 ): 

56 super().__init__() 

57 self._model_descr = model_description 

58 self._input_ids = get_member_ids(model_description.inputs) 

59 self._output_ids = get_member_ids(model_description.outputs) 

60 self._input_axes = [ 

61 tuple(a.id for a in get_axes_infos(t)) for t in model_description.inputs 

62 ] 

63 self._output_axes = [ 

64 tuple(a.id for a in get_axes_infos(t)) for t in model_description.outputs 

65 ] 

66 if isinstance(model_description, v0_4.ModelDescr): 

67 self._input_is_optional = [False] * len(model_description.inputs) 

68 else: 

69 self._input_is_optional = [ipt.optional for ipt in model_description.inputs] 

70 

71 self._devices = devices 

72 self.load() 

73 

74 @property 

75 def model_descr(self) -> AnyModelDescr: 

76 return self._model_descr 

77 

78 @abstractmethod 

79 def load(self) -> None: 

80 self._loaded = True 

81 

82 @abstractmethod 

83 def forward( 

84 self, inputs: PerMember[Optional[Tensor]] 

85 ) -> PerMember[Optional[Tensor]]: ... 

86 

87 @abstractmethod 

88 def unload(self): 

89 """Unload model from any devices, freeing their memory. 

90 

91 Note: 

92 The moder adapter should be considered unusable afterwards. 

93 """ 

94 self._loaded = False 

95 

96 def close(self): 

97 """Close the model adapter, freeing any resources. 

98 

99 Note: 

100 The moder adapter should be considered unusable afterwards. 

101 """ 

102 self.unload() 

103 

104 

105DeviceType = TypeVar("DeviceType") 

106ModelType = TypeVar("ModelType") 

107 

108 

109class LocalModelAdapter(ModelAdapter, ABC, Generic[DeviceType, ModelType]): 

110 def load(self) -> None: 

111 devices = self._devices 

112 self._model_queue: LifoQueue[Tuple[DeviceType, ModelType]] = LifoQueue() 

113 parsed_devices = self._parse_devices(devices) 

114 assert parsed_devices 

115 # prioritize devices by order specified by user 

116 device_exceptions: Dict[str, Exception] = {} 

117 self._initialized_devices: List[str] = [] 

118 for d in parsed_devices[::-1]: 

119 try: 

120 model = self._init_model_on_device(d) 

121 except Exception as e: 

122 device_exceptions[str(d)] = e 

123 else: 

124 self._model_queue.put((d, model)) 

125 self._initialized_devices.insert(0, str(d)) 

126 

127 if self._model_queue.empty(): 

128 if len(device_exceptions) == 1: 

129 raise next(iter(device_exceptions.values())) 

130 else: 

131 raise ExceptionGroup( 

132 "Failed to initialize model on any of the requested devices.", 

133 list(device_exceptions.values())[::-1], 

134 ) 

135 

136 if device_exceptions: 

137 logger.warning( 

138 "Failed to initialize model on some of the requested devices. Successfully initialized on {}, but got the following errors for other devices: {}", 

139 self._initialized_devices, 

140 device_exceptions, 

141 ) 

142 

143 super().load() 

144 

145 @abstractmethod 

146 def _parse_devices(self, devices: Optional[Sequence[str]]) -> Sequence[DeviceType]: 

147 """Parse devices 

148 

149 Note: 

150 - May not return an empty sequence. 

151 - The order of devices in the returned sequence determines the priority of device usage in the forward pass. 

152 First devices has highgest priority, last device has lowest priority. 

153 """ 

154 

155 @abstractmethod 

156 def _init_model_on_device(self, device: DeviceType) -> ModelType: ... 

157 

158 def forward( 

159 self, inputs: PerMember[Optional[Tensor]] 

160 ) -> PerMember[Optional[Tensor]]: 

161 """ 

162 Run forward pass of model to get model predictions 

163 

164 Note: sample id and stample stat attributes are passed through 

165 

166 Args: 

167 inputs: input tensors for the model, keyed by member id 

168 """ 

169 if not self._loaded: 

170 raise RuntimeError("Model must be `.load()`ed before calling forward()") 

171 

172 unexpected = [mid for mid in inputs if mid not in self._input_ids] 

173 if unexpected: 

174 warnings.warn(f"Got unexpected input tensor IDs: {unexpected}") 

175 

176 input_arrays = [ 

177 ( 

178 None 

179 if (a := inputs.get(in_id)) is None 

180 else a.transpose(in_order).to_numpy() 

181 ) 

182 for in_id, in_order in zip(self._input_ids, self._input_axes) 

183 ] 

184 

185 logger.debug( 

186 "NN input shapes: {}", 

187 [a.shape if a is not None else None for a in input_arrays], 

188 ) 

189 device, model = self._model_queue.get() 

190 try: 

191 output_arrays = self._forward_impl(device, model, input_arrays) 

192 finally: 

193 self._model_queue.put((device, model)) 

194 

195 logger.debug( 

196 "NN output shapes: {}", 

197 [a.shape if a is not None else None for a in output_arrays], 

198 ) 

199 if len(output_arrays) > len(self._output_ids): 

200 warnings.warn( 

201 f"Model produced more outputs ({len(output_arrays)}) than specified in the model description ({len(self._output_ids)}). Extra outputs will be ignored." 

202 ) 

203 output_arrays = output_arrays[: len(self._output_ids)] 

204 

205 output_tensors = [ 

206 None if a is None else Tensor(a, dims=d) 

207 for a, d in zip(output_arrays, self._output_axes) 

208 ] 

209 outputs = { 

210 tid: out 

211 for tid, out in zip( 

212 self._output_ids, 

213 output_tensors, 

214 ) 

215 if out is not None 

216 } 

217 outputs = restore_batch_multi_index(inputs, outputs) 

218 return outputs 

219 

220 @abstractmethod 

221 def _forward_impl( 

222 self, 

223 device: DeviceType, 

224 model: ModelType, 

225 input_arrays: Sequence[Optional[NDArray[Any]]], 

226 ) -> Union[List[Optional[NDArray[Any]]], Tuple[Optional[NDArray[Any]], ...]]: 

227 """framework specific forward implementation""" 

228 

229 def unload(self): 

230 for _ in range(len(self._initialized_devices)): 

231 device, model = self._model_queue.get() 

232 try: 

233 self._cleanup_pre_model_deletion(device, model) 

234 except Exception as e: 

235 logger.warning( 

236 "Got error during pre-deletion cleanup on device {}: {}", device, e 

237 ) 

238 finally: 

239 del model 

240 try: 

241 self._cleanup_post_model_deletion(device) 

242 except Exception as e: 

243 logger.warning( 

244 "Got error during post-deletion cleanup on device {}: {}", device, e 

245 ) 

246 

247 _ = gc.collect() # deallocate memory 

248 super().unload() 

249 

250 @abstractmethod 

251 def _cleanup_pre_model_deletion(self, device: DeviceType, model: ModelType) -> None: 

252 """Clean up before model reference deletion""" 

253 

254 @abstractmethod 

255 def _cleanup_post_model_deletion(self, device: DeviceType) -> None: 

256 """Clean up after model reference deletion""" 

257 

258 

259class RemoteModelAdapter(ModelAdapter, ABC, Generic[SerializedSampleBlockType]): 

260 """Model adapter to use a remote service for model inference.""" 

261 

262 def __init__( 

263 self, 

264 model_description: AnyModelDescr, 

265 server: str, 

266 sample_serializer: SampleSerializer[SerializedSampleBlockType], 

267 ): 

268 super().__init__(model_description, devices=None) 

269 self._server = server 

270 self._serializer = sample_serializer 

271 

272 @property 

273 def server(self) -> str: 

274 return self._server 

275 

276 def forward( 

277 self, inputs: PerMember[Optional[Tensor]] 

278 ) -> PerMember[Optional[Tensor]]: 

279 serialized_input = self._serializer.serialize_sample( 

280 Sample( 

281 members={k: v for k, v in inputs.items() if v is not None}, 

282 stat={}, 

283 id=None, 

284 ) 

285 ) 

286 serialized_output = self._forward_impl(serialized_input) 

287 output = self._serializer.deserialize_sample(serialized_output).members 

288 output = restore_batch_multi_index(inputs, output) 

289 return output 

290 

291 @abstractmethod 

292 def _forward_impl( 

293 self, serialized_input_sample: Iterable[SerializedSampleBlockType] 

294 ) -> Iterable[SerializedSampleBlockType]: ... 

295 

296 @abstractmethod 

297 def test(self) -> Optional[ValidationSummary]: 

298 """Run the bioimageio model test."""