Skip to content
This repository has been archived by the owner on Apr 24, 2023. It is now read-only.

Add Kubernetes pod annotations from job labels #2056

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Add Kubernetes pod annotations from job labels #2056

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
612bdcf
Add config.edn and K8s api code for pod-annotations map feature
laurameng Mar 18, 2022
18d4e0b
Add K8s api test for job-labels -> pod-annotations feature
laurameng Mar 18, 2022
ecbebea
Add duplicates check and fix some styling
laurameng Mar 18, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions scheduler/config-k8s.edn
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,10 @@
:factory-fn "cook.plugins.demo-plugin/launch-factory"}}
:hostname #config/env "COOK_HOSTNAME"
:kubernetes {:add-job-label-to-pod-prefix "platform/"
:job-label-to-pod-annotation-map { "job-label1" {"annotation-1" "value-1",
"annotation-2" "value-2"},
"job-label2" {"annotation-3" "value-3",
"annotation-4" "value-4"}}
:clobber-synthetic-pods true
:disallowed-container-paths #{"/mnt/bad"}
:disallowed-var-names #{"BADVAR"}
Expand Down
56 changes: 40 additions & 16 deletions scheduler/src/cook/kubernetes/api.clj
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1307,6 +1307,44 @@
:resolved-config resolved-config}))
resolved-config))))

(defn get-use-all-gids-annotation
"Helper method to get use-all-gids pod annotation "
[task-id]
; For non-synthetic (real job) pods, when configured to do so,
; we add a pod annotation indicating that Kubernetes should
; use all of the group IDs the user is a member of, in order
; to avoid contradictions between which group Cook thinks a
; user belongs to and which group Kubernetes thinks the user
; belongs to.
(let [[resolved-config _]
(config-incremental/resolve-incremental-config
task-id :add-use-all-gids-annotation "false")
use-all-gids-annotation-name
(:use-all-gids-annotation-name (config/kubernetes))]
(if
(and
(= "true" resolved-config)
use-all-gids-annotation-name)
{use-all-gids-annotation-name "true"}
{}))
)

