[Heuristics] Add steal rejection tracking metric and unit test

Author: prince8273

distributed/stealing.py

@@ -118,6 +118,7 @@ def __init__(self, scheduler: Scheduler):
         self.metrics = {
             "request_count_total": defaultdict(int),
             "request_cost_total": defaultdict(int),
+            "reject_count_margin_total": defaultdict(int),
         }
         self._request_counter = 0
         self.scheduler.stream_handlers["steal-response"] = self.move_task_confirm
@@ -489,10 +490,17 @@ def balance(self) -> None:
 
                     # Require at least 50% ROI on the network transfer cost to prevent thrashing
                     margin = comm_cost_thief * 0.5
-                    if (
+                    
+                    would_steal_without_margin = (
+                        occ_thief + comm_cost_thief + compute
+                        <= occ_victim - (comm_cost_victim + compute) / 2
+                    )
+                    would_steal_with_margin = (
                         occ_thief + comm_cost_thief + compute + margin
                         <= occ_victim - (comm_cost_victim + compute) / 2
-                    ):
+                    )
+
+                    if would_steal_with_margin:
                         self.move_task_request(ts, victim, thief)
                         cost = compute + comm_cost_victim
                         log.append(
@@ -523,6 +531,13 @@ def balance(self) -> None:
                         # for removing ts from stealable. If we made sure to
                         # properly clean up, we would not need this
                         stealable.discard(ts)
+                    elif would_steal_without_margin:
+                        self.metrics["reject_count_margin_total"][level] += 1
+                        logger.debug(
+                            "Work-stealing margin heuristic rejected steal of task %s "
+                            "(thief=%s, victim=%s, level=%d, margin=%.4f)",
+                            ts.key, thief.address, victim.address, level, margin,
+                        )
                 self.scheduler.check_idle_saturated(
                     victim, occ=combined_occupancy(victim)
                 )

distributed/tests/test_steal.py

@@ -2010,6 +2010,36 @@ def block(i: int, in_event: Event, block_event: Event) -> int:
         await block_event.set()
 
 
+@gen_cluster(
+    client=True,
+    nthreads=[("127.0.0.1", 1)] * 2,
+    config={"distributed.scheduler.work-stealing-interval": "100ms", **NO_AMM},
+)
+async def test_reject_count_margin_metric(c, s, a, b):
+    """
+    Verify that the margin heuristic increments reject_count_margin_total
+    when a steal is suppressed that old logic would have permitted.
+    """
+    steal = s.extensions["stealing"]
+    await steal.stop()
+
+    # Generate large data on worker A to ensure high network transfer cost
+    [x] = await c.scatter([b"0" * 50_000_000], workers=a.address)
+    
+    # Create tasks on A to saturate it and trigger stealing evaluation
+    futures = [
+        c.submit(slowidentity, x, pure=False, delay=0.01, workers=a.address, allow_other_workers=True)
+        for _ in range(10)
+    ]
+    
+    while len(a.state.tasks) < 10:
+        await asyncio.sleep(0.01)
+
+    # Balance will evaluate the cost. High comm_cost, low compute.
+    # Without margin, it would steal. With 50% ROI margin, it should reject.
+    steal.balance()
+
+    assert sum(steal.metrics["reject_count_margin_total"].values()) >= 1
 @gen_cluster(
     nthreads=[("", 1)],
     client=True,