No longer double count transfer cost in stealing (dask#7036)

Co-authored-by: fjetter <fjetter@users.noreply.github.com>
gjoseph92 · Oct 31, 2022 · 07f8265 · 07f8265
1 parent dc8d8a7
commit 07f8265
Show file tree

Hide file tree

Showing 6 changed files with 251 additions and 220 deletions.
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
@@ -2892,7 +2892,7 @@ def check_idle_saturated(self, ws: WorkerState, occ: float = -1.0):
         idle = self.idle
         saturated = self.saturated
         if (
-            (p < nc or occ < nc * avg / 2)
+            self.is_unoccupied(ws, occ, p)
             if math.isinf(self.WORKER_SATURATION)
             else not _worker_full(ws, self.WORKER_SATURATION)
         ):
@@ -2910,6 +2910,13 @@ def check_idle_saturated(self, ws: WorkerState, occ: float = -1.0):
 
             saturated.discard(ws)
 
+    def is_unoccupied(
+        self, ws: WorkerState, occupancy: float, nprocessing: int
+    ) -> bool:
+        nthreads = ws.nthreads
+        avg_occ_per_thread = self.total_occupancy / self.total_nthreads
+        return nprocessing < nthreads or occupancy < nthreads * avg_occ_per_thread / 2
+
     def get_comm_cost(self, ts: TaskState, ws: WorkerState) -> float:
         """
         Get the estimated communication cost (in s.) to compute the task