How to Checkpoint and Resume a Python Training Job to Survive Spot Instance Interruption

Cloud providers sell their spare compute capacity, the servers nobody else is renting right now, at a steep discount. AWS advertises savings of up to 90% off On-Demand prices for these “Spot Instances.” Google Cloud and Azure offer the same trade under the names “Spot VMs” and “Spot Virtual Machines.” The catch is right there in the name: it is spare capacity, and the provider can take it back with very little warning whenever a full-price customer wants it. That trade only makes sense for a training job if the job can survive losing its instance mid-run without losing its progress.

This tutorial builds that survival mechanism from scratch in Python: a training loop that saves everything it needs to resume exactly where it left off, saves it in a way that a badly timed crash cannot corrupt, and reacts to the operating system’s interruption signal by checkpointing immediately instead of just dying. Every command and script below was run end to end while writing this post, including two real interruption drills (one inside a single process, one across two separate processes) and a full test suite. No GPU, cloud account, or deep learning framework is required. Everything here is plain Python and NumPy, running on CPU, so you can see exactly what state gets saved and restored. The same pattern is a near-literal drop-in for PyTorch or any other framework once you understand it here.

What You’ll Build

By the end of this tutorial you will have:

• A small linear regression trainer whose entire state (weights, optimizer momentum, step counter, and random number generator) can be saved and restored as a single object.
• A checkpoint writer that cannot leave a corrupted file on disk, even if the process is killed in the middle of writing.
• A training script that catches the operating system’s termination signal, finishes checkpointing, and exits cleanly instead of just disappearing.
• Proof, not just an assertion, that resuming an interrupted run produces the exact same final model as a run that was never interrupted at all.

Prerequisites

• Python 3.10 or later. This tutorial used Python 3.13.14.
• pip install numpy pytest (versions used here: NumPy 2.5.2, pytest 9.1.1).
• A basic idea of what a training loop does: it repeatedly shows a model some examples, measures how wrong its predictions are, and nudges its parameters to make future predictions less wrong. If that sentence is new to you, that is genuinely all the machine learning background this tutorial assumes.
• Comfort running a script from a terminal and reading its output.

One honest platform note up front: this tutorial was developed and tested on Windows, and two of the “Common Mistakes” below are specifically about things that behave differently on Windows than on the Linux servers most cloud GPU instances actually run. Everywhere that matters, both behaviors are shown and explained, not just assumed.

Step 1: Build a Trainer Whose Entire State Fits in One Object

Start a new project folder and install the dependencies:

mkdir checkpoint-demo
cd checkpoint-demo
python -m venv .venv
.venv\Scripts\activate   # on Linux/macOS: source .venv/bin/activate
pip install numpy pytest

Create trainer.py:

"""A tiny linear-regression trainer with fully checkpointable state."""
import numpy as np


def make_dataset(n_samples=500, n_features=4, seed=42):
    rng = np.random.default_rng(seed)
    X = rng.normal(size=(n_samples, n_features))
    true_w = np.array([3.0, -2.0, 0.5, 1.5])
    true_b = -1.0
    noise = rng.normal(scale=0.1, size=n_samples)
    y = X @ true_w + true_b + noise
    return X, y, true_w, true_b


class Trainer:
    """Mini-batch gradient descent with momentum, over a fixed dataset.

    All state needed to resume training bit-for-bit lives on this object:
    weights, the momentum buffer, the step counter, and the RNG used to
    draw mini-batches.
    """

    def __init__(self, X, y, n_features, lr=0.05, momentum=0.9, batch_size=32, seed=0):
        self.X = X
        self.y = y
        self.lr = lr
        self.momentum = momentum
        self.batch_size = batch_size
        self.rng = np.random.default_rng(seed)
        self.w = np.zeros(n_features)
        self.b = 0.0
        self.v_w = np.zeros(n_features)
        self.v_b = 0.0
        self.step = 0

    def train_step(self):
        n = self.X.shape[0]
        idx = self.rng.integers(0, n, size=self.batch_size)
        Xb, yb = self.X[idx], self.y[idx]
        pred = Xb @ self.w + self.b
        err = pred - yb
        grad_w = (Xb.T @ err) / self.batch_size
        grad_b = err.mean()
        self.v_w = self.momentum * self.v_w - self.lr * grad_w
        self.v_b = self.momentum * self.v_b - self.lr * grad_b
        self.w = self.w + self.v_w
        self.b = self.b + self.v_b
        self.step += 1
        return float((err ** 2).mean())

    def state_dict(self):
        return {
            "w": self.w.copy(),
            "b": self.b,
            "step": self.step,
            "v_w": self.v_w.copy(),
            "v_b": self.v_b,
            "rng_state": self.rng.bit_generator.state,
        }

    def load_state_dict(self, state):
        self.w = np.array(state["w"], copy=True)
        self.b = float(state["b"])
        self.step = int(state["step"])
        if "v_w" in state:
            self.v_w = np.array(state["v_w"], copy=True)
            self.v_b = float(state["v_b"])
        if "rng_state" in state:
            self.rng.bit_generator.state = state["rng_state"]

A few design choices here matter more than they look. The model itself, linear regression trained with plain momentum, is deliberately boring. The point of this tutorial is checkpointing, not machine learning, so the model is simple enough that you can hold its entire state in your head. make_dataset() always builds the exact same 500-sample dataset from a fixed seed, so every run of this tutorial operates on identical data.

