Fix worker batching by ensuring env vars are passed correctly (#770)

* Ensure env vars are passed correctly * Fix small bug
dask · Jul 14, 2023 · e3b1786 · e3b1786
1 parent 3defed7
commit e3b1786
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 deletions.
diff --git a/dask_kubernetes/operator/controller/controller.py b/dask_kubernetes/operator/controller/controller.py
@@ -634,8 +634,13 @@ async def daskworkergroup_replica_update(
                 if "labels" in worker_spec["metadata"]:
                     labels.update(**worker_spec["metadata"]["labels"])
 
-            SIZE = dask.config.get("kubernetes.controller.worker-allocation.batch-size")
-            DELAY = dask.config.get("kubernetes.controller.worker-allocation.delay")
+            SIZE = int(
+                dask.config.get("kubernetes.controller.worker-allocation.batch-size")
+                or 0
+            )
+            DELAY = int(
+                dask.config.get("kubernetes.controller.worker-allocation.delay") or 0
+            )
             batch_size = min(workers_needed, SIZE) if SIZE else workers_needed
             if workers_needed > 0:
                 for _ in range(batch_size):

diff --git a/dask_kubernetes/operator/deployment/helm/dask-kubernetes-operator/templates/deployment.yaml b/dask_kubernetes/operator/deployment/helm/dask-kubernetes-operator/templates/deployment.yaml
@@ -33,12 +33,14 @@ spec:
           image: "{{ .Values.image.name }}:{{ .Values.image.tag }}"
           imagePullPolicy: {{ .Values.image.pullPolicy }}
           env:
+            {{- with .Values.workerAllocation.size }}
             - name: DASK_KUBERNETES__CONTROLLER__WORKER_ALLOCATION__BATCH_SIZE
-              value:
-                {{- toYaml .Values.workerAllocation.size | nindent 16 }}
+              value: {{ . | quote }}
+            {{- end }}
+            {{- with .Values.workerAllocation.delay }}
             - name: DASK_KUBERNETES__CONTROLLER__WORKER_ALLOCATION__DELAY
-              value:
-                {{- toYaml .Values.workerAllocation.delay | nindent 16 }}
+              value: {{ . | quote }}
+            {{- end }}
           args:
             - --liveness=http://0.0.0.0:8080/healthz
           {{- with .Values.kopfArgs }}