Fix/21994 liveness probe (#22041)

apache · Mar 8, 2022 · 99cf75d · 99cf75d
1 parent d3c168c
commit 99cf75d
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 40 deletions.
diff --git a/chart/templates/scheduler/scheduler-deployment.yaml b/chart/templates/scheduler/scheduler-deployment.yaml
@@ -165,26 +165,9 @@ spec:
             periodSeconds: {{ .Values.scheduler.livenessProbe.periodSeconds }}
             exec:
               command:
-                - sh
-                - -c
-                - exec
-                - |
-                  CONNECTION_CHECK_MAX_COUNT=0 /entrypoint python -Wignore -c "
-                  import os
-                  os.environ['AIRFLOW__CORE__LOGGING_LEVEL'] = 'ERROR'
-                  os.environ['AIRFLOW__LOGGING__LOGGING_LEVEL'] = 'ERROR'
-
-                  from airflow.jobs.scheduler_job import SchedulerJob
-                  from airflow.utils.db import create_session
-                  from airflow.utils.net import get_hostname
-                  import sys
-
-                  with create_session() as session:
-                      job = session.query(SchedulerJob).filter_by(hostname=get_hostname()).order_by(
-                          SchedulerJob.latest_heartbeat.desc()).limit(1).first()
-
-                  sys.exit(0 if job.is_alive() else 1)
-                  "
+{{- with .Values.scheduler.livenessProbe.command }}
+{{ toYaml . | indent 16 }}
+{{- end }}
           {{- if and $local (not $elasticsearch) }}
           # Serve logs if we're in local mode and we don't have elasticsearch enabled.
           ports:

diff --git a/chart/templates/triggerer/triggerer-deployment.yaml b/chart/templates/triggerer/triggerer-deployment.yaml
@@ -168,26 +168,9 @@ spec:
             periodSeconds: {{ .Values.triggerer.livenessProbe.periodSeconds }}
             exec:
               command:
-                - sh
-                - -c
-                - exec
-                - |
-                  CONNECTION_CHECK_MAX_COUNT=0 /entrypoint python -Wignore -c "
-                  import os
-                  os.environ['AIRFLOW__CORE__LOGGING_LEVEL'] = 'ERROR'
-                  os.environ['AIRFLOW__LOGGING__LOGGING_LEVEL'] = 'ERROR'
-
-                  from airflow.jobs.triggerer_job import TriggererJob
-                  from airflow.utils.db import create_session
-                  from airflow.utils.net import get_hostname
-                  import sys
-
-                  with create_session() as session:
-                      job = session.query(TriggererJob).filter_by(hostname=get_hostname()).order_by(
-                          TriggererJob.latest_heartbeat.desc()).limit(1).first()
-
-                  sys.exit(0 if job.is_alive() else 1)
-                  "
+{{- with .Values.triggerer.livenessProbe.command }}
+{{ toYaml . | indent 16 }}
+{{- end }}
         {{- if and (.Values.dags.gitSync.enabled) (not .Values.dags.persistence.enabled) }}
         {{- include "git_sync_container" . | indent 8 }}
         {{- end }}

diff --git a/chart/tests/test_scheduler.py b/chart/tests/test_scheduler.py
@@ -219,6 +219,7 @@ def test_livenessprobe_values_are_configurable(self):
                         "timeoutSeconds": 222,
                         "failureThreshold": 333,
                         "periodSeconds": 444,
+                        "command": ["sh", "-c", "echo", "wow such test"],
                     }
                 },
             },
@@ -235,6 +236,9 @@ def test_livenessprobe_values_are_configurable(self):
             "spec.template.spec.containers[0].livenessProbe.failureThreshold", docs[0]
         )
         assert 444 == jmespath.search("spec.template.spec.containers[0].livenessProbe.periodSeconds", docs[0])
+        assert ["sh", "-c", "echo", "wow such test"] == jmespath.search(
+            "spec.template.spec.containers[0].livenessProbe.exec.command", docs[0]
+        )
 
     @parameterized.expand(
         [

diff --git a/chart/tests/test_triggerer.py b/chart/tests/test_triggerer.py
@@ -219,6 +219,7 @@ def test_livenessprobe_values_are_configurable(self):
                         "timeoutSeconds": 222,
                         "failureThreshold": 333,
                         "periodSeconds": 444,
+                        "command": ["sh", "-c", "echo", "wow such test"],
                     }
                 },
             },
@@ -236,6 +237,10 @@ def test_livenessprobe_values_are_configurable(self):
         )
         assert 444 == jmespath.search("spec.template.spec.containers[0].livenessProbe.periodSeconds", docs[0])
 
+        assert ["sh", "-c", "echo", "wow such test"] == jmespath.search(
+            "spec.template.spec.containers[0].livenessProbe.exec.command", docs[0]
+        )
+
     @parameterized.expand(
         [
             ({"enabled": False}, {"emptyDir": {}}),

diff --git a/chart/values.schema.json b/chart/values.schema.json
@@ -1504,6 +1504,13 @@
                             "description": "How often (in seconds) to perform the probe. Minimum value is 1.",
                             "type": "integer",
                             "default": 60
+                        },
+                        "command": {
+                            "description": "Command for livenessProbe",
+                            "type": "array",
+                            "items": {
+                                "type": "string"
+                            }
                         }
                     }
                 },
@@ -1808,6 +1815,13 @@
                             "description": "How often (in seconds) to perform the probe. Minimum value is 1.",
                             "type": "integer",
                             "default": 60
+                        },
+                        "command": {
+                            "description": "Command for livenessProbe",
+                            "type": "array",
+                            "items": {
+                                "type": "string"
+                            }
                         }
                     }
                 },

diff --git a/chart/values.yaml b/chart/values.yaml
@@ -563,6 +563,23 @@ scheduler:
     timeoutSeconds: 20
     failureThreshold: 5
     periodSeconds: 60
+    command:
+      - sh
+      - -c
+      - |
+        CONNECTION_CHECK_MAX_COUNT=0 exec /entrypoint python -Wignore -c "
+        import os
+        os.environ['AIRFLOW__CORE__LOGGING_LEVEL'] = 'ERROR'
+        os.environ['AIRFLOW__LOGGING__LOGGING_LEVEL'] = 'ERROR'
+        from airflow.jobs.scheduler_job import SchedulerJob
+        from airflow.utils.db import create_session
+        from airflow.utils.net import get_hostname
+        import sys
+        with create_session() as session:
+            job = session.query(SchedulerJob).filter_by(hostname=get_hostname()).order_by(
+                SchedulerJob.latest_heartbeat.desc()).limit(1).first()
+        sys.exit(0 if job.is_alive() else 1)"
+
   # Airflow 2.0 allows users to run multiple schedulers,
   # However this feature is only recommended for MySQL 8+ and Postgres
   replicas: 1
@@ -943,6 +960,26 @@ triggerer:
     timeoutSeconds: 20
     failureThreshold: 5
     periodSeconds: 60
+    command:
+      - sh
+      - -c
+      - |
+        CONNECTION_CHECK_MAX_COUNT=0 exec /entrypoint python -Wignore -c "
+        import os
+        os.environ['AIRFLOW__CORE__LOGGING_LEVEL'] = 'ERROR'
+        os.environ['AIRFLOW__LOGGING__LOGGING_LEVEL'] = 'ERROR'
+
+        from airflow.jobs.triggerer_job import TriggererJob
+        from airflow.utils.db import create_session
+        from airflow.utils.net import get_hostname
+        import sys
+
+        with create_session() as session:
+            job = session.query(TriggererJob).filter_by(hostname=get_hostname()).order_by(
+                TriggererJob.latest_heartbeat.desc()).limit(1).first()
+
+        sys.exit(0 if job.is_alive() else 1)
+        "
 
   # Create ServiceAccount
   serviceAccount: