Retry k8s API requests in KubernetesPodTrigger (apache#47187)

---------

Co-authored-by: Daniel Standish <[email protected]>

Author: insomnes

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py

@@ -834,6 +834,10 @@ def trigger_reentry(self, context: Context, event: dict[str, Any]) -> Any:
             last_log_time = event.get("last_log_time")
 
             if event["status"] in ("error", "failed", "timeout"):
+                event_message = event.get("message", "No message provided")
+                self.log.error(
+                    "Trigger emitted an %s event, failing the task: %s", event["status"], event_message
+                )
                 # fetch some logs when pod is failed
                 if self.get_logs:
                     self._write_logs(self.pod, follow=follow, since_time=last_log_time)

providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/triggers/pod.py

@@ -22,7 +22,10 @@
 from collections.abc import AsyncIterator
 from enum import Enum
 from functools import cached_property
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, cast
+
+import tenacity
+from kubernetes_asyncio.client.models import V1Pod
 
 from airflow.providers.cncf.kubernetes.hooks.kubernetes import AsyncKubernetesHook
 from airflow.providers.cncf.kubernetes.utils.pod_manager import (
@@ -33,7 +36,6 @@
 from airflow.triggers.base import BaseTrigger, TriggerEvent
 
 if TYPE_CHECKING:
-    from kubernetes_asyncio.client.models import V1Pod
     from pendulum import DateTime
 
 
@@ -200,7 +202,7 @@ def _format_exception_description(self, exc: Exception) -> Any:
     async def _wait_for_pod_start(self) -> ContainerState:
         """Loops until pod phase leaves ``PENDING`` If timeout is reached, throws error."""
         while True:
-            pod = await self.hook.get_pod(self.pod_name, self.pod_namespace)
+            pod = await self._get_pod()
             if not pod.status.phase == "Pending":
                 return self.define_container_state(pod)
 
@@ -223,7 +225,7 @@ async def _wait_for_container_completion(self) -> TriggerEvent:
         if self.logging_interval is not None:
             time_get_more_logs = time_begin + datetime.timedelta(seconds=self.logging_interval)
         while True:
-            pod = await self.hook.get_pod(self.pod_name, self.pod_namespace)
+            pod = await self._get_pod()
             container_state = self.define_container_state(pod)
             if container_state == ContainerState.TERMINATED:
                 return TriggerEvent(
@@ -257,6 +259,14 @@ async def _wait_for_container_completion(self) -> TriggerEvent:
             self.log.debug("Sleeping for %s seconds.", self.poll_interval)
             await asyncio.sleep(self.poll_interval)
 
+    @tenacity.retry(stop=tenacity.stop_after_attempt(3), wait=tenacity.wait_exponential(), reraise=True)
+    async def _get_pod(self) -> V1Pod:
+        """Get the pod from Kubernetes with retries."""
+        pod = await self.hook.get_pod(name=self.pod_name, namespace=self.pod_namespace)
+        # Due to AsyncKubernetesHook overriding get_pod, we need to cast the return
+        # value to kubernetes_asyncio.V1Pod, because it's perceived as different type
+        return cast(V1Pod, pod)
+
     def _get_async_hook(self) -> AsyncKubernetesHook:
         # TODO: Remove this method when the min version of kubernetes provider is 7.12.0 in Google provider.
         return AsyncKubernetesHook(

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py

@@ -2091,10 +2091,12 @@ def test_async_create_pod_should_execute_successfully(
         ti_mock.xcom_push.assert_any_call(key="pod_namespace", value=TEST_NAMESPACE)
         assert isinstance(exc.value.trigger, KubernetesPodTrigger)
 
+    @pytest.mark.parametrize("status", ["error", "failed", "timeout"])
+    @patch(KUB_OP_PATH.format("log"))
     @patch(KUB_OP_PATH.format("cleanup"))
     @patch(HOOK_CLASS)
-    def test_async_create_pod_should_throw_exception(self, mocked_hook, mocked_cleanup):
-        """Tests that an AirflowException is raised in case of error event"""
+    def test_async_create_pod_should_throw_exception(self, mocked_hook, mocked_cleanup, mocked_log, status):
+        """Tests that an AirflowException is raised in case of error event and event is logged"""
 
         mocked_hook.return_value.get_pod.return_value = MagicMock()
         k = KubernetesPodOperator(
@@ -2111,17 +2113,21 @@ def test_async_create_pod_should_throw_exception(self, mocked_hook, mocked_clean
             deferrable=True,
         )
 
+        message = "Some message"
         with pytest.raises(AirflowException):
             k.trigger_reentry(
                 context=None,
                 event={
-                    "status": "error",
-                    "message": "Some error",
+                    "status": status,
+                    "message": message,
                     "name": TEST_NAME,
                     "namespace": TEST_NAMESPACE,
                 },
             )
 
+        log_message = "Trigger emitted an %s event, failing the task: %s"
+        mocked_log.error.assert_called_once_with(log_message, status, message)
+
     @pytest.mark.parametrize(
         "kwargs, actual_exit_code, expected_exc, pod_status, event_status",
         [

providers/cncf/kubernetes/tests/unit/cncf/kubernetes/triggers/test_pod.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import datetime
 import logging
 from asyncio import Future
@@ -386,3 +387,48 @@ async def test_run_loop_return_success_for_completed_pod_after_timeout(
             )
             == actual
         )
+
+    @pytest.mark.asyncio
+    @mock.patch(f"{TRIGGER_PATH}.hook")
+    async def test__get_pod(self, mock_hook, trigger):
+        """
+        Test that KubernetesPodTrigger _get_pod is called with the correct arguments.
+        """
+
+        mock_hook.get_pod.return_value = self._mock_pod_result(mock.MagicMock())
+
+        await trigger._get_pod()
+        mock_hook.get_pod.assert_called_with(name=POD_NAME, namespace=NAMESPACE)
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "exc_count, call_count",
+        [
+            pytest.param(0, 1, id="no exception"),
+            pytest.param(2, 3, id="2 exc, 1 success"),
+            pytest.param(3, 3, id="max retries"),
+        ],
+    )
+    @mock.patch(f"{TRIGGER_PATH}.hook")
+    async def test__get_pod_retries(
+        self,
+        mock_hook,
+        trigger,
+        exc_count,
+        call_count,
+    ):
+        """
+        Test that KubernetesPodTrigger _get_pod retries in case of an exception during
+        the hook.get_pod call.
+        """
+
+        side_effects = [Exception("Test exception") for _ in range(exc_count)] + [MagicMock()]
+
+        mock_hook.get_pod.side_effect = mock.AsyncMock(side_effect=side_effects)
+        # We expect the exception to be raised only if the number of retries is exceeded
+        context = (
+            pytest.raises(Exception, match="Test exception") if exc_count > 2 else contextlib.nullcontext()
+        )
+        with context:
+            await trigger._get_pod()
+        assert mock_hook.get_pod.call_count == call_count