Dask-cuda: Spill device memory to host memory or disk

[mrocklin]

Currently Dask spills memory to disk once it feels some pressure. Docs, Code

It would be useful to build another system like this that separately tracks device and host memory, and spills from one to the other and then to disk as needed. This is tricky because the worker can't special case different types ("oh, this is a cudf/cupy object, we should treat it as device memory") but will instead need to figure out some nice dispatching system.

@sklam I think that you had a solution to this at one point I think. Can you share? It would be helpful to see other implementations.

[sklam]

I was working on the branch at https://github.com/sklam/distributed/tree/enh/gpu_spill_to_disk. I don't remember what status it was in. I doubt it is fully working. At the very least, it does show what needs to be changed.

[mrocklin]

Thanks @sklam!

On Thu, Mar 28, 2019 at 8:15 AM Siu Kwan Lam [e-mail]
wrote:

I was working on the branch at
https://github.com/sklam/distributed/tree/enh/gpu_spill_to_disk. I don't
remember what status it was in. I doubt it is fully working. At the very
least, it does show what needs to be changed.

—
You are receiving this because you authored the thread.
Reply to this email directly, view it on GitHub
https://github.com/rapidsai/dask-cuda/issues/30#issuecomment-477639510,
or mute the thread
https://github.com/notifications/unsubscribe-auth/AASszNPai5hag46_o_2fFKuUDdxvZKprks5vbNyegaJpZM4cOV_0
.

[mrocklin]

OK, so here is a plan.

Dask workers place Python objects into Worker.data, a MutableMapping class (like a dictionary).

self.data[key_1] = my_numpy_array
self.data[key_2] = my_cudf_dataframe

This object handles all of the logic around when it should hold onto data, and when it should dump things to disk. Currently, it uses a library called zict to build a two-level LRU dictionary that keeps data in memory and then drops least recently used data to disk.

That code is here:

https://github.com/dask/distributed/blob/af1cbdb9c1b8dd56770b0e494f28eb808cdcf977/distributed/worker.py#L414-L428

That uses the Buffer class, which builds up this two-level LRU mapping. We want something similar, but slightly different. The most straightforward approach, I think, is the following:

device <-> host <-> disk

When data comes in we check to see if it is device-y, perhaps by checking for a __cuda_array_interface__ property, and if so we put it into the device LRU, otherwise we put it into the host LRU. The callback when evicting data from the device LRU places it into the host LRU. The callback evicting data from the host LRU places it into a third File dict.

from zict import LRU, Func, File, ZictBase

class DeviceHostFile(ZictBase):
    def __init__(...):
        self.device = LRU(total_device_memory, {}, weight=sizeof, on_evict=[self.push_data_to_host])
        self.host = LRU(total_host_memory, {}, weight=sizeof, on_evict=[self.push_data_to_disk])
        self.disk = Func(serialize_bytelist,
                           deserialize_bytes,
                           File(path))

    def __setitem__(self, key, value):
        if is_device_backed_object(value):  # perhaps check __cuda_array_interface__ here ?
            self.device[key] = value
        else:
            self.host[key] = value

    ...

There is plenty of work around tracking what things we've moved where, moving them back when necessary, moving other things to make space for them, and so on. Hopefully we can encapsulate all of this logic into this MutableMapping abstraction so that Dask doesn't need to think about it. When it's ready, we should be able to create a worker with ...

worker = Worker(data=DeviceHostDisk)
or 
worker = Worker(data=(DeviceHostDisk, {'host_memory': '100GiB', 'device_memory': '32GiB', ...}))

It may also be that the zict stuff is more trouble than it's worth. Figuring out the abstractions created by someone else is never fun. However, at the very least those solutions are decently solid today, and may provide some inspiration, even if folks want to make something custom here.

Serialization

We'll need to convert Python-objects-representing-device data into Python-objects-representing-host-data. One approach would be to convert cupy arrays to numpy arrays, cudf dataframes to pandas dataframes, and so on. I think that we should not do this. I think that it will be too much code and we'll have to extend it every time we want to support a new project.

Instead, as we do today for the host-disk solution, lets reuse our network serialization code. We already have to write code to efficiently serialize various objects into bytes-like objects for network communcation, so lets reuse that here. We'll need these functions

• https://distributed.dask.org/en/latest/serialization.html#distributed.protocol.serialize.serialize
• https://distributed.dask.org/en/latest/serialization.html#distributed.protocol.serialize.deserialize

Additionally, I think that we'll want to grab some implementations of these functions for common types (cupy, cudf, and so on) that are being developed as part of the UCX work.

Future work

Depending on how fast we can compress things, we might also want to consider compressing data on the device, or compressing it before we move from device to host.

Depending on relative bandwidths between device<->host and device<->disk we might also consider bypassing host memory altogether and dumping directly to disk.

cc @pentschev

[pentschev]

@mrocklin your proposal looks good to me, at least I don't see at the moment any flaws with that. Thanks for detailing it so well!

Regarding the _Future Work_ section only, I have a few comments.

Depending on how fast we can compress things, we might also want to consider compressing data on the device, or compressing it before we move from device to host.

According to my experience, compressing data can be quite time-consuming, so in general we could have that functionality given that we also allow the user to enable/disable it at will, instead of enforcing it by default.

Depending on relative bandwidths between device<->host and device<->disk we might also consider bypassing host memory altogether and dumping directly to disk.

How would we bypass host memory and dump data directly to disk? We could probably do that with GPUDirect in case there's another PCIe device connected to a SAN, for example. However, for simpler workstation setups, containing consumer-grade GPUs, GPUDirect isn't support, so it would require device memory to be spilled to disk through host memory. Do you have a different idea in mind here?

