add support for topologyspreadconstraints

wandb · Dec 3, 2024 · d05cdf8 · d05cdf8
1 parent 4103c08
commit d05cdf8
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 0 deletions.
diff --git a/charts/operator-wandb/charts/app/templates/_deployment.tpl b/charts/operator-wandb/charts/app/templates/_deployment.tpl
@@ -45,6 +45,9 @@ spec:
       {{- include "wandb.nodeSelector" . | nindent 6 }}
       {{- include "wandb.priorityClassName" . | nindent 6 }}
       {{- include "wandb.podSecurityContext" .Values.pod.securityContext | nindent 6 }}
+      {{- if .Values.topologySpreadConstraints.enabled }}
+        {{- include "wandb.topologySpreadConstraints" (dict "matchLabels" (fromYaml (include "app.labels" .))) | nindent 6 }}
+      {{- end }}
       terminationGracePeriodSeconds: 60
       initContainers:
         - name: init-db

diff --git a/charts/operator-wandb/charts/app/values.yaml b/charts/operator-wandb/charts/app/values.yaml
@@ -49,6 +49,9 @@ container:
     readOnlyRootFilesystem: false
     privileged: false
 
+topologySpreadConstraints:
+  enabled: false
+
 common:
   labels: {}
 

diff --git a/charts/operator-wandb/templates/_pods.tpl b/charts/operator-wandb/templates/_pods.tpl
@@ -88,3 +88,27 @@ securityContext:
 {{- end }}
 {{- end }}
 
+{{/*
+Return a simplified topologySpreadConstraints definition with dynamic labelSelector and matchLabels passed as a map.
+*/}}
+{{- define "wandb.topologySpreadConstraints" -}}
+{{- $matchLabels := .matchLabels -}}
+topologySpreadConstraints:
+  - maxSkew: 1
+    topologyKey: kubernetes.io/hostname
+    whenUnsatisfiable: ScheduleAnyway
+    labelSelector:
+      matchLabels:
+        {{- range $key, $value := $matchLabels }}
+        {{ $key }}: {{ $value | quote }}
+        {{- end }}
+  - maxSkew: 1
+    topologyKey: topology.kubernetes.io/zone
+    whenUnsatisfiable: ScheduleAnyway
+    labelSelector:
+      matchLabels:
+        {{- range $key, $value := $matchLabels }}
+        {{ $key }}: {{ $value | quote }}
+        {{- end }}
+{{- end }}
+
diff --git a/charts/operator-wandb/values.yaml b/charts/operator-wandb/values.yaml
@@ -204,12 +204,16 @@ parquet:
   image:
     repository: wandb/local
     tag: latest
+  topologySpreadConstraints:
+    enabled: false
 
 app:
   install: true
   image:
     repository: wandb/local
     tag: latest
+  topologySpreadConstraints:
+    enabled: false
 
 nginx:
   install: false
@@ -219,6 +223,8 @@ weave:
   image:
     repository: wandb/local
     tag: latest
+  topologySpreadConstraints:
+    enabled: false
 
 weave-trace:
   install: false
@@ -227,12 +233,16 @@ weave-trace:
     tag: latest
   datadog:
     enabled: false
+  topologySpreadConstraints:
+    enabled: false
 
 console:
   install: true
   image:
     repository: wandb/console
     tag: latest
+  topologySpreadConstraints:
+    enabled: false
 
 flat-run-fields-updater:
   install: false
@@ -241,6 +251,8 @@ flat-run-fields-updater:
     tag: latest
   pubSub:
     subscription: ""
+  topologySpreadConstraints:
+    enabled: false
 
 filestream:
   install: false
@@ -249,18 +261,24 @@ filestream:
     tag: 0.63.0-danielpanzella-filestream-local.1
   pubSub:
     subscription: ""
+  topologySpreadConstraints:
+    enabled: false
 
 executor:
   install: false
   image:
     repository: wandb/local
     tag: latest
+  topologySpreadConstraints:
+    enabled: false
 
 mysql:
   install: false
   persistence:
     size: 20Gi
     storageClass: ""
+  topologySpreadConstraints:
+    enabled: false
 
 yace:
   install: false
@@ -271,6 +289,8 @@ yace:
       prometheus.io/port: "5000"
       prometheus.io/path: "/metrics"
       prometheus.io/scheme: http
+  topologySpreadConstraints:
+    enabled: false
 
 redis:
   install: true
@@ -298,6 +318,9 @@ prometheus:
   mysql-exporter:
     install: true
 
+  topologySpreadConstraints:
+    enabled: false
+
 stackdriver:
   install: false
   pod:
@@ -306,9 +329,13 @@ stackdriver:
       prometheus.io/port: "9255"
       prometheus.io/path: "/metrics"
       prometheus.io/scheme: http
+  topologySpreadConstraints:
+    enabled: false
 
 otel:
   install: true
+  topologySpreadConstraints:
+    enabled: false
 
   # Consider deploying with a deployment rather than a daemonset to avoid
   # redundant metrics. In setups where both are scraping the same endpoints,