Chart: support for priorityClassName (#20794)

apache · Feb 4, 2022 · ec41fd5 · ec41fd5
1 parent ab762a5
commit ec41fd5
Show file tree

Hide file tree

Showing 10 changed files with 117 additions and 0 deletions.
diff --git a/chart/templates/flower/flower-deployment.yaml b/chart/templates/flower/flower-deployment.yaml
@@ -67,6 +67,9 @@ spec:
       tolerations:
 {{ toYaml $tolerations | indent 8 }}
       serviceAccountName: {{ include "flower.serviceAccountName" . }}
+      {{- if .Values.flower.priorityClassName }}
+      priorityClassName: {{ .Values.flower.priorityClassName }}
+      {{- end }}
       restartPolicy: Always
       securityContext: {{ $securityContext | nindent 8 }}
       {{- if or .Values.registry.secretName .Values.registry.connection }}

diff --git a/chart/templates/pgbouncer/pgbouncer-deployment.yaml b/chart/templates/pgbouncer/pgbouncer-deployment.yaml
@@ -59,6 +59,9 @@ spec:
         checksum/pgbouncer-config-secret: {{ include (print $.Template.BasePath "/secrets/pgbouncer-config-secret.yaml") . | sha256sum }}
         checksum/pgbouncer-certificates-secret: {{ include (print $.Template.BasePath "/secrets/pgbouncer-certificates-secret.yaml") . | sha256sum }}
     spec:
+      {{- if .Values.pgbouncer.priorityClassName }}
+      priorityClassName: {{ .Values.pgbouncer.priorityClassName }}
+      {{- end }}
       nodeSelector:
 {{ toYaml $nodeSelector | indent 8 }}
       affinity:

diff --git a/chart/templates/scheduler/scheduler-deployment.yaml b/chart/templates/scheduler/scheduler-deployment.yaml
@@ -90,6 +90,9 @@ spec:
         {{- toYaml .Values.scheduler.podAnnotations | nindent 8 }}
         {{- end }}
     spec:
+      {{- if .Values.scheduler.priorityClassName }}
+      priorityClassName: {{ .Values.scheduler.priorityClassName }}
+      {{- end }}
       nodeSelector:
 {{ toYaml $nodeSelector | indent 8 }}
       affinity:

diff --git a/chart/templates/statsd/statsd-deployment.yaml b/chart/templates/statsd/statsd-deployment.yaml
@@ -57,6 +57,9 @@ spec:
         checksum/statsd-config: {{ include (print $.Template.BasePath "/configmaps/statsd-configmap.yaml") . | sha256sum }}
 {{- end }}
     spec:
+      {{- if .Values.statsd.priorityClassName }}
+      priorityClassName: {{ .Values.statsd.priorityClassName }}
+      {{- end }}
       nodeSelector:
 {{ toYaml $nodeSelector | indent 8 }}
       affinity:

diff --git a/chart/templates/triggerer/triggerer-deployment.yaml b/chart/templates/triggerer/triggerer-deployment.yaml
@@ -73,6 +73,9 @@ spec:
         {{- toYaml .Values.triggerer.podAnnotations | nindent 8 }}
         {{- end }}
     spec:
+      {{- if .Values.triggerer.priorityClassName }}
+      priorityClassName: {{ .Values.triggerer.priorityClassName }}
+      {{- end }}
       nodeSelector:
         {{- toYaml $nodeSelector | nindent 8 }}
       affinity:

diff --git a/chart/templates/webserver/webserver-deployment.yaml b/chart/templates/webserver/webserver-deployment.yaml
@@ -87,6 +87,9 @@ spec:
         {{- end }}
     spec:
       serviceAccountName: {{ include "webserver.serviceAccountName" . }}
+      {{- if .Values.webserver.priorityClassName }}
+      priorityClassName: {{ .Values.webserver.priorityClassName }}
+      {{- end }}
       nodeSelector:
 {{ toYaml $nodeSelector | indent 8 }}
       affinity:

diff --git a/chart/templates/workers/worker-deployment.yaml b/chart/templates/workers/worker-deployment.yaml
@@ -83,6 +83,9 @@ spec:
         {{- toYaml .Values.workers.podAnnotations | nindent 8 }}
         {{- end }}
     spec:
+      {{- if .Values.workers.priorityClassName }}
+      priorityClassName: {{ .Values.workers.priorityClassName }}
+      {{- end }}
       nodeSelector:
 {{ toYaml $nodeSelector | indent 8 }}
       affinity:

diff --git a/chart/tests/test_airflow_common.py b/chart/tests/test_airflow_common.py
@@ -268,3 +268,31 @@ def test_have_all_config_mounts_on_init_containers(self):
         }
         for doc in docs:
             assert expected_mount in jmespath.search("spec.template.spec.initContainers[0].volumeMounts", doc)
