Repository: dask/distributed

Issue:
`Client.scatter` does not evenly distribute data when scattering only a few tasks
When using `Client.scatter` to distribute few but very large tasks, the load may concentrate on a single machine causing OOM errors.

This problem is caused by this logic in `scatter_to_workers`:
https://github.com/dask/distributed/blob/fa4976e3048ff37af04cb92389036b6ac1867bd9/distributed/utils_comm.py#L166-L170

There, the client scatters data in a round-robin fashion, but before moving to the next worker, it sends as many tasks to the current one as there are threads on the worker.

Instead of optimizing for heterogeneous clusters and causing inhomogeneous load while doing so, we should just apply a standard round-robin, moving to a new worker after each iteration. 


Discussion:

Resolving patch:
diff --git a/distributed/client.py b/distributed/client.py
--- a/distributed/client.py
+++ b/distributed/client.py
@@ -2450,10 +2450,9 @@ async def _scatter(
                     nthreads = await self.scheduler.ncores_running(workers=workers)
                 if not nthreads:  # pragma: no cover
                     raise ValueError("No valid workers found")
+                workers = list(nthreads.keys())
 
-                _, who_has, nbytes = await scatter_to_workers(
-                    nthreads, data2, rpc=self.rpc
-                )
+                _, who_has, nbytes = await scatter_to_workers(workers, data2, self.rpc)
 
                 await self.scheduler.update_data(
                     who_has=who_has, nbytes=nbytes, client=self.id
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -6132,16 +6132,15 @@ async def scatter(
                 raise TimeoutError("No valid workers found")
             await asyncio.sleep(0.1)
 
-        nthreads = {ws.address: ws.nthreads for ws in wss}
-
         assert isinstance(data, dict)
 
-        keys, who_has, nbytes = await scatter_to_workers(nthreads, data, rpc=self.rpc)
+        workers = list(ws.address for ws in wss)
+        keys, who_has, nbytes = await scatter_to_workers(workers, data, rpc=self.rpc)
 
         self.update_data(who_has=who_has, nbytes=nbytes, client=client)
 
         if broadcast:
-            n = len(nthreads) if broadcast is True else broadcast
+            n = len(workers) if broadcast is True else broadcast
             await self.replicate(keys=keys, workers=workers, n=n)
 
         self.log_event(
diff --git a/distributed/utils_comm.py b/distributed/utils_comm.py
--- a/distributed/utils_comm.py
+++ b/distributed/utils_comm.py
@@ -9,7 +9,7 @@
 from itertools import cycle
 from typing import Any, TypeVar
 
-from tlz import concat, drop, groupby, merge
+from tlz import drop, groupby, merge
 
 import dask.config
 from dask.optimization import SubgraphCallable
@@ -151,19 +151,16 @@ def __repr__(self):
 _round_robin_counter = [0]
 
 
-async def scatter_to_workers(nthreads, data, rpc=rpc):
+async def scatter_to_workers(workers, data, rpc=rpc):
     """Scatter data directly to workers
 
-    This distributes data in a round-robin fashion to a set of workers based on
-    how many cores they have.  nthreads should be a dictionary mapping worker
-    identities to numbers of cores.
+    This distributes data in a round-robin fashion to a set of workers.
 
     See scatter for parameter docstring
     """
-    assert isinstance(nthreads, dict)
     assert isinstance(data, dict)
 
-    workers = list(concat([w] * nc for w, nc in nthreads.items()))
+    workers = sorted(workers)
     names, data = list(zip(*data.items()))
 
     worker_iter = drop(_round_robin_counter[0] % len(workers), cycle(workers))


Test patch: