Airflow KubernetesPodOperator Set Resource Requests and Limits (CPU, Memory, GPU, Ephemeral Storage)

When using the KubernetesPodOperator (KPO) in Airflow, we can specify the resource limits and requirements for the pod containers. This can be specified across resources like CPU, memory, GPU and ephemeral storage. Let’s take a look at how the resource requests and limits have to be specified in code.

Dict

The simplest way is to pass a dictionary of the requests and limits.

KubernetesPodOperator(
    ...
    resources={
        'request_ephemeral_storage': '2000M',
        'limit_ephemeral_storage': '3000M',
        'request_memory': '5000M',
        'limit_memory': '6000M',
        'request_cpu': '800m',
        'limit_cpu': '1000m',
        'limit_gpu': 2,
    },
)

Allowed keys are – request_memory, request_cpu, limit_memory, limit_cpu, limit_gpu, request_ephemeral_storage, limit_ephemeral_storage.

V1ResourceRequirements Object

Although passing a dict (as seen above) seems like the obvious choice, in the newer versions of the Operator it will not work. The “newer” way to specify the requests and limits is by passing the V1ResourceRequirements model of the K8S Python client directly. This way, you can pretty much pass any resource that you see listed in the Kubernetes docs.

Usage:

from kubernetes.client import models as k8s
# or from kubernetes import client as k8s

KubernetesPodOperator(
    ...
    resources=k8s.V1ResourceRequirements(
        requests={
            'cpu': '800m',
            'memory': '5000M',
            'ephemeral-storage', '2000M',
        },
        limits={
            'cpu': '1000m',
            'memory': '6000M',
            'ephemeral-storage': '3000M',
            'nvidia.com/gpu': 1, # or 'amd.com/gpu'
        },
    )
)

Leave a Reply

Your email address will not be published.