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

CHANGELOG.rst

@@ -22,6 +22,7 @@ Changelog
 **Bug Fixes**
 
 - Fix EAGLE-3 training with context parallelism (``--cp_size > 1`` in ``examples/speculative_decoding``), which failed to start on ``accelerate >= 1.13`` and then raised ``got mixed torch.Tensor and DTensor``.
+- 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-17)
 ^^^^^^^^^^^^^^^^^

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,80 @@
 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
+
+    config_value = config_value.strip("'\"")
+
+    parsed = urllib.parse.urlparse(config_value)
+    hostname = parsed.hostname
+    if not hostname:
+        raise RemoteConnectionError(
+            f"Missing hostname in remote autotuning URI (scheme={parsed.scheme!r})"
+        )
+
+    try:
+        port = parsed.port
+    except ValueError as e:
+        raise RemoteConnectionError(
+            f"Invalid port in remote autotuning URI: {parsed.scheme}://{hostname} - {e}"
+        ) from e
+    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 +236,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 +248,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 +335,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):
@@ -286,23 +370,42 @@ def run(
             if result.returncode != 0:
                 self.logger.error(f"trtexec failed with return code {result.returncode}")
                 self.logger.error(f"stderr: {result.stderr}")
-                return float("inf")
+                return self._benchmark_failed()
 
             if not (match := re.search(self.latency_pattern, result.stdout, re.IGNORECASE)):
                 self.logger.warning("Could not parse median latency from trtexec output")
                 self.logger.debug(f"trtexec stdout:\n{result.stdout}")
-                return float("inf")
+                return self._benchmark_failed()
             latency = float(match.group(1))
             self.logger.info(f"TrtExec benchmark (median): {latency:.2f} ms")
             return latency
+        except RemoteConnectionError:
+            raise
         except FileNotFoundError:
             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}")
-            return float("inf")
+            return self._benchmark_failed()
+
+    def _benchmark_failed(self) -> float:
+        """Classify a trtexec failure as transient (board lost) or genuine.
+
+        The pre-check in run() passed, but the board can drop while trtexec is
+        running. Re-test connectivity: if the board is gone the failure is a
+        network event, not a property of the scheme, and must not be recorded
+        as an error in the state file.
+
+        Returns:
+            float('inf') if the board is still reachable (genuine scheme failure).
+
+        Raises:
+            RemoteConnectionError: If the remote board is no longer reachable.
+        """
+        _check_remote_connectivity(self._base_cmd, retries=self._remote_connection_retries)
+        return float("inf")
 
 
 class TensorRTPyBenchmark(Benchmark):

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:
@@ -288,7 +295,15 @@ def region_pattern_autotuning_workflow(
         baseline_path = output_dir / "baseline.onnx"
         autotuner.export_onnx(str(baseline_path), insert_qdq=False)
         baseline_log = logs_dir / "baseline.log"
-        baseline_latency = benchmark_onnx_model(str(baseline_path), str(baseline_log))
+        try:
+            baseline_latency = benchmark_onnx_model(str(baseline_path), str(baseline_log))
+        except RemoteConnectionError:
+            logger.error("Remote board connection lost during baseline, saving state before exit")
+            autotuner.save_state(str(state_path))
+            logger.info(
+                f"State saved to {state_path}, re-run the same command to resume from this checkpoint"
+            )
+            raise
         autotuner.submit(baseline_latency)
         logger.info(f"Baseline: {baseline_latency:.2f} ms")
     else:
@@ -330,9 +345,17 @@ 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))
+                logger.info(
+                    f"State saved to {state_path}, re-run the same command to resume from this checkpoint"
+                )
+                raise
 
             autotuner.submit(latency, success=(latency != float("inf")))
 
@@ -365,7 +388,17 @@ def region_pattern_autotuning_workflow(
     final_model_path = output_dir / "optimized_final.onnx"
     autotuner.export_onnx(str(final_model_path), insert_qdq=True)
     final_log = logs_dir / "final.log"
-    final_latency = benchmark_onnx_model(str(final_model_path), str(final_log))
+    try:
+        final_latency = benchmark_onnx_model(str(final_model_path), str(final_log))
+    except RemoteConnectionError:
+        logger.error(
+            "Remote board connection lost during final measurement, saving state before exit"
+        )
+        autotuner.save_state(str(state_path))
+        logger.info(
+            f"State saved to {state_path}, re-run the same command to resume from this checkpoint"
+        )
+        raise
 
     if final_latency > 0 and final_latency != float("inf"):
         speedup = baseline_latency / final_latency