Skip to main content

Configure Images and Compute Resources

When deploying Flyte tasks, workloads require specific runtime environments (operating system packages, Python dependencies, custom files) and tailored hardware resources (CPUs, system memory, shared memory, GPUs, TPUs, or hardware slices). In flyte-sdk, you define container images with Image and specify compute constraints with Resources and Device inside a TaskEnvironment.

Defining Images and Resources on a TaskEnvironment​

To execute tasks with custom dependencies and allocated compute hardware, declare a TaskEnvironment configured with an Image and a Resources object.

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Resources' does not accept keyword(s): cpu, disk, gpu, memory, shm

from pathlib import Path
from flyte import Image, Resources, TaskEnvironment

# Define a container image with system and Python dependencies
custom_image = (
    Image.from_debian_base(python_version=(3, 12))
    .with_apt_packages("curl", "git")
    .with_pip_packages("pandas", "torch")
    .with_env_vars({"MODEL_CACHE_DIR": "/tmp/cache"})
)

# Define compute resources: 2 CPUs, 8 GiB RAM, 20 GiB disk, auto shared memory, 1 A100 GPU
task_resources = Resources(
    cpu=2,
    memory="8Gi",
    disk="20Gi",
    shm="auto",
    gpu="A100:1",
)

# Create the task environment
training_env = TaskEnvironment(
    name="training_env",
    image=custom_image,
    resources=task_resources,
)

@training_env.task
async def train_model(batch_size: int) -> str:
    return f"Model trained with batch size {batch_size}"

In this setup:

  • TaskEnvironment bundles default runtime settings for all tasks registered to it. The fully qualified name for tasks registered on this environment becomes training_env.<function_name>.
  • Image uses an immutable, fluent builder starting with a base image constructor (such as from_debian_base()) and chaining layer operations (such as with_apt_packages() and with_pip_packages()).
  • Resources configures hardware requirements requested for the task container.

Building and Selecting Container Images​

Image objects cannot be instantiated directly via Image() (direct instantiation raises a TypeError). Instead, initialize an Image using a from_* base constructor, then chain with_* methods to add immutable layers.

1. Base Image Constructors​

flyte-sdk provides five base constructors:

from pathlib import Path
from flyte import Image

# 1. Official Debian base with a specific Python version (multi-arch by default)
debian_img = Image.from_debian_base(
    python_version=(3, 12),
    name="my-custom-image",
    registry="ghcr.io/my-org",
)

# 2. Existing pre-built registry image
prebuilt_img = Image.from_base("python:3.12-slim-bookworm")

# 3. Custom Dockerfile (cannot be chained with further with_* layers)
dockerfile_img = Image.from_dockerfile(
    file=Path("/path/to/Dockerfile"),
    name="dockerfile-app",
    registry="ghcr.io/my-org",
)

# 4. uv inline dependency script
uv_script_img = Image.from_uv_script(
    script=Path("scripts/run.py"),
    name="uv-script-image",
    registry="ghcr.io/my-org",
)

# 5. Pre-configured image reference name
ref_img = Image.from_ref_name("default")

2. Customizing Image Layers​

Chain modification methods to install packages, copy files, and configure the container environment:

from pathlib import Path
from flyte import Image

image = (
    Image.from_debian_base(python_version=(3, 11), name="data-pipeline")
    # System dependencies via apt-get
    .with_apt_packages("build-essential", "libpq-dev")
    # Python dependencies via pip
    .with_pip_packages("scikit-learn", "pyarrow>=14.0.0")
    # Requirements text file
    .with_requirements(Path("requirements.txt"))
    # Working directory
    .with_workdir("/app")
    # Environment variables
    .with_env_vars({"APP_ENV": "production", "LOG_LEVEL": "INFO"})
    # Arbitrary shell commands (do not prefix with RUN)
    .with_commands(["mkdir -p /tmp/models", "chmod 777 /tmp/models"])
    # Copy local files and directories into the image
    .with_source_folder(src=Path("./config"), dst="/app/config")
    .with_source_file(src=Path("./setup.cfg"), dst="/app/setup.cfg")
)

3. Managing Projects with uv or Poetry​

Install dependencies directly from project lockfiles:

from pathlib import Path
from flyte import Image

