Dask compute uses twice the expected memory

Peter9192 0
Viewed 35

I'm trying to understand and optimize the memory usage of Dask. For this simple example, I create an array that I would expect would occupy ~1GB of memory. When I realize the data with persist() this is what I find. But with compute() I see the amount of memory double.

Is this correct? And if so, why?

import dask.array as da
import resource

mem0 = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
memory = lambda: (resource.getrusage(resource.RUSAGE_SELF).ru_maxrss - mem0) / 1024

data = da.random.random((1024, 1024, 128))  # ~1 GB, not yet realized
print(f"Memory used: {memory()} MB")
# Memory used: 0.49609375 MB

data.persist()
print(f"Memory used: {memory()} MB")
# Memory used: 1023.49609375 MB

# Using .compute() instead of .persist()
# Memory used: 2047.5 MB

NB: I did not run persist and compute in the same script. I actually changed the code and ran again, to have a fair comparison.

dask python
1 Answers

When you call dask.Array.compute(), dask completes the task on the remote worker(s), then serializes the object, passes it back to the main thread, and then deserializes it. Therefore, there is likely to be one copy of the data on the workers, a copy being passed around as a string, and the version recreated on the main thread. Peak memory usage was likely more than 2x; at any given point the remote and partial local versions are in memory, plus copies of chunks in transit.

See the dask.distributed docs on managing memory. In the section on clearing data:

We remove data from distributed RAM by removing the collection from our local process. Remote data is removed once all Futures pointing to that data are removed from all client machines.

del df  # Deleting local data often deletes remote data

If this is the only copy then this will likely trigger the cluster to delete the data as well.

However, if we have multiple copies or other collections based on this one, then we’ll have to delete them all.

df2 = df[df.x < 10]
del df  # would not delete data, because df2 still tracks the futures

In your example, you never delete the reference to the remote object, meaning it must still be held on the remote cluster:

# brings the data to the local thread while keeping the remote reference
data.compute()

If instead you overwrote the remote reference, the copy on your workers could at least be dropped once the call to compute() was finished:

# drops the remote reference after bringing the data locally
data = data.compute()
Michael Delgado0
Related