Rerun known transient failures once: HF/network errors and Blackwell CUDA crashes

pyproject.toml

@@ -122,6 +122,9 @@ dev-test = [
     "pytest~=9.1.0",
     "pytest-cov~=7.1.0",
     "pytest-instafail==0.5.0",
+    # MPL-2.0 (weak copyleft) is fine: test-only and used unmodified, and its copyleft covers only
+    # its own files, not code that uses it (tqdm, a core dependency, is MPL-2.0 too)
+    "pytest-rerunfailures~=16.7", # transient-failure reruns (tests/conftest.py); 16.7 matches the exception chain
     "pytest-timeout~=2.4.0",
     "pytest-xdist~=3.8.0",
     "uv",
@@ -357,7 +360,7 @@ skips = [
 # Default additional options
 # Show a short test summary info for all except passed tests with -ra flag
 # print execution time for 50 slowest tests and generate coverage reports
-addopts = "-v -ra --instafail --cov-report=term-missing --cov-report=html --cov-report=xml:coverage.xml --cov-config=pyproject.toml --durations=50 --strict-markers"
+addopts = "-v -ra --instafail --rerun-show-tracebacks --cov-report=term-missing --cov-report=html --cov-report=xml:coverage.xml --cov-config=pyproject.toml --durations=50 --strict-markers"
 pythonpath = ["tests/"]
 # Apply per-test timeouts (see tests/conftest.py) to the test call only, not fixture setup/teardown
 timeout_func_only = true

tests/_test_utils/examples/run_command.py

@@ -19,45 +19,14 @@
 import signal
 import subprocess
 import threading
-import time
-import traceback
 import warnings
 from pathlib import Path
 
 from _test_utils.torch.distributed.utils import get_free_port
+from _test_utils.transient_errors import TransientError, find_transient_marker
 
 MODELOPT_ROOT = Path(__file__).parents[3]
 
-# Substrings that identify a *transient* failure a rerun clears: a HuggingFace Hub / dataset access
-# error, or a rendezvous port that another process bound between its selection and its use
-_TRANSIENT_MARKERS = (
-    "HfHubHTTPError",
-    "Too Many Requests",
-    "500 Server Error",
-    "502 Bad Gateway",
-    "503 Server Error",
-    "504 Server Error",
-    "Bad Gateway",
-    "Service Unavailable",
-    "Gateway Time-out",
-    "ConnectionError",
-    "ReadTimeout",
-    "ConnectTimeout",
-    "Max retries exceeded",
-    "NewConnectionError",
-    "Connection reset by peer",
-    "Connection aborted",
-    "Consistency check failed",  # partial / interrupted HF download
-    "couldn't connect to 'https://huggingface.co'",
-    "Couldn't reach",  # datasets: "Couldn't reach <repo> on the Hub"
-    "Temporary failure in name resolution",  # transient DNS
-    "Name or service not known",  # transient DNS
-    "EADDRINUSE",  # torch.distributed's TCPStore: "name: EADDRINUSE, message: address already in use"
-    "Address already in use",  # the same from Python's OSError, gloo and NCCL
-)
-_HF_MAX_RETRIES = 1
-_HF_RETRY_DELAY_S = 10
-
 
 def extend_cmd_parts(cmd_parts: list[str], **kwargs):
     for key, value in kwargs.items():
@@ -75,7 +44,7 @@ def extend_cmd_parts(cmd_parts: list[str], **kwargs):
 
 
 def _run_capturing(cmd_parts: list[str], cwd: Path, env: dict[str, str]) -> tuple[int, str]:
-    """Run a command, streaming and capturing combined stdout/stderr to catch transient HF errors.
+    """Run a command, streaming and capturing combined stdout/stderr to catch transient errors.
 
     Drained by a reader thread rather than ``subprocess.run``, which waits for EOF: a descendant
     outliving the command keeps the pipe open, hanging the read and discarding every log line.
@@ -151,15 +120,23 @@ def set_in_process_runner(runner) -> None:
     _in_process_runner = runner
 
 
+def _exception_chain_text(exc: BaseException | None) -> str:
+    """``Type: message`` of ``exc`` and of each exception it was raised from or while handling."""
+    lines, seen = [], set()
+    while exc is not None and id(exc) not in seen:
+        seen.add(id(exc))
+        lines.append(f"{type(exc).__name__}: {exc}")
+        exc = exc.__cause__ or (None if exc.__suppress_context__ else exc.__context__)
+    return "\n".join(lines)
+
+
 def run_example_command(
     cmd_parts: list[str],
     example_path: str,
     setup_free_port: bool = False,
     env: dict[str, str] | None = None,
-    hf_max_retries: int = _HF_MAX_RETRIES,
-    hf_retry_delay_s: int = _HF_RETRY_DELAY_S,
-) -> str | None:
-    """Run an example command, retrying transient errors (see ``_TRANSIENT_MARKERS``) once."""
+) -> str:
+    """Run an example command, raising ``TransientError`` on a transient failure so the test reruns."""
     print(f"[{example_path}] Running command: {cmd_parts}")
     in_process = _in_process_runner if env is None else None
     if _in_process_runner is not None and env is not None:
@@ -169,45 +146,34 @@ def run_example_command(
     env = os.environ.copy() if env is None else env
     cwd = MODELOPT_ROOT / "examples" / example_path
 
-    for attempt in range(hf_max_retries + 1):
-        if setup_free_port:
-            # Subprocess steps only: an in-process runner picks its own free port, since env is a
-            # copy it never sees.
-            env["MASTER_PORT"] = str(get_free_port())  # fresh port per attempt
-        if in_process is not None:
-            # Inside the loop so in-process steps get the same transient-error retries; these
-            # tests do hit the Hub (e.g. calib_dataset_name="cnn_dailymail").
-            try:
-                result = in_process(cmd_parts, example_path)
-            except Exception as e:
-                # Re-raise unless it looks transient, so a real failure keeps its traceback
-                # instead of being flattened into CalledProcessError.
-                # format_exc() walks __cause__/__context__: most markers are exception type
-                # names that only ever appear in a formatted traceback, never in str(e) -- a Hub
-                # ConnectionError typically surfaces wrapped in a DatasetGenerationError. The
-                # runner also attaches what the step printed, for launcher-level failures.
-                captured = getattr(e, "captured_output", "")
-                text = f"{traceback.format_exc()}\n{captured}"
-                if attempt == hf_max_retries or not any(
-                    marker in text for marker in _TRANSIENT_MARKERS
-                ):
-                    raise
-                returncode, output = 1, text
-            else:
-                assert result is not None, "an in-process runner must return output, not fall back"
-                return result
-        else:
-            returncode, output = _run_capturing(cmd_parts, cwd, env)
-        if returncode == 0:
-            return output
-        transient = any(marker in output for marker in _TRANSIENT_MARKERS)
-        if not transient or attempt == hf_max_retries:
-            raise subprocess.CalledProcessError(returncode, cmd_parts, output=output)
-        warnings.warn(
-            f"[{example_path}] transient error; retrying in "
-            f"{hf_retry_delay_s}s (attempt {attempt + 1}/{hf_max_retries})"
-        )
-        time.sleep(hf_retry_delay_s)
+    if in_process is not None:
+        try:
+            result = in_process(cmd_parts, example_path)
+        except Exception as e:
+            # Match the messages of the exception chain (a Hub ConnectionError typically surfaces
+            # wrapped in a DatasetGenerationError) and what the step printed, for launcher-level
+            # failures. Not a formatted traceback: its source lines can contain a marker.
+            marker = find_transient_marker(
+                f"{_exception_chain_text(e)}\n{getattr(e, 'captured_output', '')}"
+            )
+            if marker is None:
+                raise  # a real failure keeps its own traceback
+            raise TransientError(f"[{example_path}] transient error: {marker!r}") from e
+        assert result is not None, "an in-process runner must return output, not fall back"
+        return result
+
+    if setup_free_port:
+        # Subprocess steps only: an in-process runner picks its own free port, since env is a
+        # copy it never sees.
+        env["MASTER_PORT"] = str(get_free_port())
+    returncode, output = _run_capturing(cmd_parts, cwd, env)
+    if returncode == 0:
+        return output
+    error = subprocess.CalledProcessError(returncode, cmd_parts, output=output)
+    marker = find_transient_marker(output)
+    if marker is None:
+        raise error
+    raise TransientError(f"[{example_path}] transient error: {marker!r}") from error
 
 
 def run_hf_ptq_command(*, model: str, quant: str | None = None, vlm: bool = False, **kwargs):

tests/_test_utils/torch/misc.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import functools
 import random
 
 import pytest
@@ -53,6 +54,12 @@ def set_seed(seed_value=42):
     random.seed(seed_value)
 
 
+@functools.cache
+def is_blackwell():
+    """Whether the GPU is sm_120 or newer (e.g. RTX PRO 6000), as on the GPU CI runners."""
+    return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12
+
+
 def minimum_sm(sm):
     major, minor = sm // 10, sm % 10
     return pytest.mark.skipif(

tests/_test_utils/transient_errors.py

@@ -0,0 +1,52 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Failures that a rerun clears. ``tests/conftest.py`` reruns a test once when it fails with one."""
+
+# Substrings that identify a *transient* failure: a HuggingFace Hub / dataset access error, or a
+# rendezvous port that another process bound between its selection and its use
+TRANSIENT_MARKERS = (
+    "HfHubHTTPError",
+    "Too Many Requests",
+    "500 Server Error",
+    "502 Bad Gateway",
+    "503 Server Error",
+    "504 Server Error",
+    "Bad Gateway",
+    "Service Unavailable",
+    "Gateway Time-out",
+    "ConnectionError",
+    "ReadTimeout",
+    "ConnectTimeout",
+    "Max retries exceeded",
+    "NewConnectionError",
+    "Connection reset by peer",
+    "Connection aborted",
+    "Consistency check failed",  # partial / interrupted HF download
+    "couldn't connect to 'https://huggingface.co'",
+    "Couldn't reach",  # datasets: "Couldn't reach <repo> on the Hub"
+    "Temporary failure in name resolution",  # transient DNS
+    "Name or service not known",  # transient DNS
+    "EADDRINUSE",  # torch.distributed's TCPStore: "name: EADDRINUSE, message: address already in use"
+    "Address already in use",  # the same from Python's OSError, gloo and NCCL
+)
+
+
+class TransientError(RuntimeError):
+    """A command failed and its output shows one of ``TRANSIENT_MARKERS``."""
+
+
+def find_transient_marker(text: str) -> str | None:
+    """Return the first of ``TRANSIENT_MARKERS`` found in ``text``, if any."""
+    return next((marker for marker in TRANSIENT_MARKERS if marker in text), None)

tests/conftest.py

@@ -15,6 +15,7 @@
 
 import os
 import platform
+import re
 import sys
 from pathlib import Path
 
@@ -24,6 +25,8 @@
 from _test_utils.fs_utils import assert_unmodified_tree
 from _test_utils.mlflow import FakeMlflow, pin_tracking_env
 from _test_utils.torch.distributed.utils import init_process
+from _test_utils.torch.misc import is_blackwell
+from _test_utils.transient_errors import TRANSIENT_MARKERS, TransientError
 
 import modelopt.torch.opt as mto
 
@@ -70,9 +73,20 @@ def pytest_addoption(parser):
     "unit": 120,
 }
 
+# On Blackwell (sm_120), nemo:26.08's cuBLAS intermittently fails a GEMM once a gpu_megatron worker
+# process has run about a dozen model-building jobs. A failed job replaces the worker, so the rerun
+# gets a fresh one.
+# TODO: Remove after the CI nemo container is updated (re-check on nemo:26.10).
+_BLACKWELL_CUDA_CRASHES = [
+    "illegal memory access",
+    "CUBLAS_STATUS_INTERNAL_ERROR",
+    # The crash can also abort the worker (SIGABRT); not when another rank died otherwise, e.g. OOM
+    r"Worker process\(es\) died \(rank, exit code\): \[\(\d+, -6\)(, \(\d+, -6\))*\]",
+]
+
 
 def pytest_collection_modifyitems(config, items):
-    """Skip flag-gated tests and apply a default per-test timeout based on the test directory."""
+    """Skip flag-gated tests, apply a default per-test timeout, and rerun transient failures once."""
     skip_marks = [
         ("manual", "--run-manual"),
         ("nightly", "--run-nightly"),
@@ -103,6 +117,19 @@ def pytest_collection_modifyitems(config, items):
             )
         item.add_marker(pytest.mark.timeout(_DEFAULT_TIMEOUT[group]))
 
+    gpu_megatron_root = tests_root / "gpu_megatron"
+    for item in items:
+        only_rerun = [TransientError, *map(re.escape, TRANSIENT_MARKERS)]
+        # Only worker-pool tests: a crash in the pytest process would poison the rerun, which runs
+        # there too. is_blackwell() last, so other test groups do not initialize CUDA here.
+        if (
+            item.path.is_relative_to(gpu_megatron_root)
+            and any(name.startswith("dist_workers") for name in item.fixturenames)
+            and is_blackwell()
+        ):
+            only_rerun += _BLACKWELL_CUDA_CRASHES
+        item.add_marker(pytest.mark.flaky(reruns=1, reruns_delay=10, only_rerun=only_rerun))
+
 
 # General Fixtures #################################################################################
 @pytest.fixture

tests/gpu_megatron/torch/peft/plugins/test_megatron_peft.py

@@ -25,6 +25,7 @@
     load_distributed_checkpoint,
     save_distributed_checkpoint,
 )
+from _test_utils.torch.misc import is_blackwell
 
 import modelopt.torch.peft as mtpeft
 import modelopt.torch.quantization as mtq
@@ -40,7 +41,7 @@
 # process CUDA context and cascades into every test after it. Deterministic across CI runs; passes on
 # sm_89 with the same container and on nemo:26.06.
 skip_on_blackwell = pytest.mark.skipif(
-    torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 12,
+    is_blackwell(),
     reason="CUDA illegal memory access on Blackwell (nemo:26.08 TE/CUDA stack)",
 )