feat: add topologySpreadConstraints values

Signed-off-by: ron-damon <[email protected]> Signed-off-by: ron-damon <[email protected]>
airflow-helm · May 1, 2024 · c39c1a5 · c39c1a5
1 parent 92e871f
commit c39c1a5
Show file tree

Hide file tree

Showing 21 changed files with 180 additions and 4 deletions.
diff --git a/charts/airflow/README.md b/charts/airflow/README.md
@@ -214,6 +214,7 @@ Parameter | Description | Default
 `airflow.pools` | a list airflow pools to create | `<see values.yaml>`
 `airflow.poolsUpdate` | if we create a Deployment to perpetually sync `airflow.pools` | `true`
 `airflow.defaultNodeSelector` | default nodeSelector for airflow Pods (is overridden by pod-specific values) | `{}`
+`airflow.defaultTopologySpreadConstraints` | default topologySpreadConstraints for airflow Pods (is overridden by pod-specific values) | `[]`
 `airflow.defaultAffinity` | default affinity configs for airflow Pods (is overridden by pod-specific values) | `{}`
 `airflow.defaultTolerations` | default toleration configs for airflow Pods (is overridden by pod-specific values) | `[]`
 `airflow.defaultSecurityContext` | default securityContext configs for Pods (is overridden by pod-specific values) | `{fsGroup: 0}`
@@ -241,6 +242,7 @@ Parameter | Description | Default
 `scheduler.replicas` | the number of scheduler Pods to run | `1`
 `scheduler.resources` | resource requests/limits for the scheduler Pods | `{}`
 `scheduler.nodeSelector` | the nodeSelector configs for the scheduler Pods | `{}`
+`scheduler.topologySpreadConstraints` | the topologySpreadConstraints configs for the scheduler Pods | `[]`
 `scheduler.affinity` | the affinity configs for the scheduler Pods | `{}`
 `scheduler.tolerations` | the toleration configs for the scheduler Pods | `[]`
 `scheduler.securityContext` | the security context for the scheduler Pods | `{}`
@@ -269,6 +271,7 @@ Parameter | Description | Default
 `web.replicas` | the number of web Pods to run | `1`
 `web.resources` | resource requests/limits for the airflow web pods | `{}`
 `web.nodeSelector` | the number of web Pods to run | `{}`
+`web.topologySpreadConstraints` | the topologySpreadConstraints configs for the web Pods | `[]`
 `web.affinity` | the affinity configs for the web Pods | `{}`
 `web.tolerations` | the toleration configs for the web Pods | `[]`
 `web.securityContext` | the security context for the web Pods | `{}`
@@ -296,6 +299,7 @@ Parameter | Description | Default
 `workers.replicas` | the number of workers Pods to run | `1`
 `workers.resources` | resource requests/limits for the airflow worker Pods | `{}`
 `workers.nodeSelector` | the nodeSelector configs for the worker Pods | `{}`
+`workers.topologySpreadConstraints` | the topologySpreadConstraints configs for the worker Pods | `[]`
 `workers.affinity` | the affinity configs for the worker Pods | `{}`
 `workers.tolerations` | the toleration configs for the worker Pods | `[]`
 `workers.securityContext` | the security context for the worker Pods | `{}`
@@ -325,6 +329,7 @@ Parameter | Description | Default
 `triggerer.replicas` | the number of triggerer Pods to run | `1`
 `triggerer.resources` | resource requests/limits for the airflow triggerer Pods | `{}`
 `triggerer.nodeSelector` | the nodeSelector configs for the triggerer Pods | `{}`
+`triggerer.topologySpreadConstraints` | the topologySpreadConstraints configs for the triggerer Pods | `[]`
 `triggerer.affinity` | the affinity configs for the triggerer Pods | `{}`
 `triggerer.tolerations` | the toleration configs for the triggerer Pods | `[]`
 `triggerer.securityContext` | the security context for the triggerer Pods | `{}`
