Intra-task checkpoints

Long-running tasks — model training especially — can fail partway through: a spot instance is reclaimed, a pod is evicted, an out-of-memory error kills the process. When the task is retried, it normally starts over from the beginning.

Intra-task checkpoints let a task save in-progress state to object storage as it runs and load that state at the start of the next attempt, so a retry resumes from where the previous attempt left off instead of repeating completed work.

The checkpoint object

Inside a running task, flyte.ctx().checkpoint returns a flyte.Checkpoint (or None when checkpointing isn’t configured) bound to the action’s checkpoint location in object storage:

  • Save: await checkpoint.save(...) (async tasks) or checkpoint.save_sync(...) (sync tasks and synchronous framework callbacks). Accepts raw bytes, a file path, or a directory path — a directory is stored as a single compressed archive.
  • Load: await checkpoint.load() or checkpoint.load_sync(). Returns a local pathlib.Path to the restored file or directory tree, or None when there is no previous checkpoint (i.e. on the first attempt).
  • flyte.latest_checkpoint(root, glob_pattern="**/last.ckpt") finds the newest checkpoint file under a restored directory tree — useful for frameworks like PyTorch Lightning that write last.ckpt files into a directory.

Checkpoints only matter when the task can run more than once, so give the task retries with @env.task(retries=...). Each retry attempt sees the checkpoint saved by the attempt before it.

Checkpoints vs. caching vs. traces
  • Task caching skips an entire task when it has already run with the same inputs.
  • Traces checkpoint at the boundaries of helper functions called by a task.
  • Intra-task checkpoints save state within a single task body — mid-loop, mid-epoch — across retry attempts of the same action.

Basic usage

The simplest checkpoint is a raw byte payload. This task counts up to n_iterations, saving its progress on every iteration. A simulated failure kills it partway through; the retry loads the saved counter and continues rather than restarting from zero:

AsyncSync
generic_checkpoint.py 
import flyte

env = flyte.TaskEnvironment(name="checkpoint_generic")

RETRIES = 3

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

    # Load the previous attempt's checkpoint, if any.
    # On the first attempt there is none, so load() returns None.
    path = await checkpoint.load()
    start = int(path.read_bytes()) if path else 0

    failure_interval = n_iterations // RETRIES
    index = start
    for index in range(start, n_iterations):
        if index > start and index % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at iteration {index}")
        # Persist progress to object storage.
        await checkpoint.save(f"{index + 1}".encode())
    return index
generic_checkpoint_sync.py 
import flyte

env = flyte.TaskEnvironment(name="checkpoint_generic_sync")

RETRIES = 3

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

    # Load the previous attempt's checkpoint, if any.
    # On the first attempt there is none, so load_sync() returns None.
    path = checkpoint.load_sync()
    start = int(path.read_bytes()) if path else 0

    failure_interval = n_iterations // RETRIES
    index = start
    for index in range(start, n_iterations):
        if index > start and index % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at iteration {index}")
        # Persist progress to object storage.
        checkpoint.save_sync(f"{index + 1}".encode())
    return index

Running this with n_iterations=10 produces three failed attempts and one successful one. Each attempt fails later than the last, because each one starts from the checkpoint its predecessor saved.

Checkpointing ML training frameworks

The same pattern applies to real training loops, whatever the framework:

  1. Load the previous attempt’s checkpoint at the start of the task; if one exists, restore the model/optimizer state and work out where to resume.
  2. Save a checkpoint at a regular interval — every epoch or every N steps — as training progresses.

Frameworks with their own checkpoint files (PyTorch Lightning, Hugging Face Trainer) already write them to a local directory; there you hook their callback system and mirror that directory to the Flyte checkpoint, then feed the restored directory back to the framework’s native resume mechanism.

PyTorchPyTorch LightningHugging Face Trainerscikit-learnUnsloth
Save the model state dict, optimizer state, and epoch counter with torch.save after each epoch, and restore all three with torch.load on retry:

pytorch_checkpoint.py 
@env.task(retries=RETRIES)
async def train_linear(epochs: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    model = nn.Linear(4, 1)
    opt = torch.optim.SGD(model.parameters(), lr=0.01)

    # Resume model, optimizer, and epoch from the previous attempt, if any.
    prev = await checkpoint.load()
    if prev:
        blob = torch.load(prev, 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)

    failure_interval = epochs // RETRIES
    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()

        if epoch > start and epoch % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at epoch {epoch}")

        # Save model, optimizer, and epoch state to object storage.
        torch.save(
            {"model": model.state_dict(), "opt": opt.state_dict(), "epoch": epoch},
            wpath,
        )
        await checkpoint.save(wpath)

    with torch.no_grad():
        return float(model(torch.ones(1, 4)).squeeze().item())
Lightning already writes last.ckpt through its ModelCheckpoint callback. Subclass it to mirror the checkpoint directory to Flyte after each epoch (Lightning callbacks are synchronous, so use flyte.Checkpoint.save_sync):

pytorch_lightning_checkpoint.py 
class FlyteLightningCheckpointCallback(ModelCheckpoint):
    """A `ModelCheckpoint` that mirrors `dirpath` to the Flyte checkpoint after each epoch."""

    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:
            # Lightning callbacks are synchronous, so use save_sync
            self._flyte_checkpoint.save_sync(pathlib.Path(self.dirpath))
In the task, restore the previous tree, pick the newest last.ckpt with flyte.latest_checkpoint(restored_root, glob_pattern="**/last.ckpt"), and hand it to Trainer.fit(ckpt_path=...) — Lightning restores the model, optimizer, and epoch from there:

pytorch_lightning_checkpoint.py 
@env.task(retries=RETRIES)
def train_lightning(max_epochs: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    ckpt_dir = pathlib.Path("pl_checkpoints")
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    # Restore the previous attempt's checkpoint tree and find the newest last.ckpt.
    resume_ckpt = None
    prev = checkpoint.load_sync()
    if prev:
        last = flyte.latest_checkpoint(prev)
        if last:
            resume_ckpt = str(last)

    model = TinyModule()
    mc = FlyteLightningCheckpointCallback(
        checkpoint,
        dirpath=ckpt_dir,
        filename="last",
        save_last=True,
        save_top_k=1,
    )
    trainer = L.Trainer(
        max_epochs=max_epochs,
        enable_checkpointing=True,
        callbacks=[mc],
        enable_progress_bar=True,
        logger=False,
        accelerator="cpu",
        devices=1,
    )
    trainer.fit(model, make_loader(), ckpt_path=resume_ckpt)

    with torch.no_grad():
        return float(model(torch.ones(1, FEATURES)).squeeze().item())
transformers.Trainer writes checkpoint-<step> directories under its output_dir (here, every epoch via save_strategy="epoch"). A TrainerCallback mirrors that directory to Flyte after each save:

huggingface_trainer_checkpoint.py 
class FlyteTrainerCheckpointCallback(TrainerCallback):
    """Mirror the Trainer's `output_dir` to the Flyte checkpoint after each epoch."""

    def __init__(self, checkpoint: flyte.Checkpoint, output_dir: pathlib.Path) -> None:
        self._checkpoint = checkpoint
        self._output_dir = output_dir

    def on_epoch_end(self, args, state, control, **kwargs) -> None:
        # Trainer callbacks are synchronous, so use save_sync
        self._checkpoint.save_sync(self._output_dir)
On retry, restore the tree, locate the last Hugging Face checkpoint with get_last_checkpoint, and pass it to trainer.train(resume_from_checkpoint=...):

huggingface_trainer_checkpoint.py 
@env.task(retries=RETRIES)
def train_transformers(max_epochs: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    ckpt_dir = pathlib.Path("hf_trainer")
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    # Restore the previous attempt's checkpoint tree and find the last HF checkpoint.
    hf_resume = None
    prev = checkpoint.load_sync()
    if prev:
        hf_resume = get_last_checkpoint(str(prev))

    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID, num_labels=2)

    args = TrainingArguments(
        output_dir=str(ckpt_dir),
        num_train_epochs=max_epochs,
        per_device_train_batch_size=4,
        save_strategy="epoch",
        save_total_limit=2,
        logging_steps=1,
        report_to="none",
        seed=42,
        use_cpu=True,
    )

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=ToyTextDataset(tokenizer),
        data_collator=DataCollatorWithPadding(tokenizer),
        callbacks=[FlyteTrainerCheckpointCallback(checkpoint, ckpt_dir)],
    )
    trainer.train(resume_from_checkpoint=hf_resume)

    model.eval()
    device = next(model.parameters()).device
    with torch.no_grad():
        batch = tokenizer(
            "classification example for inference",
            return_tensors="pt",
            truncation=True,
            max_length=32,
            padding="max_length",
        )
        batch = {k: v.to(device) for k, v in batch.items()}
        logits = model(**batch).logits
        return float(logits[0, 1].item())
For estimators that support incremental training with partial_fit, pickle the estimator together with a progress counter after each training chunk. A retry unpickles the bundle and continues from the next chunk:

sklearn_partial_checkpoint.py 
@env.task(retries=RETRIES)
async def incremental_sgd(chunks: int = 10) -> float:
    checkpoint = flyte.ctx().checkpoint

    # Resume the estimator and progress from the previous attempt, if any.
    prev = await checkpoint.load()
    if prev:
        bundle = pickle.loads(prev.read_bytes())
        start = bundle["chunks_done"]
        clf = bundle["clf"]
    else:
        start = 0
        clf = SGDClassifier(max_iter=1, tol=None, random_state=0)

    bundle_path = pathlib.Path("sklearn_partial") / "sgd_bundle.pkl"
    bundle_path.parent.mkdir(parents=True, exist_ok=True)

    rng = np.random.default_rng(0)
    classes = np.array([0, 1])

    failure_interval = chunks // RETRIES
    for i in range(start, chunks):
        x = rng.standard_normal((32, 8))
        y = (x[:, 0] + x[:, 1] > 0).astype(int)
        clf.partial_fit(x, y, classes=classes)

        if i > start and i % failure_interval == 0:
            # Simulate a failure so the next attempt resumes from the checkpoint
            raise RuntimeError(f"Simulated failure at chunk {i}")

        # Pickle the estimator plus progress and save it to object storage.
        bundle_path.write_bytes(pickle.dumps({"clf": clf, "chunks_done": i + 1}))
        await checkpoint.save(bundle_path)

    x_test = rng.standard_normal((64, 8))
    y_test = (x_test[:, 0] + x_test[:, 1] > 0).astype(int)
    return float(clf.score(x_test, y_test))
LoRA fine-tuning with Unsloth and trl.SFTTrainer uses the same callback-and-resume pattern as the Hugging Face Trainer, since SFTTrainer is built on it. Unsloth requires an NVIDIA, AMD, or Intel GPU, so the task environment requests one:

unsloth_sft_checkpoint.py 
@env.task(retries=RETRIES)
def train_unsloth_sft(max_epochs: int = 10) -> float:
    from trl import SFTConfig, SFTTrainer
    from unsloth import FastLanguageModel

    checkpoint = flyte.ctx().checkpoint

    ckpt_dir = pathlib.Path("unsloth_sft")
    ckpt_dir.mkdir(parents=True, exist_ok=True)

    # Restore the previous attempt's checkpoint tree and find the last HF checkpoint.
    hf_resume = None
    prev = checkpoint.load_sync()
    if prev:
        hf_resume = get_last_checkpoint(str(prev))

    max_seq_length = 512
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_NAME,
        max_seq_length=max_seq_length,
        dtype=None,
        load_in_4bit=True,
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=16,
        target_modules=[
            "q_proj",
            "k_proj",
            "v_proj",
            "o_proj",
            "gate_proj",
            "up_proj",
            "down_proj",
        ],
        lora_alpha=16,
        lora_dropout=0.0,
        bias="none",
        use_gradient_checkpointing="unsloth",
        random_state=42,
    )

    args = SFTConfig(
        output_dir=str(ckpt_dir),
        num_train_epochs=max_epochs,
        per_device_train_batch_size=1,
        gradient_accumulation_steps=1,
        save_strategy="epoch",
        save_total_limit=2,
        logging_steps=1,
        report_to="none",
        seed=42,
        dataset_text_field="text",
        max_length=max_seq_length,
    )

    trainer = SFTTrainer(
        model=model,
        args=args,
        train_dataset=tiny_instruction_dataset(),
        processing_class=tokenizer,
        callbacks=[FlyteTrainerCheckpointCallback(checkpoint, ckpt_dir)],
    )
    trainer.train(resume_from_checkpoint=hf_resume)

    model.eval()
    device = next(model.parameters()).device
    with torch.no_grad():
        batch = tokenizer(
            "classification example for inference",
            return_tensors="pt",
            truncation=True,
            max_length=32,
            padding="max_length",
        )
        batch = {k: v.to(device) for k, v in batch.items()}
        logits = model(**batch).logits
        return float(logits[0, 1].mean().item())
Simulated failures in the runnable examples

The full example files for the basic, PyTorch, and scikit-learn cases inject a failure at a regular interval (failure_interval) so you can watch the retries resume from the checkpoint. In production code you would drop those lines — real failures (preemptions, OOMs, crashes) trigger the same resume path.

How checkpoints are stored

Each action attempt gets a checkpoint prefix in the object store configured for your cluster. flyte.Checkpoint.save uploads a file as-is, stores a directory as a gzip-compressed tarball, and accepts raw bytes as a single blob. flyte.Checkpoint.load downloads the previous attempt’s object into a local temporary workspace and returns the path — a restored directory tree, or the path to the single restored file.

Saving repeatedly overwrites the same checkpoint object, so the cost of frequent checkpointing is upload bandwidth, not unbounded storage growth. Checkpoint how often you can afford to lose work: every epoch is typical for training loops.