diff --git a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py index 7ca930b3d8490..0e63e469b49f7 100644 --- a/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py +++ b/airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py @@ -373,9 +373,10 @@ def execute(self, context: 'Context'): self.await_pod_start(pod=self.pod) if self.get_logs: - self.pod_manager.follow_container_logs( + self.pod_manager.fetch_container_logs( pod=self.pod, container_name=self.BASE_CONTAINER_NAME, + follow=True, ) else: self.pod_manager.await_container_completion( diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py index 2323ae154870b..4221ac2afd95e 100644 --- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py +++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py @@ -18,7 +18,9 @@ import json import math import time +import warnings from contextlib import closing +from dataclasses import dataclass from datetime import datetime from typing import TYPE_CHECKING, Iterable, Optional, Tuple, cast @@ -80,6 +82,14 @@ def container_is_running(pod: V1Pod, container_name: str) -> bool: return container_status.state.running is not None +@dataclass +class PodLoggingStatus: + """Used for returning the status of the pod and last log time when exiting from `fetch_container_logs`""" + + running: bool + last_log_time: Optional[DateTime] + + class PodManager(LoggingMixin): """ Helper class for creating, monitoring, and otherwise interacting with Kubernetes pods @@ -166,17 +176,23 @@ def await_pod_start(self, pod: V1Pod, startup_timeout: int = 120) -> None: raise PodLaunchFailedException(msg) time.sleep(1) - def follow_container_logs(self, pod: V1Pod, container_name: str) -> None: + def follow_container_logs(self, pod: V1Pod, container_name: str) -> PodLoggingStatus: + warnings.warn( + "Method `follow_container_logs` is deprecated. Use `fetch_container_logs` instead" + "with option `follow=True`.", + DeprecationWarning, + ) + return self.fetch_container_logs(pod=pod, container_name=container_name, follow=True) + + def fetch_container_logs( + self, pod: V1Pod, container_name: str, *, follow=False, since_time: Optional[DateTime] = None + ) -> PodLoggingStatus: """ Follows the logs of container and streams to airflow logging. Returns when container exits. - - .. note:: :meth:`read_pod_logs` follows the logs, so we shouldn't necessarily *need* to - loop as we do here. But in a long-running process we might temporarily lose connectivity. - So the looping logic is there to let us resume following the logs. """ - def follow_logs(since_time: Optional[DateTime] = None) -> Optional[DateTime]: + def consume_logs(*, since_time: Optional[DateTime] = None, follow: bool = True) -> Optional[DateTime]: """ Tries to follow container logs until container completes. For a long-running container, sometimes the log read may be interrupted @@ -193,6 +209,7 @@ def follow_logs(since_time: Optional[DateTime] = None) -> Optional[DateTime]: since_seconds=( math.ceil((pendulum.now() - since_time).total_seconds()) if since_time else None ), + follow=follow, ) for line in logs: timestamp, message = self.parse_log_line(line.decode('utf-8')) @@ -205,11 +222,16 @@ def follow_logs(since_time: Optional[DateTime] = None) -> Optional[DateTime]: ) return timestamp or since_time - last_log_time = None + # note: `read_pod_logs` follows the logs, so we shouldn't necessarily *need* to + # loop as we do here. But in a long-running process we might temporarily lose connectivity. + # So the looping logic is there to let us resume following the logs. + last_log_time = since_time while True: - last_log_time = follow_logs(since_time=last_log_time) + last_log_time = consume_logs(since_time=last_log_time, follow=follow) if not self.container_is_running(pod, container_name=container_name): - return + return PodLoggingStatus(running=False, last_log_time=last_log_time) + if not follow: + return PodLoggingStatus(running=True, last_log_time=last_log_time) else: self.log.warning( 'Pod %s log read interrupted but container %s still running', @@ -270,6 +292,7 @@ def read_pod_logs( tail_lines: Optional[int] = None, timestamps: bool = False, since_seconds: Optional[int] = None, + follow=True, ) -> Iterable[bytes]: """Reads log from the POD""" additional_kwargs = {} @@ -284,7 +307,7 @@ def read_pod_logs( name=pod.metadata.name, namespace=pod.metadata.namespace, container=container_name, - follow=True, + follow=follow, timestamps=timestamps, _preload_content=False, **additional_kwargs, diff --git a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py index 1df7ea05e7a28..23f2bc145b703 100644 --- a/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py +++ b/tests/providers/cncf/kubernetes/operators/test_kubernetes_pod.py @@ -531,9 +531,9 @@ def test_pod_template_file_kwargs_override(self, randomize_name, pod_template_fi "run_id": "test", } - @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.follow_container_logs") + @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.fetch_container_logs") @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.await_container_completion") - def test_describes_pod_on_failure(self, await_container_mock, follow_container_mock): + def test_describes_pod_on_failure(self, await_container_mock, fetch_container_mock): name_base = "test" k = KubernetesPodOperator( @@ -548,7 +548,7 @@ def test_describes_pod_on_failure(self, await_container_mock, follow_container_m do_xcom_push=False, cluster_context="default", ) - follow_container_mock.return_value = None + fetch_container_mock.return_value = None remote_pod_mock = MagicMock() remote_pod_mock.status.phase = 'Failed' self.await_pod_mock.return_value = remote_pod_mock @@ -559,9 +559,9 @@ def test_describes_pod_on_failure(self, await_container_mock, follow_container_m assert not self.client_mock.return_value.read_namespaced_pod.called - @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.follow_container_logs") + @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.fetch_container_logs") @mock.patch("airflow.providers.cncf.kubernetes.utils.pod_manager.PodManager.await_container_completion") - def test_no_handle_failure_on_success(self, await_container_mock, follow_container_mock): + def test_no_handle_failure_on_success(self, await_container_mock, fetch_container_mock): name_base = "test" k = KubernetesPodOperator( @@ -577,7 +577,7 @@ def test_no_handle_failure_on_success(self, await_container_mock, follow_contain cluster_context="default", ) - follow_container_mock.return_value = None + fetch_container_mock.return_value = None remote_pod_mock = MagicMock() remote_pod_mock.status.phase = 'Succeeded' self.await_pod_mock.return_value = remote_pod_mock diff --git a/tests/providers/cncf/kubernetes/utils/test_pod_manager.py b/tests/providers/cncf/kubernetes/utils/test_pod_manager.py index b5561fe861120..e5526a73f94df 100644 --- a/tests/providers/cncf/kubernetes/utils/test_pod_manager.py +++ b/tests/providers/cncf/kubernetes/utils/test_pod_manager.py @@ -20,6 +20,8 @@ import pendulum import pytest from kubernetes.client.rest import ApiException +from pendulum import DateTime +from pendulum.tz.timezone import Timezone from urllib3.exceptions import HTTPError as BaseHTTPError from airflow.exceptions import AirflowException @@ -186,7 +188,7 @@ def pod_state_gen(): self.mock_kube_client.read_namespaced_pod.side_effect = pod_state_gen() self.mock_kube_client.read_namespaced_pod_log.return_value = iter(()) - self.pod_manager.follow_container_logs(mock.sentinel, 'base') + self.pod_manager.fetch_container_logs(mock.sentinel, 'base') def test_monitor_pod_logs_failures_non_fatal(self): mock.sentinel.metadata = mock.MagicMock() @@ -209,7 +211,7 @@ def pod_log_gen(): self.mock_kube_client.read_namespaced_pod_log.side_effect = pod_log_gen() - self.pod_manager.follow_container_logs(mock.sentinel, 'base') + self.pod_manager.fetch_container_logs(mock.sentinel, 'base') def test_read_pod_retries_fails(self): mock.sentinel.metadata = mock.MagicMock() @@ -282,6 +284,47 @@ def test_container_is_running(self, container_is_running_mock): self.pod_manager.container_is_running(None, 'base') container_is_running_mock.assert_called_with(pod=mock_pod, container_name='base') + @pytest.mark.parametrize('follow', [True, False]) + @mock.patch('airflow.providers.cncf.kubernetes.utils.pod_manager.container_is_running') + def test_fetch_container_done(self, container_running, follow): + """If container done, should exit, no matter setting of follow.""" + mock_pod = MagicMock() + container_running.return_value = False + self.mock_kube_client.read_namespaced_pod_log.return_value = [b'2021-01-01 hi'] + ret = self.pod_manager.fetch_container_logs(pod=mock_pod, container_name='base', follow=follow) + assert ret.last_log_time == DateTime(2021, 1, 1, tzinfo=Timezone('UTC')) + assert ret.running is False + + @mock.patch('pendulum.now') + @mock.patch('airflow.providers.cncf.kubernetes.utils.pod_manager.container_is_running') + def test_fetch_container_since_time(self, container_running, mock_now): + """If given since_time, should be used.""" + mock_pod = MagicMock() + mock_now.return_value = DateTime(2020, 1, 1, 0, 0, 5, tzinfo=Timezone('UTC')) + container_running.return_value = False + self.mock_kube_client.read_namespaced_pod_log.return_value = [b'2021-01-01 hi'] + since_time = DateTime(2020, 1, 1, tzinfo=Timezone('UTC')) + self.pod_manager.fetch_container_logs(pod=mock_pod, container_name='base', since_time=since_time) + args, kwargs = self.mock_kube_client.read_namespaced_pod_log.call_args_list[0] + assert kwargs['since_seconds'] == 5 + + @pytest.mark.parametrize('follow, is_running_calls, exp_running', [(True, 3, False), (False, 1, True)]) + @mock.patch('airflow.providers.cncf.kubernetes.utils.pod_manager.container_is_running') + def test_fetch_container_running_follow( + self, container_running_mock, follow, is_running_calls, exp_running + ): + """ + When called with follow, should keep looping even after disconnections, if pod still running. + When called with follow=False, should return immediately even though still running. + """ + mock_pod = MagicMock() + container_running_mock.side_effect = [True, True, False] # only will be called once + self.mock_kube_client.read_namespaced_pod_log.return_value = [b'2021-01-01 hi'] + ret = self.pod_manager.fetch_container_logs(pod=mock_pod, container_name='base', follow=follow) + assert len(container_running_mock.call_args_list) == is_running_calls + assert ret.last_log_time == DateTime(2021, 1, 1, tzinfo=Timezone('UTC')) + assert ret.running is exp_running + def params_for_test_container_is_running(): """The `container_is_running` method is designed to handle an assortment of bad objects