add support for modifying pod image when checkpointing (#1729)

scheduler/src/cook/kubernetes/api.clj

@@ -724,6 +724,21 @@
               checkpoint))
           checkpoint)))))
 
+(defn calculate-effective-image
+  "Transform the supplied job's image as necessary. e.g. do special transformation if checkpointing is enabled
+  and an image transformation function is supplied."
+  [{:keys [calculate-effective-image-fn] :as kubernetes-config} job-submit-time image {:keys [mode]} task-id]
+  (if (and mode calculate-effective-image-fn)
+    (try
+      ((util/lazy-load-var-memo calculate-effective-image-fn) kubernetes-config job-submit-time image)
+      (catch Exception e
+        (log/error e "Error calculating effective image for checkpointing"
+                   {:calculate-effective-image-fn calculate-effective-image-fn
+                    :image image
+                    :task-id task-id})
+        image))
+    image))
+
 (defn job->pod-labels
   "Returns the dictionary of labels that should be
   added to the job's pod based on the job's labels
@@ -779,6 +794,8 @@
         {:keys [volumes volume-mounts sandbox-volume-mount-fn]} (make-volumes volumes sandbox-dir)
         {:keys [custom-shell init-container set-container-cpu-limit? sidecar]} (config/kubernetes)
         checkpoint (calculate-effective-checkpointing-config job task-id)
+        job-submit-time (tools/job->submit-time job)
+        image (calculate-effective-image (config/kubernetes) job-submit-time image checkpoint task-id)
         checkpoint-memory-overhead (:memory-overhead checkpoint)
         use-cook-init? (and init-container pod-supports-cook-init?)
         use-cook-sidecar? (and sidecar pod-supports-cook-sidecar?)

scheduler/src/cook/kubernetes/controller.clj

@@ -311,26 +311,21 @@
   ; At this point, we don't care about the launch pod or the metric timers, so toss their dictionary away.
   {:cook-expected-state :cook-expected-state/running})
 
-(def get-pod-ip->hostname-fn
-  (memoize
-    (fn [pod-ip->hostname-fn]
-      (if pod-ip->hostname-fn (util/lazy-load-var pod-ip->hostname-fn) identity))))
-
 (defn record-sandbox-url
   "Record the sandbox file server URL in datomic."
   [pod-name {:keys [pod]}]
   (when-not (api/synthetic-pod? pod-name)
     (let [task-id (-> pod .getMetadata .getName)
           pod-ip (-> pod .getStatus .getPodIP)
-          {:keys [default-workdir pod-ip->hostname-fn sidecar]} (config/kubernetes)
+          {:keys [default-workdir sidecar]} (config/kubernetes)
           sandbox-fileserver-port (:port sidecar)
           sandbox-health-check-endpoint (:health-check-endpoint sidecar)
           sandbox-url (try
                         (when (and sandbox-fileserver-port
                                    sandbox-health-check-endpoint
                                    (not (str/blank? pod-ip)))
                           (str "http://"
-                               ((get-pod-ip->hostname-fn pod-ip->hostname-fn) pod-ip)
+                               pod-ip
                                ":" sandbox-fileserver-port
                                "/files/read.json?path="
                                (URLEncoder/encode default-workdir "UTF-8")))

scheduler/src/cook/rest/api.clj

@@ -1276,8 +1276,7 @@
                                   (->> [attribute (str/upper-case (name operator)) pattern]
                                        (map str)))))
           instances (map #(fetch-instance-map db %1) (:job/instance job))
-          submit-time (when (:job/submit-time job) ; due to a bug, submit time may not exist for some jobs
-                        (.getTime (:job/submit-time job)))
+          submit-time (util/job->submit-time job)
           datasets (when (seq (:job/datasets job))
                      (dl/get-dataset-maps job))
           attempts-consumed (util/job-ent->attempts-consumed db job)

scheduler/src/cook/tools.clj

@@ -1030,3 +1030,8 @@
                   (format "%.3f" %)
                   (str %))
                resource-map))
+
+(defn job->submit-time
+  "Get submit-time for a job. due to a bug, submit time may not exist for some jobs"
+  [job]
+  (when (:job/submit-time job) (.getTime (:job/submit-time job))))

scheduler/src/cook/util.clj

@@ -161,6 +161,9 @@
         resolved
         (throw (ex-info "Unable to resolve var, is it valid?" {:var-sym var-sym}))))))
 
+(def lazy-load-var-memo
+  (memoize lazy-load-var))
+
 (def ZeroInt
   (s/both s/Int (s/pred zero? 'zero?)))