[pentschev]

One more question: how is Dask handling sliced arrays? Their pointers get adjusted to point to the beginning of the slice and its shape also gets adjusted, by default without creating a new array. How is Dask releasing data in those cases, the original array can't be fully released, as there may be other arrays referring to parts of the memory buffer. I believe this may be particularly challenging when asarray=False, which is currently necessary for all CuPy arrays wrapped in Dask arrays, for example.

[mrocklin]

Dask doesn't really think about this at all. Libraries like Numpy and CuPy handle their own memory management. I'm not very familiar with the policies that they use. Dask just holds onto Python objects until it doesn't need them, and then Python's GC kicks in.

[pentschev]

So the spilling to disk is only done for Dask partial results?

[pentschev]

Or does it write to disk even if the data can't be evicted from memory?

[mrocklin]

So the spilling to disk is only done for Dask partial results?

When you say something like x.compute() workers don't really know that you're trying to compute x. They only see the many small intermediate results that make up the computation of x. They run Python functions that produce NumPy arrays, CuPy arrays, Scikit-Learn models, and so on, and hold those objects in Worker.data. Dask workers don't know the broader objective that they're working torwards.

Currently they watch the amount of data they have in memory, both by estimating the size of each object and by watching OS-level process statistics. When they think that they're running low on memory they just start pushing some of the objects in the in-memory portion of Worker.data over to the on-disk portion of Worker.data. Hopefully this reduces the memory load, but Dask doesn't really have any direct control over that. The user code that we run can do anything, so we don't make much of an effort to reverse-engineer it.

[pentschev]

Ok, that was my intent saying "partial results", basically, intermediary computation done by the workers. So we strictly spill only Dask-created data to disk. That's much easier! Thanks for clarifying, @mrocklin !

[mrocklin]

FYI @kkraus14 you were asking about this

[kkraus14]

@mrocklin thanks for the ping and the great writeup of a plan above. From reading it one thing that comes to mind is that we should try to do this in a general way that allows for semi arbitrary tiered caching regardless of whether GPUs are in the system or not. For example, someone could possibly want:

host memory<->nvme<->ssd<->NFS

Completely not necessary for a first iteration, but designing a configurable LRU system that can be set up via a config file would allow for some really interesting things.

[mrocklin]

Yeah, that's partially the goal of the zict package. It becomes easier to build systems like that. My hope is that this task is easier because of the previous zict work, but that remains to be seen.

[kkraus14]

From looking at the zict package it looks doable, the questions that pop up is how does configuration of the zict Buffer objects get exposed to the end users. We could take a look at how systems like Alluxio: https://www.alluxio.org/docs/master/en/advanced/Alluxio-Storage-Management.html#multiple-tier-storage and Apache Crail: https://incubator-crail.readthedocs.io/en/latest/config.html#storage-tiers expose this type of configuration to end users.

[pentschev]

I was wondering now, since the mechanism for data eviction lies in https://github.com/dask/distributed/blob/b36e87d14e75c0cbdd0e3292ac65db2a836f74bd/distributed/worker.py#L2178-L2206, where would we want to have such a mechanism for device memory handling?

It seems to me that dask-cuda would be the most appropriate place to have all functionality related to device memory spilling, but it doesn't seem that there's a simple way to do it. The options I can think of now are:

1. Generalize the memory_monitor method to make no distinctions of memory types, but only treat them as levels (potentially a good solution for @kkraus14 wish to have several layers like host memory<->nvme<->ssd<->NFS); or
2. Allow passing a function to override the internal memory_monitor method, but may require additional information on Worker internal states.

Any other ideas to do this differently or better?

[mrocklin]

The memory_monitor code is fairly old and has grown organically. I wouldn't be surprised if there were ways to clean it up and make it more general, as you suggest.

Another short-term idea would be to subclass Worker and write our own CUDA specific version of that function. No idea if that's better though.

[pentschev]

I think the bigger question is: do we want to subclass Worker? If we intend to have more functionality that will be specific to CUDA workers, then it may make sense to do it, it may even be necessary.

I'm also interested in reasons _not_ to solve it in this way or that way, to prevent doing something that we want to avoid. Any particular objections to any of the ideas suggested? Ideas so far:

1. Generalize the memory_monitor method to make no distinctions of memory types, but only treat them as levels (potentially a good solution for @kkraus14 wish to have several layers like host memory<->nvme<->ssd<->NFS); or
2. Allow passing a function to override the internal memory_monitor method, but may require additional information on Worker internal states; or
3. Another short-term idea would be to subclass Worker and write our own CUDA specific version of that function. No idea if that's better though.

If there are no objections, I will then re-evaluate them and see what I can do in a timely manner to ensure we have this feature usable soon.

[mrocklin]

If I were to do this I would subclass for now to get a solution on which we could iterate quickly. After we learn what we're doing we would then find ways to generalize and integrate back into the core.

My experience so far has been that we need to try many things on the CUDA side before we get it right. I'd prefer to experiment outside of the Dask codebase when possible, which moves a bit more slowly.

[pentschev]

This is a good point, seems like a good reason. I will do it like that then, thanks for chiming in!

[pentschev]

Since #51 is merged, I'm closing this. We still have #35 which will stay open for now, the difference is that the first includes only an LRU-based mechanism, while the latter includes a memory monitor as well, which adds to complexity, and if we see we need a memory monitor, then we will get that merged later on.

[pentschev]

Actually, I don't have permission to close this, would you mind closing it @mrocklin?