[Autotuner] Add crash recovery bash script for unrecoverable CUDA errors

yf225 · yf225 · commit d2ee5a6bac72 · 2026-04-02T00:55:57.000-07:00
diff --git a/helion/autotuner/base_search.py b/helion/autotuner/base_search.py
@@ -74,6 +74,7 @@
 from helion._dist_utils import is_symm_mem_tensor
 
 if TYPE_CHECKING:
+    from collections.abc import Iterator
     from collections.abc import Sequence
 
     from ..runtime.kernel import BoundKernel
@@ -435,6 +436,7 @@ def __init__(self, kernel: _AutotunableKernel, args: Sequence[object]) -> None:
         self._precompile_tmpdir: tempfile.TemporaryDirectory[str] | None = None
         self._precompile_args_path: str | None = None
         self._precompile_result_counter = count()
+        self._bad_config_strs: set[str] = set()
 
     def _prepare(self) -> None:
         """Some initialization deferred until autotuning actually runs.
@@ -534,6 +536,41 @@ def _try_load_checkpoint(self) -> bool:
         self.log(f"Resumed at generation {self._current_generation}")
         return True
 
+    def _load_bad_configs(self) -> None:
+        """Load bad configs from _bad_configs.txt (written by crash-recovery script)."""
+        checkpoint_dir_str = self.settings.autotune_checkpoint_dir
+        if checkpoint_dir_str is None:
+            return
+        bad_configs_path = Path(checkpoint_dir_str) / "_bad_configs.txt"
+        if bad_configs_path.exists():
+            self._bad_config_strs |= {
+                line.strip()
+                for line in bad_configs_path.read_text().splitlines()
+                if line.strip()
+            }
+        if self._bad_config_strs:
+            self.log(f"Loaded {len(self._bad_config_strs)} bad config(s) to skip")
+
+    @contextlib.contextmanager
+    def _pending_config(self, config: Config) -> Iterator[None]:
+        """Write a pending-config breadcrumb before benchmarking, clear it after.
+
+        If the body raises TritonUnrecoverableRuntimeError the pending file
+        is intentionally *not* cleared so the crash-recovery script can detect it.
+        """
+        checkpoint_dir_str = self.settings.autotune_checkpoint_dir
+        if checkpoint_dir_str is None:
+            yield
+            return
+        pending_path = Path(checkpoint_dir_str) / "_pending_config.txt"
+        pending_path.write_text(str(config))
+        try:
+            yield
+        except exc.TritonUnrecoverableRuntimeError:  # noqa: TRY203
+            raise
+        else:
+            pending_path.unlink(missing_ok=True)
+
     def _compute_baseline(
         self,
     ) -> tuple[object, Sequence[int], Sequence[object] | None]:
@@ -752,6 +789,12 @@ def benchmark_function(self, config: Config, fn: CompiledConfig) -> float:
         Returns:
             The performance of the configuration in ms.
         """
+        # Skip configs that previously crashed the subprocess
+        config_str = str(config)
+        if config_str in self._bad_config_strs:
+            self.log.warning(f"Skipping known-bad config: {config}")
+            return inf
+
         self._autotune_metrics.num_configs_tested += 1
         self.counters["benchmark"] += 1
         self.log.debug(lambda: f"Running benchmark for {config!r}")
@@ -1089,7 +1132,8 @@ def _benchmark(
                     )
                 )
                 # benchmark one-by-one to avoid noisy results
