Enable native HPA for Airflow Workers

chart/templates/workers/worker-deployment.yaml

@@ -22,6 +22,7 @@
 #################################
 {{- $persistence := .Values.workers.persistence.enabled }}
 {{- $keda := .Values.workers.keda.enabled }}
+{{- $hpa := and .Values.workers.hpa.enabled (not .Values.workers.keda.enabled) }}
 {{- if or (eq .Values.executor "CeleryExecutor") (eq .Values.executor "CeleryKubernetesExecutor") }}
 {{- $nodeSelector := or .Values.workers.nodeSelector .Values.nodeSelector }}
 {{- $affinity := or .Values.workers.affinity .Values.affinity }}
@@ -59,7 +60,7 @@ spec:
   {{- if $persistence }}
   serviceName: {{ include "airflow.fullname" . }}-worker
   {{- end }}
-  {{- if not $keda }}
+  {{- if and (not $keda) (not $hpa) }}
   replicas: {{ .Values.workers.replicas }}
   {{- end }}
   {{- if $revisionHistoryLimit }}

chart/templates/workers/worker-hpa.yaml

@@ -0,0 +1,49 @@
+{{/*
+ Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements.  See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership.  The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License.  You may obtain a copy of the License at
+
+   http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied.  See the License for the
+ specific language governing permissions and limitations
+ under the License.
+*/}}
+
+################################
+## Airflow Worker HPA
+#################################
+{{- if and (and (not .Values.workers.keda.enabled) .Values.workers.hpa.enabled) (has .Values.executor (list "CeleryExecutor" "CeleryKubernetesExecutor")) }}
+apiVersion: autoscaling/v2
+kind: HorizontalPodAutoscaler
+metadata:
+  name: {{ include "airflow.fullname" . }}-worker
+  labels:
+    tier: airflow
+    component: worker-horizontalpodautoscaler
+    release: {{ .Release.Name }}
+    chart: "{{ .Chart.Name }}-{{ .Chart.Version }}"
+    heritage: {{ .Release.Service }}
+    deploymentName: {{ .Release.Name }}-worker
+    {{- if or (.Values.labels) (.Values.workers.labels) }}
+      {{- mustMerge .Values.workers.labels .Values.labels | toYaml | nindent 4 }}
+    {{- end }}
+spec:
+  scaleTargetRef:
+    apiVersion: apps/v1
+    kind: {{ ternary "StatefulSet" "Deployment" .Values.workers.persistence.enabled }}
+    name: {{ include "airflow.fullname" . }}-worker
+  minReplicas: {{ .Values.workers.hpa.minReplicaCount }}
+  maxReplicas: {{ .Values.workers.hpa.maxReplicaCount }}
+  metrics: {{- tpl (default "[]" .Values.workers.hpa.metrics) . | nindent 4 }}
+  {{- with .Values.workers.hpa.behavior }}
+  behavior: {{- toYaml . | nindent 4 }}
+  {{- end }}
+{{- end }}

chart/values.schema.json

@@ -1531,6 +1531,39 @@
                         }
                     }
                 },
