[Auto-Recovery] Add checkpoint save/load/resume

Add opt-in checkpoint support gated behind HELION_AUTOTUNE_CHECKPOINT_DIR.
When set, the autotuner saves in-progress state each generation and can
resume from a checkpoint on subsequent runs. The checkpoint file is
deleted on successful completion.

Includes pickle serialization support for BaseSearch and PopulationMember,
stable-hash-based checkpoint file naming, atomic writes, and kernel
recompilation on checkpoint load.

stack-info: PR: #1947, branch: yf225/stack/96

Author: yf225

docs/api/settings.md

@@ -209,6 +209,14 @@ def my_kernel(x: torch.Tensor) -> torch.Tensor:
    Each preset also sets a default initial population strategy (see :doc:`../deployment_autotuning` for details).
    Users can still override individual ``autotune_*`` settings; explicit values win over the preset. Controlled by ``HELION_AUTOTUNE_EFFORT``.
 
+.. autoattribute:: Settings.autotune_checkpoint_dir
+
+   Directory path for saving and resuming autotuning checkpoints. When set, the autotuner
+   saves in-progress state to ``{dir}/{stable_hash}.pt`` and auto-discovers matching
+   checkpoints on subsequent runs. The checkpoint file is deleted on successful completion.
+   When unset (default), no checkpoints are saved or loaded (opt-in).
+   Controlled by ``HELION_AUTOTUNE_CHECKPOINT_DIR``.
+
 .. autoattribute:: Settings.autotune_best_available_max_configs
 
    Maximum number of cached configs to use when seeding the initial population with the ``from_best_available`` strategy.
@@ -323,6 +331,7 @@ Built-in values for ``HELION_AUTOTUNER`` include ``"LFBOTreeSearch"`` (default),
 | ``HELION_AUTOTUNE_PROGRESS_BAR`` | ``autotune_progress_bar`` | Enable or disable the progress bar UI during autotuning. |
 | ``HELION_AUTOTUNE_IGNORE_ERRORS`` | ``autotune_ignore_errors`` | Continue autotuning even when recoverable runtime errors occur. |
 | ``HELION_AUTOTUNE_CONFIG_OVERRIDES`` | ``autotune_config_overrides`` | Supply JSON forcing particular autotuner config key/value pairs. |
+| ``HELION_AUTOTUNE_CHECKPOINT_DIR`` | ``autotune_checkpoint_dir`` | Directory path for saving/resuming autotuning checkpoints (opt-in). |
 | ``TRITON_STORE_BINARY_ONLY`` | Triton (autotuning) | Set to ``1`` during autotuning to skip Triton intermediate IRs, reducing cache size ~40%. Set to ``0`` to retain IRs for debugging. |
 | ``HELION_CACHE_DIR`` | ``LocalAutotuneCache`` | Override the on-disk directory used for cached autotuning artifacts. |
 | ``HELION_SKIP_CACHE`` | ``LocalAutotuneCache`` | When set to ``1``, skip both reading and writing the autotuning cache entirely. |

docs/deployment_autotuning.md

@@ -183,6 +183,29 @@ Related settings for `from_best_available` (see {doc}`api/settings`):
 | `autotune_best_available_max_configs` | `HELION_BEST_AVAILABLE_MAX_CONFIGS` | 20 | Maximum cached configs to seed |
 | `autotune_best_available_max_cache_scan` | `HELION_BEST_AVAILABLE_MAX_CACHE_SCAN` | 500 | Maximum cache files to scan |
 
