Checkpointing Training to Amazon S3

# Option 1: Environment variables
export AWS_ACCESS_KEY_ID=<your-access-key>
export AWS_SECRET_ACCESS_KEY=<your-secret-key>
export AWS_DEFAULT_REGION=us-east-1

# Option 2: AWS credentials file (~/.aws/credentials)
aws configure

pip install s3torchconnector

from s3torchconnector import S3Checkpoint

# Initialize the checkpoint handler
checkpoint = S3Checkpoint(region="us-east-1")

# Save model checkpoint directly to S3
checkpoint_path = "s3://my-training-checkpoints/model_epoch_10.pt"

with checkpoint.writer(checkpoint_path) as writer:
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': loss,
    }, writer)

from s3torchconnector import S3Checkpoint

checkpoint = S3Checkpoint(region="us-east-1")
checkpoint_path = "s3://my-training-checkpoints/model_epoch_10.pt"

with checkpoint.reader(checkpoint_path) as reader:
    checkpoint_data = torch.load(reader)

model.load_state_dict(checkpoint_data['model_state_dict'])
optimizer.load_state_dict(checkpoint_data['optimizer_state_dict'])
start_epoch = checkpoint_data['epoch']

pip install lightning s3fs

from lightning.pytorch import Trainer

trainer = Trainer(
    # Lightning uses fsspec to write directly to S3
    default_root_dir="s3://my-training-checkpoints/experiment-1/",
    enable_checkpointing=True,
)

trainer.fit(model)

from lightning.pytorch import Trainer

trainer = Trainer(
    default_root_dir="s3://my-training-checkpoints/experiment-1/",
    enable_checkpointing=True,
)

# Resume from a specific checkpoint
trainer.fit(
    model,
    ckpt_path="s3://my-training-checkpoints/experiment-1/checkpoints/last.ckpt",
)

# Ubuntu/Debian
sudo apt-get install s3fs

# Amazon Linux
sudo yum install s3fs-fuse

# No code changes needed - save as if it were local storage
checkpoint_dir = "/mnt/s3-checkpoints/my-experiment"

torch.save({
    'epoch': epoch,
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': loss,
}, f"{checkpoint_dir}/checkpoint_epoch_{epoch}.pt")

fusermount -u /mnt/s3-checkpoints

pip install awscli

import subprocess
import torch

def save_checkpoint_and_sync(model, optimizer, epoch, local_dir, s3_bucket):
    # Save locally first
    checkpoint_path = f"{local_dir}/checkpoint_epoch_{epoch}.pt"
    torch.save({
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }, checkpoint_path)

    # Sync to S3
    subprocess.run([
        "aws", "s3", "sync",
        local_dir,
        f"s3://{s3_bucket}/checkpoints/",
        "--exclude", "*",
        "--include", "*.pt"
    ], check=True)

# Usage in training loop
for epoch in range(start_epoch, num_epochs):
    train_one_epoch(model, dataloader, optimizer)
    save_checkpoint_and_sync(
        model, optimizer, epoch,
        local_dir="/tmp/checkpoints",
        s3_bucket="my-training-checkpoints"
    )