-                perf = self.benchmark_function(config, fn)
+                with self._pending_config(config):
+                    perf = self.benchmark_function(config, fn)
                 status = "ok" if math.isfinite(perf) else "error"
                 # Log completion after benchmarking
                 self.log.record_autotune_entry(
@@ -1194,6 +1238,7 @@ def autotune(self, *, skip_cache: bool = False) -> Config:
 
             if not self._try_load_checkpoint():
                 self._init_search()
+            self._load_bad_configs()
             try:
                 best = self._autotune()
                 self._cleanup_checkpoint()
@@ -1296,6 +1341,12 @@ def _cleanup_checkpoint(self) -> None:
             checkpoint_file.unlink()
             self.log(f"Checkpoint cleaned up: {checkpoint_file}")
 
+        # Clean up crash-recovery artifacts
+        for name in ("_pending_config.txt", "_bad_configs.txt"):
+            artifact = Path(checkpoint_dir_str) / name
+            if artifact.exists():
+                artifact.unlink()
+
     @staticmethod
     def _serialize_numpy_rng_state(
         state: tuple[str, Any, int, int, float],
diff --git a/scripts/autotune_with_crash_recovery.sh b/scripts/autotune_with_crash_recovery.sh
@@ -0,0 +1,111 @@
+#!/usr/bin/env bash
+# Autotuner crash recovery wrapper.
+#
+# Runs a command (typically a Python script that calls helion autotuning)
+# in a retry loop. When the process crashes due to an unrecoverable CUDA
+# error (illegal memory access, misaligned address, etc.), the autotuner
+# leaves a "_pending_config.txt" breadcrumb in the checkpoint directory.
+# This script detects that file, records the poison config in
+# "_bad_configs.txt", and re-runs the command. On re-run the autotuner
+# loads its checkpoint and skips the bad config.
+#
+# Progress detection:
+#   Each crash should block a different config (since blocked configs are
+#   skipped on re-run). If the same config crashes twice, the autotuner
+#   is stuck and we give up.
+#
+# Requirements:
+#   - HELION_AUTOTUNE_CHECKPOINT_DIR must be set
+#
+# Usage:
+#   HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/ckpt \
+#       scripts/autotune_with_crash_recovery.sh -- COMMAND [ARGS...]
+#
+# Examples:
+#   HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/autotune_ckpt \
+#       scripts/autotune_with_crash_recovery.sh -- python train.py
+
+set -uo pipefail
+
+# --- Argument parsing ---
+usage() {
+    cat >&2 <<'EOF'
+Usage: HELION_AUTOTUNE_CHECKPOINT_DIR=/path/to/dir \
+           autotune_with_crash_recovery.sh -- COMMAND [ARGS...]
+EOF
+    exit "${1:-1}"
+}
+
+while [[ $# -gt 0 ]]; do
+    case "$1" in
+        -h|--help)
+            usage 0
+            ;;
+        --)
+            shift
+            break
+            ;;
+        *)
+            echo "Error: unknown option '$1'" >&2
+            usage 1
+            ;;
+    esac
+done
+
+if [[ $# -eq 0 ]]; then
+    echo "Error: no command specified after --" >&2
+    usage 1
+fi
+
+if [[ -z "${HELION_AUTOTUNE_CHECKPOINT_DIR:-}" ]]; then
+    echo "Error: HELION_AUTOTUNE_CHECKPOINT_DIR must be set." >&2
+    exit 1
+fi
+
+# --- Setup ---
+checkpoint_dir="$HELION_AUTOTUNE_CHECKPOINT_DIR"
+mkdir -p "$checkpoint_dir"
+
+pending_file="$checkpoint_dir/_pending_config.txt"
+bad_configs_file="$checkpoint_dir/_bad_configs.txt"
+
+# --- Retry loop ---
+attempt=0
+last_config=""
+
+while true; do
+    attempt=$((attempt + 1))
+
+    # Run the user command (don't use set -e, capture exit code manually)
+    "$@"
+    exit_code=$?
+
+    if [[ $exit_code -eq 0 ]]; then
+        exit 0
+    fi
+
+    # Check if the autotuner left a pending config breadcrumb
+    if [[ -f "$pending_file" ]]; then
+        config=$(cat "$pending_file")
+        rm -f "$pending_file"
+        echo "$config" >> "$bad_configs_file"
+
+        echo "[crash-recovery] Process crashed (exit code $exit_code, attempt $attempt)." >&2
+        echo "[crash-recovery] Blocked config: $config" >&2
+
+        # If the same config crashed again, the bad config is not being
+        # skipped — the autotuner is stuck.
+        if [[ "$config" == "$last_config" ]]; then
+            echo "[crash-recovery] Same config crashed twice — the autotuner appears stuck." >&2
+            echo "[crash-recovery] All bad configs have been recorded. You can re-run this script and it will resume from the latest checkpoint, skipping all previously recorded bad configs." >&2
+            exit 1
+        fi
+        last_config="$config"
+
+        echo "[crash-recovery] Restarting from checkpoint..." >&2
+    else
+        # No pending file — this is not a recoverable CUDA crash.
+        # Propagate the original exit code.
+        exit "$exit_code"
+    fi
+done
diff --git a/test/data/autotune_crash_helper.py b/test/data/autotune_crash_helper.py
@@ -0,0 +1,67 @@
+"""Helper script for bash crash recovery tests.
+
+Run via:
+    HELION_AUTOTUNE_CHECKPOINT_DIR=DIR \
+        scripts/autotune_with_crash_recovery.sh -- python test/data/autotune_crash_helper.py
+
+On first run (when _CRASH_ON_FIRST_BENCHMARK is set and no counter file
+exists): patches do_bench to trigger a real CUDA illegal memory access,
+which exercises the real _pending_config context manager and
+TritonUnrecoverableRuntimeError code path.  On subsequent runs: autotuning
+resumes from checkpoint normally, skipping the bad config.
+
+Without _CRASH_ON_FIRST_BENCHMARK: runs autotuning normally (used to test
+that the bash script passes through a successful run).
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+checkpoint_dir = os.environ["HELION_AUTOTUNE_CHECKPOINT_DIR"]
+crash_on_first = os.environ.get("_CRASH_ON_FIRST_BENCHMARK", "")
+counter_file = Path(checkpoint_dir) / "_benchmark_counter"
+
+if crash_on_first and not counter_file.exists():
+    import triton
+    import triton.language as tl
+
+    import helion.autotuner.base_search as _bs
+
+    @triton.jit
+    def _ima_kernel(ptr):
+        """Triton kernel that triggers illegal memory access."""
+        bad_ptr = ptr + (1 << 40)
+        tl.store(bad_ptr, tl.full([], 42.0, dtype=tl.float32))
+
+    _original_do_bench = _bs.do_bench
+
+    def _ima_do_bench(*args, **kwargs):  # type: ignore[no-untyped-def]
+        counter_file.write_text("done")
+        # Restore original so this only fires once
+        _bs.do_bench = _original_do_bench
+        # Trigger real CUDA illegal memory access
+        x = torch.zeros(1, device="cuda")
+        _ima_kernel[(1,)](x)
+        torch.cuda.synchronize()
+        # Should not reach here — IMA raises an exception
+        return _original_do_bench(*args, **kwargs)
+
+    _bs.do_bench = _ima_do_bench
+
+# Import and run real autotuning
+from helion._testing import import_path  # noqa: E402
+
+datadir = Path(__file__).parent
+basic_kernels = import_path(datadir / "basic_kernels.py")
+
+args = (torch.randn([8, 32], device="cuda"), torch.randn([8, 32], device="cuda"))
+bound = basic_kernels.add.bind(args)
+bound.settings.autotune_checkpoint_dir = checkpoint_dir
+bound.settings.autotune_effort = "quick"
+config = bound.autotune(args, force=True)
+result = bound(*args)
+torch.testing.assert_close(result, args[0] + args[1])
diff --git a/test/test_autotuner_subprocess.py b/test/test_autotuner_subprocess.py
@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import os
+from pathlib import Path
+import subprocess
+import tempfile
+
+import pytest
+
+SCRIPT = str(
+    Path(__file__).parent.parent / "scripts" / "autotune_with_crash_recovery.sh"
+)
+HELPER = str(Path(__file__).parent / "data" / "autotune_crash_helper.py")
+
+
+class TestBashCrashRecoveryScript:
+    """Tests for scripts/autotune_with_crash_recovery.sh.
+
+    These invoke the bash script via subprocess.run(). The crash recovery
+    test uses test/data/autotune_crash_helper.py which monkey-patches
+    do_bench to trigger a real CUDA illegal memory access, exercising the
+    real _pending_config context manager and TritonUnrecoverableRuntimeError
+    code path.
+    """
+
+    def _run_script(
+        self,
+        tmp_path: Path,
+        cmd: list[str],
+        extra_env: dict[str, str] | None = None,
+    ) -> subprocess.CompletedProcess[str]:
+        """Helper to run the bash script with HELION_AUTOTUNE_CHECKPOINT_DIR set."""
+        env = {**os.environ, "HELION_AUTOTUNE_CHECKPOINT_DIR": str(tmp_path)}
+        if extra_env:
+            env.update(extra_env)
+        return subprocess.run(
+            [SCRIPT, "--", *cmd], capture_output=True, text=True, env=env
+        )
+
+    def test_normal_exit(self, tmp_path: Path) -> None:
+        """Successful command passes through exit 0."""
+        r = self._run_script(tmp_path, ["python", "-c", "pass"])
+        assert r.returncode == 0
+
+    def test_no_pending_propagates_error(self, tmp_path: Path) -> None:
+        """Non-CUDA crash (no pending file) propagates exit code."""
+        r = self._run_script(tmp_path, ["python", "-c", "import sys; sys.exit(42)"])
+        assert r.returncode == 42
+
+    @pytest.mark.timeout(120)
+    def test_real_autotune_through_bash(self) -> None:
+        """End-to-end: real autotuning succeeds through the bash script."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            r = self._run_script(
+                Path(tmpdir),
+                ["python", HELPER],
+                extra_env={
+                    "HELION_AUTOTUNE_MAX_GENERATIONS": "1",
+                    "HELION_AUTOTUNER": "PatternSearch",
+                },
+            )
+            assert r.returncode == 0, f"stderr: {r.stderr}"
+
+    @pytest.mark.timeout(120)
+    def test_real_crash_recovery_through_bash(self) -> None:
+        """End-to-end: first run crashes during real benchmarking via
+        monkey-patch, bash script detects pending file, records bad config,
+        re-runs. Second run resumes from checkpoint and succeeds."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            r = self._run_script(
+                Path(tmpdir),
+                ["python", HELPER],
+                extra_env={
+                    "_CRASH_ON_FIRST_BENCHMARK": "1",
+                    "HELION_AUTOTUNE_MAX_GENERATIONS": "1",
+                    "HELION_AUTOTUNER": "PatternSearch",
+                },
+            )
+            assert r.returncode == 0, f"stderr: {r.stderr}"
+            # Verify crash recovery happened (bad_configs.txt is cleaned
+            # up by _cleanup_checkpoint on success, so check stderr)
+            assert "[crash-recovery]" in r.stderr
+            assert "Blocked config:" in r.stderr
