diff --git a/chart/files/pod-template-file.kubernetes-helm-yaml b/chart/files/pod-template-file.kubernetes-helm-yaml index 99c617b8b292b..bad917ae2bf44 100644 --- a/chart/files/pod-template-file.kubernetes-helm-yaml +++ b/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" . }} diff --git a/chart/values.schema.json b/chart/values.schema.json index 18f2edef8b99c..0050465815b79 100644 --- a/chart/values.schema.json +++ b/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 } } } diff --git a/chart/values.yaml b/chart/values.yaml index 3c405709c7a6e..802efe6e66482 100644 --- a/chart/values.yaml +++ b/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 diff --git a/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py b/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py index 1ed5a2ad09e67..15324d6cb5108 100644 --- a/helm-tests/tests/helm_tests/airflow_aux/test_pod_template_file.py +++ b/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, ) diff --git a/helm-tests/tests/helm_tests/airflow_core/test_worker.py b/helm-tests/tests/helm_tests/airflow_core/test_worker.py index 4a9e611b92ce0..325695d9631e9 100644 --- a/helm-tests/tests/helm_tests/airflow_core/test_worker.py +++ b/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.""" diff --git a/helm-tests/tests/helm_tests/airflow_core/test_worker_sets.py b/helm-tests/tests/helm_tests/airflow_core/test_worker_sets.py index d9a6aeb5f89f3..58107647491e1 100644 --- a/helm-tests/tests/helm_tests/airflow_core/test_worker_sets.py +++ b/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"], )