Mild memory leak in dask workers

[rabernat]

Describe the issue:

I have been observing a mild but consistent increase in memory when storing large arrays to Zarr. I have reproduced this with both s3fs (my real use case) and with a dummy "Dev Null" store of my own devising.

I would expect to be able to store arrays of essentially infinite size in streaming fashion. Instead, this memory leak means that eventually I will run out of memory.

I'm aware that the root of this issue may be hard to diagnose. The ultimate cause may be upstream, in Zarr, where I am a maintainer. Very happy to work with the developers here to isolate and resolve the underlying issue. 🙏

I do not know if this issue occurs with other schedulers, as I don't know how to diagnose memory usage as conveniently as with distributed.

Minimal Complete Verifiable Example:

import dask
import dask.array as da
from dask.distributed import LocalCluster
from distributed.diagnostics import MemorySampler
import zarr


class DevNullStore(dict):
    """Dummy store in which data just vanishes."""

    def __setitem__(self, key, value):
        # only store attributes
        if key == '.zarray':
            super().__setitem__(key, value)
            
            
shape = 10_000_000
chunks = 1_000
data = da.zeros(shape, chunks=chunks)

cluster = LocalCluster(
    n_workers=8, threads_per_worker=1, host="*",
)

ms = MemorySampler()
store = zarr.storage.KVStore(DevNullStore())
with cluster.get_client() as client:
    with ms.sample("nullstore-zarr_v2-compute_true"):
        data.to_zarr(store, compressor=None)
        
ms.plot(align=True)

As you can see, memory usage increases steadily over the course of the computation. If you make the array larger, the trend continues.

Anything else we need to know?:

A likely objection to this issue might be "this is an artifact of your funky DevNullStore". However, the same behavior can be reproduced (albeit much more slowly) with s3fs. This example requires write access to s3

import s3fs
import uuid

# replace with a bucket you can write to
s3_url = "s3://earthmover-rechunk-tmp/memory-leak-test"

def make_s3_store():
    url = f'{s3_url}/{uuid.uuid4().hex}'
    store = zarr.storage.FSStore(url=url)
    return store

with cluster.get_client() as client:
    client.restart()
    with ms.sample("s3fs-zarr_v2-compute_true"):
        data.to_zarr(make_s3_store(), compressor=None)
        
ms.plot(align=True)

As you can see, the overall magnitude of the memory leak is similar: about 200 MB. This suggests that it is independent of the underlying store.

Environment:

• Dask version: 2023.7.0
• Zarr version: 2.16.1
• s3fs version: 2023.9.0
• Python version: 3.10.12
• Operating System: Linux
• Install method (conda, pip, source): conda + pip

cc @crusaderky, who helped us with an earlier iteration of this problem.

[mrocklin]

Thanks @rabernat for the well crafted example. I suspect that folks will take a look tomorrow.

[rabernat]

This issue would have been so much more difficult to describe without MemorySampler. What a great tool! 🙌

[fjetter]

I strongly suspect this is due to internal state keeping. We keep logs about internal state transitions/events and even store some actual log messages. All of these logs are bound to a maximum size and most (unfortunately not all) are configurable.

In your example, 200MB on 8 workers, i.e. 25MB each sounds about right for this kind of logging.

For instance, we have Worker.log which holds task transition logs. An entry of this log could look like

("zeroes_like-store-map-4567890", "executing", "memory", "memory", {}, "task-finished-34567898765", 1234567890.098765)

and this log deque is bound to 100_000 entries. (@crusaderky the worker transition log should likely also have a configurable transition-log-length)

In [6]: format_bytes(sizeof(("zeroes_like-store-map-4567890", "executing", "memory", "memory", {}, "task-finished-34567898765", 1234567890.098765)) * 100_000)
Out[6]: '48.07 MiB'

so that's about 50MB. Once it reaches this size, we stop storing more stuff in there (or rather throw out the oldest entries). You probably haven't reached this limit, yet.

This is a known problem for the scheduler (typically it does not matter) see #5570 or #1795

The only config that affects the worker is currently distributed.admin.log-length which you can set to zero s.t. ordinary python logs are not kept in memory.

We'll look into this and if it's indeed the transition counter we'll adjust the log length and make it configurable.

[mrocklin]

