[Autotuner] Add crash recovery script for unrecoverable CUDA errors

stack-info: PR: #1923, branch: yf225/stack/93

Author: yf225

docs/deployment_autotuning.md

@@ -196,10 +196,10 @@ automatically. On successful completion, the checkpoint file is cleaned up.
 
 ```bash
 # Enable checkpointing to a directory:
-HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/helion_checkpoints python run_kernel.py
+HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/$USER/helion_checkpoints python run_kernel.py
 
 # If interrupted, just re-run with the same directory to resume:
-HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/helion_checkpoints python run_kernel.py
+HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/$USER/helion_checkpoints python run_kernel.py
 ```
 
 Without `HELION_AUTOTUNE_CHECKPOINT_DIR`, no checkpoints are saved (opt-in).

helion/autotuner/base_search.py

@@ -367,6 +367,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.
@@ -471,6 +472,32 @@ 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 _get_pending_config_path(self) -> Path | None:
+        """Get path for pending-config sentinel, or None if checkpointing disabled."""
+        checkpoint_dir_str = self.settings.autotune_checkpoint_dir
+        if checkpoint_dir_str is None:
+            return None
+        return Path(checkpoint_dir_str) / f"{self._get_stable_hash()}.pending_config"
+
     def _compute_baseline(
         self,
     ) -> tuple[object, Sequence[int], Sequence[object] | None]:
@@ -693,6 +720,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}")
@@ -973,13 +1006,36 @@ def _benchmark(
             A list of BenchmarkResult entries containing the configuration, compiled
             callable, measured performance, status, and compilation time.
         """
+        # Filter out known-crashed configs before compilation
+        if self._crashed_config_strs:
+            original_len = len(configs)
+            configs = [c for c in configs if str(c) not in self._crashed_config_strs]
+            skipped = original_len - len(configs)
+            if skipped:
+                self.log.warning(
+                    f"Skipped {skipped} known-crashed config(s) before compilation"
+                )
+            if not configs:
+                return []
+
         fns: list[Callable[..., object]] = []
         valid_configs: list[Config] = []
         futures: list[PrecompileFuture] | None = None
+        pending_path = self._get_pending_config_path()
         for i, config in enumerate(configs):
+            # Write sentinel before compile so a hard crash (SIGKILL /
+            # CUDA IMA) leaves a trace the crash 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:
+            except Exception as e:
+                if match_unrecoverable_runtime_error(e):
+                    # Leave sentinel for crash recovery — CUDA context is
+                    # corrupted and the process cannot continue.
+                    raise
+                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
@@ -989,9 +1045,14 @@ 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
+        # NOTE: precompile runs in separate subprocesses with isolated CUDA
+        # contexts; crashes there are caught via is_working checks, not
+        # sentinels.
         if self.settings.autotune_precompile:
             futures = list(
                 starmap(
@@ -1053,7 +1114,14 @@ def _benchmark(
                     )
                 )
                 # benchmark one-by-one to avoid noisy results
+                # Write pending-config sentinel; cleared after benchmark.
+                # On crash the file stays so the crash 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(
@@ -1158,6 +1226,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()
@@ -1259,6 +1328,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],

helion/experimental/crash_recovery.py

@@ -0,0 +1,146 @@
+"""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`` sentinel 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/$USER/helion_ckpt \\
+        python -m helion.experimental.crash_recovery [--max-retries N] -- COMMAND [ARGS...]
+
+Examples
+--------
+::
+
+    HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/$USER/helion_autotune_ckpt \\
+        python -m helion.experimental.crash_recovery -- python train.py
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+import subprocess
+import sys
+
+
+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 [--max-retries N] -- COMMAND [ARGS...]"
+        ),
+    )
+    parser.add_argument(
+        "--max-retries",
+        type=int,
+        default=50,
+        help="Maximum number of crash recovery retries (default: 50)",
+    )
+    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
+    all_crashed: set[str] = set()
+
+    while True:
+        attempt += 1
+
+        result = subprocess.run(command)
+        exit_code = result.returncode
+
+        if exit_code == 0:
+            return 0
+
+        # Look for any *.pending_config sentinel left by the autotuner.
+        pending_files = sorted(checkpoint_dir.glob("*.pending_config"))
+
+        if pending_files:
+            stuck = False
+            for pending_path in pending_files:
+                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"Blocked config: {config}")
+
+                # If this config was already blocked in a previous attempt,
+                # the autotuner is not skipping it -- it's stuck.
+                if config in all_crashed:
+                    stuck = True
+                all_crashed.add(config)
+
+            _log(f"Process crashed (exit code {exit_code}, attempt {attempt}).")
+
+            if stuck:
+                _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
+
+            if attempt >= args.max_retries:
+                _log(f"Reached maximum retry limit ({args.max_retries}). Giving up.")
+                return 1
+
+            _log("Restarting from checkpoint...")
+        else:
+            # No pending file -- not a recoverable CUDA crash.
+            return exit_code
+
+
+if __name__ == "__main__":
+    sys.exit(main())