A 12-hour PyTorch training job on SageMaker costs too much, and the team wants to move it to Spot instances. A first attempt restarted from epoch 1 after an interruption and ran for 2 days. How do you set it up properly?

Turn on managed spot training with checkpoints the script really writes and reads. SageMaker syncs /opt/ml/checkpoints to S3 and copies it back after an interruption, but only your training script can save the model and optimiser state there and load it on start. Without that, every interruption starts over, which is what happened. Set MaxWaitTimeInSeconds high enough to cover waiting for capacity, and measure the saving from the billed time, not the wall clock.

1. The job settings

These are the fields of the training job (CreateTrainingJob); every SDK sets the same ones:

sm.create_training_job(
    TrainingJobName="churn-v4-spot",
    EnableManagedSpotTraining=True,
    StoppingCondition={
        "MaxRuntimeInSeconds": 14 * 3600,   # training time allowed
        "MaxWaitTimeInSeconds": 36 * 3600,  # training plus waiting for Spot, must be larger
    },
    CheckpointConfig={
        "S3Uri": "s3://ml-team/checkpoints/churn-v4/",
        "LocalPath": "/opt/ml/checkpoints",
    },
    ...  # image, role, instance type, input and output as before
)

The wait must be larger than the runtime. There is no fallback to on-demand: if Spot capacity never comes back, the job waits until MaxWaitTimeInSeconds and stops. If the result is needed by a deadline, leave room for that, or keep an on-demand run as the plan B.

2. The part the first attempt missed: the script

SageMaker copies whatever is in S3 into /opt/ml/checkpoints when the job starts or restarts. Your script must look there first:

ckpt = "/opt/ml/checkpoints/last.pt"
start_epoch = 0
if os.path.exists(ckpt):
    state = torch.load(ckpt)
    model.load_state_dict(state["model"])
    optimizer.load_state_dict(state["optimizer"])
    scheduler.load_state_dict(state["scheduler"])
    start_epoch = state["epoch"] + 1

for epoch in range(start_epoch, epochs):
    train_one_epoch(...)
    torch.save({"model": model.state_dict(), "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(), "epoch": epoch}, ckpt)

Save the optimiser and scheduler as well as the weights, or the resumed run trains with a reset learning rate and diverges from the original. How often to save is up to the script; the SDK has no setting for it. Every epoch is a fair default when an epoch takes minutes; for very long epochs, save every N steps as well.

3. More than 1 instance

With distributed training, all instances share one S3 checkpoint location. Write from rank 0 only, or put the rank in the file name, or instances overwrite each other.

4. Check what you saved

After the job, the training job's description gives 2 figures:

savings = (1 − BillableTimeInSeconds ÷ TrainingTimeInSeconds) × 100

If billable time is 3 hours against 12 hours of training time, the saving is 75%. The wall-clock time from start to end includes interruptions and waiting for capacity, so it can be much longer than either figure and is not what you pay for.

5. When Spot is a poor fit

• A built-in algorithm that does not checkpoint: the maximum wait is capped at 1 hour, so a long job may never finish.
• A hard deadline with no slack.
• Very short jobs, where checkpointing adds little and waits for capacity dominate.

What interviewers want: knowing that SageMaker moves checkpoint files but the script must write and load them, the full training state in a checkpoint, the wait limit above the runtime limit and the lack of on-demand fallback, rank-aware checkpoints in distributed training, and measuring savings from billable time.