diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py index 657404b183065..cf8c581a2a653 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/pod.py @@ -214,6 +214,8 @@ class KubernetesPodOperator(BaseOperator): will appear as part of this task's logs if get_logs is True. Defaults to None. If None, will consult the class variable BASE_CONTAINER_NAME (which defaults to "base") for the base container name to use. + :param base_container_status_polling_interval: Polling period in seconds to check for the pod base + container status. Default to 1s. :param deferrable: Run operator in the deferrable mode. :param poll_interval: Polling period in seconds to check for the status. Used only in deferrable mode. :param log_pod_spec_on_failure: Log the pod's specification if a failure occurs @@ -288,6 +290,7 @@ def __init__( startup_check_interval_seconds: int = 5, get_logs: bool = True, base_container_name: str | None = None, + base_container_status_polling_interval: float = 1, init_container_logs: Iterable[str] | str | Literal[True] | None = None, container_logs: Iterable[str] | str | Literal[True] | None = None, image_pull_policy: str | None = None, @@ -365,6 +368,7 @@ def __init__( # Fallback to the class variable BASE_CONTAINER_NAME here instead of via default argument value # in the init method signature, to be compatible with subclasses overloading the class variable value. self.base_container_name = base_container_name or self.BASE_CONTAINER_NAME + self.base_container_status_polling_interval = base_container_status_polling_interval self.init_container_logs = init_container_logs self.container_logs = container_logs or self.base_container_name self.image_pull_policy = image_pull_policy @@ -720,7 +724,11 @@ def await_pod_completion(self, pod: k8s.V1Pod): if not self.get_logs or ( self.container_logs is not True and self.base_container_name not in self.container_logs ): - self.pod_manager.await_container_completion(pod=pod, container_name=self.base_container_name) + self.pod_manager.await_container_completion( + pod=pod, + container_name=self.base_container_name, + polling_time=self.base_container_status_polling_interval, + ) except kubernetes.client.exceptions.ApiException as exc: self._handle_api_exception(exc, pod) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py index 5a55c7d04b70a..632f5faa57ffa 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/utils/pod_manager.py @@ -619,12 +619,14 @@ def fetch_requested_container_logs( pod_logging_statuses.append(status) return pod_logging_statuses - def await_container_completion(self, pod: V1Pod, container_name: str) -> None: + def await_container_completion(self, pod: V1Pod, container_name: str, polling_time: float = 1) -> None: """ Wait for the given container in the given pod to be completed. :param pod: pod spec that will be monitored :param container_name: name of the container within the pod to monitor + :param polling_time: polling time between two container status checks. + Defaults to 1s. """ while True: remote_pod = self.read_pod(pod) @@ -632,7 +634,7 @@ def await_container_completion(self, pod: V1Pod, container_name: str) -> None: if terminated: break self.log.info("Waiting for container '%s' state to be completed", container_name) - time.sleep(1) + time.sleep(polling_time) def await_pod_completion( self, pod: V1Pod, istio_enabled: bool = False, container_name: str = "base" diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py index f3c241f8cddc2..fd38de09f37ac 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/operators/test_pod.py @@ -1537,7 +1537,9 @@ def test_get_logs_but_not_for_base_container( # check that the base container is not included in the logs mock_fetch_log.assert_called_once_with(pod=pod, containers=["some_init_container"], follow_logs=True) # check that KPO waits for the base container to complete before proceeding to extract XCom - mock_await_container_completion.assert_called_once_with(pod=pod, container_name="base") + mock_await_container_completion.assert_called_once_with( + pod=pod, container_name="base", polling_time=1 + ) # check that we wait for the xcom sidecar to start before extracting XCom mock_await_xcom_sidecar.assert_called_once_with(pod=pod) @@ -1798,7 +1800,7 @@ def test_await_container_completion_refreshes_properties_on_exception( ) else: mock_await_container_completion.assert_has_calls( - [mock.call(pod=pod, container_name=k.base_container_name)] * 3 + [mock.call(pod=pod, container_name=k.base_container_name, polling_time=1)] * 3 ) mock_read_pod.assert_called() assert client != k.client @@ -1853,7 +1855,7 @@ def test_await_container_completion_retries_on_specific_exception( k.await_pod_completion(pod) expected_call_count = len(side_effect) mock_await_container_completion.assert_has_calls( - [mock.call(pod=pod, container_name=k.base_container_name)] * expected_call_count + [mock.call(pod=pod, container_name=k.base_container_name, polling_time=1)] * expected_call_count ) @pytest.mark.parametrize(