[stable/airflow] improve celery worker graceful shutdown (#23504)

Signed-off-by: Mathew Wicks <thesuperzapper@users.noreply.github.com>

stable/airflow/Chart.yaml

@@ -1,7 +1,7 @@
 apiVersion: v1
 description: Airflow is a platform to programmatically author, schedule and monitor workflows
 name: airflow
-version: 7.3.0
+version: 7.4.0
 appVersion: 1.10.10
 icon: https://airflow.apache.org/_images/pin_large.png
 home: https://airflow.apache.org/

stable/airflow/README.md

@@ -39,6 +39,7 @@ kubectl exec \
 > NOTE: for chart version numbers, see [Chart.yaml](Chart.yaml) or [helm hub](https://hub.helm.sh/charts/stable/airflow).
 
 For steps you must take when upgrading this chart, please review:
+* [v7.3.X → v7.4.0](UPGRADE.md#v73x--v740)
 * [v7.2.X → v7.3.0](UPGRADE.md#v72x--v730)
 * [v7.1.X → v7.2.0](UPGRADE.md#v71x--v720)
 * [v7.0.X → v7.1.0](UPGRADE.md#v70x--v710)
@@ -236,11 +237,12 @@ workers:
   celery:
     instances: 10
 
-    # wait until all tasks are finished before SIGTERM of Pod
+    ## wait at most 10min for running tasks to complete
     gracefullTermination: true
+    gracefullTerminationPeriod: 600
 
-  # wait AT MOST 10min for tasks on a worker to finish before SIGKILL
-  terminationPeriod: 600
+  ## how many seconds (after the 10min) to wait before SIGKILL
+  terminationPeriod: 60
 
 dags:
   git:
@@ -592,7 +594,7 @@ __Airflow Worker Values:__
 | `workers.autoscaling.*` | configs for the HorizontalPodAutoscaler of the worker Pods | `<see values.yaml>` |
 | `workers.initialStartupDelay` | the number of seconds to wait (in bash) before starting each worker container | `0` |
 | `workers.celery.*` | configs for the celery worker Pods | `<see values.yaml>` |
-| `workers.terminationPeriod` | how many seconds to wait for tasks on a worker to finish before SIGKILL | `60` |
+| `workers.terminationPeriod` | how many seconds to wait after SIGTERM before SIGKILL of the celery worker | `60` |
 | `workers.secretsDir` | directory in which to mount secrets on worker containers | `/var/airflow/secrets` |
 | `workers.secrets` | secret names which will be mounted as a file at `{workers.secretsDir}/<secret_name>` | `[]` |
 | `workers.secretsMap` | you can use secretsMap to specify a map and all the secrets will be stored within it secrets will be mounted as files at `{workers.secretsDir}/<secrets_in_map>`. If you use workers.secretsMap, then it overrides `workers.secrets`.| `""` |

stable/airflow/UPGRADE.md

@@ -1,5 +1,23 @@
 # Upgrading Steps
 
+## `v7.3.X` → `v7.4.0`
+
+__The following IMPROVEMENTS have been made:__
+
+* Reduced how likely it is for a celery worker to receive SIGKILL with graceful termination enabled.
+  * Celery worker graceful shutdown lifecycle:
+    1. prevent worker accepting new tasks
+    2. wait AT MOST `workers.celery.gracefullTerminationPeriod` for tasks to finish
+    3. send `SIGTERM` to worker
+    4. wait AT MOST `workers.terminationPeriod` for kill to finish
+    5. send `SIGKILL` to worker
+  * NOTE: 
+    * if you currently use a high value of `workers.terminationPeriod`, consider lowering it to `60` and setting a high value for `workers.celery.gracefullTerminationPeriod`
+
+__The following values have been ADDED:__
+
+* `workers.celery.gracefullTerminationPeriod`
+
 ## `v7.2.X` → `v7.3.0`
 
 __The following IMPROVEMENTS have been made:__

stable/airflow/examples/google-gke/custom-values.yaml

@@ -239,13 +239,17 @@ workers:
     ##
     instances: 10
 
-    ## if we should wait for tasks to finish on a celery worker before SIGTERM of Pod
+    ## if we should wait for tasks to finish before SIGTERM of the celery worker
     ##
     gracefullTermination: true
 
-  ## how many seconds to wait for tasks on a worker to finish before SIGKILL
+    ## how many seconds to wait for tasks to finish before SIGTERM of the celery worker
+    ##
+    gracefullTerminationPeriod: 600
+
+  ## how many seconds to wait after SIGTERM before SIGKILL of the celery worker
   ##
-  terminationPeriod: 600
+  terminationPeriod: 60
 
   ## directory in which to mount secrets on worker containers
   ##

stable/airflow/templates/configmap-scripts.yaml

@@ -36,7 +36,7 @@ data:
     # loop until all active task are finished
     echo "*** waiting for active tasks to finish"
     while (( celery inspect --broker $AIRFLOW__CELERY__BROKER_URL --destination celery@$HOSTNAME --json active | python3 -c "import json; active_tasks = json.loads(input())['celery@$HOSTNAME']; print(len(active_tasks))" > 0 )); do
-      sleep 30
+      sleep 10
     done
   preinit-db.sh: |
     #!/bin/bash

stable/airflow/templates/statefulsets-workers.yaml

@@ -60,7 +60,11 @@ spec:
         - name: {{ .Values.airflow.image.pullSecret }}
       {{- end }}
       restartPolicy: Always
+      {{- if .Values.workers.celery.gracefullTermination }}
+      terminationGracePeriodSeconds: {{ add .Values.workers.terminationPeriod .Values.workers.celery.gracefullTerminationPeriod }}
+      {{- else }}
       terminationGracePeriodSeconds: {{ .Values.workers.terminationPeriod }}
+      {{- end }}
       serviceAccountName: {{ include "airflow.serviceAccountName" . }}
       {{- if .Values.workers.nodeSelector }}
       nodeSelector:
@@ -148,7 +152,10 @@ spec:
           lifecycle:
             preStop:
               exec:
-                command: ["/home/airflow/scripts/graceful-stop-celery-worker.sh"]
+                command:
+                  - "timeout"
+                  - "{{ .Values.workers.celery.gracefullTerminationPeriod }}s"
+                  - "/home/airflow/scripts/graceful-stop-celery-worker.sh"
           {{- end}}
           envFrom:
             - configMapRef:

stable/airflow/values.yaml

@@ -530,14 +530,26 @@ workers:
     ##
     instances: 1
 
-    ## if we should wait for tasks to finish on a celery worker before SIGTERM of Pod
-    ##
-    ## NOTE:
-    ## - `workers.terminationPeriod` is still the overall timeout before worker Pods are killed using SIGKILL
+    ## if we should wait for tasks to finish before SIGTERM of the celery worker
     ##
     gracefullTermination: false
 
-  ## how many seconds to wait for tasks on a worker to finish before SIGKILL
+    ## how many seconds to wait for tasks to finish before SIGTERM of the celery worker
+    ##
+    ## graceful shutdown lifecycle:
+    ##   1. prevent worker accepting new tasks
+    ##   2. wait AT MOST `workers.celery.gracefullTerminationPeriod` for tasks to finish
+    ##   3. send SIGTERM to worker
+    ##   4. wait AT MOST `workers.terminationPeriod` for kill to finish
+    ##   5. send SIGKILL to worker
+    ##
+    gracefullTerminationPeriod: 600
+
+  ## how many seconds to wait after SIGTERM before SIGKILL of the celery worker
+  ##
+  ## WARNING:
+  ## - tasks that are still running during SIGKILL will be orphaned, this is important
+  ##   to understand with KubernetesPodOperator(), as Pods may continue running
   ##
   terminationPeriod: 60