Repository: dask/distributed

Issue:
Active Memory Manager ignores nbytes thresholds
The `acquire-replicas` event does not include a `nbytes` dict, like `compute-task` does.
This means that all tasks fetched by AMM replicate have their size set to 1kb, which in turn means that `comm_threshold_bytes` and `target_message_size` are ignored and you'll potentially have several GiBs worth of tasks in a single gather_dep command.

This in turn can kill off the receiving worker, as tasks won't have a chance to be spilled to disk. 
To be more precise, the `target` and `pause` thresholds will be completely bypassed until after all tasks have landed into RAM; the `spill` threshold *may* kick in and save the day but it's going to be a race condition between network transfer speed and disk write speed.


Discussion:

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -7091,14 +7091,19 @@ def request_acquire_replicas(
         other workers. This is a fire-and-forget operation which offers no feedback for
         success or failure, and is intended for housekeeping and not for computation.
         """
-        who_has = {key: [ws.address for ws in self.tasks[key].who_has] for key in keys}
-        if self.validate:
-            assert all(who_has.values())
+        who_has = {}
+        nbytes = {}
+        for key in keys:
+            ts = self.tasks[key]
+            assert ts.who_has
+            who_has[key] = [ws.address for ws in ts.who_has]
+            nbytes[key] = ts.nbytes
 
         self.stream_comms[addr].send(
             {
                 "op": "acquire-replicas",
                 "who_has": who_has,
+                "nbytes": nbytes,
                 "stimulus_id": stimulus_id,
             },
         )
diff --git a/distributed/worker_state_machine.py b/distributed/worker_state_machine.py
--- a/distributed/worker_state_machine.py
+++ b/distributed/worker_state_machine.py
@@ -781,8 +781,9 @@ class RefreshWhoHasEvent(StateMachineEvent):
 
 @dataclass
 class AcquireReplicasEvent(StateMachineEvent):
-    __slots__ = ("who_has",)
+    __slots__ = ("who_has", "nbytes")
     who_has: dict[str, Collection[str]]
+    nbytes: dict[str, int]
 
 
 @dataclass
@@ -2395,10 +2396,11 @@ def _handle_remove_replicas(self, ev: RemoveReplicasEvent) -> RecsInstrs:
     @_handle_event.register
     def _handle_acquire_replicas(self, ev: AcquireReplicasEvent) -> RecsInstrs:
         if self.validate:
+            assert ev.who_has.keys() == ev.nbytes.keys()
             assert all(ev.who_has.values())
 
         recommendations: Recs = {}
-        for key in ev.who_has:
+        for key, nbytes in ev.nbytes.items():
             ts = self._ensure_task_exists(
                 key=key,
                 # Transfer this data after all dependency tasks of computations with
@@ -2409,6 +2411,7 @@ def _handle_acquire_replicas(self, ev: AcquireReplicasEvent) -> RecsInstrs:
                 stimulus_id=ev.stimulus_id,
             )
             if ts.state != "memory":
+                ts.nbytes = nbytes
                 recommendations[ts] = "fetch"
 
         self._update_who_has(ev.who_has)


Test patch:
diff --git a/distributed/tests/test_worker.py b/distributed/tests/test_worker.py
--- a/distributed/tests/test_worker.py
+++ b/distributed/tests/test_worker.py
@@ -2763,7 +2763,9 @@ async def test_acquire_replicas_already_in_flight(c, s, a):
         assert b.tasks["x"].state == "flight"
 
         b.handle_stimulus(
-            AcquireReplicasEvent(who_has={"x": a.address}, stimulus_id="test")
+            AcquireReplicasEvent(
+                who_has={"x": a.address}, nbytes={"x": 1}, stimulus_id="test"
+            )
         )
         assert b.tasks["x"].state == "flight"
         b.block_gather_dep.set()
@@ -2963,6 +2965,32 @@ async def test_acquire_replicas_with_no_priority(c, s, a, b):
     assert b.tasks["x"].priority is not None
 
 
+@gen_cluster(client=True, nthreads=[("", 1)])
+async def test_acquire_replicas_large_data(c, s, a):
+    """When acquire-replicas is used to acquire multiple sizeable tasks, it respects
+    target_message_size and acquires them over multiple iterations.
+    """
+    size = a.target_message_size // 5 - 10_000
+
+    class C:
+        def __sizeof__(self):
+            return size
+
+    futs = [c.submit(C, key=f"x{i}", workers=[a.address]) for i in range(10)]
+    await wait(futs)
+
+    async with BlockedGatherDep(s.address) as b:
+        s.request_acquire_replicas(
+            b.address, [f"x{i}" for i in range(10)], stimulus_id="test"
+        )
+        await b.in_gather_dep.wait()
+        assert len(b.state.in_flight_tasks) == 5
+        assert len(b.state.data_needed) == 5
+        b.block_gather_dep.set()
+        while len(b.data) < 10:
+            await asyncio.sleep(0.01)
+
+
 @gen_cluster(client=True)
 async def test_missing_released_zombie_tasks(c, s, a, b):
     """
diff --git a/distributed/tests/test_worker_state_machine.py b/distributed/tests/test_worker_state_machine.py
--- a/distributed/tests/test_worker_state_machine.py
+++ b/distributed/tests/test_worker_state_machine.py
@@ -7,6 +7,8 @@
 import pytest
 from tlz import first
 
+from dask.sizeof import sizeof
+
 import distributed.profile as profile
 from distributed import Nanny, Worker, wait
 from distributed.protocol.serialize import Serialize
@@ -93,7 +95,9 @@ def test_WorkerState__to_dict():
     ws = WorkerState(8)
     ws.address = "127.0.0.1.1234"
     ws.handle_stimulus(
-        AcquireReplicasEvent(who_has={"x": ["127.0.0.1:1235"]}, stimulus_id="s1")
+        AcquireReplicasEvent(
+            who_has={"x": ["127.0.0.1:1235"]}, nbytes={"x": 123}, stimulus_id="s1"
+        )
     )
     ws.handle_stimulus(
         UpdateDataEvent(data={"y": object()}, report=False, stimulus_id="s2")
@@ -131,6 +135,7 @@ def test_WorkerState__to_dict():
                 "cls": "AcquireReplicasEvent",
                 "stimulus_id": "s1",
                 "who_has": {"x": ["127.0.0.1:1235"]},
+                "nbytes": {"x": 123},
             },
             {
                 "cls": "UpdateDataEvent",
@@ -142,13 +147,14 @@ def test_WorkerState__to_dict():
         "tasks": {
             "x": {
                 "key": "x",
+                "nbytes": 123,
                 "priority": [1],
                 "state": "fetch",
                 "who_has": ["127.0.0.1:1235"],
             },
             "y": {
                 "key": "y",
-                "nbytes": 16,
+                "nbytes": sizeof(object()),
                 "state": "memory",
             },
         },
@@ -664,7 +670,9 @@ async def test_missing_handle_compute_dependency(c, s, w1, w2, w3):
     await wait_for_state(f1.key, "memory", w1)
 
     w3.handle_stimulus(
-        AcquireReplicasEvent(who_has={f1.key: [w2.address]}, stimulus_id="acquire")
+        AcquireReplicasEvent(
+            who_has={f1.key: [w2.address]}, nbytes={f1.key: 1}, stimulus_id="acquire"
+        )
     )
     await wait_for_state(f1.key, "missing", w3)
 
@@ -681,7 +689,9 @@ async def test_missing_to_waiting(c, s, w1, w2, w3):
     await wait_for_state(f1.key, "memory", w1)
 
     w3.handle_stimulus(
-        AcquireReplicasEvent(who_has={f1.key: [w2.address]}, stimulus_id="acquire")
+        AcquireReplicasEvent(
+            who_has={f1.key: [w2.address]}, nbytes={f1.key: 1}, stimulus_id="acquire"
+        )
     )
     await wait_for_state(f1.key, "missing", w3)