improve test error handling

Author: FynnBe

bioimageio/core/_resource_tests.py

@@ -2,7 +2,6 @@
 import os
 import platform
 import subprocess
-import traceback
 import warnings
 from io import StringIO
 from itertools import product
@@ -47,7 +46,10 @@
     MismatchedElementsPerMillion,
     RelativeTolerance,
 )
-from bioimageio.spec._internal.validation_context import validation_context_var
+from bioimageio.spec._internal.validation_context import (
+    get_validation_context,
+    validation_context_var,
+)
 from bioimageio.spec.common import BioimageioYamlContent, PermissiveFileSource, Sha256
 from bioimageio.spec.model import v0_4, v0_5
 from bioimageio.spec.model.v0_5 import WeightsFormat
@@ -589,8 +591,17 @@ def _test_model_inference(
 ) -> None:
     test_name = f"Reproduce test outputs from test inputs ({weight_format})"
     logger.debug("starting '{}'", test_name)
-    error: Optional[str] = None
-    tb: List[str] = []
+    errors: List[ErrorEntry] = []
+
+    def add_error_entry(msg: str, with_traceback: bool = False):
+        errors.append(
+            ErrorEntry(
+                loc=("weights", weight_format),
+                msg=msg,
+                type="bioimageio.core",
+                with_traceback=with_traceback,
+            )
+        )
 
     try:
         inputs = get_test_inputs(model)
