Added autoscaling configuration for airflow workers

charts/airflow/helm-config-dp.yaml

@@ -67,6 +67,11 @@ pgbouncer:
     passwordSecret: "postgres"
     passwordSecretKey: "postgres-password"
 
+################################################################################
+# Specify KEDA configuration
+#
+keda:
+  version_2_enabled: false
 
 ################################################################################
 # Specify the airflow configuration
@@ -179,17 +184,12 @@ airflow:
       aadpodidbinding: "osdu-identity"
     podAnnotations:
       sidecar.istio.io/inject: "false"
-    autoscaling:
+    autoscale:
       enabled: false
-      ## minReplicas is picked from Values.workers.replicas and default value is 1
-      maxReplicas: 3
-      metrics:
-      - type: Resource
-        resource:
-          name: memory
-          target:
-            type: Utilization
-            averageUtilization: 60
+      minReplicas: 2
+      maxReplicas: 20
+      scaleDown:
+        coolDownPeriod: 300
     labels:
       # DO NOT DELETE THIS LABEL. SET IT TO "false" WHEN AUTOSCALING IS DISABLED, SET IT TO "true" WHEN AUTOSCALING IS ENABLED
       autoscalingEnabled: "false"
@@ -261,6 +261,7 @@ airflow:
       AIRFLOW__CORE__STORE_SERIALIZED_DAGS: True #This flag decides whether to serialise DAGs and persist them in DB
       AIRFLOW__CORE__STORE_DAG_CODE: True #This flag decides whether to persist DAG files code in DB
       AIRFLOW__WEBSERVER__WORKER_CLASS: gevent
+      AIRFLOW__CELERY__WORKER_CONCURRENCY: 16
     extraEnv:
     - name: CLOUD_PROVIDER
       value: "azure"

charts/airflow/helm-config.yaml

@@ -55,6 +55,12 @@ pgbouncer:
     passwordSecret: "postgres"
     passwordSecretKey: "postgres-password"
 
+################################################################################
+# Specify KEDA configuration
+#
+keda:
+  version_2_enabled: false
+
 
 ################################################################################
 # Specify the airflow configuration
@@ -172,17 +178,12 @@ airflow:
         memory: "5Gi"
     podLabels:
       aadpodidbinding: "osdu-identity"
-    autoscaling:
+    autoscale:
       enabled: false
-      ## minReplicas is picked from Values.workers.replicas and default value is 1
-      maxReplicas: 3
-      metrics:
-      - type: Resource
-        resource:
-          name: memory
-          target:
-            type: Utilization
-            averageUtilization: 60
+      minReplicas: 2
+      maxReplicas: 20
+      scaleDown:
+        coolDownPeriod: 300
     labels:
       # DO NOT DELETE THIS LABEL. SET IT TO "false" WHEN AUTOSCALING IS DISABLED, SET IT TO "true" WHEN AUTOSCALING IS ENABLED
       autoscalingEnabled: "false"
@@ -252,6 +253,7 @@ airflow:
       AIRFLOW__CORE__STORE_SERIALIZED_DAGS: True #This flag decides whether to serialise DAGs and persist them in DB
       AIRFLOW__CORE__STORE_DAG_CODE: True #This flag decides whether to persist DAG files code in DB
       AIRFLOW__WEBSERVER__WORKER_CLASS: gevent
+      AIRFLOW__CELERY__WORKER_CONCURRENCY: 16 # Do not remove this config as it is used for autoscaling as well
     extraEnv:
     - name: CLOUD_PROVIDER
       value: "azure"

charts/airflow/templates/autoscale/worker-config.yaml

+{{- if and .Values.airflow.web.autoscale.enabled .Values.keda.version_2_enabled -}}
+apiVersion: keda.sh/v1alpha1
+kind: ScaledObject
+metadata:
+  name: {{ printf "%s" .Release.Name | trunc 63 | trimSuffix "-" }}-worker
+spec:
+  kind: StatefulSet
+  name: {{ printf "%s" .Release.Name | trunc 63 | trimSuffix "-" }}-worker
+  pollingInterval: 30
+  minReplicaCount: {{ .Values.airflow.workers.autoscale.minReplicas }}
+  maxReplicaCount: {{ .Values.airflow.workers.autoscale.maxReplicas }}
+  advanced:
+    horizontalPodAutoscalerConfig:
+      behavior:
+        scaleUp:
+          stabilizationWindowSeconds: 180
+          policies:
+          - type: Percent
+            value: 900
+            periodSeconds: 30
+        scaleDown:
+          stabilizationWindowSeconds: 300
+          policies:
+          - type: Percent
+            value: 30
+            periodSeconds: {{ .Values.airflow.workers.autoscale.scaleDown.coolDownPeriod }}
+  triggers:
+    - type: postgresql
+      metadata:
+        host: {{ .Values.airflow.externalDatabase.host | quote }}
+        userName: {{ .Values.airflow.externalDatabase.user | quote }}
+        passwordFromEnv: DATABASE_PASSWORD
+        port: {{ .Values.airflow.externalDatabase.port | quote }}
+        dbName: {{ .Values.airflow.externalDatabase.database | quote }}
+        sslmode: "prefer"
+        targetQueryValue: 1
+        query: "SELECT ceil(COUNT(*)::decimal / {{ .Values.airflow.airflow.config.AIRFLOW__CELERY__WORKER_CONCURRENCY }}) FROM task_instance WHERE state='running' OR state='queued'"
+        metricName: "active_tasks_count"
+{{- end }}
\ No newline at end of file