[Question] A new approach to memory spilling

[madsbk]

Question: would the Dask/Distributed community be interested in an improved memory spilling model that fixes the shortcomings of the current one but make use of proxy object wrappers?

In Dask-CUDA we have introduced a new approach to memory spilling that handles object aliasing and JIT memory un-spilling: rapidsai/dask-cuda#451

The result is memory spilling that:

• Avoids double counting: Double Counting and Issues w/Spilling #4186
• Avoids spilling of the same object multiple times.
• Avoids memory spikes because of incorrect memory tally.
• Implement just-in-time un-spilling: Delay deserialization of Data in workers until actual usage. #3998
• Support communication of spilled data so that GPU data don’t have to be un-spilled just to be spilled again as part of communication: [FEA] Allow communicating spilled data rapidsai/dask-cuda#342
• First step to enable partial spilled objects such as the spilling of individual columns in a data frame: Ability to have output dask_cudf.DataFrame not necessarily all in memory BlazingDB/blazingsql#1128

The current implement in Dask-CUDA handles CUDA device objects but it is possible to generalize to also handle spilling to disk.

The disadvantage of this approach is the use of proxy objects that get exposed to the users. The inputs to a tasks might be wrapped in a proxy object, which doesn't mimic the proxied object perfectly. E.g.:

    # Type checking using instance() works as expected but direct type checking doesn't:
    >>> import numpy as np
    >>> from dask_cuda.proxy_object import asproxy
    >>> x = np.arange(3)
    >>> isinstance(asproxy(x), type(x))
    True
    >>>  type(asproxy(x)) is type(x)
    False

Because of this, the approach shouldn't be enabled by default but do you think that the Dask community would be interested in a generalization of this approach? Or is the proxy object hurdle too much of an issue?

cc. @mrocklin, @jrbourbeau, @quasiben

[quasiben]

I think having this machine optionally available would be a nice improvement. @eric-czech do many or any sgkit workloads rely on dask spilling ?

In the past, @rabernat described an example where spilling might be impactful -- though the bulk of the conversation is about scheduler decisions/improvements

[jakirkham]

Yes I think this would be great to have in Dask. I recall at one point Matt may have mentioned interest in this being upstreamed (though I could be misremembering)

[eric-czech]

@eric-czech do many or any sgkit workloads rely on dask spilling ?

Hm no in the sense that we prefer to do performance testing w/ spilling disabled but yes in the sense that we still want it to work for our users who may not know how to tune chunk sizes to avoid memory problems. We're still figuring out the latter as well.

Does this spilling typically occur when usage would exceed GPU memory or when it would exceed RAM? I can imagine that spilling may become more critical for us if it is necessary to keep workflows constrained by GPU memory << RAM from failing.

cc: @tomwhite @ravwojdyla

[quasiben]

@alimanfoo recently wrote up some thoughts on changing memory storage/spilling . @madsbk do you think a pluggable backend would be possible with the proxy object you are proposing here ?

[madsbk]

@alimanfoo recently wrote up some thoughts on changing memory storage/spilling. @madsbk do you think a pluggable backend would be possible with the proxy object you are proposing here ?

Yes, this should 100% compatible with a pluggable backend approach. Users will be able to define what should happen when spilling between layers.

BTW, the spilling issues @alimanfoo are reporting remains me of the spilling issues we were having in GPU-BDB. That was the motivation for implementing the new spilling approach :)

[mrocklin]

I think that we should talk about the problems that we're trying to solve, and then assess different solutions to those problems. This seems like one possible solution, possibly out of a few.

