Fix issue of replicated environment variables

dask · Feb 22, 2024 · fcc379a · fcc379a
1 parent 726f0f8
commit fcc379a
Showing 1 changed file with 28 additions and 27 deletions.
diff --git a/dask_kubernetes/operator/controller/controller.py b/dask_kubernetes/operator/controller/controller.py
@@ -1,4 +1,5 @@
 import asyncio
+import copy
 import time
 from collections import defaultdict
 from contextlib import suppress
@@ -79,14 +80,15 @@ def build_scheduler_deployment_spec(
         "labels": labels,
         "annotations": annotations,
     }
-    spec = {}
-    spec["replicas"] = 1
-    spec["selector"] = {
-        "matchLabels": labels,
-    }
-    spec["template"] = {
-        "metadata": metadata,
-        "spec": pod_spec,
+    spec = {
+        "replicas": 1,
+        "selector": {
+            "matchLabels": labels,
+        },
+        "template": {
+            "metadata": metadata,
+            "spec": pod_spec,
+        },
     }
     return {
         "apiVersion": "apps/v1",
@@ -132,14 +134,15 @@ def build_worker_deployment_spec(
         "labels": labels,
         "annotations": annotations,
     }
-    spec = {}
-    spec["replicas"] = 1  # make_worker_spec returns dict with a replicas key?
-    spec["selector"] = {
-        "matchLabels": labels,
-    }
-    spec["template"] = {
-        "metadata": metadata,
-        "spec": pod_spec,
+    spec = {
+        "replicas": 1,
+        "selector": {
+            "matchLabels": labels,
+        },
+        "template": {
+            "metadata": metadata,
+            "spec": copy.deepcopy(pod_spec),
+        },
     }
     deployment_spec = {
         "apiVersion": "apps/v1",
@@ -157,13 +160,11 @@ def build_worker_deployment_spec(
             "value": f"tcp://{cluster_name}-scheduler.{namespace}.svc.cluster.local:8786",
         },
     ]
-    for i in range(len(deployment_spec["spec"]["template"]["spec"]["containers"])):
-        if "env" in deployment_spec["spec"]["template"]["spec"]["containers"][i]:
-            deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"].extend(
-                env
-            )
+    for container in deployment_spec["spec"]["template"]["spec"]["containers"]:
+        if "env" in container:
+            container["env"].extend(env)
         else:
-            deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"] = env
+            container["env"] = env
     return deployment_spec
 
 
@@ -187,19 +188,19 @@ def build_job_pod_spec(job_name, cluster_name, namespace, spec, annotations, lab
             "labels": labels,
             "annotations": annotations,
         },
-        "spec": spec,
+        "spec": copy.deepcopy(spec),
     }
     env = [
         {
             "name": "DASK_SCHEDULER_ADDRESS",
             "value": f"tcp://{cluster_name}-scheduler.{namespace}.svc.cluster.local:8786",
         },
     ]
-    for i in range(len(pod_spec["spec"]["containers"])):
-        if "env" in pod_spec["spec"]["containers"][i]:
-            pod_spec["spec"]["containers"][i]["env"].extend(env)
+    for container in pod_spec["spec"]["containers"]:
+        if "env" in container:
+            container["env"].extend(env)
         else:
-            pod_spec["spec"]["containers"][i]["env"] = env
+            container["env"] = env
     return pod_spec