Say that TensorRT is missing instead of raising NameError

[v0ropaev]

What

modelopt.onnx.autocast on a model that holds any non-standard operator dies with NameError: name 'trt' is not defined instead of saying that TensorRT is not installed.

modelopt/onnx/trt_utils.py:36-41 imports TensorRT under a try/except ImportError and records the answer:

try:
    import tensorrt as trt

    TRT_PYTHON_AVAILABLE = True
except ImportError:
    TRT_PYTHON_AVAILABLE = False

load_onnx_model checks it at :414 before it goes anywhere near get_custom_layers. GraphSanitizer.find_custom_nodes (modelopt/onnx/autocast/graphsanitizer.py:136) does not, and get_custom_layers reaches for trt at its first statement.

The onnx extra in pyproject.toml does not list tensorrt, so pip install nvidia-modelopt[onnx] leaves TRT_PYTHON_AVAILABLE false, which is exactly the case the flag exists for.

Reproduced on this machine, through the public entry point, on a two-node model whose second node sits in a custom.ops domain:

--- convert_to_mixed_precision('/tmp/custom.onnx')
INFO - Found custom operators: {'MyPlugin'}
INFO - Added TRT plugin domain trt.plugins version 1
    raised NameError: name 'trt' is not defined

Traceback (most recent call last):
  File "modelopt/onnx/autocast/convert.py", line 156, in convert_to_mixed_precision
    graph_sanitizer.sanitize()
  File "modelopt/onnx/autocast/graphsanitizer.py", line 136, in find_custom_nodes
    _, all_tensor_info = get_custom_layers(self.onnx_path or self.model, self.trt_plugins)
  File "modelopt/onnx/trt_utils.py", line 131, in get_custom_layers
    trt_logger = trt.Logger(trt.Logger.WARNING)
NameError: name 'trt' is not defined

After the change, the same run says:

    raised 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.

CI does not see this because the only test over that path, test_find_custom_nodes_uses_source_model_path in tests/unit/onnx/autocast/test_graphsanitizer.py:454, mocks get_custom_layers, set_trt_plugin_domain and infer_types_shapes_tensorrt, so the real TensorRT calls are never made.

Why there and not in find_custom_nodes

get_custom_layers is where the knowledge of the dependency lives, so the check covers every caller rather than the one that happens to crash today. infer_types_shapes has the same dependency but is only reached after get_custom_layers returns, so it needs no second guard for this path.

An alternative is to make autocast skip the TensorRT half with a warning, matching what load_onnx_model does at :414. I did not, because set_trt_plugin_domain has already rewritten the custom ops into the trt.plugins domain by then, and carrying on would hand the rest of the pipeline a model with no type or shape information for those nodes. Raising says what to do; skipping would be quieter and less correct. Say if you would rather have the skip and I will send that instead.

Test

test_find_custom_nodes_without_tensorrt_says_so beside the existing custom-node test, with monkeypatch.setattr(trt_utils, "TRT_PYTHON_AVAILABLE", False), which is the same technique tests/unit/onnx/quantization/test_quantize_api.py:229 already uses. On main it fails with the NameError above; after the change it passes.

tests/unit/onnx/autocast/test_graphsanitizer.py is 17 passed. ruff check and ruff format --check (0.15.20) are clean on both files.

Unrelated, found next door

While reading convert.py I noticed convert_to_mixed_precision takes providers: list[str] = ["cpu"] and then does providers.insert(0, "trt") at :168-170, so the shared default list keeps the trt entry for every later call in the process that omits the argument. Against the real function object:

providers default at import: ['cpu']
after one model with custom ops: ['trt', 'cpu']

That is a separate change and I have left it out of this one. Happy to send it if you want it.

Summary by CodeRabbit

• Bug Fixes
  • Custom-layer detection now reports a clear error when the TensorRT Python package is unavailable, rather than proceeding with TensorRT network operations and potentially failing less clearly. This makes it easier to understand why the operation cannot continue and what environment requirement is missing.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/Model-Optimizer/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 2b4b071a-c73f-4af0-ab80-50c9b881e3e1

---

📥 Commits

Reviewing files that changed from the base of the PR and between f27b476 and 1fde24d.

---

📒 Files selected for processing (2)

• modelopt/onnx/trt_utils.py
• tests/unit/onnx/autocast/test_graphsanitizer.py

---

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 9 remain after this review.

---

---

📝 Walkthrough

Walkthrough

get_custom_layers now checks whether TensorRT Python is available and raises an explanatory ImportError if it is not. A test verifies this error through find_custom_nodes().

Changes

TensorRT availability handling

Layer / File(s)	Summary
Availability guard and test modelopt/onnx/trt_utils.py, tests/unit/onnx/autocast/test_graphsanitizer.py	get_custom_layers raises an explanatory ImportError when TensorRT Python is unavailable. A test disables TensorRT availability and checks the error raised through find_custom_nodes().

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~8 minutes

Change: Bug fix

---

Merge Risk: ⚪ Minimal · up to 1fde2

Custom-operator handling now gives a clear missing-TensorRT error, while the inspected ordinary no-custom-operator path remains unaffected. No material merge risk is evident.

Pre-merge checks | 5 | 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check		Check skipped - CodeRabbit’s high-level summary is enabled.
Title check		The title clearly and concisely describes the main change: replacing an unhelpful NameError with an explanatory message when TensorRT is unavailable.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Security Anti-Patterns		PASS. The pull request changes only modelopt/onnx/trt_utils.py and a test file. The production change adds a TensorRT availability check and raises ImportError; it does not add `torch.load(..., we…

---

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

---

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

---

_{Comment @coderabbitai help to get the list of available commands.}