Comments and test

Author: hendrikmakait

distributed/scheduler.py

@@ -3283,6 +3283,10 @@ def worker_objective(self, ts: TaskState, ws: WorkerState) -> tuple:
         """Objective function to determine which worker should get the task
 
         Minimize expected start time.  If a tie then break with data storage.
+
+        See Also
+        --------
+        WorkStealing.stealing_objective
         """
         stack_time = ws.occupancy / ws.nthreads
         start_time = stack_time + self.get_comm_cost(ts, ws)

distributed/stealing.py

@@ -529,11 +529,24 @@ def story(self, *keys_or_ts: str | TaskState) -> list:
         return out
 
     def stealing_objective(
-        self, scheduler: SchedulerState, ts: TaskState, ws: WorkerState
+        self, ts: TaskState, ws: WorkerState
     ) -> tuple[float, ...]:
+        """Objective function to determine which worker should get the task
+
+        Minimize expected start time.  If a tie then break with data storage.
+
+        Notes
+        -----
+        This method is a modified version of Scheduler.worker_objective that accounts
+        for in-flight requests. It must be kept in sync for work-stealing to work correctly.
+
+        See Also
+        --------
+        Scheduler.worker_objective
+        """
         occupancy = self._combined_occupancy(
             ws
-        ) / ws.nthreads + scheduler.get_comm_cost(ts, ws)
+        ) / ws.nthreads + self.scheduler.get_comm_cost(ts, ws)
         if ts.actor:
             return (len(ws.actors), occupancy, ws.nbytes)
         else:
@@ -553,7 +566,7 @@ def _get_thief(
             elif not ts.loose_restrictions:
                 return None
         return min(
-            potential_thieves, key=partial(self.stealing_objective, scheduler, ts)
+            potential_thieves, key=partial(self.stealing_objective, ts)
         )
 
 

distributed/tests/test_steal.py

@@ -38,6 +38,7 @@
 from distributed.utils_test import (
     NO_AMM,
     BlockedGetData,
+    async_poll_for,
     captured_logger,
     freeze_batched_send,
     gen_cluster,
@@ -1877,3 +1878,54 @@ async def test_trivial_workload_should_not_cause_work_stealing(c, s, *workers):
     await c.gather(futs)
     events = s.get_events("stealing")
     assert len(events) == 0
+
+
+@gen_cluster(
+    nthreads=[("", 1)],
+    client=True,
+    config={"distributed.scheduler.worker-saturation": "inf"},
+)
+async def test_stealing_ogjective_accounts_for_in_flight(c, s, a):
+    """Regression test that work-stealing's objective correctly accounts for in-flight data requests
+    """
+    in_event = Event()
+    block_event = Event()
+
+    def block(i: int, in_event: Event, block_event: Event) -> int:
+        in_event.set()
+        block_event.wait()
+        return i
+
+    # Stop stealing for deterministic testing
+    extension = s.extensions["stealing"]
+    await extension.stop()
+
+    try:
+        futs = c.map(block, range(20), in_event=in_event, block_event=block_event)
+        await in_event.wait()
+
+        async with Worker(s.address, nthreads=1) as b:
+            try:
+                await async_poll_for(lambda: s.idle, timeout=5)
+                wsA = s.workers[a.address]
+                wsB = s.workers[b.address]
+                ts = next(iter(wsA.processing))
+
+                # No in-flight requests, so both match
+                assert extension.stealing_objective(ts, wsA) == s.worker_objective(ts, wsA)
+                assert extension.stealing_objective(ts, wsB) == s.worker_objective(ts, wsB)
+
+                extension.balance()
+                assert extension.in_flight
+                # We move tasks from a to b
+                assert extension.stealing_objective(ts, wsA) < s.worker_objective(ts, wsA)
+                assert extension.stealing_objective(ts, wsB) > s.worker_objective(ts, wsB)
+
+                await async_poll_for(lambda: not extension.in_flight, timeout=5)
+                # No in-flight requests, so both match
+                assert extension.stealing_objective(ts, wsA) == s.worker_objective(ts, wsA)
+                assert extension.stealing_objective(ts, wsB) == s.worker_objective(ts, wsB)
+            finally:
+                await block_event.set()
+    finally:
+        await block_event.set()