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

modelopt/onnx/quantization/autotune/benchmark.py

@@ -31,7 +31,6 @@
 import os
 import re
 import shutil
-import subprocess  # nosec B404
 import tempfile
 import time
 from abc import ABC, abstractmethod
@@ -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:
@@ -159,7 +158,6 @@ def __init__(
         warmup_runs: int = 5,
         timing_runs: int = 10,
         plugin_libraries: list[str] | None = None,
-        trtexec_path: str = "trtexec",
         trtexec_args: list[str] | None = None,
     ):
         """Initialize the trtexec benchmark.
@@ -169,14 +167,11 @@ def __init__(
             warmup_runs: See :meth:`Benchmark.__init__`.
             timing_runs: See :meth:`Benchmark.__init__`.
             plugin_libraries: See :meth:`Benchmark.__init__`.
-            trtexec_path: Path to trtexec binary. Defaults to 'trtexec' which
-                         looks for the binary in PATH.
             trtexec_args: Additional command-line arguments to pass to trtexec.
                          These are appended after the standard arguments.
                          Example: ['--fp16', '--workspace=4096', '--verbose']
         """
         super().__init__(timing_cache_file, warmup_runs, timing_runs, plugin_libraries)
-        self.trtexec_path = trtexec_path
         self.trtexec_args = trtexec_args if trtexec_args is not None else []
         self.temp_dir = tempfile.mkdtemp(prefix="trtexec_benchmark_")
         self.engine_path = os.path.join(self.temp_dir, "engine.trt")
