feat: auto-checkpoint on SIGTERM/SIGINT with DDP support

- Register SIGTERM/SIGINT handlers in training loop, set stop flag on signal
- Check stop_requested at each epoch/batch boundary, break and call on_error to save checkpoint
- LocalStrategy parent forwards signal to child processes via terminate(), waits up to 600s for graceful exit
- TrainContext gains threading.Event-based stop_requested/request_stop
- Tests verify SIGTERM/SIGINT trigger checkpoint save with exit code 0, works on both CPU and GPU

astrai/parallel/setup.py

@@ -1,5 +1,8 @@
+import logging
 import os
+import signal
 import socket
+import threading
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from functools import wraps
@@ -9,6 +12,10 @@ import torch
 import torch.distributed as dist
 import torch.multiprocessing as mp
 
+from astrai.parallel.signal_handler import install_early_signal_handlers
+
+logger = logging.getLogger(__name__)
+
 
 def find_free_port() -> str:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
@@ -115,6 +122,7 @@ def _run_single_rank(
     func: Callable,
     kwargs: dict,
 ):
+    install_early_signal_handlers()
     with setup_parallel(
         rank=rank,
         world_size=world_size,
@@ -155,6 +163,7 @@ class TorchrunStrategy(LaunchStrategy):
     """External orchestrator (torchrun, SLURM, K8s) — env vars pre-set."""
 
     def launch(self, func: Callable, **kwargs):
+        install_early_signal_handlers()
         rank = int(os.environ["RANK"])
         world_size = int(os.environ["WORLD_SIZE"])
         local_rank = int(os.environ.get("LOCAL_RANK", rank))
@@ -188,6 +197,7 @@ class LocalStrategy(LaunchStrategy):
             _run_single_rank(0, *args)
             return
 
+        install_early_signal_handlers()
         ctx = mp.start_processes(
             _run_single_rank,
             args=args,
@@ -195,14 +205,46 @@ class LocalStrategy(LaunchStrategy):
             start_method=self.start_method,
             join=False,
         )
+
+        parent_stop = threading.Event()
+        original_handlers = {}
+
+        def _parent_handler(signum, frame):
+            sig = signal.Signals(signum)
+            logger.warning(
+                "Parent (pid=%d) received %s, forwarding to children...",
+                os.getpid(),
+                sig.name,
+            )
+            parent_stop.set()
+            for p in ctx.processes:
+                if p.is_alive():
+                    p.terminate()
+
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            prev = signal.signal(sig, _parent_handler)
+            if prev not in (signal.SIG_DFL, signal.SIG_IGN, None, _parent_handler):
+                original_handlers[sig] = prev
+
         try:
-            while not ctx.join():
+            while not ctx.join() and not parent_stop.is_set():
                 pass
         except BaseException:
+            logger.warning(
+                "Parent received unexpected exception, terminating children..."
+            )
             for p in ctx.processes:
-                p.terminate()
-            ctx.join()
+                if p.is_alive():
+                    p.terminate()
             raise
+        finally:
+            for sig, handler in original_handlers.items():
+                signal.signal(sig, handler)
+
+            for p in ctx.processes:
+                p.join()
+
+            ctx.join()
 
 
 def _detect_launcher() -> str:

astrai/parallel/signal_handler.py

@@ -0,0 +1,44 @@
+import logging
+import os
+import signal
+import threading
+
+logger = logging.getLogger(__name__)
+
+_early_stop = threading.Event()
+_active_context = None
+
+
+def _early_handler(signum: int, frame):
+    sig = signal.Signals(signum)
+    logger.warning(
+        "Received %s (pid=%d), requesting graceful training stop...",
+        sig.name,
+        os.getpid(),
+    )
+    _early_stop.set()
+    if _active_context is not None:
+        _active_context.request_stop()
+
+
+def install_early_signal_handlers():
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        signal.signal(sig, _early_handler)
+
+
+def register_signal_handlers(context):
+    global _active_context
+    _active_context = context
+    for sig in (signal.SIGTERM, signal.SIGINT):
+        signal.signal(sig, _early_handler)
+    if _early_stop.is_set():
+        context.request_stop()
+        logger.warning("Signal was received during initialization, stopping...")
+
+
+def unregister_signal_handlers():
+    global _active_context
+    _active_context = None
+    _early_stop.clear()
+    signal.signal(signal.SIGTERM, signal.SIG_DFL)
+    signal.signal(signal.SIGINT, signal.SIG_DFL)

astrai/trainer/train_context.py

@@ -1,3 +1,4 @@
+import threading
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Dict, Optional, Self
@@ -41,6 +42,15 @@ class TrainContext:
     rank: int = field(default=0)
     kwargs: Dict[str, Any] = field(default_factory=dict)
 
+    _stop_event: threading.Event = field(default_factory=threading.Event)
+
+    @property
+    def stop_requested(self) -> bool:
+        return self._stop_event.is_set()
+
+    def request_stop(self) -> None:
+        self._stop_event.set()
+
     @property
     def optimizer_step(self) -> int:
         return self.consumed_samples // (

astrai/trainer/trainer.py

@@ -1,8 +1,14 @@
 import logging
 from typing import List, Optional
 
+import torch.distributed as dist
+
 from astrai.config import TrainConfig
 from astrai.parallel.setup import spawn_parallel_fn
+from astrai.parallel.signal_handler import (
+    register_signal_handlers,
+    unregister_signal_handlers,
+)
 from astrai.trainer.train_callback import (
     CallbackFactory,
     TrainCallback,
@@ -58,6 +64,7 @@ class Trainer:
             .with_param_path(param_path, resume=resume)
             .build()
         )
+        register_signal_handlers(context)
         executor = context.executor
         self._call_callbacks("on_train_begin", context)
 
@@ -65,10 +72,14 @@ class Trainer:
             context.model.train()
 
             for epoch in range(context.epoch, context.config.n_epoch):
+                if context.stop_requested:
+                    break
                 context.epoch = epoch
                 self._call_callbacks("on_epoch_begin", context)
 
                 for batch in context.dataloader:
+                    if context.stop_requested:
+                        break
                     with executor.accumulate(context.model):
                         self._call_callbacks("on_batch_begin", context)
                         loss = context.strategy(batch)
@@ -91,12 +102,21 @@ class Trainer:
 
                 self._call_callbacks("on_epoch_end", context)
 
+            if context.stop_requested:
+                logger.warning(
+                    "Training interrupted by signal, saving emergency checkpoint..."
+                )
+                self._call_callbacks("on_error", context)
+
         except Exception as e:
             logger.error("Training failed: %s", str(e), exc_info=True)
             self._call_callbacks("on_error", context)
             raise
         finally:
             self._call_callbacks("on_train_end", context)
+            if executor.use_distributed and dist.is_initialized():
+                dist.barrier()
+            unregister_signal_handlers()
 
     def train(self, param_path: Optional[str] = None, resume: bool = False):
         cfg = self.train_config