Ufdjrw

I have a large data problem. The specific problem isn't super important, but I've solved it with dask. Now I have two problems.

from dask import distributed
import numpy as np

local_cluster = distributed.LocalCluster(n_workers=20, processes=True, memory_limit=0)
dask_client = distributed.Client(local_cluster)

hat_matrix = np.random.rand(1000,25000)
weight_matrix = np.random.rand(1000)
Y = np.random.rand(1000, 25000)

[scatter_hat] = dask_client.scatter([hat_matrix], broadcast=True)
[scatter_weight] = dask_client.scatter([weight_matrix], broadcast=True)

futures = [dask_client.submit(apply_function, i, scatter_hat, Y[i, :], scatter_weight)
 for i in range(Y.shape[0])]

results = dask_client.gather(futures)

I can split up Y (which is good, because I don't have enough memory to really load it all at once), but all of my workers need hat_matrix. Scattering hat_matrix and then sending Y row-wise works great. Except hat_matrix and Y are both... large, which is is fine. I have enough memory provisioned to deal with it. But I can't find any way to allow for short memory spikes (which occur during deserialization), so if I set a memory limit the nanny kills all my workers. Then all my new workers. And so on and so forth. So I have three questions:

Is there a way to set a memory limit that permits spikes as serialized data comes in and is unpacked? If I have 64 GB of memory to drive 20 processes, I'd like to set a memory limit of, say, 2.8GB per process. When I scatter 2GB of data, there's a spike to ~4GB per process for deserializing, and the nanny kills everything.

Is there a way to stagger scattering to minimize the transient memory spike?

Is there a convenient way to scatter data by disk instead of through TCP or do I have to custom write that?
(As a corollary: Is there a convenient way to load a memory-mapped dask array from a serialized file in all of my workers?)

Is there a way to set a memory limit that permits spikes as serialized data comes in and is unpacked?

In general deserialization runs arbitrary code, so dask can't really control what happens. In practice though your matrices aren't really that large though compared to typical memory available on modern hardware though, so I'm surprised that you're running into an issue. Dask is decently careful with NumPy arrays. I wouldn't expect it to use much more memory than the size of the array.

Is there a way to stagger scattering to minimize the transient memory spike?

Scattering currently goes through a broadcast tree. Your client sends to a few workers, then they send to a few more workers, and so on. By default the branching factor here is only two, so I would be surprised to see a huge blowup here.

Is there a convenient way to scatter data by disk instead of through TCP or do I have to custom write that? (As a corollary: Is there a convenient way to load a memory-mapped dask array from a serialized file in all of my workers?)

Perhaps you could use a memory mapped NumPy array of some sort rather than an in-memory NumPy array?