Repository: dask/distributed

Issue:
Alternatives for current ensure_communicating
We've recently been refactoring the worker state machine (https://github.com/dask/distributed/issues/5736)

The new system works great for single-task-based transitions but feels awkward as soon as we're trying to treat groups or batches of tasks similarly which is particularly important for fetching remote data. We are trying to not perform a network request per key but rather try to perform a request per worker and try to get as many keys of it as reasonable.

This mechanism worked so far by running `Worker.(_)ensure_communicating` frequently which iterates over all tasks and schedules a fetch/gather_dep (pseudo code below)

```python
def ensure_communicating(...):
    skip = []
    while self.data_needed and free_slots_for_fetch:
        ts = self.data_needed.pop() # highest priority task
        worker = pick_good_worker(ts) # not busy, maybe localhost, etc.
        if not worker:
            skip.add(ts)
        more_keys_on_worker = select_keys_for_gather(worker)
        schedule_gather_dep(ts, more_keys_on_worker, worker) # This schedules a coroutine and actually fetches the data
    self.data_needed.update(skip) # data needed is a heap so this is non-linear
```

This has a couple of problems
1. It's unclear when to run this function. Past iterations just ran this frequently (on every incoming request and in various code places where we guess there may be something that needs transitioning
2. We're iterating over all keys, all the time. This is bad for performance although not critical on the worker https://github.com/dask/distributed/pull/6462#issuecomment-1138764541
3. If we're calling this too soon, e.g. within a released->fetch transition, we may get poor scheduling decisions by fetching tasks prematurely since we didn't work off all fetch recommendations before calling this. For this reason, the `EnsureCommunicatingAfterTransitions ` instruction was introduced (see also https://github.com/dask/distributed/pull/6462) to mimic earlier behavior


## Proposal

I think we can get rid of `Worker.data_needed` entirely. Since we're storing everything in heaps, by now, `Worker.data_needed_per_worker` carries everything we need if we change how ensure_communicating works which might have other interesting side effects.

To illustrate this, consider the following two helper functions


```python
def get_most_important_worker_to_fetch_from(self):
    """
    Get the remote worker that holds the _most important_ data we should fetch. This is typically inferred by looking at task priorities.
    This could also be used to modify what we consider most important worker to fetch from, e.g. due to load/business, memory pressure, etc.
    """

    # Linear in num_workers
    res = None
    prio = None
    for w, heap in self.data_needed_per_worker.items():
        nprio = heap.peek().priority
        if w not in self.busy_workers and (res is None or nprio > prio):
            res = w
            prio = nprio
    return res

def get_tasks_for_worker(self, w, limit=42):
    """Get a list of tasks on worker `w` such that combined they have at most `limit` estimated bytes"""
    # This is more or less the current `Worker._select_keys_for_gather`
    # Linear in number of tasks in local state fetch on remote worker
    res = []
    size = 0
    for ts in self.data_needed_per_worker[w]:
        size += ts.get_nbytes()
        res.append(ts)
        if size > limit:
            break
    return res
```

With these two helpers we should be able to construct a new, simplified version if we switch our loop from tasks to workers

```python
def _new_ensure_communicate_with_tasks_by_worker():
    if network_limits_not_reached and worker := self.get_most_important_worker_to_fetch_from():
        tasks_to_fetch = self.get_tasks_for_worker(worker)
        instruction = GatherDep(
            worker,
            tasks_to_fetch,
        )
```

or even change the `GatherDep` instruction to worker-only and delay the decision of what tasks to fetch until this instruction is acted on, which typically only happens after all transitions in `Worker.transitions` are worked off

```python
def _new_ensure_communicate_with_tasks_by_worker():
    if network_limits_not_reached and worker := self.get_most_important_worker_to_fetch_from():
        instruction = GatherDep(worker)

def gather_dep(self, worker):
    tasks_to_fetch = self.get_tasks_for_worker(worker)
    ...
```


cc @crusaderky 


Discussion:
While I'm warm to deduplicate data_needed vs. data_needed_per_worker, I don't see how any of this removes the problem about when to trigger `_new_ensure_communicate_with_tasks_by_worker`?

My suggestion:

Remove `data_needed_per_worker` entirely, and remain only with `data_needed`.
`in_flight_workers` becomes `{worker: GatherDep}` where the GatherDep instruction contains all tasks to be fetched from that worker.
Every iteration of `_ensure_communicating` fetches the topmost task from data_needed and adds it to the matching GatherDep for that worker. Then, it adds the GatherDep to instructions (possibly the same object multiple times).
`_ensure_communicating` stops iterating when any one of the following happens:

1. there are no tasks left in data_needed
2. the total nbytes in flight across all workers exceeds a global threshold
`self.comm_nbytes >= self.comm_threshold_bytes`
3. the number of *saturated* pending GatherDep commands exceeds a threshold
`sum(gd.nbytes >= self.target_message_size for gd in self.in_flight_workers.values()) >= self.total_out_connections`

Individual GatherDep instructions stop being enriched when their individual nbytes exceed an individual threshold
`self.in_flight_workers[worker].nbytes >= self.target_message_size`; when that happens, tasks are moved to `skipped_worker_in_flight_or_busy` (read below).

At the top of `_handle_instructions`, there's a simple object deduplication:
```python
instructions = list({id(instr): instr for instr in instructions}.values())
```

`_ensure_communicating` is still called as soon as possible, e.g. at every transition to fetch. However, as a performance tweak, 
`skipped_worker_in_flight_or_busy` becomes an instance variable - a set which must lose keys every time one would otherwise remove keys from `data_needed`. Instead of being flushed back into `data_needed` at the end of each iteration of `_ensure_communicating`, it is partially flushed
1. every time a worker exits flight or busy state
2. every time update_who_has adds a new worker to a task in fetch state

flushing `skipped_worker_in_flight_or_busy` back into `data_needed` is O(n) to `len(skipped_worker_in_flight_or_busy)` and O(n*logn) to the number of tasks being added back in, but the key is that it happens a lot less frequently than now.

This closes https://github.com/dask/distributed/pull/6462#issuecomment-1138764541 for both the case of multiple fetches from the same event as well as for the case of multiple fetches from multiple commands in short succession, and will cause calling `_ensure_communicating` twice in a row to be computationally negligible.

The big difference between my proposal above and both the current system and the proposal to remove `data_needed` in the op is that my proposal respects global priority. This means that you'll end up with more, individually smaller, parallel tranfers from multiple workers as you'll be more likely to hit `self.comm_threshold_bytes` than `self.target_message_size`.

e.g. 
each key is nbytes=1
there are 99999 keys to fetch, uniformly distributed across 3 workers a, b, and c, with the same priority, and iteratively inserted in `data_needed` (a b c a b c a b c...)
self.comm_threshold_bytes=90
self.target_message_size=50

current system will fetch 50 keys from worker a and 40 keys from worker b
my proposal would fetch 30 keys from worker a, 30 from worker b, and 30 from worker c
> While I'm warm to deduplicate data_needed vs. data_needed_per_worker, I don't see how any of this removes the problem about when to trigger _new_ensure_communicate_with_tasks_by_worker?

The proposal that includes an instruction `GatherDep(worker: str)` w/out explicit tasks would delay the collection of tasks until the instruction is executed. This collection would be trivially achieved by re-using the `data_per_worker` (heap)set
Assuming that we perform all transitions before we execute any instructions, this would resolve our problem, wouldn't it?
`ensure_communicating` would effectively be reduced to simply inferring whether or not we need to communicate and not _what_ to communicate.
> The proposal that includes an instruction GatherDep(worker: str) w/out explicit tasks would delay the collection of tasks until the instruction is executed.
> ensure_communicating would effectively be reduced to simply inferring whether or not we need to communicate and not what to communicate.

I see this as a major antipattern to the current worker state machine refactoring, as it would require the whole contents of `_select_keys_for_gather` to be moved out of `WorkerState` and into `Worker.gather_dep`. This in turn would require a new hack `StartedGatheringEvent` to allow for the transition from fetch to flight.
The check for `self.comm_nbytes < self.comm_threshold_bytes` would also need to move there too.
> I see this as a major antipattern to the current worker state machine refactoring, as it would require the whole contents of _select_keys_for_gather to be moved out of WorkerState and into Worker.gather_dep. This in turn would require a new hack StartedGatheringEvent to allow for the transition from fetch to flight.
The check for self.comm_nbytes < self.comm_threshold_bytes would also need to move there too.

Good point.

---

Taking a step back, I'm wondering how we'd approach this from a TDD perspective. 

1. It should only involve `WorkerState` and not `Worker(Server)`
2. The event log should have roughly the following two items (modulo unrelevant kwargs)

This could be roughly written up as something like the following

```python
assert ComputeTaskEvent(key='f1', who_has={"f2": {"workerB"}, "f3", {"workerB"}, ...) in state.event_log
...
assert GatherDep(worker="workerB", to_gather={"f2", "f3"}) in state.instructions_log
```

and also

```python
assert AcquireReplicasEvent(who_has={"f2": {"workerB"}, "f3", {"workerB"}, ...) in state.event_log

assert GatherDep(worker="workerB", to_gather={"f2", "f3"}) in state.instructions_log

```

correct?

I'm wondering if we are already able to write such a test down, after https://github.com/dask/distributed/pull/6566


---

About the actual problem, I'm wondering if the order of transitioning cannot be used to influence this. I think right now the above requests would play out like (depth first)

- f1(released)
- f1(fetch)
- f1(flight)
- f2(released)
- f2(fetch)
- f2(flight)

Would it help if we instead performed a breadth-first transitioning?

- f1(released)
- f2(released)
- f1(fetch)
- f2(fetch)
- f1(flight)
- f2(flight)
Breadth-first would help. Not necessary with my design above though - could you comment on it?

For testing, I envision something like
```python
instructions = state.handle_stimulus(AcquireReplicasEvent(...))
assert instructions == [
    GatherDep(...),
    GatherDep(...),
    GatherDep(...),
]
assert state.tasks["x"].state == "flight"
assert state.tasks["y"].state == "fetch"
```

> could you comment on it?

sure, I'll need to give it another pass

> For testing, I envision something like

I would actually prefer to leave the task states out of this for this case. I think the events and instructions should be enough and would make for great high level tests.
IIUC your algorithm would make poor decisions in the following case

Below shows data needed as a sorted list in a simplified format. Omitting key, priority etc. merely nbytes and who_has

```python
Worker.comm_threshold_bytes = 100
[
    # nbytes, who_has
    (1, {"A"}),  # -> new GatherDep for A
    (1000, {"C"}),  # -> new GatherDep for C
    # <--- Break since sum(AllGatherDepNbytes) > Worker.comm_thresholds_bytes
    (1, {"A"}),
    (1, {"A"}),
    (1, {"A"}),
    (1, {"A"}),
    (1, {"A"}),
]
```

I think fetching all tasks from A and then fetching the task from C is better.

Your proposal about _saturated_ GatherDep requests is interesting. For the sake of completeness, if we simply limit by number of outgoing connections, we'd be exposed to a similar situation as above with the following queue

```python
Worker.total_out_connections = 3
[
    # nbytes, who_has
    (1, {"A"}),  # -> new GatherDep for A
    (1, {"B"}),  # -> new GatherDep for B
    (1, {"C"}),  # -> new GatherDep for C
    # <--- Break since count(outgoing_requests) >= Worker.total_out_connections
    (1, {"A"}),
    (1, {"A"}),
    (1, {"A"}),
    (1, {"A"}),
    (1, {"A"}),
]
```

Only filtering by _saturated_ requests would not have this problem. However, I don't feel great about this proposal since it would effectively reduce the `total_out_connections` (or it's follow up) to a mere suggestion/recommendation and not a limit which lets me wonder why we'd even want to introduce this limit in the first place. I expect our server infrastructure to be capable of handling _many_ concurrent connections and busyness is controlled on the remote / incoming request (`Worker.get_data`) end. The only reason why I'd want to throttle our data fetching is to protect CPU and memory consumption and both mostly scales with size_of_bytes (i.e. de-/serialization overhead), doesn't it?

