Introduction to Task Checkpointing

Checkpoint lifecycle: from retry paths to task code

A retry starts with two object locations: one for the checkpoint produced by the current attempt and one for the checkpoint supplied by the previous attempt. Flyte’s runtime packages those locations in CheckpointPaths, places them on TaskContext, and task code reaches the helper through flyte.ctx().checkpoint:

checkpoint = flyte.ctx().checkpoint

CheckpointPaths is only the value object carrying prev_checkpoint_path and checkpoint_path; flyte.Checkpoint performs the upload and download work. In the remote runner (src/flyte/_run.py), the paths are derived beneath the run’s raw-data path:

raw_data_path = f"{output_path}/rd/{random_id}"
checkpoint_path = f"{raw_data_path}/checkpoint"
prev_checkpoint = f"{raw_data_path}/prev_checkpoint"
checkpoint_paths = CheckpointPaths(prev_checkpoint_path=prev_checkpoint, checkpoint_path=checkpoint_path)

The runtime command in src/flyte/_bin/runtime.py similarly receives checkpoint_path and prev_checkpoint and forwards them as CheckpointPaths to load_and_run_task. For remote execution, configure a run_base_dir; the run setup uses it to derive these paths.

TaskContext.checkpoint is lazy. It returns None when checkpoint_paths is absent, when its destination is None, or when the destination contains only whitespace. Otherwise it strips the destination and previous path, constructs flyte.Checkpoint, and stores that object in TaskContext.data under flyte-sdk’s checkpoint cache key. Repeated access returns the same helper, so use the property once and retain the result when convenient:

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

TaskContext is a frozen, keyword-only dataclass, but its data dictionary is the mutable auxiliary store used for this cached helper. Its checkpoint_paths field is optional because not every context has checkpoint output configured.

The checkpoint API

BaseCheckpoint in src/flyte/_checkpoint.py defines the contract for checkpoint helpers:

• path is the local directory in which task code reads and writes checkpoint files.
• prev_exists() reports whether a previous-checkpoint source was supplied.
• load() and save() are the asynchronous pair.
• load_sync() and save_sync() are the synchronous pair.
• Save data may be a pathlib.Path, a string path, or raw bytes.

Checkpoint is flyte-sdk’s concrete implementation and is publicly available as flyte.Checkpoint. Construct it directly when you need an explicitly configured helper:

import flyte

checkpoint = flyte.Checkpoint(
    checkpoint_dest="s3://bucket/run/checkpoint",
    checkpoint_src="s3://bucket/run/prev_checkpoint",
)

The destination is one remote object, not a directory prefix containing independently uploaded files. remote_destination, remote_source, and prev_exists() expose the configured locations and whether a source URI was provided. The helper uses flyte.io.File for its blob I/O.

Choose the API pair that matches the task body. Async task code awaits load and save; ordinary def task code calls load_sync and save_sync. flyte-sdk intentionally provides these separate methods rather than converting synchronous calls with flyte.syncify.

Raw bytes for small or custom state

The generic synchronous example saves the next iteration as bytes after each iteration and starts from the value loaded on a retry:

@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:
            # Simulate a failure at a regular interval
            raise RuntimeError(f"Failed at iteration {index}, failure_interval {failure_interval}.")
        checkpoint.save_sync(f"{index + 1}".encode())
    return index

This is the complete pattern: load before the work loop, interpret the returned state in the task’s format, and save progress at the points from which a later attempt should resume. The corresponding asynchronous task uses the same state format but awaits both operations:

@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:
            # Simulate a failure at a regular interval
            raise RuntimeError(f"Failed at iteration {index}, failure_interval {failure_interval}.")
        await checkpoint.save(f"{index + 1}".encode())
    return index

For a raw file or byte object, a successful load places the object at checkpoint.path / "payload" and returns that path. The first attempt has no previous URI, so load_sync() or load() returns None. A task must therefore handle both None and the returned path rather than assuming a checkpoint exists.

Files and model state

Saving a pathlib.Path uploads a file as-is. The file can contain application-defined state, not merely model parameters. The PyTorch example stores the model state, optimizer state, and epoch together, then restores all three before continuing:

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)

The choice of serialized fields belongs to the task or framework. Checkpoint stores the resulting file as a single remote object; it does not inspect or merge the contents.

Directories as gzip tarballs

When the save input is a directory, Checkpoint creates a temporary .tar.gz containing the directory’s immediate children, recursively preserving their contents, and uploads that archive as the single destination object. On load, the archive is extracted under checkpoint.path, and the load method returns checkpoint.path so the task can scan the restored tree.

That representation fits training frameworks whose checkpoint is a directory containing several files. The PyTorch Lightning example restores the directory, locates the framework’s actual checkpoint with flyte.latest_checkpoint, and mirrors the local directory back after each framework checkpoint cycle:

class FlyteLightningCheckpointCallback(ModelCheckpoint):
    """
    `lightning.pytorch.callbacks.ModelCheckpoint` that mirrors `dirpath` to Flyte after each
    on-disk checkpoint cycle in `lightning.pytorch.callbacks.ModelCheckpoint.on_train_epoch_end`.
    """

    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 task side of that integration distinguishes a restored directory from a raw payload by using the returned path:

prev_cp_path: pathlib.Path | None = checkpoint.load_sync()
resume_ckpt: str | None = None
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

Hugging Face Trainer uses the same directory contract, but discovers its state with get_last_checkpoint and passes save_sync through a callback. Framework callbacks are commonly synchronous even when the surrounding task is asynchronous; use save_sync from such a callback unless the callback itself is async.

How restoration behaves

Checkpoint creates a temporary workspace with the prefix flyte-cp-. Its path property ensures that directory exists. The first load clears the workspace, downloads the previous object through flyte.io.File, and either extracts a tarball or moves a non-tar object to path / "payload". The temporary workspace is cleaned when the Checkpoint object is destroyed.

Loads are idempotent for a given helper. After the first attempt, load() or load_sync() returns the cached restored path instead of downloading again. If no previous source was supplied, the remote object is empty, or a recoverable missing/download error occurs, the result is None. Recoverable cases include AssertionError, FileNotFoundError, OSError, DownloadQueueEmpty, and exception groups made up only of recoverable errors. KeyboardInterrupt, SystemExit, and unrelated errors are not converted into “no checkpoint.”

These behaviors are exercised in tests/flyte/test_checkpoint.py. The tests verify directory tar round trips, raw-object round trips and the payload convention, empty remote downloads returning None, and cached TaskContext.checkpoint identity:

cp = Checkpoint(out_blob.as_uri(), prev_blob.as_uri())
payload_path = cp.load_sync()
assert payload_path == cp.path / _PAYLOAD_BASENAME
assert payload_path.read_bytes() == b"hello-checkpoint"

cp.save_sync(b"next")

They also verify that a context without checkpoint paths, or with a whitespace-only destination, exposes no helper:

assert tctx.checkpoint is None

Local and remote retry integration

For a remote run, the runtime receives the current checkpoint_path and prev_checkpoint from the execution environment and forwards them into TaskContext. For local retries, LocalController mirrors that arrangement in _stage_prev_checkpoint_for_local_retry: before retrying, it copies the current local checkpoint object to the previous-checkpoint location when the current destination is a file. The next task attempt then sees the copied object through the same TaskContext.checkpoint property.

The storage layer determines how configured URIs are handled. Tests use file:// URIs generated from temporary pathlib.Path objects; remote runs require the corresponding Flyte storage configuration and backend. Checkpoint recognizes and repairs previous-attempt paths for the remote schemes s3://, gs://, gcs://, s3a://, abfss://, and az://. The repair is a compatibility workaround for a Union executor attempt-number issue: when the embedded attempt segment matches the environment’s FLYTE_ATTEMPT_NUMBER (or that value plus one), flyte-sdk may decrement it. Local file: URIs are not modified, and mismatched segments are left unchanged.

You can inspect the current attempt independently through TaskContext.attempt_number; it reads FLYTE_ATTEMPT_NUMBER and defaults to 0 when unset. The focused checkpoint tests cover this URI-repair behavior alongside the load, save, archive, raw-payload, empty-download, and context-cache cases.

Practical constraints

• Treat the remote destination as one object. A directory save is one gzip tarball, not a collection of separately addressable remote files.
• Branch on the returned path’s meaning: a raw object returns path / "payload"; a tarball returns path.
• Use load/save in async def tasks and load_sync/save_sync in def tasks.
• Save state before a possible failure if the task should resume from that point; the checkpoint records only what the task has actually uploaded.
• Keep required results in the remote checkpoint rather than relying on the temporary workspace after task completion.
• Configure run_base_dir for remote run contexts and configure Flyte storage for remote URIs; explicit file:// paths are suitable for local checkpoint tests.