> ## Documentation Index
> Fetch the complete documentation index at: https://anaconda.com/docs/llms.txt
> Use this file to discover all available pages before exploring further.
# 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.
Advanced
**Purpose**: This tutorial teaches you how to make long-running training steps survive failures by checkpointing progress inside a running task.
**Outcome**: You'll train an XGBoost model in a flow that has built-in failures but still completes successfully, resuming from its latest checkpoint after each failure. Afterwards, you can apply the same pattern to your own training code.
**Audience**: Data scientists and ML engineers who already run flows on Anaconda Platform.
**Duration**: 30 minutes, including environment build time
## Prerequisites
Before you begin, make sure you have the following:
* A macOS or Linux machine, or a [cloud workstation](https://anaconda.com/docs/platform/concepts/what-is-a-workstation.md).
> [!NOTE]
>
> Metaflow is built for macOS and Linux. Windows machines require [Windows Subsystem for Linux (WSL)](https://learn.microsoft.com/en-us/windows/wsl/install) to run Metaflow.
* The `outerbounds` package installed and connected to your platform.
> [!TIP]
>
> For setup instructions, see [Connect to Anaconda Platform and run your first flow](https://anaconda.com/docs/platform/getting-started/connect-and-first-run.md).
* 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:
Name
What it does
`current.checkpoint.directory`
A temporary local directory where you write the files you need to resume work, such as a saved model
`current.checkpoint.save()`
Uploads the contents of that directory to your platform's storage as a new checkpoint
`current.checkpoint.is_loaded`
`True` when a retried attempt starts with an earlier checkpoint already restored to the directory
LearnMore
**Learn more**
Review the following resources to learn more about the topics covered in this section, or continue on to the next section.
* [Dealing with failures](https://docs.metaflow.org/scaling/failures)
* [Checkpointing progress](https://docs.metaflow.org/scaling/checkpoint/introduction)
## 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:
```sh theme={null}
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](https://xgboost.readthedocs.io/en/stable/python/callbacks.html) after every boosting round, which makes callbacks a natural place to save progress. Create a file named `callbacks.py` with the following contents:
```python callbacks.py expandable theme={null}
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.
> [!TIP]
>
> `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:
```python train_flow.py expandable highlight={31-32,38-45} theme={null}
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`.
LearnMore
**Learn more**
Review the following resources to learn more about the topics covered in this section, or continue on to the next section.
* [Define the environment](https://anaconda.com/docs/platform/getting-started/defining-the-environment.md)
* [Managing dependencies](https://anaconda.com/docs/platform/guides/compute/managing-dependencies.md)
* [Checkpoints in ML/AI libraries](https://docs.metaflow.org/scaling/checkpoint/checkpoint-ml-libraries)
## 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:
```sh theme={null}
python train_flow.py --environment=anaconda run
```
> [!TIP]
>
> The first run takes a few minutes longer than subsequent runs because Metaflow needs to build the environment declared in `@anaconda_base`.
Troubleshoot
The run fails with a 403 error from `repo.anaconda.cloud`
***
Cause
Your conda configuration lists authenticated Anaconda channels. Metaflow builds local environments with micromamba, which reads your conda configuration but can't use your stored Anaconda token.
***
Solution
Run the flow with micromamba's configuration files turned off:
```sh theme={null}
MAMBA_NO_RC=true python train_flow.py --environment=anaconda run
```
**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:
```text Example output (trimmed) expandable theme={null}
Workflow starting (run-id ), see it in the UI at https://ui./p/default/ResilientTrainingFlow/
[/train/ (pid 16234)] No checkpoint loaded, training from scratch
[/train/ (pid 16234)] Saved a checkpoint at round 10
...
[/train/ (pid 16234)] Saved a checkpoint at round 40
...
[/train/ (pid 16234)] RuntimeError: Simulated failure at round 40
[/train/ (pid 16234)] Task failed.
[/train/ (pid 16616)] Task is starting (retry).
[/train/ (pid 16616)] [@checkpoint] Loading the following checkpoint:
...
[/train/ (pid 16616)] Resuming from a checkpoint trained for 40 rounds
...
[/train/ (pid 17096)] Resuming from a checkpoint trained for 80 rounds
...
[/train/ (pid 17755)] Resuming from a checkpoint trained for 120 rounds
[/train/ (pid 17755)] Task finished successfully.
[/end/ (pid 23630)] Trained 128 rounds. Attempts used: 4
Done! See the run in the UI at https://ui./p/default/ResilientTrainingFlow/
```
> [!NOTE]
>
> 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.
LearnMore
**Learn more**
Review the following resources to learn more about the topics covered in this section, or continue on to the next section.
* [Visualizing results with cards](https://docs.metaflow.org/metaflow/visualizing-results)
* [Inspecting runs with the Client API](https://docs.metaflow.org/metaflow/client)
## 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:
```sh theme={null}
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.
> [!NOTE]
>
> 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](https://docs.metaflow.org/scaling/remote-tasks/spot-instances).
LearnMore
**Learn more**
Review the following resources to learn more about the topics covered in this section, or continue on to the next section.
* [Scale to the cloud](https://anaconda.com/docs/platform/getting-started/scaling-to-the-cloud.md)
* [What is compute?](https://anaconda.com/docs/platform/concepts/what-is-compute.md)
* [Cost optimization](https://anaconda.com/docs/platform/guides/monitor/cost-optimization-overview.md)
## 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](https://docs.metaflow.org/scaling/checkpoint/selecting-checkpoints) explains load policies, including how to resume training across separate runs during development.
* The [Metaflow checkpoint examples](https://github.com/outerbounds/metaflow-checkpoint-examples) repository applies `@checkpoint` to PyTorch, PyTorch Lightning, Keras, Hugging Face Transformers, and distributed training.
* [Cost optimization](https://anaconda.com/docs/platform/guides/monitor/cost-optimization-overview.md) 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!