Repository: apache/airflow
Pull request: Add Kubernetes cleanup-pods CLI command for Helm Chart
closes: https://github.com/apache/airflow/issues/11146

Adds a CLI command `airflow kubernetes cleanup-pods --namespace NAMESPACE`
to Clean up Kubernetes pods in evicted/failed/succeeded states for the Helm Chart

cc @alikhtag 

<!--
Thank you for contributing! Please make sure that your code changes
are covered with tests. And in case of new features or big changes
remember to adjust the documentation.

Feel free to ping committers for the review!

In case of existing issue, reference it using one of the following:

closes: #ISSUE
related: #ISSUE

How to write a good git commit message:
http://chris.beams.io/posts/git-commit/
-->

---
**^ Add meaningful description above**

Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/master/CONTRIBUTING.rst#pull-request-guidelines)** for more information.
In case of fundamental code change, Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvements+Proposals)) is needed.
In case of a new dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
In case of backwards incompatible changes please leave a note in [UPDATING.md](https://github.com/apache/airflow/blob/master/UPDATING.md).


Problem statement:
Cleanup CronJob fails
**Apache Airflow version**: 1.10.12

**Kubernetes version** : 1.17.9

**Environment**:

- **Cloud provider or hardware configuration**: Azure Kubernetes Service
- **OS** (e.g. from /etc/os-release): Linux
- **Install tools**: Helm chart (/chart in repo) 
- **Others**:

**What happened**:
Pod cannot run since `airflow-cleanup-pods` is not in $PATH
```
Error: failed to start container "airflow-cleanup-pods": Error response from daemon: OCI runtime create failed: container_linux.go:349: starting container process caused "exec: \"airflow-cleanup-pods\": executable file not found in $PATH": unknown
```

**What you expected to happen**: 
Cleanup job pod runs

<!-- What do you think went wrong? -->

**How to reproduce it**:
* Deploy Helm Chart
* Manually trigger `airflow-cleanup ` CronJob 
* See it cannot run since `airflow-cleanup-pods` is not in $PATH

Hints from issue discussion:
Thanks for opening your first issue here! Be sure to follow the issue template!

Found something in `github`, I assume these scripts were not migrated when `Helm` chart was donated by `Astronomer` ? 
https://github.com/astronomer/astronomer-airflow-scripts
Is it possible to add them? If not, then this `CronJob` may have to be removed. 

thanks for raising the issue, we will take a look into this one.
@kaxil I assigned you to this issue. Is it ok?
@kaxil Any progress on this ticket? I can see Helm Chart still using this script which is not available in the official image.
https://github.com/apache/airflow/blob/3970bfad4c878d99adce80a4bfd824a15132a161/chart/templates/cleanup/cleanup-cronjob.yaml#L64
I think it won't be a big problem to add this script as a command to the CLI.
https://github.com/apache/airflow/blob/3970bfad4c878d99adce80a4bfd824a15132a161/airflow/cli/cli_parser.py#L1276
What do you think about it?

Commit message: Add Kubernetes cleanup-pods CLI command for Helm Chart

closes: https://github.com/apache/airflow/issues/11146

Adds a CLI command `airflow kubernetes cleanup-pods --namespace NAMESPACE`
to Clean up Kubernetes pods in evicted/failed/succeeded states for the Helm Chart

Patch to review:
diff --git a/airflow/cli/cli_parser.py b/airflow/cli/cli_parser.py
index aa4d3b6eb6993..842b23233a0b0 100644
--- a/airflow/cli/cli_parser.py
+++ b/airflow/cli/cli_parser.py
@@ -759,6 +759,12 @@ def positive_int(value):
     help="The option name",
 )
 
