Set cluster cookie as label on workers (#106)

* Separate keywords `worker_prefix` and `manager_pod_name` * Set cluster cookie as label on workers * Add default to worker_pod_spec!
beacon-biosignals · Apr 13, 2023 · 146cd9b · 146cd9b
1 parent e8a90b3
commit 146cd9b
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/src/native_driver.jl b/src/native_driver.jl
@@ -96,7 +96,7 @@ function Distributed.launch(manager::K8sClusterManager, params::Dict, launched::
 
     cmd = `$exename $exeflags --worker=$(cluster_cookie())`
 
-    worker_manifest = worker_pod_spec(manager; cmd)
+    worker_manifest = worker_pod_spec(manager; cmd, cluster_cookie=cluster_cookie())
 
     # Note: User-defined `configure` function may or may-not be mutating
     worker_manifest = manager.configure(worker_manifest)

diff --git a/src/pod.jl b/src/pod.jl
@@ -196,11 +196,13 @@ function worker_pod_spec!(pod::AbstractDict;
                           cmd::Cmd,
                           cpu=DEFAULT_WORKER_CPU,
                           memory=DEFAULT_WORKER_MEMORY,
+                          cluster_cookie::AbstractString=cluster_cookie(),
                           service_account_name=nothing)
     pod["metadata"]["generateName"] = "$(worker_prefix)-"
 
     # Set a label with the `worker_prefix` to support easy termination of all workers
     pod["metadata"]["labels"]["worker-prefix"] = worker_prefix
+    pod["metadata"]["labels"]["cluster-cookie"] = cluster_cookie
 
     worker_container =
         rdict("name" => "worker",

diff --git a/test/pod.jl b/test/pod.jl
@@ -87,7 +87,7 @@ end
 end
 
 @testset "worker_pod_spec" begin
-    kwargs = (; worker_prefix="test-wkr", image="julia", cmd=`julia`)
+    kwargs = (; worker_prefix="test-wkr", image="julia", cmd=`julia`, cluster_cookie="🍪")
     pod = K8sClusterManagers.worker_pod_spec(; kwargs...)
 
     @test keys(pod) == Set(["apiVersion", "kind", "metadata", "spec"])
@@ -96,8 +96,9 @@ end
 
     @test keys(pod["metadata"]) == Set(["generateName", "labels"])
     @test pod["metadata"]["generateName"] == "test-wkr-"
-    @test keys(pod["metadata"]["labels"]) == Set(["worker-prefix"])
+    @test keys(pod["metadata"]["labels"]) == Set(["worker-prefix", "cluster-cookie"])
     @test pod["metadata"]["labels"]["worker-prefix"] == "test-wkr"
+    @test pod["metadata"]["labels"]["cluster-cookie"] == "🍪"
 
     @test pod["spec"]["restartPolicy"] == "Never"
     @test length(pod["spec"]["containers"]) == 1