From b670685403ba5aeb77489a65f31c6a1945e8ee9d Mon Sep 17 00:00:00 2001 From: Will Guo Date: Wed, 5 Aug 2026 03:44:14 -0700 Subject: [PATCH 1/4] fix(autotune): pre-check remote board connectivity before benchmark When using --remoteAutoTuningConfig, test TCP connectivity to the remote board before each trtexec invocation. If unreachable after configurable retries (--remote_connection_retries, default 3), save state and exit cleanly instead of running trtexec which would fail and permanently mark schemes as errored in autotune_states.yaml. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Will Guo --- CHANGELOG.rst | 1 + docs/source/guides/9_autotune.rst | 17 +++++ .../onnx/quantization/autotune/__main__.py | 12 +++ .../onnx/quantization/autotune/benchmark.py | 75 +++++++++++++++++++ modelopt/onnx/quantization/autotune/common.py | 4 + .../onnx/quantization/autotune/workflows.py | 20 ++++- 6 files changed, 125 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 632b5532ebf..91762799edc 100755 --- a/CHANGELOG.rst +++ b/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) ^^^^^^^^^^^^^^^^^ diff --git a/docs/source/guides/9_autotune.rst b/docs/source/guides/9_autotune.rst index 4c7f7b4d172..a56c78eec7c 100644 --- a/docs/source/guides/9_autotune.rst +++ b/docs/source/guides/9_autotune.rst @@ -255,6 +255,23 @@ To use remote autotuning during Q/DQ placement optimization, run with ``trtexec` Replace ```` 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 =================== diff --git a/modelopt/onnx/quantization/autotune/__main__.py b/modelopt/onnx/quantization/autotune/__main__.py index ccb826c291b..40139124da8 100644 --- a/modelopt/onnx/quantization/autotune/__main__.py +++ b/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") diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py index ba5cf1142bf..59e3efbada9 100644 --- a/modelopt/onnx/quantization/autotune/benchmark.py +++ b/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): diff --git a/modelopt/onnx/quantization/autotune/common.py b/modelopt/onnx/quantization/autotune/common.py index 31983423cd9..098bb15dc66 100644 --- a/modelopt/onnx/quantization/autotune/common.py +++ b/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. diff --git a/modelopt/onnx/quantization/autotune/workflows.py b/modelopt/onnx/quantization/autotune/workflows.py index 1059cb6cc22..689911e1300 100644 --- a/modelopt/onnx/quantization/autotune/workflows.py +++ b/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"))) From 5d007dc9547eea776fcc653f91179457a880a906 Mon Sep 17 00:00:00 2001 From: Will Guo Date: Thu, 6 Aug 2026 01:42:41 -0700 Subject: [PATCH 2/4] fix(autotune): harden remote URI parsing and mid-trtexec connection loss - Strip matching surrounding quotes from config_value before urlparse (handles embedded quotes from --remoteAutoTuningConfig=\"...\") - Raise RemoteConnectionError on missing hostname or invalid port instead of silently returning - Add _benchmark_failed() to re-check connectivity after trtexec failure; raises RemoteConnectionError if board dropped mid-execution so the scheme is not incorrectly recorded as inf - Wrap baseline and final benchmark_onnx_model calls with RemoteConnectionError handling that saves autotuner state before re-raising Co-Authored-By: Claude Opus 4.6 --- .../onnx/quantization/autotune/benchmark.py | 36 ++++++++++++++++--- .../onnx/quantization/autotune/workflows.py | 16 +++++++-- 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py index 59e3efbada9..848f2bf03fb 100644 --- a/modelopt/onnx/quantization/autotune/benchmark.py +++ b/modelopt/onnx/quantization/autotune/benchmark.py @@ -82,12 +82,19 @@ def _check_remote_connectivity(trtexec_args: list[str], retries: int = 3) -> Non if config_value is None: return + config_value = config_value.strip("'\"") + parsed = urllib.parse.urlparse(config_value) hostname = parsed.hostname if not hostname: - return + raise RemoteConnectionError(f"Missing hostname in remote autotuning URI: {config_value!r}") - port = parsed.port + try: + port = parsed.port + except ValueError as e: + raise RemoteConnectionError( + f"Invalid port in remote autotuning URI: {config_value!r} - {e}" + ) from e if port is None: port = _DEFAULT_PORTS.get(parsed.scheme, 22) @@ -361,15 +368,17 @@ 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." @@ -377,7 +386,24 @@ def run( 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): diff --git a/modelopt/onnx/quantization/autotune/workflows.py b/modelopt/onnx/quantization/autotune/workflows.py index 689911e1300..e9df38e355d 100644 --- a/modelopt/onnx/quantization/autotune/workflows.py +++ b/modelopt/onnx/quantization/autotune/workflows.py @@ -295,7 +295,12 @@ 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)) + raise autotuner.submit(baseline_latency) logger.info(f"Baseline: {baseline_latency:.2f} ms") else: @@ -377,7 +382,14 @@ 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)) + raise if final_latency > 0 and final_latency != float("inf"): speedup = baseline_latency / final_latency From 0c40465172f6cec082c37eb7b84f0d7945b0fe04 Mon Sep 17 00:00:00 2001 From: Will Guo Date: Thu, 6 Aug 2026 09:24:23 +0000 Subject: [PATCH 3/4] resolve comments Signed-off-by: Will Guo --- modelopt/onnx/quantization/autotune/benchmark.py | 6 ++++-- modelopt/onnx/quantization/autotune/workflows.py | 3 +++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/modelopt/onnx/quantization/autotune/benchmark.py b/modelopt/onnx/quantization/autotune/benchmark.py index 848f2bf03fb..28f64ccdbaf 100644 --- a/modelopt/onnx/quantization/autotune/benchmark.py +++ b/modelopt/onnx/quantization/autotune/benchmark.py @@ -87,13 +87,15 @@ def _check_remote_connectivity(trtexec_args: list[str], retries: int = 3) -> Non parsed = urllib.parse.urlparse(config_value) hostname = parsed.hostname if not hostname: - raise RemoteConnectionError(f"Missing hostname in remote autotuning URI: {config_value!r}") + 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: {config_value!r} - {e}" + f"Invalid port in remote autotuning URI: {parsed.scheme}://{hostname} - {e}" ) from e if port is None: port = _DEFAULT_PORTS.get(parsed.scheme, 22) diff --git a/modelopt/onnx/quantization/autotune/workflows.py b/modelopt/onnx/quantization/autotune/workflows.py index e9df38e355d..0a4134768a4 100644 --- a/modelopt/onnx/quantization/autotune/workflows.py +++ b/modelopt/onnx/quantization/autotune/workflows.py @@ -300,6 +300,7 @@ def region_pattern_autotuning_workflow( 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}, use --resume to continue") raise autotuner.submit(baseline_latency) logger.info(f"Baseline: {baseline_latency:.2f} ms") @@ -349,6 +350,7 @@ def region_pattern_autotuning_workflow( 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}, use --resume to continue") raise autotuner.submit(latency, success=(latency != float("inf"))) @@ -389,6 +391,7 @@ def region_pattern_autotuning_workflow( "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}, use --resume to continue") raise if final_latency > 0 and final_latency != float("inf"): From c045c3de424d2bf1dd85501f12371f9420580867 Mon Sep 17 00:00:00 2001 From: Will Guo Date: Fri, 7 Aug 2026 00:57:57 -0700 Subject: [PATCH 4/4] fix(autotune): correct recovery message to say re-run instead of --resume The parser doesn't have a --resume option; state is detected automatically on re-run. Update the user-facing message accordingly. Co-Authored-By: Claude Opus 4.6 Signed-off-by: Will Guo --- modelopt/onnx/quantization/autotune/workflows.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/modelopt/onnx/quantization/autotune/workflows.py b/modelopt/onnx/quantization/autotune/workflows.py index 0a4134768a4..a8e91037885 100644 --- a/modelopt/onnx/quantization/autotune/workflows.py +++ b/modelopt/onnx/quantization/autotune/workflows.py @@ -300,7 +300,9 @@ def region_pattern_autotuning_workflow( 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}, use --resume to continue") + 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") @@ -350,7 +352,9 @@ def region_pattern_autotuning_workflow( 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}, use --resume to continue") + 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"))) @@ -391,7 +395,9 @@ def region_pattern_autotuning_workflow( "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}, use --resume to continue") + 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"):