From 6e5f2d722b8652b0fd272a118c8f50b9834f7cb0 Mon Sep 17 00:00:00 2001 From: Haoxi Zhang Date: Fri, 11 Sep 2026 13:33:05 -0700 Subject: [PATCH 1/4] Handle legacy ONNX Runtime operator validation Signed-off-by: Haoxi Zhang --- modelopt/onnx/trt_utils.py | 29 +++++++++++++++++++++++++++-- tests/unit/onnx/test_onnx_utils.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/trt_utils.py b/modelopt/onnx/trt_utils.py index b407fdc5411..27d17f520ae 100644 --- a/modelopt/onnx/trt_utils.py +++ b/modelopt/onnx/trt_utils.py @@ -23,6 +23,7 @@ import lief import onnx +import onnx.onnx_cpp2py_export.checker as C import onnx_graphsurgeon as gs from modelopt.onnx.logging_config import logger @@ -41,6 +42,30 @@ TRT_PYTHON_AVAILABLE = False MAX_IR_VERSION = 10 +_ORT_LEGACY_ONNX_DOMAIN_OPS = {"SimplifiedLayerNormalization"} + + +def _check_onnx_model(model: onnx.ModelProto, model_path: str | None = None) -> None: + """Validate a model while tolerating legacy ONNX Runtime operators.""" + try: + onnx.checker.check_model(model_path or model) + except C.ValidationError as e: + error = str(e) + unsupported_legacy_ops = { + node.op_type + for node in model.graph.node + if not node.domain + and node.op_type in _ORT_LEGACY_ONNX_DOMAIN_OPS + and f"No Op registered for {node.op_type} with domain_version" in error + } + if not unsupported_legacy_ops: + raise + + logger.warning( + "ONNX checker does not recognize ONNX Runtime legacy operator(s) in the default " + "domain: %s. Continuing because ONNX Runtime supports these operators.", + sorted(unsupported_legacy_ops), + ) def _is_static_plugin(plugin_path: str) -> bool: @@ -453,10 +478,10 @@ def load_onnx_model( if use_external_data_format: # For large models, use the file path to avoid protobuf size limitation model_path_to_check = ir_version_onnx_path or static_shaped_onnx_path or onnx_path - onnx.checker.check_model(model_path_to_check) + _check_onnx_model(onnx_model, model_path_to_check) else: # For smaller models, checking the model object is fine - onnx.checker.check_model(onnx_model) + _check_onnx_model(onnx_model) return ( onnx_model, diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py index 36face35b90..79ed75ce8ed 100644 --- a/tests/unit/onnx/test_onnx_utils.py +++ b/tests/unit/onnx/test_onnx_utils.py @@ -333,6 +333,34 @@ def test_ir_version_support(tmp_path): ) +@pytest.mark.parametrize("use_external_data_format", [False, True]) +def test_load_onnx_model_with_ort_legacy_op(tmp_path, caplog, use_external_data_format): + node = make_node( + "SimplifiedLayerNormalization", + ["X", "scale"], + ["Y"], + name="simplified_layer_norm", + ) + graph = make_graph( + [node], + "ort_legacy_op_graph", + [make_tensor_value_info("X", onnx.TensorProto.FLOAT, [1, 4])], + [make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [1, 4])], + [make_tensor("scale", onnx.TensorProto.FLOAT, [4], [1.0] * 4)], + ) + model = make_model(graph, opset_imports=[make_opsetid("", 21)], ir_version=10) + model_path = os.path.join(tmp_path, "ort_legacy_op.onnx") + onnx.save(model, model_path) + + loaded_model, _, _, _, _ = load_onnx_model( + model_path, use_external_data_format=use_external_data_format + ) + + assert loaded_model.graph.node[0].op_type == "SimplifiedLayerNormalization" + assert loaded_model.graph.node[0].domain == "" + assert "ONNX Runtime legacy operator(s)" in caplog.text + + def _make_cast_model(cast_to, output_elem_type, with_value_info=False): """Build a tiny X -> Cast(to=cast_to) -> Y model.""" nodes = [make_node("Cast", ["X"], ["Y"], to=cast_to, name="cast")] From 7ffa341e04d6f0cef3dd3374c1f55978eceb4452 Mon Sep 17 00:00:00 2001 From: Haoxi Zhang Date: Sat, 12 Sep 2026 17:28:57 -0700 Subject: [PATCH 2/4] Fix: register legacy operator to onnx check_model Signed-off-by: Haoxi Zhang --- modelopt/onnx/quantization/quantize.py | 6 +- modelopt/onnx/trt_utils.py | 30 +------ modelopt/onnx/utils.py | 117 +++++++++++++++++++++++- tests/unit/onnx/test_onnx_utils.py | 120 +++++++++++++++++++++++-- 4 files changed, 232 insertions(+), 41 deletions(-) diff --git a/modelopt/onnx/quantization/quantize.py b/modelopt/onnx/quantization/quantize.py index fb6b662916d..e8f06639087 100755 --- a/modelopt/onnx/quantization/quantize.py +++ b/modelopt/onnx/quantization/quantize.py @@ -43,7 +43,6 @@ from typing import Any import onnx -import onnx.onnx_cpp2py_export.checker as C import onnx_graphsurgeon as gs import onnxslim from onnxruntime.quantization.calibrate import CalibrationDataReader @@ -77,6 +76,7 @@ from modelopt.onnx.utils import ( BASE_MIN_OPSET, QDQ_PRECISION_MIN_OPSET, + check_model, clear_stale_value_info, duplicate_shared_constants, get_opset_version, @@ -955,8 +955,8 @@ def quantize( # Check if the quantized model is valid try: logger.info("Validating quantized model") - onnx.checker.check_model(output_path) - except C.ValidationError as e: + check_model(onnx_model, output_path) + except onnx.checker.ValidationError as e: logger.warning("ONNX model checker failed, check your deployment status") logger.warning(e) diff --git a/modelopt/onnx/trt_utils.py b/modelopt/onnx/trt_utils.py index 27d17f520ae..cf53a411adf 100644 --- a/modelopt/onnx/trt_utils.py +++ b/modelopt/onnx/trt_utils.py @@ -23,11 +23,11 @@ import lief import onnx -import onnx.onnx_cpp2py_export.checker as C import onnx_graphsurgeon as gs from modelopt.onnx.logging_config import logger from modelopt.onnx.utils import ( + check_model, get_dynamic_graph_inputs, get_tensor_by_name, parse_shapes_spec, @@ -42,30 +42,6 @@ TRT_PYTHON_AVAILABLE = False MAX_IR_VERSION = 10 -_ORT_LEGACY_ONNX_DOMAIN_OPS = {"SimplifiedLayerNormalization"} - - -def _check_onnx_model(model: onnx.ModelProto, model_path: str | None = None) -> None: - """Validate a model while tolerating legacy ONNX Runtime operators.""" - try: - onnx.checker.check_model(model_path or model) - except C.ValidationError as e: - error = str(e) - unsupported_legacy_ops = { - node.op_type - for node in model.graph.node - if not node.domain - and node.op_type in _ORT_LEGACY_ONNX_DOMAIN_OPS - and f"No Op registered for {node.op_type} with domain_version" in error - } - if not unsupported_legacy_ops: - raise - - logger.warning( - "ONNX checker does not recognize ONNX Runtime legacy operator(s) in the default " - "domain: %s. Continuing because ONNX Runtime supports these operators.", - sorted(unsupported_legacy_ops), - ) def _is_static_plugin(plugin_path: str) -> bool: @@ -478,10 +454,10 @@ def load_onnx_model( if use_external_data_format: # For large models, use the file path to avoid protobuf size limitation model_path_to_check = ir_version_onnx_path or static_shaped_onnx_path or onnx_path - _check_onnx_model(onnx_model, model_path_to_check) + check_model(onnx_model, model_path_to_check) else: # For smaller models, checking the model object is fine - _check_onnx_model(onnx_model) + check_model(onnx_model) return ( onnx_model, diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index f8b5a41a41a..6ea22eb52b7 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -34,6 +34,89 @@ # Base minimum opset for quantization (opset 19 is the first to support fp16 scales) BASE_MIN_OPSET = 19 +# Populated permanently by the first call to check_model. +_ORT_LEGACY_ONNX_DOMAIN_OPS: frozenset[str] | None = None + + +def _convert_ort_parameter(parameter: Any) -> onnx.defs.OpSchema.FormalParameter: + """Rebuild an ORT formal parameter using ONNX's separate pybind type.""" + schema_type = onnx.defs.OpSchema + # The two bindings define equivalent enums but do not share Python types, so + # translate by stable enum name instead of relying on their numeric values. + return schema_type.FormalParameter( + parameter.name, + parameter.typeStr, + parameter.description or "", + param_option=getattr(schema_type.FormalParameterOption, parameter.option.name), + is_homogeneous=parameter.isHomogeneous, + ) + + +def _convert_ort_schema(ort_schema: Any) -> onnx.defs.OpSchema: + """Rebuild an ORT operator schema in ONNX's independent schema type. + + ORT vendors its own ONNX C++ library, so its pybind objects and registry are + not visible to the separately installed ONNX package. + """ + schema_type = onnx.defs.OpSchema + # This is the metadata used by check_model(full_check=False). Native ORT + # shape-inference functions cannot cross the binary boundary between wheels. + return schema_type( + ort_schema.name, + ort_schema.domain, + ort_schema.since_version, + ort_schema.doc or "", + inputs=[_convert_ort_parameter(parameter) for parameter in ort_schema.inputs], + outputs=[_convert_ort_parameter(parameter) for parameter in ort_schema.outputs], + type_constraints=[ + ( + constraint.type_param_str, + list(constraint.allowed_type_strs), + constraint.description or "", + ) + for constraint in ort_schema.type_constraints + ], + attributes=[ + schema_type.Attribute( + attribute.name, + getattr(schema_type.AttrType, attribute.type.name), + attribute.description or "", + required=attribute.required, + ) + for attribute in ort_schema.attributes.values() + ], + ) + + +def _register_ort_legacy_schemas() -> frozenset[str]: + """Register ORT-only default-domain schemas on the first model check.""" + global _ORT_LEGACY_ONNX_DOMAIN_OPS + if _ORT_LEGACY_ONNX_DOMAIN_OPS is not None: + return _ORT_LEGACY_ONNX_DOMAIN_OPS + + # ORT exposes its independently compiled schema registry only through this + # private binding. Import it lazily with schema registration. + from onnxruntime.capi._pybind_state import get_all_operator_schema + + onnx_standard_ops = { + schema.name for schema in onnx.defs.get_all_schemas_with_history() if not schema.domain + } + # Filter by name so ORT cannot add alternate schema histories for operators + # that ONNX already owns; only genuinely missing default-domain ops are added. + legacy_schemas = sorted( + ( + schema + for schema in get_all_operator_schema() + if not schema.domain and schema.name not in onnx_standard_ops + ), + key=lambda schema: (schema.name, schema.since_version), + ) + for schema in legacy_schemas: + onnx.defs.register_schema(_convert_ort_schema(schema)) + + _ORT_LEGACY_ONNX_DOMAIN_OPS = frozenset(schema.name for schema in legacy_schemas) + return _ORT_LEGACY_ONNX_DOMAIN_OPS + def get_input_names_from_bytes(model_bytes: bytes, external_inputs_only: bool = True) -> list[str]: """This function returns the inputs names of the given onnx model in bytes. @@ -556,8 +639,38 @@ def _get_unique_name(old_name): return onnx_model, is_modified -def check_model(model: onnx.ModelProto) -> None: - """Checks if the given model is valid.""" +def check_model(model: onnx.ModelProto, model_path: str | None = None) -> None: + """Check whether a model is structurally valid. + + ONNX Runtime legacy operators registered in the default ONNX domain are + accepted by this initial guard. Whether they can execute remains the + responsibility of the selected backend. + + Args: + model: Loaded in-memory ONNX model. Used for validation unless + model_path is supplied and for detecting legacy operators. + model_path: Optional file-backed copy to validate. Use this for models + with external data or models too large for protobuf serialization. + """ + ort_legacy_ops = _register_ort_legacy_schemas() + legacy_ops = sorted( + { + node.op_type + for node in model.graph.node + if not node.domain and node.op_type in ort_legacy_ops + } + ) + if legacy_ops: + logger.warning( + "Model uses ONNX Runtime legacy operator(s) in the default ONNX domain: %s. " + "Execution support is delegated to the selected backend.", + legacy_ops, + ) + + if model_path is not None: + onnx.checker.check_model(model_path) + return + save_as_external_data = False try: model_size = model.ByteSize() diff --git a/tests/unit/onnx/test_onnx_utils.py b/tests/unit/onnx/test_onnx_utils.py index 79ed75ce8ed..1be38d9e557 100644 --- a/tests/unit/onnx/test_onnx_utils.py +++ b/tests/unit/onnx/test_onnx_utils.py @@ -14,6 +14,9 @@ # limitations under the License. import os +import subprocess +import sys +import textwrap import numpy as np import onnx @@ -30,6 +33,7 @@ from modelopt.onnx.trt_utils import load_onnx_model from modelopt.onnx.utils import ( + check_model, clear_stale_value_info, get_input_names_from_bytes, get_output_names_from_bytes, @@ -333,24 +337,70 @@ def test_ir_version_support(tmp_path): ) -@pytest.mark.parametrize("use_external_data_format", [False, True]) -def test_load_onnx_model_with_ort_legacy_op(tmp_path, caplog, use_external_data_format): +def test_ort_legacy_schema_registration_is_lazy(): + script = textwrap.dedent( + """ + import onnx + from onnx.helper import ( + make_graph, + make_model, + make_node, + make_opsetid, + make_tensor_value_info, + ) + + import modelopt.onnx.utils as onnx_utils + + assert onnx_utils._ORT_LEGACY_ONNX_DOMAIN_OPS is None + assert not onnx.defs.has("SimplifiedLayerNormalization", 21, "") + + node = make_node("Relu", ["X"], ["Y"]) + input_value = make_tensor_value_info("X", onnx.TensorProto.FLOAT, [1]) + output_value = make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [1]) + graph = make_graph([node], "standard_graph", [input_value], [output_value]) + model = make_model(graph, opset_imports=[make_opsetid("", 21)], ir_version=10) + + onnx_utils.check_model(model) + + assert onnx_utils._ORT_LEGACY_ONNX_DOMAIN_OPS is not None + assert onnx.defs.has("SimplifiedLayerNormalization", 21, "") + """ + ) + + subprocess.run([sys.executable, "-c", script], check=True) + + +def _make_ort_legacy_op_model(following_nodes=None): node = make_node( "SimplifiedLayerNormalization", ["X", "scale"], - ["Y"], + ["normalized"], name="simplified_layer_norm", ) + following_nodes = following_nodes or [] + output_name = following_nodes[-1].output[0] if following_nodes else "normalized" graph = make_graph( - [node], + [node, *following_nodes], "ort_legacy_op_graph", [make_tensor_value_info("X", onnx.TensorProto.FLOAT, [1, 4])], - [make_tensor_value_info("Y", onnx.TensorProto.FLOAT, [1, 4])], - [make_tensor("scale", onnx.TensorProto.FLOAT, [4], [1.0] * 4)], + [make_tensor_value_info(output_name, onnx.TensorProto.FLOAT, [1, 4])], + [onnx.numpy_helper.from_array(np.ones(4, dtype=np.float32), name="scale")], ) - model = make_model(graph, opset_imports=[make_opsetid("", 21)], ir_version=10) + return make_model(graph, opset_imports=[make_opsetid("", 21)], ir_version=10) + + +@pytest.mark.parametrize("use_external_data_format", [False, True]) +def test_load_onnx_model_with_ort_legacy_op(tmp_path, caplog, use_external_data_format): + model = _make_ort_legacy_op_model() model_path = os.path.join(tmp_path, "ort_legacy_op.onnx") - onnx.save(model, model_path) + onnx.save_model( + model, + model_path, + save_as_external_data=use_external_data_format, + all_tensors_to_one_file=True, + location="ort_legacy_op.data", + size_threshold=0, + ) loaded_model, _, _, _, _ = load_onnx_model( model_path, use_external_data_format=use_external_data_format @@ -358,7 +408,59 @@ def test_load_onnx_model_with_ort_legacy_op(tmp_path, caplog, use_external_data_ assert loaded_model.graph.node[0].op_type == "SimplifiedLayerNormalization" assert loaded_model.graph.node[0].domain == "" - assert "ONNX Runtime legacy operator(s)" in caplog.text + assert "Model uses ONNX Runtime legacy operator(s)" in caplog.text + assert onnx.defs.has("SimplifiedLayerNormalization", 21, "") + + +def test_ort_legacy_op_schema_is_checked(): + model = _make_ort_legacy_op_model() + model.graph.node[0].input.pop() + + with pytest.raises(onnx.checker.ValidationError, match="SimplifiedLayerNormalization"): + check_model(model) + + assert onnx.defs.has("SimplifiedLayerNormalization", 21, "") + + +@pytest.mark.parametrize("use_model_path", [False, True]) +def test_ort_legacy_op_does_not_hide_other_validation_errors(tmp_path, use_model_path): + invalid_relu = make_node("Relu", ["normalized", "scale"], ["Y"], name="invalid_relu") + model = _make_ort_legacy_op_model([invalid_relu]) + model_path = os.path.join(tmp_path, "invalid_ort_legacy_op.onnx") + onnx.save(model, model_path) + + with pytest.raises(onnx.checker.ValidationError, match="Relu"): + check_model(model, model_path if use_model_path else None) + + assert model.graph.node[0].domain == "" + + +def test_unknown_default_domain_op_fails_validation(): + model = _make_ort_legacy_op_model() + model.graph.node[0].op_type = "UnknownModelOptOp" + + with pytest.raises(onnx.checker.ValidationError, match="No Op registered"): + check_model(model) + + +def test_ort_legacy_op_does_not_hide_missing_external_data(tmp_path): + model = _make_ort_legacy_op_model() + model_path = os.path.join(tmp_path, "missing_external_data.onnx") + data_name = "missing_external_data.bin" + onnx.save_model( + model, + model_path, + save_as_external_data=True, + all_tensors_to_one_file=True, + location=data_name, + size_threshold=0, + ) + os.remove(os.path.join(tmp_path, data_name)) + + with pytest.raises(onnx.checker.ValidationError): + check_model(model, model_path) + + assert os.listdir(tmp_path) == ["missing_external_data.onnx"] def _make_cast_model(cast_to, output_elem_type, with_value_info=False): From c1cf121eef0a5452e97485b43fc997fbac773558 Mon Sep 17 00:00:00 2001 From: Haoxi Zhang Date: Fri, 11 Sep 2026 14:08:11 -0700 Subject: [PATCH 3/4] Skip TensorRT introspection without Python bindings Signed-off-by: Haoxi Zhang --- modelopt/onnx/autocast/graphsanitizer.py | 12 +++++-- .../unit/onnx/autocast/test_graphsanitizer.py | 31 +++++++++++++++++++ 2 files changed, 40 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/autocast/graphsanitizer.py b/modelopt/onnx/autocast/graphsanitizer.py index 2154a42568e..9792c2416ca 100644 --- a/modelopt/onnx/autocast/graphsanitizer.py +++ b/modelopt/onnx/autocast/graphsanitizer.py @@ -118,13 +118,19 @@ def find_custom_nodes(self) -> None: node.op_type for node in self.model.graph.node if node.op_type not in self.standard_ops } if self.custom_ops: - from modelopt.onnx.trt_utils import infer_types_shapes_tensorrt, set_trt_plugin_domain + from modelopt.onnx import trt_utils + + if not trt_utils.TRT_PYTHON_AVAILABLE: + logger.warning( + "TensorRT Python bindings are not available; skipping custom-layer introspection." + ) + return # Set TensorRT plugin domain info in the graph for ORT compatibility - self.model = set_trt_plugin_domain(self.model, self.custom_ops) + self.model = trt_utils.set_trt_plugin_domain(self.model, self.custom_ops) # Infer types and shapes in the graph for ORT compatibility - self.model = infer_types_shapes_tensorrt(self.model, self.trt_plugins) + self.model = trt_utils.infer_types_shapes_tensorrt(self.model, self.trt_plugins) def remove_disconnected_outputs(self) -> None: """Remove disconnected outputs from the model.""" diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py index cb487b56bf0..251353967ed 100644 --- a/tests/unit/onnx/autocast/test_graphsanitizer.py +++ b/tests/unit/onnx/autocast/test_graphsanitizer.py @@ -17,9 +17,40 @@ import pytest from onnx import TensorProto, helper, numpy_helper +import modelopt.onnx.trt_utils as trt_utils from modelopt.onnx.autocast.graphsanitizer import GraphSanitizer +def test_find_custom_nodes_without_tensorrt(monkeypatch): + """Custom-op discovery should not require the optional TensorRT Python bindings.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4]) + node = helper.make_node( + "SimplifiedLayerNormalization", + ["X"], + ["Y"], + name="layer_norm", + domain="com.microsoft", + ) + graph = helper.make_graph([node], "custom_op_test", [x], [y]) + model = helper.make_model( + graph, + opset_imports=[ + helper.make_opsetid("", 18), + helper.make_opsetid("com.microsoft", 1), + ], + ) + sanitizer = GraphSanitizer(model) + + monkeypatch.setattr(trt_utils, "TRT_PYTHON_AVAILABLE", False) + + sanitizer.find_custom_nodes() + + assert sanitizer.custom_ops == {"SimplifiedLayerNormalization"} + assert sanitizer.model.graph.node[0].domain == "com.microsoft" + assert all(opset.domain != "trt.plugins" for opset in sanitizer.model.opset_import) + + def create_layernorm_model(input_shape, epsilon=1e-5, axis=-1, add_scale=True, add_bias=True): """Helper function to create an ONNX model with a decomposed LayerNorm pattern""" x = helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape) From b837b3d5f1d99b5899b6f19cb359fc24267a0239 Mon Sep 17 00:00:00 2001 From: Haoxi Zhang Date: Sat, 12 Sep 2026 18:23:37 -0700 Subject: [PATCH 4/4] Keep ORT legacy operators out of TensorRT plugin handling Signed-off-by: Haoxi Zhang --- modelopt/onnx/autocast/graphsanitizer.py | 6 ++- modelopt/onnx/utils.py | 8 ++-- .../unit/onnx/autocast/test_graphsanitizer.py | 40 ++++++++++++++++--- 3 files changed, 43 insertions(+), 11 deletions(-) diff --git a/modelopt/onnx/autocast/graphsanitizer.py b/modelopt/onnx/autocast/graphsanitizer.py index 9792c2416ca..c0a54df0717 100644 --- a/modelopt/onnx/autocast/graphsanitizer.py +++ b/modelopt/onnx/autocast/graphsanitizer.py @@ -51,6 +51,7 @@ def __init__( self.min_opset = min_opset self.max_ir_version = max_ir_version self.standard_ops = {schema.name for schema in onnx.defs.get_all_schemas()} + self.ort_legacy_ops = onnx_utils.register_ort_legacy_schemas() self.custom_ops = None self.custom_ops_low_precision_nodes = [] self.trt_plugins = trt_plugins @@ -115,7 +116,10 @@ def find_custom_nodes(self) -> None: that are not part of the standard ONNX operator set. """ self.custom_ops = { - node.op_type for node in self.model.graph.node if node.op_type not in self.standard_ops + node.op_type + for node in self.model.graph.node + if node.op_type not in self.standard_ops + and not (not node.domain and node.op_type in self.ort_legacy_ops) } if self.custom_ops: from modelopt.onnx import trt_utils diff --git a/modelopt/onnx/utils.py b/modelopt/onnx/utils.py index 6ea22eb52b7..07078a72981 100644 --- a/modelopt/onnx/utils.py +++ b/modelopt/onnx/utils.py @@ -34,7 +34,7 @@ # Base minimum opset for quantization (opset 19 is the first to support fp16 scales) BASE_MIN_OPSET = 19 -# Populated permanently by the first call to check_model. +# Populated permanently by the first legacy-schema registration. _ORT_LEGACY_ONNX_DOMAIN_OPS: frozenset[str] | None = None @@ -88,8 +88,8 @@ def _convert_ort_schema(ort_schema: Any) -> onnx.defs.OpSchema: ) -def _register_ort_legacy_schemas() -> frozenset[str]: - """Register ORT-only default-domain schemas on the first model check.""" +def register_ort_legacy_schemas() -> frozenset[str]: + """Register ORT-only default-domain schemas and return their operator names.""" global _ORT_LEGACY_ONNX_DOMAIN_OPS if _ORT_LEGACY_ONNX_DOMAIN_OPS is not None: return _ORT_LEGACY_ONNX_DOMAIN_OPS @@ -652,7 +652,7 @@ def check_model(model: onnx.ModelProto, model_path: str | None = None) -> None: model_path: Optional file-backed copy to validate. Use this for models with external data or models too large for protobuf serialization. """ - ort_legacy_ops = _register_ort_legacy_schemas() + ort_legacy_ops = register_ort_legacy_schemas() legacy_ops = sorted( { node.op_type diff --git a/tests/unit/onnx/autocast/test_graphsanitizer.py b/tests/unit/onnx/autocast/test_graphsanitizer.py index 251353967ed..ab1124a08df 100644 --- a/tests/unit/onnx/autocast/test_graphsanitizer.py +++ b/tests/unit/onnx/autocast/test_graphsanitizer.py @@ -26,18 +26,18 @@ def test_find_custom_nodes_without_tensorrt(monkeypatch): x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4]) y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4]) node = helper.make_node( - "SimplifiedLayerNormalization", + "FakeTensorRTPlugin", ["X"], ["Y"], - name="layer_norm", - domain="com.microsoft", + name="plugin", + domain="test.plugins", ) graph = helper.make_graph([node], "custom_op_test", [x], [y]) model = helper.make_model( graph, opset_imports=[ helper.make_opsetid("", 18), - helper.make_opsetid("com.microsoft", 1), + helper.make_opsetid("test.plugins", 1), ], ) sanitizer = GraphSanitizer(model) @@ -46,8 +46,36 @@ def test_find_custom_nodes_without_tensorrt(monkeypatch): sanitizer.find_custom_nodes() - assert sanitizer.custom_ops == {"SimplifiedLayerNormalization"} - assert sanitizer.model.graph.node[0].domain == "com.microsoft" + assert sanitizer.custom_ops == {"FakeTensorRTPlugin"} + assert sanitizer.model.graph.node[0].domain == "test.plugins" + assert all(opset.domain != "trt.plugins" for opset in sanitizer.model.opset_import) + + +def test_find_custom_nodes_treats_ort_legacy_op_as_known(monkeypatch): + """ORT legacy operators should not be treated as TensorRT plugins.""" + x = helper.make_tensor_value_info("X", TensorProto.FLOAT, [1, 4]) + y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [1, 4]) + scale = numpy_helper.from_array(np.ones(4, dtype=np.float32), name="scale") + node = helper.make_node( + "SimplifiedLayerNormalization", + ["X", "scale"], + ["Y"], + name="layer_norm", + ) + graph = helper.make_graph([node], "legacy_op_test", [x], [y], [scale]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 21)]) + sanitizer = GraphSanitizer(model) + + monkeypatch.setattr(trt_utils, "TRT_PYTHON_AVAILABLE", True) + + def fail_if_called(*args, **kwargs): + pytest.fail("ORT legacy operator entered TensorRT plugin handling") + + monkeypatch.setattr(trt_utils, "set_trt_plugin_domain", fail_if_called) + sanitizer.find_custom_nodes() + + assert sanitizer.custom_ops == set() + assert sanitizer.model.graph.node[0].domain == "" assert all(opset.domain != "trt.plugins" for opset in sanitizer.model.opset_import)