This approach seems interesting, but also somewhat magical/complex (although possibly that is only because I don't know much about it yet). I think that we should maybe have a longer discussion first before adopting this particular path.

@madsbk I suspect the answer is no, but do you happen to have anything that you can share where different options here were discussed that shows pros/cons, or should we start that conversation fresh?

[mrocklin]

My thoughts by the way are that we should break the MutableMapping abstraction and start explicitly including things like Disk storage in the worker. I'm also fine including device memory explicitly in the worker. Part of this is described here: #4424 I expect it to be more effective and more simple long term than proxy objects.

[madsbk]

Right, breaking the MutableMapping abstraction will make it fairly easy to fix the double counting issue (#4186) and avoid memory spikes because of incorrect memory tally. However, in order to avoid spilling of the same object multiple times, I think you need some kind of proxy object but maybe it is possible to avoid exposing the proxy objects to the task/user.

In order to support spilling/un-spilling of objects already passed to a task, I don't see a way around proxy objects.

[madsbk]

@madsbk I suspect the answer is no, but do you happen to have anything that you can share where different options here were discussed that shows pros/cons, or should we start that conversation fresh?

I think the only real disadvantage is the proxy objects. I cannot think of a more simple solution that addresses all of the mentioned issues. But if the proxy objects are too much of an issue, I can look at fixing the double counting issue by breaking the MutableMapping abstraction.

[mrocklin]

In order to support spilling/un-spilling of objects already passed to a task, I don't see a way around proxy objects.

Yeah, I think that when a task starts running we wouldn't try to spill it to disk. I'm not sure I understand this point.

More generally, I'm not sure exactly what is being proposed here. Can you say more about what a proxy object is, how it is used, and how it affects users? I get the sense that this is a topic that is already well understood to the RAPIDS team, but may not be well understood more generally.

Does this mean that we would change the object that is passed to user functions? If so, this seems like it might open up a lot of unexpected issues.

[mrocklin]

Right, breaking the MutableMapping abstraction will make it fairly easy to fix the double counting issue (#4186) and avoid memory spikes because of incorrect memory tally

This may also provide other benefits, like allowing for async reading/writing, and improved worker scheduling of tasks that is sensitive to data that is already in fast memory, or pre-fetching data from slow memory.

[madsbk]

Does this mean that we would change the object that is passed to user functions? If so, this seems like it might open up a lot of unexpected issues.

Yes, relevant objects are wrapped in a proxy object before passed to user functions.

class ProxyObject:
    """Object wrapper/proxy for serializable objects

    This is used by ProxifyHostFile to delay deserialization of returned objects.

    Objects proxied by an instance of this class will be JIT-deserialized when
    accessed. The instance behaves as the proxied object and can be accessed/used
    just like the proxied object.

    ProxyObject has some limitations and doesn't mimic the proxied object perfectly.
    Thus, if encountering problems remember that it is always possible to use unproxy()
    to access the proxied object directly or disable JIT deserialization completely
    with `jit_unspill=False`.

    Type checking using instance() works as expected but direct type checking
    doesn't:
    >>> import numpy as np
    >>> from dask_cuda.proxy_object import asproxy
    >>> x = np.arange(3)
    >>> isinstance(asproxy(x), type(x))
    True
    >>>  type(asproxy(x)) is type(x)
    False
    """

[madsbk]

Issues with the current implementation of memory spilling

Object duplication

The same memory object can live inside a task and in the worker's self.data simultaneously. If the worker decides to spill an object x in its self.data, it will not free any memory if the task holds a reference to x. Even worse, if the output of the task includes x, the worker's self.data will now contain two copies of x, a spilled and an un-spilled copy.

Spilling order

The order of which we spill objects can have a drastic effect on performance. The current implementation tracks access time of each key in self.data and not access time of objects in self.data. This means that a getitem(v, 42) task on a collection v will effectively update the access time of all items in v and not only the 42th item.

Aliases at the task level

Consider the following two tasks:

def task1():  # Create list of dataframes
    df1 = cudf.DataFrame({"a": range(10)})
    df2 = cudf.DataFrame({"a": range(10)})
    return [df1, df2]

def task2(dfs):  # Get the second item
    return dfs[1]    

Running the two task on a worker we get something like:

>>> data["k1"] = task1()
>>> data["k2"] = task2(data["k1"])
>>> data
{
    "k1": [df1, df2],
    "k2": df2,
}

Since the current implementation of spilling works on keys and handles each keys separately, it overestimate the device memory used: sizeof(df)*3. But even worse, if it decides to spill k2 no device memory is freed since k1 still holds a reference to df2!

The new spilling approach fixes this issue by wrapping identical objects in a shared ProxyObejct thus in this case df2 in both k1 and k2 will refer to the same ProxyObejct.

Sharing device buffers

Consider the following code snippet:

>>> data["df"] = cudf.DataFrame({"a": range(10)})
>>> data["grouped"] = shuffle_group(data["df"], "a", 0, 2, 2, False, 2)
>>> data["v1"] = data["grouped"][0]
>>> data["v2"] = data["grouped"][1]

In this case v1 and v2 are separate objects and are handled separately both in the current and the new spilling implementation. However, the shuffle_group() in cudf actually returns a single device memory buffer such that v1 and v2 points to the same underlying memory buffer. Thus the current implement will again overestimate the memory use and spill one of the dataframes without any effect.
The new implementation takes this into account when estimating memory usage and make sure that either both dataframes are spilled or none of them are.

[mrocklin]

Thank you for the detailed explanation @madsbk. This is quite helpful to motivate the change. I'm curious, in practice does this mostly come up in dataframe shuffle situations when we're low on memory? Are there other common use cases?

This would be a very large change in the abstractions used for Dask workers (user data would no longer be unchanged), and the introduction of a substantial amount of magic. Personally, if there are not very many problems that this solves, then I would prefer to try to think of other possible solutions to this problem before going down this path. My concern is that a change like this would negatively impact a number of other workloads in hard-to-understand ways.

[jakirkham]

Maybe this is worth discussing at the monthly meeting next week? 😉