It is interesting since it asks the question about what this local/per-request limit should actually be. Do we even need a per-request limit / `target_message_size`? Every message is already naturally capped by `comm_thresholds_bytes`, i.e. individual requests are always smaller than this, anyhow. If the local worker can handle it, why can't the remote? The remote might need to handle _many_ (`total_in_connections`) of these requests and we don't want it to die. Maybe the remote busyness should just work by inspecting size_of_bytes_requested and we'd drop the local limiting entirely. Would that simplify things?

---

Zooming out a bit, there are a couple of questions that connect our two proposals which I believe can be discussed mostly separately

1. How do we iterate over the set of tasks and workers

1.A) Maintain data_needed_per_worker

```python

to_fetch = defaultdict(list)
leftover_data_per_worker = {}
while data_per_worker and not above_global_size_threshold():
    worker, tasks = pick_global_best_worker(data_per_worker)
    if worker is None:  # E.g. all in-flight. Handled by pick_global_best_worker
        break
    while tasks:
        ts = tasks.pop()
        if above_local_threshold(worker, ts):
            leftover_data_per_worker[worker] = tasks
            break
        to_fetch[worker].append(ts)

self.data_per_worker.update(leftover_data_per_worker)
```

1B) Keep data_needed

```python

to_fetch = defaultdict(list)
self.skipped_worker_in_flight_or_busy = []
while data_needed and not above_global_size_threshold():
    ts = data_needed.pop()

    worker = pick_best_worker_for_task(ts)
    if above_local_threshold(worker, ts):
        self.skipped_worker_in_flight_or_busy.append(ts)
        continue
    to_fetch[worker].append(ts)

# TODO: IIUC your proposal would envision having skipped_worker_in_flight_or_busy to be consolidated elsewhere??
```

