From dead283d52163537e8b0f5b5d1272e6cdc7139ba Mon Sep 17 00:00:00 2001
From: invakid404 <invakid404@riseup.net>
Date: Sun, 29 Sep 2024 21:08:39 +0300
Subject: [PATCH] feat: allow deploying worker groups as statefulsets

---
 charts/windmill/templates/_helpers.tpl       | 13 ++++++++++++
 charts/windmill/templates/worker-groups.yaml | 17 ++++++++++++++-
 charts/windmill/values.yaml                  | 22 +++++++++++++++++---
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/charts/windmill/templates/_helpers.tpl b/charts/windmill/templates/_helpers.tpl
index 049df3c..f5096cd 100644
--- a/charts/windmill/templates/_helpers.tpl
+++ b/charts/windmill/templates/_helpers.tpl
@@ -60,3 +60,16 @@ Create the name of the service account to use
 {{- default "default" .Values.serviceAccount.name }}
 {{- end }}
 {{- end }}
+
+{{/*
+Validate controller kind, defaulting to "Deployment"
+*/}}
+{{- define "validateControllerKind" -}}
+{{- $validTypes := list "Deployment" "StatefulSet" -}}
+{{- $inputType := default "Deployment" . -}}
+{{- if has $inputType $validTypes -}}
+{{ $inputType }}
+{{- else -}}
+{{- fail (printf "Invalid controller type: %s. Must be either Deployment or StatefulSet" $inputType) -}}
+{{- end -}}
+{{- end -}}
diff --git a/charts/windmill/templates/worker-groups.yaml b/charts/windmill/templates/worker-groups.yaml
index d637041..e7975ad 100644
--- a/charts/windmill/templates/worker-groups.yaml
+++ b/charts/windmill/templates/worker-groups.yaml
@@ -1,8 +1,9 @@
 {{- range $v := .Values.windmill.workerGroups }}
 {{ if and $v.replicas (gt (int $v.replicas) 0)}}
 ---
+{{- $controllerType := include "validateControllerKind" $v.controller }}
 apiVersion: apps/v1
-kind: Deployment
+kind: {{ $controllerType }}
 metadata:
   name: windmill-workers-{{ $v.name }}
   labels:
@@ -14,11 +15,19 @@ metadata:
     workerGroup: {{ $v.name }}
 spec:
   replicas: {{ $v.replicas }}
+  {{- if eq $controllerType "Deployment" }}
   strategy:
     type: RollingUpdate
     rollingUpdate:
       maxSurge: 3
       maxUnavailable: 0
+  {{- else if eq $controllerType "StatefulSet" }}
+  serviceName: windmill-workers-{{ $v.name }}
+  updateStrategy:
+    type: RollingUpdate
+    rollingUpdate:
+      partition: 0
+  {{- end }}
   selector:
     matchLabels:
       app: windmill-workers
@@ -172,5 +181,11 @@ spec:
       tolerations:
 {{ toYaml . | indent 8 }}
 {{- end }}
+{{- if eq $controllerType "StatefulSet" }}
+  {{- if $v.volumeClaimTemplates }}
+  volumeClaimTemplates:
+    {{- toYaml $v.volumeClaimTemplates | nindent 2 }}
+  {{- end }}
+{{- end }}
 {{- end }}
 {{- end }}
diff --git a/charts/windmill/values.yaml b/charts/windmill/values.yaml
index 0e57be9..77d1adc 100644
--- a/charts/windmill/values.yaml
+++ b/charts/windmill/values.yaml
@@ -74,6 +74,9 @@ windmill:
     # workers configuration
     # The default worker group
     - name: "default"
+      # -- Controller to use. Valid options are "Deployment" and "StatefulSet"
+      controller: "Deployment"
+
       replicas: 3
       # -- Annotations to apply to the pods
       annotations: {}
@@ -99,7 +102,6 @@ windmill:
       # -- Security context to apply to the pod
       containerSecurityContext: {}
 
-
       # -- Affinity rules to apply to the pods
       affinity: {}
 
@@ -120,10 +122,16 @@ windmill:
       volumes: []
       volumeMounts: []
 
+      # -- Volume claim templates. Only applies when controller is "StatefulSet"
+      volumeClaimTemplates: []
+
       # -- command override
       command: []
 
     - name: "native"
+      # -- Controller to use. Valid options are "Deployment" and "StatefulSet"
+      controller: "Deployment"
+
       replicas: 1
       # -- Annotations to apply to the pods
       annotations: {}
@@ -146,7 +154,6 @@ windmill:
       # -- Security context to apply to the pod
       containerSecurityContext: {}
 
-
       # -- Affinity rules to apply to the pods
       affinity: {}
 
@@ -170,7 +177,13 @@ windmill:
       volumes: []
       volumeMounts: []
 
+      # -- Volume claim templates. Only applies when controller is "StatefulSet"
+      volumeClaimTemplates: []
+
     - name: "gpu"
+      # -- Controller to use. Valid options are "Deployment" and "StatefulSet"
+      controller: "Deployment"
+
       replicas: 0
       # -- Annotations to apply to the pods
       annotations: {}
@@ -193,7 +206,6 @@ windmill:
       # -- Security context to apply to the pod
       containerSecurityContext: {}
 
-
       # -- Affinity rules to apply to the pods
       affinity: {}
 
@@ -210,8 +222,12 @@ windmill:
       volumes: []
       volumeMounts: []
 
+      # -- Volume claim templates. Only applies when controller is "StatefulSet"
+      volumeClaimTemplates: []
+
       # -- command override
       command: []
+
   # app configuration
   app:
     # -- Annotations to apply to the pods