Stop Burning GPU Hours: Right-Size your AI Tasks with Flyte Resources

AI engineering tip of the week: Stop Burning GPU Hours: Right-Size your Tasks with Flyte Resources

Not every task in your AI pipeline needs an A100. Your data preprocessing might run fine on 1 CPU and 512MB of RAM, while your model training task needs 8 GPUs and 64GB of memory. Flyte lets you specify resources per task, so each step gets exactly what it needs.

Specify CPU, memory, and disk

import flyte

env = flyte.TaskEnvironment(
    name="data_processing",
    resources=flyte.Resources(
        cpu="2",          # 2 CPU cores
        memory="4Gi",     # 4 GiB RAM
        disk="50Gi",      # 50 GiB ephemeral storage
    ),
)

@env.task
async def preprocess(data_path: str) -> str:
    return f"preprocessed {data_path}"

CPU supports Kubernetes-style strings (`"500m"` for half a core, `"4"` for 4 cores) or plain numbers. Memory uses standard units like `"512Mi"`, `"4Gi"`, `"1Ti"`.

Request a GPU

Just add gpu= to your Resources:

training_env = flyte.TaskEnvironment(
    name="training",
    resources=flyte.Resources(
        cpu=4,
        memory="16Gi",
        gpu="A100:1",   # 1 NVIDIA A100
    ),
)

@training_env.task
async def train_model(data: str) -> float:
    import torch
    print(f"Training on: {torch.cuda.get_device_name(0)}")
    # ...
    return trained_model

Pick the right GPU for the job

Flyte supports a wide range of accelerators. Just use the name and count:

# Budget inference
flyte.Resources(gpu="T4:1")

# Standard training
flyte.Resources(gpu="A100:1")

# Large model training (80GB VRAM)
flyte.Resources(gpu="A100 80G:4")

# Frontier models
flyte.Resources(gpu="H100:8")

# Inference-optimized
flyte.Resources(gpu="L4:1")

Supported GPUs include T4, A10, A10G, A100, A100 80G, L4, L40s, H100, H200, B200, and V100.

GPU memory partitioning (MIG)

Got an A100 or H100 but don’t need the whole thing? Use GPU partitioning to slice it up:

# Use 1/7th of an A100 with 5GB VRAM
flyte.Resources(gpu=flyte.GPU(device="A100", quantity=1, partition="1g.5gb"))

# Half an H100 with 40GB VRAM
flyte.Resources(gpu=flyte.GPU(device="H100", quantity=1, partition="3g.40gb"))

This can save resources when your tasks don’t need full GPU resources, like light inference or embedding generation.

Set request and limit ranges

If your task has variable resource needs, specify a range:

flyte.Resources(
    cpu=("1", "4"),           # request 1, burst up to 4
    memory=("2Gi", "8Gi"),    # request 2 GiB, limit 8 GiB
)

Shared memory for distributed training

PyTorch DataLoader with num_workers > 0 needs shared memory (/dev/shm):

flyte.Resources(
    cpu=8,
    memory="32Gi",
    gpu="A100:2",
    shm="auto",     # auto-sized shared memory
)

Different resources for different pipeline stages

The real power is giving each task exactly what it needs:

light_env = flyte.TaskEnvironment(
    name="light",
    resources=flyte.Resources(cpu="500m", memory="512Mi"),
)

heavy_env = flyte.TaskEnvironment(
    name="heavy",
    resources=flyte.Resources(cpu=8, memory="32Gi", gpu="A100:2"),
    depends_on=[light_env],
)

@light_env.task
async def download_data() -> str:
    return "dataset_v1"

@heavy_env.task
async def train(dataset: str) -> float:
    return 0.95

@heavy_env.task
async def pipeline() -> float:
    data = await download_data()    # runs on 0.5 CPU, 512Mi
    score = await train(data)       # runs on 8 CPU, 32Gi, 2x A100
    return score

Your download step doesn’t burn GPU hours. Your training step gets the horsepower it needs. You only pay for what you use.

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

See what’s happening in the Flyte Community:

Latest from the blog

• 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

• July 28th: Seattle TwelveLabs + Qdrant: AI Systems for Video Embeddings and Search - RSVP on Luma

Releases & updates

• Flyte 2 OSS: Backend Devbox and Reimagined UI - Read on Union
• June’s release brought first-class agents with memory and tool approval, SDK-authored MCP servers, backoff retries and per-attempt timeouts, multi-pod log streaming, and beta queues and events APIs. - Read the Release notes

<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 Text-to-Speech: “Natural” Voices - RSVP on Luma
• AI Book Club: Build a Reasoning Model (From Scratch) - RSVP on Luma

That’s all for this week! —Sage Elliott