Using Checkpoints in Tasks

Resume task work with a checkpoint

To resume progress after a retry, get the task-scoped helper from flyte.ctx().checkpoint, load the previous state before doing work, and save a new state before each failure-prone progress point. flyte.ctx() is intended for task code, not module-level code, and returns the current flyte.models.TaskContext or None outside a task.

Synchronous tasks

For a regular def task, use load_sync() and save_sync(). This is the complete byte-based pattern used by examples/checkpoint/generic_data_checkpoint.py:

from __future__ import annotations

import flyte

env = flyte.TaskEnvironment(
    name="checkpoint_generic_json",
    image=flyte.Image.from_debian_base(),
)

RETRIES = 3


@env.task(retries=RETRIES)
def use_checkpoint(n_iterations: int) -> int:
    assert n_iterations > RETRIES
    checkpoint = flyte.ctx().checkpoint

    path = checkpoint.load_sync()
    prev = None if path is None else path.read_bytes()

    start = 0
    if prev:
        start = int(prev.decode())

    failure_interval = n_iterations // RETRIES
    for index in range(start, n_iterations):
        if index > start and index % failure_interval == 0:
            raise RuntimeError(f"Failed at iteration {index}, failure_interval {failure_interval}.")
        checkpoint.save_sync(f"{index + 1}".encode())
    return index

load_sync() returns a local pathlib.Path when a previous checkpoint was restored. The example reads that path as bytes, decodes the saved iteration, and starts the loop from that value. Each save_sync() call uploads the next iteration as raw bytes, so a retry can continue from the most recently uploaded value.

The helper is optional. A task that can run without checkpoint output should check both the context and the helper before using it, as the PyTorch example does:

ctx = flyte.ctx()
assert ctx is not None
cp = ctx.checkpoint
assert cp is not None

TaskContext.checkpoint returns None when checkpoint_paths is absent or when its checkpoint_path is empty or whitespace. When a destination exists, the property lazily creates a flyte.Checkpoint from the runtime-provided destination and previous-attempt source, and caches that instance in the context's data dictionary. Repeated access therefore returns the same helper.

Asynchronous tasks

For an async def task, use the asynchronous pair and await both operations. The async generic example has the same state-management logic as the synchronous version:

from __future__ import annotations

import flyte

env = flyte.TaskEnvironment(
    name="checkpoint_generic_json",
    image=flyte.Image.from_debian_base(),
)

RETRIES = 3


@env.task(retries=RETRIES)
async def use_checkpoint(n_iterations: int) -> int:
    assert n_iterations > RETRIES
    checkpoint = flyte.ctx().checkpoint

    path = await checkpoint.load()
    prev = None if path is None else path.read_bytes()

    start = 0
    if prev:
        start = int(prev.decode())

    failure_interval = n_iterations // RETRIES
    for index in range(start, n_iterations):
        if index > start and index % failure_interval == 0:
            raise RuntimeError(f"Failed at iteration {index}, failure_interval {failure_interval}.")
        await checkpoint.save(f"{index + 1}".encode())
    return index

The path returned by load() is still a local pathlib.Path, so the example reads it synchronously with read_bytes(). Only the checkpoint transfer calls are asynchronous. Use load_sync()/save_sync() in def tasks and await load()/await save() in async def tasks; the checkpoint API documents these as separate task-function-specific method pairs.

Understand the returned path

flyte.Checkpoint stores each checkpoint as one remote object. The value returned by load() or load_sync() depends on what was saved:

If no previous source was supplied, the method returns None.
If the previous remote object is empty, missing, or cannot be downloaded because of a recoverable storage error, the method also treats it as no restored state and returns None.
If the object is a raw file or raw bytes, it is placed at checkpoint.path / "payload", and that payload path is returned.
If the object is a directory checkpoint, it is a gzip-compressed tar archive. The archive is extracted beneath checkpoint.path, and the workspace directory itself is returned.

A None result should therefore be handled as a fresh start. Do not expect a raw checkpoint to retain its original filename.

Saving accepts only a pathlib.Path, a string path, or bytes:

checkpoint.save_sync(b"state")
checkpoint.save_sync("local/state.bin")
checkpoint.save_sync(pathlib.Path("local/state.bin"))

The string and Path forms must refer to an existing path; otherwise Checkpoint.save() and Checkpoint.save_sync() raise FileNotFoundError. A file is uploaded as-is. A directory is packed as a gzip tar containing its top-level entries before upload. Checkpoint does not serialize arbitrary Python objects, so serialize model or application state first with the format appropriate to that state.

The helper's path is a temporary local workspace. It is the place to read restored files and to write files that will be saved as a directory checkpoint. Consume restored paths during the task rather than treating them as durable paths after the task process ends. A Checkpoint clears the workspace before its first previous-state download and caches subsequent loads on the same instance instead of downloading again.

Save a single local file

For a framework state bundle, serialize all state into one local file and upload that file after each unit of progress. examples/checkpoint/pytorch_task_checkpoint.py saves the model, optimizer, and epoch this way:

import pathlib

import torch
from torch import nn

import flyte


