Lazily import onnx so a broken onnx does not block 'import monai' (#8455)

monai/bundle/scripts.py

@@ -57,7 +57,6 @@
 ValidationError, _ = optional_import("jsonschema.exceptions", name="ValidationError")
 Checkpoint, has_ignite = optional_import("ignite.handlers", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Checkpoint")
 requests, has_requests = optional_import("requests")
-onnx, _ = optional_import("onnx")
 huggingface_hub, _ = optional_import("huggingface_hub")
 
 logger = get_logger(module_name=__name__)
@@ -1419,6 +1418,7 @@ def onnx_export(
     converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_})
 
     def save_onnx(onnx_obj: Any, filename_prefix_or_stream: str, **kwargs: Any) -> None:
+        onnx, _ = optional_import("onnx")
         onnx.save(onnx_obj, filename_prefix_or_stream)
 
     _export(

monai/networks/utils.py

@@ -34,9 +34,6 @@
 from monai.utils.module import look_up_option, optional_import
 from monai.utils.type_conversion import convert_to_dst_type, convert_to_tensor
 
-onnx, _ = optional_import("onnx")
-onnxreference, _ = optional_import("onnx.reference")
-onnxruntime, _ = optional_import("onnxruntime")
 polygraphy, polygraphy_imported = optional_import("polygraphy")
 torch_tensorrt, _ = optional_import("torch_tensorrt", "1.4.0")
 
@@ -708,6 +705,8 @@ def convert_to_onnx(
             https://pytorch.org/docs/master/generated/torch.jit.script.html.
 
     """
+    onnx, _ = optional_import("onnx")
+
     model.eval()
     with torch.no_grad():
         torch_versioned_kwargs = {}
@@ -777,11 +776,13 @@ def convert_to_onnx(
         model_input_names = [i.name for i in onnx_model.graph.input]
         input_dict = dict(zip(model_input_names, [i.cpu().numpy() for i in inputs]))
         if use_ort:
+            onnxruntime, _ = optional_import("onnxruntime")
             ort_sess = onnxruntime.InferenceSession(
                 onnx_model.SerializeToString(), providers=ort_provider if ort_provider else ["CPUExecutionProvider"]
             )
             onnx_out = ort_sess.run(None, input_dict)
         else:
+            onnxreference, _ = optional_import("onnx.reference")
             sess = onnxreference.ReferenceEvaluator(onnx_model)
             onnx_out = sess.run(None, input_dict)
         set_determinism(seed=None)

tests/networks/test_lazy_onnx_import.py

@@ -0,0 +1,51 @@
+# Copyright (c) MONAI Consortium
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#     http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import unittest
+
+
+class TestLazyOnnxImport(unittest.TestCase):
+    """Regression test for #8455.
+
+    ``onnx``/``onnx.reference``/``onnxruntime`` used to be imported at module
+    scope in ``monai/networks/utils.py`` and ``monai/bundle/scripts.py`` via
+    ``optional_import``, which imports eagerly. Because ``import monai``
+    auto-loads ``monai.networks``, a broken or hanging onnx install (e.g.
+    onnx 1.18 on Windows) would take down ``import monai`` with no error.
+
+    The fix moves those imports inside the functions that use them, so neither
+    module binds onnx at module scope any more. Assert that directly: it is
+    deterministic and independent of which other optional packages happen to be
+    installed (some of them import onnx transitively, so checking
+    ``sys.modules`` after ``import monai`` is not a reliable signal).
+    """
+
+    def test_utils_does_not_bind_onnx_at_module_scope(self):
+        import monai.networks.utils as utils
+
+        for attr in ("onnx", "onnxreference", "onnxruntime"):
+            self.assertFalse(
+                hasattr(utils, attr),
+                f"monai.networks.utils must not import {attr} at module scope (regression for #8455)",
+            )
+
+    def test_scripts_does_not_bind_onnx_at_module_scope(self):
+        import monai.bundle.scripts as scripts
+
+        self.assertFalse(
+            hasattr(scripts, "onnx"), "monai.bundle.scripts must not import onnx at module scope (regression for #8455)"
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()