Skip to content

Centralize 'trtexec' subprocess runs in ONNX into a single function #1268

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gcunhase wants to merge 1 commit into
base: main
Choose a base branch
from
+12 −4
Open

Centralize 'trtexec' subprocess runs in ONNX into a single function #1268

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 2 additions & 3 deletions modelopt/onnx/quantization/autotune/benchmark.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,6 @@
import os
import re
import shutil
import subprocess # nosec B404
import tempfile
import time
from abc import ABC, abstractmethod
Expand All @@ -42,7 +41,7 @@
import torch

from modelopt.onnx.logging_config import logger
from modelopt.onnx.quantization.ort_utils import _check_for_trtexec
from modelopt.onnx.quantization.ort_utils import _check_for_trtexec, _run_trtexec

TRT_AVAILABLE = importlib.util.find_spec("tensorrt") is not None
if TRT_AVAILABLE:
Expand Down Expand Up @@ -270,7 +269,7 @@ def run(

cmd = [*self._base_cmd, f"--onnx={model_path}"]
self.logger.debug(f"Running: {' '.join(cmd)}")
result = subprocess.run(cmd, capture_output=True, text=True) # nosec B603
result = _run_trtexec(cmd)
self._write_log_file(
log_file,
"\n".join(
Expand Down
11 changes: 10 additions & 1 deletion modelopt/onnx/quantization/ort_utils.py
View file
Open in desktop
Copy link
Copy Markdown
Collaborator

@kevalmorabia97 kevalmorabia97 Apr 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you unify with

Model-Optimizer/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py

Line 62 in 07ae8e7

p = subprocess.Popen(cmd, stdout=log, stderr=log, cwd=str(cwd) if cwd else None) # nosec
also

Copy link
Copy Markdown
Contributor Author

@gcunhase gcunhase Apr 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I was trying to avoid coupling modelopt/torch with modelopt/onnx. What do you think is best?

Copy link
Copy Markdown
Collaborator

@kevalmorabia97 kevalmorabia97 Apr 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

modelopt.torch._deploy already has some dependencies on modelopt.onnx so its fine. We can keep this in modelopt.onnx.trt_utils or whatever names makes sense

For example this: https://github.com/NVIDIA/Model-Optimizer/blob/07ae8e71281a1c17898991b0cd76db5788f490bb/modelopt/torch/_deploy/_runtime/trt_client.py

Copy link
Copy Markdown
Contributor Author

@gcunhase gcunhase Apr 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good point, will look into this today.

Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,15 @@ def _check_lib_in_ld_library_path(ld_library_path, lib_pattern):
return False, None


def _run_trtexec(cmd, timeout=None):
"""Run a 'trtexec' command via subprocess."""
# Ensure that this command is a trtexec run
assert any("trtexec" in c for c in cmd), "Subprocess can only execute 'trtexec' commands"
Copy link
Copy Markdown
Contributor Author

@gcunhase gcunhase Apr 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@kevalmorabia97 please let me know if this is enough to ensure only trtexec runs or if this is even necessary. Thanks.

Copy link
Copy Markdown
Collaborator

@kevalmorabia97 kevalmorabia97 Apr 16, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we not pass trtexec in cmd and only pass arguments to trtexec so this function prepends trtexec to the cmd?

Please add type annotation also


# Run trtexec command
return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) # nosec B603
Comment on lines +47 to +53
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai bot Apr 15, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🔴 Critical

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Verify the problematic patterns exist at the expected locations.
rg -n 'def _run_trtexec|assert any\("trtexec"|# nosec B603|subprocess\.run\(' modelopt/onnx/quantization/ort_utils.py -C2

Repository: NVIDIA/Model-Optimizer

Length of output: 452

🏁 Script executed:

# Check for calls to _run_trtexec to understand usage context
rg -n '_run_trtexec' modelopt/onnx/quantization/ort_utils.py -A 5 -B 5

Repository: NVIDIA/Model-Optimizer

Length of output: 927

🏁 Script executed:

# Check if there are any imports of subprocess and how it's being used elsewhere
rg -n 'import subprocess|from subprocess' modelopt/onnx/quantization/ort_utils.py -B 2 -A 2

Repository: NVIDIA/Model-Optimizer

Length of output: 184

🏁 Script executed:

# Verify the file structure and get full context around the function
sed -n '1,60p' modelopt/onnx/quantization/ort_utils.py

Repository: NVIDIA/Model-Optimizer

Length of output: 2324

Replace the assert gate and remove the # nosec bypass in _run_trtexec.

Line 50 uses assert for command validation, which can be disabled with Python optimization flags (-O, -OO), making it an unreliable security control. Line 53 adds # nosec B603, which violates repo policy: "Bandit security checks are enforced via pre-commit hooks. # nosec comments are not allowed as bypasses." Replace the assert with an explicit runtime check on cmd[0] and remove the nosec suppression. Request formal codeowner approval only if the subprocess pattern genuinely cannot be restructured to satisfy Bandit without exceptions.

Proposed fix 
-def _run_trtexec(cmd, timeout=None):
+def _run_trtexec(cmd: Sequence[str], timeout: float | None = None):
     """Run a 'trtexec' command via subprocess."""
-    # Ensure that this command is a trtexec run
-    assert any("trtexec" in c for c in cmd), "Subprocess can only execute 'trtexec' commands"
+    if not cmd:
+        raise ValueError("Empty command passed to _run_trtexec")
+    exe = os.path.basename(str(cmd[0])).lower()
+    if exe not in {"trtexec", "trtexec.exe"}:
+        raise ValueError("Subprocess can only execute 'trtexec' commands")
 
     # Run trtexec command
-    return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)  # nosec B603
+    return subprocess.run(list(cmd), capture_output=True, text=True, timeout=timeout, check=False)
🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@modelopt/onnx/quantization/ort_utils.py` around lines 47 - 53, Replace the
fragile assert and the banned "# nosec" in _run_trtexec: validate that cmd is a
non-empty sequence and that cmd[0] contains "trtexec" (or equals "trtexec")
using an explicit runtime check (e.g., if not cmd or "trtexec" not in
str(cmd[0]): raise ValueError("...")), then call subprocess.run(cmd,
capture_output=True, text=True, timeout=timeout) without any "# nosec" comment;
keep the function name _run_trtexec and parameters (cmd, timeout) and raise a
clear exception (ValueError/RuntimeError) when validation fails so Bandit checks
are not bypassed.



def _check_for_trtexec(min_version: str = "10.0") -> str:
"""Check if the `trtexec` CLI tool is available in PATH and is >= min_version.

Expand Down Expand Up @@ -87,7 +96,7 @@ def _parse_version_from_string(version_str: str) -> str | None:
)

try:
result = subprocess.run([trtexec_path], capture_output=True, text=True, timeout=5) # nosec B603
result = _run_trtexec([trtexec_path], timeout=5)
banner_output = result.stdout + result.stderr
parsed_version = _parse_version_from_string(banner_output)

Expand Down
Loading