Training models that recover from interruptions

By default, Metaflow saves a step’s results only after the step finishes. If training fails three hours in—because a node goes down, a dependency times out, or the cloud provider reclaims the instance—you lose those three hours of progress, and the retry starts from the beginning. Using the @checkpoint and @retry decorators, you can build a training flow that saves its progress while it runs, so a failure costs you minutes of rework instead of hours. If the training step fails, Anaconda Platform retries it, and the new attempt picks up where the last one left off.

Prerequisites

Before you begin, make sure you have the following:

A macOS or Linux machine, or a cloud workstation.
Metaflow is built for macOS and Linux. Windows machines require Windows Subsystem for Linux (WSL) to run Metaflow.
The outerbounds package installed and connected to your platform.
For setup instructions, see Connect to Anaconda Platform and run your first flow.
Access to a compute pool for running Metaflow Tasks.

How checkpointing works

Most machine learning libraries can save a model to a file, but that file stays on the machine that wrote it. If the machine fails, a retry on a new machine can’t find it. @checkpoint solves this by saving checkpoints to your platform’s storage, so any retry can load them. By default, each task loads only the checkpoints it saved, so parallel training tasks, your colleagues’ runs, and production deployments never pick up each other’s checkpoints. A card in the Runs view shows which checkpoints each attempt loaded and saved. Adding the @checkpoint decorator to a step lets you work with checkpoints through the current.checkpoint object. This tutorial uses the following:

Writing a training step that saves checkpoints

In this section, you create two files: a module that saves checkpoints during XGBoost training and the flow that uses it.

Create a project directory

Create a directory for the tutorial and navigate to it:

mkdir checkpoint-tutorial && cd checkpoint-tutorial

Metaflow packages every Python file in this directory with your flow, so the flow can import the module you create next, both locally and in the cloud.

Create the training callbacks

XGBoost calls a training callback after every boosting round, which makes callbacks a natural place to save progress. Create a file named callbacks.py with the following contents:

callbacks.py

import os
import pickle

import xgboost
from metaflow import current


class Checkpointer(xgboost.callback.TrainingCallback):
    def __init__(self, interval=10):
        self.interval = interval

    @staticmethod
    def path():
        return os.path.join(current.checkpoint.directory, "model.pkl")

    def after_iteration(self, model, epoch, evals_log):
        rounds = model.num_boosted_rounds()
        if rounds % self.interval == 0:
            with open(self.path(), "wb") as f:
                pickle.dump(model, f)
            current.checkpoint.save()
            print(f"Saved a checkpoint at round {rounds}")
        return False

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


class SimulatedFailure(xgboost.callback.TrainingCallback):
    def __init__(self, after_rounds):
        self.after_rounds = after_rounds
        self.rounds_this_attempt = 0

    def after_iteration(self, model, epoch, evals_log):
        self.rounds_this_attempt += 1
        if self.rounds_this_attempt == self.after_rounds:
            raise RuntimeError(
                f"Simulated failure at round {model.num_boosted_rounds()}"
            )
        return False

The file defines two callbacks:

Checkpointer: Every 10 rounds, writes the model to model.pkl in the checkpoint directory and calls current.checkpoint.save(). Its load() method reads the model back after @checkpoint restores the directory on a retry.
SimulatedFailure: Raises an error after a set number of rounds in each attempt. It stands in for the real failures you’re protecting against, so you can watch recovery happen on demand. Leave it out of your own projects.

Checkpointer overwrites the same file every time. Each call to current.checkpoint.save() uploads the whole checkpoint directory, so writing new filenames each time would make every save larger and slower than the last.

Create the flow

Create a file named train_flow.py with the following contents:

train_flow.py

from metaflow import (
    FlowSpec,
    Parameter,
    anaconda_base,
    checkpoint,
    current,
    retry,
    step,
)


