Implementing a Custom Checkpoint Handler

When the built-in flyte.Checkpoint format or storage behavior is not suitable, subclass flyte.BaseCheckpoint—but implement all six abstract members and integrate the subclass explicitly, because TaskContext.checkpoint currently constructs Checkpoint, not custom subclasses.

import flyte


class CustomCheckpoint(flyte.BaseCheckpoint):
    def __init__(self, checkpoint_dest: str, checkpoint_src: str | None = None):
        self._checkpoint_dest = checkpoint_dest
        self._checkpoint_src = checkpoint_src
        # Create or configure the local workspace used by this handler.
        self._workspace = pathlib.Path(tempfile.mkdtemp(prefix="custom-cp-"))

    @property
    def path(self) -> pathlib.Path:
        self._workspace.mkdir(parents=True, exist_ok=True)
        return self._workspace

    def prev_exists(self) -> bool:
        return self._checkpoint_src is not None

    async def load(self) -> pathlib.Path | None:
        return self.load_sync()

    async def save(self, data: pathlib.Path | str | bytes) -> None:
        self.save_sync(data)

    def load_sync(self) -> pathlib.Path | None:
        if not self.prev_exists():
            return None
        # Retrieve self._checkpoint_src, restore it under self.path, and return
        # the file or directory that task code should consume.
        raise NotImplementedError

    def save_sync(self, data: pathlib.Path | str | bytes) -> None:
        # Encode data and persist it at self._checkpoint_dest.
        raise NotImplementedError

The snippet is a hypothetical extension pattern, not an existing subclass in the flyte-sdk repository. It shows the exact public base-class contract from src/flyte/_checkpoint.py; the repository contains no custom BaseCheckpoint implementation. A real handler replaces the NotImplementedError branches with its backend operations and should configure its destination, optional previous source, workspace, credentials, and serialization policy in __init__.

Implement the BaseCheckpoint contract

flyte.BaseCheckpoint is re-exported by src/flyte/__init__.py. It is an abstract class with these required members:

Member	Contract	Task code
path	Property returning a pathlib.Path for the local checkpoint workspace	Read and write the handler's format here
prev_exists()	Returns bool indicating whether a previous-checkpoint prefix was supplied	Distinguish a first attempt from retry/resume
load()	async def, returns Optional[pathlib.Path]	Use from async def tasks
save(data)	async def, accepts pathlib.Path  str  bytes, returns None	Use from async def tasks
load_sync()	Synchronous load returning Optional[pathlib.Path]	Use from ordinary def tasks and synchronous callbacks
save_sync(data)	Synchronous save accepting pathlib.Path  str  bytes, returning None	Use from ordinary def tasks and synchronous callbacks

The source annotations use the union operator (pathlib.Path | str | bytes) and Optional[pathlib.Path]; the table renders the same accepted values and return behavior. path is format-neutral: it is a directory for the subclass's local files, and the base class does not create or manage it. The subclass therefore needs to establish the directory before returning it or before restoring data into it.

Handle a first attempt and a retry

A handler should return None from both load methods when prev_exists() is false or when its backend reports that the configured previous object is absent. Callers must test the result before using it. The repository's synchronous byte-checkpoint example follows that pattern:

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

The async form is the same lifecycle with the async members:

@env.task(retries=RETRIES)
async def use_checkpoint_async(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 first example is taken from examples/checkpoint/generic_data_checkpoint.py; the second follows the async example in examples/checkpoint/generic_data_checkpoint_async.py. The saved bytes contain the progress marker, so a later attempt can decode it and choose its starting iteration. Save at a safe point after updating the state that the task needs to resume.

Restore and return the path callers need

load and load_sync should restore the previous representation into path and return the particular path task code should consume. That return value does not have to equal path: a handler can return a restored file beneath the workspace, or the workspace itself when it restores a directory tree.

The in-tree Checkpoint provides the reference behavior. It uses flyte.io.File for object I/O and treats a remote checkpoint as one object: directories are gzip-compressed tar archives, while files and bytes are uploaded directly. A non-tar object is restored as path / "payload"; a tar archive is unpacked under path, and the directory path is returned. A custom handler may choose another format, but its load operation must reverse its save operation and should document the returned path.

A file-oriented task can consume the returned path directly. The PyTorch example restores model and optimizer state from the path returned by load, then saves a local file after each epoch:

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):
    # training work omitted, as in examples/checkpoint/pytorch_task_checkpoint.py
    torch.save(
        {"model": model.state_dict(), "opt": opt.state_dict(), "epoch": epoch},
        wpath,
    )
    await cp.save(wpath)

For directory payloads, return the restored directory and preserve the directory structure expected by the consumer. The synchronous Lightning adapter demonstrates the other side of this contract: it delegates a local checkpoint directory to save_sync after the framework writes it.

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

This adapter is from examples/checkpoint/pytorch_lightning_checkpoint.py and uses the concrete flyte.Checkpoint; a custom handler can expose the same synchronous callback shape if it supports the directory format used by the framework.

Connect a custom handler to runtime paths

The runtime provides checkpoint locations through CheckpointPaths. TaskContext.checkpoint in src/flyte/models.py performs this sequence:

1. Return None when checkpoint_paths is absent or checkpoint_path is blank.
2. Read checkpoint_path as the destination and prev_checkpoint_path as the optional previous source.
3. Lazily construct Checkpoint(destination, previous_source).
4. Cache that concrete helper in TaskContext.data and return it on later accesses.

Consequently, subclassing BaseCheckpoint alone does not register or activate the subclass. The current factory imports Checkpoint as CheckpointCls and instantiates that class. Selecting a custom handler requires an explicit runtime or factory change, a wrapper around the task context, or another task-context construction path that supplies the custom instance. When changing that integration, preserve the destination/previous-source distinction: the destination is where save writes, while the optional previous source is what load reads for retry or resume.

The public import is available at the package root:

import flyte

Base = flyte.BaseCheckpoint
BuiltIn = flyte.Checkpoint

flyte.Checkpoint is the only concrete BaseCheckpoint subclass in the repository and is the standard runtime implementation. Its storage access follows Flyte's configured flyte.io.File and storage stack. A custom backend should separately document its bucket, endpoint, credentials, or client initialization; no backend-specific configuration is defined by BaseCheckpoint.

Operational checklist

• Implement path, prev_exists, load, save, load_sync, and save_sync. Leaving either async or sync pairs abstract keeps the subclass uninstantiable.
• Keep async task code on load/save and ordinary synchronous tasks or callbacks on load_sync/save_sync; the repository examples do not mix the pairs.
• Define what pathlib.Path, str, and bytes mean for your backend. The abstract signature accepts all three, but serialization is the subclass's responsibility.
• Make save and load representations reversible, and state whether a returned path is a payload file or a restored directory.
• Decide how to handle a missing previous object, empty data, and backend download errors. The built-in Checkpoint treats several missing/empty download conditions as no prior checkpoint; that behavior is not imposed by the abstract base class.
• Clear or isolate stale workspace contents before restoring a checkpoint, and decide whether repeated loads should be cached. The built-in handler clears its workspace and caches a successful restore; these are concrete Checkpoint behaviors, not BaseCheckpoint requirements.
• Test first-attempt and retry paths, bytes, files, directories, repeated loads, missing sources, and backend errors.
• If using remote URI repair or attempt-specific paths, treat it as backend-specific behavior. The built-in Checkpoint can repair certain remote previous-checkpoint URIs and consults FLYTE_ATTEMPT_NUMBER; the abstract base class has no such requirement.