From ceadc34ea976f0e69e7c8e02ff8a74c860299158 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 25 Jul 2026 19:55:54 +0800 Subject: [PATCH] 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 | 48 +++++++- astrai/parallel/signal_handler.py | 44 +++++++ astrai/trainer/train_context.py | 10 ++ astrai/trainer/trainer.py | 20 +++ tests/trainer/test_signal_handler.py | 177 +++++++++++++++++++++++++++ 5 files changed, 296 insertions(+), 3 deletions(-) create mode 100644 astrai/parallel/signal_handler.py create mode 100644 tests/trainer/test_signal_handler.py diff --git a/astrai/parallel/setup.py b/astrai/parallel/setup.py index c453dc0..2ce838c 100644 --- a/astrai/parallel/setup.py +++ b/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: diff --git a/astrai/parallel/signal_handler.py b/astrai/parallel/signal_handler.py new file mode 100644 index 0000000..5e6a480 --- /dev/null +++ b/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) diff --git a/astrai/trainer/train_context.py b/astrai/trainer/train_context.py index c089846..1cc2895 100644 --- a/astrai/trainer/train_context.py +++ b/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 // ( diff --git a/astrai/trainer/trainer.py b/astrai/trainer/trainer.py index 3f908cd..227da3e 100644 --- a/astrai/trainer/trainer.py +++ b/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 diff --git a/tests/trainer/test_signal_handler.py b/tests/trainer/test_signal_handler.py new file mode 100644 index 0000000..d219438 --- /dev/null +++ b/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}"