How to Trap SIGTERM for GPU Checkpoint Saves in Training

Distributed deep learning jobs spanning dozens or hundreds of GPUs represent substantial compute investments that run for days or weeks. In modern containerized orchestrators such as Kubernetes, Slurm, or cloud auto-scaling systems, training jobs are inevitably interrupted by node drains, spot preemptions, kernel maintenance events, or rescheduling policies. When an orchestrator initiates an eviction, it transmits a SIGTERM (Signal 15) to container processes, granting a finite grace period before issuing an unconditional, unmaskable SIGKILL (Signal 9). If your training code lacks an explicit signal trapping architecture, the job terminates abruptly, destroying hours of intermediate compute and corrupting in-flight files. On OneSource Cloud dedicated bare-metal GPU clusters, while physical single-tenancy eliminates spot preemption risks, implementing deterministic SIGTERM signal trapping remains an engineering requirement for graceful job restarts, planned maintenance, and rapid recovery.

The Distributed Signal Trap Problem: Why Naive Traps Deadlock

In standard single-GPU Python scripts, trapping SIGTERM is trivial: attach a handler via Python's standard signal.signal(signal.SIGTERM, handler) and call torch.save(). In distributed multi-GPU training environments managed by PyTorch DistributedDataParallel (DDP), Fully Sharded Data Parallel (FSDP), or Megatron-LM, this naive pattern triggers catastrophic deadlocks:

1. Asynchronous Signal Delivery: The container orchestrator does not deliver SIGTERM to all worker ranks simultaneously across nodes. A network lag of even a few hundred milliseconds means Rank 0 may receive SIGTERM while Rank 7 on another host continues executing forward-pass CUDA kernels.
2. Signal Handler Reentrancy: Executing complex CUDA operations, memory copies, or collective PyTorch calls directly inside an asynchronous Python signal handler violates signal safety rules and can cause memory corruption or runtime deadlocks.
3. Collective Communication Stalls: If Rank 0 traps SIGTERM and immediately halts its training loop to write a checkpoint file, peer ranks that have not yet trapped the signal will initiate the next iteration's dist.all_reduce() collective. Because Rank 0 is busy writing tens of gigabytes to disk, the remaining ranks freeze indefinitely waiting for Rank 0 at the collective barrier until the orchestrator's grace period expires and sends SIGKILL, resulting in half-written, corrupt state files.

Production PyTorch Distributed Signal Handler Architecture

A production-ready signal interception pattern must decouple signal capture from checkpoint serialization. Signal capture is handled asynchronously via an atomic boolean flag, while checkpoint serialization is executed synchronously at a clean step boundary across all ranks.

import os
import signal
import sys
import torch
import torch.distributed as dist
class GracefulPreemptionHandler:
    # Production-grade SIGTERM handler for distributed PyTorch training.
    def __init__(self):
        self.preemption_requested = False
        # Register both SIGTERM (orchestrator eviction) and SIGINT (Ctrl+C)
        signal.signal(signal.SIGTERM, self._signal_handler)
        signal.signal(signal.SIGINT, self._signal_handler)
    def _signal_handler(self, signum, frame):
        # Asynchronously capture signal without blocking the execution stack
        print(f"[Rank {os.environ.get('RANK', 0)}] Intercepted signal {signum}. Flagging graceful exit...")
        self.preemption_requested = True
    def should_exit(self) -> bool:
        # Synchronize preemption status across all distributed ranks.
        if not dist.is_initialized():
            return self.preemption_requested
        # Tensor flag: 1 if any rank received SIGTERM, 0 otherwise
        local_flag = torch.tensor(1 if self.preemption_requested else 0, device="cuda")
        dist.all_reduce(local_flag, op=dist.ReduceOp.MAX)
        return local_flag.item() == 1
