Skip to main content

Checkpointing and Retries

Retry lifecycle: attempts and checkpoint state​

When a task can make progress over several iterations, configure retries on the task and save the progress that a later attempt should resume from. In flyte-sdk, retries=3 means the initial attempt plus up to three retry attempts; it is not a limit of three total attempts. TaskEnvironment supplies shared environment settings, but it does not have an environment-wide retries field. Set the value on TaskEnvironment.task or replace it with task.override().

import flyte

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

@env.task(retries=3)
def use_checkpoint(n_iterations: int) -> int:
    checkpoint = flyte.ctx().checkpoint
    path = checkpoint.load_sync()
    previous = None if path is None else path.read_bytes()

    start = 0 if not previous else int(previous.decode())
    for index in range(start, n_iterations):
        checkpoint.save_sync(f"{index + 1}".encode())
    return index

This is the core pattern from examples/checkpoint/generic_data_checkpoint.py: obtain the task context with flyte.ctx(), load once before doing work, interpret the returned path, and save after each unit of progress. If an attempt fails after a save, the next attempt receives the prior checkpoint location through the runtime and can restore that state.

The access path is:

flyte.ctx()
    -> TaskContext.checkpoint
        -> Checkpoint(checkpoint_path, prev_checkpoint_path)
            -> Checkpoint.load()/load_sync()

TaskContext is the immutable (frozen=True) context dataclass in src/flyte/models.py. Its checkpoint property returns None when checkpoint_paths is absent or its destination is blank. Otherwise it constructs flyte.Checkpoint from CheckpointPaths.checkpoint_path and CheckpointPaths.prev_checkpoint_path, then caches that object in the context's mutable data dictionary under the SDK's checkpoint cache key. Repeated access therefore returns the same helper. Task code that can run without checkpoint configuration should branch on None; the framework-oriented examples assert that the helper is present.

Detecting a retry with Checkpoint.prev_exists​

Use prev_exists() when you need to know whether the runtime supplied a previous-attempt source URI:

checkpoint = flyte.ctx().checkpoint
if checkpoint is not None and checkpoint.prev_exists():
    restored = checkpoint.load_sync()
else:
    restored = None

Checkpoint.prev_exists() checks only whether a nonempty previous URI was supplied to the helper. It does not perform an object-store lookup, so True does not guarantee that a nonempty checkpoint can be downloaded. load_sync() can still return None when there is no source, the downloaded object is empty, or the load encounters a recoverable missing-data error. In most task code, calling load_sync() or load() directly and checking its returned path is sufficient; use prev_exists() when the distinction between “a previous URI was configured” and “state was successfully restored” matters.

Checkpoint.load() and load_sync() restore at most once for a given helper. They clear the temporary checkpoint workspace before downloading, ingest the object, and remember the resulting path. A subsequent load call returns the previously determined result without downloading again. TaskContext.checkpoint is independently cached at the context level.

Match the checkpoint API to the task body​

Use the asynchronous methods in an async def task and the synchronous methods in an ordinary def task. The asynchronous example in examples/checkpoint/generic_data_checkpoint_async.py is the same resume loop with the required await calls:

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

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

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

    for index in range(start, n_iterations):
        await checkpoint.save(f"{index + 1}".encode())
    return index

The API pairing is explicit:

  • async def: await checkpoint.load() and await checkpoint.save(data)
  • ordinary def: checkpoint.load_sync() and checkpoint.save_sync(data)

Do not interchange these forms. Both save and save_sync accept a bytes value, a file path, or a directory path. The generic examples use bytes to store the next loop index, so the restored path is read with Path.read_bytes() and decoded before the loop resumes.

Save application state, not only a progress number​

For a stateful task, put the model or estimator state and the progress marker in the saved object. The PyTorch example, examples/checkpoint/pytorch_task_checkpoint.py, loads a serialized model and optimizer bundle, resumes at the following epoch, and uploads the file after each completed epoch:

@env.task(retries=3)
async def train_linear(epochs: int = 3) -> float:
    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):
        # ... train one epoch ...
        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 state update, local serialization, then cp.save(wpath). Local changes made after the last successful save are not available to the retry. The scikit-learn example applies the same idea to a single pickle bundle containing both the classifier and chunks_done, allowing the retry to skip completed chunks.

Files, bytes, and directories​

Checkpoint treats the kinds of data differently:

  • bytes are uploaded as one object.
  • A file is uploaded directly.
  • A directory is converted to one gzip-compressed tar object containing its top-level entries.

