feat: add liveness probe for celery workers

charts/airflow/templates/worker/worker-statefulset.yaml

@@ -150,6 +150,44 @@ spec:
                         time.sleep(10)
                         active_tasks = i.active()[local_celery_host]
           {{- end }}
+          {{- if .Values.workers.livenessProbe.enabled }}
+          livenessProbe:
+            initialDelaySeconds: {{ .Values.workers.livenessProbe.initialDelaySeconds }}
+            timeoutSeconds: {{ .Values.workers.livenessProbe.timeoutSeconds }}
+            failureThreshold: {{ .Values.workers.livenessProbe.failureThreshold }}
+            periodSeconds: {{ .Values.workers.livenessProbe.periodSeconds }}
+            exec:
+              command:
+                {{- include "airflow.command" . | indent 16 }}
+                - "python"
+                - "-Wignore"
+                - "-c"
+                - |
+                  import os
+                  import sys
+                  import subprocess
+                  from celery import Celery
+                  from celery.app.control import Inspect
+                  from typing import List
+
+                  def run_command(cmd: List[str]) -> str:
+                      process = subprocess.Popen(cmd, stdout=subprocess.PIPE)
+                      output, error = process.communicate()
+                      if error is not None:
+                          raise Exception(error)
+                      else:
+                          return output.decode(encoding="utf-8")
+
+                  broker_url = run_command(["bash", "-c", "eval $AIRFLOW__CELERY__BROKER_URL_CMD"])
+                  local_celery_host = f"celery@{os.environ['HOSTNAME']}"
+                  app = Celery(broker=broker_url)
+
+                  # ping the local celery worker to see if it's ok
+                  i = Inspect(app=app, destination=[local_celery_host], timeout=5.0)
+                  ping_responses = i.ping()
+                  if local_celery_host not in ping_responses:
+                      sys.exit(f"celery worker '{local_celery_host}' did not respond to ping")
+          {{- end }}
           ports:
             - name: wlog
               containerPort: 8793

charts/airflow/values.yaml

@@ -992,6 +992,15 @@ workers:
     ##
     intervalSeconds: 900
 
+  ## configs for the worker Pods' liveness probe
+  ##
+  livenessProbe:
+    enabled: true
+    initialDelaySeconds: 10
+    periodSeconds: 30
+    timeoutSeconds: 60
+    failureThreshold: 5
+
   ## extra pip packages to install in the worker Pod
   ##
   ## ____ EXAMPLE _______________