# uv project with pyproject.toml and uv.lock
uv_image = Image.from_debian_base(name="uv-app").with_uv_project(
    pyproject_file=Path("pyproject.toml"),
    uvlock=Path("uv.lock"),
    project_install_mode="dependencies_only",  # or "install_project"
)

# Poetry project with pyproject.toml and poetry.lock
poetry_image = Image.from_debian_base(name="poetry-app").with_poetry_project(
    pyproject_file=Path("pyproject.toml"),
    poetry_lock=Path("poetry.lock"),
    project_install_mode="dependencies_only",
)

4. Build-Time Secrets for Private Packages​

To pull private packages from authenticated registries or repositories during image building, pass secret mounts to with_pip_packages(), with_apt_packages(), or with_requirements():

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Secret' does not accept keyword(s): key 'flyte.Secret' does not accept keyword(s): key, mount

from pathlib import Path
from flyte import Image, Secret

private_repo_image = (
    Image.from_debian_base(name="private-tools")
    .with_pip_packages(
        "git+https://$GITHUB_PAT@github.com/my-org/private-repo.git",
        secret_mounts=[Secret(key="GITHUB_PAT")],
    )
    .with_apt_packages(
        "custom-pkg",
        secret_mounts=[Secret(key="apt-secret", mount=Path("/etc/apt/apt-secret"))],
    )
)

Configuring Compute Resources​

The Resources class defines compute allocations for CPU, system memory, ephemeral disk, shared memory (shm), and hardware accelerators.

CPU and Memory Specifications​

You can pass scalar values or request/limit tuples:

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Resources' does not accept keyword(s): cpu, disk, memory, shm

from flyte import Resources

# Fixed allocations
fixed_res = Resources(
    cpu=2,             # 2 CPU cores (can also pass 0.5, "500m", "2")
    memory="4Gi",      # 4 GiB memory ("512Mi", "1Gi", "1G", "4Gi")
    disk="50Gi",       # Ephemeral local storage
    shm="4Gi",         # Shared memory for IPC/data loaders, or "auto"
)

# Request and limit ranges using 2-element tuples: (request, limit)
ranged_res = Resources(
    cpu=(1, 4),               # Request 1 core, limit up to 4 cores
    memory=("2Gi", "16Gi"),   # Request 2 GiB, limit up to 16 GiB
    disk="100Gi",
    shm="auto",               # Automatically sets /dev/shm to maximum node capacity
)

Configuring Accelerators and Hardware Partitions​

Accelerators (NVIDIA GPUs, AMD GPUs, Google Cloud TPUs, AWS Trainium/Inferentia, Habana Gaudi) can be specified on Resources using shorthand strings or explicit Device instances.

String Shorthand​

Pass the accelerator model and count formatted as "<Device>:<Quantity>":

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Resources' does not accept keyword(s): cpu, gpu, memory

from flyte import Resources

# 1 NVIDIA T4 GPU
t4_res = Resources(cpu=2, memory="8Gi", gpu="T4:1")

# 8 NVIDIA A100 80GB GPUs
a100_res = Resources(cpu=16, memory="64Gi", gpu="A100 80G:8")

# 8 NVIDIA H100 GPUs
h100_res = Resources(cpu=32, memory="128Gi", gpu="H100:8")

# AMD MI300X GPU
amd_res = Resources(cpu=8, memory="32Gi", gpu="MI300X:1")

Explicit Device Helpers and MIG / Slice Partitioning​

For hardware partitioning (such as Multi-Instance GPU partitioning on NVIDIA A100/H100/H200, or multi-dimensional slice topologies on TPUs), use helper factory functions that return Device instances:

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Resources' does not accept keyword(s): cpu, gpu, memory

from flyte import AMD_GPU, GPU, HABANA_GAUDI, Neuron, Resources, TPU, TaskEnvironment

# NVIDIA GPU with Multi-Instance GPU (MIG) partition
mig_gpu_env = TaskEnvironment(
    name="gpu_mig_env",
    resources=Resources(
        cpu=4,
        memory="16Gi",
        gpu=GPU(device="A100", quantity=1, partition="1g.5gb"),
    ),
)

# Google Cloud TPU v5p with a 2x2x1 slice topology
tpu_env = TaskEnvironment(
    name="tpu_env",
    resources=Resources(
        cpu=8,
        memory="32Gi",
        gpu=TPU(device="V5P", partition="2x2x1"),
    ),
)

