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

airflow/providers/cncf/kubernetes/sensors/spark_kubernetes.py

@@ -15,13 +15,17 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-from typing import TYPE_CHECKING, Optional, Sequence
+from __future__ import annotations
+
+from collections.abc import Sequence
+from functools import cached_property
+from typing import TYPE_CHECKING
 
 from kubernetes import client
 
 from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
-from airflow.sensors.base import BaseSensorOperator
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context
@@ -29,7 +33,7 @@ if TYPE_CHECKING:
 
 class SparkKubernetesSensor(BaseSensorOperator):
     """
-    Checks sparkApplication object in kubernetes cluster:
+    Checks sparkApplication object in kubernetes cluster.
 
     .. seealso::
         For more detail about Spark Application Object have a look at the reference:
@@ -37,6 +41,7 @@ class SparkKubernetesSensor(BaseSensorOperator):
 
     :param application_name: spark Application resource name
     :param namespace: the kubernetes namespace where the sparkApplication reside in
+    :param container_name: the kubernetes container name where the sparkApplication reside in
     :param kubernetes_conn_id: The :ref:`kubernetes connection<howto/connection:kubernetes>`
         to Kubernetes cluster.
     :param attach_log: determines whether logs for driver pod should be appended to the sensor log
@@ -53,21 +58,26 @@ class SparkKubernetesSensor(BaseSensorOperator):
         *,
         application_name: str,
         attach_log: bool = False,
-        namespace: Optional[str] = None,
+        namespace: str | None = None,
+        container_name: str = "spark-kubernetes-driver",
         kubernetes_conn_id: str = "kubernetes_default",
-        api_group: str = 'sparkoperator.k8s.io',
-        api_version: str = 'v1beta2',
+        api_group: str = "sparkoperator.k8s.io",
+        api_version: str = "v1beta2",
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
         self.application_name = application_name
         self.attach_log = attach_log
         self.namespace = namespace
+        self.container_name = container_name
         self.kubernetes_conn_id = kubernetes_conn_id
-        self.hook = KubernetesHook(conn_id=self.kubernetes_conn_id)
         self.api_group = api_group
         self.api_version = api_version
 
+    @cached_property
+    def hook(self) -> KubernetesHook:
+        return KubernetesHook(conn_id=self.kubernetes_conn_id)
+
     def _log_driver(self, application_state: str, response: dict) -> None:
         if not self.attach_log:
             return
@@ -82,7 +92,9 @@ class SparkKubernetesSensor(BaseSensorOperator):
         log_method = self.log.error if application_state in self.FAILURE_STATES else self.log.info
         try:
             log = ""
-            for line in self.hook.get_pod_logs(driver_pod_name, namespace=namespace):
+            for line in self.hook.get_pod_logs(
+                driver_pod_name, namespace=namespace, container=self.container_name
+            ):
                 log += line.decode()
             log_method(log)
         except client.rest.ApiException as e:
@@ -94,8 +106,9 @@ class SparkKubernetesSensor(BaseSensorOperator):
                 e,
             )
 
-    def poke(self, context: 'Context') -> bool:
+    def poke(self, context: Context) -> bool:
         self.log.info("Poking: %s", self.application_name)
+
         response = self.hook.get_custom_object(
             group=self.api_group,
             version=self.api_version,
@@ -103,17 +116,20 @@ class SparkKubernetesSensor(BaseSensorOperator):
             name=self.application_name,
             namespace=self.namespace,
         )
+
         try:
             application_state = response["status"]["applicationState"]["state"]
         except KeyError:
             return False
+
         if self.attach_log and application_state in self.FAILURE_STATES + self.SUCCESS_STATES:
             self._log_driver(application_state, response)
+
         if application_state in self.FAILURE_STATES:
-            raise AirflowException(f"Spark application failed with state: {application_state}")
-        elif application_state in self.SUCCESS_STATES:
+            message = f"Spark application failed with state: {application_state}"
+            raise AirflowException(message)
+        if application_state in self.SUCCESS_STATES:
             self.log.info("Spark application ended successfully")
             return True
-        else:
-            self.log.info("Spark application is still in state: %s", application_state)
-            return False
+        self.log.info("Spark application is still in state: %s", application_state)
+        return False

airflow/providers/cncf/kubernetes/template_rendering.py

@@ -0,0 +1,81 @@
+# 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
+
+from typing import TYPE_CHECKING
+
+from jinja2 import TemplateAssertionError, UndefinedError
+from kubernetes.client.api_client import ApiClient
+
+from airflow.exceptions import AirflowException
+from airflow.providers.cncf.kubernetes.kube_config import KubeConfig
+from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import create_unique_id
+from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator, generate_pod_command_args
+from airflow.utils.session import NEW_SESSION, provide_session
+
+if TYPE_CHECKING:
+    from airflow.models.taskinstance import TaskInstance
+
+
+def render_k8s_pod_yaml(task_instance: TaskInstance) -> dict | None:
+    """Render k8s pod yaml."""
+    kube_config = KubeConfig()
+    if task_instance.executor_config and task_instance.executor_config.get("pod_template_file"):
+        # If a specific pod_template_file was passed to the executor, we make
+        # sure to render the k8s pod spec using this one, and not the default one.
+        pod_template_file = task_instance.executor_config["pod_template_file"]
+    else:
+        # If no such pod_template_file override was passed, we can simply render
+        # The pod spec using the default template.
+        pod_template_file = kube_config.pod_template_file
+
+    # Generate command args using shared utility function
+    command_args = generate_pod_command_args(task_instance)
+
+    pod = PodGenerator.construct_pod(
+        dag_id=task_instance.dag_id,
+        run_id=task_instance.run_id,
+        task_id=task_instance.task_id,
+        map_index=task_instance.map_index,
+        date=None,
+        pod_id=create_unique_id(task_instance.dag_id, task_instance.task_id),
+        try_number=task_instance.try_number,
+        kube_image=kube_config.kube_image,
+        args=command_args,
+        pod_override_object=PodGenerator.from_obj(task_instance.executor_config),
+        scheduler_job_id="0",
+        namespace=kube_config.executor_namespace,
+        base_worker_pod=PodGenerator.deserialize_model_file(pod_template_file),
+        with_mutation_hook=True,
+    )
+    sanitized_pod = ApiClient().sanitize_for_serialization(pod)
+    return sanitized_pod
+
+
+@provide_session
+def get_rendered_k8s_spec(task_instance: TaskInstance, session=NEW_SESSION) -> dict | None:
+    """Fetch rendered template fields from DB."""
+    from airflow.models.renderedtifields import RenderedTaskInstanceFields
+
+    rendered_k8s_spec = RenderedTaskInstanceFields.get_k8s_pod_yaml(task_instance, session=session)
+    if not rendered_k8s_spec:
+        try:
+            rendered_k8s_spec = render_k8s_pod_yaml(task_instance)
+        except (TemplateAssertionError, UndefinedError) as e:
+            raise AirflowException(f"Unable to render a k8s spec for this taskinstance: {e}") from e
+    return rendered_k8s_spec

airflow/providers/cncf/kubernetes/triggers/__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/triggers/job.py

@@ -0,0 +1,176 @@
+# 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 asyncio
+import warnings
+from collections.abc import AsyncIterator
+from functools import cached_property
+from typing import TYPE_CHECKING, Any
+
+from airflow.exceptions import AirflowProviderDeprecationWarning
+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:
+    from kubernetes.client import V1Job
+
+
+class KubernetesJobTrigger(BaseTrigger):
+    """
+    KubernetesJobTrigger run on the trigger worker to check the state of Job.
+
+    :param job_name: The name of the job.
+    :param job_namespace: The namespace of the job.
+    :param pod_name: The name of the Pod. Parameter is deprecated, please use pod_names instead.
+    :param pod_names: The name of the Pods.
+    :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_names: list[str],
+        pod_namespace: str,
+        base_container_name: str,
+        pod_name: str | None = None,
+        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
+        if pod_name is not None:
+            self._pod_name = pod_name
+            self.pod_names = [
+                self.pod_name,
+            ]
+        else:
+            self.pod_names = pod_names
+        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
+
+    @property
+    def pod_name(self):
+        warnings.warn(
+            "`pod_name` parameter is deprecated, please use `pod_names`",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
+        return self._pod_name
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serialize KubernetesCreateJobTrigger arguments and classpath."""
+        return (
+            "airflow.providers.cncf.kubernetes.triggers.job.KubernetesJobTrigger",
+            {
+                "job_name": self.job_name,
+                "job_namespace": self.job_namespace,
+                "pod_names": self.pod_names,
+                "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]:
+        """Get current job status and yield a TriggerEvent."""
+        if self.do_xcom_push:
+            xcom_results = []
+            for pod_name in self.pod_names:
+                pod = await self.hook.get_pod(name=pod_name, namespace=self.pod_namespace)
+                await self.hook.wait_until_container_complete(
+                    name=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=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)
+                xcom_results.append(xcom_result)
+        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)
+        yield TriggerEvent(
+            {
+                "name": job.metadata.name,
+                "namespace": job.metadata.namespace,
+                "pod_names": [pod_name for pod_name in self.pod_names] if self.get_logs else None,
+                "pod_namespace": self.pod_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_results if self.do_xcom_push else None,
+            }
+        )
+
+    @cached_property
+    def hook(self) -> AsyncKubernetesHook:
+        return AsyncKubernetesHook(
+            conn_id=self.kubernetes_conn_id,
+            in_cluster=self.in_cluster,
+            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)