Scheduler memory leak / large worker footprint on simple workload

[chrisroat]

What happened:

Running an embarrassingly parallel map_overlap workload may be causing a memory leak in the scheduler. Upon completion, releasing the tasks and restarting the client does not reclaim the memory. The example below, with 200k tasks, shows a jump in scheduler memory from 100MB to 1.3GB while running the graph. After client.restart, it remains at 1.1GB.

In addition, the memory of the workers climb into the "yellow", where I believe swapping to disk begins to happen. Given the parallel nature of this workload, workers ought to be able discard pieces when they are done with them.

From a performance perspective, during client.compute, the scheduler gets unresponsive (it takes 20ish seconds to start), presumably because its loading a large graph. I have seen this cause already running computations to start erroring. I've seen lost keys and KilledWorkers.

And finally, anecdotally, it sometimes happens that one worker runs hot, getting 10x the tasks of other workers. Eventually, forward progress halts. I now watch for this, and then kill that worker, which redistributes the work and finishes the job. (I'm using dask-gateway on K8s).

What you expected to happen:

• The scheduler should not use up additional memory once a computation is done.
• Workers should shard a parallel job so that each shard can be discarded when done, keeping a low worker memory profile
• Loading a graph should not disrupt ongoing computation

Minimal Complete Verifiable Example:

import dask.array as da
import distributed
client = distributed.Client(n_workers=4, threads_per_worker=1, memory_limit='10GB')
arr = da.zeros((50, 2, 8192, 8192), chunks=(1, -1, 512, 512))
result = arr.map_overlap(lambda x: x, depth=(0,0,200,200))
store = result.to_zarr('/media/ssd/test/memory_scheduler.zarr', compute=False, overwrite=True)
future = client.compute(store)

Environment:

• Dask version: 2.18.1
• Distributed version: 2.18.0
• Python version: 3.7
• Operating System: Ubuntu 18.04
• Install method (conda, pip, source): conda

[chrisroat]

I also see a lot of this in the logs

distributed.worker - WARNING - Memory use is high but worker has no data to store to disk.  Perhaps some other process is leaking memory?  Process memory: 7.03 GB -- Worker memory limit: 10.00 GB

[mrocklin]

For memory issues on the worker I recommend that you look through dask/dask#3530

20ish seconds for 200k tasks is a bit long, but not terribly. Task overhead is in the range of 200us per task, so this is in line with expectations. There are efforts to accelerate this, but they're probably a few months out at best.

I suspect that your minimal example is intentionally simplified, but your chunk sizes are quite small. A good way to avoid task overhead is to have fewer larger tasks.

I don't know why your scheduler memory would stay at 1GB after a restart, except perhaps that your OS has decided to hold onto the memory for a while (this is common).

[chrisroat]

Thanks for following up. I did start at #3530, but was advised that since I was using a single-threaded scheduler, the problem there was not mine.

The large-graph problem is partially the chunk size, but also because overlap computations are part of the mix.

I can use larger chunks for some things, but this is one use case I have where smaller chunks are required. I'm aligned images, and the alignment varies across the larger image. Basically, I'm micro aligning 512x512 chunks. I suppose I could go the inception path: a map_overlap within a map_overlap using scheduler=synchronous on the inside.

As an aside, my real chunks are more like ~5x200x512x512 uint16, which is ~500MB. But I wanted an example that would be the same number of tasks, but quickly do-able as a self contained example.

I'm still curious as to why a purely parallel computation (other than the overlap) is not proceeding in a way that minimizes memory consumption.... it seems like it's trying to read the whole thing into all the workers memory at once.

How long could the OS hold onto things? I run both locally and on GCP on Ubuntu images, and I've never seen it go away. Each time I add a graph, it eats up more space, so it doesn't seem to be re-using a heap or anything like that. I run my dask-gateway scheduler with 15GB of memory, but generally restart before I hit that because things get touchy. (I had to turn off autoscaling, because including that on top of this seemed to bog the scheduler.)

[mrocklin]

Thanks for following up. I did start at #3530, but was advised that since I was using a single-threaded scheduler, the problem there was not mine.

To be clear, in the example above you're using the distributed scheduler. When we talk about the single-threaded scheduler we mean when you don't create a client, and you call .compute(scheduler="single-threaded")

But I wanted an example that would be the same number of tasks, but quickly do-able as a self contained example

I figured, and I appreciate it.

I'm still curious as to why a purely parallel computation (other than the overlap) is not proceeding in a way that minimizes memory consumption.... it seems like it's trying to read the whole thing into all the workers memory at once.

It should. If you're interested you can try x.visualize(color="order") and get out a pretty picture. This only works on graphs that have less than, say 500 tasks or so.

How long could the OS hold onto things?

Some OS's hold onto memory until someone else needs it. If your machine is properly running out of RAM though then that's not this.

[mrocklin]

In short, I personally don't know what's going on. I would probably have to dive in more deeply to figure out what the problem is, but it's unlikely that I'll have time to do that any time in the near future.

[quasiben]

Instead of client.restart can you call client.cancel(future) at the end of the computation ? This is usually where I start when I want to make sure everything is gone.

[chrisroat]

I do both - I cancel all futures, then restart the client. The memory still persists.

The color="order" is super neat. With optimize_graph=True, collapse_outputs=True, it is pretty fun to use.

One thing I've noticed is that the designations for overlap tasks have fractions. Is there any comparison of floating points that might happen?

From playing around with a few small examples, the orderings look fine to me. Is order fully respected globally in some way (or even just on a worker)? I could imagine that if some task is waiting due to some upstream bottleneck, later-ordered but read-to-go tasks could be processed. For example, data can always be read in and consume memory.

And with overlap tasks, is it safe to assume that just the necessary neighbor data is sliced before being transferred (and a full transfer of the neighbor chunks is not done)?

[mrocklin]

Ordering is the primary mechanism to determine the order of tasks, but as you point out the workers may choose to do other things if a task is blocked on something like communication. Yes, data is sliced before being transferred. Hopefully we only move around small pieces.

…

On Wed, Jun 17, 2020 at 1:43 PM Chris Roat ***@***.***> wrote: I do both - I cancel all futures, then restart the client. The memory still persists. The color="order" is super neat. With optimize_graph=True, collapse_outputs=True, it is pretty fun to use. One thing I've noticed is that the designations for overlap tasks have fractions. Is there any comparison of floating points that might happen? From playing around with a few small examples, the orderings look fine to me. Is order fully respected globally in some way (or even just on a worker)? I could imagine that if some task is waiting due to some upstream bottleneck, later-ordered but read-to-go tasks could be processed. For example, data can always be read in and consume memory. And with overlap tasks, is it safe to assume that just the necessary neighbor data is sliced before being transferred (and a full transfer of the neighbor chunks is not done)? — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#3898 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTGYWQJVXI5G655TSOTRXETFNANCNFSM4N6LPGGA> .

[TomAugspurger]

You'll want to double check, but I think bytes stored refers to the contents of `Worker.data`, which might include data that has been spilled to disk.

…

On Mon, Jun 22, 2020 at 6:31 PM Chris Roat ***@***.***> wrote: Not directly related to this topic, but related to the "read in data" task running: I've set DASK_DISTRIBUTED__WORKER__MEMORY__TARGET to 0.3 on a 21GB worker. The bokeh dashboard indicates workers have 16-22GB of "bytes stored", and the GCP dashboard for one worker shows 14GB used. Does "bytes stored" equate to memory, and if so, under what conditions would that go above the memory target? — You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHub <#3898 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAKAOIVG4TL4CGFWC3NR7R3RX7SVFANCNFSM4N6LPGGA> .

[chrisroat]

I was curious if any of the leak was fractional keys, and found there is some code in the overlap module that is being reported as leaking. Here is the top-20 lines reported, which totals ~400MB (of 1GB growth I see), within dask/distributed 2.19.0 (Tagging @jakirkham from the other thread)

Perhaps those knowledgable of dask internals can verify whether tracemalloc has found something real, and what possible mitigations would be.

Report:

Top 20 lines
#1: protocol/core.py:199: 32271.9 KiB
    return msgpack.loads(payload, use_list=False, **msgpack_opts)
#2: array/overlap.py:35: 26246.2 KiB
    rounded = (task[0],) + tuple(int(round(i)) for i in task[1:])
#3: array/overlap.py:112: 16850.3 KiB
    seq = list(product(*args))
#4: array/overlap.py:57: 16536.8 KiB
    if all(ind == slice(None, None, None) for ind in index):
#5: distributed/scheduler.py:2057: 15708.7 KiB
    prefix_key = key_split(key)
#6: array/overlap.py:48: 12670.9 KiB
    index.append(slice(None, None, None))
#7: <frozen importlib._bootstrap_external>:525: 11340.3 KiB
#8: distributed/scheduler.py:4723: 11140.5 KiB
    self.transition_log.append((key, start, finish2, recommendations, time()))
#9: distributed/scheduler.py:4744: 10847.7 KiB
    plugin.transition(key, start, finish2, *args, **kwargs)
#10: array/optimization.py:121: 10240.1 KiB
    dsk = dsk.copy()
#11: array/overlap.py:145: 9641.2 KiB
    frac_slice = fractional_slice((x.name,) + k, axes)
#12: distributed/scheduler.py:4713: 9206.8 KiB
    a = a.copy()
#13: array/overlap.py:52: 8693.4 KiB
    index.append(slice(-left_depth, None))
#14: array/optimization.py:191: 8232.2 KiB
    dsk[k] = (get, a, a_index)
#15: array/overlap.py:50: 7290.1 KiB
    index.append(slice(0, right_depth))
#16: distributed/scheduler.py:2055: 6600.3 KiB
    ts = TaskState(key, spec)
#17: distributed/scheduler.py:3949: 5815.5 KiB
    ts.waiters = {dts for dts in ts.dependents if dts.state == "waiting"}
#18: distributed/scheduler.py:655: 5800.7 KiB
    self.who_has = set()
#19: distributed/scheduler.py:653: 5800.3 KiB
    self.waiting_on = set()
#20: distributed/scheduler.py:650: 5800.3 KiB
    self.who_wants = set()
15712 other: 161295.1 KiB
Total allocated size: 398029.4 KiB

Code:

import tracemalloc
import linecache
import gc
import os
import tracemalloc
import dask.array as da
import distributed


tracemalloc.start()

client = distributed.Client(n_workers=4, threads_per_worker=1, memory_limit='10GB')
print(client)

arr = da.zeros((50, 2, 8192, 8192), chunks=(1, -1, 512, 512))
result = arr.map_overlap(lambda x: x, depth=(0,0,200,200))
result.to_zarr('/media/ssd/test/memory_scheduler.zarr', overwrite=True)

client.restart()  # Don't wait.  Stop immediately, reset, and see if everything is at baseline.
gc.collect()  # Really!  We want everything to go back to baseline.
# Look at System display of client -- it shows >1GB of memory usage.  Baseline is ~100MB


# Pretty print function from tracemalloc webpage
def display_top(snapshot, key_type='lineno', limit=20):
    snapshot = snapshot.filter_traces((
        tracemalloc.Filter(False, "<frozen importlib._bootstrap>"),
        tracemalloc.Filter(False, "<unknown>"),
    ))
    top_stats = snapshot.statistics(key_type)

    print("Top %s lines" % limit)
    for index, stat in enumerate(top_stats[:limit], 1):
        frame = stat.traceback[0]
        # replace "/path/to/module/file.py" with "module/file.py"
        filename = os.sep.join(frame.filename.split(os.sep)[-2:])
        print("#%s: %s:%s: %.1f KiB"
              % (index, filename, frame.lineno, stat.size / 1024))
        line = linecache.getline(frame.filename, frame.lineno).strip()
        if line:
            print('    %s' % line)

    other = top_stats[limit:]
    if other:
        size = sum(stat.size for stat in other)
        print("%s other: %.1f KiB" % (len(other), size / 1024))
    total = sum(stat.size for stat in top_stats)
    print("Total allocated size: %.1f KiB" % (total / 1024))


snapshot = tracemalloc.take_snapshot()   
display_top(snapshot)

[chrisroat]

I re-tested this in 2.25.0 and found the memory increase on the scheduler (after reset) is now even higher - 2.2GB.

The code above just shows possible leaks in the main process, which is roughly the same. I'm not sure yet how to figure out where the scheduler memory is going.

My workaround is to bring up a new cluster for each calculation.

[mrocklin]

Thanks for keeping up with this @chrisroat . I'm not sure I fully understand the printout that you gave a few months ago. What do the sizes mean there? Data that is no longer being tracked by Python? If so, I'm curious how this is being discovered. Some of those lines, like set(), seem fairly innocuous.

It's my first time seeing data like this. Can you help interpret it?

[chrisroat]

Those are coming from tracemalloc, which is supposed to detect memory leaks.

I'm realizing now, however, that the 2GB scheduler memory increase is likely not traced here, as it's likely a child process. I think memory-profiler can profile child processes, but I haven't tried that yet.

[mrocklin]

The scheduler is run in the same process by default with the local cluster. I guess I'm surprised that finding memory leaks is that easy. I'll have to learn more about how tracemalloc works.

…

On Tue, Sep 1, 2020, 7:19 PM Chris Roat ***@***.***> wrote: Those are coming from tracemalloc <https://docs.python.org/3/library/tracemalloc.html>, which is supposed to detect memory leaks. I'm realizing now, however, that the 2GB scheduler memory increase is likely not traced here, as it's likely a child process. I think memory-profiler can profile child processes, but I haven't tried that yet. — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#3898 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AACKZTEVQD5DQ5EHG2W6F33SDWTTLANCNFSM4N6LPGGA> .