From d05cdf89cb7b1c6d7836e86fadd2c80bf771d0b4 Mon Sep 17 00:00:00 2001
From: velotioaastha <aastha.gupta@velotio.com>
Date: Tue, 3 Dec 2024 17:51:14 +0530
Subject: [PATCH] add support for topologyspreadconstraints

---
 .../charts/app/templates/_deployment.tpl      |  3 +++
 charts/operator-wandb/charts/app/values.yaml  |  3 +++
 charts/operator-wandb/templates/_pods.tpl     | 24 +++++++++++++++++
 charts/operator-wandb/values.yaml             | 27 +++++++++++++++++++
 4 files changed, 57 insertions(+)

diff --git a/charts/operator-wandb/charts/app/templates/_deployment.tpl b/charts/operator-wandb/charts/app/templates/_deployment.tpl
index b1b7ec33..0a9ac284 100644
--- 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
index 65256119..e45feda3 100644
--- 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
index 76d256f9..24f72f45 100644
--- 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
index 749d6cf1..11df3277 100644
--- 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,