# AWS Trainium / Inferentia via Neuron
neuron_env = TaskEnvironment(
    name="neuron_env",
    resources=Resources(
        cpu=8,
        memory="32Gi",
        gpu=Neuron(device="Trn1"),
    ),
)

# AMD GPU
rocm_env = TaskEnvironment(
    name="rocm_env",
    resources=Resources(
        cpu=8,
        memory="32Gi",
        gpu=AMD_GPU(device="MI300X"),
    ),
)

# Habana Gaudi
gaudi_env = TaskEnvironment(
    name="gaudi_env",
    resources=Resources(
        cpu=8,
        memory="32Gi",
        gpu=HABANA_GAUDI(device="Gaudi1"),
    ),
)

Overriding and Cloning Configurations​

Runtime Resource Overrides with task.override()​

Tasks can override compute resources dynamically at invocation time (for example, when retrying after an out-of-memory error):

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Resources' does not accept keyword(s): cpu, memory 'flyte.Resources' does not accept keyword(s): cpu, gpu, memory

from flyte import Resources, TaskEnvironment

env = TaskEnvironment(
    name="dynamic_res_env",
    resources=Resources(cpu=1, memory="250Mi"),
)

@env.task
async def memory_intensive_task(data_size: int) -> int:
    return data_size * 2

@env.task
async def workflow_entrypoint() -> int:
    # Dynamically scale up CPU and memory for this invocation
    scaled_task = memory_intensive_task.override(
        resources=Resources(cpu=4, memory="8Gi", gpu="T4:1")
    )
    return await scaled_task(data_size=1000)

Reusing Base Environments with TaskEnvironment.clone_with()​

When creating multiple task environments with slight variations, use clone_with() instead of duplicating configurations:

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Resources' does not accept keyword(s): cpu, memory 'flyte.Resources' does not accept keyword(s): cpu, gpu, memory

from flyte import Image, Resources, TaskEnvironment

base_env = TaskEnvironment(
    name="base_env",
    image=Image.from_debian_base(python_version=(3, 12)).with_pip_packages("pandas"),
    resources=Resources(cpu=2, memory="4Gi"),
)

# Clone with increased resources and spot instance support
heavy_env = base_env.clone_with(
    name="heavy_env",
    resources=Resources(cpu=8, memory="32Gi", gpu="A100:1"),
    interruptible=True,
)

Troubleshooting and Edge Cases​

Direct Image() Instantiation Error​

Symptom: Calling Image() raises:

TypeError: Direct instantiation of Image not allowed, please use one of the various from_...() methods instead

Fix: Initialize using Image.from_debian_base(), Image.from_base(), Image.from_dockerfile(), Image.from_uv_script(), or Image.from_ref_name().

Layering on Dockerfile-Based Images​

Symptom: Calling .with_pip_packages() or other with_* methods on an image created with Image.from_dockerfile() raises:

ValueError: Flyte current cannot add additional layers to a Dockerfile-based Image. Please amend the dockerfile directly.

Fix: Place all package installations, environment variables, and file copies directly inside the Dockerfile when using from_dockerfile().

Invalid Tuple Lengths for CPU or Memory​

Symptom: Passing a 1-element or 3-element tuple for cpu or memory (e.g., cpu=(1, 2, 4)) raises:

ValueError: cpu tuple must have exactly two elements

Fix: Always pass exactly two elements representing (request, limit) when using tuple notation: cpu=(1, 2) or memory=("1Gi", "4Gi").

Unrecognized Partition Strings​

Symptom: Specifying an invalid partition in GPU() or TPU() raises a ValueError indicating the partition is not valid for that device. Fix: Match partition strings to supported hardware specifications (e.g., "1g.5gb", "2g.10gb", "3g.20gb", "4g.20gb", "7g.40gb" for NVIDIA A100; "1x1", "2x2x1", "2x2x2", "2x2x4" for Google Cloud TPU V5P).

Overriding Resources on Reusable Environments​

Symptom: Calling task.override(resources=...) fails when the task environment is defined with reusable=ReusePolicy(...). Fix: When overriding resources, env_vars, or secrets on a reusable environment, pass reusable="off" in the same override call:

Unverified example

This example could not be verified against this version of the codebase and may not work as shown. Validator finding: 'flyte.Resources' does not accept keyword(s): memory

from flyte import Resources

await task.override(
    resources=Resources(memory="16Gi"),
    reusable="off",
)(arg)