fix(autotune): pre-check remote board connectivity before benchmark

CHANGELOG.rst

@@ -21,6 +21,8 @@ Changelog
 
 **Bug Fixes**
 
+- Fix ONNX Autotune remote autotuning to pre-check board connectivity (configurable retries via ``--remote_connection_retries``) before each trtexec invocation. If unreachable, the autotuner saves state and exits cleanly instead of running trtexec and permanently marking schemes as errored.
+
 0.46 (2026-08-xx)
 ^^^^^^^^^^^^^^^^^
 

docs/source/guides/9_autotune.rst

@@ -255,6 +255,23 @@ To use remote autotuning during Q/DQ placement optimization, run with ``trtexec`
 
 Replace ``<remote autotuning config>`` with an actual remote autotuning configuration string (see ``trtexec --help`` for more details). Other TensorRT benchmark options (e.g. ``--timing_cache``, ``--warmup_runs``, ``--timing_runs``, ``--plugin_libraries``) are also available; run ``--help`` for details.
 
+**Connectivity pre-check:**
+
+When ``--remoteAutoTuningConfig`` is detected, the autotuner tests TCP connectivity to the remote board before each trtexec invocation. If the board is unreachable after retries, the autotuner saves state and exits cleanly — preventing transient network failures from permanently marking schemes as errored.
+
+Configure the retry count with ``--remote_connection_retries`` (default: 3):
+
+.. code-block:: bash
+
+   python -m modelopt.onnx.quantization.autotune \
+       --onnx_path model.onnx \
+       --output_dir ./model_remote_autotuned \
+       --use_trtexec \
+       --trtexec_benchmark_args "--remoteAutoTuningConfig=\"ssh://admin@192.168.1.100\" --safe --skipInference" \
+       --remote_connection_retries 5
+
+Each failed attempt is logged as a warning. If all retries fail, the process exits with an error message and preserved state. On restart (same ``--output_dir``), autotuning resumes from where it left off without re-testing already-profiled schemes.
+
 Low-Level API Usage
 ===================
 

modelopt/onnx/quantization/autotune/__main__.py

@@ -98,6 +98,10 @@ def run_autotune() -> int:
     validate_file_path(args.qdq_baseline, "QDQ baseline model")
     output_dir = Path(args.output_dir)
 
+    if not 1 <= args.remote_connection_retries <= 10:
+        logger.error("--remote_connection_retries must be between 1 and 10")
+        return 1
+
     log_benchmark_config(args)
     trtexec_args = getattr(args, "trtexec_benchmark_args", None)
     if trtexec_args and isinstance(trtexec_args, str):
@@ -109,6 +113,7 @@ def run_autotune() -> int:
         warmup_runs=args.warmup_runs,
         timing_runs=args.timing_runs,
         trtexec_args=trtexec_args,
+        remote_connection_retries=args.remote_connection_retries,
     )
 
     if benchmark_instance is None:
@@ -314,6 +319,13 @@ def get_parser() -> argparse.ArgumentParser:
         help="Additional command-line arguments to pass to trtexec as a single quoted string. "
         "Example: --trtexec_benchmark_args '--fp16 --workspace=4096 --verbose'",
     )
+    trt_group.add_argument(
+        "--remote_connection_retries",
+        type=int,
+        default=3,
+        help="Number of TCP connection attempts to the remote board before aborting (1-10). "
+        "Only relevant when --remoteAutoTuningConfig is present in trtexec args (default: 3)",
+    )
 
     # Logging
     parser.add_argument("--verbose", "-v", action="store_true", help="Enable verbose DEBUG logging")

modelopt/onnx/quantization/autotune/benchmark.py

@@ -31,8 +31,10 @@
 import os
 import re
 import shutil
+import socket
 import tempfile
 import time
+import urllib.parse
 from abc import ABC, abstractmethod
 from pathlib import Path
 from typing import Any
@@ -41,6 +43,7 @@
 import torch
 
 from modelopt.onnx.logging_config import logger
+from modelopt.onnx.quantization.autotune.common import RemoteConnectionError
 from modelopt.onnx.quantization.ort_utils import _check_for_trtexec, _run_trtexec
 
 TRT_AVAILABLE = importlib.util.find_spec("tensorrt") is not None
@@ -50,6 +53,71 @@
 TORCH_CUDA_AVAILABLE = torch.cuda.is_available()
 
 
+_DEFAULT_PORTS = {"ssh": 22, "http": 80, "https": 443}
+
+
+def _check_remote_connectivity(trtexec_args: list[str], retries: int = 3) -> None:
+    """Test TCP connectivity to the remote autotuning board before running trtexec.
+
+    Scans trtexec_args for --remoteAutoTuningConfig, parses the URI to extract
+    hostname and port, then attempts a TCP connection with a 5-second timeout.
+    Retries up to `retries` times before raising an error.
+
+    Args:
+        trtexec_args: List of trtexec command-line arguments.
+        retries: Number of connection attempts before giving up (default: 3).
+
+    Raises:
+        RemoteConnectionError: If the remote board is unreachable after all retries.
+    """
+    config_value = None
+    for i, arg in enumerate(trtexec_args):
+        if arg.startswith("--remoteAutoTuningConfig="):
+            config_value = arg.split("=", 1)[1]
+            break
+        elif arg == "--remoteAutoTuningConfig" and i + 1 < len(trtexec_args):
+            config_value = trtexec_args[i + 1]
+            break
+
+    if config_value is None:
+        return
+
+    parsed = urllib.parse.urlparse(config_value)
+    hostname = parsed.hostname
+    if not hostname:
+        return
+
+    port = parsed.port
+    if port is None:
+        port = _DEFAULT_PORTS.get(parsed.scheme, 22)
+
+    last_error = _try_connect(hostname, port, retries)
+    if last_error is not None:
+        raise RemoteConnectionError(
+            f"Cannot reach remote autotuning board at {hostname}:{port} after {retries} attempts - "
+            f"{last_error}. Exiting to avoid marking schemes as errors in state file."
+        ) from last_error
+
+
+def _try_connect(hostname: str, port: int, retries: int) -> Exception | None:
+    """Attempt TCP connection with retries. Returns None on success, last error on failure."""
+    last_error = None
+    for attempt in range(1, retries + 1):
+        try:
+            conn = socket.create_connection((hostname, port), timeout=5)
+            conn.close()
+            return None
+        except (TimeoutError, OSError) as e:  # noqa: PERF203
+            last_error = e
+            if attempt < retries:
+                logger.warning(
+                    f"Remote board connection attempt {attempt}/{retries} failed "
+                    f"({hostname}:{port}): {e}. Retrying..."
+                )
+                time.sleep(2)
+    return last_error
+
+
 def _validate_shape_range(min_shape: list, opt_shape: list, max_shape: list) -> None:
     """Raise ValueError if shape lengths differ or if min <= opt <= max fails at any dimension."""
     if len(min_shape) != len(opt_shape) or len(opt_shape) != len(max_shape):
@@ -159,6 +227,7 @@ def __init__(
         timing_runs: int = 10,
         plugin_libraries: list[str] | None = None,
         trtexec_args: list[str] | None = None,
+        remote_connection_retries: int = 3,
     ):
         """Initialize the trtexec benchmark.
 
@@ -170,8 +239,12 @@ def __init__(
             trtexec_args: Additional command-line arguments to pass to trtexec.
                          These are appended after the standard arguments.
                          Example: ['--fp16', '--workspace=4096', '--verbose']
+            remote_connection_retries: Number of TCP connection attempts to the remote
+                         board before giving up (default: 3). Only used when
+                         --remoteAutoTuningConfig is present in trtexec_args.
         """
         super().__init__(timing_cache_file, warmup_runs, timing_runs, plugin_libraries)
+        self._remote_connection_retries = remote_connection_retries
         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")
@@ -253,6 +326,8 @@ def run(
         if not os.path.exists(self.timing_cache_file):
             self.logger.debug(f"Will create timing cache: {self.timing_cache_file}")
 
+        _check_remote_connectivity(self._base_cmd, retries=self._remote_connection_retries)
+
         try:
             model_path = path_or_bytes
             if isinstance(model_path, bytes):

modelopt/onnx/quantization/autotune/common.py

@@ -47,6 +47,10 @@ class InvalidSchemeError(AutotunerError):
     """Exception raised when an invalid scheme is referenced."""
 
 
+class RemoteConnectionError(AutotunerError):
+    """Exception raised when the remote autotuning board is unreachable."""
+
+
 class RegionType(Enum):
     """Region type enumeration for hierarchical graph structure.
 

modelopt/onnx/quantization/autotune/workflows.py

@@ -29,7 +29,7 @@
 from modelopt.onnx.logging_config import logger
 from modelopt.onnx.quantization.autotune.autotuner import QDQAutotuner
 from modelopt.onnx.quantization.autotune.benchmark import TensorRTPyBenchmark, TrtExecBenchmark
-from modelopt.onnx.quantization.autotune.common import Config, PatternCache
+from modelopt.onnx.quantization.autotune.common import Config, PatternCache, RemoteConnectionError
 from modelopt.onnx.quantization.qdq_utils import get_quantized_tensors
 
 _benchmark_instance = None
@@ -75,6 +75,8 @@ def benchmark_onnx_model(
         logger.debug(f"Benchmark result: {latency:.2f} ms")
         return latency
 
+    except RemoteConnectionError:
+        raise
     except Exception as e:
         logger.error(f"Benchmark error: {e}", exc_info=True)
         return float("inf")
@@ -87,6 +89,7 @@ def init_benchmark_instance(
     warmup_runs: int = 5,
     timing_runs: int = 20,
     trtexec_args: list[str] | None = None,
+    remote_connection_retries: int = 3,
 ):
     """Initialize global TensorRT benchmark instance for model performance measurement.
 
@@ -103,6 +106,9 @@ def init_benchmark_instance(
                     Higher values give more stable median (default: 20)
         trtexec_args: Additional command-line arguments to pass to trtexec as a string (only used if use_trtexec=True).
                      Example: '--fp16 --workspace=4096 --verbose'
+        remote_connection_retries: Number of TCP connection attempts to the remote board
+                     before aborting (default: 3). Only relevant when --remoteAutoTuningConfig
+                     is present in trtexec_args.
     """
     global _benchmark_instance
     try:
@@ -113,6 +119,7 @@ def init_benchmark_instance(
                 timing_runs=timing_runs,
                 plugin_libraries=plugin_libraries,
                 trtexec_args=trtexec_args,
+                remote_connection_retries=remote_connection_retries,
             )
             logger.info("Trtexec benchmark initialized")
         else:
@@ -330,9 +337,14 @@ def region_pattern_autotuning_workflow(
             model_bytes = autotuner.export_onnx(None, insert_qdq=True)
             test_log = logs_dir / f"region_{region.id}_scheme_{scheme_idx}.log"
             flush_timing_cache = (iteration_count % 10) == 0
-            latency = benchmark_onnx_model(
-                model_bytes, str(test_log), flush_timing_cache=flush_timing_cache
-            )
+            try:
+                latency = benchmark_onnx_model(
+                    model_bytes, str(test_log), flush_timing_cache=flush_timing_cache
+                )
+            except RemoteConnectionError:
+                logger.error("Remote board connection lost, saving state before exit")
+                autotuner.save_state(str(state_path))
+                raise
 
             autotuner.submit(latency, success=(latency != float("inf")))