Repository: twosigma/Cook (language: python)

Issue:
allow preemptions for tasks with unknown status
## Changes proposed in this PR

- Fix a bug when saving preemption state for a task. Tasks with unknown status would fail to save.
- New variant of `ensure` function added `ensure-some`, where value can be one of a provided set of values.
- Added simple test of the status gating during the save.
- Remove some rebalancer code documentation that doesn't add to understanding.


## Why are we making these changes?

- Both unknown and running tasks can be preempted.





Discussion / hints:

Fix patch:
diff --git a/.gitignore b/.gitignore
index a51d55a5de..b5d8196820 100644
--- a/.gitignore
+++ b/.gitignore
@@ -44,3 +44,5 @@ dist
 .vagrant/
 jobclient/python/docs/build/
 *.egg-info/
+.clj-kondo/
+.lsp/
diff --git a/scheduler/src/cook/rebalancer.clj b/scheduler/src/cook/rebalancer.clj
index b8c62ad297..d120bb2e16 100644
--- a/scheduler/src/cook/rebalancer.clj
+++ b/scheduler/src/cook/rebalancer.clj
@@ -14,8 +14,7 @@
 ;; limitations under the License.
 ;;
 (ns cook.rebalancer
-  (:require [chime :refer [chime-at]]
-            [clojure.core.async :as async]
+  (:require [clojure.core.async :as async]
             [clojure.core.cache :as cache]
             [clojure.data.priority-map :as pm]
             [clojure.tools.logging :as log]
@@ -28,7 +27,7 @@
             [cook.scheduler.dru :as dru]
             [cook.task :as task]
             [cook.tools :as util]
-            [datomic.api :as d :refer [q]]
+            [datomic.api :as d]
             [metatransaction.core :as mt]
             [metrics.histograms :as histograms]
             [metrics.timers :as timers]
@@ -56,12 +55,8 @@
 ;;;
 ;;; Within each cycle, the rebalancer takes the first n tasks in a global DRU queue, it then 'tries' to see if they can
 ;;; (and should) preempt any existing running tasks. Only tasks over the share are eligible to be preempted. We only do
-;;; the preemption if the DRU change is sufficiently large. If there is a set of running tasks that satisfy the criteria, they are killed and the host is reserved to launch the waiting job.
-;;; be killed. 
-;;;
-;;; At the start of a cycle, Rebalancer initializes its internal state. Then, for each iteration in a given cycle,
-;;; Rebalancer processes a pending job and tries to make room for it by finding a task to preempt and updates its
-;;; internal state if such preemption is found.
+;;; the preemption if the DRU change is sufficiently large. If there is a set of running tasks that satisfy the 
+;;; criteria, they are killed and the host is reserved to launch the waiting job.
 ;;;
 ;;; Preemption Principle
 ;;; Rebalancer uses a score-based preemption algorithm.
@@ -472,26 +467,26 @@
 
 (defn transact-preemption!
   "Transacts the rebalancer preemption for the given task-ent"
-  [db conn pool-name task-ent]
+  [db conn pool-name task-ent] 
   (try
     @(d/transact
-       conn
+      conn
        ;; Make :instance/status and :instance/preempted? consistent to simplify the state machine.
        ;; We don't want to deal with {:instance/status :instance.status/running, :instance/preempted? true}
        ;; all over the place.
-       (let [job-eid (:db/id (:job/_instance task-ent))
-             task-eid (:db/id task-ent)]
-         [[:generic/ensure task-eid :instance/status (d/entid db :instance.status/running)]
-          [:generic/atomic-inc job-eid :job/preemptions 1]
+      (let [job-eid (:db/id (:job/_instance task-ent))
+            task-eid (:db/id task-ent)]
+        [[:generic/ensure-some task-eid :instance/status #{(d/entid db :instance.status/running) (d/entid db :instance.status/unknown)}]
+         [:generic/atomic-inc job-eid :job/preemptions 1]
           ;; The database can become inconsistent if we make multiple calls to :instance/update-state in a single
           ;; transaction; see the comment in the definition of :instance/update-state for more details
-          [:instance/update-state task-eid :instance.status/failed [:reason/name :preempted-by-rebalancer]]
-          [:db/add task-eid :instance/reason [:reason/name :preempted-by-rebalancer]]
-          [:db/add task-eid :instance/preempted? true]]))
+         [:instance/update-state task-eid :instance.status/failed [:reason/name :preempted-by-rebalancer]]
+         [:db/add task-eid :instance/reason [:reason/name :preempted-by-rebalancer]]
+         [:db/add task-eid :instance/preempted? true]]))
     (catch Throwable e
       (log/warn
-        e "In" pool-name "pool, failed to transact preemption"
-        {:to-preempt (prep-task-ent-for-printing task-ent)}))))
+       e "In" pool-name "pool, failed to transact preemption"
+       {:to-preempt (prep-task-ent-for-printing task-ent)}))))
 
 (defn rebalance!
   [db conn agent-attributes-cache rebalancer-reservation-atom params init-state jobs-to-make-room-for pool-name]
diff --git a/scheduler/src/cook/schema.clj b/scheduler/src/cook/schema.clj
index e9f4b64ee0..cb49faa1d5 100644
--- a/scheduler/src/cook/schema.clj
+++ b/scheduler/src/cook/schema.clj
@@ -1135,6 +1135,22 @@ for a job. E.g. {:resources {:cpus 4 :mem 3} :constraints {\"unique_host_constra
                                                                    :attribute a
                                                                    :expected v}))))}}
 
