Repository: dask/distributed

Issue:
Worker can memory overflow by fetching too much data at once
There is a possibility for a worker to suicide by fetching too much data at once

Old version:

https://github.com/dask/distributed/blob/84cbb099e0d0f1e1947d51d957f0b65e9b0635dd/distributed/worker.py#L2754-L2757

New version:

https://github.com/dask/distributed/blob/3647cfefd004281dac0ffa92349e23bc2a94d68a/distributed/worker_state_machine.py#L1425-L1429

Simple example

Let's assume that keys have on average ~500MB
Local memory available is ~10GB
total_out_connections: 50 (default)

That would cause us to fetch up to 25GB at once which would kill the worker immediately.

This problem is exacerbated if the keys data sizes are underestimated (e.g. bias/bug in dask.sizeof)


cc @crusaderky 


Discussion:
> That would cause us to fetch up to 25GB at once which would kill the worker immediately.

I think it's also likely this wouldn't always kill the worker, but freeze it like in https://github.com/dask/distributed/issues/6110. If the 25GiB come somewhat staggered in those 500MB increments, physical memory usage will creep up very close to the 10GB max, without actually requesting more than it. Once it gets close enough, the system will thrash to the point that everything stops, so more memory isn't allocated (which would put it over the limit and kill it). The smaller the key size and larger the key count, this more likely this is to happen (probably why a shuffle was a good reproducer).

Since we know the (approximate) size of all the keys we're fetching, this seems like something we can avoid doing in the first place (don't fetch more than we have memory capacity for, as a start). Maybe also a use-case for receiving data directly to disk https://github.com/dask/distributed/issues/5996?
pseudocode:
```python
if current managed_in_memory + self.comm_nbytes > threshold% * max_memory:
    stop scheduling further gather_dep's
```
or with better encapsulation:
```python
    if self.data.would_exceed_threshold(self.comm_nbytes):
        stop scheduling further gather_dep's
```
There are two ways to resume gathering:
1. a dependency fails to arrive, thus reducing comm_nbytes
2. enough dependencies arrive to send the worker over threshold. Older keys are automatically spilled, thus reducing managed_in_memory until it's below threshold.

Note that the above mitigates the issue, but doesn't actually prevent the use case that @fjetter showed above in case of a task that individually requires more memory than what's available:

```python
CHUNK_SIZE = 500 // 8 * 2**20  # 500 MiB
a = da.random.random(CHUNK_SIZE * 20, CHUNK_SIZE)  # 10 GiB total
a = a.persist()
f = client.submit(lambda _: None, futures_of(a))
```
What will happen is that whatever worker runs f will start accumulating all chunks of a locally.
In main, if there are 20+ workers in the cluster, this means immediate death from gather_dep.
With the change above, gather_dep will be successful and all the data will be copied onto the worker and progressively spilled.
However, when the time comes to transition f from ready to running, then
1. chunks of a will be moved from slow to fast
2. before f can start running, the first chunks are already being sent back from fast to slow; however they also exist in memory. The worker goes to 100% and if we're lucky it's terminated, if not it just freezes due to swap thrashing.

And this is not considering unrelated unmanaged memory - e.g. on a multithreaded worker, there may be other tasks running.



Two thoughts:

1.  How often does this happen?  Presumably this was related to the Florian Shuffle Deadlock (I need a better name for this).  Have we seen this happen in other common situations?

> Let's assume that keys have on average ~500MB
> Local memory available is ~10GB
> total_out_connections: 50 (default)

> That would cause us to fetch up to 25GB at once which would kill the worker immediately.

If planning for worst case would cripple things (You must have 25GB of memory free before you do any work!) then I would not want to plan for worst case.  I think that we could reduce this burden in a few ways:

1. We could try to learn the size of transfers on the worker (creating something like a `worker.py::TaskPrefix`) 
2. We could ask the sender's side to cancel for us.  "Hey other worker!  I want data "x" but I'm running kinda low on memory.  I only have 4 GB free.   Please only send something if it's less than 2GB?  If it's greater, let me know how large it is instead so that I can plan accordingly"
3. We could start to reduce the number of concurrent inbound connetions as our memory starts to dry up.

Just some thoughts to reduce the burden here.  Please don't let these distract too much from the main conversation.
> 1. We could try to learn the size of transfers on the worker (creating something like a worker.py::TaskPrefix)
> 2. We could ask the sender's side to cancel for us. "Hey other worker! I want data "x" but I'm running kinda low on memory. I only have 4 GB free. Please only send something if it's less than 2GB? If it's greater, let me know how large it is instead so that I can plan accordingly"

The receiving worker already knows in advance how large the data will be before it asks for it. That's returned by ``select_keys_for_gather`` which in turn is reflected into ``Worker.comm_nbytes``.

> 3. We could start to reduce the number of concurrent inbound connetions as our memory starts to dry up.

I would not want to count workers towards the maximum if they are going to send me only a few kilobytes .
Ah, great then.  
In coiled/coiled-runtime#229, we have found a reliable reproducer for a worker getting `oom-killed` by fetching too much data. The issue is exacerbated by a misconfigured `Worker.memory_limit` (coiled/feedback#185), but I am confident that we can also reproduce this issue if that is fixed by simply scaling `test_tensordot_stress` to more workers. With a default chunk size of 100 MiB and the default of 50 outgoing and 10 incoming connections per worker, we might use up to ~2.9 GiB for communication, which is ~90 % of all available memory available to the worker on a `t3.medium`.

EDIT: Apart from 100 MiB being the default chosen by `dask.array` itself in coiled/coiled-runtime#229, we recommend chunk sizes between 10(0) MiB and 1 GiB in several places:
* https://blog.dask.org/2021/11/02/choosing-dask-chunk-sizes#rough-rules-of-thumb
* https://docs.dask.org/en/stable/array-chunks.html#specifying-chunk-shapes
One quick yet imperfect fix might be to introduce a parameter that limits the amount of memory we want to dedicate to communications, effectively adjusting the above if-clause to
```python3
if ( 
    (
        len(self.in_flight_workers) >= self.total_out_connections 
         and self.comm_nbytes >= self.comm_threshold_bytes
    ) or self.comm_nbytes >= self.comm_memory_limit
 ): 
     return {}, [] 
```
This value could be set as a percentage of `Worker.memory_limit` to account for different machines. As a downside, this might have negative effects on situations where all we want to do for now is fetching an initial set of data to work with.

To better understand the burden comms put onto our workers, we may also want to add metrics about the number of open connections and the size of `self.comm_nbytes`. XREF: #6892

In addition, we may want to drop `self.comm_threshold_bytes` as well as the limits on incoming/outgoing connections as a parameters and fully rely on a mechanism that looks at currently used memory to schedule comms limited by the size of the connection pool. 

cc @fjetter 
It's also worth noting that the same issue exists for workers _sending_ data. There's a [limit to the number of incoming connections](https://github.com/dask/distributed/blob/0d21c7860888bae72e6ce20c38f93bc34b0454e5/distributed/worker.py#L1660-L1663) a worker can have at once (default 10), but not the number of bytes it can be sending at once. If a worker is paused, it abruptly [limits itself to 1 connection](https://github.com/dask/distributed/blob/0d21c7860888bae72e6ce20c38f93bc34b0454e5/distributed/worker.py#L1654-L1655), but that message could still be large.

Serialization may not be zero-copy. If compression is used, it's never zero-copy. TLS comms will also require a copy in the SSL buffer. So even if the sent data is already in memory, we could potentially be doubling (or more) memory usage of that data until the other worker confirms it _received_ the message.

The worst case is when the data being sent is currently spilled to disk. `get_data` will then read it into memory (which stores a strong reference to it in the `data.fast` dict, so even after it's been sent, it will remain in memory at least until the next memory monitor interval, which could a half-second away—months away, compared to the speed of memory allocations). Then, it might have even more copies in serialization, compression, and encryption. Sending a spilled key can be a very costly operation, as well as slow, which is why I advocate for https://github.com/dask/distributed/issues/5996 so much.

Ideally, a worker shouldn't be willing to send more than `memory_target - current_memory` bytes. Of course, `current_memory` isn't something we can measure all the time, so a fixed threshold like @hendrikmakait suggested above might be a good place to start.
I wrote this a big ago, but I'd like to point to https://github.com/dask/distributed/issues/6212 as a more holistic approach to this problem. The overall idea is that:

- We do many things concurrently that require non-trivial amounts of memory (run user tasks, receive keys, send keys, etc.).
- If we have 1gb of memory left, and we want to both receive 600mb of data from another worker, and load 600mb of spilled data back into memory to run a task, wouldn't it be nice if those operations could coordinate to wait for each other, instead of running concurrently and probably killing the worker?

It's common to use a synchronization primitive like a semaphore or [`CapacityLimiter`](https://trio.readthedocs.io/en/stable/reference-core.html#trio.CapacityLimiter) to manage concurrent code's access to a limited resource. Maybe we could have a similar structure around memory?
> The worst case is when the data being sent is currently spilled to disk. `get_data` will then read it into memory (which stores a strong reference to it in the `data.fast` dict, so even after it's been sent, it will remain in memory at least until the next memory monitor interval, which could a half-second away

Not quite correct. If unspilling a key causes the reported managed memory to get above `target` (60%), then other keys will immediately be spilled out, until you are again below `target`.

The mechanism you mention happens specifically and only when managed memory is below 60% and process memory is above `spill` (70%). In other words, you either have at least 10% unmanaged memory or `sizeof()` is severely under-reporting memory usage.
The 10% unmanaged memory use case is interesting, as it is a lot easier to reach on smaller workers. A brand new worker (with pandas imported) occupies a baseline of 138 MiB on my machine.
> It's common to use a synchronization primitive like a semaphore or [CapacityLimiter](https://trio.readthedocs.io/en/stable/reference-core.html#trio.CapacityLimiter) to manage concurrent code's access to a limited resource. Maybe we could have a similar structure around memory?

We do not know how much data sending/receiving actually requires, nor do we know how much a user task requires. Without reliable measurements I would prefer sticking to a pragmatic, low complexity approach for now.

I do believe that simply introducing a threshold as proposed by @hendrikmakait  in https://github.com/dask/distributed/issues/6208#issuecomment-1216710657 will improve things significantly already.

I'm not sure if an absolute threshold, a relative (to memory_limit) threshold or a dynamic threshold (percentage of free/available memory) fits best. My gut feeling tells me to start with a relative threshold (based on memory_limit). Together with the added instrumentation, we can iterate on this choice incrementally, if necessary. If there are still significant shortcomings afterwards, we can think about a more complex solution.

Resolving patch:
diff --git a/distributed/worker.py b/distributed/worker.py
--- a/distributed/worker.py
+++ b/distributed/worker.py
@@ -745,6 +745,18 @@ def __init__(
             memory_spill_fraction=memory_spill_fraction,
             memory_pause_fraction=memory_pause_fraction,
         )
+
+        transfer_incoming_bytes_limit = None
+        transfer_incoming_bytes_fraction = dask.config.get(
+            "distributed.worker.memory.transfer"
+        )
+        if (
+            self.memory_manager.memory_limit is not None
+            and transfer_incoming_bytes_fraction is not False
+        ):
+            transfer_incoming_bytes_limit = int(
+                self.memory_manager.memory_limit * transfer_incoming_bytes_fraction
+            )
         state = WorkerState(
             nthreads=nthreads,
             data=self.memory_manager.data,
@@ -754,6 +766,7 @@ def __init__(
             transfer_incoming_count_limit=transfer_incoming_count_limit,
             validate=validate,
             transition_counter_max=transition_counter_max,
+            transfer_incoming_bytes_limit=transfer_incoming_bytes_limit,
         )
         BaseWorker.__init__(self, state)
 
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
@@ -1232,6 +1232,9 @@ class WorkerState:
     #: In production, it should always be set to False.
     transition_counter_max: int | Literal[False]
 
+    #: Limit of bytes for incoming data transfers; this is used for throttling.
+    transfer_incoming_bytes_limit: int | None
+
     #: Statically-seeded random state, used to guarantee determinism whenever a
     #: pseudo-random choice is required
     rng: random.Random
@@ -1250,6 +1253,7 @@ def __init__(
         transfer_incoming_count_limit: int = 9999,
         validate: bool = True,
         transition_counter_max: int | Literal[False] = False,
+        transfer_incoming_bytes_limit: int | None = None,
     ):
         self.nthreads = nthreads
 
@@ -1293,6 +1297,7 @@ def __init__(
         self.stimulus_log = deque(maxlen=10_000)
         self.transition_counter = 0
         self.transition_counter_max = transition_counter_max
+        self.transfer_incoming_bytes_limit = transfer_incoming_bytes_limit
         self.actors = {}
         self.rng = random.Random(0)
 
@@ -1468,17 +1473,36 @@ def _purge_state(self, ts: TaskState) -> None:
         self.long_running.discard(ts)
         self.in_flight_tasks.discard(ts)
 
+    def _should_throttle_incoming_transfers(self) -> bool:
+        """Decides whether the WorkerState should throttle data transfers from other workers.
+
+        Returns
+        -------
+        * True if the number of incoming data transfers reached its limit
+        and the size of incoming data transfers reached the minimum threshold for throttling
+        * True if the size of incoming data transfers reached its limit
+        * False otherwise
+        """
+        reached_count_limit = (
+            self.transfer_incoming_count >= self.transfer_incoming_count_limit
+        )
+        reached_throttle_threshold = (
+            self.transfer_incoming_bytes
+            >= self.transfer_incoming_bytes_throttle_threshold
+        )
+        reached_bytes_limit = (
+            self.transfer_incoming_bytes_limit is not None
+            and self.transfer_incoming_bytes >= self.transfer_incoming_bytes_limit
+        )
+        return reached_count_limit and reached_throttle_threshold or reached_bytes_limit
+
     def _ensure_communicating(self, *, stimulus_id: str) -> RecsInstrs:
         """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 (
-            self.transfer_incoming_count >= self.transfer_incoming_count_limit
-            and self.transfer_incoming_bytes
-            >= self.transfer_incoming_bytes_throttle_threshold
-        ):
+        if self._should_throttle_incoming_transfers():
             return {}, []
 
         recommendations: Recs = {}
@@ -1489,7 +1513,12 @@ def _ensure_communicating(self, *, stimulus_id: str) -> RecsInstrs:
             to_gather_tasks, total_nbytes = self._select_keys_for_gather(
                 available_tasks
             )
-            assert to_gather_tasks
+            # We always load at least one task
+            assert to_gather_tasks or self.transfer_incoming_bytes
+            # ...but that task might be selected in the previous iteration of the loop
+            if not to_gather_tasks:
+                break
+
             to_gather_keys = {ts.key for ts in to_gather_tasks}
 
             logger.debug(
@@ -1531,11 +1560,7 @@ def _ensure_communicating(self, *, stimulus_id: str) -> RecsInstrs:
             self.in_flight_workers[worker] = to_gather_keys
             self.transfer_incoming_count_total += 1
             self.transfer_incoming_bytes += total_nbytes
-            if (
-                self.transfer_incoming_count >= self.transfer_incoming_count_limit
-                and self.transfer_incoming_bytes
-                >= self.transfer_incoming_bytes_throttle_threshold
-            ):
+            if self._should_throttle_incoming_transfers():
                 break
 
         return recommendations, instructions
@@ -1610,18 +1635,28 @@ def _select_keys_for_gather(
         """Helper of _ensure_communicating.
 
         Fetch all tasks that are replicated on the target worker within a single
-        message, up to transfer_message_target_bytes.
+        message, up to transfer_message_target_bytes or until we reach the limit
+        for the size of incoming data transfers.
         """
         to_gather: list[TaskState] = []
         total_nbytes = 0
 
+        if self.transfer_incoming_bytes_limit is not None:
+            bytes_left_to_fetch = min(
+                self.transfer_incoming_bytes_limit - self.transfer_incoming_bytes,
+                self.transfer_message_target_bytes,
+            )
+        else:
+            bytes_left_to_fetch = self.transfer_message_target_bytes
+
         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.transfer_message_target_bytes
-            ):
+                # When there is no other traffic, the top-priority task is fetched
+                # regardless of its size to ensure progress
+                self.transfer_incoming_bytes
+                or to_gather
+            ) and total_nbytes + ts.get_nbytes() > bytes_left_to_fetch:
                 break
             for worker in ts.who_has:
                 # This also effectively pops from available


Test patch:
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
@@ -3,6 +3,7 @@
 import asyncio
 import gc
 import pickle
+from collections import defaultdict
 from collections.abc import Iterator
 
 import pytest
@@ -1351,3 +1352,189 @@ def test_transfer_incoming_metrics(ws):
     assert ws.transfer_incoming_bytes == 0
     assert ws.transfer_incoming_count == 0
     assert ws.transfer_incoming_count_total == 4
+
+
+def test_throttling_does_not_affect_first_transfer(ws):
+    ws.transfer_incoming_count_limit = 100
+    ws.transfer_incoming_bytes_limit = 100
+    ws.transfer_incoming_bytes_throttle_threshold = 1
+    ws2 = "127.0.0.1:2"
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "c",
+            who_has={"a": [ws2]},
+            nbytes={"a": 200},
+            stimulus_id="s1",
+        )
+    )
+    assert ws.tasks["a"].state == "flight"
+
+
+def test_throttle_incoming_transfers_on_count_limit(ws):
+    ws.transfer_incoming_count_limit = 1
+    ws.transfer_incoming_bytes_limit = 100_000
+    ws.transfer_incoming_bytes_throttle_threshold = 1
+    ws2 = "127.0.0.1:2"
+    ws3 = "127.0.0.1:3"
+    who_has = {"a": [ws2], "b": [ws3]}
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "c",
+            who_has=who_has,
+            nbytes={"a": 100, "b": 100},
+            stimulus_id="s1",
+        )
+    )
+    tasks_by_state = defaultdict(list)
+    for ts in ws.tasks.values():
+        tasks_by_state[ts.state].append(ts)
+    assert len(tasks_by_state["flight"]) == 1
+    assert len(tasks_by_state["fetch"]) == 1
+    assert ws.transfer_incoming_bytes == 100
+
+    in_flight_task = tasks_by_state["flight"][0]
+    ws.handle_stimulus(
+        GatherDepSuccessEvent(
+            worker=who_has[in_flight_task.key][0],
+            data={in_flight_task.key: 123},
+            total_nbytes=100,
+            stimulus_id="s2",
+        )
+    )
+    assert tasks_by_state["flight"][0].state == "memory"
+    assert tasks_by_state["fetch"][0].state == "flight"
+    assert ws.transfer_incoming_bytes == 100
+
+
+def test_throttling_incoming_transfer_on_transfer_bytes_same_worker(ws):
+    ws.transfer_incoming_count_limit = 100
+    ws.transfer_incoming_bytes_limit = 250
+    ws.transfer_incoming_bytes_throttle_threshold = 1
+    ws2 = "127.0.0.1:2"
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "d",
+            who_has={"a": [ws2], "b": [ws2], "c": [ws2]},
+            nbytes={"a": 100, "b": 100, "c": 100},
+            stimulus_id="s1",
+        )
+    )
+    tasks_by_state = defaultdict(list)
+    for ts in ws.tasks.values():
+        tasks_by_state[ts.state].append(ts)
+    assert ws.transfer_incoming_bytes == 200
+    assert len(tasks_by_state["flight"]) == 2
+    assert len(tasks_by_state["fetch"]) == 1
+
+    ws.handle_stimulus(
+        GatherDepSuccessEvent(
+            worker=ws2,
+            data={ts.key: 123 for ts in tasks_by_state["flight"]},
+            total_nbytes=200,
+            stimulus_id="s2",
+        )
+    )
+    assert all(ts.state == "memory" for ts in tasks_by_state["flight"])
+    assert all(ts.state == "flight" for ts in tasks_by_state["fetch"])
+
+
+def test_throttling_incoming_transfer_on_transfer_bytes_different_workers(ws):
+    ws.transfer_incoming_count_limit = 100
+    ws.transfer_incoming_bytes_limit = 150
+    ws.transfer_incoming_bytes_throttle_threshold = 1
+    ws2 = "127.0.0.1:2"
+    ws3 = "127.0.0.1:3"
+    who_has = {"a": [ws2], "b": [ws3]}
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "c",
+            who_has=who_has,
+            nbytes={"a": 100, "b": 100},
+            stimulus_id="s1",
+        )
+    )
+    tasks_by_state = defaultdict(list)
+    for ts in ws.tasks.values():
+        tasks_by_state[ts.state].append(ts)
+    assert ws.transfer_incoming_bytes == 100
+    assert len(tasks_by_state["flight"]) == 1
+    assert len(tasks_by_state["fetch"]) == 1
+
+    in_flight_task = tasks_by_state["flight"][0]
+    ws.handle_stimulus(
+        GatherDepSuccessEvent(
+            worker=who_has[in_flight_task.key][0],
+            data={in_flight_task.key: 123},
+            total_nbytes=100,
+            stimulus_id="s2",
+        )
+    )
+    assert tasks_by_state["flight"][0].state == "memory"
+    assert tasks_by_state["fetch"][0].state == "flight"
+
+
+def test_do_not_throttle_connections_while_below_threshold(ws):
+    ws.transfer_incoming_count_limit = 1
+    ws.transfer_incoming_bytes_limit = 200
+    ws.transfer_incoming_bytes_throttle_threshold = 50
+    ws2 = "127.0.0.1:2"
+    ws3 = "127.0.0.1:3"
+    ws4 = "127.0.0.1:4"
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "b",
+            who_has={"a": [ws2]},
+            nbytes={"a": 1},
+            stimulus_id="s1",
+        )
+    )
+    assert ws.tasks["a"].state == "flight"
+
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "d",
+            who_has={"c": [ws3]},
+            nbytes={"c": 1},
+            stimulus_id="s2",
+        )
+    )
+    assert ws.tasks["c"].state == "flight"
+
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "f",
+            who_has={"e": [ws4]},
+            nbytes={"e": 100},
+            stimulus_id="s3",
+        )
+    )
+    assert ws.tasks["e"].state == "flight"
+    assert ws.transfer_incoming_bytes == 102
+
+
+def test_throttle_on_transfer_bytes_regardless_of_threshold(ws):
+    ws.transfer_incoming_count_limit = 1
+    ws.transfer_incoming_bytes_limit = 100
+    ws.transfer_incoming_bytes_throttle_threshold = 50
+    ws2 = "127.0.0.1:2"
+    ws3 = "127.0.0.1:3"
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "b",
+            who_has={"a": [ws2]},
+            nbytes={"a": 1},
+            stimulus_id="s1",
+        )
+    )
+    assert ws.tasks["a"].state == "flight"
+
+    ws.handle_stimulus(
+        ComputeTaskEvent.dummy(
+            "d",
+            who_has={"c": [ws3]},
+            nbytes={"c": 100},
+            stimulus_id="s2",
+        )
+    )
+    assert ws.tasks["c"].state == "fetch"
+    assert ws.transfer_incoming_bytes == 1