Improve work stealing for scaling situations

Author: fjetter

distributed/stealing.py

@@ -1,5 +1,7 @@
 import logging
+import random
 from collections import defaultdict, deque
+from functools import partial
 from math import log2
 from time import time
 
@@ -113,27 +115,21 @@ def steal_time_ratio(self, ts):
         For example a result of zero implies a task without dependencies.
         level: The location within a stealable list to place this value
         """
-        if not ts.dependencies:  # no dependencies fast path
-            return 0, 0
 
         split = ts.prefix.name
         if split in fast_tasks:
             return None, None
 
         ws = ts.processing_on
         compute_time = ws.processing[ts]
-        if compute_time < 0.005:  # 5ms, just give up
-            return None, None
 
         nbytes = ts.get_nbytes_deps()
         transfer_time = nbytes / self.scheduler.bandwidth + LATENCY
         cost_multiplier = transfer_time / compute_time
-        if cost_multiplier > 100:
-            return None, None
 
         level = int(round(log2(cost_multiplier) + 6))
-        if level < 1:
-            level = 1
+
+        level = min(len(self.cost_multipliers) - 1, level)
 
         return cost_multiplier, level
 
@@ -344,7 +340,10 @@ def maybe_move_task(level, ts, sat, idl, duration, cost_multiplier):
                             thieves = idle
                         if not thieves:
                             break
-                        thief = thieves[i % len(thieves)]
+
+                        thief = self._maybe_pick_thief(ts, thieves)
+                        if not thief:
+                            thief = thieves[i % len(thieves)]
 
                         duration = sat.processing.get(ts)
                         if duration is None:
@@ -380,7 +379,10 @@ def maybe_move_task(level, ts, sat, idl, duration, cost_multiplier):
                             thieves = idle
                         if not thieves:
                             continue
-                        thief = thieves[i % len(thieves)]
+                        thief = self._maybe_pick_thief(ts, thieves)
+                        if not thief:
+                            thief = thieves[i % len(thieves)]
+
                         duration = sat.processing[ts]
 
                         maybe_move_task(
@@ -394,6 +396,32 @@ def maybe_move_task(level, ts, sat, idl, duration, cost_multiplier):
             if s.digests:
                 s.digests["steal-duration"].add(stop - start)
 
+    def _maybe_pick_thief(self, ts, thieves):
+        """Try to be smart about picking a thief given some options. We're
+        trying to pick a thief which has dependencies for a given task, if
+        possible and will pick the one which works best for us given the
+        Scheduler.worker_objective
+
+        If no idle worker with dependencies is found, this returns None.
+        """
+        if ts._dependencies:
+            who_has = set()
+            for dep in ts._dependencies:
+                who_has.update(dep.who_has)
+
+            thieves_with_data = who_has & set(thieves)
+
+            # If there are potential thieves with dependencies we
+            # should prefer them and pick the one which works best.
+            # Otherwise just random/round robin
+            if thieves_with_data:
+                if len(thieves_with_data) > 10:
+                    thieves_with_data = random.sample(thieves_with_data, 10)
+                return min(
+                    thieves_with_data,
+                    key=partial(self.scheduler.worker_objective, ts),
+                )
+
     def restart(self, scheduler):
         for stealable in self.stealable.values():
             for s in stealable: