1
+ import hashlib
2
+ import platform
3
+ import subprocess
1
4
import traceback
2
5
import warnings
6
+ from io import StringIO
3
7
from itertools import product
4
- from typing import Dict , Hashable , List , Literal , Optional , Sequence , Set , Tuple , Union
8
+ from pathlib import Path
9
+ from tempfile import TemporaryDirectory
10
+ from typing import (
11
+ Callable ,
12
+ Dict ,
13
+ Hashable ,
14
+ List ,
15
+ Literal ,
16
+ Optional ,
17
+ Sequence ,
18
+ Set ,
19
+ Tuple ,
20
+ Union ,
21
+ )
5
22
6
23
import numpy as np
7
24
from loguru import logger
25
+ from typing_extensions import assert_never , get_args
8
26
9
27
from bioimageio .spec import (
28
+ BioimageioCondaEnv ,
10
29
InvalidDescr ,
11
30
ResourceDescr ,
12
31
build_description ,
13
32
dump_description ,
33
+ get_conda_env ,
14
34
load_description ,
35
+ save_bioimageio_package ,
15
36
)
16
37
from bioimageio .spec ._internal .common_nodes import ResourceDescrBase
17
- from bioimageio .spec .common import BioimageioYamlContent , PermissiveFileSource , Sha256
18
- from bioimageio .spec .get_conda_env import get_conda_env
38
+ from bioimageio .spec ._internal . io import is_yaml_value
39
+ from bioimageio .spec ._internal . io_utils import read_yaml , write_yaml
19
40
from bioimageio .spec .model import v0_4 , v0_5
20
41
from bioimageio .spec .model .v0_5 import WeightsFormat
21
42
from bioimageio .spec .summary import (
@@ -81,11 +102,11 @@ def enable_determinism(mode: Literal["seed_only", "full"]):
81
102
82
103
try :
83
104
try :
84
- import tensorflow as tf # pyright: ignore[reportMissingImports]
105
+ import tensorflow as tf
85
106
except ImportError :
86
107
pass
87
108
else :
88
- tf .random .seed (0 )
109
+ tf .random .set_seed (0 )
89
110
if mode == "full" :
90
111
tf .config .experimental .enable_op_determinism ()
91
112
# TODO: find possibility to switch it off again??
@@ -94,7 +115,7 @@ def enable_determinism(mode: Literal["seed_only", "full"]):
94
115
95
116
96
117
def test_model (
97
- source : Union [v0_5 .ModelDescr , PermissiveFileSource ],
118
+ source : Union [v0_4 . ModelDescr , v0_5 .ModelDescr , PermissiveFileSource ],
98
119
weight_format : Optional [WeightsFormat ] = None ,
99
120
devices : Optional [List [str ]] = None ,
100
121
absolute_tolerance : float = 1.5e-4 ,
@@ -118,6 +139,11 @@ def test_model(
118
139
)
119
140
120
141
142
+ def default_run_command (args : Sequence [str ]):
143
+ logger .info ("running '{}'..." , " " .join (args ))
144
+ _ = subprocess .run (args , shell = True , text = True , check = True )
145
+
146
+
121
147
def test_description (
122
148
source : Union [ResourceDescr , PermissiveFileSource , BioimageioYamlContent ],
123
149
* ,
@@ -130,21 +156,194 @@ def test_description(
130
156
determinism : Literal ["seed_only" , "full" ] = "seed_only" ,
131
157
expected_type : Optional [str ] = None ,
132
158
sha256 : Optional [Sha256 ] = None ,
159
+ runtime_env : Union [
160
+ Literal ["currently-active" , "as-described" ], Path , BioimageioCondaEnv
161
+ ] = ("currently-active" ),
162
+ run_command : Callable [[Sequence [str ]], None ] = default_run_command ,
133
163
) -> ValidationSummary :
134
- """Test a bioimage.io resource dynamically, e.g. prediction of test tensors for models"""
135
- rd = load_description_and_test (
136
- source ,
137
- format_version = format_version ,
138
- weight_format = weight_format ,
139
- devices = devices ,
140
- absolute_tolerance = absolute_tolerance ,
141
- relative_tolerance = relative_tolerance ,
142
- decimal = decimal ,
143
- determinism = determinism ,
144
- expected_type = expected_type ,
164
+ """Test a bioimage.io resource dynamically, e.g. prediction of test tensors for models.
165
+
166
+ Args:
167
+ source: model description source.
168
+ weight_format: Weight format to test.
169
+ Default: All weight formats present in **source**.
170
+ devices: Devices to test with, e.g. 'cpu', 'cuda'.
171
+ Default (may be weight format dependent): ['cuda'] if available, ['cpu'] otherwise.
172
+ absolute_tolerance: Maximum absolute tolerance of reproduced output tensors.
173
+ relative_tolerance: Maximum relative tolerance of reproduced output tensors.
174
+ determinism: Modes to improve reproducibility of test outputs.
175
+ runtime_env: (Experimental feature!) The Python environment to run the tests in
176
+ - `"currently-active"`: Use active Python interpreter.
177
+ - `"as-described"`: Use `bioimageio.spec.get_conda_env` to generate a conda
178
+ environment YAML file based on the model weights description.
179
+ - A `BioimageioCondaEnv` or a path to a conda environment YAML file.
180
+ Note: The `bioimageio.core` dependency will be added automatically if not present.
181
+ run_command: (Experimental feature!) Function to execute (conda) terminal commands in a subprocess
182
+ (ignored if **runtime_env** is `"currently-active"`).
183
+ """
184
+ if runtime_env == "currently-active" :
185
+ rd = load_description_and_test (
186
+ source ,
187
+ format_version = format_version ,
188
+ weight_format = weight_format ,
189
+ devices = devices ,
190
+ absolute_tolerance = absolute_tolerance ,
191
+ relative_tolerance = relative_tolerance ,
192
+ decimal = decimal ,
193
+ determinism = determinism ,
194
+ expected_type = expected_type ,
145
195
sha256 = sha256 ,
196
+ )
197
+ return rd .validation_summary
198
+
199
+ if runtime_env == "as-described" :
200
+ conda_env = None
201
+ elif isinstance (runtime_env , (str , Path )):
202
+ conda_env = BioimageioCondaEnv .model_validate (read_yaml (Path (runtime_env )))
203
+ elif isinstance (runtime_env , BioimageioCondaEnv ):
204
+ conda_env = runtime_env
205
+ else :
206
+ assert_never (runtime_env )
207
+
208
+ with TemporaryDirectory (ignore_cleanup_errors = True ) as _d :
209
+ working_dir = Path (_d )
210
+ if isinstance (source , (dict , ResourceDescrBase )):
211
+ file_source = save_bioimageio_package (
212
+ source , output_path = working_dir / "package.zip"
213
+ )
214
+ else :
215
+ file_source = source
216
+
217
+ return _test_in_env (
218
+ file_source ,
219
+ working_dir = working_dir ,
220
+ weight_format = weight_format ,
221
+ conda_env = conda_env ,
222
+ devices = devices ,
223
+ absolute_tolerance = absolute_tolerance ,
224
+ relative_tolerance = relative_tolerance ,
225
+ determinism = determinism ,
226
+ run_command = run_command ,
227
+ )
228
+
229
+
230
+ def _test_in_env (
231
+ source : PermissiveFileSource ,
232
+ * ,
233
+ working_dir : Path ,
234
+ weight_format : Optional [WeightsFormat ],
235
+ conda_env : Optional [BioimageioCondaEnv ],
236
+ devices : Optional [Sequence [str ]],
237
+ absolute_tolerance : float ,
238
+ relative_tolerance : float ,
239
+ determinism : Literal ["seed_only" , "full" ],
240
+ run_command : Callable [[Sequence [str ]], None ],
241
+ ) -> ValidationSummary :
242
+ descr = load_description (source )
243
+
244
+ if not isinstance (descr , (v0_4 .ModelDescr , v0_5 .ModelDescr )):
245
+ raise NotImplementedError ("Not yet implemented for non-model resources" )
246
+
247
+ if weight_format is None :
248
+ all_present_wfs = [
249
+ wf for wf in get_args (WeightsFormat ) if getattr (descr .weights , wf )
250
+ ]
251
+ ignore_wfs = [wf for wf in all_present_wfs if wf in ["tensorflow_js" ]]
252
+ logger .info (
253
+ "Found weight formats {}. Start testing all{}..." ,
254
+ all_present_wfs ,
255
+ f" (except: { ', ' .join (ignore_wfs )} ) " if ignore_wfs else "" ,
256
+ )
257
+ summary = _test_in_env (
258
+ source ,
259
+ working_dir = working_dir / all_present_wfs [0 ],
260
+ weight_format = all_present_wfs [0 ],
261
+ devices = devices ,
262
+ absolute_tolerance = absolute_tolerance ,
263
+ relative_tolerance = relative_tolerance ,
264
+ determinism = determinism ,
265
+ conda_env = conda_env ,
266
+ run_command = run_command ,
267
+ )
268
+ for wf in all_present_wfs [1 :]:
269
+ additional_summary = _test_in_env (
270
+ source ,
271
+ working_dir = working_dir / wf ,
272
+ weight_format = wf ,
273
+ devices = devices ,
274
+ absolute_tolerance = absolute_tolerance ,
275
+ relative_tolerance = relative_tolerance ,
276
+ determinism = determinism ,
277
+ conda_env = conda_env ,
278
+ run_command = run_command ,
279
+ )
280
+ for d in additional_summary .details :
281
+ # TODO: filter reduntant details; group details
282
+ summary .add_detail (d )
283
+ return summary
284
+
285
+ if weight_format == "pytorch_state_dict" :
286
+ wf = descr .weights .pytorch_state_dict
287
+ elif weight_format == "torchscript" :
288
+ wf = descr .weights .torchscript
289
+ elif weight_format == "keras_hdf5" :
290
+ wf = descr .weights .keras_hdf5
291
+ elif weight_format == "onnx" :
292
+ wf = descr .weights .onnx
293
+ elif weight_format == "tensorflow_saved_model_bundle" :
294
+ wf = descr .weights .tensorflow_saved_model_bundle
295
+ elif weight_format == "tensorflow_js" :
296
+ raise RuntimeError (
297
+ "testing 'tensorflow_js' is not supported by bioimageio.core"
298
+ )
299
+ else :
300
+ assert_never (weight_format )
301
+
302
+ assert wf is not None
303
+ if conda_env is None :
304
+ conda_env = get_conda_env (entry = wf )
305
+
306
+ # remove name as we crate a name based on the env description hash value
307
+ conda_env .name = None
308
+
309
+ dumped_env = conda_env .model_dump (mode = "json" , exclude_none = True )
310
+ if not is_yaml_value (dumped_env ):
311
+ raise ValueError (f"Failed to dump conda env to valid YAML { conda_env } " )
312
+
313
+ env_io = StringIO ()
314
+ write_yaml (dumped_env , file = env_io )
315
+ encoded_env = env_io .getvalue ().encode ()
316
+ env_name = hashlib .sha256 (encoded_env ).hexdigest ()
317
+
318
+ try :
319
+ run_command (["where" if platform .system () == "Windows" else "which" , "conda" ])
320
+ except Exception as e :
321
+ raise RuntimeError ("Conda not available" ) from e
322
+
323
+ working_dir .mkdir (parents = True , exist_ok = True )
324
+ try :
325
+ run_command (["conda" , "activate" , env_name ])
326
+ except Exception :
327
+ path = working_dir / "env.yaml"
328
+ _ = path .write_bytes (encoded_env )
329
+ logger .debug ("written conda env to {}" , path )
330
+ run_command (["conda" , "env" , "create" , f"--file={ path } " , f"--name={ env_name } " ])
331
+ run_command (["conda" , "activate" , env_name ])
332
+
333
+ summary_path = working_dir / "summary.json"
334
+ run_command (
335
+ [
336
+ "conda" ,
337
+ "run" ,
338
+ "-n" ,
339
+ env_name ,
340
+ "bioimageio" ,
341
+ "test" ,
342
+ str (source ),
343
+ f"--summary-path={ summary_path } " ,
344
+ ]
146
345
)
147
- return rd . validation_summary
346
+ return ValidationSummary . model_validate_json ( summary_path . read_bytes ())
148
347
149
348
150
349
def load_description_and_test (
0 commit comments