# 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](https://www.union.ai/docs/v2/flyte/user-guide/api-reference/flyte-sdk/packages/flyte/resources).

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](#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"`:

```python
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`](https://kubernetes.io/docs/concepts/configuration/manage-resources-containers/#local-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"
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/resources/resources.py*

### 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)
```

*Source: https://github.com/unionai/unionai-examples/blob/main/v2/user-guide/task-configuration/resources/resources.py*

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](https://www.union.ai/docs/v2/flyte/user-guide/api-reference/flyte-sdk/packages/flyte/resources).

## 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:

```python
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`:

```python
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:

```python
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()`:

```python
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:

```python
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:

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

> [!NOTE]
> Which accelerator types are actually available depends on your deployment and the node pools configured in your cluster. Requesting a type that no node provides will leave the task's pod `Pending`.

For the full list of accepted accelerator strings and device-configuration options, see the [`Resources` API reference](https://www.union.ai/docs/v2/flyte/user-guide/api-reference/flyte-sdk/packages/flyte/resources).

---
**Source**: https://github.com/unionai/unionai-docs/blob/main/content/user-guide/task-configuration/resources.md
**HTML**: https://www.union.ai/docs/v2/flyte/user-guide/task-configuration/resources/