@rabernat have you actually run into problems where workers run out of memory, or are you seeing the increase and afraid that it will continue forever?

[rabernat]

@rabernat have you actually run into problems where workers run out of memory

No I have not. This emerged from a broader investigation into memory usage. By the time we made the graph big enough to really cause problems, we would likely first hit other known scaling limitations due to the number of tasks. It would be perfectly reasonable to say, "yes, there is a small memory leak, but it doesn't affect any real-world workflows in practice," and prioritize this issue accordingly.

[mrocklin]

we would likely first hit other known scaling limitations due to the number of tasks

Which limitation did you run into, I'm curious? (hearing about these again is useful to help prioritize them).

[rabernat]

This graph had 10_000 tasks and leaked 200 MB of memory. If it had 100_000 tasks, it would leak 2GB of memory. However, at this point, our graph would start becoming large, in violation of the best practices. Symptoms of this would be very slow serialization time and high memory load on the scheduler.

[crusaderky]

Reproduced.

This is just the core behaviour of dask workers - nothing to do with zarr specifically.

import asyncio
import dask.bag as db
from distributed.diagnostics import MemorySampler
from distributed import Client, wait

WIDTH = 80_000
NITER = 10
data = db.from_sequence(range(WIDTH), npartitions=WIDTH)

client = Client(n_workers=16, threads_per_worker=2)

async def wait_cleanup(dask_worker):
    while dask_worker.state.tasks:
        await asyncio.sleep(0.1)

ms = MemorySampler()
with ms.sample():
    for _ in range(NITER):
        wait(data.persist())
        client.run(wait_cleanup)

ms.plot(grid=True, align=True)

After forcing all deques to maxlen=0 - both the configurable ones and the hardcoded ones - the plot changes as follows:

as you can see, there's around 25~30 MiB worth of deques per worker, plus a genuine memory leak of roughly 400 bytes per task (2.3 GiB to 2.6 GiB over 80k tasks x 10 runs) - which is substantial in the long run.

Pending actions

• ensure that all deques have configurable maxlen
• simplify deque length config. There are currently 8 different settings in the config which feels like way too many. How would people feel if I reduced some deques (namely, transaction logs) from 100k entries to 10k?
  Ideally I would love to have a single setting that controls them all, and for it to be the first thing that we tell users to switch off when they walk in complaining about leaks.
• investigate the remaining leak

[fjetter]

How would people feel if I reduced some deques (namely, transaction logs) from 100k entries to 10k?

Should be fine. I would even think zero is a good default since most users don't care about this. In the rare case where we instruct folks to create dumps we can also tell them to set these numbers

[crusaderky]

50 runs, no deques

[crusaderky]

I've reran the test, downstream of #8173 (log-length=0), this time without the MemoryMonitor.

import asyncio
import gc
import dask.bag as db
from dask.utils import format_bytes
from distributed.diagnostics import MemorySampler
from distributed import Client, wait

WIDTH = 80_000
data = db.from_sequence(range(WIDTH), npartitions=WIDTH)

client = Client(n_workers=16, threads_per_worker=2)

async def wait_cleanup(dask_worker):
    while dask_worker.state.tasks:
        await asyncio.sleep(0.1)

    gc.collect()
    await asyncio.sleep(0.25)
    gc.collect()
    await asyncio.sleep(0.25)
    return dask_worker.monitor.get_process_memory()

measures = []
for _ in range(100):
    wait(data.persist())
    mems = client.run(wait_cleanup)
    total = sum(mems.values())
    measures.append(total)

import pandas

df = pandas.DataFrame({"rss": measures})
df["delta"] = df.rss.diff()
df = df.iloc[1:]
df.index *= 80_000 // 16
df["delta"] = df["delta"].astype(int)
df["delta per task"] = df["delta"] // 80_000
df["delta per task"].rolling(10).mean().plot(grid=True)

On the X axis, you have the number of tasks per worker since the start of the test.
On the Y axis, you have how many bytes worth of leak each task caused (rolling average).

>>> df["delta per task"].loc[200_000:].mean()
28.661290322580644

So while there is still a leak in the long term, it's less than 30 bytes per task, which I don't think we should spend time investigating as the only users impacted are those that keep the same workers online for weeks at a time running flat out.