Metaflow Checkpoint

0.2.10 · active · verified Thu Apr 16

Metaflow-checkpoint is an experimental extension for Metaflow that provides in-task checkpointing capabilities. It allows users to periodically save the progress of long-running Metaflow steps, such as machine learning model training, ensuring recovery from failures without losing significant work. The library, currently at version 0.2.10, is released as an independent extension to core Metaflow.

Common errors

```
NameError: name 'checkpoint' is not defined
```
cause The `metaflow-checkpoint` extension was not installed, or the `checkpoint` decorator was not imported correctly from `metaflow`.

fix First, ensure the library is installed: `pip install metaflow-checkpoint`. Then, import it alongside other Metaflow components: `from metaflow import FlowSpec, step, current, checkpoint`.
```
My flow didn't resume from the last saved state when I restarted it with `python my_flow.py run`.
```
cause The default `load_policy='fresh'` for the `@checkpoint` decorator is designed for recovery within the *same* run (e.g., after a `@retry`). It deliberately does not load checkpoints from previous runs.

fix If you intend to resume from a previous run, particularly during development, change the decorator to `@checkpoint(load_policy='eager')`. For more complex scenarios, consider `load_policy=None` and implement custom loading via `current.checkpoint.load()`.
```
TypeError: 'MetaflowCheckpoint' object is not callable
```
cause Attempting to call `current.checkpoint` as a function, e.g., `current.checkpoint()` or incorrectly assuming `current.checkpoint` itself is the decorator.

fix The `checkpoint` symbol imported from `metaflow` is the decorator (`@checkpoint`). `current.checkpoint` is an object available *inside* a decorated step, providing methods like `current.checkpoint.save()` and properties like `current.checkpoint.directory`.

Warnings

breaking The `metaflow-checkpoint` library is explicitly labeled as EXPERIMENTAL. Its APIs may change in future versions, and it does not offer the same backwards compatibility guarantees as core Metaflow APIs.
Fix: Always review the latest Metaflow documentation and `metaflow-checkpoint` release notes when upgrading to new versions. Be prepared for potential code adjustments.
gotcha The default `load_policy='fresh'` for `@checkpoint` only loads task-specific checkpoints for retries within the same run. It explicitly disregards existing checkpoints when a *new* run is initiated.
Fix: For iterative development and resuming across different runs (e.g., stopping a flow and restarting later), use `@checkpoint(load_policy='eager')`. For custom loading logic, use `@checkpoint(load_policy=None)` and manually call `current.checkpoint.load()`.
gotcha Files saved to `current.checkpoint.directory` can accumulate across invocations if not managed, potentially leading to performance degradation over time as `current.checkpoint.save()` processes more data.
Fix: Ensure that your checkpointing logic overwrites existing files within `current.checkpoint.directory` or explicitly cleans up the directory between checkpoint saves to prevent excessive file accumulation. For example, always use the same filename for your latest model state.

Install

pip install metaflow-checkpoint Install latest version

Imports

checkpoint
wrong:
```
from metaflow_checkpoint import checkpoint
```
correct:
```
from metaflow import FlowSpec, step, current, checkpoint
```
The `checkpoint` decorator is designed to be imported directly from `metaflow` after `metaflow-checkpoint` is installed, not from a separate module path.
current.checkpoint.save
wrong:
```
checkpoint.save()
```
correct:
```
current.checkpoint.save()
```
`save()` is a method on the `current.checkpoint` object within a Metaflow step, not a standalone function or method on the `checkpoint` decorator itself.

Quickstart

This quickstart demonstrates a Metaflow flow using the `@checkpoint` decorator. The `start` step simulates a long-running, flaky process that increments a counter. It saves the counter value to `current.checkpoint.directory` and calls `current.checkpoint.save()` periodically. Upon restart (due to `@retry` or `resume` command), it loads the last saved counter using `current.checkpoint.is_loaded` and `current.checkpoint.directory`. The `load_policy='eager'` allows checkpoints to be reused across different runs, aiding iterative development. Run with `python your_flow.py run` and try `python your_flow.py resume start` after an interruption.

import os
import random
from metaflow import FlowSpec, step, current, checkpoint, retry

class CheckpointCounterFlow(FlowSpec):
    @retry(times=2, minutes_between_retries=1)
    @checkpoint(load_policy='eager') # Use 'eager' for development across runs
    @step
    def start(self):
        self.counter = 0
        if current.checkpoint.is_loaded:
            print(f"Resuming from checkpoint. Counter was {self.counter}")
            with open(os.path.join(current.checkpoint.directory, 'counter'), 'r') as f:
                self.counter = int(f.read())
            print(f"Successfully loaded counter: {self.counter}")
        else:
            print("Starting from scratch.")

        for i in range(5):
            self.counter += 1
            print(f"Processing iteration {i+1}, counter is {self.counter}")
            # Save progress periodically
            with open(os.path.join(current.checkpoint.directory, 'counter'), 'w') as f:
                f.write(str(self.counter))
            current.checkpoint.save()

            # Simulate a flaky operation
            if random.random() < 0.3:
                raise Exception("Simulated failure!")

        self.next(self.end)

    @step
    def end(self):
        print(f"Flow finished. Final counter value: {self.counter}")

if __name__ == '__main__':
    CheckpointCounterFlow()

view raw JSON →