(defn job-label->pod-annotations
"Given a job, return all pod annotations configured based on the job's labels"
[job]
(let [{:keys [job-label-to-pod-annotation-map]} (config/kubernetes)
requested-pod-annotations
(-> job
(tools/job-ent->label)
(get "add-pod-annotation" "")
; the user can pass us multiple comma-separated values
(str/split #",")
)]
(->> requested-pod-annotations
(select-keys job-label-to-pod-annotation-map)
(vals)
(into {}))))

(defn ^V1Pod task-metadata->pod
"Given a task-request and other data generate the kubernetes V1Pod to launch that task."
[namespace {:keys [cook-pool-taint-name cook-pool-taint-prefix cook-pool-label-name] compute-cluster-name :name}
Expand Down Expand Up @@ -1445,22 +1483,8 @@
(let [pod-annotations'
(if (synthetic-pod? pod-name)
pod-annotations
; For non-synthetic (real job) pods, when configured to do so,
; we add a pod annotation indicating that Kubernetes should
; use all of the group IDs the user is a member of, in order
; to avoid contradictions between which group Cook thinks a
; user belongs to and which group Kubernetes thinks the user
; belongs to.
(let [[resolved-config _]
(config-incremental/resolve-incremental-config
task-id :add-use-all-gids-annotation "false")
use-all-gids-annotation-name
(:use-all-gids-annotation-name (config/kubernetes))]
(cond-> pod-annotations
(and
(= "true" resolved-config)
use-all-gids-annotation-name)
(assoc use-all-gids-annotation-name "true"))))]
; add additional annotations for real pods
(merge (job-label->pod-annotations job) (get-use-all-gids-annotation task-id) pod-annotations))]
(when (seq pod-annotations')
(.setAnnotations metadata pod-annotations')))

Expand Down
90 changes: 89 additions & 1 deletion scheduler/test/cook/test/kubernetes/api.clj
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -582,7 +582,95 @@
(is (= "required-cook-job-container" (.getName container)))
(is (not (set/subset?
#{"TEST_AGENT"}
(->> container-env (map #(.getName %)) set))))))))))
(->> container-env (map #(.getName %)) set)))))))

(testing "job-labels->pod-annotations"
(with-redefs [config/kubernetes (constantly {:job-label-to-pod-annotation-map {"label1" {"k1" "v1", "k2" "v2"},
"label2" {"k3" "v3", "k4" "v4"},
"label3" {"ka" "va", "kb" "vb", "kc" "vc"}}})]
; No labels
(let [task-metadata {:command {:user "test-user"}
:task-request {:job {:job/label []}
:scalar-requests {"mem" 512 "cpus" 1.0}}}]
(let [^V1Pod pod (api/task-metadata->pod "test-namespace"
fake-cc-config
task-metadata)
pod-annotations (-> pod .getMetadata .getAnnotations)]
(is (empty? (select-keys pod-annotations ["k1", "k2", "k3", "k4", "ka", "kb", "kc"])))))

; Single label
(let [task-metadata {:command {:user "test-user"}
:task-request {:job {:job/label [{:label/key "add-pod-annotation"
:label/value "label1"}
{:label/key "platform/baz"
:label/value "qux"}
{:label/key "platform/another"
:label/value "included"}]}
:scalar-requests {"mem" 512 "cpus" 1.0}}}]
; Simple match
(let [^V1Pod pod (api/task-metadata->pod "test-namespace"
fake-cc-config
task-metadata)
pod-annotations (-> pod .getMetadata .getAnnotations)]
(is (find pod-annotations "k1"))
(is (find pod-annotations "k2"))
(is (empty? (select-keys pod-annotations ["k3", "k4", "ka", "kb", "kc"]))))
; No pod-annotations
(with-redefs [config/kubernetes (constantly {})]
(let [^V1Pod pod (api/task-metadata->pod "test-namespace"
fake-cc-config
task-metadata)
pod-annotations (-> pod .getMetadata .getAnnotations)]
(is (empty? (select-keys pod-annotations ["k1", "k2", "k3", "k4", "ka", "kb", "kc"]))))))
; Comma-delimited multi-match
(let [task-metadata {:command {:user "test-user"}
:task-request {:job {:job/label [{:label/key "add-pod-annotation"
:label/value "label3,label1"}]}
:scalar-requests {"mem" 512 "cpus" 1.0}}}]
(let [^V1Pod pod (api/task-metadata->pod "test-namespace"
fake-cc-config
task-metadata)
pod-annotations (-> pod .getMetadata .getAnnotations)]
(is (find pod-annotations "k1"))
(is (find pod-annotations "k2"))
(is (find pod-annotations "ka"))
(is (find pod-annotations "kb"))
(is (find pod-annotations "kc"))
(is (empty? (select-keys pod-annotations ["k3", "k4"])))))
; Comma-delimited partial-match
(let [task-metadata {:command {:user "test-user"}
:task-request {:job {:job/label [{:label/key "add-pod-annotation"
:label/value "label2,not-yet-defined"}]}
:scalar-requests {"mem" 512 "cpus" 1.0}}}]
(let [^V1Pod pod (api/task-metadata->pod "test-namespace"
fake-cc-config
task-metadata)
pod-annotations (-> pod .getMetadata .getAnnotations)]
(is (find pod-annotations "k3"))
(is (find pod-annotations "k4"))
(is (empty? (select-keys pod-annotations ["k1", "k2", "ka", "kb", "kc"])))))
; Comma-delimited duplicates
(let [task-metadata {:command {:user "test-user"}
:task-request {:job {:job/label [{:label/key "add-pod-annotation"
:label/value "label1,label1"}]}
:scalar-requests {"mem" 512 "cpus" 1.0}}}]
(let [^V1Pod pod (api/task-metadata->pod "test-namespace"
fake-cc-config
task-metadata)
pod-annotations (-> pod .getMetadata .getAnnotations)]
(is (find pod-annotations "k1"))
(is (find pod-annotations "k2"))
(is (empty? (select-keys pod-annotations ["k3", "k4", "ka", "kb", "kc"])))))
; No-matches
(let [task-metadata {:command {:user "test-user"}
:task-request {:job {:job/label [{:label/key "add-pod-annotation"
:label/value "not-yet-defined"}]}
:scalar-requests {"mem" 512 "cpus" 1.0}}}]
(let [^V1Pod pod (api/task-metadata->pod "test-namespace"
fake-cc-config
task-metadata)
pod-annotations (-> pod .getMetadata .getAnnotations)]
(is (empty? (select-keys pod-annotations ["k1", "k2", "k3", "k4", "ka", "kb", "kc"]))))))))))

(defn- k8s-volume->clj [^V1Volume volume]
{:name (.getName volume)
Expand Down