diff --git a/scheduler/config-k8s.edn b/scheduler/config-k8s.edn index 846bb85842..ae0ec12dce 100644 --- a/scheduler/config-k8s.edn +++ b/scheduler/config-k8s.edn @@ -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"} diff --git a/scheduler/src/cook/kubernetes/api.clj b/scheduler/src/cook/kubernetes/api.clj index d947b10856..14db4f9b7c 100644 --- a/scheduler/src/cook/kubernetes/api.clj +++ b/scheduler/src/cook/kubernetes/api.clj @@ -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} @@ -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'))) diff --git a/scheduler/test/cook/test/kubernetes/api.clj b/scheduler/test/cook/test/kubernetes/api.clj index 57356c060f..37a4275ba5 100644 --- a/scheduler/test/cook/test/kubernetes/api.clj +++ b/scheduler/test/cook/test/kubernetes/api.clj @@ -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)