@anaconda_base(
    python="3.12",
    packages={"xgboost": "3.4.1", "scikit-learn": "1.9.1"},
)
class ResilientTrainingFlow(FlowSpec):
    rounds = Parameter("rounds", help="Total boosting rounds", default=128)
    crash_every = Parameter(
        "crash_every",
        help="Rounds per attempt before a simulated failure (0 disables)",
        default=40,
    )

    @step
    def start(self):
        from sklearn.datasets import load_breast_cancer

        self.X, self.y = load_breast_cancer(return_X_y=True)
        self.next(self.train)

    @checkpoint
    @retry(times=3)
    @step
    def train(self):
        import xgboost as xgb
        from callbacks import Checkpointer, SimulatedFailure

        if current.checkpoint.is_loaded:
            start_model = Checkpointer.load()
            done = start_model.num_boosted_rounds()
            print(f"Resuming from a checkpoint trained for {done} rounds")
        else:
            start_model = None
            done = 0
            print("No checkpoint loaded, training from scratch")

        callbacks = [Checkpointer(interval=10)]
        if self.crash_every:
            callbacks.append(SimulatedFailure(after_rounds=self.crash_every))

        self.model = xgb.XGBClassifier(
            n_estimators=self.rounds - done, callbacks=callbacks
        )
        self.model.fit(self.X, self.y, xgb_model=start_model)
        self.model.set_params(callbacks=None)
        self.total_rounds = self.model.get_booster().num_boosted_rounds()
        self.attempts = current.retry_count + 1
        self.next(self.end)

    @step
    def end(self):
        print(f"Trained {self.total_rounds} rounds. Attempts used: {self.attempts}")


if __name__ == "__main__":
    ResilientTrainingFlow()

The highlighted lines are where checkpointing happens:

@retry(times=3) gives the train step up to three more attempts after a failure, for four attempts in total.
@checkpoint saves progress within each attempt and, on a retry, loads the latest checkpoint into current.checkpoint.directory before your code runs.
The if current.checkpoint.is_loaded block decides where training starts. On a retry, it loads the saved model and counts the rounds already completed. XGBoost then trains only the remaining rounds, continuing from the loaded model through the xgb_model argument.

The start step loads a small sample dataset from scikit-learn, and @anaconda_base declares the packages both steps need, sourced from Anaconda’s main channel.Order matters in the callbacks list. Checkpointer runs before SimulatedFailure, so the checkpoint at round 40 is saved before the simulated failure at round 40.After training finishes, set_params(callbacks=None) removes the callbacks from the trained model before Metaflow stores it as the model artifact. Without this line, loading the artifact later, for example in a notebook, would fail unless callbacks.py was importable.

Your checkpoint-tutorial directory now contains callbacks.py and train_flow.py.

Watching the flow recover

Run the flow on your machine first. It’s the quickest way to see the full cycle of failures, retries, and resumed training.

Run the flow

From the checkpoint-tutorial directory, run the flow:

python train_flow.py --environment=anaconda run

The first run takes a few minutes longer than subsequent runs because Metaflow needs to build the environment declared in @anaconda_base.

Follow the attempts in the output

The train step fails three times, and each failure prints a full traceback ending in RuntimeError: Simulated failure. These errors are expected. Between the tracebacks, look for lines like the following:

Example output (trimmed)

Workflow starting (run-id <RUN_ID>), see it in the UI at https://ui.<YOUR_DOMAIN>/p/default/ResilientTrainingFlow/<RUN_ID>
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 16234)] No checkpoint loaded, training from scratch
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 16234)] Saved a checkpoint at round 10
...
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 16234)] Saved a checkpoint at round 40
...
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 16234)] RuntimeError: Simulated failure at round 40
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 16234)] Task failed.
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 16616)] Task is starting (retry).
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 16616)] [@checkpoint] Loading the following checkpoint:
...
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 16616)] Resuming from a checkpoint trained for 40 rounds
...
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 17096)] Resuming from a checkpoint trained for 80 rounds
...
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 17755)] Resuming from a checkpoint trained for 120 rounds
[<RUN_ID>/train/<TRAIN_TASK_ID> (pid 17755)] Task finished successfully.
[<RUN_ID>/end/<END_TASK_ID> (pid 23630)] Trained 128 rounds. Attempts used: 4
Done! See the run in the UI at https://ui.<YOUR_DOMAIN>/p/default/ResilientTrainingFlow/<RUN_ID>