def save_emergency_checkpoint(model, optimizer, step, checkpoint_dir):
    rank = dist.get_rank() if dist.is_initialized() else 0
    temp_path = os.path.join(checkpoint_dir, f"checkpoint_step_{step}.pt.tmp")
    final_path = os.path.join(checkpoint_dir, f"checkpoint_step_{step}.pt")
    if rank == 0:
        print(f"Writing emergency checkpoint at step {step}...")
        state_dict = {
            "step": step,
            "model_state": model.state_dict(),
            "optimizer_state": optimizer.state_dict(),
        }
        # Save to temporary file first to prevent partial checkpoint corruption
        torch.save(state_dict, temp_path)
        os.replace(temp_path, final_path)
        print(f"Emergency checkpoint safely committed: {final_path}")
    # Synchronize all ranks with timeout before exiting
    if dist.is_initialized():
        dist.barrier()

In the main training loop, evaluate handler.should_exit() at the end of each step, after gradient clipping and optimizer updates:

handler = GracefulPreemptionHandler()
for step, batch in enumerate(dataloader):
    optimizer.zero_grad()
    outputs = model(batch)
    loss = criterion(outputs, batch["labels"])
    loss.backward()
    optimizer.step()
    # Synchronous preemption checkpoint boundary
    if handler.should_exit():
        save_emergency_checkpoint(model, optimizer, step, "/mnt/storage/checkpoints")
        if dist.is_initialized():
            dist.destroy_process_group()
        sys.exit(0)  # Exit with 0 to indicate graceful completion to orchestrator

Storage Throughput and Grace Period Calibration

The success of emergency checkpointing depends on whether storage write bandwidth allows the model state to be flushed before the orchestrator delivers SIGKILL. The required orchestrator grace period is calculated using the following sizing equation:

$$T_{\text{grace}} \ge \left( \frac{\text{Checkpoint Size (GB)}}{\text{Storage Write Bandwidth (GB/s)}} \right) + T_{\text{step\_drain}} + T_{\text{barrier}} + T_{\text{safety}}$$

Where Checkpoint Size includes model weights, FP32 master weights, and optimizer momentum/variance tensors (typically 140GB–280GB for a 70B parameter model under standard FSDP sharding). If your shared storage throughput is limited to 500 MB/s, writing 200GB requires 400 seconds. An orchestrator configured with a standard 120-second grace period will inevitably kill the process mid-write.

OneSource Cloud eliminates storage I/O starvation by provisioning dedicated high-throughput NVMe-oF scratch storage alongside multi-terabyte local Gen5 NVMe arrays on every bare-metal node. Even during emergency eviction simulations, checkpoints serialize at multi-gigabyte-per-second line rates, ensuring zero lost progress.

FAQ

Why should checkpoints be saved to a temporary file before renaming?

If an orchestrator delivers SIGKILL while torch.save() is actively streaming bytes to disk, the target file is left in a truncated, unparseable state. When the training job restarts and attempts to resume from the corrupted file, PyTorch throws an EOFError or unpickling exception, aborting the restart. Saving to .tmp and using an atomic filesystem rename (os.replace) ensures that downstream loaders only read fully committed, intact checkpoints.

How do I test SIGTERM signal trapping before deploying to production?

Run a multi-GPU training job inside your staging environment and issue a manual termination signal to the container root process: kill -15 <PID>. Inspect the output logs to confirm that all ranks cleanly break out of the iteration loop, write the checkpoint file, report successful barrier completion, and exit with status code 0 rather than exit code 143 (killed by SIGTERM) or 137 (killed by SIGKILL).

What happens if Rank 0 fails while saving an emergency checkpoint?

In distributed training, if Rank 0 encounters an I/O exception while writing the checkpoint, peer ranks waiting at dist.barrier() will hang until the NCCL watchdog timer (NCCL_ASYNC_ERROR_HANDLING) triggers a timeout exception. To handle this, wrap the barrier in a timeout block or configure asynchronous error handling so that peer ranks fail gracefully and allow the orchestrator to reschedule the pod cleanly.