Notice that each training step draws a fresh random mini-batch of indices using self.rng, which is an instance attribute rather than part of the global state in numpy.random. That choice is what makes resuming reproducible at all: if two different Trainer objects have the RNG in the same internal state, calling train_step() on either one will draw the exact same “random” batch next. You’ll rely on that fact directly in Step 7.

Step 2: See What Happens With No Fault Tolerance at All

Before building the fix, it’s worth feeling the actual problem. Create naive_train.py:

"""No checkpointing, no signal handling: what most first training scripts look like."""
import sys
from trainer import make_dataset, Trainer

TOTAL_STEPS = 200
INTERRUPT_AT_STEP = int(sys.argv[1]) if len(sys.argv) > 1 else 150

X, y, _, _ = make_dataset()
trainer = Trainer(X, y, n_features=4, seed=1)
print("[naive] starting from step 0", flush=True)

while trainer.step < TOTAL_STEPS:
    loss = trainer.train_step()
    if trainer.step % 50 == 0:
        print(f"[naive] step {trainer.step:4d}  loss {loss:.4f}", flush=True)
    if trainer.step == INTERRUPT_AT_STEP:
        print(f"[naive] step {trainer.step}: the process gets killed right here "
              f"(no handler registered, so this is where it really ends)", flush=True)
        import signal
        signal.raise_signal(signal.SIGTERM)

print(f"[naive] training complete at step {trainer.step}", flush=True)

The signal.raise_signal(signal.SIGTERM) call stands in for a real interruption arriving from outside the process. Run it, and watch it die partway through:

$ python naive_train.py 150
[naive] starting from step 0
[naive] step   50  loss 0.1564
[naive] step  100  loss 0.0143
[naive] step  150  loss 0.0127
[naive] step 150: the process gets killed right here (no handler registered, so this is where it really ends)

Nothing was ever written to disk, so nothing survives. Run the exact same command again:

$ python naive_train.py 999999
[naive] starting from step 0
[naive] step   50  loss 0.1564
[naive] step  100  loss 0.0143
[naive] step  150  loss 0.0127
[naive] step  200  loss 0.0127
[naive] training complete at step 200

It starts over from step 0 and redoes all 150 steps of work that already happened once. On a rented GPU, that is not just an annoyance, it is money spent twice for the same result. Everything from here on is about closing that gap.

Step 3: Decide What Actually Belongs in a Checkpoint

It’s tempting to assume “checkpointing” just means saving the model’s weights. It doesn’t, and the gap between “saves weights” and “saves everything needed to resume correctly” is where most homegrown checkpointing breaks in ways that are hard to notice. A complete checkpoint for this trainer needs:

• The weights and bias. The obvious part: the actual learned parameters.
• The momentum buffer (v_w, v_b). Momentum-based optimizers keep a running average of recent gradients so they can build up speed in a consistent direction. If you reset that buffer to zero on every resume, you are not continuing training, you are restarting the optimizer’s memory while keeping its destination. Step 7 measures exactly how much this changes the outcome.
• The step counter. Needed for logging, learning rate schedules, and knowing when training is actually finished.
• The random number generator’s internal state. This is the one people skip most often, and the one with the largest measurable effect here. Every mini-batch is drawn from self.rng. If a resumed run gets a fresh RNG instead of the exact one from the moment of interruption, every mini-batch it draws afterward will differ from what the uninterrupted run would have seen, even with the same seed, because the seed only fixes where the sequence starts, not where it currently is.

Frameworks like PyTorch bundle the first two into what they call a “state dict,” which is exactly the naming convention Trainer.state_dict() and load_state_dict() borrow above. The Next Steps section at the end covers what changes when you move to a real framework.

Step 4: Save and Load Checkpoints Without Corrupting Them

The obvious way to save a checkpoint is to open the target file and write to it. The obvious way is also wrong, because if the process is killed while that write is still in progress, the file is left truncated on disk, and a truncated pickle is not a smaller checkpoint, it is a checkpoint that no longer loads at all.

The fix is a pattern used by databases and version control systems alike: write the new data to a temporary file first, then atomically rename that temporary file over the real target path. Python’s os module documentation states it plainly for the underlying rename operation: “If successful, the renaming will be an atomic operation (this is a POSIX requirement).” os.replace() extends that same guarantee to the case where the destination file already exists, silently replacing it instead of raising an error, which is exactly what a repeatedly-overwritten checkpoint needs. Because a rename either fully completes or does not happen at all, there is no window in which a reader can observe a half-written file at the final path.

Create checkpoint_io.py:

"""Checkpoint save/load, plus a crash simulator to prove atomic writes matter."""
import os
import pickle
import tempfile


class SimulatedCrash(Exception):
    """Raised to stand in for the process being killed mid-write."""


class CrashingWriter:
    """Wraps a file object; raises SimulatedCrash after N bytes are written."""

    def __init__(self, fileobj, crash_after_bytes=None):
        self._f = fileobj
        self._written = 0
        self._crash_after = crash_after_bytes

    def write(self, data):
        if self._crash_after is not None and self._written + len(data) > self._crash_after:
            allowed = max(0, self._crash_after - self._written)
            self._f.write(data[:allowed])
            self._f.flush()
            self._written += allowed
            raise SimulatedCrash(f"process died after {self._written} bytes")
        self._f.write(data)
        self._written += len(data)

    def flush(self):
        self._f.flush()


def save_checkpoint_naive(path, state, crash_after_bytes=None):
    """The wrong way: write straight to the final path."""
    with open(path, "wb") as f:
        writer = CrashingWriter(f, crash_after_bytes)
        pickle.dump(state, writer)


def save_checkpoint_atomic(path, state, crash_after_bytes=None):
    """The right way: write to a temp file, then atomically replace."""
    directory = os.path.dirname(os.path.abspath(path)) or "."
    fd, tmp_path = tempfile.mkstemp(dir=directory, prefix=".ckpt-tmp-")
    try:
        with os.fdopen(fd, "wb") as f:
            writer = CrashingWriter(f, crash_after_bytes)
            pickle.dump(state, writer)
        os.replace(tmp_path, path)
    except SimulatedCrash:
        try:
            os.remove(tmp_path)
        except FileNotFoundError:
            pass
        raise


def load_checkpoint(path):
    with open(path, "rb") as f:
        return pickle.load(f)

CrashingWriter is a small, deliberate piece of test scaffolding: it wraps a real file object and raises SimulatedCrash after a chosen number of bytes, standing in for the process being killed mid-write at a precise, reproducible point. Trying to reproduce a real kill by racing a signal against a write’s actual timing would be flaky by nature. This isn’t.

Now prove the difference. This first script simulates a crash while saving the very first checkpoint, when there is nothing on disk yet to lose:

from trainer import make_dataset, Trainer
from checkpoint_io import save_checkpoint_naive, save_checkpoint_atomic, load_checkpoint, SimulatedCrash
import os, pickle

X, y, _, _ = make_dataset()
t = Trainer(X, y, n_features=4, seed=1)
for _ in range(50):
    t.train_step()
state = t.state_dict()

serialized_len = len(pickle.dumps(state))
crash_point = int(serialized_len * 0.6)
print(f"full checkpoint is {serialized_len} bytes; simulating a crash after {crash_point} bytes\n")

naive_path = "naive_ckpt.pkl"
try:
    save_checkpoint_naive(naive_path, state, crash_after_bytes=crash_point)
except SimulatedCrash as e:
    print("NAIVE SAVE:", e)
print(f"naive checkpoint file left on disk: {os.path.getsize(naive_path)} bytes "
      f"(expected full size: {serialized_len})")
try:
    load_checkpoint(naive_path)
except Exception as e:
    print(f"naive checkpoint FAILED TO LOAD: {type(e).__name__}: {e}")

print()
atomic_path = "atomic_ckpt.pkl"
try:
    save_checkpoint_atomic(atomic_path, state, crash_after_bytes=crash_point)
except SimulatedCrash as e:
    print("ATOMIC SAVE (no prior checkpoint):", e)
print("does atomic_ckpt.pkl exist after the crash?", os.path.exists(atomic_path))

full checkpoint is 454 bytes; simulating a crash after 272 bytes

NAIVE SAVE: process died after 272 bytes
naive checkpoint file left on disk: 272 bytes (expected full size: 454)
naive checkpoint FAILED TO LOAD: UnpicklingError: pickle data was truncated

ATOMIC SAVE (no prior checkpoint): process died after 272 bytes
does atomic_ckpt.pkl exist after the crash? False

The naive path leaves a 272-byte fragment of a 454-byte file that cannot be unpickled at all. The atomic path leaves nothing at the target path, because the crash happened while writing the temporary file, before the rename that would have made it visible ever ran. No corrupted file beats a missing file, since a missing file is unambiguous: you know you have to fall back to whatever checkpoint you had before.

That “whatever you had before” case is the one that matters most in practice, because by the time you’re checkpointing repeatedly during a long training run, you always have an older good checkpoint sitting at that path already. Test what a crash does when it happens while overwriting one, continuing in the same script with training carried further:

good_state = state  # the step-50 checkpoint saved successfully above
for _ in range(50):
    t.train_step()
newer_state = t.state_dict()  # step 100, about to be interrupted mid-save
newer_crash_point = int(len(pickle.dumps(newer_state)) * 0.6)

for label, save_fn, path in [("naive", save_checkpoint_naive, "naive_ckpt2.pkl"),
                              ("atomic", save_checkpoint_atomic, "atomic_ckpt2.pkl")]:
    save_fn(path, good_state)  # a real, complete checkpoint at step 50
    try:
        save_fn(path, newer_state, crash_after_bytes=newer_crash_point)  # crashes saving step 100
    except SimulatedCrash as e:
        print(f"[{label}] crash while saving newer checkpoint: {e}")
    try:
        loaded = load_checkpoint(path)
        print(f"[{label}] checkpoint on disk after crash is READABLE, step={loaded['step']}")
    except Exception as e:
        print(f"[{label}] checkpoint on disk after crash FAILED TO LOAD: {type(e).__name__}: {e}")

[naive] good checkpoint saved at step 50
[naive] crash while saving newer checkpoint: process died after 271 bytes
[naive] checkpoint on disk after crash FAILED TO LOAD: UnpicklingError: pickle data was truncated

[atomic] good checkpoint saved at step 50
[atomic] crash while saving newer checkpoint: process died after 271 bytes
[atomic] checkpoint on disk after crash is READABLE, step=50

This is the case that actually costs you. The naive writer opens the target file in write mode immediately, which truncates the existing good checkpoint the instant the write starts, before any new data has even landed. A crash partway through does not just fail to update the checkpoint, it destroys the one you already had, leaving you with nothing at all. The atomic writer never touches the real path until the new file is completely written and ready, so the exact same crash at the exact same byte leaves the old, complete, step-50 checkpoint sitting there untouched.

Step 5: Catch the Interruption Signal

On Linux, the operating system most cloud GPU instances run, a process shutdown is delivered as a signal, and the specific signal used for “please stop soon” is SIGTERM. AWS documents exactly how this applies to Spot Instances in its Spot Instance interruption notices guide: “A Spot Instance interruption notice is a warning that is issued two minutes before Amazon EC2 stops or terminates your Spot Instance.” AWS delivers that notice two ways: as an EventBridge event, and as an item in the instance’s own metadata, retrievable via IMDSv2 at http://169.254.169.254/latest/meta-data/spot/instance-action, which AWS recommends polling “every 5 seconds.” Google Cloud’s Spot VMs and Azure’s Spot Virtual Machines use their own equivalent notice mechanisms with their own notice windows. Whichever cloud you’re on, the pattern is the same: something outside your training process learns the interruption is coming and needs to tell your process to wrap up, and on Linux that “telling” ultimately arrives as a signal your code can catch.

Python’s signal module lets you register a handler function that runs when a given signal arrives, instead of letting the operating system’s default behavior (typically, immediate termination) happen. Add this to a new file, train.py:

"""A training script that checkpoints periodically and shuts down cleanly on SIGTERM."""
import os
import signal
import sys

from trainer import make_dataset, Trainer
from checkpoint_io import save_checkpoint_atomic, load_checkpoint

CHECKPOINT_PATH = "training_ckpt.pkl"
TOTAL_STEPS = 200
CHECKPOINT_EVERY = 10

_shutdown_requested = False


def _handle_shutdown(signum, frame):
    global _shutdown_requested
    _shutdown_requested = True
    print(f"[trainer] received signal {signum}: will checkpoint and stop at the next safe point", flush=True)


def install_signal_handlers():
    signal.signal(signal.SIGTERM, _handle_shutdown)
    if sys.platform == "win32" and hasattr(signal, "SIGBREAK"):
        signal.signal(signal.SIGBREAK, _handle_shutdown)


def run():
    install_signal_handlers()
    X, y, _, _ = make_dataset()
    trainer = Trainer(X, y, n_features=4, seed=1)

    if os.path.exists(CHECKPOINT_PATH):
        state = load_checkpoint(CHECKPOINT_PATH)
        trainer.load_state_dict(state)
        print(f"[trainer] resumed from checkpoint at step {trainer.step}", flush=True)
    else:
        print("[trainer] no checkpoint found, starting from step 0", flush=True)

    while trainer.step < TOTAL_STEPS:
        loss = trainer.train_step()

        if trainer.step % CHECKPOINT_EVERY == 0:
            save_checkpoint_atomic(CHECKPOINT_PATH, trainer.state_dict())

        if trainer.step % 50 == 0:
            print(f"[trainer] step {trainer.step:4d}  loss {loss:.4f}", flush=True)

        if _shutdown_requested:
            save_checkpoint_atomic(CHECKPOINT_PATH, trainer.state_dict())
            print(f"[trainer] checkpointed at step {trainer.step}, exiting cleanly", flush=True)
            sys.exit(0)

    save_checkpoint_atomic(CHECKPOINT_PATH, trainer.state_dict())
    print(f"[trainer] training complete at step {trainer.step}, final loss {loss:.4f}", flush=True)


if __name__ == "__main__":
    run()

The handler itself does almost nothing: it just sets a flag and logs that a signal arrived. That’s intentional. Signal handlers in Python run on the main thread in between bytecode instructions, and can interrupt your code at points where it is not safe to do heavy work like writing a large file. The safe pattern is always the same: let the handler flip a flag, and let your normal loop check that flag at a point where it knows it’s safe to act, in this case, right after a training step completes.

install_signal_handlers() also registers signal.SIGBREAK, but only on Windows, and only if it exists. That line exists for Step 8, and the Gotchas section explains exactly why Windows needs it.

Step 6: Run the Interruption Drill

To make this reproducible without racing against wall-clock timing or thread scheduling, simulate the interruption notice arriving by calling signal.raise_signal() from inside the loop at an exact step. Python added raise_signal() specifically for this: the official documentation describes it plainly as sending “a signal to the calling process,” added in Python 3.8. As far as the registered handler is concerned, a self-raised signal and one delivered by an outside process are indistinguishable; both go through the exact same code path. Create demo_interrupted_run.py:

"""Reproduces exactly what happens when the training process in train.py
gets a SIGTERM partway through. In production, the cloud platform's own
agent sends that signal from a separate process; here we raise it on
ourselves at an exact step so the demo is reproducible instead of racing
against wall-clock timing. Either way, the training loop's own handler
can't tell the difference: a caught SIGTERM is a caught SIGTERM.
"""
import os
import signal
import sys

from trainer import make_dataset, Trainer
from checkpoint_io import save_checkpoint_atomic, load_checkpoint