@env.task(retries=RETRIES)
async def train_linear(epochs: int = 3) -> float:
    assert epochs > RETRIES
    ctx = flyte.ctx()
    assert ctx is not None
    cp = ctx.checkpoint
    assert cp is not None

    model = nn.Linear(4, 1)
    opt = torch.optim.SGD(model.parameters(), lr=0.01)

    prev_cp_path: pathlib.Path | None = await cp.load()
    if prev_cp_path:
        blob = torch.load(prev_cp_path, map_location="cpu", weights_only=False)
        model.load_state_dict(blob["model"])
        opt.load_state_dict(blob["opt"])
        start = int(blob["epoch"]) + 1
    else:
        start = 0

    wpath = pathlib.Path("pytorch_linear") / "training.pt"
    wpath.parent.mkdir(parents=True, exist_ok=True)

    for epoch in range(start, epochs):
        x = torch.randn(8, 4)
        y = torch.randn(8, 1)
        loss = torch.nn.functional.mse_loss(model(x), y)
        opt.zero_grad()
        loss.backward()
        opt.step()

        torch.save(
            {"model": model.state_dict(), "opt": opt.state_dict(), "epoch": epoch},
            wpath,
        )
        await cp.save(wpath)

    with torch.no_grad():
        return float(model(torch.ones(1, 4)).squeeze().item())

The important ordering is restore, perform work, serialize the updated state, then call cp.save(wpath). The saved file is uploaded as one raw object; on the next attempt, the path returned by load() is the payload file that torch.load() reads.

Save and restore a directory

Use a directory when the framework produces several files or a layout that must be preserved. The Lightning example restores the directory returned by load_sync(), finds the framework's last.ckpt, and mirrors the local checkpoint directory after Lightning writes it:

import pathlib

import flyte
from lightning.pytorch.callbacks import ModelCheckpoint


class FlyteLightningCheckpointCallback(ModelCheckpoint):
    def __init__(self, flyte_checkpoint: flyte.Checkpoint, *, dirpath: str | pathlib.Path, **kwargs) -> None:
        super().__init__(dirpath=str(dirpath), **kwargs)
        self._flyte_checkpoint = flyte_checkpoint

    def on_train_epoch_end(self, trainer, pl_module) -> None:
        super().on_train_epoch_end(trainer, pl_module)
        if self.dirpath:
            self._flyte_checkpoint.save_sync(pathlib.Path(self.dirpath))


@env.task(retries=RETRIES)
def train_lightning(max_epochs: int = 3) -> float:
    assert max_epochs > RETRIES
    ctx = flyte.ctx()
    assert ctx is not None
    checkpoint = ctx.checkpoint
    assert checkpoint is not None

    ckpt_dir = pathlib.Path("pl_checkpoints")
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    prev_cp_path: pathlib.Path | None = checkpoint.load_sync()
    resume_ckpt: str | None = None
    epoch_start = 0
    if prev_cp_path:
        last = flyte.latest_checkpoint(prev_cp_path)
        last_ckpt = str(last) if last else None
        if last_ckpt:
            ck = torch.load(last_ckpt, map_location="cpu", weights_only=False)
            epoch_start = int(ck.get("epoch", 0))
            resume_ckpt = last_ckpt

    mc = FlyteLightningCheckpointCallback(
        checkpoint,
        dirpath=ckpt_dir,
        filename="last",
        save_last=True,
        save_top_k=1,
    )
    # Pass `mc` to the framework's trainer together with the model and data.

Here prev_cp_path is the restored workspace directory, not a file. flyte.latest_checkpoint(prev_cp_path) searches that restored tree for the newest matching last.ckpt; the local callback then calls save_sync() on dirpath, causing the directory to be uploaded as a tarred single object.

Runtime and retry requirements

Task bodies normally do not construct checkpoint URIs themselves. The runtime supplies checkpoint_path and the optional prev_checkpoint_path through CheckpointPaths in TaskContext. The task runner receives these values from the checkpoint output and previous-checkpoint prefixes and passes them into the task context. Checkpoint recovery is useful across attempts only when the task has retries configured and state was saved before the attempt failed; the generic examples use @env.task(retries=RETRIES) and deliberately raise after earlier saves.

BaseCheckpoint is the abstract contract behind the helper: it defines path, prev_exists(), asynchronous load()/save(), and synchronous load_sync()/save_sync(). User task code should normally use the concrete top-level flyte.Checkpoint obtained from TaskContext.checkpoint rather than constructing or subclassing BaseCheckpoint.

Troubleshoot checkpoint behavior

checkpoint is None: The task context has no checkpoint_paths, or its destination is blank. Check the task's checkpoint output configuration and keep the ctx is not None / checkpoint is not None checks when the configuration is optional.
The task always starts at zero: load*() returns None for a missing, empty, or recoverably unavailable previous object. Also verify that the task saved state before the failure and that retries are configured.
A directory's expected filename is missing: Directory checkpoints are extracted under checkpoint.path; raw objects use checkpoint.path / "payload". For framework directories, inspect or scan the returned tree rather than reading the return value as a single file.
A save fails with FileNotFoundError: The supplied string or pathlib.Path does not exist at save time. Create the parent directory and write the framework file before calling save*().
The wrong API is being called from an async or sync task: Match def with load_sync()/save_sync() and async def with await load()/await save().
A previous remote URI contains an attempt suffix: For supported remote schemes, Checkpoint includes a compatibility repair for known previous-attempt URI layouts, using the current action identity and FLYTE_ATTEMPT_NUMBER. Local file: paths are not rewritten; the repair is not a substitute for supplying the runtime's checkpoint paths.

Resume task work with a checkpoint​

Synchronous tasks​

Asynchronous tasks​

Understand the returned path​

Save a single local file​

Save and restore a directory​

Runtime and retry requirements​

Troubleshoot checkpoint behavior​