[AIRFLOW-5413] Allow K8S worker pod to be configured from JSON/YAML file (#6230)

* [AIRFLOW-5413] enable pod config from file

* Update airflow/kubernetes/pod_generator.py

Co-Authored-By: Ash Berlin-Taylor <[email protected]>

* Update airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py

Co-Authored-By: Ash Berlin-Taylor <[email protected]>

Co-authored-by: Ash Berlin-Taylor <[email protected]>
(cherry picked from commit 967930c)

Author: 2 people

airflow/config_templates/config.yml

@@ -1736,6 +1736,13 @@
       type: string
       example: ~
       default: ""
+    - name: pod_template_file
+      description: |
+        Path to the YAML pod file. If set, all other kubernetes-related fields are ignored.
+      version_added: ~
+      type: string
+      example: ~
+      default: ""
     - name: worker_container_tag
       description: ~
       version_added: ~

airflow/config_templates/default_airflow.cfg

@@ -801,6 +801,9 @@ verify_certs = True
 [kubernetes]
 # The repository, tag and imagePullPolicy of the Kubernetes Image for the Worker to Run
 worker_container_repository =
+
+# Path to the YAML pod file. If set, all other kubernetes-related fields are ignored.
+pod_template_file =
 worker_container_tag =
 worker_container_image_pull_policy = IfNotPresent
 

airflow/contrib/operators/kubernetes_pod_operator.py

@@ -130,8 +130,18 @@ class KubernetesPodOperator(BaseOperator):  # pylint: disable=too-many-instance-
     :type schedulername: str
     :param full_pod_spec: The complete podSpec
     :type full_pod_spec: kubernetes.client.models.V1Pod
+    :param init_containers: init container for the launched Pod
+    :type init_containers: list[kubernetes.client.models.V1Container]
+    :param log_events_on_failure: Log the pod's events if a failure occurs
+    :type log_events_on_failure: bool
+    :param do_xcom_push: If True, the content of the file
+        /airflow/xcom/return.json in the container will also be pushed to an
+        XCom when the container completes.
+    :type do_xcom_push: bool
+    :param pod_template_file: path to pod template file
+    :type pod_template_file: str
     """
-    template_fields = ('cmds', 'arguments', 'env_vars', 'config_file')
+    template_fields = ('cmds', 'arguments', 'env_vars', 'config_file', 'pod_template_file')
 
     @apply_defaults
     def __init__(self,  # pylint: disable=too-many-arguments,too-many-locals
@@ -215,8 +225,8 @@ def __init__(self,  # pylint: disable=too-many-arguments,too-many-locals
         self.full_pod_spec = full_pod_spec
         self.init_containers = init_containers or []
         self.log_events_on_failure = log_events_on_failure
-        self.priority_class_name = priority_class_name
         self.pod_template_file = pod_template_file
+        self.priority_class_name = priority_class_name
         self.name = self._set_name(name)
 
     @staticmethod
@@ -348,6 +358,7 @@ def create_new_pod_for_operator(self, labels, launcher):
             init_containers=self.init_containers,
             restart_policy='Never',
             schedulername=self.schedulername,
+            pod_template_file=self.pod_template_file,
             priority_class_name=self.priority_class_name,
             pod=self.full_pod_spec,
         ).gen_pod()

airflow/example_dags/example_kubernetes_executor_config.py

@@ -83,14 +83,14 @@ def test_volume_mount():
         }
     )
 
-    # Test that we can run tasks as a normal user
+    # Test that we can add labels to pods
     third_task = PythonOperator(
         task_id="non_root_task",
         python_callable=print_stuff,
         executor_config={
             "KubernetesExecutor": {
-                "securityContext": {
-                    "runAsUser": 1000
+                "labels": {
+                    "release": "stable"
                 }
             }
         }

airflow/executors/kubernetes_executor.py

@@ -20,7 +20,6 @@
 import multiprocessing
 import time
 from queue import Empty
-from uuid import uuid4
 
 import kubernetes
 from dateutil import parser
@@ -72,6 +71,9 @@ def __init__(self):
         )
         self.kube_node_selectors = configuration_dict.get('kubernetes_node_selectors', {})
         self.kube_annotations = configuration_dict.get('kubernetes_annotations', {}) or None
+        self.pod_template_file = conf.get(self.kubernetes_section, 'pod_template_file',
+                                          fallback=None)
+
         self.kube_labels = configuration_dict.get('kubernetes_labels', {})
         self.delete_worker_pods = conf.getboolean(
             self.kubernetes_section, 'delete_worker_pods')
@@ -220,6 +222,8 @@ def _get_security_context_val(self, scontext):
             return int(val)
 
     def _validate(self):
+        if self.pod_template_file:
+            return
         # TODO: use XOR for dags_volume_claim and git_dags_folder_mount_point
         if not self.dags_volume_claim \
             and not self.dags_volume_host \
@@ -498,10 +502,7 @@ def _create_pod_id(dag_id, task_id):
             dag_id)
         safe_task_id = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
             task_id)
-        safe_uuid = AirflowKubernetesScheduler._strip_unsafe_kubernetes_special_chars(
-            uuid4().hex)
-        return AirflowKubernetesScheduler._make_safe_pod_id(safe_dag_id, safe_task_id,
-                                                            safe_uuid)
+        return safe_dag_id + safe_task_id
 
     @staticmethod
     def _label_safe_datestring_to_datetime(string):

airflow/kubernetes/pod_generator.py

@@ -24,21 +24,35 @@
 import copy
 import hashlib
 import re
+try:
+    from inspect import signature
+except ImportError:
+    # Python 2.7
+    from funcsigs import signature  # type: ignore
+import os
 import uuid
+from functools import reduce
 
 import kubernetes.client.models as k8s
+import yaml
+from kubernetes.client.api_client import ApiClient
 
+from airflow.exceptions import AirflowConfigException
 from airflow.version import version as airflow_version
 
 MAX_LABEL_LEN = 63
 
 MAX_POD_ID_LEN = 253
 
 
-class PodDefaults:
+class PodDefaults(object):
     """
     Static defaults for the PodGenerator
     """
+
+    def __init__(self):
+        pass
+
     XCOM_MOUNT_PATH = '/airflow/xcom'
     SIDECAR_CONTAINER_NAME = 'airflow-xcom-sidecar'
     XCOM_CMD = 'trap "exit 0" INT; while true; do sleep 30; done;'
@@ -82,7 +96,7 @@ def make_safe_label_value(string):
     return safe_label
 
 
-class PodGenerator:
+class PodGenerator(object):
     """
     Contains Kubernetes Airflow Worker configuration logic
 
@@ -147,9 +161,11 @@ class PodGenerator:
     :param dnspolicy: Specify a dnspolicy for the pod
     :type dnspolicy: str
     :param schedulername: Specify a schedulername for the pod
-    :type schedulername: str
-    :param pod: The fully specified pod.
-    :type pod: kubernetes.client.models.V1Pod
+    :type schedulername: Optional[str]
+    :param pod: The fully specified pod. Mutually exclusive with `path_or_string`
+    :type pod: Optional[kubernetes.client.models.V1Pod]
+    :param pod_template_file: Path to YAML file. Mutually exclusive with `pod`
+    :type pod_template_file: Optional[str]
     :param extract_xcom: Whether to bring up a container for xcom
     :type extract_xcom: bool
     """
@@ -167,8 +183,8 @@ def __init__(
         node_selectors=None,
         ports=None,
         volumes=None,
-        image_pull_policy='IfNotPresent',
-        restart_policy='Never',
+        image_pull_policy=None,
+        restart_policy=None,
         image_pull_secrets=None,
         init_containers=None,
         service_account_name=None,
@@ -183,9 +199,16 @@ def __init__(
         schedulername=None,
         priority_class_name=None,
         pod=None,
+        pod_template_file=None,
         extract_xcom=False,
     ):
-        self.ud_pod = pod
+        self.validate_pod_generator_args(locals())
+
+        if pod_template_file:
+            self.ud_pod = self.deserialize_model_file(pod_template_file)
+        else:
+            self.ud_pod = pod
+
         self.pod = k8s.V1Pod()
         self.pod.api_version = 'v1'
         self.pod.kind = 'Pod'
@@ -348,37 +371,7 @@ def from_obj(obj):
                 'iam.cloud.google.com/service-account': gcp_service_account_key
             })
 
-        pod_spec_generator = PodGenerator(
-            image=namespaced.get('image'),
-            envs=namespaced.get('env'),
-            cmds=namespaced.get('cmds'),
-            args=namespaced.get('args'),
-            labels=namespaced.get('labels'),
-            node_selectors=namespaced.get('node_selectors'),
-            name=namespaced.get('name'),
-            ports=namespaced.get('ports'),
-            volumes=namespaced.get('volumes'),
-            volume_mounts=namespaced.get('volume_mounts'),
-            namespace=namespaced.get('namespace'),
-            image_pull_policy=namespaced.get('image_pull_policy'),
-            restart_policy=namespaced.get('restart_policy'),
-            image_pull_secrets=namespaced.get('image_pull_secrets'),
-            init_containers=namespaced.get('init_containers'),
-            service_account_name=namespaced.get('service_account_name'),
-            resources=resources,
-            annotations=namespaced.get('annotations'),
-            affinity=namespaced.get('affinity'),
-            hostnetwork=namespaced.get('hostnetwork'),
-            tolerations=namespaced.get('tolerations'),
-            security_context=namespaced.get('security_context'),
-            configmaps=namespaced.get('configmaps'),
-            dnspolicy=namespaced.get('dnspolicy'),
-            schedulername=namespaced.get('schedulername'),
-            pod=namespaced.get('pod'),
-            extract_xcom=namespaced.get('extract_xcom'),
-        )
-
-        return pod_spec_generator.gen_pod()
+        return PodGenerator(**namespaced).gen_pod()
 
     @staticmethod
     def reconcile_pods(base_pod, client_pod):
@@ -495,12 +488,73 @@ def construct_pod(
             name=pod_id
         ).gen_pod()
 
-        # Reconcile the pod generated by the Operator and the Pod
-        # generated by the .cfg file
-        pod_with_executor_config = PodGenerator.reconcile_pods(worker_config,
-                                                               kube_executor_config)
-        # Reconcile that pod with the dynamic fields.
-        return PodGenerator.reconcile_pods(pod_with_executor_config, dynamic_pod)
+        # Reconcile the pods starting with the first chronologically,
+        # Pod from the airflow.cfg -> Pod from executor_config arg -> Pod from the K8s executor
+        pod_list = [worker_config, kube_executor_config, dynamic_pod]
+
+        return reduce(PodGenerator.reconcile_pods, pod_list)
+
+    @staticmethod
+    def deserialize_model_file(path):
+        """
+        :param path: Path to the file
+        :return: a kubernetes.client.models.V1Pod
+
+        Unfortunately we need access to the private method
+        ``_ApiClient__deserialize_model`` from the kubernetes client.
+        This issue is tracked here; https://github.com/kubernetes-client/python/issues/977.
+        """
+        api_client = ApiClient()
+        if os.path.exists(path):
+            with open(path) as stream:
+                pod = yaml.safe_load(stream)
+        else:
+            pod = yaml.safe_load(path)
+
+        # pylint: disable=protected-access
+        return api_client._ApiClient__deserialize_model(pod, k8s.V1Pod)
+
+    @staticmethod
+    def validate_pod_generator_args(given_args):
+        """
+        :param given_args: The arguments passed to the PodGenerator constructor.
+        :type given_args: dict
+        :return: None
+
+        Validate that if `pod` or `pod_template_file` are set that the user is not attempting
+        to configure the pod with the other arguments.
+        """
+        pod_args = list(signature(PodGenerator).parameters.items())
+
+        def predicate(k, v):
+            """
+            :param k: an arg to PodGenerator
+            :type k: string
+            :param v: the parameter of the given arg
+            :type v: inspect.Parameter
+            :return: bool
+
+            returns True if the PodGenerator argument has no default arguments
+            or the default argument is None, and it is not one of the listed field
+            in `non_empty_fields`.
+            """
+            non_empty_fields = {
+                'pod', 'pod_template_file', 'extract_xcom', 'service_account_name', 'image_pull_policy',
+                'restart_policy'
+            }
+
+            return (v.default is None or v.default is v.empty) and k not in non_empty_fields
+
+        args_without_defaults = {k: given_args[k] for k, v in pod_args if predicate(k, v) and given_args[k]}
+
+        if given_args['pod'] and given_args['pod_template_file']:
+            raise AirflowConfigException("Cannot pass both `pod` and `pod_template_file` arguments")
+        if args_without_defaults and (given_args['pod'] or given_args['pod_template_file']):
+            raise AirflowConfigException(
+                "Cannot configure pod and pass either `pod` or `pod_template_file`. Fields {} passed.".format(
+                    list(args_without_defaults.keys())
+                )
+            )
 
 
 def merge_objects(base_obj, client_obj):

airflow/kubernetes/worker_configuration.py

@@ -28,7 +28,12 @@
 
 
 class WorkerConfiguration(LoggingMixin):
-    """Contains Kubernetes Airflow Worker configuration logic"""
+    """
+    Contains Kubernetes Airflow Worker configuration logic
+
+    :param kube_config: the kubernetes configuration from airflow.cfg
+    :type kube_config: airflow.executors.kubernetes_executor.KubeConfig
+    """
 
     dags_volume_name = 'airflow-dags'
     logs_volume_name = 'airflow-logs'
@@ -424,9 +429,12 @@ def generate_dag_volume_mount_path(self):
 
     def as_pod(self):
         """Creates POD."""
-        pod_generator = PodGenerator(
+        if self.kube_config.pod_template_file:
+            return PodGenerator(pod_template_file=self.kube_config.pod_template_file).gen_pod()
+
+        pod = PodGenerator(
             image=self.kube_config.kube_image,
-            image_pull_policy=self.kube_config.kube_image_pull_policy,
+            image_pull_policy=self.kube_config.kube_image_pull_policy or 'IfNotPresent',
             image_pull_secrets=self.kube_config.image_pull_secrets,
             volumes=self._get_volumes(),
             volume_mounts=self._get_volume_mounts(),
@@ -436,10 +444,10 @@ def as_pod(self):
             tolerations=self.kube_config.kube_tolerations,
             envs=self._get_environment(),
             node_selectors=self.kube_config.kube_node_selectors,
-            service_account_name=self.kube_config.worker_service_account_name,
-        )
+            service_account_name=self.kube_config.worker_service_account_name or 'default',
+            restart_policy='Never'
+        ).gen_pod()
 
-        pod = pod_generator.gen_pod()
         pod.spec.containers[0].env_from = pod.spec.containers[0].env_from or []
         pod.spec.containers[0].env_from.extend(self._get_env_from())
         pod.spec.security_context = self._get_security_context()