Self-Healing Agents and Pipelines: Adjust infra at runtime with .override()

AI engineering tip of the week: Self-Healing Agents and Pipelines: Adjust infra at runtime with .override()

You've defined a task with 4GB of memory. But sometimes you need to run it on a larger dataset that needs 32GB. Or you want to disable caching for a one-off debug run. Or you need to give a specific invocation a custom display name.

The .override() method lets you change task configuration at the call site without touching the task definition. Same task, different settings, one line of code.

You can see in the example image above we override both the memory resources and task name as the pipeline self-heals from OOM errors.

Override resources

import flyte

env = flyte.TaskEnvironment(
    name="training",
    resources=flyte.Resources(cpu=2, memory="4Gi"),
)

@env.task
async def train(dataset: str) -> float:
    return model

@env.task
async def pipeline(dataset: str) -> float:
    if dataset == "large":
        # Same task, more resources
        return await train.override(
            resources=flyte.Resources(cpu=8, memory="32Gi", gpu="A100:1")
        )(dataset)
    else:
        return await train(dataset)

The task code stays the same. Only the infrastructure changes.

Auto-recover from OOM

This is one of the most practical patterns: catch an OOM error and retry with more memory.

import flyte.errors

@env.task
async def process_data(path: str) -> str:
    return f"processed {path}"

@env.task
async def smart_processing(path: str) -> str:
    try:
        return await process_data(path)
    except flyte.errors.OOMError:
        print("OOM! Trying with 4x memory...")
        return await process_data.override(
            resources=flyte.Resources(cpu=4, memory="16Gi")
        )(path)

Progressive resource scaling

For unpredictable workloads, try multiple resource levels:

MEMORY_LEVELS = ["1Gi", "4Gi", "16Gi", "64Gi"]

@env.task
async def auto_scale(path: str) -> str:
    for memory in MEMORY_LEVELS:
        try:
            return await process_data.override(
                resources=flyte.Resources(cpu=2, memory=memory)
            )(path)
        except flyte.errors.OOMError:
            print(f"OOM at {memory}, trying next level...")
    raise RuntimeError("Failed at all memory levels")

Override retries and timeouts

from datetime import timedelta

@env.task(retries=1, timeout=60)
async def fetch_data(url: str) -> str:
    return f"data from {url}"

@env.task
async def pipeline() -> list[str]:
    # Quick call with defaults
    fast = await fetch_data("https://fast-api.example.com")

    # Slow endpoint: more retries, longer timeout
    slow = await fetch_data.override(
        retries=5,
        timeout=timedelta(minutes=10),
    )("https://slow-api.example.com")

    return [fast, slow]

Custom display names for parallel runs

When you fan out the same task many times, `.override(short_name=...)` makes each invocation identifiable in the Flyte UI:

import asyncio

@env.task
async def train_model(batch_size: int) -> float:
    return 1.0 / batch_size

@env.task
async def grid_search(batch_sizes: list[int]) -> list[float]:
    results = await asyncio.gather(*(
        train_model.override(
            short_name=f"train-bs-{bs}"
        )(batch_size=bs)
        for bs in batch_sizes
    ))
    return list(results)

Instead of seeing "train_model" five times in the UI, you see "train-bs-4", "train-bs-8", "train-bs-16", etc.

Override cache behavior

Disable caching for a specific call without changing the task definition:

@env.task(cache="auto")
async def expensive_computation(data: str) -> str:
    return f"result for {data}"

@env.task
async def debug_pipeline(data: str) -> str:
    # Normal run: uses cache
    cached_result = await expensive_computation(data)

    # Debug run: bypass cache to force re-execution
    fresh_result = await expensive_computation.override(
        cache="disable"
    )(data)

    return fresh_result

What you can override

• resources: CPU, memory, GPU, disk, shared memory
• retries: Number of retry attempts
• timeout: Max runtime and queue time
• short_name: Display name in the Flyte UI
• cache: Cache behavior ("auto", "disable", or a Cache object)
• env_vars: Environment variables
• secrets: Secret injection

Why this matters

`.override()` means you write one task and adapt it to different situations at the call site. No duplicate task definitions, no configuration flags, no if/else inside the task. The task stays clean. The caller decides the infrastructure.

The flyte errors combined with override also lets you to provide infrastructure as context allowing you to build duerable self-healing agentic systems.

Full override docs: https://www.union.ai/docs/v2/flyte/user-guide/tasks/task-configuration/resources/

See what's happening in the Flyte Community:

Latest from the blog

• Durable Execution for Any AI Agent Framework - Read on Union
• A Modern Alternative to Slurm for ML Workloads - Read on Union.ai
• Anatomy of a Durable Run - Read on Union.ai
• Flyte 2 Is Generally Available: The Durable, Open-Source AI Runtime - Read on Union.ai
• Why Untrusted Kernel Evaluation Needs Process Isolation (and How We Built It) - Read on Union
• A Memory Store Built on Flyte and Cognee - Read on Union.ai

Recent talks & recordings

• Flyte 2: The Durable Runtime Built for AI - Watch on YouTube
• 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

• more events will be posted soon, subscribe to the luma calendar: https://luma.com/unionai

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

• Sebastian Raschka | Build a Reasoning Model (From Scratch) - Author Discussion - RSVP on Luma
• World Models with DreamerV3 - AI Build & Learn - RSVP on Luma

That's all for this week! - Sage Elliott