cleanup trtexec based benchmarking code

Signed-off-by: dmoodie <dmoodie@nvidia.com>

Author: dthienan-nv

modelopt/onnx/quantization/autotune/benchmark.py

@@ -36,6 +36,7 @@
 import tempfile
 import time
 from abc import ABC, abstractmethod
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Any
 from urllib.parse import parse_qs, urlparse
@@ -153,6 +154,178 @@ def _write_log_file(self, file: Path | str | None, content: str) -> None:
 )
 _STD_PATTERN = r"\[I\]\s+GPU Compute Time:.*?median\s*=\s*([\d.]+)\s*ms"
 
+_URL_PASSWORD_RE = re.compile(r"(://[^:/?#@]+):[^@/?#]+@")
+
+
+def _redact_url_password(s: str) -> str:
+    """Replace any ``scheme://user:password@host`` substring with ``user:******@host``.
+
+    Used so SSH passwords supplied via ``--remoteAutoTuningConfig`` don't leak
+    into log messages or exception strings.
+    """
+    return _URL_PASSWORD_RE.sub(r"\1:******@", s)
+
+
+def _build_base_trtexec_cmd(
+    *,
+    timing_runs: int,
+    warmup_runs: int,
+    engine_path: str,
+    timing_cache_file: str,
+    plugin_libraries: list[str] | None = None,
+    log: Any = None,
+) -> list[str]:
+    """Build the static portion of the trtexec command line (no ``--onnx=`` yet).
+
+    Plugin libraries that don't exist on disk are skipped with a warning if a
+    logger is supplied. The leading ``trtexec`` binary path is not included —
+    the caller is responsible for prepending it.
+
+    Args:
+        timing_runs: Value for ``--avgRuns`` and ``--iterations``.
+        warmup_runs: Value for ``--warmUp``.
+        engine_path: Path used for ``--saveEngine=``.
+        timing_cache_file: Path used for ``--timingCacheFile=``.
+        plugin_libraries: Paths to ``.so`` libraries for ``--staticPlugins``.
+        log: Optional logger used to warn about missing plugins and trace adds.
+    """
+    cmd = [
+        f"--avgRuns={timing_runs}",
+        f"--iterations={timing_runs}",
+        f"--warmUp={warmup_runs}",
+        "--stronglyTyped",
+        f"--saveEngine={engine_path}",
+        f"--timingCacheFile={timing_cache_file}",
+    ]
+    for plugin_lib in plugin_libraries or []:
+        plugin_path = Path(plugin_lib).resolve()
+        if not plugin_path.exists():
+            if log is not None:
+                log.warning(f"Plugin library not found: {plugin_path}")
+            continue
+        cmd.append(f"--staticPlugins={plugin_path}")
+        if log is not None:
+            log.debug(f"Added plugin library: {plugin_path}")
+    return cmd
+
+
+def _extract_remote_config_value(trtexec_args: list[str], *, log: Any = None) -> str | None:
+    """Find the value of ``--remoteAutoTuningConfig`` in ``trtexec_args``.
+
+    Supports both inline (``--remoteAutoTuningConfig=value``) and split
+    (``--remoteAutoTuningConfig value``) forms.
+
+    Returns:
+        The value as a string, or ``None`` if the flag is absent. Returning
+        an empty string is possible (e.g. ``--remoteAutoTuningConfig=``); the
+        caller decides whether to treat that as an error.
+
+    Raises:
+        ValueError: If the flag appears more than once, has no value at the
+            end of the list, or is malformed (e.g. missing the ``=``
+            separator). SSH passwords in malformed args are redacted before
+            being included in the error or debug log.
+    """
+    matches = [a for a in trtexec_args if "--remoteAutoTuningConfig" in a]
+    if not matches:
+        return None
+    if len(matches) != 1:
+        raise ValueError("Exactly one --remoteAutoTuningConfig argument is required")
+
+    for i, arg in enumerate(trtexec_args):
+        if not arg.startswith("--remoteAutoTuningConfig"):
+            continue
+        if arg == "--remoteAutoTuningConfig":
+            if i + 1 >= len(trtexec_args):
+                raise ValueError("Missing value for --remoteAutoTuningConfig")
+            return trtexec_args[i + 1]
+        if arg.startswith("--remoteAutoTuningConfig="):
+            return arg.split("=", 1)[1]
+        # Malformed: starts with the flag name but neither uses ``=`` nor is
+        # the bare flag. Redact any embedded SSH password before surfacing.
+        redacted_arg = _redact_url_password(arg)
+        if log is not None:
+            log.debug(f"Parsing remoteAutoTuningConfig arg: {redacted_arg}")
+        raise ValueError(f"Malformed --remoteAutoTuningConfig argument: {redacted_arg}")
+    return None  # pragma: no cover — unreachable; ``matches`` proved presence
+
+
+@dataclass(frozen=True)
+class _RemoteAutotuningConfig:
+    """Resolved remote-autotuning destination parsed from a ``ssh://`` URL."""
+
+    user: str
+    password: str  # may be empty when no password was supplied
+    ip: str
+    port: int
+    options: dict[str, str]
+    bin_path: str  # dirname of ``remote_exec_path``
+    lib_path: str  # value of ``remote_lib_path``
+
+
+def _parse_remote_autotuning_url(url: str) -> _RemoteAutotuningConfig:
+    """Parse a ``--remoteAutoTuningConfig`` URL into structured fields.
+
+    Required URL form::
+
+        ssh://user[:password]@host[:port]?remote_exec_path=PATH&remote_lib_path=PATH
+
+    Raises:
+        ValueError: If the scheme is not ``ssh://``; if user or host are
+            missing; or if required query parameters are missing or
+            duplicated. Duplicate keys are rejected explicitly because
+            silently collapsing them would produce empty remote paths
+            downstream.
+    """
+    if not url.startswith("ssh://"):
+        raise ValueError("Only 'ssh://' remote autotuning config URLs are supported")
+    parsed = urlparse(url)
+    if parsed.username is None:
+        raise ValueError("Unable to parse remote user from --remoteAutoTuningConfig")
+    if parsed.hostname is None:
+        raise ValueError("Unable to parse remote IP from --remoteAutoTuningConfig")
+
+    parsed_query = parse_qs(parsed.query)
+    duplicates = sorted(k for k, v in parsed_query.items() if len(v) > 1)
+    if duplicates:
+        raise ValueError(f"Duplicate query parameters in --remoteAutoTuningConfig: {duplicates}")
+    options = {k: v[0] for k, v in parsed_query.items()}
+
+    required_params = ["remote_exec_path", "remote_lib_path"]
+    missing = [p for p in required_params if p not in options]
+    if missing:
+        raise ValueError(
+            f"Missing required query parameters in --remoteAutoTuningConfig: {missing}"
+        )
+
+    return _RemoteAutotuningConfig(
+        user=parsed.username,
+        password=parsed.password or "",
+        ip=parsed.hostname,
+        port=parsed.port if parsed.port is not None else 22,
+        options=options,
+        bin_path=os.path.dirname(options["remote_exec_path"]),
+        lib_path=options["remote_lib_path"],
+    )
+
+
+def _ensure_remote_autotuning_flags(trtexec_args: list[str], *, log: Any = None) -> list[str]:
+    """Return ``trtexec_args`` with ``--safe`` and ``--skipInference`` appended if missing.
+
+    Remote autotuning requires both flags. A warning is emitted for each flag
+    that has to be injected so the user sees that their argv was modified.
+    """
+    result = list(trtexec_args)
+    for flag in ("--safe", "--skipInference"):
+        if flag in result:
+            continue
+        if log is not None:
+            log.warning(
+                f"Remote autotuning requires '{flag}' to be set. Adding it to trtexec arguments."
+            )
+        result.append(flag)
+    return result
+
 
 class TrtExecBenchmark(Benchmark):
     """TensorRT benchmark using trtexec command-line tool.
@@ -182,120 +355,58 @@ def __init__(
                          Example: ['--fp16', '--workspace=4096', '--verbose']
         """
         super().__init__(timing_cache_file, warmup_runs, timing_runs, plugin_libraries)
