Repository: dask/distributed

Issue:
`ReduceReplicas.run` requires non-trivial amount of time on the scheduler
While profiling a large scale computation (1.6MM tasks) I noticed how `ReduceReplicas.run` requires a non-trivial amount of time.

In the profile I am looking at it it runs for about 200ms every 2s, i.e. it is taking about 10% of the entire CPU time on the scheduler.

![image](https://github.com/dask/distributed/assets/8629629/416dd84c-09e1-4344-8ebf-0944749fe20b)

Almost all of this time is spent in a single set comprehension

![image](https://github.com/dask/distributed/assets/8629629/c58ee994-db8e-4ea1-a81b-c76bf6f0a0a2)

https://github.com/dask/distributed/blob/ec9aa8547070f3441f6dbe0b7dbc2a25ca70b601/distributed/active_memory_manager.py#L539

which is counting the number of waiters for every task that is replicated. For large graphs, this can be quite expensive and for a policy that is scheduled to run every 2s we should try to reduce the footprint of this call.

cc @crusaderky 



Discussion:

Resolving patch:
diff --git a/distributed/active_memory_manager.py b/distributed/active_memory_manager.py
--- a/distributed/active_memory_manager.py
+++ b/distributed/active_memory_manager.py
@@ -533,11 +533,17 @@ def run(self) -> SuggestionGenerator:
         for ts in self.manager.scheduler.replicated_tasks:
             desired_replicas = 1  # TODO have a marker on TaskState
 
-            # If a dependent task has not been assigned to a worker yet, err on the side
-            # of caution and preserve an additional replica for it.
-            # However, if two dependent tasks have been already assigned to the same
-            # worker, don't double count them.
-            nwaiters = len({waiter.processing_on or waiter for waiter in ts.waiters})
+            nwaiters = len(ts.waiters)
+            if desired_replicas < nwaiters < 20:
+                # If a dependent task has not been assigned to a worker yet, err on the
+                # side of caution and preserve an additional replica for it.
+                # However, if two dependent tasks have been already assigned to the same
+                # worker, don't double count them.
+                # This calculation is quite CPU-intensive, so it's disabled for tasks
+                # with lots of waiters.
+                nwaiters = len(
+                    {waiter.processing_on or waiter for waiter in ts.waiters}
+                )
 
             ndrop_key = len(ts.who_has) - max(desired_replicas, nwaiters)
             if ts in self.manager.pending:


Test patch:
diff --git a/distributed/tests/test_active_memory_manager.py b/distributed/tests/test_active_memory_manager.py
--- a/distributed/tests/test_active_memory_manager.py
+++ b/distributed/tests/test_active_memory_manager.py
@@ -23,6 +23,7 @@
     NO_AMM,
     BlockedGatherDep,
     assert_story,
+    async_poll_for,
     captured_logger,
     gen_cluster,
     gen_test,
@@ -739,9 +740,72 @@ async def test_ReduceReplicas(c, s, *workers):
         ],
     ):
         s.extensions["amm"].run_once()
+    await async_poll_for(lambda: len(s.tasks["x"].who_has) == 1, timeout=5)
+
+
+@pytest.mark.parametrize(
+    "nwaiters_w1,nwaiters_w2,nwaiters_nonproc,nreplicas",
+    [
+        (0, 0, 0, 1),
+        (1, 0, 0, 1),
+        (0, 0, 1, 1),
+        (2, 0, 0, 1),
+        (1, 1, 0, 2),
+        (1, 1, 1, 3),
+        (1, 1, 2, 4),
+        (2, 1, 1, 3),
+        (1, 1, 17, 4),
+        (17, 1, 1, 3),
+        # Fast code path: if there are 20+ waiters, don't check processing_on
+        (18, 1, 1, 4),
+    ],
+)
+@gen_cluster(
+    nthreads=[("", 1)] * 4,
+    client=True,
+    config={
+        "distributed.scheduler.active-memory-manager.start": False,
+        "distributed.scheduler.active-memory-manager.policies": [
+            {"class": "distributed.active_memory_manager.ReduceReplicas"},
+        ],
+    },
+)
+async def test_ReduceReplicas_with_waiters(
+    c, s, w1, w2, w3, w4, nwaiters_w1, nwaiters_w2, nwaiters_nonproc, nreplicas
+):
+    """If there are waiters, even if they are not processing on a worker yet, preserve
+    extra replicas.
+    If there are between 2 and 19 waiters, consider which workers they've been assigned
+    to and don't double count waiters that have been assigned to the same worker.
+    """
+    ev = Event()
+    x = (await c.scatter({"x": 123}, broadcast=True))["x"]
+    y = c.submit(lambda ev: ev.wait(), ev, key="y", workers=[w4.address])
+    waiters_w1 = [
+        c.submit(lambda x, ev: ev.wait(), x, ev, key=("zw1", i), workers=[w1.address])
+        for i in range(nwaiters_w1)
+    ]
+    waiters_w2 = [
+        c.submit(lambda x, ev: ev.wait(), x, ev, key=("zw2", i), workers=[w2.address])
+        for i in range(nwaiters_w2)
+    ]
+    waiters_nonproc = [
+        c.submit(lambda x, y: None, x, y, key=("znp", i))
+        for i in range(nwaiters_nonproc)
+    ]
+    nwaiters = nwaiters_w1 + nwaiters_w2 + nwaiters_nonproc
+    await async_poll_for(lambda: len(s.tasks) == nwaiters + 2, timeout=5)
+    for fut in waiters_w1:
+        assert s.tasks[fut.key].processing_on == s.workers[w1.address]
+    for fut in waiters_w2:
+        assert s.tasks[fut.key].processing_on == s.workers[w2.address]
+    for fut in waiters_nonproc:
+        assert s.tasks[fut.key].processing_on is None
 
-    while len(s.tasks["x"].who_has) > 1:
-        await asyncio.sleep(0.01)
+    s.extensions["amm"].run_once()
+    await asyncio.sleep(0.2)  # Test no excessive drops
+    await async_poll_for(lambda: len(s.tasks["x"].who_has) == nreplicas, timeout=5)
+    await ev.set()
 
 
 @pytest.mark.parametrize("start_amm", [False, True])