[REVIEW] Out-of-memory callback resource adaptor

[madsbk]

This PR implements a new resource adaptor that calls a callback function when an allocation fails. The idea being that the callback function can free up memory (e.g. by spilling) and ask rmm to retry the allocation.

/**
 * @brief Resource that uses `Upstream` to allocate memory and calls `callback`
 * when allocations throws `std::bad_alloc`.
 *
 * An instance of this resource can be constructed with an existing, upstream
 * resource in order to satisfy allocation requests.
 *
 * The callback function takes an allocation size and a closure and returns
 * whether to retry the allocation or throw `std::bad_alloc`.
 *
 * @tparam Upstream Type of the upstream resource used for
 * allocation/deallocation.
 */
template <typename Upstream>
class oom_callback_resource_adaptor final : public device_memory_resource

This is motivated by the fairly primitive spilling in Dask. Currently, Dask and Dask-CUDA has no way of handling OOM errors other then restarting tasks or workers. Instead they spill preemptively based on some very conservative memory thresholds. For instance, most workflows in Dask-CUDA starts spilling when half the GPU memory is in use.
This PR makes it possible for projects like Dask and Dask-CUDA to trigger spilling on demand instead of preemptively.

cc. @jrhemstad @shwina

[shwina]

Cython looks fantastic!

[jrhemstad]

Wow, that was fast.

[jrhemstad]

Can we also get a test written in C++?

[rongou]

We do this for Spark as well, but from Scala. Perhaps we can reuse this. @jlowe @revans2 @abellina

[abellina]

Yeap @rongou we do it here: https://github.com/rapidsai/cudf/blob/branch-21.12/java/src/main/native/src/RmmJni.cpp#L139. We hook at the mr level, and call a JNI function in our case.

Here's where we handle an oom: https://github.com/rapidsai/cudf/blob/branch-21.12/java/src/main/native/src/RmmJni.cpp#L212

Our resource had handling for threshold-based OOM (i.e. not real OOM from RMM, but instead a low/high watermark for some preemptive spilling). We are not using the low/high watermark at the moment. I am sure part of that could be refactored into its own memory resource.

[madsbk]

Can we also get a test written in C++?

@jrhemstad thanks for the review, I have added a C++ test, renamed the closure argument, and added some more doc.

[harrism]

I think the naming should be more general and not use an acronym. Other than that and a few doc improvements, looks like a great contribution. Thanks!

[madsbk]

I think the naming should be more general and not use an acronym. Other than that and a few doc improvements, looks like a great contribution. Thanks!

Thanks for the review @harrism, I think I have addressed all of your suggestions.

[harrism]

Nearly there.

[harrism]

@gpucibot merge