apache-airflow-providers-cncf-kubernetes 8.4.0__tar.gz → 8.4.1__tar.gz

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: apache-airflow-providers-cncf-kubernetes
-Version: 8.4.0
+Version: 8.4.1
 Summary: Provider package apache-airflow-providers-cncf-kubernetes for Apache Airflow
 Keywords: airflow-provider,cncf.kubernetes,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -29,8 +29,8 @@ Requires-Dist: google-re2>=1.0
 Requires-Dist: kubernetes>=29.0.0,<=30.1.0
 Requires-Dist: kubernetes_asyncio>=29.0.0,<=30.1.0
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.0
+Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.1/changelog.html
+Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.1
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: Twitter, https://twitter.com/ApacheAirflow
@@ -80,7 +80,7 @@ Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 
 Package ``apache-airflow-providers-cncf-kubernetes``
 
-Release: ``8.4.0``
+Release: ``8.4.1``
 
 
 `Kubernetes <https://kubernetes.io/>`__
@@ -93,7 +93,7 @@ This is a provider package for ``cncf.kubernetes`` provider. All classes for thi
 are in ``airflow.providers.cncf.kubernetes`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.1/>`_.
 
 Installation
 ------------
@@ -120,4 +120,4 @@ PIP package             Version required
 ======================  =====================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.1/changelog.html>`_.

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/README.rst

@@ -42,7 +42,7 @@
 
 Package ``apache-airflow-providers-cncf-kubernetes``
 
-Release: ``8.4.0``
+Release: ``8.4.1``
 
 
 `Kubernetes <https://kubernetes.io/>`__
@@ -55,7 +55,7 @@ This is a provider package for ``cncf.kubernetes`` provider. All classes for thi
 are in ``airflow.providers.cncf.kubernetes`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.1/>`_.
 
 Installation
 ------------
@@ -82,4 +82,4 @@ PIP package             Version required
 ======================  =====================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.1/changelog.html>`_.

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/airflow/providers/cncf/kubernetes/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "8.4.0"
+__version__ = "8.4.1"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.8.0"

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py

@@ -220,6 +220,16 @@ class KubernetesJobWatcher(multiprocessing.Process, LoggingMixin):
             # However, need to free the executor slot from the current executor.
             self.log.info("Event: pod %s adopted, annotations: %s", pod_name, annotations_string)
             self.watcher_queue.put((pod_name, namespace, ADOPTED, annotations, resource_version))
+        elif hasattr(pod.status, "reason") and pod.status.reason == "ProviderFailed":
+            # Most likely this happens due to Kubernetes setup (virtual kubelet, virtual nodes, etc.)
+            self.log.error(
+                "Event: %s failed to start with reason ProviderFailed, annotations: %s",
+                pod_name,
+                annotations_string,
+            )
+            self.watcher_queue.put(
+                (pod_name, namespace, TaskInstanceState.FAILED, annotations, resource_version)
+            )
         elif status == "Pending":
             # deletion_timestamp is set by kube server when a graceful deletion is requested.
             # since kube server have received request to delete pod set TI state failed

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/airflow/providers/cncf/kubernetes/executors/local_kubernetes_executor.py

@@ -103,6 +103,11 @@ class LocalKubernetesExecutor(LoggingMixin):
         """Number of new tasks this executor instance can accept."""
         return self.local_executor.slots_available
 
+    @property
+    def slots_occupied(self):
+        """Number of tasks this executor instance is currently managing."""
+        return len(self.running) + len(self.queued_tasks)
+
     def queue_command(
         self,
         task_instance: TaskInstance,

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/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": 1723969891,
+        "source-date-epoch": 1724520459,
         "versions": [
+            "8.4.1",
             "8.4.0",
             "8.3.4",
             "8.3.3",

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/airflow/providers/cncf/kubernetes/hooks/kubernetes.py

@@ -37,7 +37,11 @@ from airflow.hooks.base import BaseHook
 from airflow.models import Connection
 from airflow.providers.cncf.kubernetes.kube_client import _disable_verify_ssl, _enable_tcp_keepalive
 from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import should_retry_creation
-from airflow.providers.cncf.kubernetes.utils.pod_manager import PodOperatorHookProtocol
+from airflow.providers.cncf.kubernetes.utils.pod_manager import (
+    PodOperatorHookProtocol,
+    container_is_completed,
+    container_is_running,
+)
 from airflow.utils import yaml
 
 if TYPE_CHECKING:
@@ -682,6 +686,7 @@ class AsyncKubernetesHook(KubernetesHook):
 
         if self.config_dict:
             self.log.debug(LOADING_KUBE_CONFIG_FILE_RESOURCE.format("config dictionary"))
+            self._is_in_cluster = False
             await async_config.load_kube_config_from_dict(self.config_dict)
             return async_client.ApiClient()
 
@@ -830,3 +835,37 @@ class AsyncKubernetesHook(KubernetesHook):
                 return job
             self.log.info("The job '%s' is incomplete. Sleeping for %i sec.", name, poll_interval)
             await asyncio.sleep(poll_interval)
+
+    async def wait_until_container_complete(
+        self, name: str, namespace: str, container_name: str, poll_interval: float = 10
+    ) -> None:
+        """
+        Wait for the given container in the given pod to be completed.
+
+        :param name: Name of Pod to fetch.
+        :param namespace: Namespace of the Pod.
+        :param container_name: name of the container within the pod to monitor
+        """
+        while True:
+            pod = await self.get_pod(name=name, namespace=namespace)
+            if container_is_completed(pod=pod, container_name=container_name):
+                break
+            self.log.info("Waiting for container '%s' state to be completed", container_name)
+            await asyncio.sleep(poll_interval)
+
+    async def wait_until_container_started(
+        self, name: str, namespace: str, container_name: str, poll_interval: float = 10
+    ) -> None:
+        """
+        Wait for the given container in the given pod to be started.
+
+        :param name: Name of Pod to fetch.
+        :param namespace: Namespace of the Pod.
+        :param container_name: name of the container within the pod to monitor
+        """
+        while True:
+            pod = await self.get_pod(name=name, namespace=namespace)
+            if container_is_running(pod=pod, container_name=container_name):
+                break
+            self.log.info("Waiting for container '%s' state to be running", container_name)
+            await asyncio.sleep(poll_interval)

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/airflow/providers/cncf/kubernetes/operators/job.py

@@ -19,6 +19,7 @@
 from __future__ import annotations
 
 import copy
+import json
 import logging
 import os
 from functools import cached_property
@@ -39,6 +40,7 @@ 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.providers.cncf.kubernetes.utils.pod_manager import EMPTY_XCOM_RESULT, PodNotFoundException
 from airflow.utils import yaml
 from airflow.utils.context import Context
 
@@ -135,7 +137,7 @@ class KubernetesJobOperator(KubernetesPodOperator):
         return hook
 
     @cached_property
-    def client(self) -> BatchV1Api:
+    def job_client(self) -> BatchV1Api:
         return self.hook.batch_v1_client
 
     def create_job(self, job_request_obj: k8s.V1Job) -> k8s.V1Job:
@@ -150,6 +152,11 @@ class KubernetesJobOperator(KubernetesPodOperator):
                 "Deferrable mode is available only with parameter `wait_until_job_complete=True`. "
                 "Please, set it up."
             )
+        if (self.get_logs or self.do_xcom_push) and not self.wait_until_job_complete:
+            self.log.warning(
+                "Getting Logs and pushing to XCom are 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
@@ -159,16 +166,34 @@ 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.pod is None:
+            self.pod = self.get_or_create_pod(  # must set `self.pod` for `on_kill`
+                pod_request_obj=self.pod_request_obj,
+                context=context,
+            )
+
         if self.wait_until_job_complete and self.deferrable:
             self.execute_deferrable()
             return
 
         if self.wait_until_job_complete:
+            if self.do_xcom_push:
+                self.pod_manager.await_container_completion(
+                    pod=self.pod, container_name=self.base_container_name
+                )
+                self.pod_manager.await_xcom_sidecar_container_start(pod=self.pod)
+                xcom_result = self.extract_xcom(pod=self.pod)
             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,
             )
+            if self.get_logs:
+                self.pod_manager.fetch_requested_container_logs(
+                    pod=self.pod,
+                    containers=self.container_logs,
+                    follow_logs=True,
+                )
 
         ti.xcom_push(key="job", value=self.job.to_dict())
         if self.wait_until_job_complete:
@@ -176,17 +201,24 @@ class KubernetesJobOperator(KubernetesPodOperator):
                 raise AirflowException(
                     f"Kubernetes job '{self.job.metadata.name}' is failed with error '{error_message}'"
                 )
+            if self.do_xcom_push:
+                return xcom_result
 
     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]
+                pod_name=self.pod.metadata.name,  # type: ignore[union-attr]
+                pod_namespace=self.pod.metadata.namespace,  # type: ignore[union-attr]
+                base_container_name=self.base_container_name,
                 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,
+                get_logs=self.get_logs,
+                do_xcom_push=self.do_xcom_push,
             ),
             method_name="execute_complete",
         )
@@ -197,6 +229,22 @@ class KubernetesJobOperator(KubernetesPodOperator):
         if event["status"] == "error":
             raise AirflowException(event["message"])
 
+        if self.get_logs:
+            pod_name = event["pod_name"]
+            pod_namespace = event["pod_namespace"]
+            self.pod = self.hook.get_pod(pod_name, pod_namespace)
+            if not self.pod:
+                raise PodNotFoundException("Could not find pod after resuming from deferral")
+            self._write_logs(self.pod)
+
+        if self.do_xcom_push:
+            xcom_result = event["xcom_result"]
+            if isinstance(xcom_result, str) and xcom_result.rstrip() == EMPTY_XCOM_RESULT:
+                self.log.info("xcom result file is empty.")
+                return None
+            self.log.info("xcom result: \n%s", xcom_result)
+            return json.loads(xcom_result)
+
     @staticmethod
     def deserialize_job_template_file(path: str) -> k8s.V1Job:
         """
@@ -229,7 +277,9 @@ class KubernetesJobOperator(KubernetesPodOperator):
             }
             if self.termination_grace_period is not None:
                 kwargs.update(grace_period_seconds=self.termination_grace_period)
-            self.client.delete_namespaced_job(**kwargs)
+            self.job_client.delete_namespaced_job(**kwargs)
+        if self.pod:
+            super().on_kill()
 
     def build_job_request_obj(self, context: Context | None = None) -> k8s.V1Job:
         """
@@ -254,6 +304,7 @@ class KubernetesJobOperator(KubernetesPodOperator):
             metadata=pod_template.metadata,
             spec=pod_template.spec,
         )
+        self.pod_request_obj = pod_template
 
         job = k8s.V1Job(
             api_version="batch/v1",

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/airflow/providers/cncf/kubernetes/triggers/job.py

@@ -16,10 +16,13 @@
 # under the License.
 from __future__ import annotations
 
+import asyncio
 from functools import cached_property
 from typing import TYPE_CHECKING, Any, AsyncIterator
 
-from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook
+from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook, KubernetesHook
+from airflow.providers.cncf.kubernetes.utils.pod_manager import PodManager
+from airflow.providers.cncf.kubernetes.utils.xcom_sidecar import PodDefaults
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
@@ -32,32 +35,49 @@ class KubernetesJobTrigger(BaseTrigger):
 
     :param job_name: The name of the job.
     :param job_namespace: The namespace of the job.
+    :param pod_name: The name of the Pod.
+    :param pod_namespace: The namespace of the Pod.
+    :param base_container_name: The name of the base container in the pod.
     :param kubernetes_conn_id: The :ref:`kubernetes connection id <howto/connection:kubernetes>`
         for the Kubernetes cluster.
     :param cluster_context: Context that points to kubernetes cluster.
     :param config_file: Path to kubeconfig file.
     :param poll_interval: Polling period in seconds to check for the status.
     :param in_cluster: run kubernetes client with in_cluster configuration.
+    :param get_logs: get the stdout of the base container as logs of the tasks.
+    :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.
     """
 
     def __init__(
         self,
         job_name: str,
         job_namespace: str,
+        pod_name: str,
+        pod_namespace: str,
+        base_container_name: str,
         kubernetes_conn_id: str | None = None,
         poll_interval: float = 10.0,
         cluster_context: str | None = None,
         config_file: str | None = None,
         in_cluster: bool | None = None,
+        get_logs: bool = True,
+        do_xcom_push: bool = False,
     ):
         super().__init__()
         self.job_name = job_name
         self.job_namespace = job_namespace
+        self.pod_name = pod_name
+        self.pod_namespace = pod_namespace
+        self.base_container_name = base_container_name
         self.kubernetes_conn_id = kubernetes_conn_id
         self.poll_interval = poll_interval
         self.cluster_context = cluster_context
         self.config_file = config_file
         self.in_cluster = in_cluster
+        self.get_logs = get_logs
+        self.do_xcom_push = do_xcom_push
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serialize KubernetesCreateJobTrigger arguments and classpath."""
@@ -66,16 +86,36 @@ class KubernetesJobTrigger(BaseTrigger):
             {
                 "job_name": self.job_name,
                 "job_namespace": self.job_namespace,
+                "pod_name": self.pod_name,
+                "pod_namespace": self.pod_namespace,
+                "base_container_name": self.base_container_name,
                 "kubernetes_conn_id": self.kubernetes_conn_id,
                 "poll_interval": self.poll_interval,
                 "cluster_context": self.cluster_context,
                 "config_file": self.config_file,
                 "in_cluster": self.in_cluster,
+                "get_logs": self.get_logs,
+                "do_xcom_push": self.do_xcom_push,
             },
         )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
         """Get current job status and yield a TriggerEvent."""
+        if self.get_logs or self.do_xcom_push:
+            pod = await self.hook.get_pod(name=self.pod_name, namespace=self.pod_namespace)
+        if self.do_xcom_push:
+            await self.hook.wait_until_container_complete(
+                name=self.pod_name, namespace=self.pod_namespace, container_name=self.base_container_name
+            )
+            self.log.info("Checking if xcom sidecar container is started.")
+            await self.hook.wait_until_container_started(
+                name=self.pod_name,
+                namespace=self.pod_namespace,
+                container_name=PodDefaults.SIDECAR_CONTAINER_NAME,
+            )
+            self.log.info("Extracting result from xcom sidecar container.")
+            loop = asyncio.get_running_loop()
+            xcom_result = await loop.run_in_executor(None, self.pod_manager.extract_xcom, pod)
         job: V1Job = await self.hook.wait_until_job_complete(name=self.job_name, namespace=self.job_namespace)
         job_dict = job.to_dict()
         error_message = self.hook.is_job_failed(job=job)
@@ -83,11 +123,14 @@ class KubernetesJobTrigger(BaseTrigger):
             {
                 "name": job.metadata.name,
                 "namespace": job.metadata.namespace,
+                "pod_name": pod.metadata.name if self.get_logs else None,
+                "pod_namespace": pod.metadata.namespace if self.get_logs else None,
                 "status": "error" if error_message else "success",
                 "message": f"Job failed with error: {error_message}"
                 if error_message
                 else "Job completed successfully",
                 "job": job_dict,
+                "xcom_result": xcom_result if self.do_xcom_push else None,
             }
         )
 
@@ -99,3 +142,13 @@ class KubernetesJobTrigger(BaseTrigger):
             config_file=self.config_file,
             cluster_context=self.cluster_context,
         )
+
+    @cached_property
+    def pod_manager(self) -> PodManager:
+        sync_hook = KubernetesHook(
+            conn_id=self.kubernetes_conn_id,
+            in_cluster=self.in_cluster,
+            config_file=self.config_file,
+            cluster_context=self.cluster_context,
+        )
+        return PodManager(kube_client=sync_hook.core_v1_client)

{apache_airflow_providers_cncf_kubernetes-8.4.0 → apache_airflow_providers_cncf_kubernetes-8.4.1}/pyproject.toml

@@ -28,7 +28,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-cncf-kubernetes"
-version = "8.4.0"
+version = "8.4.1"
 description = "Provider package apache-airflow-providers-cncf-kubernetes for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -66,8 +66,8 @@ dependencies = [
 ]
 
 [project.urls]
-"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.0"
-"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.0/changelog.html"
+"Documentation" = "https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.1"
+"Changelog" = "https://airflow.apache.org/docs/apache-airflow-providers-cncf-kubernetes/8.4.1/changelog.html"
 "Bug Tracker" = "https://github.com/apache/airflow/issues"
 "Source Code" = "https://github.com/apache/airflow"
 "Slack Chat" = "https://s.apache.org/airflow-slack"