+   {:db/id (d/tempid :db.part/user)
+    :db/ident :generic/ensure-some
+    :db/doc "Ensures an attribute of an entity has at least one of the expected values. Throws exception otherwise"
+    :db/fn #db/fn {:lang "clojure"
+                   :params [db e a values]
+                   :requires [[metatransaction.core :as mt]]
+                   :code
+                   (let [db (mt/filter-committed db)]
+                     (if (not-every? nil? (flatten
+                                           (for [v values]
+                                             (seq (d/datoms db :eavt e a v)))))
+                       nil
+                       (throw (ex-info "Fail to ensure attribute" {:entity e
+                                                                   :attribute a
+                                                                   :expected values}))))}}
+
    {:db/id (d/tempid :db.part/user)
     :db/ident :job/reasons->attempts-consumed
     :db/doc "Determines the amount of attempts consumed by a collection of failure reasons."
@@ -1384,14 +1400,14 @@ for a job. E.g. {:resources {:cpus 4 :mem 3} :constraints {\"unique_host_constra
                         oldvalues (into #{} (map :v) (datomic.api/datoms db :eavt eid aid))]
                     (-> []
                         (into (comp
-                                (remove newvalues)
-                                (map (if is-component
-                                       #(do [:db.fn/retractEntity %])
-                                       #(do [:db/retract eid aid %]))))
+                               (remove newvalues)
+                               (map (if is-component
+                                      #(do [:db.fn/retractEntity %])
+                                      #(do [:db/retract eid aid %]))))
                               oldvalues)
                         (into (comp
-                                (remove oldvalues)
-                                (map #(do [:db/add eid aid %])))
+                               (remove oldvalues)
+                               (map #(do [:db/add eid aid %])))
                               newvalues)))}}])
 
 (def reason-entities


Test patch:
diff --git a/integration/tests/cook/test_basic.py b/integration/tests/cook/test_basic.py
index 47d6d380b0..1e539330bf 100644
--- a/integration/tests/cook/test_basic.py
+++ b/integration/tests/cook/test_basic.py
@@ -662,6 +662,7 @@ def progress_string(a):
         self.assertEqual(80, instance['progress'], message)
         self.assertEqual('80%', instance['progress_message'], message)
 
+    @unittest.skipUnless(util.is_job_progress_supported(), 'Test depends on progress reporting')
     def test_progress_update_rest(self):
         job_uuid, resp = util.submit_job(self.cook_url)
         self.assertEqual(201, resp.status_code, msg=resp.content)
diff --git a/integration/tests/cook/test_multi_user.py b/integration/tests/cook/test_multi_user.py
index 6ce0072fc4..0b03d66f95 100644
--- a/integration/tests/cook/test_multi_user.py
+++ b/integration/tests/cook/test_multi_user.py
@@ -9,8 +9,7 @@
 
 import pytest
 from retrying import retry
-
-from tests.cook import mesos, util, reasons
+from tests.cook import mesos, reasons, util
 
 
 @pytest.mark.multi_user
@@ -425,15 +424,23 @@ def job_was_preempted(job):
                     self.logger.debug(f'Checking if instance was preempted: {instance}')
                     # Rebalancing marks the instance failed eagerly, so also wait for end_time to ensure it was
                     # actually killed
-                    if instance.get('reason_string') == 'Preempted by rebalancer' and instance.get(
-                            'end_time') is not None:
-                        return True
+                    if instance.get('end_time') is not None:
+                        if instance.get('reason_string') == 'Preempted by rebalancer' and instance.get('preempted'):
+                            return True
+                        else:
+                            # If the instance has an end_time but was not preempted, it means the preemption
+                            # failed. In this case, the job is killed but not correctly marked as preempted.
+                            # We set fail_on_error to true in the wait_until utility method. This means the 
+                            # resulting AssertionError will properly fail the test and not wait to fail at
+                            # the end of the timeout.
+                            self.fail(f'Instance failed but was not marked as preempted: {instance}')
+
                 self.logger.info(f'Job has not been preempted: {job}')
                 return False
 
             max_wait_ms = util.rebalancer_interval_seconds() * 1000 * 2.5
             self.logger.info(f'Waiting up to {max_wait_ms} milliseconds for preemption to happen')
-            util.wait_until(low_priority_job, job_was_preempted, max_wait_ms=max_wait_ms, wait_interval_ms=5000)
+            util.wait_until(low_priority_job, job_was_preempted, max_wait_ms=max_wait_ms, wait_interval_ms=5000, fail_on_error=True)
         finally:
             with self.user_factory.admin():
                 util.kill_jobs(self.cook_url, all_job_uuids, assert_response=False)
diff --git a/integration/tests/cook/util.py b/integration/tests/cook/util.py
index 5cdabc8cee..ecca3233d2 100644
--- a/integration/tests/cook/util.py
+++ b/integration/tests/cook/util.py
@@ -18,7 +18,6 @@
 import numpy
 import requests
 from retrying import retry
-
 from tests.cook import mesos
 
 logger = logging.getLogger(__name__)
@@ -811,16 +810,26 @@ def load_instance(cook_url, instance_uuid, assert_response=True):
     return load_resource(cook_url, 'instances', instance_uuid, assert_response)
 
 
-def wait_until(query, predicate, max_wait_ms=DEFAULT_TIMEOUT_MS, wait_interval_ms=DEFAULT_WAIT_INTERVAL_MS):
+def wait_until(query, predicate, max_wait_ms=DEFAULT_TIMEOUT_MS, wait_interval_ms=DEFAULT_WAIT_INTERVAL_MS, fail_on_error=False):
     """
     Block until the predicate is true for the result of the provided query.
     `query` is a thunk (nullary callable) that may be called multiple times.
     `predicate` is a unary callable that takes the result value of `query`
     and returns True if the condition is met, or False otherwise.
+
+    If `fail_on_error` is true and a non-RuntimeError is thrown, the method 
+    will stop retrying.
+
     See `wait_for_job` for an example of using this method.
     """
 
-    @retry(stop_max_delay=max_wait_ms, wait_fixed=wait_interval_ms)
+    def _retry_if_not_fail_on_error(exception):
+        if fail_on_error:
+            return isinstance(exception, RuntimeError)
+        else: 
+            return True
+
+    @retry(stop_max_delay=max_wait_ms, wait_fixed=wait_interval_ms, retry_on_exception=_retry_if_not_fail_on_error)
     def wait_until_inner():
         response = query()
         if not predicate(response):
@@ -839,7 +848,7 @@ def wait_until_inner():
             details = final_response.content
         except AttributeError:
             details = str(final_response)
-        logger.info(f"Timeout exceeded waiting for condition. Details: {details}")
+        logger.info(f"Timeout exceeded or error waiting for condition. Details: {details}")
         raise
 
 
diff --git a/scheduler/test/cook/test/rebalancer.clj b/scheduler/test/cook/test/rebalancer.clj
index df11c1f4cc..9430f38d64 100644
--- a/scheduler/test/cook/test/rebalancer.clj
+++ b/scheduler/test/cook/test/rebalancer.clj
@@ -20,14 +20,16 @@
             [cook.cache :as ccache]
             [cook.queries :as queries]
             [cook.quota :as quota]
-            [cook.test.postgres]
             [cook.rebalancer :as rebalancer :refer [->State]]
             [cook.scheduler.dru :as dru]
             [cook.scheduler.scheduler :as sched]
             [cook.scheduler.share :as share]
-            [cook.test.testutil :refer [create-dummy-group create-dummy-instance create-dummy-job init-agent-attributes-cache restore-fresh-database! setup]]
+            [cook.test.postgres]
+            [cook.test.testutil :refer [create-dummy-group
+                                        create-dummy-instance create-dummy-job
+                                        init-agent-attributes-cache restore-fresh-database! setup]]
             [cook.tools :as util]
-            [datomic.api :as d :refer [q]]))
+            [datomic.api :as d]))
 
 (use-fixtures :once cook.test.postgres/with-pg-db)
 
@@ -985,6 +987,32 @@
             (doall (for [[scored expected] scored->expected]
               (is (= scored expected))))))))))
 
+(deftest test-transact-preemption
+  (setup)
+  (testing "testing conditional to preempt only running and unknown instances"
+    (let [datomic-uri "datomic:mem://test-init-state"
+          conn (restore-fresh-database! datomic-uri)
+          db (d/db conn)
+          job1 (create-dummy-job conn :user "alexh" :memory 10.0 :ncpus 10.0)
+          job2 (create-dummy-job conn :user "alexh" :memory 10.0 :ncpus 10.0)
+          job3 (create-dummy-job conn :user "alexh" :memory 10.0 :ncpus 10.0)
+
+          task1 (create-dummy-instance conn job1 :instance-status :instance.status/running)
+          task2 (create-dummy-instance conn job2 :instance-status :instance.status/unknown)
+          task3 (create-dummy-instance conn job3 :instance-status :instance.status/success)
+
+          task-ent1 (d/entity (d/db conn) task1)
+          task-ent2 (d/entity (d/db conn) task2)
+          task-ent3 (d/entity (d/db conn) task3)]
+
+      ;; Transaction exceptions are wrapped and converted to a log, so test for nil return.
+      (is (not (nil? (rebalancer/transact-preemption! db conn "pool1" task-ent1))))
+      (is (not (nil? (rebalancer/transact-preemption! db conn "pool1" task-ent2))))
+
+      (is (nil? (rebalancer/transact-preemption! db conn "pool1" task-ent1)))
+      (is (nil? (rebalancer/transact-preemption! db conn "pool1" task-ent2)))
+      (is (nil? (rebalancer/transact-preemption! db conn "pool1" task-ent3))))))
+
 (defn rebalance
   "Calculates the jobs to make for and the initial state, and then delegates to rebalancer/rebalance"
   [db agent-attributes-cache pending-job-ents host->spare-resources rebalancer-reservation-atom