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.

1. 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.

2. 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.

3. 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.

1. 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.

2. 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.

3. Inspect the checkpoint card

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

   

   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!