Handle Failed jobs when wait_until_complete=True for KubernetesJobOpe…

…rator

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

@@ -34,6 +34,7 @@
 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.utils.jobs import JOB_FINAL_STATUS_CONDITION_TYPES
 from airflow.providers.cncf.kubernetes.utils.pod_manager import PodOperatorHookProtocol
 from airflow.utils import yaml
 
@@ -525,7 +526,7 @@ def get_job_status(self, job_name: str, namespace: str) -> V1Job:
         )
 
     def wait_job_until_complete(self, job_name: str, namespace: str, job_poll_interval: float = 10) -> V1Job:
-        """Block job of specified name from Google Cloud until it is complete.
+        """Block job of specified name from Google Cloud until it is complete or failed.
 
         :param job_name: Name of Job to fetch.
         :param namespace: Namespace of the Job.
@@ -535,11 +536,8 @@ def wait_job_until_complete(self, job_name: str, namespace: str, job_poll_interv
         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 conditions := job.status.conditions:
-                if condition_complete := next((c for c in conditions if c.type == "Complete"), None):
-                    if condition_complete.status:
-                        self.log.info("The job '%s' is complete.", job_name)
-                        return job
+            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)
 
@@ -558,6 +556,34 @@ def list_jobs_from_namespace(self, namespace: str) -> V1JobList:
         """
         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 conditions := job.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) -> bool:
+        """Check whether the given job is failed.
+
+        :return: Boolean indicating that the given job is failed.
+        """
+        conditions = job.status.conditions or []
+        return bool(next((c for c in conditions if c.type == "Failed" and c.status), None))
+
 
 def _get_bool(val) -> bool | None:
     """Convert val to bool if can be done with certainty; if we cannot infer intention we return None."""

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

@@ -26,6 +26,7 @@
 from kubernetes.client import BatchV1Api, models as k8s
 from kubernetes.client.api_client import ApiClient
 
+from airflow.exceptions import AirflowException
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import KubernetesHook
 from airflow.providers.cncf.kubernetes.kubernetes_helper_functions import (
     add_unique_suffix,
@@ -65,7 +66,8 @@ 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 is complete. Default is False.
+    :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.
     """
@@ -143,11 +145,16 @@ def execute(self, context: Context):
         ti.xcom_push(key="job_namespace", value=self.job.metadata.namespace)
 
         if self.wait_job_until_complete:
-            self.hook.wait_job_until_complete(
+            self.job = self.hook.wait_job_until_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.hook.batch_v1_client.api_client.sanitize_for_serialization(self.job)
+        )
+        if self.hook.is_job_failed(job=self.job):
+            raise AirflowException(f"Kubernetes job '{self.job.metadata.name}' is failed")
 
     @staticmethod
     def deserialize_job_template_file(path: str) -> k8s.V1Job:

airflow/providers/cncf/kubernetes/utils/jobs.py

@@ -0,0 +1,25 @@
+# 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.
+"""Launches JOBs."""
+from __future__ import annotations
+
+JOB_FINAL_STATUS_CONDITION_TYPES = {
+    "Complete",
+    "Failed",
+}
+
+JOB_STATUS_CONDITION_TYPES = JOB_FINAL_STATUS_CONDITION_TYPES | {"Suspended"}

tests/providers/cncf/kubernetes/hooks/test_kubernetes.py

@@ -462,6 +462,60 @@ def test_get_job_status(self, mock_client, mock_kube_config_merger, mock_kube_co
         )
         assert job_actual == job_expected
 
