Template container_resources of KubernetesPod

[Atlaoui]

What do you see as an issue?

I have DAG that uses a KubernetesPodOperator to run a containerized task. I would like to be able to parameterize the resources (memory and CPU) allocated to the container, so that I can change them depending on the specific DAG run.

Solving the problem

I tried using Airflow's params to define the container_resources argument in the KubernetesPodOperator, but I couldn't find a way to reference the parameters in the dictionary. Here's a simplified version of my code:

from airflow import DAG
from airflow.contrib.operators import KubernetesPodOperator

default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2023, 3, 7),
    'email': ['airflow@example.com'],
    'email_on_failure': False,
    'email_on_retry': False,
}

with DAG(
    'my_dag',
    default_args=default_args,
    schedule_interval='@daily',
    catchup=False,
    params={
        "request_memory": Param("2000", type="string"),
        "request_cpu": Param("4000", type="string"),
        "limit_memory": Param("4G", type="string"),
        "limit_cpu": Param("5000", type="string"),
    }
) as dag:

    task = KubernetesPodOperator(
        task_id='my_task',
        name='my_task',
        image='my_image:latest',
        namespace='my_namespace',
        image_pull_policy='Always',
        is_delete_operator_pod=True,
        container_resources={
            'request_memory': '{{ params.request_memory }}',# tried dag.params["request_memory"],
            'limit_memory': '{{ params.limit_memory }}' 
            'request_cpu': '{{ params.request_cpu }}',
            'limit_cpu': '{{ params.limit_cpu }}',
        },
    )
    
 # I also tried to  use k8s.V1ResourceRequirements


ressources = k8s.V1ResourceRequirements(
    limits={
        'cpu': '{{ params.limit_cpu }}',
        'memory': '{{ params.limit_memory }}'
    },
    requests={
        'cpu': '{{ params.request_cpu }}',
        'memory': '{{ params.request_memory }}'
    }
)

task = KubernetesPodOperator(
    container_resources=resources
)

Anything else

In the doc we can read that :

container_resources (k8s.V1ResourceRequirements | None) – resources for the launched pod. (templated)

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[josh-fell]

@Atlaoui Trying to understand what's going on. Is this a documentation issue or are you saying that you are unable to pass a templated arg for container_resources?

The container_resources parameter is a template field like you mentioned so your example should work properly. Although, your import of the KubernetesPodOperator is quite outdated (pre-Airflow 2.0). Do you/have you tried installing the Kubernetes provider package (latest is apache-airflow-providers-cncf-kubernetes==5.1.1) and updating the operator import to airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator?

[josh-fell]

Or is the second container_resources arg the issue?

ressources = k8s.V1ResourceRequirements(
    limits={
        'cpu': '{{ params.limit_cpu }}',
        'memory': '{{ params.limit_memory }}'
    },
    requests={
        'cpu': '{{ params.request_cpu }}',
        'memory': '{{ params.request_memory }}'
    }
)

[eladkal]

Something does look odd here but i'm not sure if it's related to templating as both cases it looks like wrong syntax?

from datetime import datetime

from kubernetes import client as k8s

from airflow import DAG
from airflow.models import Param
from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator

with DAG(
    dag_id="my_kube_dag",
    start_date=datetime(2023, 3, 7),
    schedule_interval=None,
    catchup=False,
    params={
        "request_memory": Param("2000", type="string"),
        "request_cpu": Param("4000", type="string"),
        "limit_memory": Param("4G", type="string"),
        "limit_cpu": Param("5000", type="string"),
    }
) as dag:
    task1 = KubernetesPodOperator(
        namespace="default",
        image="ubuntu:16.04",
        cmds=["bash", "-cx"],
        arguments=["lalala"],
        task_id="task1",
        in_cluster=False,
        do_xcom_push=False,
        is_delete_operator_pod=False,
        container_resources=k8s.V1ResourceRequirements(
            limits={
                "cpu": "1",
                "memory": "1G",
            },
            requests={
                "cpu": "1",
                "memory": "1",
            },
        )
    )

    task2 = KubernetesPodOperator(
        namespace="default",
        image="ubuntu:16.04",
        cmds=["bash", "-cx"],
        arguments=["lalala"],
        task_id="task2",
        in_cluster=False,
        do_xcom_push=False,
        is_delete_operator_pod=False,
        container_resources=k8s.V1ResourceRequirements(
            limits={
                "cpu": "{{ params.limit_cpu }}",
                "memory": "{{ params.limit_memory }}",
            },
            requests={
                "cpu": "{{ params.request_cpu }}",
                "memory": "{{ params.request_memory }}",
            }
        )
    )

task1:

task2:

[Atlaoui]

Ah sorry I did not realized that I was on the wrong doc version , I was using the apache-airflow-providers-cncf-kubernetes==4.2.2, I updated to 5.2.2 and it's worked
Thank you