Support strategy/updateStrategy on scheduler (#16069)

Allow the schedulers strategy (for Deployment) or
updateStrategy (for StatefulSet) to be set via helm parameters.

chart/templates/scheduler/scheduler-deployment.yaml

@@ -50,6 +50,14 @@ spec:
   serviceName: {{ .Release.Name }}-scheduler
 {{- end }}
   replicas: {{ .Values.scheduler.replicas }}
+  {{- if and $stateful .Values.scheduler.updateStrategy }}
+  updateStrategy:
+    {{- toYaml .Values.scheduler.updateStrategy | nindent 4 }}
+  {{- end }}
+  {{- if and (not $stateful) .Values.scheduler.strategy }}
+  strategy:
+    {{- toYaml .Values.scheduler.strategy | nindent 4 }}
+  {{- end }}
   selector:
     matchLabels:
       tier: airflow

chart/tests/test_scheduler.py

@@ -215,3 +215,56 @@ def test_airflow_local_settings(self):
             "subPath": "airflow_local_settings.py",
             "readOnly": True,
         } in jmespath.search("spec.template.spec.containers[0].volumeMounts", docs[0])
+
+    @parameterized.expand(
+        [
+            ("CeleryExecutor", False, {"rollingUpdate": {"partition": 0}}, None),
+            ("CeleryExecutor", True, {"rollingUpdate": {"partition": 0}}, None),
+            ("LocalExecutor", False, {"rollingUpdate": {"partition": 0}}, None),
+            ("LocalExecutor", True, {"rollingUpdate": {"partition": 0}}, {"rollingUpdate": {"partition": 0}}),
+            ("LocalExecutor", True, None, None),
+        ]
+    )
+    def test_scheduler_update_strategy(
+        self, executor, persistence, update_strategy, expected_update_strategy
+    ):
+        """updateStrategy should only be used when we have LocalExecutor and workers.persistence"""
+        docs = render_chart(
+            values={
+                "executor": executor,
+                "workers": {"persistence": {"enabled": persistence}},
+                "scheduler": {"updateStrategy": update_strategy},
+            },
+            show_only=["templates/scheduler/scheduler-deployment.yaml"],
+        )
+
+        assert expected_update_strategy == jmespath.search("spec.updateStrategy", docs[0])
+
+    @parameterized.expand(
+        [
+            ("LocalExecutor", False, None, None),
+            ("LocalExecutor", False, {"type": "Recreate"}, {"type": "Recreate"}),
+            ("LocalExecutor", True, {"type": "Recreate"}, None),
+            ("CeleryExecutor", True, None, None),
+            ("CeleryExecutor", False, None, None),
+            ("CeleryExecutor", True, {"type": "Recreate"}, {"type": "Recreate"}),
+            (
+                "CeleryExecutor",
+                False,
+                {"rollingUpdate": {"maxSurge": "100%", "maxUnavailable": "50%"}},
+                {"rollingUpdate": {"maxSurge": "100%", "maxUnavailable": "50%"}},
+            ),
+        ]
+    )
+    def test_scheduler_strategy(self, executor, persistence, strategy, expected_strategy):
+        """strategy should be used when we aren't using both LocalExecutor and workers.persistence"""
+        docs = render_chart(
+            values={
+                "executor": executor,
+                "workers": {"persistence": {"enabled": persistence}},
+                "scheduler": {"strategy": strategy},
+            },
+            show_only=["templates/scheduler/scheduler-deployment.yaml"],
+        )
+
+        assert expected_strategy == jmespath.search("spec.strategy", docs[0])

chart/values.schema.json

@@ -1088,6 +1088,22 @@
                     "type": "integer",
                     "default": 1
                 },
+                "updateStrategy": {
+                    "description": "Specifies the strategy used to replace old Pods by new ones when deployed as a StatefulSet (when using LocalExecutor and workers.persistence).",
+                    "type": [
+                        "null",
+                        "object"
+                    ],
+                    "default": null
+                },
+                "strategy": {
+                    "description": "Specifies the strategy used to replace old Pods by new ones when deployed as a Deployment (when not using LocalExecutor and workers.persistence).",
+                    "type": [
+                        "null",
+                        "object"
+                    ],
+                    "default": null
+                },
                 "serviceAccount": {
                     "description": "Create ServiceAccount.",
                     "type": "object",

chart/values.yaml

@@ -429,6 +429,12 @@ scheduler:
   # Airflow 2.0 allows users to run multiple schedulers,
   # However this feature is only recommended for MySQL 8+ and Postgres
   replicas: 1
+  # Update Strategy when scheduler is deployed as a StatefulSet
+  # (when using LocalExecutor and workers.persistence)
+  updateStrategy: ~
+  # Update Strategy when scheduler is deployed as a Deployment
+  # (when not using LocalExecutor and workers.persistence)
+  strategy: ~
 
   # Create ServiceAccount
   serviceAccount: