From 1fde24dc2565cc2464ce80ea53752d17d200a309 Mon Sep 17 00:00:00 2001 From: Dmitry Voropaev Date: Sun, 11 Oct 2026 12:45:08 +0300 Subject: [PATCH] Say that TensorRT is missing instead of raising NameError trt_utils imports tensorrt under a try/except and keeps the answer in TRT_PYTHON_AVAILABLE, which load_onnx_model checks before it reaches get_custom_layers. GraphSanitizer.find_custom_nodes does not, and the onnx extra does not install tensorrt, so autocast on a model holding any non-standard op dies with NameError: name 'trt' is not defined from inside trt_utils rather than saying what is missing. The guard goes in get_custom_layers, where the knowledge lives, so every caller gets the message. Signed-off-by: Dmitry Voropaev --- modelopt/onnx/trt_utils.py | 10 ++++++++++ tests/unit/onnx/autocast/test_graphsanitizer.py | 17 +++++++++++++++++ 2 files changed, 27 insertions(+) diff --git a/modelopt/onnx/trt_utils.py b/modelopt/onnx/trt_utils.py index b407fdc5411..e862d93f97d 100644 --- a/modelopt/onnx/trt_utils.py +++ b/modelopt/onnx/trt_utils.py @@ -127,6 +127,16 @@ def get_custom_layers( """ logger.debug("Checking for custom TensorRT ops") + # The `onnx` extra does not install the TensorRT Python package, and reading a model's + # custom layers goes through a TensorRT network. Say so, rather than letting the `trt` + # name that the import at the top of this module never bound raise a NameError. + if not TRT_PYTHON_AVAILABLE: + raise ImportError( + "The TensorRT Python package is required to read the custom layers of an ONNX " + "model, and it is not installed. Install `tensorrt` to work with a model that " + "contains custom operators." + ) + # Initialize TensorRT logger and plugins trt_logger = trt.Logger(trt.Logger.WARNING) trt.init_libnvinfer_plugins(trt_logger, "") diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py index f42324a6338..6b56f745ab2 100644 --- a/tests/unit/onnx/autocast/test_graphsanitizer.py +++ b/tests/unit/onnx/autocast/test_graphsanitizer.py @@ -21,6 +21,7 @@ from onnx import TensorProto, helper, numpy_helper import modelopt.onnx.autocast.graphsanitizer as graphsanitizer +import modelopt.onnx.trt_utils as trt_utils from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer @@ -471,3 +472,19 @@ def test_find_custom_nodes_uses_source_model_path(tmp_path, monkeypatch): get_custom_layers.assert_called_once_with(str(model_path.resolve()), []) infer_types_shapes.assert_called_once_with(model, [], all_tensor_info=tensor_info) + + +def test_find_custom_nodes_without_tensorrt_says_so(monkeypatch): + """The `onnx` extra does not install tensorrt, so this has to be readable.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1]) + custom_node = helper.make_node("CustomOp", [x.name], [y.name], name="custom") + graph = helper.make_graph([custom_node], "custom_graph", [x], [y]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 22)]) + + monkeypatch.setattr(trt_utils, "TRT_PYTHON_AVAILABLE", False) + monkeypatch.setattr(graphsanitizer, "set_trt_plugin_domain", Mock(return_value=model)) + + sanitizer = GraphSanitizer(model, min_opset=22) + with pytest.raises(ImportError, match="TensorRT Python package"): + sanitizer.find_custom_nodes()