Resources

Task resources specify the computational limits and requests (CPU, memory, GPU, storage) that will be allocated to each task’s container during execution.

To specify resource requirements for your task, instantiate a Resources object with the desired parameters and assign it to either the resources parameter of the TaskEnvironment or the resources parameter of the override function (for invocation overrides).

Every task defined using that TaskEnvironment will run with the specified resources. If a specific task has its own resources defined in the decorator, it will override the environment’s resources for that task only.

If neither TaskEnvironment nor the task decorator specifies resources, the default resource allocation will be used.

Resources data class

For the full class definition, parameter types, and accepted formats, see the Resources API reference.

The main parameters are:

• cpu: CPU allocation, as a number, string ("500m"), or (request, limit) tuple.
• memory: Memory with Kubernetes units, such as "4Gi", or a (request, limit) tuple. Leave headroom below a node’s total RAM: its allocatable memory is smaller (the kubelet reserves overhead for the OS and system daemons), so a request near a node’s nominal capacity can leave the pod stuck Pending.
• gpu: GPU or other accelerator allocation, as "A100:2", an integer count, or flyte.GPU(), flyte.TPU(), flyte.AMD_GPU(), or flyte.Device() for advanced config. See Accelerators below.
• disk: Ephemeral storage, such as "10Gi".
• shm: Shared memory, such as "1Gi" or "auto".

Ephemeral storage

The disk parameter requests ephemeral storage — node-local scratch disk for the task’s container. Set it as a string with Kubernetes units, for example "50Gi":

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

Under the hood, disk maps to the Kubernetes ephemeral-storage resource on the task’s container.

What it covers. Ephemeral storage is the local disk a task writes to while it runs: the container’s writable filesystem and any temporary files your code creates on the local filesystem during execution (downloaded datasets, intermediate outputs, model checkpoints staged before offload). It is distinct from the offloaded storage backing flyte.io.File and flyte.io.Dir, which lives in the blob store rather than on the node.

Lifecycle. Ephemeral storage is tied to the task’s pod: it is provisioned when the task starts and reclaimed when the pod terminates. Nothing written to it survives beyond the task run, so use it for scratch work — persist anything you need to keep to a flyte.io.File or flyte.io.Dir in object storage.

Single value, not a request/limit range. Unlike cpu and memory, disk takes a single string, not a (request, limit) tuple.

Default behavior. Flyte does not set an ephemeral-storage request or limit when disk is unset. (A cluster-level Kubernetes LimitRange, if configured, may still inject a default.) The pod can still write to node-local disk, but it may be evicted if the node comes under storage pressure. Tasks doing heavy local data processing should set disk explicitly.

Examples

Usage in TaskEnvironment

Here’s a complete example of defining a TaskEnvironment with resource specifications for a machine learning training workload:

import flyte

# Define a TaskEnvironment for ML training tasks
env = flyte.TaskEnvironment(
    name="ml-training",
    resources=flyte.Resources(
        cpu=("2", "4"),        # Request 2 cores, allow up to 4 cores for scaling
        memory=("2Gi", "12Gi"), # Request 2 GiB, allow up to 12 GiB for large datasets
        disk="50Gi",           # 50 GiB ephemeral storage for checkpoints
        shm="8Gi"              # 8 GiB shared memory for efficient data loading
    )
)

# Use the environment for tasks
@env.task
async def train_model(dataset_path: str) -> str:
    # This task will run with flexible resource allocation
    return "model trained"

Usage in a task-specific override

# Demonstrate resource override at task invocation level
@env.task
async def heavy_training_task() -> str:
    return "heavy model trained with overridden resources"

@env.task
async def main():
    # Task using environment-level resources
    result = await train_model("data.csv")
    print(result)

    # Task with overridden resources at invocation time
    result = await heavy_training_task.override(
        resources=flyte.Resources(
            cpu="4",
            memory="24Gi",
            disk="100Gi",
            shm="16Gi"
        )
    )()
    print(result)

For complete format specifications for each resource type (CPU, memory, GPU/TPU/Device, disk, shared memory), including accepted string formats, request/limit ranges, GPU partitioning, and supported accelerator types, see the Resources API reference.

Accelerators

To run a task on a hardware accelerator — an NVIDIA or AMD GPU, a Google TPU, an AWS Trainium/Inferentia (Neuron) chip, or an Intel Habana Gaudi device — set the gpu parameter of Resources. Despite its name, gpu selects any supported accelerator type.

You can request an accelerator in three ways:

• Type-and-count string — "<device>:<count>", for example "T4:1" or "A100:8". This is the most common form.
• Integer count — gpu=2 requests that many GPUs of whatever type is available.
• Device object — flyte.GPU(), flyte.TPU(), flyte.AMD_GPU(), or flyte.Device() for advanced configuration such as GPU partitioning (MIG) and TPU slice topologies.

The gpu value can be set on the TaskEnvironment (applying to every task in it) or on a per-task override, exactly like the other resource fields shown above.

NVIDIA GPUs

Request an NVIDIA GPU by type and count:

import flyte

env = flyte.TaskEnvironment(
    name="nvidia",
    resources=flyte.Resources(gpu="T4:1"),  # one NVIDIA T4
)

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

GPU partitioning (MIG)

To request a Multi-Instance GPU (MIG) partition, use flyte.GPU() with a partition:

resources=flyte.Resources(gpu=flyte.GPU(device="A100", quantity=1, partition="1g.5gb"))

Partitioning is available on A100, A100 80G, H100, and H200.

Google TPUs

Use flyte.TPU() with the device type and, optionally, a slice topology:

resources=flyte.Resources(gpu=flyte.TPU(device="V5P", partition="2x2x1"))

Supported TPU device types are V5P and V6E.

AMD GPUs

Request an AMD GPU by type and count, or with flyte.AMD_GPU():

resources=flyte.Resources(gpu="MI300X:1")

Supported AMD types include MI100, MI210, MI250, MI250X, MI300A, MI300X, MI325X, MI350X, and MI355X.

AWS Trainium and Inferentia (Neuron)

Request AWS Neuron accelerators — Trainium (Trn) or Inferentia (Inf) — by type and count:

resources=flyte.Resources(gpu="Trn1:1")

Supported types include Trn1, Trn1n, Trn2, Trn2u, Inf1, and Inf2.

Intel Habana Gaudi

Request an Intel Habana Gaudi accelerator by type and count:

resources=flyte.Resources(gpu="Gaudi1:1")

For the full list of accepted accelerator strings and device-configuration options, see the Resources API reference.