+### Checkpointing Long-Running Autotuning
+
+For very long autotuning sessions, you can save and resume state using
+checkpoints. This is useful when tuning might be interrupted (e.g., preemptible
+instances) or when you want to continue tuning from a previous unfinished run.
+
+Set the `HELION_AUTOTUNE_CHECKPOINT_DIR` environment variable to a directory
+path. The autotuner will periodically save checkpoints there, keyed by the
+kernel's stable hash. If interrupted, re-run with the same directory to resume
+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
+
+# If interrupted, just re-run with the same directory to resume:
+HELION_AUTOTUNE_CHECKPOINT_DIR=/tmp/helion_checkpoints python run_kernel.py
+```
+
+Without `HELION_AUTOTUNE_CHECKPOINT_DIR`, no checkpoints are saved (opt-in).
+Multiple kernels can safely use the same directory — each kernel writes to a
+file named by its unique stable hash.
+
 ## Deploy a Single Config
 
 If one configuration wins for every production call, bake it into the decorator:

helion/autotuner/base_search.py

@@ -12,6 +12,8 @@
 import math
 from math import inf
 import os
+from pathlib import Path
+import pickle
 import pprint
 import random
 import re
@@ -407,6 +409,114 @@ def cleanup(self) -> None:
         self._precompile_args_path = None
         self._precompile_result_counter = count()
 
+    # Fields excluded from pickle checkpoints: unpicklable infrastructure,
+    # fields recomputed by _prepare(), and fields loaded separately.
+    _CHECKPOINT_EXCLUDE = frozenset(
+        {
+            # Unpicklable infrastructure
+            "kernel",
+            "args",
+            "log",
+            "settings",
+            "config_spec",
+            "_precompile_tmpdir",
+            "_precompile_args_path",
+            "_precompile_result_counter",
+            # Recomputed by _prepare() before checkpoint load
+            "_baseline_output",
+            "_baseline_post_args",
+            "_mutated_arg_indices",
+            "_effective_atol",
+            "_effective_rtol",
+            "_jobs",
+            "_autotune_metrics",
+            "_prepared",
+            "_skip_cache",
+            # Loaded separately via _load_crashed_configs()
+            "_crashed_config_strs",
+        }
+    )
+
+    def __getstate__(self) -> dict[str, Any]:
+        return {
+            k: v for k, v in self.__dict__.items() if k not in self._CHECKPOINT_EXCLUDE
+        }
+
+    _stable_hash: str | None = None
+
+    def _get_stable_hash(self) -> str:
+        """Get the full stable hash for this kernel's cache key (cached)."""
+        if self._stable_hash is None:
+            from .local_cache import LocalAutotuneCache
+
+            self._stable_hash = LocalAutotuneCache(self)._generate_key().stable_hash()
+        return self._stable_hash
+
+    def _try_load_checkpoint(self) -> bool:
+        """Attempt to load checkpoint from checkpoint dir. Returns True if successful."""
+        checkpoint_dir_str = self.settings.autotune_checkpoint_dir
+        if checkpoint_dir_str is None:
+            return False
+
+        checkpoint_dir = Path(checkpoint_dir_str)
+        stable_hash = self._get_stable_hash()
+        checkpoint_file = checkpoint_dir / f"{stable_hash}.pt"
+
+        if not checkpoint_file.exists():
+            return False  # No matching checkpoint; start fresh
+
+        # Matching file exists, attempt to load
+        self.log(f"Resuming from checkpoint: {checkpoint_file}")
+        try:
+            with open(checkpoint_file, "rb") as f:
+                loaded = pickle.load(f)
+        except Exception as e:
+            raise exc.CheckpointError(
+                f"Failed to load checkpoint file '{checkpoint_file}': {e}\n"
+                f"The file may be corrupted. Delete it to start fresh."
+            ) from e
+
+        # Validate stable hash matches (guards against renamed/copied files)
+        loaded_hash = getattr(loaded, "_stable_hash", None)
+        if loaded_hash is not None and loaded_hash != self._get_stable_hash():
+            raise exc.CheckpointError(
+                "Checkpoint is incompatible: kernel, hardware, or input shapes "
+                "may have changed."
+            )
+
+        # Copy loaded search state into self (self already has kernel, args,
+        # log, etc. from __init__ and _prepare())
+        self.__dict__.update(loaded.__dict__)
+
+        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]:
@@ -1091,9 +1201,13 @@ def autotune(self, *, skip_cache: bool = False) -> Config:
                 self._precompile_args_path = args_path
             exit_stack.callback(self.cleanup)
 
-            self._init_search()
+            checkpoint_enabled = self.settings.autotune_checkpoint_dir is not None
+            if not (checkpoint_enabled and self._try_load_checkpoint()):
+                self._init_search()
             try:
                 best = self._autotune()
+                if checkpoint_enabled:
+                    self._cleanup_checkpoint()
             finally:
                 self._finalize_autotune_metrics()
         end = time.perf_counter()
@@ -1119,6 +1233,7 @@ def _init_search(self) -> None:
         """
         Initialize the search state for a fresh autotuning run.
 
+        This method is called when starting autotuning without a checkpoint.
         Subclasses should override this to set up initial population and state.
         After this method, _current_generation should be set to the generation
         that _autotune() should start its loop from.
@@ -1135,6 +1250,68 @@ def _autotune(self) -> Config:
         """
         raise NotImplementedError
 
+    def save_checkpoint(self) -> Path | None:
+        """
+        Save current autotuner state to checkpoint file.
+
+        Only saves when autotune_checkpoint_dir is set (opt-in).
+        Overwrites the same file each generation (keyed by stable hash).
+        Uses pickle to serialize the entire autotuner object (minus unpicklable
+        fields excluded by __getstate__).
+
+        Returns:
+            Path to saved checkpoint file, or None if not saved
+        """
+        from ..runtime.kernel import BoundKernel
+
+        # External kernels don't support caching/checkpointing
+        if not isinstance(self.kernel, BoundKernel):
+            return None
+
+        if not self.kernel.is_cacheable():
+            return None
+
+        checkpoint_dir_str = self.settings.autotune_checkpoint_dir
+        if checkpoint_dir_str is None:
+            return None  # Opt-in: no dir set, no saving
+
+        stable_hash = self._get_stable_hash()
+        checkpoint_dir = Path(checkpoint_dir_str)
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        checkpoint_path = checkpoint_dir / f"{stable_hash}.pt"
+
+        # Atomic write using temp file + rename
+        tmp = checkpoint_dir / f".tmp.{stable_hash}.{os.getpid()}"
+        with open(tmp, "wb") as f:
+            pickle.dump(self, f)
+        os.replace(tmp, checkpoint_path)
+
+        self.log(f"Checkpoint saved: {checkpoint_path}")
+        return checkpoint_path
+
+    def _cleanup_checkpoint(self) -> None:
+        """Delete checkpoint file on successful autotune completion.
+
+        Checkpoints are ephemeral in-progress state. Once autotuning
+        completes successfully, the result is cached normally and the
+        checkpoint is no longer needed.
+        """
+        checkpoint_dir_str = self.settings.autotune_checkpoint_dir
+        if checkpoint_dir_str is None:
+            return
+
+        stable_hash = self._get_stable_hash()
+        checkpoint_file = Path(checkpoint_dir_str) / f"{stable_hash}.pt"
+        if checkpoint_file.exists():
+            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()
+
     def set_generation(self, generation: int) -> None:
         self._autotune_metrics.num_generations = generation
 
@@ -1189,6 +1366,15 @@ class PopulationMember:
     def perf(self) -> float:
         return self.perfs[-1]
 
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state["fn"] = None  # compiled functions are not picklable
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self.fn = _unset_fn
+
 
 def performance(member: PopulationMember) -> float:
     """
@@ -1587,6 +1773,8 @@ def set_generation(self, generation: int) -> None:
             return
         self._current_generation = generation
         super().set_generation(generation)
+        if generation > 0 and self.settings.autotune_checkpoint_dir is not None:
+            self.save_checkpoint()
 
     def statistics(self) -> str:
         """
@@ -1597,6 +1785,30 @@ def statistics(self) -> str:
         """
         return population_statistics(self.population)
 
+    def _try_load_checkpoint(self) -> bool:
+        if not super()._try_load_checkpoint():
+            return False
+        # Recompile kernel functions for population members after checkpoint load
+        recompile_failures: list[tuple[PopulationMember, str]] = []
+        for member in self.population:
+            if member.fn is _unset_fn and member.status == "ok":
+                try:
+                    member.fn = self.kernel.compile_config(
+                        member.config, allow_print=False
+                    )
+                except Exception as e:
+                    member.fn = _unset_fn
+                    member.status = "error"
+                    member.perfs.append(inf)  # Ensure member won't be selected as best
+                    recompile_failures.append((member, str(e)))
+
+        if recompile_failures:
+            self.log(
+                f"Warning: {len(recompile_failures)} config(s) failed to recompile "
+                f"and will be skipped. First failure: {recompile_failures[0][1]}"
+            )
+        return True
+
     def run_finishing_phase(
         self, best: PopulationMember, rounds: int
     ) -> PopulationMember:

helion/exc.py

@@ -58,6 +58,12 @@ class AutotuneError(BaseError):
     message = "{0}"
 
 
+class CheckpointError(AutotuneError):
+    """Exception raised when checkpoint loading/saving fails."""
+
+    message = "{0}"
+
+
 class BackendImplementationMissing(BaseError):
     message = "Backend '{backend}' is missing required implementation: {detail}"
 

helion/runtime/settings.py

@@ -492,6 +492,11 @@ class _Settings:
     autotune_baseline_rtol: float | None = None
     autotune_baseline_accuracy_check_fn: Callable[[object, object], None] | None = None
     autotune_benchmark_fn: Callable[..., list[float]] | None = None
+    autotune_checkpoint_dir: str | None = dataclasses.field(
+        default_factory=functools.partial(
+            os.environ.get, "HELION_AUTOTUNE_CHECKPOINT_DIR"
+        )
+    )
     autotune_best_available_max_configs: int = dataclasses.field(
         default_factory=functools.partial(
             _env_get_int, "HELION_BEST_AVAILABLE_MAX_CONFIGS", 20
@@ -640,6 +645,13 @@ class Settings(_Settings):
             "(fns: list[Callable[[], object]], *, repeat: int, desc: str | None = None) -> list[float]. "
             "If None (default), uses the built-in benchmark function."
         ),
+        "autotune_checkpoint_dir": (
+            "Directory path for saving and resuming autotuning checkpoints. "
+            "When set, the autotuner saves in-progress state to this directory using the "
+            "kernel's stable hash as the filename, and auto-discovers matching checkpoints "
+            "on subsequent runs. The checkpoint file is deleted on successful completion. "
+            "When unset (default), no checkpoints are saved or loaded."
+        ),
         "autotune_best_available_max_configs": (
             "Maximum number of cached configs to use for FROM_BEST_AVAILABLE initial population strategy. "
             "Set HELION_BEST_AVAILABLE_MAX_CONFIGS=N to override. Default is 20."