Accelerators#

Union allows you to specify requests and limits for the number of GPUs available for a given task. However, in some cases, you may want to be more specific about the type of GPU to be used.

You can use the accelerator parameter to specify specific GPU types.

Union Serverless comes with three GPU types available:

NVIDIA T4 Tensor Core GPU
NVIDIA L4 Tensor Core GPU
NVIDIA A100 GPU

Pricing for these GPUs can found on the Union Pricing page.

NVIDIA T4 Tensor Core GPU#

The NVIDIA T4 Tensor Core GPU is the default. To use it for a task, specify the number of GPUs required in the limits parameter:

    @task(
        limits=Resources(gpu="1")
    )
    def my_task():
        ...

Or, you can explicitly specify the accelerator parameter as follows:

    @task(
        limits=Resources(gpu="1"),
        accelerator=GPUAccelerator("nvidia-tesla-t4")
    )
    def my_task():
        ...

NVIDIA L4 Tensor Core GPU#

To use the NVIDIA L4 Tensor Core GPU for a task, you must specify the number of GPUs required in the limits parameter, and also specify the accelerator parameter as follows:

from flytekit.extras.accelerators import L4

@task(
    requests=Resources(gpu="1"),
    accelerator=L4,
)
def my_task():
    ...

NVIDIA A100 GPU#

To use the NVIDIA A100 GPU for a task you must specify the number of GPUs required in the limits parameter, and also specify the accelerator parameter as follows:

@task(
    requests=Resources(gpu="1"),
    accelerator=GPUAccelerator("nvidia-tesla-a100"),
)
def my_task():
    ...