Checkpoint Data Formats

The flyte.Checkpoint class enables tasks to persist state across retries by saving data to object storage. When a task fails and is retried, it can load the checkpoint from the previous attempt to resume work instead of starting from scratch. flyte-sdk supports three primary data formats for checkpoints: raw bytes, single files, and entire directories.

Accessing the Checkpoint Helper

You do not instantiate the Checkpoint class directly. Instead, access it through the checkpoint property of the current Flyte context. Tasks must be defined within a TaskEnvironment to enable features like retries and checkpointing.

import flyte

env = flyte.TaskEnvironment("checkpoint-example")

@env.task(retries=3)
def my_task():
    checkpoint = flyte.ctx().checkpoint
    # Use checkpoint.load_sync() and checkpoint.save_sync()

Saving and Loading Raw Bytes

For simple state like counters or small serialized objects, you can save bytes directly. This is useful for tracking progress in iterative loops.

In an async task, use the asynchronous save and load methods:

@env.task(retries=3)
async def process_iterations(n: int):
    cp = flyte.ctx().checkpoint
    
    # load() returns a pathlib.Path to the restored data, or None if no checkpoint exists
    path = await cp.load()
    start_index = 0
    if path:
        # For bytes/files, the data is restored to a file named "payload"
        start_index = int(path.read_bytes().decode())

    for i in range(start_index, n):
        # Perform work...
        
        # Save the next starting point as bytes
        await cp.save(str(i + 1).encode())

Internally, flyte.Checkpoint.save detects the bytes type and uploads it. When you call load(), flyte-sdk downloads the object and, because it is not a tarball, moves it to self.path / "payload" (where self.path is a unique temporary directory).

Saving and Loading Files

For larger state, such as a trained model weights file, you can pass a pathlib.Path or a string path to a file.

import torch
import pathlib

@env.task(retries=2)
def train_model():
    cp = flyte.ctx().checkpoint
    model = MyModel()
    
    prev_path = cp.load_sync()
    if prev_path:
        # prev_path points to the "payload" file containing the model weights
        model.load_state_dict(torch.load(prev_path))
    
    # ... training logic ...
    
    save_path = pathlib.Path("model.pt")
    torch.save(model.state_dict(), save_path)
    cp.save_sync(save_path)

When a single file is saved, flyte-sdk uploads it directly to the remote destination. Upon loading, the file is restored to the local temporary workspace as a file named payload. The load() method returns the path to this specific file.

Saving and Loading Directories

If your task produces multiple files or a complex directory structure, you can save the entire directory. flyte-sdk automatically packages the directory into a gzipped tarball before uploading.

@env.task(retries=2)
def complex_task():
    cp = flyte.ctx().checkpoint
    # Use the checkpoint's managed path for local work
    working_dir = cp.path / "outputs"
    working_dir.mkdir(exist_ok=True)
    
    # Restore previous directory state if it exists
    restored_path = cp.load_sync()
    if restored_path:
        # restored_path will be the directory containing the unpacked files
        pass

    # Save the entire directory
    cp.save_sync(working_dir)

Internal Handling of Directories

When you pass a directory path to save() or save_sync():

1. The _save_directory_as_tarball method is triggered.
2. It uses _tar_directory_to_file to create a .tar.gz archive of the directory's immediate children.
3. The archive is uploaded to the remote object store.

When load() or load_sync() is called for a tarball checkpoint:

1. The archive is downloaded to a temporary location.
2. flyte-sdk detects it is a tarfile and extracts the contents directly into flyte.Checkpoint.path.
3. The method returns flyte.Checkpoint.path itself, allowing you to use standard tools like pathlib.Path.rglob to find your files.

Summary of Return Values

The path returned by load() or load_sync() depends on the format of the saved data:

Saved Format	Remote Storage	Restored Local Path	Return Value
Bytes	Single Object	.../flyte-cp-xxx/payload	pathlib.Path to payload
File	Single Object	.../flyte-cp-xxx/payload	pathlib.Path to payload
Directory	.tar.gz Object	.../flyte-cp-xxx/ (unpacked)	pathlib.Path to the directory

Note: The temporary directory flyte-cp-xxx is managed by the Checkpoint instance and is automatically cleaned up when the object is deleted.