NVIDIA / NVIDIA/cudf

[FEA] provide external libraries a way of getting a `DeviceBuffer` pointer that can become spillable again

Open
#14,029 3 comments 0 reactions 0 assignees View on GitHub
0 - Backlog feature request improvement Performance Python
Dominant language
C++
Stars
9.8k
Forks
1.1k
Avg merge
3d 6m
Merged PRs (30d)
278

Description

**Is your feature request related to a problem? Please describe.**

When running in a multi-gpu setting, message passing with ucx-py takes a `DeviceBuffer` and obtains the device memory pointer through the `__cuda_array_interface__`. This, correctly, marks the buffer as unspillable.

It would be nice if there were a way to expose a pointer that is marked as unspillable until the external library drops the reference (kind of like `acquire_spill_lock`). `ucx-py` could then use it, and scope the pointer use to the lifetime of the message request (once the request is completed, the pointer can be dropped and is available for spilling again).

**Describe the solution you'd like**

If we were to hand back an object that had a `weakref.finalize(obj, unmark_spillable)` callback, when it was dropped, we could let the buffer be spillable again.

**Describe alternatives you've considered**

Making ucx-py aware of cudf and using `acquire_spill_lock`.

cc @madsbk / @vyasr / @galipremsagar

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.