diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py index df6dbc877d0..b87478a1572 100644 --- a/modelopt/onnx/quantization/autotune/benchmark.py +++ b/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}") diff --git a/modelopt/onnx/quantization/ort_utils.py b/modelopt/onnx/quantization/ort_utils.py index 2c5a0b7d2da..f7799c634f0 100755 --- a/modelopt/onnx/quantization/ort_utils.py +++ b/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) diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/constants.py b/modelopt/torch/_deploy/_runtime/tensorrt/constants.py index c4f387482e9..d9ace1645a3 100644 --- a/modelopt/torch/_deploy/_runtime/tensorrt/constants.py +++ b/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 diff --git a/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py b/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py index 055a1f26b27..bb8bbd292b8 100644 --- a/modelopt/torch/_deploy/_runtime/tensorrt/engine_builder.py +++ b/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 diff --git a/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py b/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py index 38fce51f4aa..ff7f77cf617 100755 --- a/tests/unit/torch/deploy/_runtime/tensorrt/test_engine_builder.py +++ b/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"