@@ -602,13 +613,15 @@ def _test_model_inference(
             results = prediction_pipeline.predict_sample_without_blocking(inputs)
 
         if len(results.members) != len(expected.members):
-            error = f"Expected {len(expected.members)} outputs, but got {len(results.members)}"
+            add_error_entry(
+                f"Expected {len(expected.members)} outputs, but got {len(results.members)}"
+            )
 
         else:
             for m, expected in expected.members.items():
                 actual = results.members.get(m)
                 if actual is None:
-                    error = "Output tensors for test case may not be None"
+                    add_error_entry("Output tensors for test case may not be None")
                     break
 
                 rtol, atol, mismatched_tol = _get_tolerance(
@@ -627,7 +640,7 @@ def _test_model_inference(
                     a_max = abs_diff[a_max_idx].item()
                     a_actual = actual[a_max_idx].item()
                     a_expected = expected[a_max_idx].item()
-                    error = (
+                    add_error_entry(
                         f"Output '{m}' disagrees with {mismatched_elements} of"
                         + f" {expected.size} expected values."
                         + f"\n Max relative difference: {r_max:.2e}"
@@ -638,30 +651,18 @@ def _test_model_inference(
                     )
                     break
     except Exception as e:
-        if validation_context_var.get().raise_errors:
+        if get_validation_context().raise_errors:
             raise e
 
-        error = str(e)
-        tb = traceback.format_exception(type(e), e, e.__traceback__, chain=True)
+        add_error_entry(str(e), with_traceback=True)
 
     model.validation_summary.add_detail(
         ValidationDetail(
             name=test_name,
             loc=("weights", weight_format),
-            status="passed" if error is None else "failed",
+            status="failed" if errors else "passed",
             recommended_env=get_conda_env(entry=dict(model.weights)[weight_format]),
-            errors=(
-                []
-                if error is None
-                else [
-                    ErrorEntry(
-                        loc=("weights", weight_format),
-                        msg=error,
-                        type="bioimageio.core",
-                        traceback=tb,
-                    )
-                ]
-            ),
+            errors=errors,
         )
     )
 
@@ -816,11 +817,9 @@ def get_ns(n: int):
                 if stop_early and error is not None:
                     break
     except Exception as e:
-        if validation_context_var.get().raise_errors:
+        if get_validation_context().raise_errors:
             raise e
 
-        error = str(e)
-        tb = traceback.format_tb(e.__traceback__)
         model.validation_summary.add_detail(
             ValidationDetail(
                 name=f"Run {weight_format} inference for parametrized inputs",
@@ -829,9 +828,9 @@ def get_ns(n: int):
                 errors=[
                     ErrorEntry(
                         loc=("weights", weight_format),
-                        msg=error,
+                        msg=str(e),
                         type="bioimageio.core",
-                        traceback=tb,
+                        with_traceback=True,
                     )
                 ],
             )
@@ -854,7 +853,7 @@ def _test_expected_resource_type(
                     ErrorEntry(
                         loc=("type",),
                         type="type",
-                        msg=f"expected type {expected_type}, found {rd.type}",
+                        msg=f"Expected type {expected_type}, found {rd.type}",
                     )
                 ]
             ),

bioimageio/core/backends/_model_adapter.py

@@ -1,3 +1,4 @@
+import sys
 import warnings
 from abc import ABC, abstractmethod
 from typing import (
@@ -87,7 +88,7 @@ def create(
             )
 
         weights = model_description.weights
-        errors: List[Tuple[SupportedWeightsFormat, Exception]] = []
+        errors: List[Exception] = []
         weight_format_priority_order = (
             DEFAULT_WEIGHT_FORMAT_PRIORITY_ORDER
             if weight_format_priority_order is None
@@ -112,7 +113,7 @@ def create(
                         model_description=model_description, devices=devices
                     )
                 except Exception as e:
-                    errors.append((wf, e))
+                    errors.append(e)
             elif wf == "tensorflow_saved_model_bundle":
                 assert weights.tensorflow_saved_model_bundle is not None
                 try:
@@ -122,7 +123,7 @@ def create(
                         model_description=model_description, devices=devices
                     )
                 except Exception as e:
-                    errors.append((wf, e))
+                    errors.append(e)
             elif wf == "onnx":
                 assert weights.onnx is not None
                 try:
@@ -132,7 +133,7 @@ def create(
                         model_description=model_description, devices=devices
                     )
                 except Exception as e:
-                    errors.append((wf, e))
+                    errors.append(e)
             elif wf == "torchscript":
                 assert weights.torchscript is not None
                 try:
@@ -142,7 +143,7 @@ def create(
                         model_description=model_description, devices=devices
                     )
                 except Exception as e:
-                    errors.append((wf, e))
+                    errors.append(e)
             elif wf == "keras_hdf5":
                 assert weights.keras_hdf5 is not None
                 # keras can either be installed as a separate package or used as part of tensorflow
@@ -158,27 +159,24 @@ def create(
                         model_description=model_description, devices=devices
                     )
                 except Exception as e:
-                    errors.append((wf, e))
+                    errors.append(e)
             else:
                 assert_never(wf)
 
         assert errors
         if len(weight_format_priority_order) == 1:
             assert len(errors) == 1
-            wf, e = errors[0]
-            raise ValueError(
-                f"The '{wf}' model adapter could not be created"
-                + f" in this environment:\n{e.__class__.__name__}({e}).\n\n"
-            ) from e
+            raise errors[0]
 
         else:
-            error_list = "\n - ".join(
-                f"{wf}: {e.__class__.__name__}({e})" for wf, e in errors
-            )
-            raise ValueError(
+            msg = (
                 "None of the weight format specific model adapters could be created"
-                + f" in this environment. Errors are:\n\n{error_list}.\n\n"
+                + " in this environment."
             )
+            if sys.version_info[:2] >= (3, 11):
+                raise ExceptionGroup(msg, errors)
+            else:
+                raise ValueError(msg) from Exception(errors)
 
     @final
     def load(self, *, devices: Optional[Sequence[str]] = None) -> None:

bioimageio/core/backends/pytorch_backend.py

@@ -97,7 +97,7 @@ def load_torch_model(
     load_state: bool = True,
     devices: Optional[Sequence[Union[str, torch.device]]] = None,
 ) -> nn.Module:
-    arch = import_callable(
+    custom_callable = import_callable(
         weight_spec.architecture,
         sha256=(
             weight_spec.architecture_sha256
@@ -110,27 +110,29 @@ def load_torch_model(
         if isinstance(weight_spec, v0_4.PytorchStateDictWeightsDescr)
         else weight_spec.architecture.kwargs
     )
-    try:
-        # calling custom user code
-        network = arch(**model_kwargs)
-    except Exception as e:
-        raise RuntimeError("Failed to initialize PyTorch model") from e
-
-    if not isinstance(network, nn.Module):
-        raise ValueError(
-            f"calling {weight_spec.architecture.callable_name if isinstance(weight_spec.architecture, (v0_4.CallableFromFile, v0_4.CallableFromDepencency)) else weight_spec.architecture.callable} did not return a torch.nn.Module"
-        )
+    torch_model = custom_callable(**model_kwargs)
+
+    if not isinstance(torch_model, nn.Module):
+        if isinstance(
+            weight_spec.architecture,
+            (v0_4.CallableFromFile, v0_4.CallableFromDepencency),
+        ):
+            callable_name = weight_spec.architecture.callable_name
+        else:
+            callable_name = weight_spec.architecture.callable
+
+        raise ValueError(f"Calling {callable_name} did not return a torch.nn.Module.")
 
     if load_state or devices:
         use_devices = get_devices(devices)
-        network = network.to(use_devices[0])
+        torch_model = torch_model.to(use_devices[0])
         if load_state:
-            network = load_torch_state_dict(
-                network,
+            torch_model = load_torch_state_dict(
+                torch_model,
                 path=download(weight_spec).path,
                 devices=use_devices,
             )
-    return network
+    return torch_model
 
 
 def load_torch_state_dict(

bioimageio/core/digest_spec.py

@@ -111,22 +111,23 @@ def _import_from_file_impl(
             module_spec = importlib.util.spec_from_loader(module_name, loader=None)
             assert module_spec is not None
             module = importlib.util.module_from_spec(module_spec)
-            exec(source_code, module.__dict__)
+            source_compiled = compile(
+                source_code, str(local_source.path), "exec"
+            )  # compile source to attach file name
+            exec(source_compiled, module.__dict__)
             sys.modules[module_spec.name] = module  # cache this module
         except Exception as e:
-            raise ImportError(
-                f"Failed to import {module_name[:-58]}... from {source}"
-            ) from e
+            raise ImportError(f"Failed to import {source} .") from e
 
     try:
         callable_attr = getattr(module, callable_name)
     except AttributeError as e:
         raise AttributeError(
-            f"Imported custom module `{module_name[:-58]}...` has no `{callable_name}` attribute"
+            f"Imported custom module from {source} has no `{callable_name}` attribute."
         ) from e
     except Exception as e:
         raise AttributeError(
-            f"Failed to access `{callable_name}` attribute from imported custom module `{module_name[:-58]}...`"
+            f"Failed to access `{callable_name}` attribute from custom module imported from {source} ."
         ) from e
 
     else: