[CI] Enable retry logic for flaky CI tests (#14983)

.github/workflows/pr-test.yml

@@ -38,6 +38,7 @@ jobs:
       multimodal_gen: ${{ steps.filter.outputs.multimodal_gen || steps.scheduled.outputs.multimodal_gen }}
       max_parallel: ${{ steps.set-parallel.outputs.max_parallel }}
       b200_runner: ${{ steps.set-runner.outputs.b200_runner }}
+      enable_retry: ${{ steps.set-retry.outputs.enable_retry }}
     steps:
       - name: Checkout code
         uses: actions/checkout@v4
@@ -91,6 +92,12 @@ jobs:
             echo "b200_runner=4-gpu-b200" >> $GITHUB_OUTPUT
           fi
 
+      - name: Enable retry for CI
+        id: set-retry
+        run: |
+          echo "enable_retry=true" >> $GITHUB_OUTPUT
+          echo "Retry logic enabled for CI"
+
       - name: Show filter results in summary (table)
         run: |
           {
@@ -103,6 +110,7 @@ jobs:
             echo "| multimodal_gen | ${{ steps.filter.outputs.multimodal_gen || steps.scheduled.outputs.multimodal_gen }} |"
             echo "| max_parallel   | ${{ steps.set-parallel.outputs.max_parallel }} |"
             echo "| b200_runner    | ${{ steps.set-runner.outputs.b200_runner }} |"
+            echo "| enable_retry   | ${{ steps.set-retry.outputs.enable_retry }} |"
           } >> $GITHUB_STEP_SUMMARY
 
   # =============================================== PR Gate ====================================================
@@ -676,7 +684,11 @@ jobs:
           timeout-minutes: 30
           run: |
             cd test/srt
-            python3 run_suite.py --suite quantization_test
+            RETRY_FLAG=""
+            if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+              RETRY_FLAG="--enable-retry"
+            fi
+            python3 run_suite.py --suite quantization_test $RETRY_FLAG
 
   unit-test-backend-1-gpu:
     needs: [check-changes, call-gate, stage-a-test-1]
@@ -718,7 +730,11 @@ jobs:
         timeout-minutes: 30
         run: |
           cd test/srt
-          python3 run_suite.py --suite per-commit-1-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 15
+          RETRY_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+            RETRY_FLAG="--enable-retry"
+          fi
+          python3 run_suite.py --suite per-commit-1-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 15 $RETRY_FLAG
 
 
   stage-b-test-4-gpu-b200:
@@ -801,7 +817,11 @@ jobs:
         timeout-minutes: 30
         run: |
           cd test/srt
-          python3 run_suite.py --suite per-commit-2-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 2
+          RETRY_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+            RETRY_FLAG="--enable-retry"
+          fi
+          python3 run_suite.py --suite per-commit-2-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 $RETRY_FLAG
 
   unit-test-backend-4-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu, stage-b-test-4-gpu-b200]
@@ -842,7 +862,11 @@ jobs:
         timeout-minutes: 20
         run: |
           cd test/srt
-          python3 run_suite.py --suite per-commit-4-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 3
+          RETRY_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+            RETRY_FLAG="--enable-retry"
+          fi
+          python3 run_suite.py --suite per-commit-4-gpu --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 $RETRY_FLAG
 
   unit-test-backend-8-gpu-h200:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu, stage-b-test-4-gpu-b200]
@@ -889,7 +913,11 @@ jobs:
         timeout-minutes: 20
         run: |
           cd test/srt
-          python3 run_suite.py --suite per-commit-8-gpu-h200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 3
+          RETRY_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+            RETRY_FLAG="--enable-retry"
+          fi
+          python3 run_suite.py --suite per-commit-8-gpu-h200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 $RETRY_FLAG
 
   unit-test-backend-8-gpu-h20:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu, stage-b-test-4-gpu-b200]
@@ -931,7 +959,11 @@ jobs:
         timeout-minutes: 20
         run: |
           cd test/srt
-          python3 run_suite.py --suite per-commit-8-gpu-h20 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2
+          RETRY_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+            RETRY_FLAG="--enable-retry"
+          fi
+          python3 run_suite.py --suite per-commit-8-gpu-h20 --auto-partition-id ${{ matrix.part }} --auto-partition-size 2 $RETRY_FLAG
 
   performance-test-1-gpu-part-1:
     needs: [check-changes, call-gate, stage-a-test-1]
@@ -1220,10 +1252,10 @@ jobs:
           pip install -e .
 
       - name: Evaluate accuracy
-        timeout-minutes: 20
+        timeout-minutes: 25
         run: |
           cd test/srt
-          python3 test_eval_accuracy_large.py
+          python3 -m sglang.test.ci.run_with_retry test_eval_accuracy_large.py
 
   accuracy-test-2-gpu:
     needs: [check-changes, call-gate, accuracy-test-1-gpu]
@@ -1260,10 +1292,10 @@ jobs:
           pip install -e .
 
       - name: Evaluate accuracy (TP=2)
-        timeout-minutes: 20
+        timeout-minutes: 25
         run: |
           cd test/srt
-          python3 test_moe_eval_accuracy_large.py
+          python3 -m sglang.test.ci.run_with_retry test_moe_eval_accuracy_large.py
 
   unit-test-deepep-4-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu, stage-b-test-4-gpu-b200]
@@ -1300,7 +1332,11 @@ jobs:
         timeout-minutes: 20
         run: |
           cd test/srt
-          python3 run_suite.py --suite per-commit-4-gpu-deepep
+          RETRY_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+            RETRY_FLAG="--enable-retry"
+          fi
+          python3 run_suite.py --suite per-commit-4-gpu-deepep $RETRY_FLAG
 
   unit-test-deepep-8-gpu:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu, stage-b-test-4-gpu-b200]
@@ -1337,7 +1373,11 @@ jobs:
         timeout-minutes: 20
         run: |
           cd test/srt
-          python3 run_suite.py --suite per-commit-8-gpu-h200-deepep
+          RETRY_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+            RETRY_FLAG="--enable-retry"
+          fi
+          python3 run_suite.py --suite per-commit-8-gpu-h200-deepep $RETRY_FLAG
 
   unit-test-backend-4-gpu-b200:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu, stage-b-test-4-gpu-b200]
@@ -1379,7 +1419,11 @@ jobs:
         timeout-minutes: 30
         run: |
           cd test/srt
-          IS_BLACKWELL=1 python3 run_suite.py --suite per-commit-4-gpu-b200 --auto-partition-id ${{ matrix.part }}  --auto-partition-size 3 --timeout-per-file 1800
+          RETRY_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+            RETRY_FLAG="--enable-retry"
+          fi
+          IS_BLACKWELL=1 python3 run_suite.py --suite per-commit-4-gpu-b200 --auto-partition-id ${{ matrix.part }} --auto-partition-size 3 --timeout-per-file 1800 $RETRY_FLAG
 
   unit-test-backend-4-gpu-gb200:
     needs: [check-changes, call-gate, unit-test-backend-2-gpu, stage-b-test-4-gpu-b200, sgl-kernel-build-wheels-arm]
@@ -1418,7 +1462,11 @@ jobs:
         timeout-minutes: 45
         run: |
           cd test/srt
-          python3 run_suite.py --suite per-commit-4-gpu-gb200 --auto-partition-id 0 --auto-partition-size 1 --timeout-per-file 3600
+          RETRY_FLAG=""
+          if [[ "${{ needs.check-changes.outputs.enable_retry }}" == "true" ]]; then
+            RETRY_FLAG="--enable-retry"
+          fi
+          python3 run_suite.py --suite per-commit-4-gpu-gb200 --auto-partition-id 0 --auto-partition-size 1 --timeout-per-file 3600 $RETRY_FLAG
 
   pr-test-finish:
     needs:

python/sglang/test/ci/ci_utils.py

@@ -1,4 +1,6 @@
+import logging
 import os
+import re
 import subprocess
 import threading
 import time
@@ -8,6 +10,10 @@ from typing import Callable, List, Optional, Union
 from sglang.srt.utils.common import kill_process_tree
 from sglang.test.ci.ci_register import CIRegistry
 
+# Configure logger to output to stdout
+logging.basicConfig(level=logging.INFO, format="%(message)s")
+logger = logging.getLogger(__name__)
+
 
 @dataclass
 class TestFile:
@@ -15,6 +21,62 @@ class TestFile:
     estimated_time: float = 60
 
 
+# Patterns that indicate retriable accuracy/performance failures
+RETRIABLE_PATTERNS = [
+    r"AssertionError:.*not greater than",
+    r"AssertionError:.*not less than",
+    r"AssertionError:.*not equal to",
+    r"AssertionError:.*!=.*expected",
+    r"accuracy",
+    r"score",
+    r"latency",
+    r"throughput",
+    r"timeout",
+]
+
+# Patterns that indicate non-retriable failures (real code errors)
+NON_RETRIABLE_PATTERNS = [
+    r"SyntaxError",
+    r"ImportError",
+    r"ModuleNotFoundError",
+    r"NameError",
+    r"TypeError",
+    r"AttributeError",
+    r"RuntimeError",
+    r"CUDA out of memory",
+    r"OOM",
+    r"Segmentation fault",
+    r"core dumped",
+    r"ConnectionRefusedError",
+    r"FileNotFoundError",
+]
+
+
+def is_retriable_failure(output: str) -> tuple[bool, str]:
+    """
+    Determine if a test failure is retriable based on output patterns.
+
+    Returns:
+        tuple: (is_retriable, reason)
+    """
+    # Check for non-retriable patterns first
+    for pattern in NON_RETRIABLE_PATTERNS:
+        if re.search(pattern, output, re.IGNORECASE):
+            return False, f"non-retriable error: {pattern}"
+
+    # Check for retriable patterns
+    for pattern in RETRIABLE_PATTERNS:
+        if re.search(pattern, output, re.IGNORECASE):
+            return True, f"retriable pattern: {pattern}"
+
+    # If we have an AssertionError but didn't match non-retriable, assume retriable
+    if re.search(r"AssertionError", output):
+        return True, "AssertionError (assuming retriable)"
+
+    # Default: not retriable
+    return False, "unknown failure type"
+
+
 def run_with_timeout(
     func: Callable,
     args: tuple = (),
@@ -39,10 +101,21 @@ def run_with_timeout(
     return ret_value[0]
 
 
+def write_github_step_summary(content: str):
+    """Write content to GitHub Step Summary if available."""
+    summary_file = os.environ.get("GITHUB_STEP_SUMMARY")
+    if summary_file:
+        with open(summary_file, "a") as f:
+            f.write(content)
+
+
 def run_unittest_files(
     files: Union[List[TestFile], List[CIRegistry]],
     timeout_per_file: float,
     continue_on_error: bool = False,
+    enable_retry: bool = False,
+    max_attempts: int = 2,
+    retry_wait_seconds: int = 60,
 ):
     """
     Run a list of test files.
@@ -52,11 +125,16 @@ def run_unittest_files(
         timeout_per_file: Timeout in seconds for each test file
         continue_on_error: If True, continue running remaining tests even if one fails.
                           If False, stop at first failure (default behavior for PR tests).
+        enable_retry: If True, retry failed tests that appear to be accuracy/performance
+                     assertion failures (not code errors).
+        max_attempts: Maximum number of attempts per file including initial run (default: 2).
+        retry_wait_seconds: Seconds to wait between retries (default: 60).
     """
     tic = time.perf_counter()
     success = True
     passed_tests = []
     failed_tests = []
+    retried_tests = []  # Track which tests were retried
 
     for i, file in enumerate(files):
         if isinstance(file, CIRegistry):
@@ -66,77 +144,154 @@ def run_unittest_files(
             filename, estimated_time = file.name, file.estimated_time
 
         process = None
+        output_lines = []
 
-        def run_one_file(filename):
-            nonlocal process
+        def run_one_file(filename, capture_output=False):
+            nonlocal process, output_lines
 
-            filename = os.path.join(os.getcwd(), filename)
-            print(
-                f".\n.\nBegin ({i}/{len(files) - 1}):\npython3 {filename}\n.\n.\n",
-                flush=True,
+            full_path = os.path.join(os.getcwd(), filename)
+            logger.info(
+                f".\n.\nBegin ({i}/{len(files) - 1}):\npython3 {full_path}\n.\n.\n"
             )
-            tic = time.perf_counter()
+            file_tic = time.perf_counter()
 
-            process = subprocess.Popen(
-                ["python3", filename], stdout=None, stderr=None, env=os.environ
-            )
-            process.wait()
-            elapsed = time.perf_counter() - tic
+            if capture_output:
+                # Capture output for retry decision
+                process = subprocess.Popen(
+                    ["python3", full_path],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.STDOUT,
+                    env=os.environ,
+                    text=True,
+                    errors="ignore",  # Ignore non-UTF-8 bytes to prevent UnicodeDecodeError
+                )
+                output_lines = []
+                for line in process.stdout:
+                    logger.info(line.rstrip())
+                    output_lines.append(line)
+                process.wait()
+            else:
+                process = subprocess.Popen(
+                    ["python3", full_path], stdout=None, stderr=None, env=os.environ
+                )
+                process.wait()
 
-            print(
-                f".\n.\nEnd ({i}/{len(files) - 1}):\n{filename=}, {elapsed=:.0f}, {estimated_time=}\n.\n.\n",
-                flush=True,
+            elapsed = time.perf_counter() - file_tic
+
+            logger.info(
+                f".\n.\nEnd ({i}/{len(files) - 1}):\n{filename=}, {elapsed=:.0f}, {estimated_time=}\n.\n.\n"
             )
             return process.returncode
 
-        try:
-            ret_code = run_with_timeout(
-                run_one_file, args=(filename,), timeout=timeout_per_file
-            )
-            if ret_code != 0:
-                print(
-                    f"\n✗ FAILED: {filename} returned exit code {ret_code}\n",
-                    flush=True,
+        # Retry loop for each file
+        attempt = 1
+        file_passed = False
+        was_retried = False
+
+        while attempt <= (max_attempts if enable_retry else 1):
+            if attempt > 1:
+                logger.info(
+                    f"\n[CI Retry] Attempt {attempt}/{max_attempts} for {filename}\n"
                 )
-                success = False
-                failed_tests.append((filename, f"exit code {ret_code}"))
-                if not continue_on_error:
-                    # Stop at first failure for PR tests
+                was_retried = True
+
+            try:
+                ret_code = run_with_timeout(
+                    run_one_file,
+                    args=(filename,),
+                    kwargs={"capture_output": enable_retry},
+                    timeout=timeout_per_file,
+                )
+
+                if ret_code == 0:
+                    file_passed = True
+                    if was_retried:
+                        logger.info(
+                            f"\n✓ PASSED on retry (attempt {attempt}): {filename}\n"
+                        )
+                        retried_tests.append((filename, attempt, "passed"))
+                    passed_tests.append(filename)
                     break
-                # Otherwise continue to next test for nightly tests
-            else:
-                passed_tests.append(filename)
-        except TimeoutError:
-            kill_process_tree(process.pid)
-            time.sleep(5)
-            print(
-                f"\n✗ TIMEOUT: {filename} after {timeout_per_file} seconds\n",
-                flush=True,
-            )
-            success = False
-            failed_tests.append((filename, f"timeout after {timeout_per_file}s"))
-            if not continue_on_error:
-                # Stop at first timeout for PR tests
+                else:
+                    # Check if we should retry
+                    if enable_retry and attempt < max_attempts:
+                        output = "".join(output_lines)
+                        is_retriable, reason = is_retriable_failure(output)
+
+                        if is_retriable:
+                            logger.info(f"\n[CI Retry] {filename} failed with {reason}")
+                            logger.info(
+                                f"[CI Retry] Waiting {retry_wait_seconds}s before retry...\n"
+                            )
+                            time.sleep(retry_wait_seconds)
+                            attempt += 1
+                            continue
+                        else:
+                            logger.info(
+                                f"\n[CI Retry] {filename} failed with {reason} - not retrying\n"
+                            )
+
+                    # No retry or not retriable
+                    logger.info(
+                        f"\n✗ FAILED: {filename} returned exit code {ret_code}\n"
+                    )
+                    if was_retried:
+                        retried_tests.append((filename, attempt, "failed"))
+                    failed_tests.append((filename, f"exit code {ret_code}"))
+                    break
+
+            except TimeoutError:
+                kill_process_tree(process.pid)
+                time.sleep(5)
+                logger.info(
+                    f"\n✗ TIMEOUT: {filename} after {timeout_per_file} seconds\n"
+                )
+                if was_retried:
+                    retried_tests.append((filename, attempt, "timeout"))
+                failed_tests.append((filename, f"timeout after {timeout_per_file}s"))
                 break
-            # Otherwise continue to next test for nightly tests
+
+        if not file_passed:
+            success = False
+            if not continue_on_error:
+                break
+
+    elapsed_total = time.perf_counter() - tic
 
     if success:
-        print(f"Success. Time elapsed: {time.perf_counter() - tic:.2f}s", flush=True)
+        logger.info(f"Success. Time elapsed: {elapsed_total:.2f}s")
     else:
-        print(f"Fail. Time elapsed: {time.perf_counter() - tic:.2f}s", flush=True)
+        logger.info(f"Fail. Time elapsed: {elapsed_total:.2f}s")
 
     # Print summary
-    print(f"\n{'='*60}", flush=True)
-    print(f"Test Summary: {len(passed_tests)}/{len(files)} passed", flush=True)
-    print(f"{'='*60}", flush=True)
+    logger.info(f"\n{'='*60}")
+    logger.info(f"Test Summary: {len(passed_tests)}/{len(files)} passed")
+    if enable_retry and retried_tests:
+        logger.info(f"Retries: {len(retried_tests)} test(s) were retried")
+    logger.info(f"{'='*60}")
     if passed_tests:
-        print("✓ PASSED:", flush=True)
+        logger.info("✓ PASSED:")
         for test in passed_tests:
-            print(f"  {test}", flush=True)
+            logger.info(f"  {test}")
     if failed_tests:
-        print("\n✗ FAILED:", flush=True)
+        logger.info("\n✗ FAILED:")
         for test, reason in failed_tests:
-            print(f"  {test} ({reason})", flush=True)
-    print(f"{'='*60}\n", flush=True)
+            logger.info(f"  {test} ({reason})")
+    if retried_tests:
+        logger.info("\n↻ RETRIED:")
+        for test, attempts, result in retried_tests:
+            logger.info(f"  {test} ({attempts} attempts, {result})")
+    logger.info(f"{'='*60}\n")
+
+    # Write GitHub Step Summary only if retries occurred
+    if retried_tests:
+        passed_on_retry = [t for t, _, r in retried_tests if r == "passed"]
+        failed_after_retry = [t for t, _, r in retried_tests if r != "passed"]
+        summary = f"**↻ Retried {len(retried_tests)} test(s):**\n"
+        if passed_on_retry:
+            summary += f"- ✓ Passed on retry: {', '.join(passed_on_retry)}\n"
+        if failed_after_retry:
+            summary += f"- ✗ Still failed: {', '.join(failed_after_retry)}\n"
+        write_github_step_summary(summary)
 
     return 0 if success else -1

python/sglang/test/ci/run_with_retry.py

@@ -0,0 +1,53 @@
+"""
+Simple wrapper to run a test file with retry logic.
+
+Usage:
+    python3 -m sglang.test.ci.run_with_retry test_file.py [--max-attempts 2] [--retry-wait 60]
+"""
+
+import argparse
+import sys
+
+from sglang.test.ci.ci_utils import TestFile, run_unittest_files
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Run a test file with retry logic")
+    parser.add_argument("test_file", help="The test file to run")
+    parser.add_argument(
+        "--max-attempts",
+        type=int,
+        default=2,
+        help="Maximum number of attempts (default: 2)",
+    )
+    parser.add_argument(
+        "--retry-wait",
+        type=int,
+        default=60,
+        help="Seconds to wait between retries (default: 60)",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=int,
+        default=1200,
+        help="Timeout per attempt in seconds (default: 1200)",
+    )
+    args = parser.parse_args()
+
+    # Create a TestFile with a reasonable estimated time
+    test_file = TestFile(name=args.test_file, estimated_time=args.timeout)
+
+    exit_code = run_unittest_files(
+        files=[test_file],
+        timeout_per_file=args.timeout,
+        continue_on_error=False,
+        enable_retry=True,
+        max_attempts=args.max_attempts,
+        retry_wait_seconds=args.retry_wait,
+    )
+
+    sys.exit(exit_code)
+
+
+if __name__ == "__main__":
+    main()

test/run_suite.py

@@ -168,10 +168,18 @@ def run_a_suite(args):
 
     pretty_print_tests(args, ci_tests, skipped_tests)
 
+    # Add extra timeout when retry is enabled
+    timeout = args.timeout_per_file
+    if args.enable_retry:
+        timeout += args.retry_timeout_increase
+
     return run_unittest_files(
         ci_tests,
-        timeout_per_file=args.timeout_per_file,
+        timeout_per_file=timeout,
         continue_on_error=args.continue_on_error,
+        enable_retry=args.enable_retry,
+        max_attempts=args.max_attempts,
+        retry_wait_seconds=args.retry_wait_seconds,
     )
 
 
@@ -214,6 +222,30 @@ def main():
         type=int,
         help="Use auto load balancing. The number of parts.",
     )
+    parser.add_argument(
+        "--enable-retry",
+        action="store_true",
+        default=False,
+        help="Enable smart retry for accuracy/performance assertion failures (not code errors)",
+    )
+    parser.add_argument(
+        "--max-attempts",
+        type=int,
+        default=2,
+        help="Maximum number of attempts per file including initial run (default: 2)",
+    )
+    parser.add_argument(
+        "--retry-wait-seconds",
+        type=int,
+        default=60,
+        help="Seconds to wait between retries (default: 60)",
+    )
+    parser.add_argument(
+        "--retry-timeout-increase",
+        type=int,
+        default=600,
+        help="Additional timeout in seconds when retry is enabled (default: 600)",
+    )
     args = parser.parse_args()
 
     # Validate auto-partition arguments

test/srt/run_suite.py

@@ -491,6 +491,30 @@ def main():
         default=False,
         help="Continue running remaining tests even if one fails (useful for nightly tests)",
     )
+    arg_parser.add_argument(
+        "--enable-retry",
+        action="store_true",
+        default=False,
+        help="Enable smart retry for accuracy/performance assertion failures (not code errors)",
+    )
+    arg_parser.add_argument(
+        "--max-attempts",
+        type=int,
+        default=2,
+        help="Maximum number of attempts per file including initial run (default: 2)",
+    )
+    arg_parser.add_argument(
+        "--retry-wait-seconds",
+        type=int,
+        default=60,
+        help="Seconds to wait between retries (default: 60)",
+    )
+    arg_parser.add_argument(
+        "--retry-timeout-increase",
+        type=int,
+        default=600,
+        help="Additional timeout in seconds when retry is enabled (default: 600)",
+    )
     args = arg_parser.parse_args()
     print(f"{args=}")
 
@@ -506,7 +530,19 @@ def main():
 
     print("The running tests are ", [f.name for f in files])
 
-    exit_code = run_unittest_files(files, args.timeout_per_file, args.continue_on_error)
+    # Add extra timeout when retry is enabled
+    timeout = args.timeout_per_file
+    if args.enable_retry:
+        timeout += args.retry_timeout_increase
+
+    exit_code = run_unittest_files(
+        files,
+        timeout,
+        args.continue_on_error,
+        args.enable_retry,
+        args.max_attempts,
+        args.retry_wait_seconds,
+    )
     exit(exit_code)