We could engage in a thorough runtime analysis and try to estimate which is faster. My gut feeling tells me mine should fare better even with the amortized cost of maintaining `data_needed_per_worker` (which we're currently paying already). I think this shouldn't matter in practice since on the worker the number of tasks and the number of remote workers should be relatively small. However, you are suggesting to introduce `skipped_worker_in_flight_or_busy` as a performance optimization which worries me a bit.
The big difference I see is that we have different opportunities to optimize the control flow via `pick_global_best_worker` vs `pick_best_worker_for_task`. My gut tells me 1A is more powerful to eventually include load, busyness, etc. in this decision. I am also a bit worried about the additional complexity with `skipped_worker_in_flight_or_busy` becoming an instance variable. Nonlocal state is scary and prone to errors.

2. What kind of limits do we want to impose

We may not need local limits. Not sure. I see room for this in both proposals and we can kick this down the road. I just had the feeling it was more important in your proposal than in mine because you talked about `skipped_worker_in_flight_or_busy` being an instance variable. I had the impression that if we dropped request-specific limiting, this would not be necessary

3. How do we propagate this information from the transition system back to the instruction w/out moving transition logic to the server class

I think this is the big one. I see a couple of options

3A) Mutate things in place like what you described. I think there are nuances but overall this would amount to either mutating issued instructions or mutating some state the instruction references (e.g. instruction points to a dict/set on the WorkerState we keep updating). Honestly, this feels like cheating and equally opposed to the new design as letting the server perform transition. This proposal would harm replayability / lineage and I would like to avoid this

3B) Delay `ensure_communicating` artificially after all. We could just call `_ensure_communicating` every time after `self._transitions` in `handle_stimulus` and not worry about this problem anymore. I think this would still be better than before when we called it after every stream comm and occasionally whenever we felt the need for it.
We might still want to refactor some of the loop logic to clean up the batching logic / `_select_keys_for_gather` but that is less urgent.

3C) Use a breadth first transition. this might benefit us in other areas as well since it would allow us to consider priorities more strictly. FWIW, I encountered an ordering issue about transitions before on scheduler side, see https://github.com/dask/distributed/pull/6248#discussion_r862436822

3D) Perform some aggregation before `_handle_instructions`, i.e.

```python
merge_instructions([
    GatherDepDecision(worker="A", key="foo"),
    GatherDepDecision(worker="A", key="bar"),
    GatherDepDecision(worker="B", key="baz"),
])
>> {GatherDep(worker="A", keys={"foo", "bar"}), GatherDep(worker="B", keys={"baz"})
```

Of course, these three questions are related but I think we can make any combination work with appropriate adjustments

## Conclusion

- I don't have a strong preference about which way we iterate (worker vs tasks). I slightly favor per-worker because I think we have more room for improvement there. If we need a performance optimization for per-tasks just to make it work, I would again pick the per-worker case.
- I am worried about mutating instructions and would like to avoid this (See questions 3 above)

---

> At the top of _handle_instructions, there's a simple object deduplication:

FWIW, I've come across use cases for an insertion ordered set a couple of times now. We might want to consider adding such a collection instead of doing any deduplication. Should also be trivial to implement if we re-use the HeapSet.


> 3C) Use a breadth first transition. this might benefit us in other areas as well since it would allow us to consider priorities more strictly. FWIW, I encountered an ordering issue about transitions before on scheduler side, see https://github.com/dask/distributed/pull/6248#discussion_r862436822

I played with this and it doesn't work: https://github.com/crusaderky/distributed/tree/WSMR/breadth_first

