Skip to content

Experimental reservation aware resource adapter and reservation - #174

Open
nirandaperera wants to merge 6 commits into
NVIDIA:mainfrom
nirandaperera:exp-res-aware-resourece-adapter
Open

Experimental reservation aware resource adapter and reservation#174
nirandaperera wants to merge 6 commits into
NVIDIA:mainfrom
nirandaperera:exp-res-aware-resourece-adapter

Conversation

@nirandaperera

@nirandaperera nirandaperera commented Jul 30, 2026

Copy link
Copy Markdown
Contributor

[C++] Add experimental reservation-aware memory resource adaptor

Related to rapidsai/rapidsmpf#1142

Callers could only discover that a memory budget was exhausted at allocation time, with no way to claim a slice of it up front. This adds cucascade::memory::experimental, where you reserve bytes against a limit and then allocate through that reservation, so an over-budget request is refused at reserve() rather than failing deep inside a cudf call. Ported from rapidsmpf and extended so that a single reservation type covers device, host, and host+device memory.

  • reserve(size, allow_overbooking) yields a memory_reservation that charges each allocation against its balance() and refunds the unspent remainder once the last reference dies. Overbooking is opt-in: permitted, the full size is granted and the shortfall reported by overbooking(); refused, a zero-sized reservation comes back and throws rmm::out_of_memory on first use.
  • Three adaptors — device_adaptor, host_adaptor, and host_device_adaptor — forward their upstream's accessibility to every reservation they grant, so a device adaptor cannot hand out host-usable memory.
  • memory_reservation is one type no matter where its memory lives. To hand it to cudf or RMM, project it with as_device(), as_host(), or as_host_device(); each throws cucascade::logic_error when the reservation lacks that accessibility, and accessibility() lets callers branch without exceptions.
  • limit(), current_allocated(), total_reserved(), and available() expose the budget, with get_main_record() for peak and allocation-count statistics. The adaptor is itself a memory resource, so allocations made directly through it are still tracked against available(), just not capped.
device_adaptor adaptor{any_device_resource{rmm::mr::cuda_memory_resource{}}, 1GB};

auto res = adaptor.reserve(4MB, allow_overbooking::NO);

rmm::device_buffer buf{1MB, stream, res.as_device()};
assert(res.balance() == 4MB - 1MB);

auto table = cudf::groupby(..., stream, res.as_device()); // all allocations are capped against res

// The adaptor is a resource too: tracked against available(), but not capped.
rmm::device_buffer scratch{1ull << 20, stream, adaptor};

Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Signed-off-by: niranda perera <niranda.perera@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant