rework annotation map for opt b

twosigma · Apr 1, 2022 · 2a6c7f6 · 2a6c7f6
1 parent 6d0abb6
commit 2a6c7f6
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 8 deletions.
diff --git a/scheduler/config-k8s.edn b/scheduler/config-k8s.edn
@@ -82,6 +82,7 @@
                                :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-lookup-key "add-pod-annotation"
               :job-label-to-pod-annotation-map { "job-label1" {"annotation-1" "value-1",
                                                              "annotation-2" "value-2"},
                                                 "job-label2" {"annotation-3" "value-3",

diff --git a/scheduler/src/cook/kubernetes/api.clj b/scheduler/src/cook/kubernetes/api.clj
@@ -1335,14 +1335,15 @@
 (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)
+  (let [{:keys [job-label-to-pod-annotation-map job-label-to-pod-annotation-lookup-key]} (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 #",")
-          )]
+        (if job-label-to-pod-annotation-lookup-key
+          (-> job
+            (tools/job-ent->label)
+            (get job-label-to-pod-annotation-lookup-key "")
+            ; the user can pass us multiple comma-separated values
+            (str/split #","))
+          "")]
     (->> requested-pod-annotations
          (select-keys job-label-to-pod-annotation-map)
          (vals)

diff --git a/scheduler/test/cook/test/kubernetes/api.clj b/scheduler/test/cook/test/kubernetes/api.clj
@@ -585,7 +585,8 @@
                        (->> 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"},
+        (with-redefs [config/kubernetes (constantly {:job-label-to-pod-annotation-lookup-key "add-pod-annotation"
+                                                     :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
@@ -615,7 +616,33 @@
               (is (find pod-annotations "k1"))
               (is (find pod-annotations "k2"))
               (is (empty? (select-keys pod-annotations ["k3", "k4", "ka", "kb", "kc"]))))
+            ; No lookup-key
+            (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"}}})]
+              (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"])))))
+            ; Empty lookup-key
+            (with-redefs [config/kubernetes (constantly {:job-label-to-pod-annotation-lookup-key ""
+                                                         :job-label-to-pod-annotation-map {"label1" {"k1" "v1", "k2" "v2"},
+                                                                                           "label2" {"k3" "v3", "k4" "v4"},
+                                                                                           "label3" {"ka" "va", "kb" "vb", "kc" "vc"}}})]
+              (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"])))))
             ; No pod-annotations
+            (with-redefs [config/kubernetes (constantly {:job-label-to-pod-annotation-lookup-key "add-pod-annotation"})]
+              (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"])))))
+            ; Neither lookup-key nor pod-annotations
             (with-redefs [config/kubernetes (constantly {})]
               (let [^V1Pod pod (api/task-metadata->pod "test-namespace"
                                                        fake-cc-config