Story of the new test `test_cluster_gather_deps`:
```
- ('x1', 'ensure-task-exists', 'released', 'test', 1655388950.2006707)
- ('x2', 'ensure-task-exists', 'released', 'test', 1655388950.2006752)
- ('x3', 'ensure-task-exists', 'released', 'test', 1655388950.2006783)
- ('gather-dependencies', '127.0.0.1:1235', {'x1'}, 'test', 1655388950.2007372)
- ('x1', 'released', 'fetch', 'fetch', {'x1': ('flight', '127.0.0.1:1235')}, 'test', 1655388950.2007432)
- ('x2', 'released', 'fetch', 'fetch', {}, 'test', 1655388950.2007544)
- ('x3', 'released', 'fetch', 'fetch', {}, 'test', 1655388950.2007656)
- ('x1', 'fetch', 'flight', 'flight', {}, 'test', 1655388950.2007716)
```
Depth-first (without the change to `_transitions`), it would be:
```
- ('x1', 'ensure-task-exists', 'released', 'test', 1655388950.2006707)
- ('x2', 'ensure-task-exists', 'released', 'test', 1655388950.2006752)
- ('x3', 'ensure-task-exists', 'released', 'test', 1655388950.2006783)
- ('gather-dependencies', '127.0.0.1:1235', {'x1'}, 'test', 1655388950.2007372)
- ('x1', 'released', 'fetch', 'fetch', {'x1': ('flight', '127.0.0.1:1235')}, 'test', 1655388950.2007432)
- ('x1', 'fetch', 'flight', 'flight', {}, 'test', 1655388950.2007716)
- ('x2', 'released', 'fetch', 'fetch', {}, 'test', 1655388950.2007544)
- ('x3', 'released', 'fetch', 'fetch', {}, 'test', 1655388950.2007656)
```


Breadth-first actually doesn't change anything. We're calling `transition_released_fetch` sequentially on the 3 tasks; each of them invokes `_ensure_communicating`, so the worker is moved to `in_flight_workers` by the first task.

> 3B) Delay ensure_communicating artificially after all. We could just call _ensure_communicating every time after self._transitions in handle_stimulus and not worry about this problem anymore. I think this would still be better than before when we called it after every stream comm and occasionally whenever we felt the need for it.

I appreciate the pragmaticity of this one. However it must be accompanied by a solution to the O(nlogn) issue for unschedulable tasks, since on its own it would further exacerbate it.

> 3D) Perform some aggregation before `_handle_instructions`

This is also reasonable.

I've implemented 3B in #6462, and removed data_needed in #6587.

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -859,8 +859,7 @@ def data(self) -> MutableMapping[str, Any]:
     comm_nbytes = DeprecatedWorkerStateAttribute()
     comm_threshold_bytes = DeprecatedWorkerStateAttribute()
     constrained = DeprecatedWorkerStateAttribute()
-    data_needed = DeprecatedWorkerStateAttribute()
-    data_needed_per_worker = DeprecatedWorkerStateAttribute()
+    data_needed_per_worker = DeprecatedWorkerStateAttribute(target="data_needed")
     executed_count = DeprecatedWorkerStateAttribute()
     executing_count = DeprecatedWorkerStateAttribute()
     generation = DeprecatedWorkerStateAttribute()
@@ -883,6 +882,15 @@ def data(self) -> MutableMapping[str, Any]:
     validate_task = DeprecatedWorkerStateAttribute()
     waiting_for_data_count = DeprecatedWorkerStateAttribute()
 
+    @property
+    def data_needed(self) -> set[TaskState]:
+        warnings.warn(
+            "The `Worker.data_needed` attribute has been removed; "
+            "use `Worker.state.data_needed[address]`",
+            FutureWarning,
+        )
+        return {ts for tss in self.state.data_needed.values() for ts in tss}
+
     ##################
     # Administrative #
     ##################
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
@@ -949,14 +949,10 @@ class WorkerState:
     has_what: defaultdict[str, set[str]]
 
     #: The tasks which still require data in order to execute and are in memory on at
-    #: least another worker, prioritized as a heap. All and only tasks with
-    #: ``TaskState.state == 'fetch'`` are in this collection.
-    data_needed: HeapSet[TaskState]
-
-    #: Same as :attr:`data_needed`, individually for every peer worker. A
-    #: :class:`TaskState` with multiple entries in :attr:`~TaskState.who_has` will
-    #: appear multiple times here.
-    data_needed_per_worker: defaultdict[str, HeapSet[TaskState]]
+    #: least another worker, prioritized as per-worker heaps. All and only tasks with
+    #: ``TaskState.state == 'fetch'`` are in this collection. A :class:`TaskState` with
+    #: multiple entries in :attr:`~TaskState.who_has` will appear multiple times here.
+    data_needed: defaultdict[str, HeapSet[TaskState]]
 
     #: Number of bytes to fetch from the same worker in a single call to
     #: :meth:`BaseWorker.gather_dep`. Multiple small tasks that can be fetched from the
@@ -1043,6 +1039,10 @@ class WorkerState:
     #: In production, it should always be set to False.
     transition_counter_max: int | Literal[False]
 
+    #: Statically-seeded random state, used to guarantee determinism whenever a
+    #: pseudo-random choice is required
+    rng: random.Random
+
     __slots__ = tuple(__annotations__)
 
     def __init__(
@@ -1077,8 +1077,7 @@ def __init__(
         self.running = True
         self.waiting_for_data_count = 0
         self.has_what = defaultdict(set)
-        self.data_needed = HeapSet(key=operator.attrgetter("priority"))
-        self.data_needed_per_worker = defaultdict(
+        self.data_needed = defaultdict(
             partial(HeapSet[TaskState], key=operator.attrgetter("priority"))
         )
         self.in_flight_workers = {}
@@ -1100,6 +1099,7 @@ def __init__(
         self.transition_counter = 0
         self.transition_counter_max = transition_counter_max
         self.actors = {}
+        self.rng = random.Random(0)
 
     def handle_stimulus(self, *stims: StateMachineEvent) -> Instructions:
         """Process one or more external events, transition relevant tasks to new states,
@@ -1196,12 +1196,12 @@ def _update_who_has(self, who_has: Mapping[str, Collection[str]]) -> None:
             for worker in ts.who_has - workers:
                 self.has_what[worker].remove(key)
                 if ts.state == "fetch":
-                    self.data_needed_per_worker[worker].remove(ts)
+                    self.data_needed[worker].remove(ts)
 
             for worker in workers - ts.who_has:
                 self.has_what[worker].add(key)
                 if ts.state == "fetch":
-                    self.data_needed_per_worker[worker].add(ts)
+                    self.data_needed[worker].add(ts)
 
             ts.who_has = workers
 
@@ -1218,9 +1218,8 @@ def _purge_state(self, ts: TaskState) -> None:
 
         for worker in ts.who_has:
             self.has_what[worker].discard(ts.key)
-            self.data_needed_per_worker[worker].discard(ts)
+            self.data_needed[worker].discard(ts)
         ts.who_has.clear()
-        self.data_needed.discard(ts)
 
         self.threads.pop(key, None)
 
@@ -1238,66 +1237,49 @@ def _purge_state(self, ts: TaskState) -> None:
         self.in_flight_tasks.discard(ts)
 
     def _ensure_communicating(self, *, stimulus_id: str) -> RecsInstrs:
-        if not self.running:
+        """Transition tasks from fetch to flight, until there are no more tasks in fetch
+        state or a threshold has been reached.
+        """
+        if not self.running or not self.data_needed:
+            return {}, []
+        if (
+            len(self.in_flight_workers) >= self.total_out_connections
+            and self.comm_nbytes >= self.comm_threshold_bytes
+        ):
             return {}, []
-
-        skipped_worker_in_flight_or_busy = []
 
         recommendations: Recs = {}
         instructions: Instructions = []
 
-        while self.data_needed and (
-            len(self.in_flight_workers) < self.total_out_connections
-            or self.comm_nbytes < self.comm_threshold_bytes
-        ):
+        for worker, available_tasks in self._select_workers_for_gather():
+            assert worker != self.address
+            to_gather_tasks, total_nbytes = self._select_keys_for_gather(
+                available_tasks
+            )
+            assert to_gather_tasks
+            to_gather_keys = {ts.key for ts in to_gather_tasks}
+
             logger.debug(
-                "Ensure communicating. Pending: %d. Connections: %d/%d. Busy: %d",
+                "Gathering %d tasks from %s; %d more remain. "
+                "Pending workers: %d; connections: %d/%d; busy: %d",
+                len(to_gather_tasks),
+                worker,
+                len(available_tasks),
                 len(self.data_needed),
                 len(self.in_flight_workers),
                 self.total_out_connections,
                 len(self.busy_workers),
             )
-
-            ts = self.data_needed.pop()
-
-            if self.validate:
-                assert ts.state == "fetch"
-                assert self.address not in ts.who_has
-
-            if not ts.who_has:
-                recommendations[ts] = "missing"
-                continue
-
-            workers = [
-                w
-                for w in ts.who_has
-                if w not in self.in_flight_workers and w not in self.busy_workers
-            ]
-            if not workers:
-                skipped_worker_in_flight_or_busy.append(ts)
-                continue
-
-            for w in ts.who_has:
-                self.data_needed_per_worker[w].remove(ts)
-
-            host = get_address_host(self.address)
-            local = [w for w in workers if get_address_host(w) == host]
-            worker = random.choice(local or workers)
-
-            to_gather_tasks, total_nbytes = self._select_keys_for_gather(worker, ts)
-            to_gather_keys = {ts.key for ts in to_gather_tasks}
-
             self.log.append(
                 ("gather-dependencies", worker, to_gather_keys, stimulus_id, time())
             )
 
-            self.comm_nbytes += total_nbytes
-            self.in_flight_workers[worker] = to_gather_keys
-            for d_ts in to_gather_tasks:
+            for ts in to_gather_tasks:
                 if self.validate:
-                    assert d_ts.state == "fetch"
-                    assert d_ts not in recommendations
-                recommendations[d_ts] = ("flight", worker)
+                    assert ts.state == "fetch"
+                    assert worker in ts.who_has
+                    assert ts not in recommendations
+                recommendations[ts] = ("flight", worker)
 
             # A single invocation of _ensure_communicating may generate up to one
             # GatherDep instruction per worker. Multiple tasks from the same worker may
@@ -1313,11 +1295,104 @@ def _ensure_communicating(self, *, stimulus_id: str) -> RecsInstrs:
                 )
             )
 
-        for ts in skipped_worker_in_flight_or_busy:
-            self.data_needed.add(ts)
+            self.in_flight_workers[worker] = to_gather_keys
+            self.comm_nbytes += total_nbytes
+            if (
+                len(self.in_flight_workers) >= self.total_out_connections
+                and self.comm_nbytes >= self.comm_threshold_bytes
+            ):
+                break
 
         return recommendations, instructions
 
+    def _select_workers_for_gather(self) -> Iterator[tuple[str, HeapSet[TaskState]]]:
+        """Helper of _ensure_communicating.
+
+        Yield the peer workers and tasks in data_needed, sorted by:
+
+        1. By highest-priority task available across all workers
+        2. If tied, first by local peer workers, then remote. Note that, if a task is
+           replicated across multiple host, it may go in a tie with itself.
+        3. If still tied, by number of tasks available to be fetched from the host
+           (see note below)
+        4. If still tied, by a random element. This is statically seeded to guarantee
+           reproducibility.
+
+           FIXME https://github.com/dask/distributed/issues/6620
+                 You won't get determinism when a single task is replicated on multiple
+                 workers, because TaskState.who_has changes order at every interpreter
+                 restart.
+
+        Omit workers that are either busy or in flight.
+        Remove peer workers with no tasks from data_needed.
+
+        Note
+        ----
+        Instead of number of tasks, we could've measured total nbytes and/or number of
+        tasks that only exist on the worker. Raw number of tasks is cruder but simpler.
+        """
+        host = get_address_host(self.address)
+        heap = []
+
+        for worker, tasks in list(self.data_needed.items()):
+            if not tasks:
+                del self.data_needed[worker]
+                continue
+            if worker in self.in_flight_workers or worker in self.busy_workers:
+                continue
+            heap.append(
+                (
+                    tasks.peek().priority,
+                    get_address_host(worker) != host,  # False < True
+                    -len(tasks),
+                    self.rng.random(),
+                    worker,
+                    tasks,
+                )
+            )
+
+        heapq.heapify(heap)
+        while heap:
+            _, is_remote, ntasks_neg, rnd, worker, tasks = heapq.heappop(heap)
+            # The number of tasks and possibly the top priority task may have changed
+            # since the last sort, since _select_keys_for_gather may have removed tasks
+            # that are also replicated on a higher-priority worker.
+            if not tasks:
+                del self.data_needed[worker]
+            elif -ntasks_neg != len(tasks):
+                heapq.heappush(
+                    heap,
+                    (tasks.peek().priority, is_remote, -len(tasks), rnd, worker, tasks),
+                )
+            else:
+                yield worker, tasks
+                if not tasks:  # _select_keys_for_gather just emptied it
+                    del self.data_needed[worker]
+
+    def _select_keys_for_gather(
+        self, available: HeapSet[TaskState]
+    ) -> tuple[list[TaskState], int]:
+        """Helper of _ensure_communicating.
+
+        Fetch all tasks that are replicated on the target worker within a single
+        message, up to target_message_size.
+        """
+        to_gather: list[TaskState] = []
+        total_nbytes = 0
+
+        while available:
+            ts = available.peek()
+            # The top-priority task is fetched regardless of its size
+            if to_gather and total_nbytes + ts.get_nbytes() > self.target_message_size:
+                break
+            for worker in ts.who_has:
+                # This also effectively pops from available
+                self.data_needed[worker].remove(ts)
+            to_gather.append(ts)
+            total_nbytes += ts.get_nbytes()
+
+        return to_gather, total_nbytes
+
     def _ensure_computing(self) -> RecsInstrs:
         if not self.running:
             return {}, []
@@ -1461,37 +1536,6 @@ def _put_key_in_memory(
         self.log.append((ts.key, "put-in-memory", stimulus_id, time()))
         return recommendations
 
-    def _select_keys_for_gather(
-        self, worker: str, ts: TaskState
-    ) -> tuple[set[TaskState], int]:
-        """``_ensure_communicating`` decided to fetch a single task from a worker,
-        following priority. In order to minimise overhead, request fetching other tasks
-        from the same worker within the message, following priority for the single
-        worker but ignoring higher priority tasks from other workers, up to
-        ``target_message_size``.
-        """
-        tss = {ts}
-        total_bytes = ts.get_nbytes()
-        tasks = self.data_needed_per_worker[worker]
-
-        while tasks:
-            ts = tasks.peek()
-            if self.validate:
-                assert ts.state == "fetch"
-                assert worker in ts.who_has
-            if total_bytes + ts.get_nbytes() > self.target_message_size:
-                break
-            tasks.pop()
-            self.data_needed.remove(ts)
-            for other_worker in ts.who_has:
-                if other_worker != worker:
-                    self.data_needed_per_worker[other_worker].remove(ts)
-
-            tss.add(ts)
-            total_bytes += ts.get_nbytes()
-
-        return tss, total_bytes
-
     ###############
     # Transitions #
     ###############
@@ -1503,9 +1547,8 @@ def _transition_generic_fetch(self, ts: TaskState, stimulus_id: str) -> RecsInst
         ts.state = "fetch"
         ts.done = False
         assert ts.priority
-        self.data_needed.add(ts)
         for w in ts.who_has:
-            self.data_needed_per_worker[w].add(ts)
+            self.data_needed[w].add(ts)
         return {}, []
 
     def _transition_missing_waiting(
@@ -1613,9 +1656,8 @@ def _transition_fetch_flight(
             assert ts.state == "fetch"
             assert ts.who_has
             # The task has already been removed by _ensure_communicating
-            assert ts not in self.data_needed
             for w in ts.who_has:
-                assert ts not in self.data_needed_per_worker[w]
+                assert ts not in self.data_needed[w]
 
         ts.done = False
         ts.state = "flight"
@@ -1623,13 +1665,6 @@ def _transition_fetch_flight(
         self.in_flight_tasks.add(ts)
         return {}, []
 
-    def _transition_fetch_missing(
-        self, ts: TaskState, *, stimulus_id: str
-    ) -> RecsInstrs:
-        # _ensure_communicating could have just popped this task out of data_needed
-        self.data_needed.discard(ts)
-        return self._transition_generic_missing(ts, stimulus_id=stimulus_id)
-
     def _transition_memory_released(
         self, ts: TaskState, *, stimulus_id: str
     ) -> RecsInstrs:
@@ -2151,7 +2186,7 @@ def _transition_released_forgotten(
         ("executing", "released"): _transition_executing_released,
         ("executing", "rescheduled"): _transition_executing_rescheduled,
         ("fetch", "flight"): _transition_fetch_flight,
-        ("fetch", "missing"): _transition_fetch_missing,
+        ("fetch", "missing"): _transition_generic_missing,
         ("fetch", "released"): _transition_generic_released,
         ("flight", "error"): _transition_flight_error,
         ("flight", "fetch"): _transition_flight_fetch,
@@ -2500,21 +2535,22 @@ def _handle_compute_task(self, ev: ComputeTaskEvent) -> RecsInstrs:
 
             if self.validate:
                 assert ev.who_has.keys() == ev.nbytes.keys()
-                assert all(ev.who_has.values())
+                for dep_workers in ev.who_has.values():
+                    assert dep_workers
+                    assert len(dep_workers) == len(set(dep_workers))
 
-            for dep_key, dep_workers in ev.who_has.items():
+            for dep_key, nbytes in ev.nbytes.items():
                 dep_ts = self._ensure_task_exists(
                     key=dep_key,
                     priority=priority,
                     stimulus_id=ev.stimulus_id,
                 )
+                self.tasks[dep_key].nbytes = nbytes
+
                 # link up to child / parents
                 ts.dependencies.add(dep_ts)
                 dep_ts.dependents.add(ts)
 
-            for dep_key, value in ev.nbytes.items():
-                self.tasks[dep_key].nbytes = value
-
             self._update_who_has(ev.who_has)
         else:
             raise RuntimeError(  # pragma: nocover
@@ -2549,7 +2585,7 @@ def _handle_gather_dep_success(self, ev: GatherDepSuccessEvent) -> RecsInstrs:
                 self.log.append((ts.key, "missing-dep", ev.stimulus_id, time()))
                 if self.validate:
                     assert ts.state != "fetch"
-                    assert ts not in self.data_needed_per_worker[ev.worker]
+                    assert ts not in self.data_needed[ev.worker]
                 ts.who_has.discard(ev.worker)
                 self.has_what[ev.worker].discard(ts.key)
                 recommendations[ts] = "fetch"
@@ -2600,16 +2636,24 @@ def _handle_gather_dep_network_failure(
         either retry a different worker, or ask the scheduler to inform us of a new
         worker if no other worker is available.
         """
-        self.data_needed_per_worker.pop(ev.worker)
-        for key in self.has_what.pop(ev.worker):
-            ts = self.tasks[key]
-            ts.who_has.discard(ev.worker)
-
         recommendations: Recs = {}
+
         for ts in self._gather_dep_done_common(ev):
             self.log.append((ts.key, "missing-dep", ev.stimulus_id, time()))
             recommendations[ts] = "fetch"
 
+        for ts in self.data_needed.pop(ev.worker, ()):
+            if self.validate:
+                assert ts.state == "fetch"
+                assert ev.worker in ts.who_has
+            if ts.who_has == {ev.worker}:
+                # This can override a recommendation from the previous for loop
+                recommendations[ts] = "missing"
+
+        for key in self.has_what.pop(ev.worker):
+            ts = self.tasks[key]
+            ts.who_has.remove(ev.worker)
+
         return recommendations, []
 
     @_handle_event.register
@@ -2821,10 +2865,9 @@ def _to_dict(self, *, exclude: Container[str] = ()) -> dict:
             "ready": self.ready,
             "constrained": self.constrained,
             "data": dict.fromkeys(self.data),
-            "data_needed": [ts.key for ts in self.data_needed.sorted()],
-            "data_needed_per_worker": {
+            "data_needed": {
                 w: [ts.key for ts in tss.sorted()]
-                for w, tss in self.data_needed_per_worker.items()
+                for w, tss in self.data_needed.items()
             },
             "executing": {ts.key for ts in self.executing},
             "long_running": self.long_running,
@@ -2888,11 +2931,10 @@ def _validate_task_fetch(self, ts: TaskState) -> None:
         assert ts.key not in self.data
         assert self.address not in ts.who_has
         assert not ts.done
-        assert ts in self.data_needed
-        # Note: ts.who_has may be empty; see GatherDepNetworkFailureEvent
+        assert ts.who_has
         for w in ts.who_has:
             assert ts.key in self.has_what[w]
-            assert ts in self.data_needed_per_worker[w]
+            assert ts in self.data_needed[w]
 
     def _validate_task_missing(self, ts: TaskState) -> None:
         assert ts.key not in self.data
@@ -2916,8 +2958,7 @@ def _validate_task_released(self, ts: TaskState) -> None:
         assert ts.key not in self.data
         assert not ts._next
         assert not ts._previous
-        assert ts not in self.data_needed
-        for tss in self.data_needed_per_worker.values():
+        for tss in self.data_needed.values():
             assert ts not in tss
         assert ts not in self.executing
         assert ts not in self.in_flight_tasks
@@ -3001,19 +3042,15 @@ def validate_state(self) -> None:
                 assert k in self.tasks, self.story(k)
                 assert worker in self.tasks[k].who_has
 
-        for ts in self.data_needed:
-            assert ts.state == "fetch", self.story(ts)
-        for worker, tss in self.data_needed_per_worker.items():
+        for worker, tss in self.data_needed.items():
             for ts in tss:
                 assert ts.state == "fetch"
-                assert ts in self.data_needed
                 assert worker in ts.who_has
 
         # Test that there aren't multiple TaskState objects with the same key in any
         # Set[TaskState]. See note in TaskState.__hash__.
         for ts in chain(
-            self.data_needed,
-            *self.data_needed_per_worker.values(),
+            *self.data_needed.values(),
             self.missing_dep_flight,
             self.in_flight_tasks,
             self.executing,
@@ -3027,8 +3064,7 @@ def validate_state(self) -> None:
             assert self.transition_counter < self.transition_counter_max
 
         # Test that there aren't multiple TaskState objects with the same key in data_needed
-        assert len({ts.key for ts in self.data_needed}) == len(self.data_needed)
-        for tss in self.data_needed_per_worker.values():
+        for tss in self.data_needed.values():
             assert len({ts.key for ts in tss}) == len(tss)
 
 


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
@@ -2992,7 +2992,7 @@ def __sizeof__(self):
         )
         await b.in_gather_dep.wait()
         assert len(b.state.in_flight_tasks) == 5
-        assert len(b.state.data_needed) == 5
+        assert len(b.state.data_needed[a.address]) == 5
         b.block_gather_dep.set()
         while len(b.data) < 10:
             await asyncio.sleep(0.01)
@@ -3348,7 +3348,6 @@ async def test_Worker__to_dict(c, s, a):
         "transition_counter",
         "tasks",
         "data_needed",
-        "data_needed_per_worker",
     }
     assert d["tasks"]["x"]["key"] == "x"
     assert d["data"] == {"x": None}
diff --git a/distributed/tests/test_worker_memory.py b/distributed/tests/test_worker_memory.py
--- a/distributed/tests/test_worker_memory.py
+++ b/distributed/tests/test_worker_memory.py
@@ -527,7 +527,7 @@ def pause_on_unpickle():
     # - w and z respectively make x and y go into fetch state.
     #   w has a higher priority than z, therefore w's dependency x has a higher priority
     #   than z's dependency y.
-    #   a.state.data_needed = ["x", "y"]
+    #   a.state.data_needed[b.address] = ["x", "y"]
     # - ensure_communicating decides to fetch x but not to fetch y together with it, as
     #   it thinks x is 1TB in size
     # - x fetch->flight; a is added to in_flight_workers
@@ -543,7 +543,7 @@ def pause_on_unpickle():
     await asyncio.sleep(0.1)
     assert a.state.tasks["y"].state == "fetch"
     assert "y" not in a.data
-    assert [ts.key for ts in a.state.data_needed] == ["y"]
+    assert [ts.key for ts in a.state.data_needed[b.address]] == ["y"]
 
     # Unpausing kicks off ensure_communicating again
     a.status = Status.running
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
@@ -31,6 +31,7 @@
     ExecuteSuccessEvent,
     GatherDep,
     Instruction,
+    PauseEvent,
     RecommendationsConflict,
     RefreshWhoHasEvent,
     ReleaseWorkerDataMsg,
@@ -39,6 +40,7 @@
     SerializedTask,
     StateMachineEvent,
     TaskState,
+    UnpauseEvent,
     UpdateDataEvent,
     merge_recs_instructions,
 )
@@ -117,8 +119,7 @@ def test_WorkerState__to_dict(ws):
         "busy_workers": [],
         "constrained": [],
         "data": {"y": None},
-        "data_needed": [],
-        "data_needed_per_worker": {"127.0.0.1:1235": []},
+        "data_needed": {},
         "executing": [],
         "in_flight_tasks": ["x"],
         "in_flight_workers": {"127.0.0.1:1235": ["x"]},
@@ -860,6 +861,9 @@ async def test_deprecated_worker_attributes(s, a, b):
     with pytest.warns(FutureWarning, match=msg):
         assert a.in_flight_tasks == 0
 
+    with pytest.warns(FutureWarning, match="attribute has been removed"):
+        assert a.data_needed == set()
+
 
 @pytest.mark.parametrize(
     "nbytes,n_in_flight",
@@ -885,3 +889,92 @@ def test_aggregate_gather_deps(ws, nbytes, n_in_flight):
     assert len(instructions) == 1
     assert isinstance(instructions[0], GatherDep)
     assert len(ws.in_flight_tasks) == n_in_flight
+
+
+def test_gather_priority(ws):
+    """Test that tasks are fetched in the following order:
+
+    1. by task priority
+    2. in case of tie, from local workers first
+    3. in case of tie, from the worker with the most tasks queued
+    4. in case of tie, from a random worker (which is actually deterministic).
+    """
+    ws.total_out_connections = 4
+
+    instructions = ws.handle_stimulus(
+        PauseEvent(stimulus_id="pause"),
+        # Note: tasks fetched by acquire-replicas always have priority=(1, )
+        AcquireReplicasEvent(
+            who_has={
+                # Remote + local
+                "x1": ["127.0.0.2:1", "127.0.0.1:2"],
+                # Remote. After getting x11 from .1, .2  will have less tasks than .3
+                "x2": ["127.0.0.2:1"],
+                "x3": ["127.0.0.3:1"],
+                "x4": ["127.0.0.3:1"],
+                # It will be a random choice between .2, .4, .5, .6, and .7
+                "x5": ["127.0.0.4:1"],
+                "x6": ["127.0.0.5:1"],
+                "x7": ["127.0.0.6:1"],
+                # This will be fetched first because it's on the same worker as y
+                "x8": ["127.0.0.7:1"],
+            },
+            # Substantial nbytes prevents total_out_connections to be overridden by
+            # comm_threshold_bytes, but it's less than target_message_size
+            nbytes={f"x{i}": 4 * 2**20 for i in range(1, 9)},
+            stimulus_id="compute1",
+        ),
+        # A higher-priority task, even if scheduled later, is fetched first
+        ComputeTaskEvent(
+            key="z",
+            who_has={"y": ["127.0.0.7:1"]},
+            nbytes={"y": 1},
+            priority=(0,),
+            duration=1.0,
+            run_spec=None,
+            function=None,
+            args=None,
+            kwargs=None,
+            resource_restrictions={},
+            actor=False,
+            annotations={},
+            stimulus_id="compute2",
+        ),
+        UnpauseEvent(stimulus_id="unpause"),
+    )
+
+    assert instructions == [
+        # Highest-priority task first. Lower priority tasks from the same worker are
+        # shoved into the same instruction (up to 50MB worth)
+        GatherDep(
+            stimulus_id="unpause",
+            worker="127.0.0.7:1",
+            to_gather={"y", "x8"},
+            total_nbytes=1 + 4 * 2**20,
+        ),
+        # Followed by local workers
+        GatherDep(
+            stimulus_id="unpause",
+            worker="127.0.0.1:2",
+            to_gather={"x1"},
+            total_nbytes=4 * 2**20,
+        ),
+        # Followed by remote workers with the most tasks
+        GatherDep(
+            stimulus_id="unpause",
+            worker="127.0.0.3:1",
+            to_gather={"x3", "x4"},
+            total_nbytes=8 * 2**20,
+        ),
+        # Followed by other remote workers, randomly.
+        # Determinism is guaranteed by a statically-seeded random number generator.
+        # FIXME It would have not been deterministic if we instead of multiple keys we
+        #       had used a single key with multiple workers, because sets
+        #       (like TaskState.who_has) change order at every interpreter restart.
+        GatherDep(
+            stimulus_id="unpause",
+            worker="127.0.0.4:1",
+            to_gather={"x5"},
+            total_nbytes=4 * 2**20,
+        ),
+    ]