Repository: apache/airflow

Issue:
Incorrect await_container_completion in KubernetesPodOperator
### Apache Airflow version

2.3.4

### What happened

The [await_container_completion](
https://github.com/apache/airflow/blob/2.4.0/airflow/providers/cncf/kubernetes/utils/pod_manager.py#L259) function has a negated condition

```
while not self.container_is_running(pod=pod, container_name=container_name):
```
that causes our Airflow tasks running <1 s to never be completed, causing an infinite loop. 

I see this was addressed and released in https://github.com/apache/airflow/pull/23883, but later reverted in https://github.com/apache/airflow/pull/24474. How come it was reverted? The thread on that revert PR with comments from @jedcunningham and @potiuk didn't really address why the fix was reverted.

### What you think should happen instead

Pods finishing within 1s should be properly handled.

### How to reproduce

_No response_

### Operating System

Linux

### Versions of Apache Airflow Providers

apache-airflow-providers-cncf-kubernetes==4.3.0

### Deployment

Other 3rd-party Helm chart

### Deployment details

_No response_

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)

Incorrect await_container_completion in KubernetesPodOperator
### Apache Airflow version

2.3.4

### What happened

The [await_container_completion](
https://github.com/apache/airflow/blob/2.4.0/airflow/providers/cncf/kubernetes/utils/pod_manager.py#L259) function has a negated condition

```
while not self.container_is_running(pod=pod, container_name=container_name):
```
that causes our Airflow tasks running <1 s to never be completed, causing an infinite loop. 

I see this was addressed and released in https://github.com/apache/airflow/pull/23883, but later reverted in https://github.com/apache/airflow/pull/24474. How come it was reverted? The thread on that revert PR with comments from @jedcunningham and @potiuk didn't really address why the fix was reverted.

### What you think should happen instead

Pods finishing within 1s should be properly handled.

### How to reproduce

_No response_

### Operating System

Linux

### Versions of Apache Airflow Providers

apache-airflow-providers-cncf-kubernetes==4.3.0

### Deployment

Other 3rd-party Helm chart

### Deployment details

_No response_

### Anything else

_No response_

### Are you willing to submit PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)



Discussion:
It caused failure of tests and we were unable to point the source so we had to revert
https://github.com/apache/airflow/pull/24479

https://apache-airflow.slack.com/archives/CCPRP7943/p1655306336391429
https://apache-airflow.slack.com/archives/CCPRP7943/p1655300699968739

Quoting Ash from slack:
> That'll learn us for merging a bug fix without any new tests  

Feel free to raise PR with tests :)
> It caused failure of tests and we were unable to point the source so we had to revert #24479
> 
> https://apache-airflow.slack.com/archives/CCPRP7943/p1655306336391429 https://apache-airflow.slack.com/archives/CCPRP7943/p1655300699968739
> 
> Quoting Ash from slack:
> 
> > That'll learn us for merging a bug fix without any new tests
> 
> Feel free to raise PR with tests :)

I see, thanks for the links. Though it's quite a severe bug if the condition incorrectly is negated when waiting for completed pods. Let me read up on the threads.
It caused failure of tests and we were unable to point the source so we had to revert
https://github.com/apache/airflow/pull/24479

https://apache-airflow.slack.com/archives/CCPRP7943/p1655306336391429
https://apache-airflow.slack.com/archives/CCPRP7943/p1655300699968739

Quoting Ash from slack:
> That'll learn us for merging a bug fix without any new tests  

Feel free to raise PR with tests :)
> It caused failure of tests and we were unable to point the source so we had to revert #24479
> 
> https://apache-airflow.slack.com/archives/CCPRP7943/p1655306336391429 https://apache-airflow.slack.com/archives/CCPRP7943/p1655300699968739
> 
> Quoting Ash from slack:
> 
> > That'll learn us for merging a bug fix without any new tests
> 
> Feel free to raise PR with tests :)

I see, thanks for the links. Though it's quite a severe bug if the condition incorrectly is negated when waiting for completed pods. Let me read up on the threads.

Resolving patch:
diff --git a/airflow/providers/cncf/kubernetes/utils/pod_manager.py b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
--- a/airflow/providers/cncf/kubernetes/utils/pod_manager.py
+++ b/airflow/providers/cncf/kubernetes/utils/pod_manager.py
@@ -265,7 +265,13 @@ def consume_logs(*, since_time: DateTime | None = None, follow: bool = True) ->
                 time.sleep(1)
 
     def await_container_completion(self, pod: V1Pod, container_name: str) -> None:
-        while not self.container_is_running(pod=pod, container_name=container_name):
+        """
+        Waits 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
+        """
+        while self.container_is_running(pod=pod, container_name=container_name):
             time.sleep(1)
 
     def await_pod_completion(self, pod: V1Pod) -> V1Pod:


Test patch:
diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py
--- a/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -1025,7 +1025,6 @@ def test_pod_name(self):
                 do_xcom_push=False,
             )
 
-    @mock.patch(f"{POD_MANAGER_CLASS}.await_pod_completion", new=MagicMock)
     def test_on_kill(self):
         hook = KubernetesHook(conn_id=None, in_cluster=False)
         client = hook.core_v1_client
@@ -1045,8 +1044,20 @@ def test_on_kill(self):
             termination_grace_period=0,
         )
         context = create_context(k)
-        with pytest.raises(AirflowException):
-            k.execute(context)
+
+        class ShortCircuitException(Exception):
+            pass
+
+        # use this mock to short circuit and NOT wait for container completion
+        with mock.patch.object(
+            k.pod_manager, "await_container_completion", side_effect=ShortCircuitException()
+        ):
+            # cleanup will be upset since the pod should not be completed.. so skip it
+            with mock.patch.object(k, "cleanup"):
+                with pytest.raises(ShortCircuitException):
+                    k.execute(context)
+
+        # when we get here, the pod should still be running
         name = k.pod.metadata.name
         pod = client.read_namespaced_pod(name=name, namespace=namespace)
         assert pod.status.phase == "Running"