+    @pytest.mark.parametrize(
+        "conditions, expected_result",
+        [
+            (None, False),
+            ([], False),
+            ([mock.MagicMock(type="Complete", status=True)], False),
+            ([mock.MagicMock(type="Complete", status=False)], False),
+            ([mock.MagicMock(type="Failed", status=False)], False),
+            ([mock.MagicMock(type="Failed", status=True)], True),
+            (
+                [mock.MagicMock(type="Complete", status=False), mock.MagicMock(type="Failed", status=True)],
+                True,
+            ),
+            (
+                [mock.MagicMock(type="Complete", status=True), mock.MagicMock(type="Failed", status=True)],
+                True,
+            ),
+        ],
+    )
+    @patch("kubernetes.config.kube_config.KubeConfigLoader")
+    @patch("kubernetes.config.kube_config.KubeConfigMerger")
+    def test_is_job_failed(self, mock_merger, mock_loader, conditions, expected_result):
+        mock_job = mock.MagicMock()
+        mock_job.status.conditions = conditions
+
+        hook = KubernetesHook()
+        actual_result = hook.is_job_failed(mock_job)
+
+        assert actual_result == expected_result
+
+    @pytest.mark.parametrize(
+        "condition_type, status, expected_result",
+        [
+            ("Complete", False, False),
+            ("Complete", True, True),
+            ("Failed", False, False),
+            ("Failed", True, True),
+            ("Suspended", False, False),
+            ("Suspended", True, False),
+            ("Unknown", False, False),
+            ("Unknown", True, False),
+        ],
+    )
+    @patch("kubernetes.config.kube_config.KubeConfigLoader")
+    @patch("kubernetes.config.kube_config.KubeConfigMerger")
+    def test_is_job_complete(self, mock_merger, mock_loader, condition_type, status, expected_result):
+        mock_job = mock.MagicMock()
+        mock_job.status.conditions = [mock.MagicMock(type=condition_type, status=status)]
+
+        hook = KubernetesHook()
+        actual_result = hook.is_job_complete(mock_job)
+
+        assert actual_result == expected_result
+
     @patch("kubernetes.config.kube_config.KubeConfigLoader")
     @patch("kubernetes.config.kube_config.KubeConfigMerger")
     @patch(f"{HOOK_MODULE}.KubernetesHook.get_job_status")

tests/providers/cncf/kubernetes/operators/test_job.py

@@ -24,6 +24,7 @@
 import pytest
 from kubernetes.client import ApiClient, models as k8s
 
+from airflow.exceptions import AirflowException
 from airflow.models import DAG, DagModel, DagRun, TaskInstance
 from airflow.providers.cncf.kubernetes.operators.job import KubernetesJobOperator
 from airflow.utils import timezone
@@ -458,6 +459,7 @@ def test_task_id_as_name_dag_id_is_ignored(self):
     @patch(JOB_OPERATORS_PATH.format("KubernetesJobOperator.create_job"))
     @patch(HOOK_CLASS)
     def test_execute(self, mock_hook, mock_create_job, mock_build_job_request_obj):
+        mock_hook.return_value.is_job_failed.return_value = False
         mock_job_request_obj = mock_build_job_request_obj.return_value
         mock_job_expected = mock_create_job.return_value
         mock_ti = mock.MagicMock()
@@ -483,6 +485,19 @@ def test_execute(self, mock_hook, mock_create_job, mock_build_job_request_obj):
         assert execute_result is None
         assert not mock_hook.wait_job_until_complete.called
 
+    @patch(JOB_OPERATORS_PATH.format("KubernetesJobOperator.build_job_request_obj"))
+    @patch(JOB_OPERATORS_PATH.format("KubernetesJobOperator.create_job"))
+    @patch(HOOK_CLASS)
+    def test_execute_fail(self, mock_hook, mock_create_job, mock_build_job_request_obj):
+        mock_hook.return_value.is_job_failed.return_value = True
+
+        op = KubernetesJobOperator(
+            task_id="test_task_id",
+        )
+
+        with pytest.raises(AirflowException):
+            op.execute(context=dict(ti=mock.MagicMock()))
+
     @patch(JOB_OPERATORS_PATH.format("KubernetesJobOperator.build_job_request_obj"))
     @patch(JOB_OPERATORS_PATH.format("KubernetesJobOperator.create_job"))
     @patch(f"{HOOK_CLASS}.wait_job_until_complete")