From 6b9bb1d108bd75c98cc477110632f5eb18f853d8 Mon Sep 17 00:00:00 2001 From: acsierra Date: Sun, 28 Apr 2024 14:46:22 +0200 Subject: [PATCH 1/6] Handling exception getting logs when pods finish success --- .../providers/cncf/kubernetes/operators/pod.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py index 0b387352ae8a6..129dc1566e85a 100644 --- a/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/airflow/providers/cncf/kubernetes/operators/pod.py @@ -36,6 +36,7 @@ from deprecated import deprecated from kubernetes.client import CoreV1Api, V1Pod, models as k8s from kubernetes.stream import stream +from kubernetes.client.exceptions import ApiException from urllib3.exceptions import HTTPError from airflow.configuration import conf @@ -766,9 +767,15 @@ def _clean(self, event: dict[str, Any]): # Skip await_pod_completion when the event is 'timeout' due to the pod can hang # on the ErrImagePull or ContainerCreating step and it will never complete if event["status"] != "timeout": - self.pod = self.pod_manager.await_pod_completion( - self.pod, istio_enabled, self.base_container_name - ) + try: + self.pod = self.pod_manager.await_pod_completion( + self.pod, istio_enabled, self.base_container_name + ) + except ApiException as e: + if e.status == 404: + self.pod = None + else: + raise e if self.pod is not None: self.post_complete_action( pod=self.pod, @@ -796,11 +803,11 @@ def write_logs(self, pod: k8s.V1Pod, follow: bool = False, since_time: DateTime line = raw_line.decode("utf-8", errors="backslashreplace").rstrip("\n") if line: self.log.info("Container logs: %s", line) - except HTTPError as e: + except (HTTPError, ApiException) as e: self.log.warning( "Reading of logs interrupted with error %r; will retry. " "Set log level to DEBUG for traceback.", - e, + e if not isinstance(e, ApiException) else e.reason, ) def post_complete_action(self, *, pod, remote_pod, **kwargs): From 628480af1f4568258267a7888fca6c6e21a897b4 Mon Sep 17 00:00:00 2001 From: acsierra Date: Mon, 29 Apr 2024 12:56:02 +0200 Subject: [PATCH 2/6] Include test --- .../cncf/kubernetes/operators/test_pod.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py index db10bff330b01..b90f4492b39e2 100644 --- a/tests/providers/cncf/kubernetes/operators/test_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_pod.py @@ -25,6 +25,7 @@ import pendulum import pytest from kubernetes.client import ApiClient, V1Pod, V1PodSecurityContext, V1PodStatus, models as k8s +from kubernetes.client.exceptions import ApiException from urllib3 import HTTPResponse from airflow.exceptions import AirflowException, AirflowSkipException, TaskDeferred @@ -1998,6 +1999,29 @@ def test_async_write_logs_should_execute_successfully( else: mock_manager.return_value.read_pod_logs.assert_not_called() + @patch(KUB_OP_PATH.format("post_complete_action")) + @patch(KUB_OP_PATH.format("extract_xcom")) + @patch(HOOK_CLASS) + @patch(KUB_OP_PATH.format("pod_manager")) + def test_async_write_logs_handler_api_exception( + self, mock_manager, mocked_hook, mock_extract_xcom + ): + mock_manager.read_pod_logs.return_value = ApiException(status=404) + mocked_hook.return_value.get_pod.return_value = k8s.V1Pod( + metadata=k8s.V1ObjectMeta(name=TEST_NAME, namespace=TEST_NAMESPACE) + ) + mock_extract_xcom.return_value = "{}" + k = KubernetesPodOperator( + task_id="task", + get_logs=True, + deferrable=True, + ) + self.run_pod_async(k) + mock_k = k + mock_k.log.warning = MagicMock(return_value='test') + + mock_k.log.warning.assert_called_once() + @pytest.mark.parametrize( "log_pod_spec_on_failure,expect_match", [ From fb216bc1e0d101c98a487a2671b53a49abf0808e Mon Sep 17 00:00:00 2001 From: acsierra Date: Mon, 29 Apr 2024 17:22:17 +0200 Subject: [PATCH 3/6] Add log --- airflow/providers/cncf/kubernetes/operators/pod.py | 5 ++++- tests/providers/cncf/kubernetes/operators/test_pod.py | 6 ++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py index 129dc1566e85a..5152202d71c21 100644 --- a/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/airflow/providers/cncf/kubernetes/operators/pod.py @@ -35,8 +35,8 @@ import kubernetes from deprecated import deprecated from kubernetes.client import CoreV1Api, V1Pod, models as k8s -from kubernetes.stream import stream from kubernetes.client.exceptions import ApiException +from kubernetes.stream import stream from urllib3.exceptions import HTTPError from airflow.configuration import conf @@ -774,6 +774,9 @@ def _clean(self, event: dict[str, Any]): except ApiException as e: if e.status == 404: self.pod = None + self.log.info( + "Pod not found while waiting for completion. The last status was %r", event["status"] + ) else: raise e if self.pod is not None: diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py index b90f4492b39e2..8bd8c9aeb2f31 100644 --- a/tests/providers/cncf/kubernetes/operators/test_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_pod.py @@ -2003,9 +2003,7 @@ def test_async_write_logs_should_execute_successfully( @patch(KUB_OP_PATH.format("extract_xcom")) @patch(HOOK_CLASS) @patch(KUB_OP_PATH.format("pod_manager")) - def test_async_write_logs_handler_api_exception( - self, mock_manager, mocked_hook, mock_extract_xcom - ): + def test_async_write_logs_handler_api_exception(self, mock_manager, mocked_hook, mock_extract_xcom): mock_manager.read_pod_logs.return_value = ApiException(status=404) mocked_hook.return_value.get_pod.return_value = k8s.V1Pod( metadata=k8s.V1ObjectMeta(name=TEST_NAME, namespace=TEST_NAMESPACE) @@ -2018,7 +2016,7 @@ def test_async_write_logs_handler_api_exception( ) self.run_pod_async(k) mock_k = k - mock_k.log.warning = MagicMock(return_value='test') + mock_k.log.warning = MagicMock(return_value="test") mock_k.log.warning.assert_called_once() From ce7c310dab8df29a8b1f3272eec634db5cd9fb2e Mon Sep 17 00:00:00 2001 From: acsierra Date: Tue, 30 Apr 2024 09:03:04 +0200 Subject: [PATCH 4/6] Fix tests --- tests/providers/cncf/kubernetes/operators/test_pod.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py index 8bd8c9aeb2f31..9f60d16e27085 100644 --- a/tests/providers/cncf/kubernetes/operators/test_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_pod.py @@ -1711,7 +1711,7 @@ def run_pod_async(self, operator: KubernetesPodOperator, map_index: int = -1): remote_pod_mock.metadata.namespace = TEST_NAMESPACE self.await_pod_mock.return_value = remote_pod_mock - operator.execute_complete( + operator.trigger_reentry( context=context, event={ "status": "success", @@ -1999,12 +1999,15 @@ def test_async_write_logs_should_execute_successfully( else: mock_manager.return_value.read_pod_logs.assert_not_called() - @patch(KUB_OP_PATH.format("post_complete_action")) + @pytest.mark.parametrize("evaluate_status", [404, None]) @patch(KUB_OP_PATH.format("extract_xcom")) @patch(HOOK_CLASS) @patch(KUB_OP_PATH.format("pod_manager")) - def test_async_write_logs_handler_api_exception(self, mock_manager, mocked_hook, mock_extract_xcom): - mock_manager.read_pod_logs.return_value = ApiException(status=404) + def test_async_write_logs_handler_api_exception( + self, mock_manager, mocked_hook, mock_extract_xcom, evaluate_status + ): + mock_manager.read_pod_logs.return_value = ApiException(status=evaluate_status) + mock_manager.await_pod_completion.return_value = ApiException(status=404) mocked_hook.return_value.get_pod.return_value = k8s.V1Pod( metadata=k8s.V1ObjectMeta(name=TEST_NAME, namespace=TEST_NAMESPACE) ) From cd8c1baadaf72cd414e32bc0957c787aec2bc5ea Mon Sep 17 00:00:00 2001 From: acsierra Date: Wed, 1 May 2024 11:51:24 +0200 Subject: [PATCH 5/6] Change log type --- airflow/providers/cncf/kubernetes/operators/pod.py | 2 +- .../providers/cncf/kubernetes/operators/test_pod.py | 12 +++++------- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/airflow/providers/cncf/kubernetes/operators/pod.py b/airflow/providers/cncf/kubernetes/operators/pod.py index 5152202d71c21..5e8de6bebdacb 100644 --- a/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/airflow/providers/cncf/kubernetes/operators/pod.py @@ -774,7 +774,7 @@ def _clean(self, event: dict[str, Any]): except ApiException as e: if e.status == 404: self.pod = None - self.log.info( + self.log.warning( "Pod not found while waiting for completion. The last status was %r", event["status"] ) else: diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py index 9f60d16e27085..9a2b5ba6a16d1 100644 --- a/tests/providers/cncf/kubernetes/operators/test_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_pod.py @@ -2000,14 +2000,15 @@ def test_async_write_logs_should_execute_successfully( mock_manager.return_value.read_pod_logs.assert_not_called() @pytest.mark.parametrize("evaluate_status", [404, None]) + @patch(KUB_OP_PATH.format("post_complete_action")) @patch(KUB_OP_PATH.format("extract_xcom")) @patch(HOOK_CLASS) @patch(KUB_OP_PATH.format("pod_manager")) def test_async_write_logs_handler_api_exception( - self, mock_manager, mocked_hook, mock_extract_xcom, evaluate_status + self, mock_manager, mocked_hook, mock_extract_xcom, evaluate_status, post_complete_action ): - mock_manager.read_pod_logs.return_value = ApiException(status=evaluate_status) - mock_manager.await_pod_completion.return_value = ApiException(status=404) + mock_manager.read_pod_logs.side_effect = ApiException(status=evaluate_status) + mock_manager.await_pod_completion.side_effect = ApiException(status=404) mocked_hook.return_value.get_pod.return_value = k8s.V1Pod( metadata=k8s.V1ObjectMeta(name=TEST_NAME, namespace=TEST_NAMESPACE) ) @@ -2018,10 +2019,7 @@ def test_async_write_logs_handler_api_exception( deferrable=True, ) self.run_pod_async(k) - mock_k = k - mock_k.log.warning = MagicMock(return_value="test") - - mock_k.log.warning.assert_called_once() + post_complete_action.assert_not_called() @pytest.mark.parametrize( "log_pod_spec_on_failure,expect_match", From 2226c724c452733d45ec5a4768415f3f05d40aff Mon Sep 17 00:00:00 2001 From: acsierra Date: Mon, 27 May 2024 15:03:59 +0200 Subject: [PATCH 6/6] Fix test --- tests/providers/cncf/kubernetes/operators/test_pod.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/providers/cncf/kubernetes/operators/test_pod.py b/tests/providers/cncf/kubernetes/operators/test_pod.py index 4b020164d7a36..f7a8e24d6973d 100644 --- a/tests/providers/cncf/kubernetes/operators/test_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_pod.py @@ -2062,15 +2062,15 @@ def test_async_write_logs_should_execute_successfully( else: mock_manager.return_value.read_pod_logs.assert_not_called() - @pytest.mark.parametrize("evaluate_status", [404, None]) @patch(KUB_OP_PATH.format("post_complete_action")) + @patch(KUB_OP_PATH.format("client")) @patch(KUB_OP_PATH.format("extract_xcom")) @patch(HOOK_CLASS) @patch(KUB_OP_PATH.format("pod_manager")) def test_async_write_logs_handler_api_exception( - self, mock_manager, mocked_hook, mock_extract_xcom, evaluate_status, post_complete_action + self, mock_manager, mocked_hook, mock_extract_xcom, post_complete_action, mocked_client ): - mock_manager.read_pod_logs.side_effect = ApiException(status=evaluate_status) + mocked_client.read_namespaced_pod_log.side_effect = ApiException(status=404) mock_manager.await_pod_completion.side_effect = ApiException(status=404) mocked_hook.return_value.get_pod.return_value = k8s.V1Pod( metadata=k8s.V1ObjectMeta(name=TEST_NAME, namespace=TEST_NAMESPACE)