[Autotuner] Add crash recovery script for unrecoverable CUDA errors

yf225 · yf225 · commit 12b3b3b2aca5 · 2026-04-02T02:18:15.000-07:00
stack-info: PR: #1923, branch: yf225/stack/93
diff --git a/helion/autotuner/base_search.py b/helion/autotuner/base_search.py
@@ -435,6 +435,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._crashed_config_strs: set[str] = set()
 
     def _prepare(self) -> None:
         """Some initialization deferred until autotuning actually runs.
@@ -534,6 +535,26 @@ def _try_load_checkpoint(self) -> bool:
         self.log(f"Resumed at generation {self._current_generation}")
         return True
 
+    def _load_crashed_configs(self) -> None:
+        """Load crashed configs from {hash}.crashed_configs (written by crash-recovery script)."""
+        checkpoint_dir_str = self.settings.autotune_checkpoint_dir
+        if checkpoint_dir_str is None:
+            return
+        crashed_configs_path = (
+            Path(checkpoint_dir_str)
+            / f"{self._get_stable_hash()}.crashed_configs"
+        )
+        if crashed_configs_path.exists():
+            self._crashed_config_strs |= {
+                line.strip()
+                for line in crashed_configs_path.read_text().splitlines()
+                if line.strip()
+            }
+        if self._crashed_config_strs:
+            self.log(
+                f"Loaded {len(self._crashed_config_strs)} crashed config(s) to skip"
+            )
+
     def _compute_baseline(
         self,
     ) -> tuple[object, Sequence[int], Sequence[object] | None]:
@@ -752,6 +773,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._crashed_config_strs:
+            self.log.warning(f"Skipping known-crashed 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}")
@@ -1016,10 +1043,23 @@ def _benchmark(
         fns: list[Callable[..., object]] = []
         valid_configs: list[Config] = []
         futures: list[PrecompileFuture] | None = None
+        # Compute pending config path once for breadcrumb writes.
+        checkpoint_dir_str = self.settings.autotune_checkpoint_dir
+        pending_path = (
+            Path(checkpoint_dir_str) / f"{self._get_stable_hash()}.pending_config"
+            if checkpoint_dir_str is not None
+            else None
+        )
         for i, config in enumerate(configs):
+            # Write breadcrumb before compile so a hard crash (SIGKILL /
+            # CUDA IMA) leaves a trace the bash recovery script can find.
+            if pending_path is not None:
+                pending_path.write_text(str(config))
             try:
                 fn = self.kernel.compile_config(config, allow_print=False)
             except Exception:
+                if pending_path is not None:
+                    pending_path.unlink(missing_ok=True)
                 # If all configs failed, raise error
                 if not valid_configs and i == len(configs) - 1:
                     raise
@@ -1029,6 +1069,8 @@ def _benchmark(
                     exc_info=True,
                 )
                 continue
+            if pending_path is not None:
+                pending_path.unlink(missing_ok=True)
             fns.append(fn)
             valid_configs.append(config)
         configs = valid_configs
@@ -1089,7 +1131,14 @@ def _benchmark(
                     )
                 )
                 # benchmark one-by-one to avoid noisy results
+                # Write pending-config breadcrumb; cleared after benchmark.
+                # On crash the file stays so the bash recovery script can
+                # detect which config caused the failure.
+                if pending_path is not None:
+                    pending_path.write_text(str(config))
                 perf = self.benchmark_function(config, fn)
+                if pending_path is not None:
+                    pending_path.unlink(missing_ok=True)
                 status = "ok" if math.isfinite(perf) else "error"
                 # Log completion after benchmarking
                 self.log.record_autotune_entry(
@@ -1194,6 +1243,7 @@ def autotune(self, *, skip_cache: bool = False) -> Config:
 
             if not self._try_load_checkpoint():
                 self._init_search()
+            self._load_crashed_configs()
             try:
                 best = self._autotune()
                 self._cleanup_checkpoint()
@@ -1296,6 +1346,12 @@ def _cleanup_checkpoint(self) -> None:
             checkpoint_file.unlink()
             self.log(f"Checkpoint cleaned up: {checkpoint_file}")
 
+        # Clean up crash-recovery artifacts
+        for suffix in (".pending_config", ".crashed_configs"):
+            artifact = Path(checkpoint_dir_str) / f"{stable_hash}{suffix}"
+            if artifact.exists():
+                artifact.unlink()
+
     @staticmethod
     def _serialize_numpy_rng_state(
         state: tuple[str, Any, int, int, float],
diff --git a/helion/experimental/crash_recovery.py b/helion/experimental/crash_recovery.py
@@ -0,0 +1,136 @@
+"""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
+``{hash}.pending_config`` breadcrumb in the checkpoint directory.  This script
+detects that file, records the poison config in ``{hash}.crashed_configs``, and
+re-runs the command.  On re-run the autotuner loads its checkpoint and skips
+the crashed 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 in the environment.
+
+Usage
+-----
+::
+
+    HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/ckpt \\
+        python -m helion.experimental.crash_recovery -- COMMAND [ARGS...]
+
+Examples
+--------
+::
+
+    HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/autotune_ckpt \\
+        python -m helion.experimental.crash_recovery -- python train.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import glob
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+
+def _log(msg: str) -> None:
+    print(f"[crash-recovery] {msg}", file=sys.stderr)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(
+        description="Autotuner crash recovery wrapper.",
+        usage=(
+            "HELION_AUTOTUNE_CHECKPOINT_DIR=/path/to/dir\n"
+            "           %(prog)s -- COMMAND [ARGS...]"
+        ),
+    )
+    parser.add_argument(
+        "command",
+        nargs=argparse.REMAINDER,
+        help="Command to run (after '--' separator)",
+    )
+    args = parser.parse_args(argv)
+
+    # argparse.REMAINDER absorbs '--' as first element when present.
+    command: list[str] = args.command
+    if command and command[0] == "--":
+        command = command[1:]
+    if not command:
+        parser.error("no command specified after --")
+
+    checkpoint_dir_str = os.environ.get("HELION_AUTOTUNE_CHECKPOINT_DIR", "")
+    if not checkpoint_dir_str:
+        print(
+            "Error: HELION_AUTOTUNE_CHECKPOINT_DIR must be set.",
+            file=sys.stderr,
+        )
+        return 1
+
+    checkpoint_dir = Path(checkpoint_dir_str)
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+    attempt = 0
+    last_config = ""
+
+    while True:
+        attempt += 1
+
+        result = subprocess.run(command)
+        exit_code = result.returncode
+
+        if exit_code == 0:
+            return 0
+
+        # Look for any *.pending_config breadcrumb left by the autotuner.
+        pending_files = sorted(glob.glob(str(checkpoint_dir / "*.pending_config")))
+
+        if pending_files:
+            # Process the first pending file found (typically only one exists
+            # because the process crashed while working on a single config).
+            pending_path = Path(pending_files[0])
+            hash_prefix = pending_path.stem  # {hash} without .pending_config
+            crashed_configs_path = checkpoint_dir / f"{hash_prefix}.crashed_configs"
+
+            config = pending_path.read_text().strip()
+            pending_path.unlink()
+
+            with open(crashed_configs_path, "a") as f:
+                f.write(config + "\n")
+
+            _log(f"Process crashed (exit code {exit_code}, attempt {attempt}).")
+            _log(f"Blocked config: {config}")
+
+            # If the same config crashed again, the crashed config is not
+            # being skipped -- the autotuner is stuck.
+            if config == last_config:
+                _log(
+                    "Same config crashed twice \u2014 the autotuner appears stuck."
+                )
+                _log(
+                    "All crashed configs have been recorded. You can re-run "
+                    "this script and it will resume from the latest "
+                    "checkpoint, skipping all previously recorded crashed "
+                    "configs."
+                )
+                return 1
+            last_config = config
+
+            _log("Restarting from checkpoint...")
+        else:
+            # No pending file -- not a recoverable CUDA crash.
+            return exit_code
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/test/data/autotune_crash_helper.py b/test/data/autotune_crash_helper.py
@@ -0,0 +1,90 @@
+"""Helper script for crash recovery tests.
+
+Run via:
+    HELION_AUTOTUNE_CHECKPOINT_DIR=DIR \
+        python -m helion.experimental.crash_recovery -- python test/data/autotune_crash_helper.py
+
+On first run (when _CRASH_ON_FIRST_BENCHMARK or _CRASH_ON_FIRST_COMPILE is
+set and no counter file exists): patches do_bench / compile_config to trigger
+a hard crash, which exercises the pending_config breadcrumb and the crash
+recovery script.  On subsequent runs: autotuning resumes from checkpoint
+normally, skipping the crashed config.
+
+Without the crash env vars: 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_benchmark = os.environ.get("_CRASH_ON_FIRST_BENCHMARK", "")
+crash_on_first_compile = os.environ.get("_CRASH_ON_FIRST_COMPILE", "")
+counter_file = Path(checkpoint_dir) / "_crash_counter"
+
+if crash_on_first_benchmark 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
+
+if crash_on_first_compile and not counter_file.exists():
+    import helion.autotuner.base_search as _bs
+
+    # Wrap _benchmark so the real breadcrumb-writing code runs, but
+    # compile_config triggers a hard crash (os._exit) on first call.
+    _original_benchmark = _bs.BaseSearch._benchmark
+
+    def _crashing_benchmark(self, configs, **kwargs):  # type: ignore[no-untyped-def]
+        _orig_compile = self.kernel.compile_config
+
+        def _crash_compile(*args, **kw):  # type: ignore[no-untyped-def]
+            counter_file.write_text("done")
+            # Simulate a hard crash (SIGKILL / CUDA IMA) during
+            # compile_config. os._exit bypasses all Python exception
+            # handling including try/except.
+            os._exit(1)
+
+        self.kernel.compile_config = _crash_compile
+        return _original_benchmark(self, configs, **kwargs)
+
+    _bs.BaseSearch._benchmark = _crashing_benchmark  # type: ignore[assignment]
+
+# 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
