# 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.

> [!NOTE] Checkpoints vs. caching vs. traces
> - **Task caching** skips an entire task when it has already run with the same inputs.
> - **[Traces](https://www.union.ai/docs/v2/union/user-guide/task-programming/intra-task-checkpoints/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:

### Async

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

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/generic_checkpoint.py*

### Sync

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

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/generic_checkpoint_sync.py*

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.

### PyTorch

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:

<br>

```
@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())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/pytorch_checkpoint.py*

### PyTorch Lightning

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

<br>

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

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/pytorch_lightning_checkpoint.py*

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:

<br>

```
@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())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/pytorch_lightning_checkpoint.py*

### Hugging Face Trainer

`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:

<br>

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

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/huggingface_trainer_checkpoint.py*

On retry, restore the tree, locate the last Hugging Face checkpoint with
`get_last_checkpoint`, and pass it to `trainer.train(resume_from_checkpoint=...)`:

<br>

```
@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())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/huggingface_trainer_checkpoint.py*

### scikit-learn

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:

<br>

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

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/sklearn_partial_checkpoint.py*

### Unsloth

LoRA fine-tuning with [Unsloth](https://unsloth.ai/) 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:

<br>

```
@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())
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-programming/intra-task-checkpoints/unsloth_sft_checkpoint.py*

> [!NOTE] 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.

---
**Source**: https://github.com/unionai/unionai-docs/blob/main/content/user-guide/task-programming/intra-task-checkpoints.md
**HTML**: https://www.union.ai/docs/v2/union/user-guide/task-programming/intra-task-checkpoints/