Cython MR ownership and resource_refs

[harrism]

@wence- wrote the following as a comment on a PR. I wanted to move it here where we can have a better discussion.

On the Python side, we need someone to own concrete objects (we can't just provide classes that wrap device_async_resource_ref since that is non-owning).

Adapters also need a way to keep their upstream alive.

We are removing the device_memory_resource inheritance on the C++ side, but I think until we have an owning any_resource available, we probably need to keep something on the Cython side.

Here is a very slide-ware sketch approach for discussion:

cdef extern from * nogil:
    """
    struct hack {
        std::function<rmm::device_async_resource_ref> _to_ref;
        hack() = default; // unsafe, but this is all hacky anyway
        rmm::device_async_resource_ref to_ref() { return _to_ref(); }
        template<typename T>
        hack(std::unique_ptr<T> obj) : _to_ref{[&obj]() -> rmm::device_async_resource_ref { return *obj; }} {};
    };
    """
       
cdef class DeviceMemoryResource:
   cdef hack obj;
   cdef device_async_resource_ref as_ref(self):
       return self.obj.to_ref()

cdef class CudaMemoryResource(DeviceMemoryResource):
   def __cinit__(self):
       self.obj = hack(move(make_unique[cuda_memory_resource]()))
   

cdef class PoolMemoryResource(DeviceMemoryResource):
   cdef DeviceMemoryResource upstream
   def __cinit__(self, DeviceMemoryResource upstream):
      self.upstream = upstream
      self.obj = hack(move(make_unique[pool_memory_resource](upstream.to_ref())))

This is effectively working around not (yet) having any_resource available by doing a minimal type-erased owning wrapper.

[harrism]

In spite of the "hack", @wence-, I do like the way you've encapsulated the hack in an object so that each derived MR type just has to copy and modify a single line to implement ownership and to_ref()!

I think we still have the problem of Cython not being able to return a resource_ref since it has no default ctor, but I think if RMM Python continues to have its own get/set_current_device_resource() then only code inside that will be touching the resource_ref anyway.

[wence-]

any_resource is now in the cudax experimental namespace in CCCL. That still leaves a few things to think about:

• If set_current_device_resource stores an any_resource in the global map, then the ownership is clear, but how would one use it from python? We need it to be safe that if we grab the reference that if someone subsequently does set_current_device_resource again, we don't have use-after-free
• If we decide the use any_resource, how should we design the cython classes for adaptors? Should they own their upstreams, or should that be the responsibility of the stored any_resource? If we want that, then the C++ adaptors need to accept an any_resource and not just a _ref.
• If we don't, do we want to try and solve the problem of set_current_device_resource not being safe to call from both python and C++?
• your thought here.

[harrism]

If set_current_device_resource stores an any_resource in the global map, then the ownership is clear,

Is it? Is any_resource like a unique_ptr?

If not, I'm not sure we want to use any_resource -- we don't want pools to be copied around. :)

[harrism]

adaptors? Should they own their upstreams

I don't think adaptors should own their upstreams. One might want a pool and to have multiple adapted versions of it.

[harrism]

If we don't, do we want to try and solve the problem of set_current_device_resource not being safe to call from both python and C++?

Again, I don't think libraries should be calling it, only application code. But perhaps we can make it safe.

[wence-]

I don't think adaptors should own their upstreams. One might want a pool and to have multiple adapted versions of it.

OK, let's have a tighter question: should adaptors keep their upstreams alive? I think on the python side we want the answer to be "yes". On the C++ side I can understand having RAII-style liveness semantics.

Is it? Is any_resource like a unique_ptr?
If not, I'm not sure we want to use any_resource -- we don't want pools to be copied around. :)

I think (and tagging @miscco since I think I didn't fully understand our discussions the other week), that any_resource is not like unique_ptr: AFAICT it has a copy constructor defined, but I don't know what this "does" if the resource in question is a pool (say).

[miscco]

So any_resource does own the resource that it contains.

This means that copying an any_resource creates a copy of e.g. the pool

This is most often not what you want.

However, what you can do is either own the upstream by move constructing it into the adaptor or create a resource_ref from the upstream and pass that into the adaptor. In that case the any_resource will own the resource_ref so you get your reference semantics. But at the same time it wont keep the resource alive.

If you want your cake and eat it too, you would need something like a shared_resource

template<class Resource>
// requires cuda::mr::resource<Resource>
struct shared_resource {

    shared_resource() : resource_(new Resource()) {};

    // The resource interface    
    void allocate(...) {
        resource_->allocate(...);
    }    
    void deallocate(...) {
        resource_->deallocate(...);
    }    

    std::shared_ptr<Resource> resource_{};
};

[harrism]

The reason that the issue of current device resource changing underneath a library has never worried me is because this is how I think about set/get_current_device_resource:

1. Only application code should ever call set_current_device_resource(). Library code should never call it unless told to do so by application code.
2. Library code should only ever call get_current_device_resource().

This is not particularly well documented.

BTW, I noticed that PyTorch has an even more strict (and enforced) If you set the external CUDA allocator for Pytorch, it is an error to set it again.

[wence-]

Only application code should ever call set_current_device_resource(). Library code should never call it unless told to do so by application code.

I think this is reasonable, but what counts as an "application"? We can argue whether or not it should be the case, but right now cudf (python) calls set_current_device_resource: is cudf an application (in which case fine?), or a library (in which case, dubious).

[harrism]

That's a good question.

libcudf is definitely a library. It does not call set_current_device_resource() except in examples and tests (which are applications).

I can see that this is more grey when it comes to Python. Especially in cases like cudf.pandas initialization where we want to provide sensible defaults. We should keep discussing this.