The pid changes with each attempt because every retry runs as a new process with nothing left in memory from the failed attempt. Only the checkpoint carries progress forward.

Expected result: The run ends with Trained 128 rounds. Attempts used: 4. Without @checkpoint, every retry would start from round zero and fail at round 40 again, and the run would fail after its fourth attempt.

Inspect the checkpoint card

Follow the link to view your run in the UI, select the train task, then select the Cards tab.

The task view for the train task, with the task selected in the task list and the Cards tab open, showing a checkpoint card for each attempt

Expected result: The task has one checkpoint card for each attempt. Select the last task_checkpoints tab to see the final attempt’s card, which shows the checkpoint it loaded from the previous attempt and the checkpoints it saved.

Running the flow on cloud compute

Real training jobs run in the cloud, where failures are more likely. Nodes get replaced, and a retried task often lands on a different instance from the one that failed. Because @checkpoint stores checkpoints in your platform’s datastore rather than on the instance’s disk, the new instance can still load them. Run the same flow on cloud compute. You don’t need to change any code:

python train_flow.py --environment=fast-bakery run --with kubernetes

Anaconda Platform builds the environment into a container image and runs each attempt in its own container on the cluster. The first run takes a few minutes while the image builds. Expected result: The output shows the same sequence of failures and resumed attempts, and the run ends with Trained 128 rounds. Attempts used: 4.

Checkpoints and spot instances

Checkpointing lets you reliably train you models on spot instances, which cloud providers sell from their spare capacity at a steep discount. The tradeoff is that the provider can reclaim a spot instance whenever it needs the capacity back, with as little as two minutes’ notice. To your flow, a reclaimed instance is just another failure, like the simulated failures in this tutorial. @retry starts a new attempt on a different instance, and @checkpoint resumes training from the latest checkpoint, so a long training job can run on cheaper capacity and lose only the work since its last save.

On AWS spot instances, Metaflow also signals the two-minute warning through current.spot_termination_notice, so your training code can save one last checkpoint before the instance shuts down. For an example, see Using spot instances.

Tutorial wrap up

You’ve built a training flow that survives failures. Each time the train step failed, Anaconda Platform retried it, and the new attempt picked up from the latest checkpoint instead of starting over. What you learned

Why artifacts alone can’t protect progress inside a long-running step
How to save progress during a task with current.checkpoint.directory and current.checkpoint.save()
How @checkpoint and @retry work together to resume a failed task from its latest checkpoint
How to inspect loaded and saved checkpoints in the checkpoint card
Why checkpoints make lower-cost spot instances practical for long training jobs

Next steps

Selecting a checkpoint to use explains load policies, including how to resume training across separate runs during development.
The Metaflow checkpoint examples repository applies @checkpoint to PyTorch, PyTorch Lightning, Keras, Hugging Face Transformers, and distributed training.
Cost optimization covers other ways to lower compute costs on Anaconda Platform.

Have feedback? Use the feedback buttons below to let us know what you thought about this tutorial and help us improve!

​Prerequisites

​How checkpointing works

​Writing a training step that saves checkpoints

​Watching the flow recover

​Running the flow on cloud compute

​Checkpoints and spot instances

​Tutorial wrap up

Prerequisites

How checkpointing works

Writing a training step that saves checkpoints

Watching the flow recover

Running the flow on cloud compute

Checkpoints and spot instances

Tutorial wrap up