feat: Add liveness probes for celery workers

Currently pods can enter a 'zombie' state where they become disconnected from celery and will no longer pick up new jobs. We create a liveness probe to detect such pods so they can be killed by k8s and recreated Signed-off-by: Nick Wood <nwood@cloudflare.com>
airflow-helm · Aug 2, 2023 · 40dbe4e · 40dbe4e
1 parent c3339af
commit 40dbe4e
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/charts/airflow/templates/worker/worker-statefulset.yaml b/charts/airflow/templates/worker/worker-statefulset.yaml
@@ -150,6 +150,22 @@ spec:
                         time.sleep(10)
                         active_tasks = i.active()[local_celery_host]
           {{- end }}
+          {{- if .Values.workers.celery.livenessProbe.enabled }}
+          livenessProbe:
+            initialDelaySeconds: {{ .Values.workers.celery.livenessProbe.initialDelaySeconds }}
+            timeoutSeconds: {{ .Values.workers.celery.livenessProbe.timeoutSeconds }}
+            failureThreshold: {{ .Values.workers.celery.livenessProbe.failureThreshold }}
+            periodSeconds: {{ .Values.workers.celery.livenessProbe.periodSeconds }}
+            exec:
+              command:
+                {{- if .Values.workers.celery.livenessProbe.command }}
+                {{ toYaml .Values.workers.celery.livenessProbe.command | nindent 16 }}
+                {{- else}}
+                - sh
+                - -c
+                - exec /entrypoint python -m celery --app airflow.executors.celery_executor.app inspect ping -d celery@${HOSTNAME}
+                {{- end }}
+          {{- end }}
           ports:
             - name: wlog
               containerPort: 8793

diff --git a/charts/airflow/values.yaml b/charts/airflow/values.yaml
@@ -963,6 +963,14 @@ workers:
     ##
     gracefullTerminationPeriod: 600
 
+    ## config for celery worker liveness probes
+    livenessProbe:
+      enabled: false
+      initialDelaySeconds: 300
+      periodSeconds: 60
+      timeoutSeconds: 20
+      failureThreshold: 5
+
   ## 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