From ff4261d07c01f54214caa3a1461549a39dc9377b Mon Sep 17 00:00:00 2001 From: sbias <63045405+s-ducks@users.noreply.github.com> Date: Mon, 13 Nov 2023 18:10:29 +0000 Subject: [PATCH 1/3] add change --- dask_kubernetes/operator/controller/controller.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/dask_kubernetes/operator/controller/controller.py b/dask_kubernetes/operator/controller/controller.py index 00c47aeed..1dc6b3416 100644 --- a/dask_kubernetes/operator/controller/controller.py +++ b/dask_kubernetes/operator/controller/controller.py @@ -64,6 +64,17 @@ def _get_labels(meta): } +def _consolidate_env_vars(existing_env_vars, additional_env_vars): + existing_env_names = {d["name"] for d in existing_env_vars} + additional_env_names = {d["name"] for d in additional_env_vars} + + overlapping_env_names = existing_env_names.intersection(additional_env_names) + + additional_env_vars_to_keep = [d for d in additional_env_vars if d["name"] not in overlapping_env_names] + + return [*existing_env_vars, *additional_env_vars_to_keep] + + def build_scheduler_deployment_spec( cluster_name, namespace, pod_spec, annotations, labels ): @@ -159,6 +170,8 @@ def build_worker_deployment_spec( ] for i in range(len(deployment_spec["spec"]["template"]["spec"]["containers"])): if "env" in deployment_spec["spec"]["template"]["spec"]["containers"][i]: + existing_env_vars = deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"] + _consolidate_env_vars(existing_env_vars, env) deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"].extend( env ) @@ -197,6 +210,8 @@ def build_job_pod_spec(job_name, cluster_name, namespace, spec, annotations, lab ] for i in range(len(pod_spec["spec"]["containers"])): if "env" in pod_spec["spec"]["containers"][i]: + existing_env_vars = deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"] + _consolidate_env_vars(existing_env_vars, env) pod_spec["spec"]["containers"][i]["env"].extend(env) else: pod_spec["spec"]["containers"][i]["env"] = env From e842e3cabf5e18d45fc6390aa066aa61b0ae592e Mon Sep 17 00:00:00 2001 From: sbias <63045405+s-ducks@users.noreply.github.com> Date: Tue, 21 Nov 2023 15:21:27 +0000 Subject: [PATCH 2/3] actually add the env vars --- dask_kubernetes/operator/controller/controller.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/dask_kubernetes/operator/controller/controller.py b/dask_kubernetes/operator/controller/controller.py index 1dc6b3416..db1dd5479 100644 --- a/dask_kubernetes/operator/controller/controller.py +++ b/dask_kubernetes/operator/controller/controller.py @@ -171,10 +171,8 @@ def build_worker_deployment_spec( for i in range(len(deployment_spec["spec"]["template"]["spec"]["containers"])): if "env" in deployment_spec["spec"]["template"]["spec"]["containers"][i]: existing_env_vars = deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"] - _consolidate_env_vars(existing_env_vars, env) - deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"].extend( - env - ) + all_env_vars = _consolidate_env_vars(existing_env_vars, env) + deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"] = all_env_vars else: deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"] = env return deployment_spec @@ -211,8 +209,8 @@ def build_job_pod_spec(job_name, cluster_name, namespace, spec, annotations, lab for i in range(len(pod_spec["spec"]["containers"])): if "env" in pod_spec["spec"]["containers"][i]: existing_env_vars = deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"] - _consolidate_env_vars(existing_env_vars, env) - pod_spec["spec"]["containers"][i]["env"].extend(env) + all_env_vars = _consolidate_env_vars(existing_env_vars, env) + pod_spec["spec"]["containers"][i]["env"] = all_env_vars else: pod_spec["spec"]["containers"][i]["env"] = env return pod_spec From 255e70229de71238e14c97a7755960a9bd32749d Mon Sep 17 00:00:00 2001 From: sbias <63045405+s-ducks@users.noreply.github.com> Date: Mon, 27 Nov 2023 22:21:05 +0000 Subject: [PATCH 3/3] use deepcopy --- dask_kubernetes/operator/controller/controller.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/dask_kubernetes/operator/controller/controller.py b/dask_kubernetes/operator/controller/controller.py index db1dd5479..e954f79a6 100644 --- a/dask_kubernetes/operator/controller/controller.py +++ b/dask_kubernetes/operator/controller/controller.py @@ -2,6 +2,7 @@ import time from collections import defaultdict from contextlib import suppress +from copy import deepcopy from datetime import datetime from uuid import uuid4 @@ -208,7 +209,7 @@ def build_job_pod_spec(job_name, cluster_name, namespace, spec, annotations, lab ] for i in range(len(pod_spec["spec"]["containers"])): if "env" in pod_spec["spec"]["containers"][i]: - existing_env_vars = deployment_spec["spec"]["template"]["spec"]["containers"][i]["env"] + existing_env_vars = pod_spec["spec"]["template"]["spec"]["containers"][i]["env"] all_env_vars = _consolidate_env_vars(existing_env_vars, env) pod_spec["spec"]["containers"][i]["env"] = all_env_vars else: @@ -594,7 +595,7 @@ async def daskworkergroup_replica_update( namespace=namespace, cluster_name=cluster_name, uuid=uuid4().hex[:10], - pod_spec=worker_spec["spec"], + pod_spec=deepcopy(worker_spec["spec"]), annotations=annotations, labels=labels, )