Add workers.celery.terminationGracePeriodSeconds & workers.kubernetes.terminationGracePeriodSeconds

Author: Miretpl

chart/files/pod-template-file.kubernetes-helm-yaml

@@ -233,7 +233,7 @@ spec:
   {{- if $schedulerName }}
   schedulerName: {{ $schedulerName }}
   {{- end }}
-  terminationGracePeriodSeconds: {{ .Values.workers.terminationGracePeriodSeconds }}
+  terminationGracePeriodSeconds: {{ .Values.workers.kubernetes.terminationGracePeriodSeconds | default .Values.workers.terminationGracePeriodSeconds }}
   tolerations: {{- toYaml $tolerations | nindent 4 }}
   topologySpreadConstraints: {{- toYaml $topologySpreadConstraints | nindent 4 }}
   serviceAccountName: {{ include "worker.serviceAccountName" . }}

chart/values.schema.json

@@ -2282,7 +2282,7 @@
                     "$ref": "#/definitions/io.k8s.api.core.v1.ResourceRequirements"
                 },
                 "terminationGracePeriodSeconds": {
-                    "description": "Grace period for tasks to finish after SIGTERM is sent from Kubernetes. It is used by Airflow Celery workers and pod-template-file.",
+                    "description": "Grace period for tasks to finish after SIGTERM is sent from Kubernetes. It is used by Airflow Celery workers and pod-template-file. Use ``workers.celery.terminationGracePeriodSeconds`` and/or ``workers.kubernetes.terminationGracePeriodSeconds`` to separate value between Celery workers and pod-template-file",
                     "type": "integer",
                     "default": 600
                 },
@@ -3275,6 +3275,14 @@
                                 }
                             ],
                             "$ref": "#/definitions/io.k8s.api.core.v1.ResourceRequirements"
+                        },
+                        "terminationGracePeriodSeconds": {
+                            "description": "Grace period for tasks to finish after SIGTERM is sent from Kubernetes.",
+                            "type": [
+                                "integer",
+                                "null"
+                            ],
+                            "default": null
                         }
                     }
                 },
@@ -3548,6 +3556,14 @@
                                 }
                             ],
                             "$ref": "#/definitions/io.k8s.api.core.v1.ResourceRequirements"
+                        },
+                        "terminationGracePeriodSeconds": {
+                            "description": "Grace period for tasks to finish after SIGTERM is sent from Kubernetes.",
+                            "type": [
+                                "integer",
+                                "null"
+                            ],
+                            "default": null
                         }
                     }
                 }

chart/values.yaml

@@ -909,6 +909,8 @@ workers:
 
   # Grace period for tasks to finish after SIGTERM is sent from kubernetes.
   # It is used by Airflow Celery workers and pod-template-file.
+  # Use workers.celery.terminationGracePeriodSeconds and/or workers.kubernetes.terminationGracePeriodSeconds
+  # to separate value between Celery workers and pod-template-file
   terminationGracePeriodSeconds: 600
 
   # This setting tells kubernetes that its ok to evict when it wants to scale a node down.
@@ -1262,6 +1264,9 @@ workers:
     #   cpu: 100m
     #   memory: 128Mi
 
+    # Grace period for tasks to finish after SIGTERM is sent from kubernetes
+    terminationGracePeriodSeconds: ~
+
   kubernetes:
     # Command to use in pod-template-file (templated)
     command: ~
@@ -1326,6 +1331,9 @@ workers:
     #   cpu: 100m
     #   memory: 128Mi
 
+    # Grace period for tasks to finish after SIGTERM is sent from kubernetes
+    terminationGracePeriodSeconds: ~
+
 # Airflow scheduler settings
 scheduler:
   enabled: true

helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py

@@ -1114,13 +1114,17 @@ def test_workers_container_lifecycle_webhooks_are_configurable(self, workers_val
             "preStop": {"exec": {"command": ["echo", "preStop", "test-release"]}}
         }
 
-    def test_termination_grace_period_seconds(self):
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"terminationGracePeriodSeconds": 123},
+            {"kubernetes": {"terminationGracePeriodSeconds": 123}},
+            {"terminationGracePeriodSeconds": 1, "kubernetes": {"terminationGracePeriodSeconds": 123}},
+        ],
+    )
+    def test_termination_grace_period_seconds(self, workers_values):
         docs = render_chart(
-            values={
-                "workers": {
-                    "terminationGracePeriodSeconds": 123,
-                },
-            },
+            values={"workers": workers_values},
             show_only=["templates/pod-template-file.yaml"],
             chart_dir=self.temp_chart_dir,
         )

helm-tests/tests/helm_tests/airflow_core/test_worker.py

@@ -1721,6 +1721,21 @@ def test_pod_management_policy_not_valid_value(self, workers_values):
                 values={"workers": workers_values}, show_only=["templates/workers/worker-deployment.yaml"]
             )
 
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {"terminationGracePeriodSeconds": 5},
+            {"celery": {"terminationGracePeriodSeconds": 5}},
+            {"terminationGracePeriodSeconds": 10, "celery": {"terminationGracePeriodSeconds": 5}},
+        ],
+    )
+    def test_termination_grace_period(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values}, show_only=["templates/workers/worker-deployment.yaml"]
+        )
+
+        assert jmespath.search("spec.template.spec.terminationGracePeriodSeconds", docs[0]) == 5
+
 
 class TestWorkerLogGroomer(LogGroomerTestBase):
     """Worker groomer."""

helm-tests/tests/helm_tests/airflow_core/test_worker_sets.py

@@ -2177,21 +2177,34 @@ def test_overwrite_resources(self, values):
             "limits": {"cpu": "3m", "memory": "4Mi"},
         }
 
-    def test_overwrite_termination_grace_period_seconds(self):
-        docs = render_chart(
-            values={
-                "workers": {
-                    "celery": {
-                        "enableDefault": False,
-                        "sets": [
-                            {
-                                "name": "test",
-                                "terminationGracePeriodSeconds": 5,
-                            }
-                        ],
-                    },
+    @pytest.mark.parametrize(
+        "workers_values",
+        [
+            {
+                "celery": {
+                    "enableDefault": False,
+                    "sets": [{"name": "test", "terminationGracePeriodSeconds": 5}],
                 }
             },
+            {
+                "terminationGracePeriodSeconds": 20,
+                "celery": {
+                    "enableDefault": False,
+                    "sets": [{"name": "test", "terminationGracePeriodSeconds": 5}],
+                },
+            },
+            {
+                "celery": {
+                    "terminationGracePeriodSeconds": 20,
+                    "enableDefault": False,
+                    "sets": [{"name": "test", "terminationGracePeriodSeconds": 5}],
+                }
+            },
+        ],
+    )
+    def test_overwrite_termination_grace_period_seconds(self, workers_values):
+        docs = render_chart(
+            values={"workers": workers_values},
             show_only=["templates/workers/worker-deployment.yaml"],
         )