example_kubernetes_executor_config.py

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""
This is an example dag for using a Kubernetes Executor Configuration.
"""
import logging
import os

from airflow import DAG
from airflow.example_dags.libs.helper import print_stuff
from airflow.operators.python import PythonOperator
from airflow.settings import AIRFLOW_HOME
from airflow.utils.dates import days_ago

default_args = {
    'owner': 'airflow',
}

log = logging.getLogger(__name__)

try:
    from kubernetes.client import models as k8s

    with DAG(
        dag_id='example_kubernetes_executor_config',
        default_args=default_args,
        schedule_interval=None,
        start_date=days_ago(2),
        tags=['example3'],
    ) as dag:

        def test_sharedvolume_mount():
            """
            Tests whether the volume has been mounted.
            """
            for i in range(5):
                try:
                    return_code = os.system("cat /shared/test.txt")
                    if return_code != 0:
                        raise ValueError(f"Error when checking volume mount. Return code {return_code}")
                except ValueError as e:
                    if i > 4:
                        raise e

        def test_volume_mount():
            """
            Tests whether the volume has been mounted.
            """
            with open('/foo/volume_mount_test.txt', 'w') as foo:
                foo.write('Hello')

            return_code = os.system("cat /foo/volume_mount_test.txt")
            if return_code != 0:
                raise ValueError(f"Error when checking volume mount. Return code {return_code}")

        # You can use annotations on your kubernetes pods!
        start_task = PythonOperator(
            task_id="start_task",
            python_callable=print_stuff,
            executor_config={
                "pod_override": k8s.V1Pod(metadata=k8s.V1ObjectMeta(annotations={"test": "annotation"}))
            },
        )

        # [START task_with_volume]
        volume_task = PythonOperator(
            task_id="task_with_volume",
            python_callable=test_volume_mount,
            executor_config={
                "pod_override": k8s.V1Pod(
                    spec=k8s.V1PodSpec(
                        containers=[
                            k8s.V1Container(
                                name="base",
                                volume_mounts=[
                                    k8s.V1VolumeMount(
                                        mount_path="/foo/", name="example-kubernetes-test-volume"
                                    )
                                ],
                            )
                        ],
                        volumes=[
                            k8s.V1Volume(
                                name="example-kubernetes-test-volume",
                                host_path=k8s.V1HostPathVolumeSource(path="/tmp/"),
                            )
                        ],
                    )
                ),
            },
        )
        # [END task_with_volume]

        # [START task_with_template]
        task_with_template = PythonOperator(
            task_id="task_with_template",
            python_callable=print_stuff,
            executor_config={
                "pod_template_file": os.path.join(AIRFLOW_HOME, "pod_templates/basic_template.yaml"),
                "pod_override": k8s.V1Pod(metadata=k8s.V1ObjectMeta(labels={"release": "stable"})),
            },
        )
        # [END task_with_template]

        # [START task_with_sidecar]
        sidecar_task = PythonOperator(
            task_id="task_with_sidecar",
            python_callable=test_sharedvolume_mount,
            executor_config={
                "pod_override": k8s.V1Pod(
                    spec=k8s.V1PodSpec(
                        containers=[
                            k8s.V1Container(
                                name="base",
                                volume_mounts=[
                                    k8s.V1VolumeMount(mount_path="/shared/", name="shared-empty-dir")
                                ],
                            ),
                            k8s.V1Container(
                                name="sidecar",
                                image="ubuntu",
                                args=["echo \"retrieved from mount\" > /shared/test.txt"],
                                command=["bash", "-cx"],
                                volume_mounts=[
                                    k8s.V1VolumeMount(mount_path="/shared/", name="shared-empty-dir")
                                ],
                            ),
                        ],
                        volumes=[
                            k8s.V1Volume(name="shared-empty-dir", empty_dir=k8s.V1EmptyDirVolumeSource()),
                        ],
                    )
                ),
            },
        )
        # [END task_with_sidecar]

        # Test that we can add labels to pods
        third_task = PythonOperator(
            task_id="non_root_task",
            python_callable=print_stuff,
            executor_config={
                "pod_override": k8s.V1Pod(metadata=k8s.V1ObjectMeta(labels={"release": "stable"}))
            },
        )

        other_ns_task = PythonOperator(
            task_id="other_namespace_task",
            python_callable=print_stuff,
            executor_config={
                "KubernetesExecutor": {"namespace": "test-namespace", "labels": {"release": "stable"}}
            },
        )

        start_task >> volume_task >> third_task
        start_task >> other_ns_task
        start_task >> sidecar_task
        start_task >> task_with_template
except ImportError as e:
    log.warning("Could not import DAGs in example_kubernetes_executor_config.py: %s", str(e))
    log.warning("Install kubernetes dependencies with: pip install apache-airflow['cncf.kubernetes']")