From bac4acd6dca42a9550c7e5e084de8014304ab40e Mon Sep 17 00:00:00 2001 From: Sumit Maheshwari Date: Mon, 21 Sep 2020 11:53:23 +0530 Subject: [PATCH] [TWTR][EWT-472] Add lifecycle support while launching worker pods (#59) --- .../kubernetes_request_factory/kubernetes_request_factory.py | 5 +++++ .../kubernetes_request_factory/pod_request_factory.py | 2 ++ airflow/contrib/kubernetes/pod.py | 4 +++- airflow/version.py | 2 +- 4 files changed, 11 insertions(+), 2 deletions(-) diff --git a/airflow/contrib/kubernetes/kubernetes_request_factory/kubernetes_request_factory.py b/airflow/contrib/kubernetes/kubernetes_request_factory/kubernetes_request_factory.py index 14d63e0df1bfe..fade44e49cb39 100644 --- a/airflow/contrib/kubernetes/kubernetes_request_factory/kubernetes_request_factory.py +++ b/airflow/contrib/kubernetes/kubernetes_request_factory/kubernetes_request_factory.py @@ -262,3 +262,8 @@ def _apply_env_from(pod, req): def extract_priority_class(pod, req): if pod.priority_class: req['spec']['priorityClassName'] = pod.priority_class + + @staticmethod + def extract_lifecycle(pod, req): + if pod.lifecycle: + req['spec']['containers'][0]['lifecycle'] = pod.lifecycle \ No newline at end of file diff --git a/airflow/contrib/kubernetes/kubernetes_request_factory/pod_request_factory.py b/airflow/contrib/kubernetes/kubernetes_request_factory/pod_request_factory.py index d8fed58f594ba..53540bfac641f 100644 --- a/airflow/contrib/kubernetes/kubernetes_request_factory/pod_request_factory.py +++ b/airflow/contrib/kubernetes/kubernetes_request_factory/pod_request_factory.py @@ -66,6 +66,7 @@ def create(self, pod): self.extract_security_context(pod, req) self.extract_dnspolicy(pod, req) self.extract_priority_class(pod, req) + self.extract_lifecycle(pod, req) return req @@ -137,4 +138,5 @@ def create(self, pod): self.extract_security_context(pod, req) self.extract_dnspolicy(pod, req) self.extract_priority_class(pod, req) + self.extract_lifecycle(pod, req) return req diff --git a/airflow/contrib/kubernetes/pod.py b/airflow/contrib/kubernetes/pod.py index b889562580651..65c313755f22d 100644 --- a/airflow/contrib/kubernetes/pod.py +++ b/airflow/contrib/kubernetes/pod.py @@ -114,7 +114,8 @@ def __init__( configmaps=None, pod_runtime_info_envs=None, dnspolicy=None, - priority_class=None + priority_class=None, + lifecycle=None ): self.image = image self.envs = envs or {} @@ -143,3 +144,4 @@ def __init__( self.pod_runtime_info_envs = pod_runtime_info_envs or [] self.dnspolicy = dnspolicy self.priority_class = priority_class + self.lifecycle = lifecycle or {} diff --git a/airflow/version.py b/airflow/version.py index 038a1058dcdff..b8a4b08d49281 100644 --- a/airflow/version.py +++ b/airflow/version.py @@ -18,5 +18,5 @@ # under the License. # -version = '1.10.4+twtr18' +version = '1.10.4+twtr19'