-        self.trtexec_args = trtexec_args if trtexec_args is not None else []
+        self.trtexec_args = list(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")
         self.temp_model_path = os.path.join(self.temp_dir, "temp_model.onnx")
         self.logger.debug(f"Created temporary engine directory: {self.temp_dir}")
         self.logger.debug(f"Temporary model path: {self.temp_model_path}")
 
-        self._base_cmd = [
-            f"--avgRuns={self.timing_runs}",
-            f"--iterations={self.timing_runs}",
-            f"--warmUp={self.warmup_runs}",
-            "--stronglyTyped",
-            f"--saveEngine={self.engine_path}",
-            f"--timingCacheFile={self.timing_cache_file}",
-        ]
-
-        for plugin_lib in self.plugin_libraries:
-            plugin_path = Path(plugin_lib).resolve()
-            if not plugin_path.exists():
-                self.logger.warning(f"Plugin library not found: {plugin_path}")
-                continue
-            self._base_cmd.append(f"--staticPlugins={plugin_path}")
-            self.logger.debug(f"Added plugin library: {plugin_path}")
+        self._base_cmd = _build_base_trtexec_cmd(
+            timing_runs=self.timing_runs,
+            warmup_runs=self.warmup_runs,
+            engine_path=self.engine_path,
+            timing_cache_file=self.timing_cache_file,
+            plugin_libraries=self.plugin_libraries,
+            log=self.logger,
+        )
 
-        trtexec_args = self.trtexec_args or []
-        self.has_remote_config = any("--remoteAutoTuningConfig" in arg for arg in trtexec_args)
+        # Defaults for remote-autotuning fields; overwritten when configured.
+        self.has_remote_config: bool = False
         self.remote_ip: str | None = None
         self.remote_port: int = 22
         self.remote_user: str = "root"
         self.remote_password: str = ""
         self.remote_engine_path: str = "trtexec_benchmark_model.trt"
         self.remote_bin_path: str = "trtexec"
+        self.remote_lib_path: str = ""
+        self.remote_options: dict[str, str] = {}
 
-        if self.has_remote_config:
-            remote_config = [arg for arg in trtexec_args if "--remoteAutoTuningConfig" in arg]
-            if len(remote_config) != 1:
-                raise ValueError("Exactly one --remoteAutoTuningConfig argument is required")
-            # Parse --remoteAutoTuningConfig argument, which may be given as:
-            #  ('--remoteAutoTuningConfig=ssh://user:pass@host:port?...') or
-            #  ('--remoteAutoTuningConfig', 'ssh://user:pass@host:port?...')
-            #
-            # The logic: find the arg starting with '--remoteAutoTuningConfig'
-            # If formatted as '--remoteAutoTuningConfig=...', split off the '='
-            # Otherwise, grab the next argument.
-            config_arg_value: str | None = None
-            for i, arg in enumerate(trtexec_args):
-                if arg.startswith("--remoteAutoTuningConfig"):
-                    if arg == "--remoteAutoTuningConfig":
-                        # Value should be the next argument
-                        if i + 1 < len(trtexec_args):
-                            config_arg_value = trtexec_args[i + 1]
-                        else:
-                            raise ValueError("Missing value for --remoteAutoTuningConfig")
-                    elif arg.startswith("--remoteAutoTuningConfig="):
-                        config_arg_value = arg.split("=", 1)[1]
-                    else:
-                        raise ValueError(f"Malformed --remoteAutoTuningConfig argument: {arg}")
-                    break
-            if not config_arg_value:
+        remote_value = _extract_remote_config_value(self.trtexec_args, log=self.logger)
+        if remote_value is not None:
+            self.has_remote_config = True
+            if not remote_value:
                 raise ValueError("Could not parse --remoteAutoTuningConfig argument")
-            remote_config_str: str = config_arg_value
-
-            if not remote_config_str.startswith("ssh://"):
-                raise ValueError("Only 'ssh://' remote autotuning config URLs are supported")
-            parsed = urlparse(remote_config_str)
-            self.remote_user = parsed.username
-            self.remote_password = parsed.password
-            self.remote_ip = parsed.hostname
-            self.remote_port = parsed.port
-            if self.remote_user is None:
-                raise ValueError("Unable to parse remote user from --remoteAutoTuningConfig")
-            if self.remote_ip is None:
-                raise ValueError("Unable to parse remote IP from --remoteAutoTuningConfig")
-            if self.remote_port is None:
-                self.remote_port = 22
-            # Parse query options into a dict. Reject duplicate keys: silently
-            # collapsing them would let a stray ``?remote_exec_path=a&remote_exec_path=b``
-            # produce ``os.path.dirname(str(['a', 'b']))`` == '' downstream.
-            parsed_query = parse_qs(parsed.query)
-            duplicates = sorted(k for k, v in parsed_query.items() if len(v) > 1)
-            if duplicates:
-                raise ValueError(
-                    f"Duplicate query parameters in --remoteAutoTuningConfig: {duplicates}"
-                )
-            self.remote_options = {k: v[0] for k, v in parsed_query.items()}
-            required_params = ["remote_exec_path", "remote_lib_path"]
-            missing = [p for p in required_params if p not in self.remote_options]
-            if missing:
-                raise ValueError(
-                    f"Missing required query parameters in --remoteAutoTuningConfig: {missing}"
-                )
-            self.remote_bin_path = os.path.dirname(str(self.remote_options["remote_exec_path"]))
-            self.remote_lib_path = str(self.remote_options["remote_lib_path"])
+            config = _parse_remote_autotuning_url(remote_value)
+            self.remote_user = config.user
+            self.remote_password = config.password
+            self.remote_ip = config.ip
+            self.remote_port = config.port
+            self.remote_options = config.options
+            self.remote_bin_path = config.bin_path
+            self.remote_lib_path = config.lib_path
             try:
                 _check_for_trtexec(min_version="10.15")
                 self.logger.debug("TensorRT Python API version >= 10.15 detected")
-                if "--safe" not in trtexec_args:
-                    self.logger.warning(
-                        "Remote autotuning requires '--safe' to be set. Adding it to trtexec arguments."
-                    )
-                    self.trtexec_args.append("--safe")
-                if "--skipInference" not in trtexec_args:
-                    self.logger.warning(
-                        "Remote autotuning requires '--skipInference' to be set. Adding it to trtexec arguments."
-                    )
-                    self.trtexec_args.append("--skipInference")
-            except ImportError as e:
+            except ImportError:
                 self.logger.warning(
                     "Remote autotuning is not supported with TensorRT version < 10.15."
                 )
-                raise e
+                raise
+            self.trtexec_args = _ensure_remote_autotuning_flags(self.trtexec_args, log=self.logger)
 
         self.is_safe = "--safe" in self.trtexec_args
-        self._base_cmd.extend(trtexec_args)
+        self._base_cmd.extend(self.trtexec_args)
 
         self.logger.debug(f"Base command template: {' '.join(self._base_cmd)}")