@@ -186,7 +181,6 @@ def __init__(
         self.latency_pattern = r"\[I\]\s+Latency:.*?median\s*=\s*([\d.]+)\s*ms"
 
         self._base_cmd = [
-            self.trtexec_path,
             f"--avgRuns={self.timing_runs}",
             f"--iterations={self.timing_runs}",
             f"--warmUp={self.warmup_runs}",
@@ -268,13 +262,14 @@ def run(
                 self.logger.debug(f"Wrote model bytes to temporary file: {model_path}")
 
             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
+            full_cmd = ["trtexec", *cmd]
+            self.logger.debug(f"Running: {' '.join(full_cmd)}")
+            result = _run_trtexec(cmd)
             self._write_log_file(
                 log_file,
                 "\n".join(
                     [
-                        f"Command: {' '.join(cmd)}",
+                        f"Command: {' '.join(full_cmd)}",
                         f"Return code: {result.returncode}",
                         "=" * 80,
                         "STDOUT:",
@@ -301,8 +296,9 @@ def run(
             self.logger.info(f"TrtExec benchmark (median): {latency:.2f} ms")
             return latency
         except FileNotFoundError:
-            self.logger.error(f"trtexec binary not found: {self.trtexec_path}")
-            self.logger.error("Please ensure TensorRT is installed and trtexec path is correct")
+            self.logger.error(
+                "'trtexec' binary not found. Please ensure TensorRT is installed and 'trtexec' is in PATH."
+            )
             return float("inf")
         except Exception as e:
             self.logger.error(f"Benchmark failed: {e}")

modelopt/onnx/quantization/ort_utils.py

@@ -46,6 +46,30 @@ def _check_lib_in_ld_library_path(ld_library_path, lib_pattern):
     return False, None
 
 
+def _run_trtexec(
+    args: list[str] | None = None, timeout: float | None = None
+) -> subprocess.CompletedProcess:
+    """Run a 'trtexec' command via subprocess.
+
+    Args:
+        args: Arguments to pass to trtexec (without the 'trtexec' command itself).
+        timeout: Optional subprocess timeout in seconds.
+
+    Returns:
+        The completed subprocess result.
+
+    Raises:
+        FileNotFoundError: If the 'trtexec' binary is not found in PATH.
+    """
+    cmd = ["trtexec", *(args or [])]
+    try:
+        return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)  # nosec B603
+    except FileNotFoundError as e:
+        raise FileNotFoundError(
+            "'trtexec' binary not found. Please ensure TensorRT is installed and 'trtexec' is in PATH."
+        ) from e
+
+
 def _check_for_trtexec(min_version: str = "10.0") -> str:
     """Check if the `trtexec` CLI tool is available in PATH and is >= min_version.
 
@@ -89,7 +113,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(timeout=5)
         banner_output = result.stdout + result.stderr
         parsed_version = _parse_version_from_string(banner_output)
 

modelopt/torch/_deploy/_runtime/tensorrt/constants.py

@@ -32,10 +32,6 @@
 ONE_GIBI_IN_BYTES = 1 << 30
 
 # TensorRT conversion tool names
-TRTEXEC = "trtexec"
-
-# trtexec path within docker
-TRTEXEC_PATH = "trtexec"
 DEFAULT_ARTIFACT_DIR = "modelopt_build/trt_artifacts"
 
 # Default conversion params

modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py

@@ -15,10 +15,9 @@
 
 import logging
 import shutil
-import subprocess  # nosec
 import sys
 from pathlib import Path
-from tempfile import NamedTemporaryFile, TemporaryDirectory, gettempdir
+from tempfile import TemporaryDirectory, gettempdir
 
 from ..._runtime.common import read_bytes, timeit, write_bytes, write_string
 from ..._runtime.tensorrt.layerwise_profiling import process_layerwise_result
@@ -28,7 +27,6 @@
     DEFAULT_NUM_INFERENCE_PER_RUN,
     SHA_256_HASH_LENGTH,
     TRT_MODE_FLAGS,
-    TRTEXEC_PATH,
     WARMUP_TIME_MS,
     TRTMode,
 )
@@ -41,31 +39,50 @@
 )
 
 
-# TODO: Get rid of this function or get approval for `# nosec` usage if we want to include this
-#   as a non-compiled python file in the release.
-def _run_command(cmd: list[str], cwd: Path | None = None) -> tuple[int, bytes]:
-    """Util function to execute a command.
+def _run_trtexec_with_logging(args: list[str], cwd: Path | None = None) -> tuple[int, bytes]:
+    """Run a 'trtexec' command via subprocess, logging the cmd and any failure output.
 
-    This util will not direct stdout and stderr to console if the cmd succeeds.
+    The 'trtexec' binary is hardcoded as the executable; only its arguments may be supplied
+    by the caller. This restricts the function to trtexec invocations.
+
+    Output handling: stdout and stderr are merged and captured in memory.
+    On failure (non-zero returncode) or timeout, the captured output is logged at ERROR level;
+    on success, this function emits nothing to the console.
 
     Args:
-        cmd: the command line list
-        cwd: current working directory
+        args: Arguments to pass to trtexec (without the 'trtexec' command itself).
+        cwd: Optional working directory for the subprocess.
 
     Returns:
-        return code: 0 means successful, otherwise means failed
-        log_string: the stdout and stderr output as a string
+        A tuple of (returncode, output) where output is the combined stdout/stderr bytes.
 
+    Raises:
+        FileNotFoundError: If the 'trtexec' binary is not found in PATH.
+        subprocess.TimeoutExpired: If trtexec does not finish within 60 minutes.
+            The captured output is logged before re-raising.
     """
+    import subprocess  # nosec
+
+    cmd = ["trtexec", *args]
     logging.info(" ".join(cmd))
-    with NamedTemporaryFile("w+b") as log:
-        p = subprocess.Popen(cmd, stdout=log, stderr=log, cwd=str(cwd) if cwd else None)  # nosec
-        p.wait()
-        log.seek(0)
-        output = log.read()
-        if p.returncode != 0:
-            logging.error(output.decode(errors="ignore"))
-        return p.returncode, output
+    try:
+        result = subprocess.run(  # nosec B603 - cmd[0] is hardcoded "trtexec"
+            cmd,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.STDOUT,
+            cwd=str(cwd) if cwd else None,
+            timeout=3600,
+        )
+    except FileNotFoundError as e:
+        raise FileNotFoundError(
+            "'trtexec' binary not found. Please ensure TensorRT is installed and 'trtexec' is in PATH."
+        ) from e
+    except subprocess.TimeoutExpired as e:
+        logging.error((e.stdout or b"").decode(errors="ignore"))
+        raise
+    if result.returncode != 0:
+        logging.error(result.stdout.decode(errors="ignore"))
+    return result.returncode, result.stdout
 
 
 def _get_profiling_params(profiling_runs: int) -> list[str]:
@@ -181,7 +198,7 @@ def _build_command(
         calib_cache_path: Path | None = None,
         timing_cache_path: Path | None = None,
     ) -> list[str]:
-        cmd = [TRTEXEC_PATH, f"--onnx={onnx_path}"]
+        cmd = [f"--onnx={onnx_path}"]
         cmd.extend(TRT_MODE_FLAGS[trt_mode])
 
         if trt_mode == TRTMode.INT8 and calib_cache and calib_cache_path:
@@ -235,7 +252,7 @@ def _setup_files_and_paths(
         cmd = _build_command(onnx_path, engine_path, calib_cache_path, timing_cache_path)
 
         try:
-            ret_code, out = _run_command(cmd)
+            ret_code, out = _run_trtexec_with_logging(cmd)
             if ret_code != 0:
                 return None, out
 
@@ -284,7 +301,7 @@ def profile_engine(
     """
 
     def _build_command(engine_path: Path, profile_path: Path, layer_info_path: Path) -> list[str]:
-        cmd = [TRTEXEC_PATH, f"--loadEngine={engine_path}"]
+        cmd = [f"--loadEngine={engine_path}"]
         cmd += _get_profiling_params(profiling_runs)
 
         if enable_layerwise_profiling:
@@ -320,7 +337,7 @@ def _setup_files_and_paths(tmp_dir_path: Path, engine_hash: str) -> tuple[Path,
         cmd = _build_command(engine_path, profile_path, layer_info_path)
 
         try:
-            ret_code, out = _run_command(cmd)
+            ret_code, out = _run_trtexec_with_logging(cmd)
             if ret_code != 0:
                 return None, out
 

tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py

@@ -55,7 +55,7 @@ def setup_mocks():
 
     with (
         mock.patch(
-            "modelopt.torch._deploy._runtime.tensorrt.engine_builder._run_command"
+            "modelopt.torch._deploy._runtime.tensorrt.engine_builder._run_trtexec_with_logging"
         ) as mock_run,
         mock.patch(
             "modelopt.torch._deploy._runtime.tensorrt.engine_builder.TemporaryDirectory"