Helm Chart: Add startupProbe to flower deployment (#45012)

* tests * schema * chart * template * pre-commit ruff --------- Co-authored-by: Jed Cunningham <[email protected]>
apache · Jan 27, 2025 · 0c4bbab · 0c4bbab
1 parent 9285cc7
commit 0c4bbab
Show file tree

Hide file tree

Showing 4 changed files with 72 additions and 0 deletions.
diff --git a/chart/templates/flower/flower-deployment.yaml b/chart/templates/flower/flower-deployment.yaml
@@ -141,6 +141,18 @@ spec:
             initialDelaySeconds: {{ .Values.flower.readinessProbe.initialDelaySeconds }}
             periodSeconds: {{ .Values.flower.readinessProbe.periodSeconds }}
             timeoutSeconds: {{ .Values.flower.readinessProbe.timeoutSeconds }}
+          startupProbe:
+            failureThreshold: {{ .Values.flower.startupProbe.failureThreshold }}
+            exec:
+              command:
+                - curl
+                {{- if (or .Values.flower.secretName (and .Values.flower.username .Values.flower.password))}}
+                - "--user"
+                - $AIRFLOW__CELERY__FLOWER_BASIC_AUTH
+                {{- end }}
+                - {{ printf "localhost:%s" (.Values.ports.flowerUI | toString) }}
+            periodSeconds: {{ .Values.flower.startupProbe.periodSeconds }}
+            timeoutSeconds: {{ .Values.flower.startupProbe.timeoutSeconds }}
           envFrom:
           {{- include "custom_airflow_environment_from" . | default "\n  []" | indent 10 }}
           env:

diff --git a/chart/values.schema.json b/chart/values.schema.json
@@ -6244,6 +6244,28 @@
                         }
                     }
                 },
+                "startupProbe": {
+                    "description": "Startup probe configuration.",
+                    "type": "object",
+                    "additionalProperties": false,
+                    "properties": {
+                        "timeoutSeconds": {
+                            "description": "Flower Startup probe timeout seconds.",
+                            "type": "integer",
+                            "default": 20
+                        },
+                        "failureThreshold": {
+                            "description": "Flower Startup probe failure threshold.",
+                            "type": "integer",
+                            "default": 6
+                        },
+                        "periodSeconds": {
+                            "description": "Flower Startup probe period seconds.",
+                            "type": "integer",
+                            "default": 10
+                        }
+                    }
+                },
                 "revisionHistoryLimit": {
                     "description": "Number of old replicasets to retain.",
                     "type": [

diff --git a/chart/values.yaml b/chart/values.yaml
@@ -1956,6 +1956,13 @@ flower:
     failureThreshold: 10
     periodSeconds: 5
 
+  # Wait for at most 1 minute (6*10s) for the flower container to startup.
+  # livenessProbe kicks in after the first successful startupProbe
+  startupProbe:
+    timeoutSeconds: 20
+    failureThreshold: 6
+    periodSeconds: 10
+
   # Max number of old replicasets to retain
   revisionHistoryLimit: ~
 

diff --git a/helm_tests/other/test_flower.py b/helm_tests/other/test_flower.py
@@ -153,6 +153,12 @@ def test_should_create_flower_deployment_with_authorization(self):
             "$AIRFLOW__CELERY__FLOWER_BASIC_AUTH",
             "localhost:7777",
         ]
+        assert jmespath.search("spec.template.spec.containers[0].startupProbe.exec.command", docs[0]) == [
+            "curl",
+            "--user",
+            "$AIRFLOW__CELERY__FLOWER_BASIC_AUTH",
+            "localhost:7777",
+        ]
 
     def test_should_create_flower_deployment_without_authorization(self):
         docs = render_chart(
@@ -175,6 +181,10 @@ def test_should_create_flower_deployment_without_authorization(self):
             "curl",
             "localhost:7777",
         ]
+        assert jmespath.search("spec.template.spec.containers[0].startupProbe.exec.command", docs[0]) == [
+            "curl",
+            "localhost:7777",
+        ]
 
     def test_scheduler_name(self):
         docs = render_chart(
@@ -429,6 +439,27 @@ def test_probe_values_are_configurable(self, probe):
         assert jmespath.search(f"spec.template.spec.containers[0].{probe}.failureThreshold", docs[0]) == 333
         assert jmespath.search(f"spec.template.spec.containers[0].{probe}.periodSeconds", docs[0]) == 444
 
+    def test_startup_probe_values_are_configurable(self):
+        docs = render_chart(
+            values={
+                "flower": {
+                    "enabled": True,
+                    "startupProbe": {
+                        "timeoutSeconds": 222,
+                        "failureThreshold": 333,
+                        "periodSeconds": 444,
+                    },
+                },
+            },
+            show_only=["templates/flower/flower-deployment.yaml"],
+        )
+
+        assert jmespath.search("spec.template.spec.containers[0].startupProbe.timeoutSeconds", docs[0]) == 222
+        assert (
+            jmespath.search("spec.template.spec.containers[0].startupProbe.failureThreshold", docs[0]) == 333
+        )
+        assert jmespath.search("spec.template.spec.containers[0].startupProbe.periodSeconds", docs[0]) == 444
+
 
 class TestFlowerService:
     """Tests flower service."""