apache-airflow-providers-cncf-kubernetes 8.0.1rc1__py3-none-any.whl → 8.1.0__py3-none-any.whl

airflow/providers/cncf/kubernetes/__init__.py

@@ -27,7 +27,7 @@ import packaging.version
 
 __all__ = ["__version__"]
 
-__version__ = "8.0.1"
+__version__ = "8.1.0"
 
 try:
     from airflow import __version__ as airflow_version

airflow/providers/cncf/kubernetes/backcompat/backwards_compat_converters.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Executes task in a Kubernetes POD."""
+
 from __future__ import annotations
 
 from kubernetes.client import ApiClient, models as k8s

airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py

@@ -21,6 +21,7 @@ KubernetesExecutor.
     For more information on how the KubernetesExecutor works, take a look at the guide:
     :doc:`/kubernetes_executor`
 """
+
 from __future__ import annotations
 
 import contextlib

airflow/providers/cncf/kubernetes/get_provider_info.py

@@ -28,8 +28,9 @@ def get_provider_info():
         "name": "Kubernetes",
         "description": "`Kubernetes <https://kubernetes.io/>`__\n",
         "state": "ready",
-        "source-date-epoch": 1709555165,
+        "source-date-epoch": 1712665231,
         "versions": [
+            "8.1.0",
             "8.0.1",
             "8.0.0",
             "7.14.0",
@@ -135,7 +136,10 @@ def get_provider_info():
         "triggers": [
             {
                 "integration-name": "Kubernetes",
-                "python-modules": ["airflow.providers.cncf.kubernetes.triggers.pod"],
+                "python-modules": [
+                    "airflow.providers.cncf.kubernetes.triggers.pod",
+                    "airflow.providers.cncf.kubernetes.triggers.job",
+                ],
             }
         ],
         "connection-types": [

airflow/providers/cncf/kubernetes/hooks/kubernetes.py

@@ -16,10 +16,12 @@
 # under the License.
 from __future__ import annotations
 
+import asyncio
 import contextlib
 import json
 import tempfile
 from functools import cached_property
+from time import sleep
 from typing import TYPE_CHECKING, Any, Generator
 
 import aiofiles
@@ -37,10 +39,18 @@ from airflow.providers.cncf.kubernetes.utils.pod_manager import PodOperatorHookP
 from airflow.utils import yaml
 
 if TYPE_CHECKING:
+    from kubernetes.client import V1JobList
     from kubernetes.client.models import V1Deployment, V1Job, V1Pod
 
 LOADING_KUBE_CONFIG_FILE_RESOURCE = "Loading Kubernetes configuration file kube_config from {}..."
 
+JOB_FINAL_STATUS_CONDITION_TYPES = {
+    "Complete",
+    "Failed",
+}
+
+JOB_STATUS_CONDITION_TYPES = JOB_FINAL_STATUS_CONDITION_TYPES | {"Suspended"}
+
 
 def _load_body_to_dict(body: str) -> dict:
     try:
@@ -502,6 +512,114 @@ class KubernetesHook(BaseHook, PodOperatorHookProtocol):
             raise e
         return resp
 
+    def get_job(self, job_name: str, namespace: str) -> V1Job:
+        """Get Job of specified name and namespace.
+
+        :param job_name: Name of Job to fetch.
+        :param namespace: Namespace of the Job.
+        :return: Job object
+        """
+        return self.batch_v1_client.read_namespaced_job(name=job_name, namespace=namespace, pretty=True)
+
+    def get_job_status(self, job_name: str, namespace: str) -> V1Job:
+        """Get job with status of specified name and namespace.
+
+        :param job_name: Name of Job to fetch.
+        :param namespace: Namespace of the Job.
+        :return: Job object
+        """
+        return self.batch_v1_client.read_namespaced_job_status(
+            name=job_name, namespace=namespace, pretty=True
+        )
+
+    def wait_until_job_complete(self, job_name: str, namespace: str, job_poll_interval: float = 10) -> V1Job:
+        """Block job of specified name and namespace until it is complete or failed.
+
+        :param job_name: Name of Job to fetch.
+        :param namespace: Namespace of the Job.
+        :param job_poll_interval: Interval in seconds between polling the job status
+        :return: Job object
+        """
+        while True:
+            self.log.info("Requesting status for the job '%s' ", job_name)
+            job: V1Job = self.get_job_status(job_name=job_name, namespace=namespace)
+            if self.is_job_complete(job=job):
+                return job
+            self.log.info("The job '%s' is incomplete. Sleeping for %i sec.", job_name, job_poll_interval)
+            sleep(job_poll_interval)
+
+    def list_jobs_all_namespaces(self) -> V1JobList:
+        """Get list of Jobs from all namespaces.
+
+        :return: V1JobList object
+        """
+        return self.batch_v1_client.list_job_for_all_namespaces(pretty=True)
+
+    def list_jobs_from_namespace(self, namespace: str) -> V1JobList:
+        """Get list of Jobs from dedicated namespace.
+
+        :param namespace: Namespace of the Job.
+        :return: V1JobList object
+        """
+        return self.batch_v1_client.list_namespaced_job(namespace=namespace, pretty=True)
+
+    def is_job_complete(self, job: V1Job) -> bool:
+        """Check whether the given job is complete (with success or fail).
+
+        :return: Boolean indicating that the given job is complete.
+        """
+        if status := job.status:
+            if conditions := status.conditions:
+                if final_condition_types := list(
+                    c for c in conditions if c.type in JOB_FINAL_STATUS_CONDITION_TYPES and c.status
+                ):
+                    s = "s" if len(final_condition_types) > 1 else ""
+                    self.log.info(
+                        "The job '%s' state%s: %s",
+                        job.metadata.name,
+                        s,
+                        ", ".join(f"{c.type} at {c.last_transition_time}" for c in final_condition_types),
+                    )
+                    return True
+        return False
+
+    @staticmethod
+    def is_job_failed(job: V1Job) -> str | bool:
+        """Check whether the given job is failed.
+
+        :return: Error message if the job is failed, and False otherwise.
+        """
+        if status := job.status:
+            conditions = status.conditions or []
+            if fail_condition := next((c for c in conditions if c.type == "Failed" and c.status), None):
+                return fail_condition.reason
+        return False
+
+    @staticmethod
+    def is_job_successful(job: V1Job) -> str | bool:
+        """Check whether the given job is completed successfully..
+
+        :return: Error message if the job is failed, and False otherwise.
+        """
+        if status := job.status:
+            conditions = status.conditions or []
+            return bool(next((c for c in conditions if c.type == "Complete" and c.status), None))
+        return False
+
+    def patch_namespaced_job(self, job_name: str, namespace: str, body: object) -> V1Job:
+        """
+        Update the specified Job.
+
+        :param job_name: name of the Job
+        :param namespace: the namespace to run within kubernetes
+        :param body: json object with parameters for update
+        """
+        return self.batch_v1_client.patch_namespaced_job(
+            name=job_name,
+            namespace=namespace,
+            body=body,
+        )
+
 
 def _get_bool(val) -> bool | None:
     """Convert val to bool if can be done with certainty; if we cannot infer intention we return None."""
@@ -667,3 +785,34 @@ class AsyncKubernetesHook(KubernetesHook):
             except HTTPError:
                 self.log.exception("There was an error reading the kubernetes API.")
                 raise
+
+    async def get_job_status(self, name: str, namespace: str) -> V1Job:
+        """
+        Get job's status object.
+
+        :param name: Name of the pod.
+        :param namespace: Name of the pod's namespace.
+        """
+        async with self.get_conn() as connection:
+            v1_api = async_client.BatchV1Api(connection)
+            job: V1Job = await v1_api.read_namespaced_job_status(
+                name=name,
+                namespace=namespace,
+            )
+        return job
+
+    async def wait_until_job_complete(self, name: str, namespace: str, poll_interval: float = 10) -> V1Job:
+        """Block job of specified name and namespace until it is complete or failed.
+
+        :param name: Name of Job to fetch.
+        :param namespace: Namespace of the Job.
+        :param poll_interval: Interval in seconds between polling the job status
+        :return: Job object
+        """
+        while True:
+            self.log.info("Requesting status for the job '%s' ", name)
+            job: V1Job = await self.get_job_status(name=name, namespace=namespace)
+            if self.is_job_complete(job=job):
+                return job
+            self.log.info("The job '%s' is incomplete. Sleeping for %i sec.", name, poll_interval)
+            await asyncio.sleep(poll_interval)

airflow/providers/cncf/kubernetes/k8s_model.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Classes for interacting with Kubernetes API."""
+
 from __future__ import annotations
 
 from abc import ABC, abstractmethod

airflow/providers/cncf/kubernetes/kube_client.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Client for kubernetes communication."""
+
 from __future__ import annotations
 
 import logging

airflow/providers/cncf/kubernetes/kubernetes_helper_functions.py

@@ -104,7 +104,7 @@ def create_unique_id(
         name += task_id
     base_name = slugify(name, lowercase=True)[:max_length].strip(".-")
     if unique:
-        return add_pod_suffix(pod_name=base_name, rand_len=8, max_len=max_length)
+        return add_unique_suffix(name=base_name, rand_len=8, max_len=max_length)
     else:
         return base_name
 

airflow/providers/cncf/kubernetes/operators/custom_object_launcher.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Launches Custom object."""
+
 from __future__ import annotations
 
 import time
@@ -59,9 +60,9 @@ class SparkJobSpec:
         if self.spec.get("dynamicAllocation", {}).get("enabled"):
             if not all(
                 [
-                    self.spec["dynamicAllocation"]["initialExecutors"],
-                    self.spec["dynamicAllocation"]["minExecutors"],
-                    self.spec["dynamicAllocation"]["maxExecutors"],
+                    self.spec["dynamicAllocation"].get("initialExecutors"),
+                    self.spec["dynamicAllocation"].get("minExecutors"),
+                    self.spec["dynamicAllocation"].get("maxExecutors"),
                 ]
             ):
                 raise AirflowException("Make sure initial/min/max value for dynamic allocation is passed")

airflow/providers/cncf/kubernetes/operators/job.py

@@ -15,6 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """Executes a Kubernetes Job."""
+
 from __future__ import annotations
 
 import copy
@@ -25,7 +26,11 @@ from typing import TYPE_CHECKING, Sequence
 
 from kubernetes.client import BatchV1Api, models as k8s
 from kubernetes.client.api_client import ApiClient
+from kubernetes.client.rest import ApiException
 
+from airflow.configuration import conf
+from airflow.exceptions import AirflowException
+from airflow.models import BaseOperator
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
 from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
     add_unique_suffix,
@@ -33,7 +38,9 @@ from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
 )
 from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
 from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator, merge_objects
+from airflow.providers.cncf.kubernetes.triggers.job import KubernetesJobTrigger
 from airflow.utils import yaml
+from airflow.utils.context import Context
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -65,6 +72,12 @@ class KubernetesJobOperator(KubernetesPodOperator):
     :param selector: The selector of this V1JobSpec.
     :param suspend: Suspend specifies whether the Job controller should create Pods or not.
     :param ttl_seconds_after_finished: ttlSecondsAfterFinished limits the lifetime of a Job that has finished execution (either Complete or Failed).
+    :param wait_until_job_complete: Whether to wait until started job finished execution (either Complete or
+        Failed). Default is False.
+    :param job_poll_interval: Interval in seconds between polling the job status. Default is 10.
+        Used if the parameter `wait_until_job_complete` set True.
+    :param deferrable: Run operator in the deferrable mode. Note that the parameter
+        `wait_until_job_complete` must be set True.
     """
 
     template_fields: Sequence[str] = tuple({"job_template_file"} | set(KubernetesPodOperator.template_fields))
@@ -82,6 +95,9 @@ class KubernetesJobOperator(KubernetesPodOperator):
         selector: k8s.V1LabelSelector | None = None,
         suspend: bool | None = None,
         ttl_seconds_after_finished: int | None = None,
+        wait_until_job_complete: bool = False,
+        job_poll_interval: float = 10,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -97,6 +113,9 @@ class KubernetesJobOperator(KubernetesPodOperator):
         self.selector = selector
         self.suspend = suspend
         self.ttl_seconds_after_finished = ttl_seconds_after_finished
+        self.wait_until_job_complete = wait_until_job_complete
+        self.job_poll_interval = job_poll_interval
+        self.deferrable = deferrable
 
     @cached_property
     def _incluster_namespace(self):
@@ -126,6 +145,11 @@ class KubernetesJobOperator(KubernetesPodOperator):
         return job_request_obj
 
     def execute(self, context: Context):
+        if self.deferrable and not self.wait_until_job_complete:
+            self.log.warning(
+                "Deferrable mode is available only with parameter `wait_until_job_complete=True`. "
+                "Please, set it up."
+            )
         self.job_request_obj = self.build_job_request_obj(context)
         self.job = self.create_job(  # must set `self.job` for `on_kill`
             job_request_obj=self.job_request_obj
@@ -135,6 +159,44 @@ class KubernetesJobOperator(KubernetesPodOperator):
         ti.xcom_push(key="job_name", value=self.job.metadata.name)
         ti.xcom_push(key="job_namespace", value=self.job.metadata.namespace)
 
+        if self.wait_until_job_complete and self.deferrable:
+            self.execute_deferrable()
+            return
+
+        if self.wait_until_job_complete:
+            self.job = self.hook.wait_until_job_complete(
+                job_name=self.job.metadata.name,
+                namespace=self.job.metadata.namespace,
+                job_poll_interval=self.job_poll_interval,
+            )
+
+        ti.xcom_push(key="job", value=self.job.to_dict())
+        if self.wait_until_job_complete:
+            if error_message := self.hook.is_job_failed(job=self.job):
+                raise AirflowException(
+                    f"Kubernetes job '{self.job.metadata.name}' is failed with error '{error_message}'"
+                )
+
+    def execute_deferrable(self):
+        self.defer(
+            trigger=KubernetesJobTrigger(
+                job_name=self.job.metadata.name,  # type: ignore[union-attr]
+                job_namespace=self.job.metadata.namespace,  # type: ignore[union-attr]
+                kubernetes_conn_id=self.kubernetes_conn_id,
+                cluster_context=self.cluster_context,
+                config_file=self.config_file,
+                in_cluster=self.in_cluster,
+                poll_interval=self.job_poll_interval,
+            ),
+            method_name="execute_complete",
+        )
+
+    def execute_complete(self, context: Context, event: dict, **kwargs):
+        ti = context["ti"]
+        ti.xcom_push(key="job", value=event["job"])
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
+
     @staticmethod
     def deserialize_job_template_file(path: str) -> k8s.V1Job:
         """
@@ -163,6 +225,7 @@ class KubernetesJobOperator(KubernetesPodOperator):
             kwargs = {
                 "name": job.metadata.name,
                 "namespace": job.metadata.namespace,
+                "job": self.hook.batch_v1_client.api_client.sanitize_for_serialization(self.job),
             }
             if self.termination_grace_period is not None:
                 kwargs.update(grace_period_seconds=self.termination_grace_period)
@@ -284,3 +347,179 @@ class KubernetesJobOperator(KubernetesPodOperator):
             return merge_objects(base_spec, client_spec)
 
         return None
+
+
+class KubernetesDeleteJobOperator(BaseOperator):
+    """
+    Delete a Kubernetes Job.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KubernetesDeleteJobOperator`
+
+    :param name: name of the Job.
+    :param namespace: the namespace to run within kubernetes.
+    :param kubernetes_conn_id: The :ref:`kubernetes connection id <howto/connection:kubernetes>`
+        for the Kubernetes cluster.
+    :param config_file: The path to the Kubernetes config file. (templated)
+        If not specified, default value is ``~/.kube/config``
+    :param in_cluster: run kubernetes client with in_cluster configuration.
+    :param cluster_context: context that points to kubernetes cluster.
+        Ignored when in_cluster is True. If None, current-context is used. (templated)
+    :param delete_on_status: Condition for performing delete operation depending on the job status. Values:
+        ``None`` - delete the job regardless of its status, "Complete" - delete only successfully completed
+        jobs, "Failed" - delete only failed jobs. (default: ``None``)
+    :param wait_for_completion: Whether to wait for the job to complete. (default: ``False``)
+    :param poll_interval: Interval in seconds between polling the job status. Used when the `delete_on_status`
+        parameter is set. (default: 10.0)
+    """
+
+    template_fields: Sequence[str] = (
+        "config_file",
+        "name",
+        "namespace",
+        "cluster_context",
+    )
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        namespace: str,
+        kubernetes_conn_id: str | None = KubernetesHook.default_conn_name,
+        config_file: str | None = None,
+        in_cluster: bool | None = None,
+        cluster_context: str | None = None,
+        delete_on_status: str | None = None,
+        wait_for_completion: bool = False,
+        poll_interval: float = 10.0,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.name = name
+        self.namespace = namespace
+        self.kubernetes_conn_id = kubernetes_conn_id
+        self.config_file = config_file
+        self.in_cluster = in_cluster
+        self.cluster_context = cluster_context
+        self.delete_on_status = delete_on_status
+        self.wait_for_completion = wait_for_completion
+        self.poll_interval = poll_interval
+
+    @cached_property
+    def hook(self) -> KubernetesHook:
+        return KubernetesHook(
+            conn_id=self.kubernetes_conn_id,
+            in_cluster=self.in_cluster,
+            config_file=self.config_file,
+            cluster_context=self.cluster_context,
+        )
+
+    @cached_property
+    def client(self) -> BatchV1Api:
+        return self.hook.batch_v1_client
+
+    def execute(self, context: Context):
+        try:
+            if self.delete_on_status not in ("Complete", "Failed", None):
+                raise AirflowException(
+                    "The `delete_on_status` parameter must be one of 'Complete', 'Failed' or None. "
+                    "The current value is %s",
+                    str(self.delete_on_status),
+                )
+
+            if self.wait_for_completion:
+                job = self.hook.wait_until_job_complete(
+                    job_name=self.name, namespace=self.namespace, job_poll_interval=self.poll_interval
+                )
+            else:
+                job = self.hook.get_job_status(job_name=self.name, namespace=self.namespace)
+
+            if (
+                self.delete_on_status is None
+                or (self.delete_on_status == "Complete" and self.hook.is_job_successful(job=job))
+                or (self.delete_on_status == "Failed" and self.hook.is_job_failed(job=job))
+            ):
+                self.log.info("Deleting kubernetes Job: %s", self.name)
+                self.client.delete_namespaced_job(name=self.name, namespace=self.namespace)
+                self.log.info("Kubernetes job was deleted.")
+            else:
+                self.log.info(
+                    "Deletion of the job %s was skipped due to settings of on_status=%s",
+                    self.name,
+                    self.delete_on_status,
+                )
+        except ApiException as e:
+            if e.status == 404:
+                self.log.info("The Kubernetes job %s does not exist.", self.name)
+            else:
+                raise e
+
+
+class KubernetesPatchJobOperator(BaseOperator):
+    """
+    Update a Kubernetes Job.
+
+    .. seealso::
+        For more information on how to use this operator, take a look at the guide:
+        :ref:`howto/operator:KubernetesPatchJobOperator`
+
+    :param name: name of the Job
+    :param namespace: the namespace to run within kubernetes
+    :param body: Job json object with parameters for update
+        https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.25/#job-v1-batch
+        e.g. ``{"spec": {"suspend": True}}``
+    :param kubernetes_conn_id: The :ref:`kubernetes connection id <howto/connection:kubernetes>`
+        for the Kubernetes cluster.
+    :param config_file: The path to the Kubernetes config file. (templated)
+        If not specified, default value is ``~/.kube/config``
+    :param in_cluster: run kubernetes client with in_cluster configuration.
+    :param cluster_context: context that points to kubernetes cluster.
+        Ignored when in_cluster is True. If None, current-context is used. (templated)
+    """
+
+    template_fields: Sequence[str] = (
+        "config_file",
+        "name",
+        "namespace",
+        "body",
+        "cluster_context",
+    )
+
+    def __init__(
+        self,
+        *,
+        name: str,
+        namespace: str,
+        body: object,
+        kubernetes_conn_id: str | None = KubernetesHook.default_conn_name,
+        config_file: str | None = None,
+        in_cluster: bool | None = None,
+        cluster_context: str | None = None,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.name = name
+        self.namespace = namespace
+        self.body = body
+        self.kubernetes_conn_id = kubernetes_conn_id
+        self.config_file = config_file
+        self.in_cluster = in_cluster
+        self.cluster_context = cluster_context
+
+    @cached_property
+    def hook(self) -> KubernetesHook:
+        return KubernetesHook(
+            conn_id=self.kubernetes_conn_id,
+            in_cluster=self.in_cluster,
+            config_file=self.config_file,
+            cluster_context=self.cluster_context,
+        )
+
+    def execute(self, context: Context) -> dict:
+        self.log.info("Updating existing Job: %s", self.name)
+        job_object = self.hook.patch_namespaced_job(
+            job_name=self.name, namespace=self.namespace, body=self.body
+        )
+        self.log.info("Job was updated.")
+        return k8s.V1Job.to_dict(job_object)

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

@@ -16,6 +16,7 @@
 # specific language governing permissions and limitations
 # under the License.
 """This module is deprecated. Please use :mod:`airflow.providers.cncf.kubernetes.operators.pod` instead."""
+
 from __future__ import annotations
 
 import warnings