[airflow] branch v1-10-test updated: Make grace_period_seconds option on K8sPodOperator (#10727)

[di...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

dimberman pushed a commit to branch v1-10-test
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/v1-10-test by this push:
     new 6b86255  Make grace_period_seconds option on K8sPodOperator (#10727)
6b86255 is described below

commit 6b86255e5952f2ff3a37cb832457132c99ee40f9
Author: Daniel Imberman <da...@gmail.com>
AuthorDate: Fri Sep 4 11:35:50 2020 -0700

    Make grace_period_seconds option on K8sPodOperator (#10727)
    
    * Make grace_period_seconds option on K8sPodOperator
    
    This PR allows users to choose whether they want to gracefully kill
    pods when they delete tasks in the UI or if they would like to
    immediately kill them.
    
    * Update airflow/providers/cncf/kubernetes/operators/kubernetes_pod.py
    
    Co-authored-by: Kaxil Naik <ka...@gmail.com>
    
    Co-authored-by: Kaxil Naik <ka...@gmail.com>
---
 airflow/contrib/operators/kubernetes_pod_operator.py | 14 ++++++++++++--
 kubernetes_tests/test_kubernetes_pod_operator.py     |  4 ++--
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/airflow/contrib/operators/kubernetes_pod_operator.py b/airflow/contrib/operators/kubernetes_pod_operator.py
index fb984c0..55856143 100644
--- a/airflow/contrib/operators/kubernetes_pod_operator.py
+++ b/airflow/contrib/operators/kubernetes_pod_operator.py
@@ -141,6 +141,11 @@ class KubernetesPodOperator(BaseOperator):  # pylint: disable=too-many-instance-
     :type do_xcom_push: bool
     :param pod_template_file: path to pod template file
     :type pod_template_file: str
+    :param priority_class_name: priority class name for the launched Pod
+    :type priority_class_name: str
+    :param termination_grace_period: Termination grace period if task killed in UI,
+        defaults to kubernetes default
+    :type termination_grace_period: int
     """
     template_fields = ('image', 'cmds', 'arguments', 'env_vars', 'config_file', 'pod_template_file')
 
@@ -184,6 +189,7 @@ class KubernetesPodOperator(BaseOperator):  # pylint: disable=too-many-instance-
                  do_xcom_push=False,
                  pod_template_file=None,
                  priority_class_name=None,
+                 termination_grace_period=None,
                  *args,
                  **kwargs):
         if kwargs.get('xcom_push') is not None:
@@ -211,7 +217,7 @@ class KubernetesPodOperator(BaseOperator):  # pylint: disable=too-many-instance-
         self.node_selectors = node_selectors or {}
         self.annotations = annotations or {}
         self.affinity = affinity or {}
-        self.resources = self._set_resources(resources)
+        self.resources = self._set_resources(resources)  # noqa
         self.config_file = config_file
         self.image_pull_secrets = image_pull_secrets
         self.service_account_name = service_account_name
@@ -229,6 +235,7 @@ class KubernetesPodOperator(BaseOperator):  # pylint: disable=too-many-instance-
         self.pod_template_file = pod_template_file
         self.priority_class_name = priority_class_name
         self.name = self._set_name(name)
+        self.termination_grace_period = termination_grace_period
         self.client = None
 
     @staticmethod
@@ -453,4 +460,7 @@ class KubernetesPodOperator(BaseOperator):  # pylint: disable=too-many-instance-
             pod = self.pod
             namespace = pod.metadata.namespace
             name = pod.metadata.name
-            self.client.delete_namespaced_pod(name=name, namespace=namespace, grace_period_seconds=0)
+            kwargs = {}
+            if self.termination_grace_period is not None:
+                kwargs = {"grace_period_seconds": self.termination_grace_period}
+            self.client.delete_namespaced_pod(name=name, namespace=namespace, **kwargs)
diff --git a/kubernetes_tests/test_kubernetes_pod_operator.py b/kubernetes_tests/test_kubernetes_pod_operator.py
index efc48b8..4364885 100644
--- a/kubernetes_tests/test_kubernetes_pod_operator.py
+++ b/kubernetes_tests/test_kubernetes_pod_operator.py
@@ -1027,6 +1027,7 @@ class TestKubernetesPodOperatorSystem(unittest.TestCase):
             task_id=name,
             in_cluster=False,
             do_xcom_push=False,
+            termination_grace_period=0,
         )
         context = create_context(k)
         monitor_mock.return_value = (State.SUCCESS, None)
@@ -1036,7 +1037,6 @@ class TestKubernetesPodOperatorSystem(unittest.TestCase):
         self.assertEqual(pod.status.phase, "Running")
         k.on_kill()
         with self.assertRaises(ApiException):
-            # pod should be deleted
-            client.read_namespaced_pod(name=name, namespace=namespace)
+            pod = client.read_namespaced_pod(name=name, namespace=namespace)
 
 # pylint: enable=unused-argument