CHECKPOINT_PATH = "training_ckpt.pkl"
TOTAL_STEPS = 200
CHECKPOINT_EVERY = 10
INTERRUPT_AT_STEP = int(sys.argv[1]) if len(sys.argv) > 1 else 73

_shutdown_requested = False


def _handle_shutdown(signum, frame):
    global _shutdown_requested
    _shutdown_requested = True
    print(f"[trainer] received signal {signum}: will checkpoint and stop at the next safe point", flush=True)


def main():
    signal.signal(signal.SIGTERM, _handle_shutdown)

    X, y, _, _ = make_dataset()
    trainer = Trainer(X, y, n_features=4, seed=1)

    if os.path.exists(CHECKPOINT_PATH):
        state = load_checkpoint(CHECKPOINT_PATH)
        trainer.load_state_dict(state)
        print(f"[trainer] resumed from checkpoint at step {trainer.step}", flush=True)
    else:
        print("[trainer] no checkpoint found, starting from step 0", flush=True)

    while trainer.step < TOTAL_STEPS:
        loss = trainer.train_step()

        if trainer.step == INTERRUPT_AT_STEP:
            print(f"[demo] simulating the interruption notice at step {trainer.step}", flush=True)
            signal.raise_signal(signal.SIGTERM)

        if trainer.step % CHECKPOINT_EVERY == 0:
            save_checkpoint_atomic(CHECKPOINT_PATH, trainer.state_dict())

        if _shutdown_requested:
            save_checkpoint_atomic(CHECKPOINT_PATH, trainer.state_dict())
            print(f"[trainer] checkpointed at step {trainer.step}, exiting cleanly", flush=True)
            sys.exit(0)

    save_checkpoint_atomic(CHECKPOINT_PATH, trainer.state_dict())
    print(f"[trainer] training complete at step {trainer.step}, final loss {loss:.4f}", flush=True)


if __name__ == "__main__":
    main()

Run it once to get interrupted, then run it again to resume:

$ python demo_interrupted_run.py 73
[trainer] no checkpoint found, starting from step 0
[demo] simulating the interruption notice at step 73
[trainer] received signal 15: will checkpoint and stop at the next safe point
[trainer] checkpointed at step 73, exiting cleanly

$ python demo_interrupted_run.py 999
[trainer] resumed from checkpoint at step 73
[trainer] training complete at step 200, final loss 0.0127

Signal 15 is the numeric value of SIGTERM, and it matches the “step 73” the demo asked for exactly, since raising the signal from inside the loop, rather than from an external timer, removes any timing uncertainty. The second run picks up at step 73 and finishes at step 200 with the same final loss (0.0127) an uninterrupted run reaches. That’s a good sign, but “the loss number looks the same” isn’t proof. The next step gets an actual proof.

Step 7: Prove the Resume Is Actually Correct

A resumed run that merely gets a reasonable-looking loss isn’t good enough for a serious claim of correctness. What you actually want to know is: does interrupting and resuming produce the identical model you would have gotten by never stopping at all? Test it directly by replaying both paths from the same seed and comparing the final weights:

import numpy as np
from trainer import make_dataset, Trainer

X, y, _, _ = make_dataset()

# Straight-through: 200 steps, never stopped.
straight = Trainer(X, y, n_features=4, seed=1)
for _ in range(200):
    straight.train_step()

# Interrupted-and-resumed: 73 steps, checkpoint, load into a brand new
# Trainer object (as if it were a fresh process), finish the remaining steps.
first_leg = Trainer(X, y, n_features=4, seed=1)
for _ in range(73):
    first_leg.train_step()
saved_state = first_leg.state_dict()

resumed = Trainer(X, y, n_features=4, seed=1)
resumed.load_state_dict(saved_state)
for _ in range(200 - 73):
    resumed.train_step()

print("straight-through final w:", straight.w)
print("resumed        final w:", resumed.w)
print("weights identical (np.array_equal):", np.array_equal(straight.w, resumed.w))

straight-through final w: [ 3.00093105 -2.00165165  0.4963147   1.52859812]
resumed        final w: [ 3.00093105 -2.00165165  0.4963147   1.52859812]
weights identical (np.array_equal): True

Not “close.” Identical, down to every bit, because load_state_dict() restores the RNG to the exact internal state it had at step 73, so the very next call to self.rng.integers(...) draws the exact same mini-batch the uninterrupted run would have drawn at that point, and every step after that follows the same deterministic path.

Now break it on purpose, the same two ways Step 3 warned about, and measure what actually happens rather than assuming:

broken_state = dict(saved_state)
del broken_state["rng_state"]
broken = Trainer(X, y, n_features=4, seed=1)
broken.load_state_dict(broken_state)  # rng falls back to a fresh seed=1 generator
for _ in range(200 - 73):
    broken.train_step()
print("max abs diff (no RNG state):", np.max(np.abs(straight.w - broken.w)))

broken2_state = dict(saved_state)
del broken2_state["v_w"]
del broken2_state["v_b"]
broken2 = Trainer(X, y, n_features=4, seed=1)
broken2.load_state_dict(broken2_state)
broken2.rng.bit_generator.state = saved_state["rng_state"]  # keep RNG correct this time
for _ in range(200 - 73):
    broken2.train_step()
print("max abs diff (no momentum state):", np.max(np.abs(straight.w - broken2.w)))

max abs diff (no RNG state): 0.03256805502829496
max abs diff (no momentum state): 8.017916268143743e-05

Dropping the RNG state produces a weight vector that still looks like a perfectly reasonable fit to a human eye, and would pass almost any casual sanity check, but it is measurably a different model: the largest single weight is off by about 0.033, roughly 400 times larger than the error caused by the other mistake. That’s the RNG state’s effect: every mini-batch drawn from that point forward is a different random sample than the original run would have seen, so the entire rest of training walks a different, though still valid, path.

Dropping the momentum state produces a much smaller effect here, a maximum difference of about 0.00008, small enough that you would probably never notice it by eye. It is still not zero. Resetting the velocity buffer to zero mid-training means the next few updates move slightly less far than they should, since momentum’s whole job is to keep building speed in a consistent direction. With only 127 steps left to run and a momentum coefficient of 0.9, that reset mostly washes out by the end, but a different momentum coefficient, a longer remaining run, or a more sensitive loss landscape can make this gap far more visible than it is here. The honest lesson isn’t “momentum state barely matters,” it’s “both of these are real, silent sources of non-reproducibility, and only one of them happened to be small in this particular setup.”

Step 8: Test It the Way a Real Interruption Would Actually Arrive

Every demo so far has one process signaling itself. That’s honest for proving the checkpoint math is correct, but a real Spot Instance interruption comes from an entirely separate process, the cloud platform’s own agent, running outside your training job. Prove the same handler works when the signal genuinely arrives from somewhere else. Create worker.py, a copy of the training loop with a small per-step delay so there’s time for an outside process to interrupt it:

"""The same training loop as train.py, with a small per-step delay so a
separate controller process has time to deliver a real interruption signal.
"""
import os
import signal
import sys
import time

from trainer import make_dataset, Trainer
from checkpoint_io import save_checkpoint_atomic, load_checkpoint

CHECKPOINT_PATH = "worker_ckpt.pkl"
TOTAL_STEPS = 200
CHECKPOINT_EVERY = 10

_shutdown_requested = False


def _handle_shutdown(signum, frame):
    global _shutdown_requested
    _shutdown_requested = True
    print(f"[worker pid={os.getpid()}] received signal {signum}: will checkpoint and stop", flush=True)


def install_signal_handlers():
    signal.signal(signal.SIGTERM, _handle_shutdown)
    if sys.platform == "win32" and hasattr(signal, "SIGBREAK"):
        signal.signal(signal.SIGBREAK, _handle_shutdown)


def main():
    install_signal_handlers()
    print(f"[worker pid={os.getpid()}] starting", flush=True)
    X, y, _, _ = make_dataset()
    trainer = Trainer(X, y, n_features=4, seed=1)

    if os.path.exists(CHECKPOINT_PATH):
        state = load_checkpoint(CHECKPOINT_PATH)
        trainer.load_state_dict(state)
        print(f"[worker pid={os.getpid()}] resumed from checkpoint at step {trainer.step}", flush=True)

    while trainer.step < TOTAL_STEPS:
        loss = trainer.train_step()
        time.sleep(0.03)  # slow the loop down so the demo has time to interrupt it

        if trainer.step % CHECKPOINT_EVERY == 0:
            save_checkpoint_atomic(CHECKPOINT_PATH, trainer.state_dict())

        if _shutdown_requested:
            save_checkpoint_atomic(CHECKPOINT_PATH, trainer.state_dict())
            print(f"[worker pid={os.getpid()}] checkpointed at step {trainer.step}, exiting cleanly", flush=True)
            sys.exit(0)

    save_checkpoint_atomic(CHECKPOINT_PATH, trainer.state_dict())
    print(f"[worker pid={os.getpid()}] training complete at step {trainer.step}, final loss {loss:.4f}", flush=True)


if __name__ == "__main__":
    main()

Then create agent.py, standing in for the cloud platform’s own preemption agent:

"""Stands in for the cloud platform's own preemption agent: launches the
training job, waits, then delivers a real interruption from a separate
process. On Linux this would be proc.send_signal(signal.SIGTERM); Windows
has no way to deliver an asynchronous signal to another process, so this
uses CTRL_BREAK_EVENT, its closest real analog (see the Gotchas section
for what happens if you reach for SIGTERM here instead).
"""
import signal
import subprocess
import sys
import time

CREATE_NEW_PROCESS_GROUP = 0x00000200

proc = subprocess.Popen([sys.executable, "worker.py"], creationflags=CREATE_NEW_PROCESS_GROUP)
print(f"[agent] launched worker (Popen pid={proc.pid})", flush=True)
time.sleep(1.5)
print("[agent] sending interruption notice now", flush=True)
proc.send_signal(signal.CTRL_BREAK_EVENT)
proc.wait(timeout=10)
print(f"[agent] worker exited with code {proc.returncode}", flush=True)

On Linux, replace creationflags=CREATE_NEW_PROCESS_GROUP and signal.CTRL_BREAK_EVENT with a plain subprocess.Popen([sys.executable, "worker.py"]) and proc.send_signal(signal.SIGTERM); the rest of the script, and the signal.signal(signal.SIGTERM, ...) handler already registered in worker.py, need no changes at all. Run the Windows version:

$ python agent.py
[agent] launched worker (Popen pid=2292)
[worker pid=18288] starting
[agent] sending interruption notice now
[worker pid=18288] received signal 21: will checkpoint and stop
[worker pid=18288] checkpointed at step 44, exiting cleanly
[agent] worker exited with code 0

$ python worker.py
[worker pid=15136] starting
[worker pid=15136] resumed from checkpoint at step 44
[worker pid=15136] training complete at step 200, final loss 0.0127

Signal 21 is SIGBREAK, the handler this exact script registered for Windows back in Step 5. The worker checkpoints at whatever step it happened to be on when the agent’s notice arrived (44, in this run; the exact step varies slightly with system timing) and picks up from there on the next run, landing on the same final loss as every earlier demo.

Common Mistakes and Gotchas

A crash mid-write can destroy your last good checkpoint, not just the new one

Step 4 already proved this, but it’s worth stating as its own rule: writing directly to your checkpoint’s final path is not “usually fine.” It truncates the existing file the moment the write starts, so any crash during that write can leave you with nothing, even though you had a perfectly good checkpoint one line of code earlier. Always write to a temporary file and rename it into place.

On Windows, sending SIGTERM to another process doesn’t run your handler, it hard-kills it

This one is easy to hit by accident if you develop on Windows and assume signals behave the same everywhere. Python’s own subprocess documentation says it directly: “On Windows, SIGTERM is an alias for terminate(),” and the documentation for Popen.terminate() explains what that means: “On Windows the Win32 API function TerminateProcess() is called to stop the child.” TerminateProcess() is an operating-system-level hard kill. It does not deliver anything a Python signal.signal() handler can catch, runs no cleanup code in the target process, and gives it no chance to checkpoint.

Proving it against the exact worker.py from Step 8: sending signal.SIGTERM instead of signal.CTRL_BREAK_EVENT looks almost identical from the agent’s side, but produces a silently worse outcome:

$ python agent_wrong.py   # same as agent.py, but proc.send_signal(signal.SIGTERM)
[agent] launched worker (Popen pid=5408)
[worker pid=11960] starting
[agent] sending SIGTERM (the wrong way on Windows)
[agent] worker exited with code 1

Notice what’s missing: no “received signal” line, no “checkpointed at step N, exiting cleanly” line, and an exit code of 1 instead of the clean 0. The worker’s registered handler never ran at all. A checkpoint file does exist afterward, but only because of the routine periodic save every 10 steps, not because of any graceful shutdown; loading it shows it stopped at step 40, whatever the last periodic checkpoint happened to be before the hard kill landed, silently discarding any progress made since. There is no error, no warning, nothing to tell you this happened differently than the graceful path. The fix is what Step 8 already does: register a handler for signal.SIGBREAK as well as signal.SIGTERM, and have the controlling process send signal.CTRL_BREAK_EVENT to a child launched with creationflags=subprocess.CREATE_NEW_PROCESS_GROUP. Python’s signal module documentation confirms CTRL_BREAK_EVENT is Windows-only and requires exactly that process group setup. On Linux, none of this applies: SIGTERM sent to a child process is delivered as a real signal and runs your handler exactly as documented.

A Windows virtual environment’s python.exe is not the process you think it is

While building the demos above, launching worker.py through a virtual environment’s Scripts\python.exe on Windows and printing proc.pid from the launching side consistently showed a different number than what os.getpid() reported from inside the worker script itself. In one run, the launching side saw pid 2292, while the worker printed pid 18288. This isn’t a fluke: on Windows, a virtual environment’s python.exe is a small launcher that starts the real interpreter as its own child process, so the PID subprocess.Popen hands you belongs to that launcher, not to the process actually running your code. It didn’t break anything in this tutorial’s demos, since CTRL_BREAK_EVENT is delivered to the whole process group rather than one specific PID, but it matters the moment you try to target a specific process ID directly, for example with os.kill(pid, ...), and get a PID that doesn’t point where you expect.

Checkpointing every single step is its own mistake

None of this is a reason to checkpoint as often as possible. Every checkpoint write costs real I/O time that isn’t spent training. train.py above checkpoints every 10 steps as a reasonable illustrative default, not a universal rule. The right frequency depends on the ratio between how long one training step takes and how long a checkpoint write takes: a step that takes minutes on a large model can afford to checkpoint often, while a step that takes milliseconds, like the one in this tutorial, would spend more time checkpointing than training if you did it on every step. Size the interval against your cloud provider’s actual notice window instead: if a checkpoint reliably takes 10 seconds to write and your platform gives you two minutes of warning, checkpointing every few minutes leaves ample margin; if it takes two minutes to write, you have a much bigger problem than checkpoint frequency.

How to Confirm It All Works End to End

Turn Step 7’s correctness proof and Step 4’s crash resistance proof into a test suite you can rerun any time you touch this code. Create test_checkpointing.py:

import numpy as np
import pytest

from trainer import make_dataset, Trainer
from checkpoint_io import (
    save_checkpoint_atomic,
    save_checkpoint_naive,
    load_checkpoint,
    SimulatedCrash,
)


@pytest.fixture
def dataset():
    return make_dataset()


def test_resume_matches_straight_through_exactly(dataset):
    X, y, _, _ = dataset
    straight = Trainer(X, y, n_features=4, seed=1)
    for _ in range(200):
        straight.train_step()

    first_leg = Trainer(X, y, n_features=4, seed=1)
    for _ in range(73):
        first_leg.train_step()
    state = first_leg.state_dict()

    resumed = Trainer(X, y, n_features=4, seed=1)
    resumed.load_state_dict(state)
    for _ in range(200 - 73):
        resumed.train_step()

    assert np.array_equal(straight.w, resumed.w)
    assert straight.b == resumed.b
    assert straight.step == resumed.step == 200


def test_resume_without_rng_state_diverges(dataset):
    X, y, _, _ = dataset
    straight = Trainer(X, y, n_features=4, seed=1)
    for _ in range(200):
        straight.train_step()

    first_leg = Trainer(X, y, n_features=4, seed=1)
    for _ in range(73):
        first_leg.train_step()
    state = first_leg.state_dict()
    del state["rng_state"]

    broken = Trainer(X, y, n_features=4, seed=1)
    broken.load_state_dict(state)
    for _ in range(200 - 73):
        broken.train_step()

    assert not np.array_equal(straight.w, broken.w)


def test_atomic_save_survives_a_crash_mid_write(tmp_path, dataset):
    X, y, _, _ = dataset
    trainer = Trainer(X, y, n_features=4, seed=1)
    for _ in range(50):
        trainer.train_step()
    good_state = trainer.state_dict()
    path = tmp_path / "ckpt.pkl"

    save_checkpoint_atomic(str(path), good_state)
    assert load_checkpoint(str(path))["step"] == 50

    for _ in range(50):
        trainer.train_step()
    newer_state = trainer.state_dict()

    with pytest.raises(SimulatedCrash):
        save_checkpoint_atomic(str(path), newer_state, crash_after_bytes=100)

    reloaded = load_checkpoint(str(path))
    assert reloaded["step"] == 50


def test_naive_save_corrupts_the_checkpoint_on_a_crash(tmp_path, dataset):
    X, y, _, _ = dataset
    trainer = Trainer(X, y, n_features=4, seed=1)
    for _ in range(50):
        trainer.train_step()
    state = trainer.state_dict()
    path = tmp_path / "ckpt.pkl"

    save_checkpoint_naive(str(path), state)
    assert load_checkpoint(str(path))["step"] == 50

    for _ in range(50):
        trainer.train_step()
    newer_state = trainer.state_dict()

    with pytest.raises(SimulatedCrash):
        save_checkpoint_naive(str(path), newer_state, crash_after_bytes=100)

    with pytest.raises(Exception):
        load_checkpoint(str(path))


def test_state_dict_round_trip_preserves_everything(dataset):
    X, y, _, _ = dataset
    trainer = Trainer(X, y, n_features=4, seed=7)
    for _ in range(10):
        trainer.train_step()
    state = trainer.state_dict()

    fresh = Trainer(X, y, n_features=4, seed=999)  # different seed on purpose
    fresh.load_state_dict(state)

    assert np.array_equal(fresh.w, trainer.w)
    assert fresh.b == trainer.b
    assert fresh.step == trainer.step
    assert np.array_equal(fresh.v_w, trainer.v_w)
    assert fresh.rng.bit_generator.state == trainer.rng.bit_generator.state

$ pytest test_checkpointing.py -v
test_checkpointing.py::test_resume_matches_straight_through_exactly PASSED
test_checkpointing.py::test_resume_without_rng_state_diverges PASSED
test_checkpointing.py::test_atomic_save_survives_a_crash_mid_write PASSED
test_checkpointing.py::test_naive_save_corrupts_the_checkpoint_on_a_crash PASSED
test_checkpointing.py::test_state_dict_round_trip_preserves_everything PASSED

5 passed in 0.13s

Each test guards one specific claim this tutorial made: that a correct resume is bit-for-bit identical to an uninterrupted run, that dropping RNG state measurably breaks that guarantee, that an atomic write survives a crash without losing the previous good checkpoint, that a naive write does not, and that every field on the trainer’s state dict round-trips through save and load without being silently dropped.

Next Steps

Everything above uses plain NumPy so the mechanics stay visible, but the same pattern applies almost unchanged to real training frameworks:

• In PyTorch, model.state_dict() and optimizer.state_dict() replace this tutorial’s w/b/v_w/v_b dictionary directly. You still need to save and restore RNG state explicitly, with torch.get_rng_state() and torch.random.set_rng_state() (plus torch.cuda.get_rng_state_all() if you’re training on GPU), and it’s just as easy to forget as it was here.
• If you’re using Hugging Face’s Trainer, it already implements most of this pattern for you through resume_from_checkpoint, including RNG state, as long as you let it manage checkpointing rather than writing your own save calls alongside it.
• For multi-GPU or multi-node training, a single pickled file stops scaling. Look at torch.distributed.checkpoint, which shards a checkpoint across ranks and can save and load them in parallel.
• Rather than waiting passively for a signal, you can poll your cloud provider’s own interruption-notice endpoint directly (AWS’s IMDSv2 spot/instance-action path used in this tutorial, or GCP’s and Azure’s equivalents) and trigger your own shutdown routine the moment a notice appears, which gives you the platform’s full warning window instead of however much of it survives by the time an external agent gets around to signaling your process.
• Consider checkpointing on a schedule that leaves real margin against your platform’s notice window rather than right up against it. AWS’s own two-minute warning is generous for a fast checkpoint write, but a large model’s checkpoint can itself take minutes to write to slow storage, and you want that write to finish well before the window closes.