Change runtime.py to choose tasks with lowest (instead of highest) pr…

…iority (#505) Currently, the priority is set to the timestamp of the earliest undispatched task. Choosing earliest tasks will reduce the maximum waiting time when queue is nonempty Co-authored-by: Max Ryabinin <[email protected]> Co-authored-by: Pavel Samygin <[email protected]> (cherry picked from commit 6395e89)
learning-at-home · Sep 13, 2022 · 2ba000b · 2ba000b
1 parent a2f2407
commit 2ba000b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/hivemind/moe/server/runtime.py b/hivemind/moe/server/runtime.py
@@ -143,8 +143,8 @@ def iterate_minibatches_from_pools(self, timeout=None):
                 if self.SHUTDOWN_TRIGGER in ready_objects:
                     break  # someone asked us to shutdown, break from the loop
 
-                logger.debug("Choosing the pool with highest priority")
-                pool = max(ready_objects, key=lambda pool: pool.priority)
+                logger.debug("Choosing the pool with first priority")
+                pool = min(ready_objects, key=lambda pool: pool.priority)
 
                 logger.debug(f"Loading batch from {pool.name}")
                 batch_index, batch_tensors = pool.load_batch_to_runtime(timeout, self.device)

diff --git a/hivemind/moe/server/task_pool.py b/hivemind/moe/server/task_pool.py
@@ -27,7 +27,7 @@ class TaskPoolBase(mp.context.ForkProcess, metaclass=ABCMeta):
     def __init__(self, process_func: callable, daemon=True, **kwargs):
         super().__init__(daemon=daemon, **kwargs)
         self.process_func = process_func
-        self._priority = mp.Value(ctypes.c_double, 1.0)  # higher priority = the more urgent to process this pool
+        self._priority = mp.Value(ctypes.c_double, 1.0)  # lower priority = the more urgent to process this pool
 
     @abstractmethod
     def run(self):
@@ -170,7 +170,7 @@ def _pool_input_loop(self, pending_batches: Dict[Any, List[Task]], *args, **kwar
             for skip_i in range(prev_num_tasks):
                 finished_task_timestamp = (
                     self.undispatched_task_timestamps.get()
-                )  # earlier timestamp = higher priority
+                )  # earlier timestamp = smaller (better) priority, earlier processing
                 if skip_i == prev_num_tasks - 1:
                     self.priority = finished_task_timestamp