Run the KubernetesPodOperator on Astro

The KubernetesPodOperator is one of the most customizable Apache Airflow operators. A task using the KubernetesPodOperator runs in a dedicated, isolated Kubernetes Pod that terminates after the task completes. To learn more about the benefits and usage of the KubernetesPodOperator, see the KubernetesPodOperator Learn guide.

On Astro, the infrastructure required to run the KubernetesPodOperator is built into every Deployment and is managed by Astronomer. Astro supports setting a default Pod configuration so that any task Pods without specific resource requests and limits cannot exceed your expected resource usage for the Deployment.

Some task-level configurations will differ on Astro compared to other Airflow environments. Use this document to learn how to configure individual task Pods for different use cases on Astro. To configure the default Pod resources for all KubernetesPodOperator Pods, see Configure Kubernetes Pod resources.

Known limitations​

• Cross-account service accounts are not supported on Pods launched in an Astro cluster. To allow access to external data sources, you can provide credentials and secrets to tasks.

• PersistentVolumes (PVs) are not supported on Pods launched in an Astro cluster.

• You can't use an image built for an ARM architecture in the KubernetesPodOperator. To build images using the x86 architecture on a Mac with an Apple chip, include the --platform flag in the FROM command of the Dockerfile that constructs your custom image. For example:

  FROM --platform=linux/amd64 postgres:latest
  

  If you use an ARM image, your KPO task will fail with the error: base] exec /usr/bin/psql: exec format error.

Prerequisites​

• An Astro project.
• An Astro Deployment.

Set up the KubernetesPodOperator on Astro​

The following snippet is the minimum configuration you'll need to create a KubernetesPodOperator task on Astro:

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


namespace = conf.get("kubernetes", "NAMESPACE")

KubernetesPodOperator(
    namespace=namespace,
    image="<your-docker-image>",
    cmds=["<commands-for-image>"],
    arguments=["<arguments-for-image>"],
    labels={"<pod-label>": "<label-name>"},
    name="<pod-name>",
    task_id="<task-name>",
    get_logs=True,
    in_cluster=True,
)

For each instantiation of the KubernetesPodOperator, you must specify the following values:

• namespace = conf.get("kubernetes", "NAMESPACE"): Every Deployment runs on its own Kubernetes namespace within a cluster. Information about this namespace can be programmatically imported as long as you set this variable.
• image: This is the Docker image that the operator will use to run its defined task, commands, and arguments. Astro assumes that this value is an image tag that's publicly available on Docker Hub. To pull an image from a private registry, see Pull images from a Private Registry.
• in_cluster: If a Connection object is not passed to the KubernetesPodOperator's kubernetes_conn_id parameter, specify in_cluster=True to run the task in the Deployment's Astro cluster.

Configure task-level Pod resources​

Astro automatically allocates resources to Pods created by the KubernetesPodOperator. Unless otherwise specified in your task-level configuration, the amount of resources your task Pod can use is defined by your default Pod resource configuration. To further optimize your resource usage, Astronomer recommends specifying compute resource requests and limits for each task.

To do so, define a kubernetes.client.models.V1ResourceRequirements object and provide that to the container_resources argument of the KubernetesPodOperator. For example:

from airflow.configuration import conf
from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
from kubernetes.client import models as k8s

compute_resources = k8s.V1ResourceRequirements(
    limits={"cpu": "800m", "memory": "3Gi"},
    requests={"cpu": "800m", "memory": "3Gi"}
)

namespace = conf.get("kubernetes", "NAMESPACE")

KubernetesPodOperator(
    namespace=namespace,
    image="<your-docker-image>",
    cmds=["<commands-for-image>"],
    arguments=["<arguments-for-image>"],
    labels={"<pod-label>": "<label-name>"},
    name="<pod-name>",
    container_resources=compute_resources,
    task_id="<task-name>",
    get_logs=True,
    in_cluster=True,
)

Applying the previous code example ensures that when this DAG runs, it launches a Kubernetes Pod with exactly 800m of CPU and 3Gi of memory as long as that infrastructure is available in your Deployment. After the task finishes, the Pod will terminate gracefully.

Mount a temporary directory​

To run a task run the KubernetesPodOperator that utilizes your Deployment's ephemeral storage, mount an emptyDir volume to the KubernetesPodOperator. For example:

from airflow.configuration import conf
from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
from kubernetes.client import models as k8s

volume = k8s.V1Volume(
    name="cache-volume",
    emptyDir={},
)

volume_mounts = [
    k8s.V1VolumeMount(
        mount_path="/cache", name="cache-volume"
    )
]

example_volume_test = KubernetesPodOperator(
    namespace=namespace,
    image="<your-docker-image>",
    cmds=["<commands-for-image>"],
    arguments=["<arguments-for-image>"],
    labels={"<pod-label>": "<label-name>"},
    name="<pod-name>",
    task_id="<task-name>",
    get_logs=True,
    in_cluster=True,
    volume_mounts=volume_mounts,
    volumes=[volume],
)

Run the current Airflow image​

You can run your Deployment's current Airflow image with the KubernetesPodOperator by using the environment variable, ASTRONOMER_AIRFLOW_IMAGE. This environment variable allows you to run KPO tasks using the same Runtime image that your Deployment uses, in a way that won't be affected by underlying cluster infrastructure changes that might change the base image repository URL. The ASTRONOMER_AIRFLOW_IMAGE environment variable allows you to ensure that your Deployment retrieves the correct image URL.

The following code example shows how you can configure your KPO with ASTRONOMER_AIRFLOW_IMAGE:

import os

from airflow.configuration import conf
from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator
from kubernetes.client import models as k8s


KubernetesPodOperator(
    namespace=conf.get("kubernetes", "NAMESPACE"),
    image=os.getenv("ASTRONOMER_AIRFLOW_IMAGE"),
    image_pull_secrets=[
        k8s.V1LocalObjectReference("image-pull-secret"),
        k8s.V1LocalObjectReference("per-dp-registry-pull-secret"),
    ],
    cmds=["<commands-for-image>"],
    arguments=["<arguments-for-image>"],
    labels={"<pod-label>": "<label-name>"},
    name="<pod-name>",
    task_id="<task-name>",
    get_logs=True,
    in_cluster=True,
)

Run images from a private registry​

By default, the KubernetesPodOperator expects to pull a Docker image that's hosted publicly. If your images are hosted on the container registry native to your cloud provider, you can grant access to the images directly. Otherwise, if you are using any other private registry, you need to create a Kubernetes Secret containing credentials to the registry, then specify the Kubernetes Secret in your DAG.

{
    "auths": {
        "https://index.docker.io/v1/": {
            "auth": "c3R...zE2"
        }
    }
}

from kubernetes.client import models as k8s

KubernetesPodOperator(
    namespace=namespace,
    image_pull_secrets=[k8s.V1LocalObjectReference("<your-secret-name>")],
    image="<your-docker-image>",
    cmds=["<commands-for-image>"],
    arguments=["<arguments-for-image>"],
    labels={"<pod-label>": "<label-name>"},
    name="<pod-name>",
    task_id="<task-name>",
    get_logs=True,
    in_cluster=True,
)

Use secret environment variables with the KubernetesPodOperator​

Astro environment variables marked as secrets are stored in a Kubernetes secret called env-secrets. To use a secret value in a task running on the Kubernetes executor, you pull the value from env-secrets and mount it to the Pod running your task as a new Kubernetes Secret.

1. Add the following import to your DAG file:

   from airflow.kubernetes.secret import Secret
   

2. Define a Kubernetes Secret in your DAG instantiation using the following format:

   secret_env = Secret(deploy_type="env", deploy_target="<VARIABLE_KEY>", secret="env-secrets", key="<VARIABLE_KEY>")
   namespace = conf.get("kubernetes", "NAMESPACE")
   

3. Reference the key for the environment variable, formatted as $VARIABLE_KEY in the task using the KubernetesPodOperator.

In the following example, a secret named MY_SECRET is pulled from env-secrets and printed to logs.

import pendulum
from airflow.kubernetes.secret import Secret

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

with DAG(
        dag_id='test-kube-pod-secret',
        start_date=pendulum.datetime(2022, 1, 1, tz="UTC"),
        end_date=pendulum.datetime(2022, 1, 5, tz="UTC"),
        schedule_interval="@once",
) as dag:

    secret_env = Secret(deploy_type="env", deploy_target="MY_SECRET", secret="env-secrets", key="MY_SECRET")

    namespace = conf.get("kubernetes", "NAMESPACE")

    k = KubernetesPodOperator(
        namespace=namespace,
        image="ubuntu:16.04",
        cmds=["bash", "-cx"],
        arguments=["echo $MY_SECRET && sleep 150"],
        name="test-name",
        task_id="test-task",
        get_logs=True,
        in_cluster=True,
        secrets=[secret_env],
    )

Use the @task.kubernetes decorator​

The @task.kubernetes decorator provides a TaskFlow alternative to the traditional KubernetesPodOperator, which allows you to run a specified task in its own Kubernetes pod. Note that the Docker image provided to the @task.kubernetes decorator's image parameter must support executing Python scripts in order to leverage the KubernetesPodOperator decorator.

Like regular @task decorated functions, XComs can be passed to the Python script running in the dedicated Kubernetes pod. If do_xcom_push is set to True in the decorator parameters, the value returned by the decorated function is pushed to XCom.

Astronomer recommends using the @task.kubernetes decorator instead of the KubernetesPodOperator when using XCom with Python scripts in a dedicated Kubernetes pod.

from pendulum import datetime
from airflow.configuration import conf
from airflow.decorators import dag, task
import random

# get the current Kubernetes namespace Airflow is running in
namespace = conf.get("kubernetes", "NAMESPACE")


@dag(
    start_date=datetime(2023, 1, 1),
    catchup=False,
    schedule="@daily",
)
def kubernetes_decorator_example_dag():
    @task
    def extract_data():
        # simulating querying from a database
        data_point = random.randint(0, 100)
        return data_point

    @task.kubernetes(
        # specify the Docker image to launch, it needs to be able to run a Python script
        image="python",
        # launch the Pod on the same cluster as Airflow is running on
        in_cluster=True,
        # launch the Pod in the same namespace as Airflow is running in
        namespace=namespace,
        # Pod configuration
        # naming the Pod
        name="my_pod",
        # log stdout of the container as task logs
        get_logs=True,
        # log events in case of Pod failure
        log_events_on_failure=True,
        # enable pushing to XCom
        do_xcom_push=True,
    )
    def transform(data_point):
        multiplied_data_point = 23 * int(data_point)
        return multiplied_data_point

    @task
    def load_data(**context):
        # pull the XCom value that has been pushed by the KubernetesPodOperator
        transformed_data_point = context["ti"].xcom_pull(
            task_ids="transform", key="return_value"
        )
        print(transformed_data_point)

    load_data(transform(extract_data()))


kubernetes_decorator_example_dag()

Related documentation​

• How to use cluster ConfigMaps, Secrets, and Volumes with Pods
• KubernetesPodOperator Airflow Guide