Let callers reserve partitioning memory - #23833
Conversation
madsbk
added
improvement
github-actions
Bot
added
libcudf
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughSummary by CodeRabbit
WalkthroughThe change adds cost-estimation APIs and caller-provided memory reservations for partitioning, splitting, and unpacking in C++ and Python. It adds reservation validation, unspill accounting, public exports, type declarations, bindings, and regression tests. ChangesPartition memory APIs
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: 🟡 Moderate · up to The PR adds caller-supplied reservations and memory-cost APIs, but moved-from partition data can still be dereferenced during unpacking or cost estimation, risking a runtime crash for affected callers. Merge should wait for this correctness issue to be fixed or explicitly accepted. Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…e-partitioning-memory
…e-partitioning-memory
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
|
|
||
| } // namespace | ||
|
|
||
| std::pair<std::vector<cudf::table_view>, std::unique_ptr<cudf::table>> partition_and_split( |
There was a problem hiding this comment.
Q: why wouldnt this require an external reservation variant?
There was a problem hiding this comment.
No reason in principle, it just has no caller that needs one. partition_and_split is used only from the benchmarks and test_shuffler.cpp, it isn't bound in Python, and nothing in the cudf-polars shuffle path calls it, so there is nothing to hand it a reservation.
| RAPIDSMPF_NVTX_FUNC_RANGE(); | ||
| RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); | ||
| RAPIDSMPF_EXPECTS(num_partitions > 0, "Need to split to at least one partition"); | ||
| if (table.num_rows() == 0) { | ||
| auto splits = | ||
| std::vector<cudf::size_type>(rapidsmpf::safe_cast<std::uint64_t>(num_partitions - 1), 0); | ||
| return split_and_pack(table, splits, stream, br, allow_overbooking); | ||
| return split_and_pack_impl(table, splits, stream, br, allow_overbooking, reservation); | ||
| } | ||
|
|
||
| // hash_partition does a deep-copy. Therefore, we need to reserve memory for | ||
| // at least the size of the table. | ||
| auto reservation = | ||
| br->reserve_device_memory_and_spill(estimated_memory_usage(table, stream), allow_overbooking); | ||
| // at least the size of the table. `packed_size()` measures the same bytes as the | ||
| // copy needs, rounded up to the packing alignment, so it is a safe over-estimate. | ||
| auto const reorder_bytes = cudf::packed_size(table, stream, br->device_mr()); | ||
| check_reservation(reservation, reorder_bytes); | ||
| auto own_reservation = reserve_unless_provided(reservation, reorder_bytes, br, allow_overbooking); | ||
| auto [reordered, split_points] = cudf::hash_partition( | ||
| table, columns_to_hash, num_partitions, hash_function, seed, stream, br->device_mr()); | ||
| reservation.clear(); | ||
| consume_reservation(br, reservation, own_reservation, reorder_bytes); | ||
| std::vector<cudf::size_type> splits(split_points.begin() + 1, split_points.end() - 1); | ||
| return split_and_pack(reordered->view(), splits, stream, br, allow_overbooking); | ||
| return split_and_pack_impl(reordered->view(), splits, stream, br, allow_overbooking, reservation); |
There was a problem hiding this comment.
Can I suggest an alternative RAII here?
struct scoped_reservation{
scoped_reservation(BufferResource* br, MemoryReservation* res, size_t size): br(br), res(res), size(size){
// res != nullptr
// move check reservation logic here
}
~scoped_reservation(){
br->release(*res, size);
}
BufferResource* br;
MemoryReservation* res;
size_t size;
}Now, let's accept reservation as a optional<MemoryReservation>. If its nullopt, assign the new reservation sized reorder_bytes. then create a scoped_reservation. I think its more idiomatic
There was a problem hiding this comment.
maybe we can call it release_reservation instead
There was a problem hiding this comment.
MemoryReservation is already RAII: its destructor releases the remaining balance, so scope-exit cleanup is covered and the current code is exception-safe either way.
What consume_reservation() does is a partial release when an allocation lands, so the claim shrinks before the second phase runs on the remainder. That is not a scope-exit operation.
Also, MemoryReservation is move-only, and the public API takes it by reference so the caller retains ownership and can inspect it afterwards. It therefore cannot be an optional.
Finally, having a destructor call release() would make it potentially throwing, since release() throws when the reservation is too small.
There was a problem hiding this comment.
Using a ReservationHandle class instead.
| rapidsmpf::BufferResource* br, | ||
| rapidsmpf::MemoryReservation& reservation) | ||
| { | ||
| RAPIDSMPF_EXPECTS(reservation.mem_type() == rapidsmpf::MemoryType::DEVICE, |
There was a problem hiding this comment.
DEVICE only.
These functions don't choose where the memory goes, everything allocates through br->device_mr(), so the reservation type is determined rather than a preference. consume_reservation() calls br->release(), which decrements memory_reserved_ for the reservation's own memory type, so a PINNED_HOST reservation would decrement the pinned counter while the allocation lands in the device-tracked adaptor.
For unpack_and_concat it is already enforced downstream by move_to_device_buffer(). Checking at the boundary just gives a better error message before any partition has been moved.
|
|
||
| std::vector<cudf::table_view> unpacked; | ||
| std::vector<cudf::packed_columns> references; | ||
| std::vector<rmm::cuda_stream_view> packed_data_streams; |
There was a problem hiding this comment.
shouldnt this be a set/unordered_set? Otherwise we calling duplicated stream joins, isnt it? 🤔
There was a problem hiding this comment.
rmm::cuda_stream_view doesn't have a hash or < operator though, and adding a custom one here is overkill. Maybe RMM should add it.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
| * | ||
| * @param br Buffer resource holding the reservation. | ||
| * @param reservation The caller's reservation, or `nullptr`. | ||
| * @param own The reservation made by `reserve_unless_provided()`, if any. |
There was a problem hiding this comment.
The reservation can come from anywhere, right? Ah, no, because we've got this weird split between the caller-provided reservation and the optional reservation that we made.
There was a problem hiding this comment.
Agreed, and the code was already inconsistent about it: unpack_and_concat_impl() resolved both into a single reference while the other two carried both around. Replaced check_reservation(), reserve_unless_provided() and consume_reservation() with a small ReservationHandle that resolves the two cases once, privately, so no signature mentions the split any more. The size check moves into its constructor, so a handle cannot be built without validating. release() stays explicit rather than running in the destructor, since BufferResource::release() throws.
| [[nodiscard]] std::optional<rapidsmpf::MemoryReservation> reserve_unless_provided( | ||
| rapidsmpf::MemoryReservation* reservation, | ||
| std::size_t size, | ||
| rapidsmpf::BufferResource* br, |
There was a problem hiding this comment.
question: What happens/should happen if the caller provides a reservation but it is not big enough?
There was a problem hiding this comment.
Now, ReservationHandle's constructor throws reservation_error before anything is allocated, so nothing has been allocated, the reservation is untouched, and the input table or partitions are unchanged.
| auto const reorder_bytes = cudf::packed_size(table, stream, br->device_mr()); | ||
| check_reservation(reservation, reorder_bytes); |
There was a problem hiding this comment.
This check just throws if the provided reservation is not big enough. Is that what we want? It seems like there's no way to recover at that point?
There was a problem hiding this comment.
It is recoverable, precisely because the check runs before the allocation, in ReservationHandle's constructor. Catch, reserve more, call again, nothing was consumed on the failed attempt. That is why it is up front rather than relying on release() to fail later, which was the earlier behaviour and left an unreserved allocation already on the device.
| br->reserve_device_memory_and_spill(estimated_memory_usage(table, stream), allow_overbooking); | ||
| // at least the size of the table. `packed_size()` measures the same bytes as the | ||
| // copy needs, rounded up to the packing alignment, so it is a safe over-estimate. | ||
| auto const reorder_bytes = cudf::packed_size(table, stream, br->device_mr()); |
There was a problem hiding this comment.
question: In the case that the caller provided a reservation they presumably did so by calling packed_size to get the size they need. This calls packed_size again on the same table (incurring another stream sync). Can we avoid that somehow?
There was a problem hiding this comment.
Not cleanly, I think. Deriving the size from reservation.size() breaks as soon as a caller over-reserves, which is legal. The only sound alternative is to thread the cost breakdown through the API, which adds a parameter to every overload. I am not sure that is worth it?
| { | ||
| // The reorder and the packed partitions are each about one packed table, and an | ||
| // empty table skips the reorder entirely. | ||
| auto const packed_size = cudf::packed_size(table, stream, temp_mr); |
There was a problem hiding this comment.
nit: packed_size incurs a stream sync in the general case, should we skip it if table.num_rows() is zero (and just return zero?)
| // Covers the unspill below and the concatenation at the end. | ||
| check_reservation(reservation, total_size + non_device_size); | ||
| auto own_reservation = | ||
| reserve_unless_provided(reservation, total_size + non_device_size, br, allow_overbooking); |
There was a problem hiding this comment.
Same queries as above about recovery.
There was a problem hiding this comment.
The handle is built before the loop that moves the partitions, so an undersized reservation throws with nothing allocated and the caller's vector still intact. Catch, reserve more, retry.
3 participants
partition_and_pack(),split_and_pack()andunpack_and_concat()can now take a caller-providedMemoryReservation&instead of reserving and spilling internally. That lets a caller reserve before it starts, so the point where it might block is explicit rather than buried inside the call.partition_and_pack_cost(),split_and_pack_cost()andunpack_and_concat_cost()return the peak device memory each function needs. All six are bound incudf_streaming.partition_utils, where the reservation is an optional trailing argument, so existing callers are unaffected.Nothing calls the new overloads yet beyond the tests. cudf-polars picks them up in the follow-up #23834, which gives the shuffle memory backpressure on both the insert and the extract side.
Notes
Only the unspill share of
unpack_and_concat_cost()is exact, since the buffer resource consumes it while moving each partition. The rest is an estimate, because libcudf allocates againstBufferResource::device_mr()and never sees the reservation.split_and_pack_cost()in particular under-reports for more than one partition, sincecontiguous_split()aligns every column buffer of every partition. That was true of the internal reservations before this change too.unpack_and_concat_cost()has an overload taking a vector of pointers, so Cython can compute the cost without moving the partitions out of their Python owners.