Recover from Pipeline Failures with Retries and Self-Healing Tasks

AI engineering tip of the week: Recover from Pipeline Failures with Retries and Self-Healing Tasks

ML pipelines fail. APIs time out. Spot instances get preempted. GPUs run out of memory. The question isn't whether your pipeline will eventually fail, it's how gracefully it recovers.

Flyte 2 handles transient failures with automatic retries, and can turn others into Python exceptions you can catch. Once a failure is a Python exception, the code that hit it can also fix it. Together they cover everything from a flaky API call to self-healing agents that catch their own OOM and ask for more resources.

Automatic retries

Add `retries=N` to your task and Flyte automatically re-runs it up to N times if it fails:

import flyte

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

@env.task(retries=3)
async def call_external_api(url: str) -> str:
    # If this fails, Flyte retries up to 3 more times (4 total attempts)
    response = await fetch(url)
    return response

No retry loops, no exponential backoff boilerplate. Flyte handles it at the infrastructure level. Each retry gets a fresh container.

Catch errors from other tasks with try/except

In Flyte 2, task errors are just Python exceptions. Catch them with standard `try`/`except`:

import flyte.errors

@env.task
async def risky_task(x: int) -> int:
    if x < 0:
        raise ValueError(f"Invalid input: {x}")
    return x * 2

@env.task
async def safe_pipeline(x: int) -> int:
    try:
        result = await risky_task(x)
    except flyte.errors.RuntimeUserError as e:
        print(f"Task failed with: {e.code} - recovering...")
        result = await risky_task(abs(x))  # retry with fixed input
    return result

The `e.code` field contains the exception class name (like `"ValueError"`), so you can handle different error types differently.

Handle out-of-memory errors

OOM errors are common in ML. Flyte gives you a specific exception for them:

@env.task
async def train(data: str) -> float:
    # might OOM on large datasets
    return model

@env.task
async def resilient_training(data: str) -> float:
    try:
        return await train(data)
    except flyte.errors.OOMError:
        print("OOM! Retrying with more memory...")
        return await train.override(
            resources=flyte.Resources(memory="16Gi")
        )(data)

Catch the OOM, bump the resources with `.override()`, and try again. No manual intervention needed.

This is what makes self-healing pipelines and agents possible: your infrastructure becomes part of the context your code can reason about. The failure isn't a log line someone reads the next morning, it's a typed Python exception in the same function that can fix it. An agent that hits an OOM can request more memory. One that hits a timeout can swap to a smaller model or a bigger GPU. One that gets a `NonRecoverableError` knows not to waste attempts and can escalate instead.

@env.task
async def self_healing_inference(prompt: str) -> str:
    for memory, gpu in [("8Gi", "L4:1"), ("32Gi", "A100:1"), ("80Gi", "H100:1")]:
        try:
            return await run_model.override(
                resources=flyte.Resources(memory=memory, gpu=gpu)
            )(prompt)
        except flyte.errors.OOMError:
            print(f"OOM at {memory}/{gpu} - escalating")
    raise flyte.errors.NonRecoverableError("Exhausted all resource tiers")

Same idea scales past resources. Because the handler is ordinary Python, the recovery path can call another task, ask an LLM what to do next, or write the failure to a store the next run reads. Your pipeline no longer a static DAG that dies on the first bad node and starts being something that adapts to what the infrastructure tells it.

Skip retries for non-recoverable errors

Some errors should fail immediately. Retrying a bad input 3 times just wastes time. Use `NonRecoverableError`:

@env.task(retries=3)
async def validate_and_process(config: dict) -> str:
    if "model_path" not in config:
        raise flyte.errors.NonRecoverableError(
            "Missing required field 'model_path'. Retrying won't help."
        )
    return f"processed with {config['model_path']}"

Even though the task has 3 retries configured, `NonRecoverableError` causes immediate failure without consuming any retry attempts.

Add timeouts to prevent runaway tasks

Combine retries with timeouts so a stuck task doesn't block your pipeline forever:

from datetime import timedelta

@env.task(
    retries=2,
    timeout=flyte.Timeout(
        max_runtime=timedelta(minutes=30),
        max_queued_time=timedelta(minutes=10),
    ),
)
async def bounded_training(data: str) -> float:
    # Fails if it takes more than 30 minutes per attempt
    # Also fails if it sits in the queue for more than 10 minutes
    return model

Combine retries with traces for maximum efficiency

When you pair retries with `@flyte.trace` (from Flyte Log #3), failed retries skip work that already completed:

@flyte.trace
async def step_1(data: str) -> str:
    return data.upper()

@flyte.trace
async def step_2(data: str) -> str:
    return data + "_processed"

@env.task(retries=3)
async def multi_step_pipeline(data: str) -> str:
    result = await step_1(data)    # checkpointed
    result = await step_2(result)  # checkpointed
    return result
    # If step_2 fails, the retry skips step_1 entirely

Key things to know

• `retries=3` means up to 4 total attempts (1 initial + 3 retries)
• Task errors become `flyte.errors.RuntimeUserError` with the original exception name in `e.code`
• `OOMError` and `TaskTimeoutError` are specific subclasses you can catch individually
• `NonRecoverableError` skips all remaining retries and fails immediately
• Retries, timeouts, and traces all compose together naturally

Full retries docs: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/retries-and-timeouts/

See what's happening in the Flyte Community:

Latest from the blog

• Flyte 2 Is Generally Available: The Durable, Open-Source AI Runtime - Read on Union.ai
• A Memory Store Built on Flyte and Cognee - Read on Union.ai
• Building Grounded Agents on Fresh Web Data - Read on Union.ai
• From DNA to 3D Fold: Compare a Gene Across Six Species with Carbon and ESMFold - Read on Union.ai
• Run Models, Agents and Apps on Infrastructure You Own - Read on union.ai
• Agents That Survive Production: Rebuilding 21 Design Patterns on Flyte - Read on union.ai
• Introducing Queues and Cluster Controls: Durable Workloads Under Contention - Read on union.ai
• Fine-tune an LLM with LoRA & QLoRA in a Flyte Pipeline - Read on union.ai

Recent talks & recordings

• When the Pipeline Breaks: Building ML Infrastructure for Biotech R&D | Session 1 - Watch on YouTube
• Building Code Mode Agents - Watch on YouTube
• LLM fine-tuning with GRPO - Watch on YouTube
• LLM fine-tuning with LoRA & QLoRA - Watch on YouTube

Upcoming events

• Aug 18th: Flyte 2: The Durable Runtime Built for AI - RSVP on Luma
• Aug 20th: Seattle RAG & Agent Context with Vector Stores | AI Hacknight - RSVP on Luma

Releases & updates

• Flyte 2 Is Generally Available: The Durable, Open-Source AI Runtime - Read on Union.ai

<div class="button-group is-center"><a class="button" target="_blank" rel="noopener noreferrer" href="https://www.union.ai/docs/v2/flyte/user-guide/run-modes/running-devbox/">Download Devbox</a></div>

From the community

• Open-Source Music Generation: Text-to-Music & Lyrics-to-Song - RSVP on Luma
• AI Book Club: Build a Reasoning Model (From Scratch) - RSVP on Luma

That's all for this week! - Sage Elliott