apache-airflow-providers-cncf-kubernetes 3.1.0__py3-none-any.whl → 10.10.0rc1__py3-none-any.whl

airflow/providers/cncf/kubernetes/pod_generator.py

@@ -0,0 +1,592 @@
+# 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.
+"""
+Pod generator.
+
+This module provides an interface between the previous Pod
+API and outputs a kubernetes.client.models.V1Pod.
+The advantage being that the full Kubernetes API
+is supported and no serialization need be written.
+"""
+
+from __future__ import annotations
+
+import copy
+import logging
+import os
+import re
+import warnings
+from functools import reduce
+from typing import TYPE_CHECKING
+
+from dateutil import parser
+from kubernetes.client import models as k8s
+from kubernetes.client.api_client import ApiClient
+
+from airflow.exceptions import (
+    AirflowConfigException,
+)
+from airflow.providers.cncf.kubernetes.backcompat import get_logical_date_key
+from airflow.providers.cncf.kubernetes.exceptions import PodMutationHookException, PodReconciliationError
+from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
+    POD_NAME_MAX_LENGTH,
+    add_unique_suffix,
+)
+from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_0_PLUS
+from airflow.utils import yaml
+from airflow.utils.hashlib_wrapper import md5
+from airflow.version import version as airflow_version
+
+if TYPE_CHECKING:
+    import datetime
+
+    from airflow.executors import workloads
+    from airflow.models.taskinstance import TaskInstance
+
+log = logging.getLogger(__name__)
+
+MAX_LABEL_LEN = 63
+
+
+def workload_to_command_args(workload: workloads.ExecuteTask) -> list[str]:
+    """
+    Convert a workload object to Task SDK command arguments.
+
+    :param workload: The ExecuteTask workload to convert
+    :return: List of command arguments for the Task SDK
+    """
+    ser_input = workload.model_dump_json()
+    return [
+        "python",
+        "-m",
+        "airflow.sdk.execution_time.execute_workload",
+        "--json-string",
+        ser_input,
+    ]
+
+
+def generate_pod_command_args(task_instance: TaskInstance) -> list[str]:
+    """
+    Generate command arguments for a ``TaskInstance`` to be used in a Kubernetes pod.
+
+    This function handles backwards compatibility between Airflow 2.x and 3.x:
+    - In Airflow 2.x: Uses the existing ``command_as_list()`` method
+    - In Airflow 3.x: Uses the Task SDK workload approach with serialized workload
+    """
+    if AIRFLOW_V_3_0_PLUS:
+        # In Airflow 3+, use the Task SDK workload approach
+        from airflow.executors import workloads
+
+        workload = workloads.ExecuteTask.make(task_instance)
+        return workload_to_command_args(workload)
+    # In Airflow 2.x, use the existing method
+    return task_instance.command_as_list()
+
+
+def make_safe_label_value(string: str) -> str:
+    """
+    Normalize a provided label to be of valid length and characters.
+
+    Valid label values must be 63 characters or less and must be empty or begin and
+    end with an alphanumeric character ([a-z0-9A-Z]) with dashes (-), underscores (_),
+    dots (.), and alphanumerics between.
+
+    If the label value is greater than 63 chars once made safe, or differs in any
+    way from the original value sent to this function, then we need to truncate to
+    53 chars, and append it with a unique hash.
+    """
+    safe_label = re.sub(r"^[^a-z0-9A-Z]*|[^a-zA-Z0-9_\-\.]|[^a-z0-9A-Z]*$", "", string)
+
+    if len(safe_label) > MAX_LABEL_LEN or string != safe_label:
+        safe_hash = md5(string.encode()).hexdigest()[:9]
+        safe_label = safe_label[: MAX_LABEL_LEN - len(safe_hash) - 1] + "-" + safe_hash
+
+    return safe_label
+
+
+def datetime_to_label_safe_datestring(datetime_obj: datetime.datetime) -> str:
+    """
+    Transform a datetime string to use as a label.
+
+    Kubernetes doesn't like ":" in labels, since ISO datetime format uses ":" but
+    not "_" let's
+    replace ":" with "_"
+
+    :param datetime_obj: datetime.datetime object
+    :return: ISO-like string representing the datetime
+    """
+    return datetime_obj.isoformat().replace(":", "_").replace("+", "_plus_")
+
+
+def label_safe_datestring_to_datetime(string: str) -> datetime.datetime:
+    """
+    Transform a label back to a datetime object.
+
+    Kubernetes doesn't permit ":" in labels. ISO datetime format uses ":" but not
+    "_", let's
+    replace ":" with "_"
+
+    :param string: str
+    :return: datetime.datetime object
+    """
+    return parser.parse(string.replace("_plus_", "+").replace("_", ":"))
+
+
+class PodGenerator:
+    """
+    Contains Kubernetes Airflow Worker configuration logic.
+
+    Represents a kubernetes pod and manages execution of a single pod.
+    Any configuration that is container specific gets applied to
+    the first container in the list of containers.
+
+    :param pod: The fully specified pod. Mutually exclusive with `pod_template_file`
+    :param pod_template_file: Path to YAML file. Mutually exclusive with `pod`
+    :param extract_xcom: Whether to bring up a container for xcom
+    """
+
+    def __init__(
+        self,
+        pod: k8s.V1Pod | None = None,
+        pod_template_file: str | None = None,
+        extract_xcom: bool = True,
+    ):
+        if not pod_template_file and not pod:
+            raise AirflowConfigException(
+                "Podgenerator requires either a `pod` or a `pod_template_file` argument"
+            )
+        if pod_template_file and pod:
+            raise AirflowConfigException("Cannot pass both `pod` and `pod_template_file` arguments")
+
+        if pod_template_file:
+            self.ud_pod = self.deserialize_model_file(pod_template_file)
+        else:
+            self.ud_pod = pod
+
+        # Attach sidecar
+        self.extract_xcom = extract_xcom
+
+    @staticmethod
+    def from_obj(obj) -> dict | k8s.V1Pod | None:
+        """Convert to pod from obj."""
+        if obj is None:
+            return None
+
+        k8s_legacy_object = obj.get("KubernetesExecutor", None)
+        k8s_object = obj.get("pod_override", None)
+
+        if k8s_legacy_object and k8s_object:
+            raise AirflowConfigException(
+                "Can not have both a legacy and new"
+                "executor_config object. Please delete the KubernetesExecutor"
+                "dict and only use the pod_override kubernetes.client.models.V1Pod"
+                "object."
+            )
+        if not k8s_object and not k8s_legacy_object:
+            return None
+
+        if isinstance(k8s_object, k8s.V1Pod):
+            return k8s_object
+        raise TypeError(
+            "Cannot convert a non-kubernetes.client.models.V1Pod object into a KubernetesExecutorConfig"
+        )
+
+    @staticmethod
+    def reconcile_pods(base_pod: k8s.V1Pod, client_pod: k8s.V1Pod | None) -> k8s.V1Pod:
+        """
+        Merge Kubernetes Pod objects.
+
+        :param base_pod: has the base attributes which are overwritten if they exist
+            in the client pod and remain if they do not exist in the client_pod
+        :param client_pod: the pod that the client wants to create.
+        :return: the merged pods
+
+        This can't be done recursively as certain fields are overwritten and some are concatenated.
+        """
+        if client_pod is None:
+            return base_pod
+
+        client_pod_cp = copy.deepcopy(client_pod)
+        client_pod_cp.spec = PodGenerator.reconcile_specs(base_pod.spec, client_pod_cp.spec)
+        client_pod_cp.metadata = PodGenerator.reconcile_metadata(base_pod.metadata, client_pod_cp.metadata)
+        client_pod_cp = merge_objects(base_pod, client_pod_cp)
+
+        return client_pod_cp
+
+    @staticmethod
+    def reconcile_metadata(base_meta, client_meta):
+        """
+        Merge Kubernetes Metadata objects.
+
+        :param base_meta: has the base attributes which are overwritten if they exist
+            in the client_meta and remain if they do not exist in the client_meta
+        :param client_meta: the spec that the client wants to create.
+        :return: the merged specs
+        """
+        if base_meta and not client_meta:
+            return base_meta
+        if not base_meta and client_meta:
+            return client_meta
+        if client_meta and base_meta:
+            client_meta.labels = merge_objects(base_meta.labels, client_meta.labels)
+            client_meta.annotations = merge_objects(base_meta.annotations, client_meta.annotations)
+            extend_object_field(base_meta, client_meta, "managed_fields")
+            extend_object_field(base_meta, client_meta, "finalizers")
+            extend_object_field(base_meta, client_meta, "owner_references")
+            return merge_objects(base_meta, client_meta)
+
+        return None
+
+    @staticmethod
+    def reconcile_specs(
+        base_spec: k8s.V1PodSpec | None, client_spec: k8s.V1PodSpec | None
+    ) -> k8s.V1PodSpec | None:
+        """
+        Merge Kubernetes PodSpec objects.
+
+        :param base_spec: has the base attributes which are overwritten if they exist
+            in the client_spec and remain if they do not exist in the client_spec
+        :param client_spec: the spec that the client wants to create.
+        :return: the merged specs
+        """
+        if base_spec and not client_spec:
+            return base_spec
+        if not base_spec and client_spec:
+            return client_spec
+        if client_spec and base_spec:
+            client_spec.containers = PodGenerator.reconcile_containers(
+                base_spec.containers, client_spec.containers
+            )
+            merged_spec = extend_object_field(base_spec, client_spec, "init_containers")
+            merged_spec = extend_object_field(base_spec, merged_spec, "volumes")
+            return merge_objects(base_spec, merged_spec)
+
+        return None
+
+    @staticmethod
+    def reconcile_containers(
+        base_containers: list[k8s.V1Container], client_containers: list[k8s.V1Container]
+    ) -> list[k8s.V1Container]:
+        """
+        Merge Kubernetes Container objects.
+
+        :param base_containers: has the base attributes which are overwritten if they exist
+            in the client_containers and remain if they do not exist in the client_containers
+        :param client_containers: the containers that the client wants to create.
+        :return: the merged containers
+
+        The runs recursively over the list of containers.
+        """
+        if not base_containers:
+            return client_containers
+        if not client_containers:
+            return base_containers
+
+        client_container = client_containers[0]
+        base_container = base_containers[0]
+        client_container = extend_object_field(base_container, client_container, "volume_mounts")
+        client_container = extend_object_field(base_container, client_container, "env")
+        client_container = extend_object_field(base_container, client_container, "env_from")
+        client_container = extend_object_field(base_container, client_container, "ports")
+        client_container = extend_object_field(base_container, client_container, "volume_devices")
+        client_container = merge_objects(base_container, client_container)
+
+        return [
+            client_container,
+            *PodGenerator.reconcile_containers(base_containers[1:], client_containers[1:]),
+        ]
+
+    @classmethod
+    def construct_pod(
+        cls,
+        dag_id: str,
+        task_id: str,
+        pod_id: str,
+        try_number: int,
+        kube_image: str,
+        date: datetime.datetime | None,
+        args: list[str],
+        pod_override_object: k8s.V1Pod | None,
+        base_worker_pod: k8s.V1Pod,
+        namespace: str,
+        scheduler_job_id: str,
+        run_id: str | None = None,
+        map_index: int = -1,
+        *,
+        with_mutation_hook: bool = False,
+    ) -> k8s.V1Pod:
+        """
+        Create a Pod.
+
+        Construct a pod by gathering and consolidating the configuration from 3 places:
+            - airflow.cfg
+            - executor_config
+            - dynamic arguments
+        """
+        if len(pod_id) > POD_NAME_MAX_LENGTH:
+            warnings.warn(
+                f"pod_id supplied is longer than {POD_NAME_MAX_LENGTH} characters; "
+                f"truncating and adding unique suffix.",
+                UserWarning,
+                stacklevel=2,
+            )
+            pod_id = add_unique_suffix(name=pod_id, max_len=POD_NAME_MAX_LENGTH)
+        try:
+            image = pod_override_object.spec.containers[0].image  # type: ignore
+            if not image:
+                image = kube_image
+        except Exception:
+            image = kube_image
+
+        annotations = {
+            "dag_id": dag_id,
+            "task_id": task_id,
+            "try_number": str(try_number),
+        }
+        if map_index >= 0:
+            annotations["map_index"] = str(map_index)
+        if date:
+            annotations[get_logical_date_key()] = date.isoformat()
+        if run_id:
+            annotations["run_id"] = run_id
+
+        main_container = k8s.V1Container(
+            name="base",
+            args=args,
+            image=image,
+            env=[
+                k8s.V1EnvVar(name="AIRFLOW_IS_K8S_EXECUTOR_POD", value="True"),
+            ],
+        )
+        dynamic_pod = k8s.V1Pod(
+            metadata=k8s.V1ObjectMeta(
+                namespace=namespace,
+                annotations=annotations,
+                name=pod_id,
+                labels=cls.build_labels_for_k8s_executor_pod(
+                    dag_id=dag_id,
+                    task_id=task_id,
+                    try_number=try_number,
+                    airflow_worker=scheduler_job_id,
+                    map_index=map_index,
+                    logical_date=date,
+                    run_id=run_id,
+                ),
+            ),
+        )
+
+        podspec = k8s.V1PodSpec(
+            containers=[main_container],
+        )
+
+        dynamic_pod.spec = podspec
+
+        # Reconcile the pods starting with the first chronologically,
+        # Pod from the pod_template_File -> Pod from the K8s executor -> Pod from executor_config arg
+        pod_list = [base_worker_pod, dynamic_pod, pod_override_object]
+
+        try:
+            pod = reduce(PodGenerator.reconcile_pods, pod_list)
+        except Exception as e:
+            raise PodReconciliationError from e
+
+        if with_mutation_hook:
+            from airflow.settings import pod_mutation_hook
+
+            try:
+                pod_mutation_hook(pod)
+            except Exception as e:
+                raise PodMutationHookException from e
+
+        return pod
+
+    @classmethod
+    def build_selector_for_k8s_executor_pod(
+        cls,
+        *,
+        dag_id,
+        task_id,
+        try_number,
+        map_index=None,
+        logical_date=None,
+        run_id=None,
+        airflow_worker=None,
+        include_version=False,
+    ):
+        """
+        Generate selector for kubernetes executor pod.
+
+        :meta private:
+        """
+        labels = cls.build_labels_for_k8s_executor_pod(
+            dag_id=dag_id,
+            task_id=task_id,
+            try_number=try_number,
+            map_index=map_index,
+            logical_date=logical_date,
+            run_id=run_id,
+            airflow_worker=airflow_worker,
+            include_version=include_version,
+        )
+        label_strings = [f"{label_id}={label}" for label_id, label in sorted(labels.items())]
+        selector = ",".join(label_strings)
+        if not airflow_worker:  # this filters out KPO pods even when we don't know the scheduler job id
+            selector += ",airflow-worker"
+        return selector
+
+    @classmethod
+    def build_labels_for_k8s_executor_pod(
+        cls,
+        *,
+        dag_id,
+        task_id,
+        try_number,
+        airflow_worker=None,
+        map_index=None,
+        logical_date=None,
+        run_id=None,
+        include_version=True,
+    ):
+        """
+        Generate labels for kubernetes executor pod.
+
+        :meta private:
+        """
+        labels = {
+            "dag_id": make_safe_label_value(dag_id),
+            "task_id": make_safe_label_value(task_id),
+            "try_number": str(try_number),
+            "kubernetes_executor": "True",
+        }
+        if include_version:
+            labels["airflow_version"] = airflow_version.replace("+", "-")
+        if airflow_worker is not None:
+            labels["airflow-worker"] = make_safe_label_value(str(airflow_worker))
+        if map_index is not None and map_index >= 0:
+            labels["map_index"] = str(map_index)
+        if logical_date:
+            labels[get_logical_date_key()] = datetime_to_label_safe_datestring(logical_date)
+        if run_id:
+            labels["run_id"] = make_safe_label_value(run_id)
+        return labels
+
+    @staticmethod
+    def serialize_pod(pod: k8s.V1Pod) -> dict:
+        """
+        Convert a k8s.V1Pod into a json serializable dictionary.
+
+        :param pod: k8s.V1Pod object
+        :return: Serialized version of the pod returned as dict
+        """
+        api_client = ApiClient()
+        return api_client.sanitize_for_serialization(pod)
+
+    @staticmethod
+    def deserialize_model_file(path: str) -> k8s.V1Pod:
+        """
+        Generate a Pod from a file.
+
+        :param path: Path to the file
+        :return: a kubernetes.client.models.V1Pod
+        """
+        if os.path.exists(path):
+            with open(path) as stream:
+                pod = yaml.safe_load(stream)
+        else:
+            pod = None
+            log.warning("Model file %s does not exist", path)
+
+        return PodGenerator.deserialize_model_dict(pod)
+
+    @staticmethod
+    def deserialize_model_dict(pod_dict: dict | None) -> k8s.V1Pod:
+        """
+        Deserializes a Python dictionary to k8s.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.
+
+        :param pod_dict: Serialized dict of k8s.V1Pod object
+        :return: De-serialized k8s.V1Pod
+        """
+        api_client = ApiClient()
+        return api_client._ApiClient__deserialize_model(pod_dict, k8s.V1Pod)
+
+
+def merge_objects(base_obj, client_obj):
+    """
+    Merge objects.
+
+    :param base_obj: has the base attributes which are overwritten if they exist
+        in the client_obj and remain if they do not exist in the client_obj
+    :param client_obj: the object that the client wants to create.
+    :return: the merged objects
+    """
+    if not base_obj:
+        return client_obj
+    if not client_obj:
+        return base_obj
+
+    client_obj_cp = copy.deepcopy(client_obj)
+
+    if isinstance(base_obj, dict) and isinstance(client_obj_cp, dict):
+        base_obj_cp = copy.deepcopy(base_obj)
+        base_obj_cp.update(client_obj_cp)
+        return base_obj_cp
+
+    for base_key in base_obj.to_dict():
+        base_val = getattr(base_obj, base_key, None)
+        if not getattr(client_obj, base_key, None) and base_val is not None:
+            if not isinstance(client_obj_cp, dict):
+                setattr(client_obj_cp, base_key, base_val)
+            else:
+                client_obj_cp[base_key] = base_val
+    return client_obj_cp
+
+
+def extend_object_field(base_obj, client_obj, field_name):
+    """
+    Add field values to existing objects.
+
+    :param base_obj: an object which has a property `field_name` that is a list
+    :param client_obj: an object which has a property `field_name` that is a list.
+        A copy of this object is returned with `field_name` modified
+    :param field_name: the name of the list field
+    :return: the client_obj with the property `field_name` being the two properties appended
+    """
+    client_obj_cp = copy.deepcopy(client_obj)
+    base_obj_field = getattr(base_obj, field_name, None)
+    client_obj_field = getattr(client_obj, field_name, None)
+
+    if (not isinstance(base_obj_field, list) and base_obj_field is not None) or (
+        not isinstance(client_obj_field, list) and client_obj_field is not None
+    ):
+        raise ValueError(
+            f"The chosen field must be a list. Got {type(base_obj_field)} base_object_field "
+            f"and {type(client_obj_field)} client_object_field."
+        )
+
+    if not base_obj_field:
+        return client_obj_cp
+    if not client_obj_field:
+        setattr(client_obj_cp, field_name, base_obj_field)
+        return client_obj_cp
+
+    appended_fields = base_obj_field + client_obj_field
+    setattr(client_obj_cp, field_name, appended_fields)
+    return client_obj_cp

airflow/providers/cncf/kubernetes/pod_template_file_examples/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/cncf/kubernetes/pod_template_file_examples/dags_in_image_template.yaml

@@ -0,0 +1,68 @@
+# 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 file is used for documentation purposes. Example can be found at docs/executor/kubernetes.rst
+#
+# [START template_with_dags_in_image]
+---
+apiVersion: v1
+kind: Pod
+metadata:
+  name: placeholder-name
+spec:
+  containers:
+    - env:
+        - name: AIRFLOW__CORE__EXECUTOR
+          value: LocalExecutor
+        # Hard Coded Airflow Envs
+        - name: AIRFLOW__CORE__FERNET_KEY
+          valueFrom:
+            secretKeyRef:
+              name: RELEASE-NAME-fernet-key
+              key: fernet-key
+        - name: AIRFLOW__DATABASE__SQL_ALCHEMY_CONN
+          valueFrom:
+            secretKeyRef:
+              name: RELEASE-NAME-airflow-metadata
+              key: connection
+        - name: AIRFLOW_CONN_AIRFLOW_DB
+          valueFrom:
+            secretKeyRef:
+              name: RELEASE-NAME-airflow-metadata
+              key: connection
+      image: dummy_image
+      imagePullPolicy: IfNotPresent
+      name: base
+      volumeMounts:
+        - mountPath: "/opt/airflow/logs"
+          name: airflow-logs
+        - mountPath: /opt/airflow/airflow.cfg
+          name: airflow-config
+          readOnly: true
+          subPath: airflow.cfg
+  restartPolicy: Never
+  securityContext:
+    runAsUser: 50000
+    fsGroup: 50000
+  serviceAccountName: "RELEASE-NAME-worker-serviceaccount"
+  volumes:
+    - emptyDir: {}
+      name: airflow-logs
+    - configMap:
+        name: RELEASE-NAME-airflow-config
+      name: airflow-config
+# [END template_with_dags_in_image]