+                "hpa": {
+                    "description": "HPA configuration.",
+                    "type": "object",
+                    "additionalProperties": false,
+                    "properties": {
+                        "enabled": {
+                            "description": "Allow HPA autoscaling (KEDA must be disabled).",
+                            "type": "boolean",
+                            "default": false
+                        },
+                        "minReplicaCount": {
+                            "description": "Minimum number of workers created by KEDA.",
+                            "type": "integer",
+                            "default": 0
+                        },
+                        "maxReplicaCount": {
+                            "description": "Maximum number of workers created by KEDA.",
+                            "type": "integer",
+                            "default": 10
+                        },
+                        "metrics": {
+                            "description": "Specifications for which to use to calculate the desired replica count.",
+                            "type": "string",
+                            "default": "See values.yaml"
+                        },
+                        "behavior": {
+                            "description": "HorizontalPodAutoscalerBehavior configures the scaling behavior of the target.",
+                            "type": "object",
+                            "default": {},
+                            "$ref": "#/definitions/io.k8s.api.autoscaling.v2beta2.HorizontalPodAutoscalerBehavior"
+                        }
+                    }
+                },
                 "persistence": {
                     "description": "Persistence configuration.",
                     "type": "object",

chart/values.yaml

@@ -592,6 +592,28 @@ workers:
     # This configuration will be ignored if PGBouncer is not enabled
     usePgbouncer: true
 
+  # Allow HPA (KEDA must be disabled).
+  hpa:
+    enabled: false
+
+    # Minimum number of workers created by HPA
+    minReplicaCount: 0
+
+    # Maximum number of workers created by HPA
+    maxReplicaCount: 10
+
+    # Specifications for which to use to calculate the desired replica count
+    metrics: |-
+      - type: Resource
+        resource:
+          name: cpu
+          target:
+            type: Utilization
+            averageUtilization: 80
+
+    # Scaling behavior of the target in both Up and Down directions
+    behavior: {}
+
   persistence:
     # Enable persistent volumes
     enabled: true

helm_tests/airflow_core/test_worker.py

@@ -977,6 +977,94 @@ def test_should_use_keda_query(self, query, executor, expected_query):
         assert expected_query == jmespath.search("spec.triggers[0].metadata.query", docs[0])
 
 
+class TestWorkerHPAAutoScaler:
+    """Tests worker HPA auto scaler."""
+
+    def test_should_be_disabled_on_keda_enabled(self):
+        docs = render_chart(
+            values={
+                "executor": "CeleryExecutor",
+                "workers": {
+                    "keda": {"enabled": True},
+                    "hpa": {"enabled": True},
+                    "labels": {"test_label": "test_label_value"},
+                },
+            },
+            show_only=[
+                "templates/workers/worker-kedaautoscaler.yaml",
+                "templates/workers/worker-hpa.yaml",
+            ],
+        )
+        assert "test_label" in jmespath.search("metadata.labels", docs[0])
+        assert jmespath.search("metadata.labels", docs[0])["test_label"] == "test_label_value"
+        assert len(docs) == 1
+
+    def test_should_add_component_specific_labels(self):
+        docs = render_chart(
+            values={
+                "executor": "CeleryExecutor",
+                "workers": {
+                    "hpa": {"enabled": True},
+                    "labels": {"test_label": "test_label_value"},
+                },
+            },
+            show_only=["templates/workers/worker-hpa.yaml"],
+        )
+
+        assert "test_label" in jmespath.search("metadata.labels", docs[0])
+        assert jmespath.search("metadata.labels", docs[0])["test_label"] == "test_label_value"
+
+    def test_should_remove_replicas_field(self):
+        docs = render_chart(
+            values={
+                "executor": "CeleryExecutor",
+                "workers": {
+                    "hpa": {"enabled": True},
+                },
+            },
+            show_only=["templates/workers/worker-deployment.yaml"],
+        )
+        assert "replicas" not in jmespath.search("spec", docs[0])
+
+    @pytest.mark.parametrize(
+        "metrics, executor, expected_metrics",
+        [
+            # default metrics
+            (
+                None,
+                "CeleryExecutor",
+                {
+                    "type": "Resource",
+                    "resource": {"name": "cpu", "target": {"type": "Utilization", "averageUtilization": 80}},
+                },
+            ),
+            # custom metric
+            (
+                '[{"type":"Pods","pods":{"metric":{"name":"custom_prometheus"},"target":{"type":"AverageValue","averageValue":"20"}}}]',
+                "CeleryKubernetesExecutor",
+                {
+                    "type": "Pods",
+                    "pods": {
+                        "metric": {"name": "custom_prometheus"},
+                        "target": {"type": "AverageValue", "averageValue": "20"},
+                    },
+                },
+            ),
+        ],
+    )
+    def test_should_use_hpa_metrics(self, metrics, executor, expected_metrics):
+        docs = render_chart(
+            values={
+                "executor": executor,
+                "workers": {
+                    "hpa": {"enabled": True, **({"metrics": metrics} if metrics else {})},
+                },
+            },
+            show_only=["templates/workers/worker-hpa.yaml"],
+        )
+        assert expected_metrics == jmespath.search("spec.metrics[0]", docs[0])
+
+
 class TestWorkerNetworkPolicy:
     """Tests worker network policy."""
 

helm_tests/other/test_hpa.py

@@ -0,0 +1,97 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import jmespath
+import pytest
+
+from tests.charts.helm_template_generator import render_chart
+
+
+class TestHPA:
+    """Tests HPA."""
+
+    def test_hpa_disabled_by_default(self):
+        """Disabled by default."""
+        docs = render_chart(
+            values={},
+            show_only=["templates/workers/worker-hpa.yaml"],
+        )
+        assert docs == []
+
+    @pytest.mark.parametrize(
+        "executor, is_created",
+        [
+            ("CeleryExecutor", True),
+            ("CeleryKubernetesExecutor", True),
+        ],
+    )
+    def test_hpa_enabled(self, executor, is_created):
+        """HPA should only be created when enabled and executor is Celery or CeleryKubernetes."""
+        docs = render_chart(
+            values={
+                "workers": {"hpa": {"enabled": True}, "persistence": {"enabled": False}},
+                "executor": executor,
+            },
+            show_only=["templates/workers/worker-hpa.yaml"],
+        )
+        if is_created:
+            assert jmespath.search("metadata.name", docs[0]) == "release-name-worker"
+        else:
+            assert docs == []
+
+    @pytest.mark.parametrize("executor", ["CeleryExecutor", "CeleryKubernetesExecutor"])
+    def test_hpa_behavior(self, executor):
+        """Verify HPA behavior."""
+        expected_behavior = {
+            "scaleDown": {
+                "stabilizationWindowSeconds": 300,
+                "policies": [{"type": "Percent", "value": 100, "periodSeconds": 15}],
+            }
+        }
+        docs = render_chart(
+            values={
+                "workers": {
+                    "hpa": {
+                        "enabled": True,
+                        "behavior": expected_behavior,
+                    },
+                },
+                "executor": executor,
+            },
+            show_only=["templates/workers/worker-hpa.yaml"],
+        )
+        assert jmespath.search("spec.behavior", docs[0]) == expected_behavior
+
+    @pytest.mark.parametrize(
+        "enabled, kind",
+        [
+            ("enabled", "StatefulSet"),
+            ("not_enabled", "Deployment"),
+        ],
+    )
+    def test_persistence(self, enabled, kind):
+        """If worker persistence is enabled, scaleTargetRef should be StatefulSet else Deployment."""
+        is_enabled = enabled == "enabled"
+        docs = render_chart(
+            values={
+                "workers": {"hpa": {"enabled": True}, "persistence": {"enabled": is_enabled}},
+                "executor": "CeleryExecutor",
+            },
+            show_only=["templates/workers/worker-hpa.yaml"],
+        )
+        assert jmespath.search("spec.scaleTargetRef.kind", docs[0]) == kind