After loading a raw file or bytes checkpoint, the payload is available at checkpoint.path / "payload", and load()/load_sync() returns that payload path. After loading a directory checkpoint, the archive is extracted below checkpoint.path, and the load method returns the restored root directory. The workspace is temporary, so consume or copy restored state during the task rather than treating that path as durable storage.

Directory checkpoints are useful when a framework creates several files. In examples/checkpoint/pytorch_lightning_checkpoint.py, the synchronous task restores the directory, asks flyte.latest_checkpoint to select the framework's newest checkpoint, and passes the selected file back 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

# The callback mirrors the local checkpoint directory to Flyte.
mc = FlyteLightningCheckpointCallback(
    checkpoint,
    dirpath=pathlib.Path("pl_checkpoints"),
    filename="last",
    save_last=True,
    save_top_k=1,
)
trainer.fit(model, loader, ckpt_path=resume_ckpt)

The callback and trainer setup shown here come from the same example. latest_checkpoint is needed because a restored directory can contain more than one framework-managed checkpoint; Checkpoint restores the tree, while the framework-specific selector chooses which file to resume.

Retry configuration and serialization​

Use an integer for the common case:

@env.task(retries=3)
def process() -> str:
    return "done"

Use flyte.RetryStrategy when you want the retry policy represented explicitly:

@env.task(retries=flyte.RetryStrategy(count=5))
def process() -> str:
    return "done"

RetryStrategy in src/flyte/_retry.py is a dataclass with one field, count: int. The public top-level import is flyte.RetryStrategy. TaskEnvironment.task accepts either an integer or a RetryStrategy; the generated task template normalizes an integer in TaskTemplate.__post_init__. Later, get_proto_retry_strategy in src/flyte/_internal/runtime/task_serde.py serializes the normalized object as protobuf metadata with retries=retries.count. That serializer expects a RetryStrategy, not a raw integer, so normalization must occur before serialization.

A task override can replace the retry policy as well. The user-API test uses this form and verifies that the serialized task metadata contains count 5:

new_td = td.override(
    short_name="new_oomer",
    resources=flyte.Resources(cpu=3, memory="100Mi"),
    retries=flyte.RetryStrategy(5),
    timeout=100,
    env_vars={"FOO": "BAR"},
    secrets=secrets,
)

RetryStrategy itself performs no additional validation in its implementation; count=0 is the documented no-retry value. The task-level API defaults to 0.

Runtime paths and attempt context​

The runtime supplies CheckpointPaths to TaskContext. Its destination is the current attempt's checkpoint prefix, and its previous path identifies the checkpoint source staged from the preceding attempt. Local retry execution stages the current destination to the previous-checkpoint path before retrying, while remote execution supplies analogous prefixes. TaskContext.checkpoint uses those paths; there is no separate checkpoint option that task code needs to construct manually.

For diagnostics, TaskContext.attempt_number reads FLYTE_ATTEMPT_NUMBER as an integer and defaults to 0 when the variable is unset:

ctx = flyte.ctx()
if ctx is not None:
    print(ctx.attempt_number)

The variable is also used when Checkpoint repairs certain remote previous-checkpoint URIs. That repair applies only to matching remote attempt-directory paths under the runtime's conditions; it is skipped for local file: URIs and cases such as a missing task context or an inconsistent attempt number. Ordinary checkpoint code should use the context-provided helper rather than constructing or editing these URIs.

Checkpoint object I/O delegates to flyte.io.File. Remote schemes therefore use the normal Flyte storage configuration; explicit file:// paths are also supported and are used by checkpoint tests without requiring flyte.init.

Operational boundaries​

  • Save before intentionally raising or before any failure point whose completed work you want to retain. Only the last uploaded checkpoint is available to the next attempt.
  • Treat load() returning None as a valid starting-state branch. prev_exists() only reports a configured source URI, and a source can still be missing or empty.
  • Keep the sync/async pairing consistent, including framework callbacks: Lightning's synchronous task uses load_sync, while the PyTorch task uses await cp.load() and await cp.save(...).
  • Remember that the restored workspace is cleared before each actual restore and cleaned up with the helper's temporary directory. A saved directory is one archive object, not a persistent remote directory tree.
  • A task can stop before consuming all configured retries. examples/basics/non_recoverable_error.py raises flyte.errors.NonRecoverableError with retries=3 and documents that the task fails on its first attempt without consuming the remaining retries.
  • flyte.ctx() is task-scoped. Outside an active task execution it returns None, so checkpoint access belongs in the task body rather than module-level initialization.