@@ -350,6 +355,7 @@ Parameter | Description | Default
 `flower.enabled` | if the Flower UI should be deployed | `true`
 `flower.resources` | resource requests/limits for the flower Pods | `{}`
 `flower.nodeSelector` | the nodeSelector configs for the flower Pods | `{}`
+`flower.topologySpreadConstraints` | the topologySpreadConstraints configs for the flower Pods | `[]`
 `flower.affinity` | the affinity configs for the flower Pods | `{}`
 `flower.tolerations` | the toleration configs for the flower Pods | `[]`
 `flower.securityContext` | the security context for the flower Pods | `{}`
@@ -440,6 +446,7 @@ Parameter | Description | Default
 `pgbouncer.image.*` | configs for the pgbouncer container image | `<see values.yaml>`
 `pgbouncer.resources` | resource requests/limits for the pgbouncer Pods | `{}`
 `pgbouncer.nodeSelector` | the nodeSelector configs for the pgbouncer Pods | `{}`
+`pgbouncer.topologySpreadConstraints` | the topologySpreadConstraints configs for the pgbouncer Pods | `[]`
 `pgbouncer.affinity` | the affinity configs for the pgbouncer Pods | `{}`
 `pgbouncer.tolerations` | the toleration configs for the pgbouncer Pods | `[]`
 `pgbouncer.securityContext` | the security context for the pgbouncer Pods | `{}`

diff --git a/charts/airflow/docs/faq/kubernetes/affinity-node-selectors-tolerations.md b/charts/airflow/docs/faq/kubernetes/affinity-node-selectors-tolerations.md
@@ -5,13 +5,13 @@
 # Configure Pod Affinity/Selectors/Tolerations
 
 If your environment needs to use Pod [affinity](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#affinity-and-anti-affinity), 
-[nodeSelector](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#nodeselector), 
+[nodeSelector](https://kubernetes.io/docs/concepts/scheduling-eviction/assign-pod-node/#nodeselector), [topologySpreadConstraints](https://kubernetes.io/docs/concepts/scheduling-eviction/topology-spread-constraints/#topologyspreadconstraints-field),
 or [tolerations](https://kubernetes.io/docs/concepts/scheduling-eviction/taint-and-toleration/), 
 we provide many values that allow fine-grained control over the Pod definitions.
 
 ## Global Configs
 
-To set affinity, nodeSelector, and tolerations for all airflow Pods, you may use the `airflow.{defaultNodeSelector,defaultAffinity,defaultTolerations}` values:
+To set affinity, nodeSelector, topologySpreadConstraints, and tolerations for all airflow Pods, you may use the `airflow.{defaultNodeSelector,defaultTopologySpreadConstraints,defaultAffinity,defaultTolerations}` values:
 
 ```yaml
 airflow:
@@ -20,6 +20,16 @@ airflow:
     # my_node_label_1: value1
     # my_node_label_2: value2
 
+  ## https://kubernetes.io/docs/concepts/scheduling-eviction/topology-spread-constraints/#topologyspreadconstraints-field
+  defaultTopologySpreadConstraints: []
+    # - maxSkew: 1
+    #   topologyKey: topology.kubernetes.io/zone
+    #   whenUnsatisfiable: DoNotSchedule
+    #   labelSelector:
+    #     matchLabels:
+    #       my_label_1: value1
+    #       my_label_2: value2
+
   ## https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.20/#affinity-v1-core
   defaultAffinity: {}
     # podAffinity:
@@ -69,53 +79,60 @@ redis:
 
 ## Per-Resource Configs
 
-To set affinity, nodeSelector, and tolerations for specific pods, you may use the following values:
+To set affinity, nodeSelector, topologySpreadConstraints, and tolerations for specific pods, you may use the following values:
 
 ```yaml
 airflow:
   ## airflow KubernetesExecutor pod_template
   kubernetesPodTemplate:
     nodeSelector: {}
+    topologySpreadConstraints: []
     affinity: {}
     tolerations: []
 
   ## sync deployments
   sync:
     nodeSelector: {}
+    topologySpreadConstraints: []
     affinity: {}
     tolerations: []
 
 ## airflow schedulers
 scheduler:
   nodeSelector: {}
+  topologySpreadConstraints: []
   affinity: {}
   tolerations: []
 
 ## airflow webserver
 web:
   nodeSelector: {}
+  topologySpreadConstraints: []
   affinity: {}
   tolerations: []
 
 ## airflow workers
 workers:
   nodeSelector: {}
+  topologySpreadConstraints: []
   affinity: {}
   tolerations: []
 
 ## airflow triggerer
 triggerer:
   nodeSelector: {}
+  topologySpreadConstraints: []
   affinity: {}
   tolerations: []
 
 ## airflow workers
 flower:
   nodeSelector: {}
+  topologySpreadConstraints: []
   affinity: {}
   tolerations: []
 ```
 
 > 🟦 __Tip__ 🟦
 >
-> The `airflow.{defaultNodeSelector,defaultAffinity,defaultTolerations}` values are overridden by the per-resource values like `scheduler.{nodeSelector,affinity,tolerations}`.
+> The `airflow.{defaultNodeSelector,defaultTopologySpreadConstraints,defaultAffinity,defaultTolerations}` values are overridden by the per-resource values like `scheduler.{nodeSelector,topologySpreadConstraints,affinity,tolerations}`.
diff --git a/charts/airflow/files/pod_template.kubernetes-helm-yaml b/charts/airflow/files/pod_template.kubernetes-helm-yaml
@@ -1,4 +1,5 @@
 {{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.airflow.kubernetesPodTemplate.nodeSelector) }}
+{{- $podTopologySpreadConstraints := include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" .Values.airflow.kubernetesPodTemplate.topologySpreadConstraints) }}
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.airflow.kubernetesPodTemplate.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.airflow.kubernetesPodTemplate.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.airflow.kubernetesPodTemplate.securityContext) }}
@@ -31,6 +32,10 @@ spec:
   nodeSelector:
     {{- $podNodeSelector | nindent 4 }}
   {{- end }}
+  {{- if $podTopologySpreadConstraints }}
+  topologySpreadConstraints:
+    {{- $podTopologySpreadConstraints | nindent 4 }}
+  {{- end }}
   {{- if $podAffinity }}
   affinity:
     {{- $podAffinity | nindent 4 }}

diff --git a/charts/airflow/templates/_helpers/pods.tpl b/charts/airflow/templates/_helpers/pods.tpl
@@ -29,6 +29,14 @@ EXAMPLE USAGE: {{ include "airflow.nodeSelector" (dict "Release" .Release "Value
 {{- .nodeSelector | default .Values.airflow.defaultNodeSelector | toYaml }}
 {{- end }}
 
+{{/*
+Define the topologySpreadConstraints for airflow pods
+EXAMPLE USAGE: {{ include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" $topologySpreadConstraints) }}
+*/}}
+{{- define "airflow.podTopologySpreadConstraints" }}
+{{- .topologySpreadConstraints | default .Values.airflow.defaultTopologySpreadConstraints | toYaml }}
+{{- end }}
+
 {{/*
 Define the Affinity for airflow pods
 EXAMPLE USAGE: {{ include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" $affinity) }}

diff --git a/charts/airflow/templates/db-migrations/db-migrations-deployment.yaml b/charts/airflow/templates/db-migrations/db-migrations-deployment.yaml
@@ -1,5 +1,6 @@
 {{- if and (.Values.airflow.dbMigrations.enabled) (not .Values.airflow.dbMigrations.runAsJob) }}
 {{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.airflow.dbMigrations.nodeSelector) }}
+{{- $podTopologySpreadConstraints := include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" .Values.airflow.dbMigrations.topologySpreadConstraints) }}
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.airflow.dbMigrations.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.airflow.dbMigrations.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.airflow.dbMigrations.securityContext) }}
@@ -65,6 +66,10 @@ spec:
       nodeSelector:
         {{- $podNodeSelector | nindent 8 }}
       {{- end }}
+      {{- if $podTopologySpreadConstraints }}
+      topologySpreadConstraints:
+        {{- $podTopologySpreadConstraints | nindent 8 }}
+      {{- end }}
       {{- if $podAffinity }}
       affinity:
         {{- $podAffinity | nindent 8 }}

diff --git a/charts/airflow/templates/db-migrations/db-migrations-job.yaml b/charts/airflow/templates/db-migrations/db-migrations-job.yaml
@@ -1,5 +1,6 @@
 {{- if and (.Values.airflow.dbMigrations.enabled) (.Values.airflow.dbMigrations.runAsJob) }}
 {{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.airflow.dbMigrations.nodeSelector) }}
+{{- $podTopologySpreadConstraints := include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" .Values.airflow.dbMigrations.topologySpreadConstraints) }}
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.airflow.dbMigrations.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.airflow.dbMigrations.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.airflow.dbMigrations.securityContext) }}
@@ -59,6 +60,10 @@ spec:
       nodeSelector:
         {{- $podNodeSelector | nindent 8 }}
       {{- end }}
+      {{- if $podTopologySpreadConstraints }}
+      topologySpreadConstraints:
+        {{- $podTopologySpreadConstraints | nindent 8 }}
+      {{- end }}
       {{- if $podAffinity }}
       affinity:
         {{- $podAffinity | nindent 8 }}

diff --git a/charts/airflow/templates/flower/flower-deployment.yaml b/charts/airflow/templates/flower/flower-deployment.yaml
@@ -1,5 +1,6 @@
 {{- if .Values.flower.enabled }}
 {{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.flower.nodeSelector) }}
+{{- $podTopologySpreadConstraints := include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" .Values.flower.topologySpreadConstraints) }}
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.flower.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.flower.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.flower.securityContext) }}
@@ -69,6 +70,10 @@ spec:
       nodeSelector:
         {{- $podNodeSelector | nindent 8 }}
       {{- end }}
+      {{- if $podTopologySpreadConstraints }}
+      topologySpreadConstraints:
+        {{- $podTopologySpreadConstraints | nindent 8 }}
+      {{- end }}
       {{- if $podAffinity }}
       affinity:
         {{- $podAffinity | nindent 8 }}

diff --git a/charts/airflow/templates/pgbouncer/pgbouncer-deployment.yaml b/charts/airflow/templates/pgbouncer/pgbouncer-deployment.yaml
@@ -52,6 +52,7 @@
 
 {{- if include "airflow.pgbouncer.should_use" . }}
 {{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.pgbouncer.nodeSelector) }}
+{{- $podTopologySpreadConstraints := include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" .Values.pgbouncer.topologySpreadConstraints) }}
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.pgbouncer.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.pgbouncer.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.pgbouncer.securityContext) }}
@@ -115,6 +116,10 @@ spec:
       nodeSelector:
         {{- $podNodeSelector | nindent 8 }}
       {{- end }}
+      {{- if $podTopologySpreadConstraints }}
+      topologySpreadConstraints:
+        {{- $podTopologySpreadConstraints | nindent 8 }}
+      {{- end }}
       {{- if $podAffinity }}
       affinity:
         {{- $podAffinity | nindent 8 }}

diff --git a/charts/airflow/templates/scheduler/scheduler-deployment.yaml b/charts/airflow/templates/scheduler/scheduler-deployment.yaml
@@ -1,4 +1,5 @@
 {{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.scheduler.nodeSelector) }}
+{{- $podTopologySpreadConstraints := include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" .Values.scheduler.topologySpreadConstraints) }}
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.scheduler.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.scheduler.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.scheduler.securityContext) }}
@@ -77,6 +78,10 @@ spec:
       nodeSelector:
         {{- $podNodeSelector | nindent 8 }}
       {{- end }}
+      {{- if $podTopologySpreadConstraints }}
+      topologySpreadConstraints:
+        {{- $podTopologySpreadConstraints | nindent 8 }}
+      {{- end }}
       {{- if $podAffinity }}
       affinity:
         {{- $podAffinity | nindent 8 }}

diff --git a/charts/airflow/templates/sync/sync-connections-deployment.yaml b/charts/airflow/templates/sync/sync-connections-deployment.yaml
@@ -1,5 +1,6 @@
 {{- if and (.Values.airflow.connections) (.Values.airflow.connectionsUpdate) }}
 {{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.airflow.sync.nodeSelector) }}
+{{- $podTopologySpreadConstraints := include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" .Values.airflow.sync.topologySpreadConstraints) }}
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.airflow.sync.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.airflow.sync.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.airflow.sync.securityContext) }}
@@ -65,6 +66,10 @@ spec:
       nodeSelector:
         {{- $podNodeSelector | nindent 8 }}
       {{- end }}
+      {{- if $podTopologySpreadConstraints }}
+      topologySpreadConstraints:
+        {{- $podTopologySpreadConstraints | nindent 8 }}
+      {{- end }}
       {{- if $podAffinity }}
       affinity:
         {{- $podAffinity | nindent 8 }}

diff --git a/charts/airflow/templates/sync/sync-connections-job.yaml b/charts/airflow/templates/sync/sync-connections-job.yaml
@@ -1,5 +1,6 @@
 {{- if and (.Values.airflow.connections) (not .Values.airflow.connectionsUpdate) }}
 {{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.scheduler.nodeSelector) }}
+{{- $podTopologySpreadConstraints := include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" .Values.scheduler.topologySpreadConstraints) }}
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.scheduler.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.scheduler.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.airflow.sync.securityContext) }}
@@ -59,6 +60,10 @@ spec:
       nodeSelector:
         {{- $podNodeSelector | nindent 8 }}
       {{- end }}
+      {{- if $podTopologySpreadConstraints }}
+      topologySpreadConstraints:
+        {{- $podTopologySpreadConstraints | nindent 8 }}
+      {{- end }}
       {{- if $podAffinity }}
       affinity:
         {{- $podAffinity | nindent 8 }}

diff --git a/charts/airflow/templates/sync/sync-pools-deployment.yaml b/charts/airflow/templates/sync/sync-pools-deployment.yaml
@@ -1,5 +1,6 @@
 {{- if and (.Values.airflow.pools) (.Values.airflow.poolsUpdate) }}
 {{- $podNodeSelector := include "airflow.podNodeSelector" (dict "Release" .Release "Values" .Values "nodeSelector" .Values.airflow.sync.nodeSelector) }}
+{{- $podTopologySpreadConstraints := include "airflow.podTopologySpreadConstraints" (dict "Release" .Release "Values" .Values "topologySpreadConstraints" .Values.airflow.sync.topologySpreadConstraints) }}
 {{- $podAffinity := include "airflow.podAffinity" (dict "Release" .Release "Values" .Values "affinity" .Values.airflow.sync.affinity) }}
 {{- $podTolerations := include "airflow.podTolerations" (dict "Release" .Release "Values" .Values "tolerations" .Values.airflow.sync.tolerations) }}
 {{- $podSecurityContext := include "airflow.podSecurityContext" (dict "Release" .Release "Values" .Values "securityContext" .Values.airflow.sync.securityContext) }}
@@ -65,6 +66,10 @@ spec:
       nodeSelector:
         {{- $podNodeSelector | nindent 8 }}
       {{- end }}
+      {{- if $podTopologySpreadConstraints }}
+      topologySpreadConstraints:
+        {{- $podTopologySpreadConstraints | nindent 8 }}
+      {{- end }}
       {{- if $podAffinity }}
       affinity:
         {{- $podAffinity | nindent 8 }}