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

airflow/providers/cncf/kubernetes/k8s_model.py

@@ -0,0 +1,62 @@
+# 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.
+"""Classes for interacting with Kubernetes API."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from functools import reduce
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from kubernetes.client import models as k8s
+
+
+class K8SModel(ABC):
+    """
+    Airflow Kubernetes models are here for backwards compatibility reasons only.
+
+    Ideally clients should use the kubernetes API
+    and the process of
+
+        client input -> Airflow k8s models -> k8s models
+
+    can be avoided. All of these models implement the
+    `attach_to_pod` method so that they integrate with the kubernetes client.
+    """
+
+    @abstractmethod
+    def attach_to_pod(self, pod: k8s.V1Pod) -> k8s.V1Pod:
+        """
+        Attaches to pod.
+
+        :param pod: A pod to attach this Kubernetes object to
+        :return: The pod with the object attached
+        """
+
+
+def append_to_pod(pod: k8s.V1Pod, k8s_objects: list[K8SModel] | None):
+    """
+    Attach additional specs to an existing pod object.
+
+    :param pod: A pod to attach a list of Kubernetes objects to
+    :param k8s_objects: a potential None list of K8SModels
+    :return: pod with the objects attached if they exist
+    """
+    if not k8s_objects:
+        return pod
+    return reduce(lambda p, o: o.attach_to_pod(p), k8s_objects, pod)

airflow/providers/cncf/kubernetes/kube_client.py

@@ -0,0 +1,156 @@
+# 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.
+"""Client for kubernetes communication."""
+
+from __future__ import annotations
+
+import logging
+
+import urllib3.util
+
+from airflow.configuration import conf
+
+log = logging.getLogger(__name__)
+
+try:
+    from kubernetes import client, config
+    from kubernetes.client import Configuration
+    from kubernetes.client.rest import ApiException
+
+    has_kubernetes = True
+
+    def _get_default_configuration() -> Configuration:
+        if hasattr(Configuration, "get_default_copy"):
+            return Configuration.get_default_copy()
+        return Configuration()
+
+    def _disable_verify_ssl() -> None:
+        configuration = _get_default_configuration()
+        configuration.verify_ssl = False
+        Configuration.set_default(configuration)
+
+except ImportError as e:
+    # We need an exception class to be able to use it in ``except`` elsewhere
+    # in the code base
+    ApiException = BaseException
+    has_kubernetes = False
+    _import_err = e
+
+
+def _enable_tcp_keepalive() -> None:
+    """
+    Enable TCP keepalive mechanism.
+
+    This prevents urllib3 connection to hang indefinitely when idle connection
+    is time-outed on services like cloud load balancers or firewalls.
+
+    See https://github.com/apache/airflow/pull/11406 for detailed explanation.
+
+    Please ping @michalmisiewicz or @dimberman in the PR if you want to modify this function.
+    """
+    import socket
+
+    from urllib3.connection import HTTPConnection, HTTPSConnection
+
+    tcp_keep_idle = conf.getint("kubernetes_executor", "tcp_keep_idle")
+    tcp_keep_intvl = conf.getint("kubernetes_executor", "tcp_keep_intvl")
+    tcp_keep_cnt = conf.getint("kubernetes_executor", "tcp_keep_cnt")
+
+    socket_options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
+
+    if hasattr(socket, "TCP_KEEPIDLE"):
+        socket_options.append((socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, tcp_keep_idle))
+    else:
+        log.debug("Unable to set TCP_KEEPIDLE on this platform")
+
+    if hasattr(socket, "TCP_KEEPINTVL"):
+        socket_options.append((socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, tcp_keep_intvl))
+    else:
+        log.debug("Unable to set TCP_KEEPINTVL on this platform")
+
+    if hasattr(socket, "TCP_KEEPCNT"):
+        socket_options.append((socket.IPPROTO_TCP, socket.TCP_KEEPCNT, tcp_keep_cnt))
+    else:
+        log.debug("Unable to set TCP_KEEPCNT on this platform")
+
+    # Cast both the default options and our socket options
+    socket_options_cast: list[tuple[int, int, int | bytes]] = [
+        (level, opt, val) for level, opt, val in socket_options
+    ]
+    default_options_cast: list[tuple[int, int, int | bytes]] = [
+        (level, opt, val) for level, opt, val in HTTPSConnection.default_socket_options
+    ]
+
+    # Then use the cast versions for both HTTPS and HTTP
+    HTTPSConnection.default_socket_options = default_options_cast + socket_options_cast
+    HTTPConnection.default_socket_options = default_options_cast + socket_options_cast
+
+
+def get_kube_client(
+    in_cluster: bool | None = None,
+    cluster_context: str | None = None,
+    config_file: str | None = None,
+) -> client.CoreV1Api:
+    """
+    Retrieve Kubernetes client.
+
+    :param in_cluster: whether we are in cluster
+    :param cluster_context: context of the cluster
+    :param config_file: configuration file
+    :return: kubernetes client
+    """
+    if in_cluster is None:
+        in_cluster = conf.getboolean("kubernetes_executor", "in_cluster")
+    if not has_kubernetes:
+        raise _import_err
+
+    if conf.getboolean("kubernetes_executor", "enable_tcp_keepalive"):
+        _enable_tcp_keepalive()
+
+    configuration = _get_default_configuration()
+    api_client_retry_configuration = conf.getjson(
+        "kubernetes_executor", "api_client_retry_configuration", fallback={}
+    )
+
+    if not conf.getboolean("kubernetes_executor", "verify_ssl"):
+        _disable_verify_ssl()
+
+    if isinstance(api_client_retry_configuration, dict):
+        configuration.retries = urllib3.util.Retry(**api_client_retry_configuration)
+    else:
+        raise ValueError("api_client_retry_configuration should be a dictionary")
+
+    if in_cluster:
+        config.load_incluster_config(client_configuration=configuration)
+    else:
+        if cluster_context is None:
+            cluster_context = conf.get("kubernetes_executor", "cluster_context", fallback=None)
+        if config_file is None:
+            config_file = conf.get("kubernetes_executor", "config_file", fallback=None)
+        config.load_kube_config(
+            config_file=config_file, context=cluster_context, client_configuration=configuration
+        )
+
+    if not conf.getboolean("kubernetes_executor", "verify_ssl"):
+        configuration.verify_ssl = False
+
+    ssl_ca_cert = conf.get("kubernetes_executor", "ssl_ca_cert")
+    if ssl_ca_cert:
+        configuration.ssl_ca_cert = ssl_ca_cert
+
+    api_client = client.ApiClient(configuration=configuration)
+    return client.CoreV1Api(api_client)

airflow/providers/cncf/kubernetes/kube_config.py

@@ -0,0 +1,125 @@
+# 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.
+from __future__ import annotations
+
+import warnings
+
+from airflow.configuration import conf
+from airflow.exceptions import AirflowConfigException, AirflowProviderDeprecationWarning
+from airflow.settings import AIRFLOW_HOME
+
+
+class KubeConfig:
+    """Configuration for Kubernetes."""
+
+    core_section = "core"
+    kubernetes_section = "kubernetes_executor"
+    logging_section = "logging"
+
+    def __init__(self):
+        configuration_dict = conf.as_dict(display_sensitive=True)
+        self.core_configuration = configuration_dict[self.core_section]
+        self.airflow_home = AIRFLOW_HOME
+        self.dags_folder = conf.get(self.core_section, "dags_folder")
+        self.parallelism = conf.getint(self.core_section, "parallelism")
+        self.pod_template_file = conf.get(self.kubernetes_section, "pod_template_file", fallback=None)
+
+        self.delete_worker_pods = conf.getboolean(self.kubernetes_section, "delete_worker_pods")
+        self.delete_worker_pods_on_failure = conf.getboolean(
+            self.kubernetes_section, "delete_worker_pods_on_failure"
+        )
+        self.worker_pod_pending_fatal_container_state_reasons = []
+        if conf.get(self.kubernetes_section, "worker_pod_pending_fatal_container_state_reasons", fallback=""):
+            self.worker_pod_pending_fatal_container_state_reasons = [
+                r.strip()
+                for r in conf.get(
+                    self.kubernetes_section, "worker_pod_pending_fatal_container_state_reasons"
+                ).split(",")
+            ]
+
+        self.worker_pods_creation_batch_size = conf.getint(
+            self.kubernetes_section, "worker_pods_creation_batch_size"
+        )
+        self.worker_container_repository = conf.get(self.kubernetes_section, "worker_container_repository")
+        if self.worker_container_repository:
+            warnings.warn(
+                "Configuration 'worker_container_repository' is deprecated. "
+                "Use 'pod_template_file' to specify the container image repository instead.",
+                AirflowProviderDeprecationWarning,
+                stacklevel=2,
+            )
+        self.worker_container_tag = conf.get(self.kubernetes_section, "worker_container_tag")
+        if self.worker_container_tag:
+            warnings.warn(
+                "Configuration 'worker_container_tag' is deprecated. "
+                "Use 'pod_template_file' to specify the container image tag instead.",
+                AirflowProviderDeprecationWarning,
+                stacklevel=2,
+            )
+        if self.worker_container_repository and self.worker_container_tag:
+            self.kube_image = f"{self.worker_container_repository}:{self.worker_container_tag}"
+        else:
+            self.kube_image = None
+
+        # The Kubernetes Namespace in which the Scheduler and Webserver reside. Note
+        # that if your
+        # cluster has RBAC enabled, your scheduler may need service account permissions to
+        # create, watch, get, and delete pods in this namespace.
+        self.kube_namespace = conf.get(self.kubernetes_section, "namespace")
+        if self.kube_namespace and self.kube_namespace != "default":
+            warnings.warn(
+                "Configuration 'namespace' is deprecated. "
+                "Use 'pod_template_file' to specify the namespace instead.",
+                AirflowProviderDeprecationWarning,
+                stacklevel=2,
+            )
+        self.multi_namespace_mode = conf.getboolean(self.kubernetes_section, "multi_namespace_mode")
+        if self.multi_namespace_mode and conf.get(
+            self.kubernetes_section, "multi_namespace_mode_namespace_list"
+        ):
+            self.multi_namespace_mode_namespace_list = conf.get(
+                self.kubernetes_section, "multi_namespace_mode_namespace_list"
+            ).split(",")
+        else:
+            self.multi_namespace_mode_namespace_list = None
+        # The Kubernetes Namespace in which pods will be created by the executor. Note
+        # that if your
+        # cluster has RBAC enabled, your workers may need service account permissions to
+        # interact with cluster components.
+        self.executor_namespace = conf.get(self.kubernetes_section, "namespace")
+
+        self.kube_client_request_args = conf.getjson(
+            self.kubernetes_section, "kube_client_request_args", fallback={}
+        )
+        if not isinstance(self.kube_client_request_args, dict):
+            raise AirflowConfigException(
+                f"[{self.kubernetes_section}] 'kube_client_request_args' expected a JSON dict, got "
+                + type(self.kube_client_request_args).__name__
+            )
+        if self.kube_client_request_args:
+            if "_request_timeout" in self.kube_client_request_args and isinstance(
+                self.kube_client_request_args["_request_timeout"], list
+            ):
+                self.kube_client_request_args["_request_timeout"] = tuple(
+                    self.kube_client_request_args["_request_timeout"]
+                )
+        self.delete_option_kwargs = conf.getjson(self.kubernetes_section, "delete_option_kwargs", fallback={})
+        if not isinstance(self.delete_option_kwargs, dict):
+            raise AirflowConfigException(
+                f"[{self.kubernetes_section}] 'delete_option_kwargs' expected a JSON dict, got "
+                + type(self.delete_option_kwargs).__name__
+            )

airflow/providers/cncf/kubernetes/kubernetes_executor_templates/__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/kubernetes_executor_templates/basic_template.yaml

@@ -0,0 +1,79 @@
+# 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.
+---
+kind: Pod
+apiVersion: v1
+metadata:
+  name: placeholder-name-dont-delete
+  namespace: placeholder-name-dont-delete
+  labels:
+    mylabel: foo
+spec:
+  containers:
+    - name: base
+      image: placeholder-name-dont-delete
+      env:
+        - name: AIRFLOW__CORE__EXECUTOR
+          value: LocalExecutor
+        - name: AIRFLOW_HOME
+          value: /opt/airflow
+        - name: AIRFLOW__CORE__DAGS_FOLDER
+          value: /opt/airflow/dags
+        - name: AIRFLOW__CORE__FERNET_KEY
+          valueFrom:
+            secretKeyRef:
+              name: airflow-fernet-key
+              key: fernet-key
+        - name: AIRFLOW__DATABASE__SQL_ALCHEMY_CONN
+          valueFrom:
+            secretKeyRef:
+              name: airflow-airflow-metadata
+              key: connection
+        - name: foo
+          value: bar
+      resources: {}
+      volumeMounts:
+        - name: airflow-logs
+          mountPath: /opt/airflow/logs
+        - name: airflow-config
+          readOnly: true
+          mountPath: /opt/airflow/airflow.cfg
+          subPath: airflow.cfg
+        - name: airflow-config
+          readOnly: true
+          mountPath: /opt/airflow/config/airflow_local_settings.py
+          subPath: airflow_local_settings.py
+      terminationMessagePath: /dev/termination-log
+      terminationMessagePolicy: File
+      imagePullPolicy: IfNotPresent
+  volumes:
+    - name: airflow-logs
+      emptyDir: {}
+    - name: airflow-config
+      configMap:
+        name: airflow-airflow-config
+        defaultMode: 420
+  restartPolicy: Never
+  terminationGracePeriodSeconds: 30
+  serviceAccountName: airflow-worker
+  serviceAccount: airflow-worker
+  securityContext:
+    runAsUser: 50000
+    fsGroup: 50000
+  imagePullSecrets:
+    - name: airflow-registry
+  schedulerName: default-scheduler

airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py

@@ -0,0 +1,165 @@
+# 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.
+from __future__ import annotations
+
+import logging
+import secrets
+import string
+from functools import cache
+from typing import TYPE_CHECKING
+
+import pendulum
+from kubernetes.client.rest import ApiException
+from slugify import slugify
+
+from airflow.configuration import conf
+from airflow.providers.cncf.kubernetes.backcompat import get_logical_date_key
+
+if TYPE_CHECKING:
+    from airflow.models.taskinstancekey import TaskInstanceKey
+
+log = logging.getLogger(__name__)
+
+alphanum_lower = string.ascii_lowercase + string.digits
+
+POD_NAME_MAX_LENGTH = 63  # Matches Linux kernel's HOST_NAME_MAX default value minus 1.
+
+
+def rand_str(num):
+    """
+    Generate random lowercase alphanumeric string of length num.
+
+    :meta private:
+    """
+    return "".join(secrets.choice(alphanum_lower) for _ in range(num))
+
+
+def add_unique_suffix(*, name: str, rand_len: int = 8, max_len: int = POD_NAME_MAX_LENGTH) -> str:
+    """
+    Add random string to pod or job name while staying under max length.
+
+    :param name: name of the pod or job
+    :param rand_len: length of the random string to append
+    :param max_len: maximum length of the pod name
+    :meta private:
+    """
+    suffix = "-" + rand_str(rand_len)
+    return name[: max_len - len(suffix)].strip("-.") + suffix
+
+
+def create_unique_id(
+    dag_id: str | None = None,
+    task_id: str | None = None,
+    *,
+    max_length: int = POD_NAME_MAX_LENGTH,
+    unique: bool = True,
+) -> str:
+    """
+    Generate unique pod or job ID given a dag_id and / or task_id.
+
+    :param dag_id: DAG ID
+    :param task_id: Task ID
+    :param max_length: max number of characters
+    :param unique: whether a random string suffix should be added
+    :return: A valid identifier for a kubernetes pod name
+    """
+    if not (dag_id or task_id):
+        raise ValueError("Must supply either dag_id or task_id.")
+    name = ""
+    if dag_id:
+        name += dag_id
+    if task_id:
+        if name:
+            name += "-"
+        name += task_id
+    base_name = slugify(name, lowercase=True)[:max_length].strip(".-")
+    if unique:
+        return add_unique_suffix(name=base_name, rand_len=8, max_len=max_length)
+    return base_name
+
+
+def annotations_to_key(annotations: dict[str, str]) -> TaskInstanceKey:
+    """Build a TaskInstanceKey based on pod annotations."""
+    log.debug("Creating task key for annotations %s", annotations)
+    dag_id = annotations["dag_id"]
+    task_id = annotations["task_id"]
+    try_number = int(annotations["try_number"])
+    annotation_run_id = annotations.get("run_id")
+    map_index = int(annotations.get("map_index", -1))
+
+    # Compat: Look up the run_id from the TI table!
+    from airflow.models.dagrun import DagRun
+    from airflow.models.taskinstance import TaskInstance, TaskInstanceKey
+    from airflow.settings import Session
+
+    logical_date_key = get_logical_date_key()
+
+    if not annotation_run_id and logical_date_key in annotations:
+        logical_date = pendulum.parse(annotations[logical_date_key])
+        # Do _not_ use create-session, we don't want to expunge
+        if Session is None:
+            raise RuntimeError("Session not configured. Call configure_orm() first.")
+        session = Session()
+
+        task_instance_run_id = (
+            session.query(TaskInstance.run_id)
+            .join(TaskInstance.dag_run)
+            .filter(
+                TaskInstance.dag_id == dag_id,
+                TaskInstance.task_id == task_id,
+                getattr(DagRun, logical_date_key) == logical_date,
+            )
+            .scalar()
+        )
+    else:
+        task_instance_run_id = annotation_run_id
+
+    return TaskInstanceKey(
+        dag_id=dag_id,
+        task_id=task_id,
+        run_id=task_instance_run_id,
+        try_number=try_number,
+        map_index=map_index,
+    )
+
+
+@cache
+def get_logs_task_metadata() -> bool:
+    return conf.getboolean("kubernetes_executor", "logs_task_metadata")
+
+
+def annotations_for_logging_task_metadata(annotation_set):
+    if get_logs_task_metadata():
+        annotations_for_logging = annotation_set
+    else:
+        annotations_for_logging = "<omitted>"
+    return annotations_for_logging
+
+
+def should_retry_creation(exception: BaseException) -> bool:
+    """
+    Check if an Exception indicates a transient error and warrants retrying.
+
+    This function is needed for preventing 'No agent available' error. The error appears time to time
+    when users try to create a Resource or Job. This issue is inside kubernetes and in the current moment
+    has no solution. Like a temporary solution we decided to retry Job or Resource creation request each
+    time when this error appears.
+    More about this issue here: https://github.com/cert-manager/cert-manager/issues/6457
+    """
+    if isinstance(exception, ApiException):
+        return str(exception.status) == "500"
+    return False