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

tests/trainer/test_signal_handler.py

@@ -0,0 +1,177 @@
+import json
+import multiprocessing as mp
+import os
+import signal
+import time
+
+import pytest
+import torch
+import torch.optim as optim
+from torch.utils.data import Dataset
+
+from astrai.config import TrainConfig
+from astrai.config.model_config import AutoRegressiveLMConfig
+from astrai.model.transformer import AutoRegressiveLM
+from astrai.parallel.signal_handler import register_signal_handlers
+from astrai.trainer import Trainer
+from astrai.trainer.schedule import SchedulerFactory
+from astrai.trainer.train_context import TrainContext
+
+
+class _PicklableDataset(Dataset):
+    def __init__(self, length=200, max_length=64, vocab_size=1000):
+        self.length = length
+        self.max_length = max_length
+        self.vocab_size = vocab_size
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, idx):
+        return {
+            "input_ids": torch.randint(0, self.vocab_size, (self.max_length,)),
+            "target_ids": torch.randint(0, self.vocab_size, (self.max_length,)),
+        }
+
+
+def _build_model():
+    config = AutoRegressiveLMConfig(
+        vocab_size=1000,
+        hidden_size=8,
+        num_attention_heads=2,
+        num_key_value_heads=1,
+        intermediate_size=16,
+        max_position_embeddings=64,
+        num_hidden_layers=2,
+        rms_norm_eps=1e-5,
+    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    return AutoRegressiveLM(config).to(device=device)
+
+
+class _ReadyCallback:
+    def __init__(self, ready_file):
+        self._ready_file = ready_file
+
+    def on_train_begin(self, context):
+        with open(self._ready_file, "w") as f:
+            f.write("ready")
+            f.flush()
+            os.fsync(f.fileno())
+
+
+def _inner_run(batch_per_device, ckpt_interval, ckpt_dir, log_dir, ready_file):
+    dataset = _PicklableDataset()
+
+    def model_fn():
+        return _build_model()
+
+    def optimizer_fn(m):
+        return optim.AdamW(m.parameters(), lr=0.001)
+
+    def scheduler_fn(optim):
+        return SchedulerFactory.create(
+            "cosine", optim, warmup_steps=10, lr_decay_steps=10, min_rate=0.05
+        )
+
+    train_config = TrainConfig(
+        strategy="seq",
+        model_fn=model_fn,
+        dataset=dataset,
+        optimizer_fn=optimizer_fn,
+        scheduler_fn=scheduler_fn,
+        ckpt_dir=ckpt_dir,
+        log_dir=log_dir,
+        n_epoch=1,
+        batch_per_device=batch_per_device,
+        ckpt_interval=ckpt_interval,
+        grad_accum_steps=1,
+        random_seed=42,
+        device_type="cuda" if torch.cuda.is_available() else "cpu",
+    )
+
+    trainer = Trainer(train_config)
+    trainer.callbacks.insert(0, _ReadyCallback(ready_file))
+    trainer.train()
+
+
+def _spawn_train_and_signal(ckpt_dir, sig, timeout=120):
+    log_dir = os.path.join(ckpt_dir, "logs")
+    ready_file = os.path.join(ckpt_dir, "ready.txt")
+
+    ctx = mp.get_context("spawn")
+    p = ctx.Process(
+        target=_inner_run,
+        args=(2, 1000, ckpt_dir, log_dir, ready_file),
+    )
+    p.start()
+
+    deadline = time.time() + 30
+    while time.time() < deadline:
+        if os.path.exists(ready_file):
+            with open(ready_file) as f:
+                if f.read().strip() == "ready":
+                    break
+        if not p.is_alive():
+            break
+        time.sleep(0.5)
+
+    assert p.is_alive(), "Training process died before becoming ready"
+
+    os.kill(p.pid, sig)
+    p.join(timeout=timeout)
+
+    if p.is_alive():
+        p.kill()
+        p.join(timeout=5)
+
+    return p.exitcode
+
+
+def test_context_stop_flag():
+    ctx = TrainContext()
+    assert not ctx.stop_requested
+    ctx.request_stop()
+    assert ctx.stop_requested
+
+
+def test_register_signal_handlers():
+    ctx = TrainContext()
+    register_signal_handlers(ctx)
+    assert not ctx.stop_requested
+    os.kill(os.getpid(), signal.SIGTERM)
+    assert ctx.stop_requested
+
+
+def test_sigterm_triggers_checkpoint_save(base_test_env):
+    exitcode = _spawn_train_and_signal(base_test_env["test_dir"], signal.SIGTERM)
+    assert exitcode == 0, f"Training process exited with code {exitcode} (expected 0)"
+
+    ckpt_dir = base_test_env["test_dir"]
+    meta_files = []
+    for root, dirs, files in os.walk(ckpt_dir):
+        for f in files:
+            if f == "meta.json":
+                meta_files.append(os.path.join(root, f))
+
+    assert len(meta_files) > 0, f"No checkpoint meta.json found in {ckpt_dir}"
+
+    with open(meta_files[-1]) as f:
+        meta = json.load(f)
+    assert "consumed_samples" in meta
+    assert meta["consumed_samples"] >= 0
+
+
+@pytest.mark.slow
+def test_sigint_triggers_checkpoint_save(base_test_env):
+    exitcode = _spawn_train_and_signal(base_test_env["test_dir"], signal.SIGINT)
+    assert exitcode == 0, f"Training process exited with code {exitcode} (expected 0)"
+
+    ckpt_dir = base_test_env["test_dir"]
+    meta_files = []
+    for root, dirs, files in os.walk(ckpt_dir):
+        for f in files:
+            if f == "meta.json":
+                meta_files.append(os.path.join(root, f))
+
+    assert len(meta_files) > 0, f"No checkpoint meta.json found in {ckpt_dir}"