Skip to content

Commit a8a50ec

Browse files
committed
Merge branch 'install_conda_envs' into fix_torch_load
# Conflicts: # bioimageio/core/_resource_tests.py
2 parents 7f6fdf1 + 376507f commit a8a50ec

File tree

5 files changed

+255
-29
lines changed

5 files changed

+255
-29
lines changed

README.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -375,6 +375,11 @@ The model specification and its validation tools can be found at <https://github
375375

376376
## Changelog
377377

378+
### 0.7.1 (to be released)
379+
380+
- New test function `bioimageio.core.test_description_in_conda_env` that uses conda
381+
in subprocesses to test a resource in a dedicated conda environment.
382+
378383
### 0.7.0
379384

380385
- breaking:

bioimageio/core/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@
104104
"Stat",
105105
"tensor",
106106
"Tensor",
107+
"test_description_in_conda_env",
107108
"test_description",
108109
"test_model",
109110
"test_resource",

bioimageio/core/_resource_tests.py

Lines changed: 217 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,42 @@
1+
import hashlib
2+
import platform
3+
import subprocess
14
import traceback
25
import warnings
6+
from io import StringIO
37
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+
)
522

623
import numpy as np
724
from loguru import logger
25+
from typing_extensions import assert_never, get_args
826

927
from bioimageio.spec import (
28+
BioimageioCondaEnv,
1029
InvalidDescr,
1130
ResourceDescr,
1231
build_description,
1332
dump_description,
33+
get_conda_env,
1434
load_description,
35+
save_bioimageio_package,
1536
)
1637
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
1940
from bioimageio.spec.model import v0_4, v0_5
2041
from bioimageio.spec.model.v0_5 import WeightsFormat
2142
from bioimageio.spec.summary import (
@@ -81,11 +102,11 @@ def enable_determinism(mode: Literal["seed_only", "full"]):
81102

82103
try:
83104
try:
84-
import tensorflow as tf # pyright: ignore[reportMissingImports]
105+
import tensorflow as tf
85106
except ImportError:
86107
pass
87108
else:
88-
tf.random.seed(0)
109+
tf.random.set_seed(0)
89110
if mode == "full":
90111
tf.config.experimental.enable_op_determinism()
91112
# TODO: find possibility to switch it off again??
@@ -94,7 +115,7 @@ def enable_determinism(mode: Literal["seed_only", "full"]):
94115

95116

96117
def test_model(
97-
source: Union[v0_5.ModelDescr, PermissiveFileSource],
118+
source: Union[v0_4.ModelDescr, v0_5.ModelDescr, PermissiveFileSource],
98119
weight_format: Optional[WeightsFormat] = None,
99120
devices: Optional[List[str]] = None,
100121
absolute_tolerance: float = 1.5e-4,
@@ -118,6 +139,11 @@ def test_model(
118139
)
119140

120141

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+
121147
def test_description(
122148
source: Union[ResourceDescr, PermissiveFileSource, BioimageioYamlContent],
123149
*,
@@ -130,21 +156,194 @@ def test_description(
130156
determinism: Literal["seed_only", "full"] = "seed_only",
131157
expected_type: Optional[str] = None,
132158
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,
133163
) -> 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,
145195
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+
]
146345
)
147-
return rd.validation_summary
346+
return ValidationSummary.model_validate_json(summary_path.read_bytes())
148347

149348

150349
def load_description_and_test(

bioimageio/core/cli.py

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
Dict,
1919
Iterable,
2020
List,
21+
Literal,
2122
Mapping,
2223
Optional,
2324
Sequence,
@@ -113,14 +114,14 @@ def descr_id(self) -> str:
113114

114115

115116
class ValidateFormatCmd(CmdBase, WithSource):
116-
"""validate the meta data format of a bioimageio resource."""
117+
"""Validate the meta data format of a bioimageio resource."""
117118

118119
def run(self):
119120
sys.exit(validate_format(self.descr))
120121

121122

122123
class TestCmd(CmdBase, WithSource):
123-
"""Test a bioimageio resource (beyond meta data formatting)"""
124+
"""Test a bioimageio resource (beyond meta data formatting)."""
124125

125126
weight_format: WeightFormatArgAll = "all"
126127
"""The weight format to limit testing to.
@@ -133,19 +134,36 @@ class TestCmd(CmdBase, WithSource):
133134
decimal: int = 4
134135
"""Precision for numerical comparisons"""
135136

137+
runtime_env: Union[Literal["currently-active", "as-described"], Path] = Field(
138+
"currently-active", alias="runtime-env"
139+
)
140+
"""The python environment to run the tests in
141+
142+
- `"currently-active"`: use active Python interpreter
143+
- `"as-described"`: generate a conda environment YAML file based on the model
144+
weights description.
145+
- A path to a conda environment YAML.
146+
Note: The `bioimageio.core` dependency will be added automatically if not present.
147+
"""
148+
149+
summary_path: Optional[Path] = Field(None, alias="summary-path")
150+
"""Path to save validation summary as JSON file."""
151+
136152
def run(self):
137153
sys.exit(
138154
test(
139155
self.descr,
140156
weight_format=self.weight_format,
141157
devices=self.devices,
142158
decimal=self.decimal,
159+
summary_path=self.summary_path,
160+
runtime_env=self.runtime_env,
143161
)
144162
)
145163

146164

147165
class PackageCmd(CmdBase, WithSource):
148-
"""save a resource's metadata with its associated files."""
166+
"""Save a resource's metadata with its associated files."""
149167

150168
path: CliPositionalArg[Path]
151169
"""The path to write the (zipped) package to.
@@ -551,10 +569,10 @@ def input_dataset(stat: Stat):
551569

552570
class Bioimageio(
553571
BaseSettings,
572+
cli_implicit_flags=True,
554573
cli_parse_args=True,
555574
cli_prog_name="bioimageio",
556575
cli_use_class_docs_for_groups=True,
557-
cli_implicit_flags=True,
558576
use_attribute_docstrings=True,
559577
):
560578
"""bioimageio - CLI for bioimage.io resources 🦒"""

0 commit comments

Comments
 (0)