Run the KubernetesPodOperator on Astro

The KubernetesPodOperator is one of the most powerful Apache Airflow operators. Similar to the Kubernetes executor, this operator dynamically launches a Pod in Kubernetes for each task and terminates each Pod once the task is complete. This results in an isolated, containerized execution environment for each task that is separate from tasks otherwise being executed by Celery workers.

Benefits

You can use the KubernetesPodOperator to:

• Execute a custom Docker image per task with Python packages and dependencies that would otherwise conflict with the rest of your Deployment's dependencies. This includes Docker images in a private registry or repository.
• Run tasks in a Kubernetes cluster outside of the Astro data plane. This can be helpful when you need to run individual tasks on infrastructure that isn't currently supported by Astro, such as GPU nodes or other third-party services.
• Specify CPU and memory as task-level limits or minimums to optimize performance and reduce costs.
• Write task logic in a language other than Python. This gives you flexibility and can enable new use cases across teams.
• Scale task growth horizontally in a way that is cost-effective, dynamic, and minimally dependent on worker resources.
• Set Kubernetes-native configurations in a YAML file, including volumes, secrets, and affinities.

On Astro, the Kubernetes infrastructure required to run the KubernetesPodOperator is built into every cluster in the data plane and is managed by Astronomer.

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 cannot run a KubernetesPodOperator task in a worker queue or node pool that is different than the worker queue of its parent worker. For example, a KubernetesPodOperator task that is triggered by an m5.4xlarge worker on AWS will also be run on an m5.4xlarge node. To run a task on a different node instance type, you must launch it in a Kubernetes cluster outside of the Astro data plane. If you need assistance launching KubernetesPodOperator tasks in external Kubernetes clusters, contact Astronomer support.

Prerequisites

• An Astro project.
• An Astro Deployment.

Set Up the KubernetesPodOperator

To use the KubernetesPodOperator in a DAG, add the following import statements and instantiation to your DAG file:

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,
)

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. The value you specify is assumed to be 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.

This is the minimum configuration required to run tasks with the KubernetesPodOperator on Astro. To further customize the way your tasks are run, see the topics below.

Configure Task-Level Pod Resources

Astro automatically allocates resources to Pods created by the KubernetesPodOperator. Resources used by the KubernetesPodOperator are not technically limited, meaning that the operator could theoretically use any CPU and memory that's available in your cluster to complete a task. Because of this, Astronomer recommends specifying compute resource requests and limits per task.

To do so, define a kubernetes.client.models.V1ResourceRequirements object and provide that to the 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>",
    resources=compute_resources,
    task_id="<task-name>",
    get_logs=True,
)

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

Run images from a Private Registry

By default, the KubernetesPodOperator expects to pull a Docker image that's hosted publicly on Docker Hub. If you want to execute a Docker image that's hosted in a private registry, you'll need to create a Kubernetes Secret and then specify the Kubernetes Secret in your DAG. If your Docker image is hosted in an Amazon Elastic Container Registry (ECR) repository, see Docker images hosted in private Amazon ECR repositories

Prerequisites

• An Astro project.
• An Astro Deployment.
• Access to a private Docker registry.
• kubectl, the command line tool for Kubernetes.

Step 1: Create a Kubernetes Secret

To run Docker images from a private registry on Astro, a Kubernetes Secret that contains credentials to your registry must be created. Injecting this secret into your Deployment's namespace will give your tasks access to Docker images within your private registry.

1. Log in to your Docker registry and follow the Kubernetes documentation to produce a /.docker/config.json file.
2. In the Cloud UI, select a Workspace and then select the Deployment you want to use the KubernetesPodOperator with.
3. Copy the value in the NAMESPACE field.
4. Contact Astronomer support and provide the namespace of the Deployment.

Astronomer Support will give you instructions on how to securely send the output of your /.docker/config.json file. Do not send this file by email, as it contains sensitive credentials to your registry. Astronomer will use this file to create a Kubernetes secret and inject it into your Deployment's namespace.

Step 2: Specify the Kubernetes Secret in your DAG

Once Astronomer has added the Kubernetes secret to your Deployment, you will be notified and provided with the name of the secret.

After you receive the name of your Kubernetes secret from Astronomer, you can run images from your private registry by importing models from kubernetes.client and configuring image_pull_secrets in your KubernetesPodOperator instantiation:

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,
)

Docker images hosted in private Amazon ECR repositories

If your Docker image is hosted in an Amazon ECR repository, add a permissions policy to the repository to allow the KubernetesPodOperator to pull the Docker image. You don't need to create a Kubernetes secret, or specify the Kubernetes secret in your DAG. Docker images hosted on Amazon ECR repositories can only be pulled from AWS clusters.

1. Log in to the Amazon ECR Dashboard and then select Menu > Repositories.

2. Click the Private tab and then click the name of the repository that hosts the Docker image.

3. Click Permissions in the left menu.

4. Click Edit policy JSON.

5. Copy and paste the following policy into the Edit JSON pane:

   {
   "Version": "2008-10-17",
   "Statement": [
       {
       "Sid": "AllowImagePullAstro",
       "Effect": "Allow",
       "Principal": {
           "AWS": "arn:aws:iam::<AstroAccountID>:root"
       },
       "Action": [
           "ecr:GetDownloadUrlForLayer",
           "ecr:BatchGetImage"
       ]
       }
   ]
   }

6. Replace <AstroAccountID> with your Astro AWS account ID.

7. Click Save to create a new permissions policy named AllowImagePullAstro.

8. Set up the KubernetesPodOperator.

9. Replace <your-docker-image> in the instantiation of the KubernetesPodOperator with the Amazon ECR repository URI that hosts the Docker image. To locate the URI:

   • In the Amazon ECR Dashboard, click Repositories in the left menu.
   • Click the Private tab and then copy the URI of the repository that hosts the Docker image.

Related documentation

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