+# kubernetes cleanup-pods
+ARG_NAMESPACE = Arg(
+    ("--namespace",),
+    default='default',
+    help="Kubernetes Namespace",
+)
 
 ALTERNATIVE_CONN_SPECS_ARGS = [
     ARG_CONN_TYPE, ARG_CONN_HOST, ARG_CONN_LOGIN, ARG_CONN_PASSWORD, ARG_CONN_SCHEMA, ARG_CONN_PORT
@@ -1273,6 +1279,12 @@ class GroupCommand(NamedTuple):
 )
 
 KUBERNETES_COMMANDS = (
+    ActionCommand(
+        name='cleanup-pods',
+        help="Clean up Kubernetes pods in evicted/failed/succeeded states",
+        func=lazy_load_command('airflow.cli.commands.kubernetes_command.cleanup_pods'),
+        args=(ARG_NAMESPACE, ),
+    ),
     ActionCommand(
         name='generate-dag-yaml',
         help="Generate YAML files for all tasks in DAG. Useful for debugging tasks without "
diff --git a/airflow/cli/commands/kubernetes_command.py b/airflow/cli/commands/kubernetes_command.py
index 71517868e4277..615d0cf45d523 100644
--- a/airflow/cli/commands/kubernetes_command.py
+++ b/airflow/cli/commands/kubernetes_command.py
@@ -16,9 +16,12 @@
 # under the License.
 """Kubernetes sub-commands"""
 import os
+import sys
 
 import yaml
+from kubernetes import client, config
 from kubernetes.client.api_client import ApiClient
+from kubernetes.client.rest import ApiException
 
 from airflow.executors.kubernetes_executor import KubeConfig, create_pod_id
 from airflow.kubernetes import pod_generator
@@ -41,8 +44,7 @@ def generate_pod_yaml(args):
         pod = PodGenerator.construct_pod(
             dag_id=args.dag_id,
             task_id=ti.task_id,
-            pod_id=create_pod_id(
-                args.dag_id, ti.task_id),
+            pod_id=create_pod_id(args.dag_id, ti.task_id),
             try_number=ti.try_number,
             kube_image=kube_config.kube_image,
             date=ti.execution_date,
@@ -50,7 +52,7 @@ def generate_pod_yaml(args):
             pod_override_object=PodGenerator.from_obj(ti.executor_config),
             scheduler_job_id="worker-config",
             namespace=kube_config.executor_namespace,
-            base_worker_pod=PodGenerator.deserialize_model_file(kube_config.pod_template_file)
+            base_worker_pod=PodGenerator.deserialize_model_file(kube_config.pod_template_file),
         )
         pod_mutation_hook(pod)
         api_client = ApiClient()
@@ -61,3 +63,64 @@ def generate_pod_yaml(args):
             sanitized_pod = api_client.sanitize_for_serialization(pod)
             output.write(yaml.dump(sanitized_pod))
     print(f"YAML output can be found at {yaml_output_path}/airflow_yaml_output/")
+
+
+@cli_utils.action_logging
+def cleanup_pods(args):
+    """Clean up k8s pods in evicted/failed/succeeded states"""
+    namespace = args.namespace
+
+    # https://kubernetes.io/docs/concepts/workloads/pods/pod-lifecycle/
+    # All Containers in the Pod have terminated in success, and will not be restarted.
+    pod_succeeded = 'succeeded'
+
+    # All Containers in the Pod have terminated, and at least one Container has terminated in failure.
+    # That is, the Container either exited with non-zero status or was terminated by the system.
+    pod_failed = 'failed'
+
+    # https://kubernetes.io/docs/tasks/administer-cluster/out-of-resource/
+    pod_reason_evicted = 'evicted'
+    # If pod is failed and restartPolicy is:
+    # * Always: Restart Container; Pod phase stays Running.
+    # * OnFailure: Restart Container; Pod phase stays Running.
+    # * Never: Pod phase becomes Failed.
+    pod_restart_policy_never = 'never'
+
+    print('Loading Kubernetes configuration')
+    config.load_incluster_config()
+    core_v1 = client.CoreV1Api()
+    print(f'Listing pods in namespace {namespace}')
+    pod_list = core_v1.list_namespaced_pod(namespace)
+
+    for pod in pod_list.items:
+        pod_name = pod.metadata.name
+        print(f'Inspecting pod {pod_name}')
+        pod_phase = pod.status.phase.lower()
+        pod_reason = pod.status.reason.lower() if pod.status.reason else ''
+        pod_restart_policy = pod.spec.restart_policy.lower()
+
+        if (
+            pod_phase == pod_succeeded
+            or (pod_phase == pod_failed and pod_restart_policy == pod_restart_policy_never)
+            or (pod_reason == pod_reason_evicted)
+        ):
+
+            print(
+                f'Deleting pod "{pod_name}" phase "{pod_phase}" and reason "{pod_reason}", '
+                f'restart policy "{pod_restart_policy}"'
+            )
+            try:
+                _delete_pod(pod.metadata.name, namespace)
+            except ApiException as e:
+                print(f"can't remove POD: {e}", file=sys.stderr)
+            continue
+        print(f'No action taken on pod {pod_name}')
+
+
+def _delete_pod(name, namespace):
+    """Helper Function for cleanup_pods"""
+    core_v1 = client.CoreV1Api()
+    delete_options = client.V1DeleteOptions()
+    print(f'Deleting POD "{name}" from "{namespace}" namespace')
+    api_response = core_v1.delete_namespaced_pod(name=name, namespace=namespace, body=delete_options)
+    print(api_response)
diff --git a/chart/templates/cleanup/cleanup-cronjob.yaml b/chart/templates/cleanup/cleanup-cronjob.yaml
index d1268aeac9def..fd253d2b27dbb 100644
--- a/chart/templates/cleanup/cleanup-cronjob.yaml
+++ b/chart/templates/cleanup/cleanup-cronjob.yaml
@@ -61,7 +61,7 @@ spec:
               image: {{ template "default_airflow_image" . }}
               imagePullPolicy: {{ .Values.images.airflow.pullPolicy }}
               # Don't use entry point here, we don't need to wait on pg-bouncer etc being available.
-              command: ["airflow-cleanup-pods", "--namespace={{ .Release.Namespace }}"]
+              args: ["kubernetes", "cleanup-pods", "--namespace={{ .Release.Namespace }}"]
               env:
               {{- include "standard_airflow_environment" . | indent 12 }}
 {{- end }}
diff --git a/tests/cli/commands/test_kubernetes_command.py b/tests/cli/commands/test_kubernetes_command.py
index 0cb6069692152..fdc15d087d224 100644
--- a/tests/cli/commands/test_kubernetes_command.py
+++ b/tests/cli/commands/test_kubernetes_command.py
@@ -18,13 +18,16 @@
 import os
 import tempfile
 import unittest
+from unittest import mock
+from unittest.mock import MagicMock
+
+import kubernetes
 
 from airflow.cli import cli_parser
 from airflow.cli.commands import kubernetes_command
 
 
 class TestGeneratteDagYamlCommand(unittest.TestCase):
-
     @classmethod
     def setUpClass(cls):
         cls.parser = cli_parser.get_parser()
@@ -40,3 +43,90 @@ def test_generate_dag_yaml(self):
             self.assertEqual(len(os.listdir(out_dir)), 6)
             self.assertTrue(os.path.isfile(out_dir + file_name))
             self.assertGreater(os.stat(out_dir + file_name).st_size, 0)
+
+    @mock.patch('kubernetes.client.CoreV1Api.delete_namespaced_pod')
+    def test_delete_pod(self, delete_namespaced_pod):
+        kubernetes_command._delete_pod('dummy', 'awesome-namespace')
+        delete_namespaced_pod.assert_called_with(body=mock.ANY, name='dummy', namespace='awesome-namespace')
+
+    @mock.patch('airflow.cli.commands.kubernetes_command._delete_pod')
+    @mock.patch('kubernetes.client.CoreV1Api.list_namespaced_pod')
+    @mock.patch('kubernetes.config.load_incluster_config')
+    def test_cleanup_succeeded_pods(self, load_incluster_config, list_namespaced_pod, delete_pod):
+        pod1 = MagicMock()
+        pod1.metadata.name = 'dummy'
+        pod1.status.phase = 'Succeeded'
+        pod1.status.reason = None
+        list_namespaced_pod().items = [pod1]
+        kubernetes_command.cleanup_pods(
+            self.parser.parse_args(['kubernetes', 'cleanup-pods', '--namespace', 'awesome-namespace'])
+        )
+        delete_pod.assert_called_with('dummy', 'awesome-namespace')
+        load_incluster_config.assert_called_once()
+
+    @mock.patch('airflow.cli.commands.kubernetes_command._delete_pod')
+    @mock.patch('kubernetes.client.CoreV1Api.list_namespaced_pod')
+    @mock.patch('kubernetes.config.load_incluster_config')
+    def test_no_cleanup_failed_pods_wo_restart_policy_never(
+        self, load_incluster_config, list_namespaced_pod, delete_pod
+    ):
+        pod1 = MagicMock()
+        pod1.metadata.name = 'dummy2'
+        pod1.status.phase = 'Failed'
+        pod1.status.reason = None
+        pod1.spec.restart_policy = 'Always'
+        list_namespaced_pod().items = [pod1]
+        kubernetes_command.cleanup_pods(
+            self.parser.parse_args(['kubernetes', 'cleanup-pods', '--namespace', 'awesome-namespace'])
+        )
+        delete_pod.assert_not_called()
+        load_incluster_config.assert_called_once()
+
+    @mock.patch('airflow.cli.commands.kubernetes_command._delete_pod')
+    @mock.patch('kubernetes.client.CoreV1Api.list_namespaced_pod')
+    @mock.patch('kubernetes.config.load_incluster_config')
+    def test_cleanup_failed_pods_w_restart_policy_never(
+        self, load_incluster_config, list_namespaced_pod, delete_pod
+    ):
+        pod1 = MagicMock()
+        pod1.metadata.name = 'dummy3'
+        pod1.status.phase = 'Failed'
+        pod1.status.reason = None
+        pod1.spec.restart_policy = 'Never'
+        list_namespaced_pod().items = [pod1]
+        kubernetes_command.cleanup_pods(
+            self.parser.parse_args(['kubernetes', 'cleanup-pods', '--namespace', 'awesome-namespace'])
+        )
+        delete_pod.assert_called_with('dummy3', 'awesome-namespace')
+        load_incluster_config.assert_called_once()
+
+    @mock.patch('airflow.cli.commands.kubernetes_command._delete_pod')
+    @mock.patch('kubernetes.client.CoreV1Api.list_namespaced_pod')
+    @mock.patch('kubernetes.config.load_incluster_config')
+    def test_cleanup_evicted_pods(self, load_incluster_config, list_namespaced_pod, delete_pod):
+        pod1 = MagicMock()
+        pod1.metadata.name = 'dummy4'
+        pod1.status.phase = 'Failed'
+        pod1.status.reason = 'Evicted'
+        pod1.spec.restart_policy = 'Never'
+        list_namespaced_pod().items = [pod1]
+        kubernetes_command.cleanup_pods(
+            self.parser.parse_args(['kubernetes', 'cleanup-pods', '--namespace', 'awesome-namespace'])
+        )
+        delete_pod.assert_called_with('dummy4', 'awesome-namespace')
+        load_incluster_config.assert_called_once()
+
+    @mock.patch('airflow.cli.commands.kubernetes_command._delete_pod')
+    @mock.patch('kubernetes.client.CoreV1Api.list_namespaced_pod')
+    @mock.patch('kubernetes.config.load_incluster_config')
+    def test_cleanup_api_exception_continue(self, load_incluster_config, list_namespaced_pod, delete_pod):
+        delete_pod.side_effect = kubernetes.client.rest.ApiException(status=0)
+        pod1 = MagicMock()
+        pod1.metadata.name = 'dummy'
+        pod1.status.phase = 'Succeeded'
+        pod1.status.reason = None
+        list_namespaced_pod().items = [pod1]
+        kubernetes_command.cleanup_pods(
+            self.parser.parse_args(['kubernetes', 'cleanup-pods', '--namespace', 'awesome-namespace'])
+        )
+        load_incluster_config.assert_called_once()


Code review comments:
diff_hunk: @@ -61,3 +63,64 @@ def generate_pod_yaml(args):
             sanitized_pod = api_client.sanitize_for_serialization(pod)
             output.write(yaml.dump(sanitized_pod))
     print(f"YAML output can be found at {yaml_output_path}/airflow_yaml_output/")
+
+
+@cli_utils.action_logging
+def cleanup_pods(args):
+    """Clean up k8s pods in evicted/failed/succeeded states"""
+    namespace = args.namespace
+
+    # https://kubernetes.io/docs/concepts/workloads/pods/pod-lifecycle/
+    # All Containers in the Pod have terminated in success, and will not be restarted.
+    pod_succeeded = 'succeeded'
+
+    # All Containers in the Pod have terminated, and at least one Container has terminated in failure.
+    # That is, the Container either exited with non-zero status or was terminated by the system.
+    pod_failed = 'failed'
+
+    # https://kubernetes.io/docs/tasks/administer-cluster/out-of-resource/
+    pod_reason_evicted = 'evicted'
+    # If pod is failed and restartPolicy is:
+    # * Always: Restart Container; Pod phase stays Running.
+    # * OnFailure: Restart Container; Pod phase stays Running.
+    # * Never: Pod phase becomes Failed.
+    pod_restart_policy_never = 'never'
+
+    print('Loading Kubernetes configuration')
+    config.load_incluster_config()
path: airflow/cli/commands/kubernetes_command.py
text: @user1:
Should we use external config also? in particular, I would like to use this command in Breeze

@user2:
Can we have whatever the default behaviour of `kubectl` is?

@user1:
I thought about using this function. 
https://github.com/apache/airflow/blob/b6ab0b591459e7efbd4a45e268de499b9c175bdf/airflow/kubernetes/kube_client.py#L95

@author:
Done in https://github.com/apache/airflow/pull/11802/commits/5e4a0785ea10d2cb9295423959e53f91b8a5469d

diff_hunk: @@ -61,3 +63,64 @@ def generate_pod_yaml(args):
             sanitized_pod = api_client.sanitize_for_serialization(pod)
             output.write(yaml.dump(sanitized_pod))
     print(f"YAML output can be found at {yaml_output_path}/airflow_yaml_output/")
+
+
+@cli_utils.action_logging
+def cleanup_pods(args):
+    """Clean up k8s pods in evicted/failed/succeeded states"""
+    namespace = args.namespace
+
+    # https://kubernetes.io/docs/concepts/workloads/pods/pod-lifecycle/
+    # All Containers in the Pod have terminated in success, and will not be restarted.
+    pod_succeeded = 'succeeded'
+
+    # All Containers in the Pod have terminated, and at least one Container has terminated in failure.
+    # That is, the Container either exited with non-zero status or was terminated by the system.
+    pod_failed = 'failed'
+
+    # https://kubernetes.io/docs/tasks/administer-cluster/out-of-resource/
+    pod_reason_evicted = 'evicted'
+    # If pod is failed and restartPolicy is:
+    # * Always: Restart Container; Pod phase stays Running.
+    # * OnFailure: Restart Container; Pod phase stays Running.
+    # * Never: Pod phase becomes Failed.
+    pod_restart_policy_never = 'never'
+
+    print('Loading Kubernetes configuration')
+    config.load_incluster_config()
+    core_v1 = client.CoreV1Api()
+    print(f'Listing pods in namespace {namespace}')
+    pod_list = core_v1.list_namespaced_pod(namespace)
path: airflow/cli/commands/kubernetes_command.py
text: @user1:
On large clusters, retrieving the collection of some resource types may result in very large responses that can impact the server and client. For instance, a cluster may have tens of thousands of pods, each of which is 1-2kb of encoded JSON. Retrieving all pods across all namespaces may result in a very large response (10-20MB) and consume a large amount of server resources. 
Should we use chunking?

@user2:
This is only for a specific namespace, but yes chunking would be good - do you have any examples to hand?

@user1:
```python
continue_token = None
while True:
    pod_list = core_v1.list_namespaced_pod(namespace, limit=500, _continue=continue_token)
    for pod in pod__list.items:
        # Process item
    continue_token = pod_list.metadata._continue
    if not continue_token:
        break    
```
Not tested, but should work.

@author:
Done in https://github.com/apache/airflow/pull/11802/commits/5e4a0785ea10d2cb9295423959e53f91b8a5469d