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.

from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator

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 to the container_resources keyword argument instead of resources. 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
from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator

KubernetesPodOperator(
    ...
    container_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'
        },
    )
)