+
+    def test_priority_class_name(self):
+        docs = render_chart(
+            values={
+                "flower": {"priorityClassName": "low-priority-flower"},
+                "pgbouncer": {"priorityClassName": "low-priority-pgbouncer"},
+                "scheduler": {"priorityClassName": "low-priority-scheduler"},
+                "statsd": {"priorityClassName": "low-priority-statsd"},
+                "triggerer": {"priorityClassName": "low-priority-triggerer"},
+                "webserver": {"priorityClassName": "low-priority-webserver"},
+                "workers": {"priorityClassName": "low-priority-worker"},
+            },
+            show_only=[
+                "templates/flower/flower-deployment.yaml",
+                "templates/pgbouncer/pgbouncer-deployment.yaml",
+                "templates/scheduler/scheduler-deployment.yaml",
+                "templates/statsd/statsd-deployment.yaml",
+                "templates/triggerer/triggerer-deployment.yaml",
+                "templates/webserver/webserver-deployment.yaml",
+                "templates/workers/worker-deployment.yaml",
+            ],
+        )
+
+        for doc in docs:
+            component = doc['metadata']['labels']['component']
+            priority = doc['spec']['template']['spec']['priorityClassName']
+
+            assert priority == f"low-priority-{component}"
diff --git a/chart/values.schema.json b/chart/values.schema.json
@@ -1351,6 +1351,14 @@
                         "type": "string"
                     }
                 },
+                "priorityClassName": {
+                    "description": "Specify priority for worker pods.",
+                    "type": [
+                        "string",
+                        "null"
+                    ],
+                    "default": null
+                },
                 "affinity": {
                     "description": "Specify scheduling constraints for worker pods.",
                     "type": "object",
@@ -1660,6 +1668,14 @@
                         "type": "string"
                     }
                 },
+                "priorityClassName": {
+                    "description": "Specify priority for scheduler pods.",
+                    "type": [
+                        "string",
+                        "null"
+                    ],
+                    "default": null
+                },
                 "affinity": {
                     "description": "Specify scheduling constraints for scheduler pods.",
                     "type": "object",
@@ -1930,6 +1946,14 @@
                         "type": "string"
                     }
                 },
+                "priorityClassName": {
+                    "description": "Specify priority for triggerer pods.",
+                    "type": [
+                        "string",
+                        "null"
+                    ],
+                    "default": null
+                },
                 "affinity": {
                     "description": "Specify scheduling constraints for triggerer pods.",
                     "type": "object",
@@ -2632,6 +2656,14 @@
                         "type": "string"
                     }
                 },
+                "priorityClassName": {
+                    "description": "Specify priority for webserver pods.",
+                    "type": [
+                        "string",
+                        "null"
+                    ],
+                    "default": null
+                },
                 "affinity": {
                     "description": "Specify scheduling constraints for webserver pods.",
                     "type": "object",
@@ -2915,6 +2947,14 @@
                         "type": "string"
                     }
                 },
+                "priorityClassName": {
+                    "description": "Specify priority for Flower pods.",
+                    "type": [
+                        "string",
+                        "null"
+                    ],
+                    "default": null
+                },
                 "affinity": {
                     "description": "Specify scheduling constraints for Flower pods.",
                     "type": "object",
@@ -3046,6 +3086,14 @@
                         "type": "string"
                     }
                 },
+                "priorityClassName": {
+                    "description": "Specify priority for StatsD pods.",
+                    "type": [
+                        "string",
+                        "null"
+                    ],
+                    "default": null
+                },
                 "affinity": {
                     "description": "Specify scheduling constraints for StatsD pods.",
                     "type": "object",
@@ -3347,6 +3395,14 @@
                         "type": "string"
                     }
                 },
+                "priorityClassName": {
+                    "description": "Specify priority for PgBouncer pods.",
+                    "type": [
+                        "string",
+                        "null"
+                    ],
+                    "default": null
+                },
                 "affinity": {
                     "description": "Specify scheduling constraints for PgBouncer pods.",
                     "type": "object",

diff --git a/chart/values.yaml b/chart/values.yaml
@@ -514,6 +514,7 @@ workers:
 
   # Select certain nodes for airflow worker pods.
   nodeSelector: {}
+  priorityClassName: ~
   affinity: {}
   # default worker affinity is:
   #  podAntiAffinity:
@@ -638,6 +639,8 @@ scheduler:
   #      weight: 100
   tolerations: []
 
+  priorityClassName: ~
+
   podAnnotations: {}
 
   logGroomerSidecar:
@@ -866,6 +869,7 @@ webserver:
 
   # Select certain nodes for airflow webserver pods.
   nodeSelector: {}
+  priorityClassName: ~
   affinity: {}
   # default webserver affinity is:
   #  podAntiAffinity:
@@ -960,6 +964,8 @@ triggerer:
   #      weight: 100
   tolerations: []
 
+  priorityClassName: ~
+
   podAnnotations: {}
 
 # Flower settings
@@ -1049,6 +1055,8 @@ flower:
   affinity: {}
   tolerations: []
 
+  priorityClassName: ~
+
   podAnnotations: {}
 
 # Statsd settings
@@ -1090,6 +1098,8 @@ statsd:
   affinity: {}
   tolerations: []
 
+  priorityClassName: ~
+
   # Additional mappings for statsd exporter.
   extraMappings: []
 
@@ -1196,6 +1206,8 @@ pgbouncer:
   affinity: {}
   tolerations: []
 
+  priorityClassName: ~
+
   uid: 65534
 
   metricsExporterSidecar: