Worker can memory overflow by fetching too much data at once

[fjetter]

There is a possibility for a worker to suicide by fetching too much data at once

Old version:

distributed/distributed/worker.py

Lines 2754 to 2757 in 84cbb09

	while self.data_needed and (
	len(self.in_flight_workers) < self.total_out_connections
	or self.comm_nbytes < self.comm_threshold_bytes
	):

New version:

distributed/distributed/worker_state_machine.py

Lines 1425 to 1429 in 3647cfe

	if (
	len(self.in_flight_workers) >= self.total_out_connections
	and self.comm_nbytes >= self.comm_threshold_bytes
	):
	return {}, []

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

[gjoseph92]

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 #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 #5996?

[crusaderky]

pseudocode:

if current managed_in_memory + self.comm_nbytes > threshold% * max_memory:
    stop scheduling further gather_dep's

or with better encapsulation:

    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:

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.

[mrocklin]

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.

[crusaderky]

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 .

[mrocklin]

Ah, great then.

[hendrikmakait]

In coiled/benchmarks#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/benchmarks#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

[hendrikmakait]

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

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