Task Checkpointing

When a task fails after making progress, use flyte.ctx().checkpoint to save the completed state and load it at the beginning of the next retry. The smallest synchronous and asynchronous patterns are:

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

For an async def task, use the awaited methods instead:

@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

These examples are from examples/checkpoint/generic_data_checkpoint.py and examples/checkpoint/generic_data_checkpoint_async.py. Each successful iteration saves the next iteration number before continuing. The injected exception therefore causes the next retry to start from the last value that was successfully saved.

Choose the checkpoint API for the task function

flyte.BaseCheckpoint defines the checkpoint contract:

Operation	def task	async def task	Result
Restore prior state	load_sync()	await load()	A restored pathlib.Path, or None
Persist state	save_sync(data)	await save(data)	None
Workspace	path	path	Local directory for restored or locally prepared files
Prior-attempt availability	prev_exists()	prev_exists()	Whether a previous checkpoint source was provided

Use load_sync() and save_sync() in ordinary synchronous task bodies, and load() and save() in asynchronous task bodies. Checkpoint implements both variants. A synchronous task should not replace its _sync calls with the asynchronous methods, and an asynchronous task should not block on the synchronous API.

flyte.ctx().checkpoint can be None when the runtime did not configure checkpoint paths or when its destination path is blank. The executable examples assert that the task context and checkpoint are present when they need them:

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

Save bytes, files, and directories

flyte.Checkpoint represents one remote checkpoint object. Its constructor accepts the destination URI and, optionally, the previous-attempt source URI:

checkpoint = flyte.Checkpoint(
    checkpoint_dest="file:///tmp/checkpoint",
    checkpoint_src="file:///tmp/previous-checkpoint",
)

In normal task code, obtain it from flyte.ctx().checkpoint rather than constructing it directly. Checkpoint.save() and Checkpoint.save_sync() accept bytes, a pathlib.Path, or a string interpreted as a local filesystem path:

• bytes are uploaded directly as one raw object.
• A file is uploaded as one raw object.
• A directory is serialized as a gzip-compressed tar containing its immediate children, then uploaded as one object.
• A missing file or directory raises FileNotFoundError.

A raw object is restored beneath the temporary workspace as checkpoint.path / "payload"; load() and load_sync() return that payload path. A tarred directory is extracted beneath checkpoint.path; those methods return the workspace directory itself.

The PyTorch example uses a local file to package model, optimizer, and loop state before uploading it after each epoch:

import pathlib

import torch
import torch.nn as nn

import flyte

env = flyte.TaskEnvironment(
    name="checkpoint_pytorch",
    image=flyte.Image.from_debian_base().with_pip_packages("torch"),
)

RETRIES = 3


@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())

Here prev_cp_path is the raw payload path because wpath is a file. The example deserializes that path, resumes at the saved epoch plus one, and overwrites the local file before each cp.save(wpath).

Restore and mirror a checkpoint directory

Use a directory when a framework produces several related files. Checkpoint.save_sync(directory) archives the directory, while a later load returns the restored workspace directory—not the original directory name. Search the returned tree for the framework file you need.

flyte.latest_checkpoint selects one file below a root. Its default pattern is recursive **/last.ckpt, ranked by modification time. Supply glob_pattern and key when the framework uses another layout or when a training step is more authoritative than mtime.

The Lightning example restores a tree, selects its newest last.ckpt, and passes that path to Lightning:

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
        print(f"Resuming from epoch {epoch_start} from checkpoint {ck}")

To mirror Lightning's on-disk directory after the framework has written its checkpoint, the example calls super().on_train_epoch_end(...) first and then saves the directory:

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

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

The complete task creates pl_checkpoints, loads it with checkpoint.load_sync(), installs this callback with filename="last" and save_last=True, and calls trainer.fit(..., ckpt_path=resume_ckpt). The same directory semantics apply to other framework integrations: restore the returned directory, locate the framework's actual checkpoint file, and save the framework output directory after its own save operation.

Runtime behavior and retries

Task execution supplies checkpoint locations through CheckpointPaths. The runtime constructs a destination checkpoint location and may provide a previous-attempt location. TaskContext.checkpoint then:

1. Returns None if checkpoint_paths is absent or checkpoint_path is empty.
2. Strips surrounding whitespace from the configured destination and previous source.
3. Lazily constructs a flyte.Checkpoint.
4. Caches that object in TaskContext.data under the checkpoint cache key, so repeated access returns the same helper.

CheckpointPaths only stores the platform-provided URI strings; it does not perform downloads or uploads. I/O is delegated to flyte.io.File by Checkpoint.

Checkpoint recovery requires a later task attempt. Configure task retries, as the examples do with @env.task(retries=RETRIES). For local retries, the local controller stages the current checkpoint object at the previous-checkpoint location before the retry. Remote execution supplies the corresponding previous-attempt URI. This lets the same Checkpoint.load() or load_sync() code run in both modes.

Troubleshooting checkpoint restores

load() returns None

load() and load_sync() return None on the first attempt when no previous source URI was supplied. They also treat an empty or missing previous object, and recognized recoverable storage failures, as no restored state. Recoverable failures include AssertionError, FileNotFoundError, OSError, and DownloadQueueEmpty; exception groups made entirely from those failures are also recoverable. Other exceptions are raised.

Handle None as a fresh start, as the generic and PyTorch examples do. A checkpoint only contains the state from the last successful save; saving after a failure check means the failed unit is retried.

The restored path is not the expected filename

For raw files and bytes, use the path returned by load() or load_sync(); it is checkpoint.path / "payload". For directory checkpoints, use the returned directory and search it with pathlib.Path.rglob, flyte.latest_checkpoint, or a framework utility. Do not assume that the original local directory name is returned.

A second load behaves differently from the first

A Checkpoint downloads at most once. The first load clears its temporary workspace before fetching and then caches the restored result. Files placed in checkpoint.path before that first load are discarded. The workspace is temporary and is cleaned up when the Checkpoint is destroyed, so consume or copy restored data during the task.

A remote previous URI appears to point at the wrong attempt

For supported remote schemes (s3://, gs://, gcs://, s3a://, abfss://, and az://), Checkpoint applies repair_union_prev_checkpoint_uri to the previous URI when applicable. The helper can decrement a matching attempt-directory segment using the run/action metadata; FLYTE_ATTEMPT_NUMBER provides a conservative check when set. Local file: paths are not modified. This is a compatibility workaround for the Union executor's previous-checkpoint attempt URI and is documented in the source for removal when the backend behavior is fixed.

Public checkpoint API

flyte.__init__ re-exports the supported checkpoint symbols, so task code can use the public flyte namespace:

Symbol	Use
flyte.BaseCheckpoint	Abstract interface for async and sync load/save helpers, workspace access, and prior-state detection.
flyte.Checkpoint	Concrete helper backed by flyte.io.File; handles raw objects and gzip-tarred directories.
flyte.latest_checkpoint	Selects a matching checkpoint file under a restored directory, by default the newest **/last.ckpt by mtime.
flyte.ctx().checkpoint	Runtime-created, lazily cached Checkpoint for the current task's checkpoint paths.