The Art of Checkpointing: Surviving the Eviction

Don't Lose Your Progress

To make spot instances usable, your training code must be fault-tolerant. The key technique is checkpointing — periodically saving a complete snapshot of your model's current state to durable, persistent storage that survives instance termination.

A checkpoint file contains everything needed to resume training from that exact point:
- Model weights: The learned parameters of the neural network.
- Optimizer state: Gradient accumulators, momentum terms, and learning rate scheduler state. Without this, resuming training can cause the optimizer to behave erratically for many steps.
- Current epoch and step number: So the training loop knows where to restart.
- Random number generator state: For reproducibility of data shuffling.

You save checkpoints to durable storage like Amazon S3, Google Cloud Storage, or a network-attached filesystem — somewhere that persists independently of the instance's local disk.

The Restoration Flow

When your spot instance receives the termination warning, it should immediately attempt to save one final checkpoint. When you spin up a replacement instance, your training script checks persistent storage for the latest checkpoint, loads it, and resumes from that point. Instead of losing days of work, you only lose the progress made since the last checkpoint — typically 10–30 minutes depending on your checkpoint frequency.

Here is the core pattern in PyTorch:

import torch
import os

CHECKPOINT_PATH = "/mnt/checkpoints/checkpoint_latest.pt"

def save_checkpoint(model, optimizer, epoch, step, path):
    torch.save({
        'epoch': epoch,
        'step': step,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }, path)
    print(f"Checkpoint saved at epoch {epoch}, step {step}")

def load_checkpoint(model, optimizer, path):
    if not os.path.exists(path):
        print("No checkpoint found. Starting from scratch.")
        return 0, 0  # start_epoch, start_step
    checkpoint = torch.load(path)
    model.load_state_dict(checkpoint['model_state_dict'])
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    print(f"Resumed from epoch {checkpoint['epoch']}, step {checkpoint['step']}")
    return checkpoint['epoch'], checkpoint['step']

# --- At training startup ---
start_epoch, start_step = load_checkpoint(model, optimizer, CHECKPOINT_PATH)

# --- Inside the training loop ---
for epoch in range(start_epoch, NUM_EPOCHS):
    for step, batch in enumerate(dataloader):
        if epoch == start_epoch and step < start_step:
            continue  # Skip already-processed steps
        # ... training logic ...
        if step % CHECKPOINT_EVERY_N_STEPS == 0:
            save_checkpoint(model, optimizer, epoch, step, CHECKPOINT_PATH)

Handling the Termination Signal

On AWS, you can detect the 2-minute warning before it arrives by running a background thread that polls the EC2 Instance Metadata Service (IMDS) every few seconds for a spot/termination-time field; the moment that field appears, the thread triggers an emergency checkpoint save before the instance is reclaimed.

Key Point: Checkpoint frequency is a trade-off. Checkpointing too frequently (every step) adds significant I/O overhead and can slow training by 10–30%. Checkpointing too infrequently (every 10 epochs) risks losing hours of work per interruption. A practical rule of thumb: checkpoint every 15–30 minutes of wall-clock time, which balances overhead against recovery cost. For very large models where saving a checkpoint takes minutes, consider asynchronous checkpointing — saving to a background thread while training continues.