From 4f392388030e9d013891f98e292fc48568095b84 Mon Sep 17 00:00:00 2001 From: Jonas Dedden Date: Wed, 3 Apr 2024 09:03:14 +0200 Subject: [PATCH 1/2] Implement new container env mechanism also for Dask Jobs --- .../operator/controller/controller.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/dask_kubernetes/operator/controller/controller.py b/dask_kubernetes/operator/controller/controller.py index 13702d97..cd52a9eb 100644 --- a/dask_kubernetes/operator/controller/controller.py +++ b/dask_kubernetes/operator/controller/controller.py @@ -197,17 +197,19 @@ def build_job_pod_spec(job_name, cluster_name, namespace, spec, annotations, lab }, "spec": copy.deepcopy(spec), } - env = [ - { + scheduler_env = { "name": "DASK_SCHEDULER_ADDRESS", "value": f"tcp://{cluster_name}-scheduler.{namespace}.svc.cluster.local:8786", - }, - ] + } for container in pod_spec["spec"]["containers"]: - if "env" in container: - container["env"].extend(env) - else: - container["env"] = env + if "env" not in container: + container["env"] = [scheduler_env] + continue + + container_env_names = [env_item["name"] for env_item in container["env"]] + + if "DASK_SCHEDULER_ADDRESS" not in container_env_names: + container["env"].append(scheduler_env) return pod_spec From 2c4aa0aa0283171bf6f93626bd3d937ac0994e40 Mon Sep 17 00:00:00 2001 From: Jonas Dedden Date: Wed, 3 Apr 2024 09:20:25 +0200 Subject: [PATCH 2/2] Black fixes --- dask_kubernetes/operator/controller/controller.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dask_kubernetes/operator/controller/controller.py b/dask_kubernetes/operator/controller/controller.py index cd52a9eb..1254ee52 100644 --- a/dask_kubernetes/operator/controller/controller.py +++ b/dask_kubernetes/operator/controller/controller.py @@ -198,8 +198,8 @@ def build_job_pod_spec(job_name, cluster_name, namespace, spec, annotations, lab "spec": copy.deepcopy(spec), } scheduler_env = { - "name": "DASK_SCHEDULER_ADDRESS", - "value": f"tcp://{cluster_name}-scheduler.{namespace}.svc.cluster.local:8786", + "name": "DASK_SCHEDULER_ADDRESS", + "value": f"tcp://{cluster_name}-scheduler.{namespace}.svc.cluster.local:8786", } for container in pod_spec["spec"]["containers"]: if "env" not in container: