Let callers reserve partitioning memory by madsbk · Pull Request #23833 · NVIDIA/cudf

madsbk · 2026-08-26T12:41:38Z

partition_and_pack(), split_and_pack() and unpack_and_concat() can now take a caller-provided MemoryReservation& 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() and unpack_and_concat_cost() return the peak device memory each function needs. All six are bound in cudf_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 against BufferResource::device_mr() and never sees the reservation. split_and_pack_cost() in particular under-reports for more than one partition, since contiguous_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.

…e-partitioning-memory

coderabbitai · 2026-08-26T14:57:21Z

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
No new commits to review - use @coderabbitai full review for a full pass

📝 Walkthrough

Summary by CodeRabbit

New Features
- Added memory-cost estimates for partitioning, splitting, and unpacking operations.
- Added optional memory reservations for predictable device-memory usage.
- Added support for specifying packed data sizes during reservation-aware partitioning and splitting.
- Exposed cost-estimation helpers through the Python package.
Bug Fixes
- Improved reservation validation, accounting, and reuse after insufficient-memory attempts.
- Cost estimates no longer consume partition data.
- Improved handling of empty tables and consumed partitions.
Tests
- Added coverage for reservations, round trips, and error handling.

Walkthrough

The change adds cost-estimation APIs and caller-provided memory reservations for partitioning, splitting, and unpacking in C++ and Python. It updates stream types, reservation handling, bindings, public declarations, exports, and regression tests.

Changes

Partition memory APIs

Layer / File(s)	Summary
C++ partition and split reservation flow `cpp/libcudf_streaming/include/cudf_streaming/partition_utils.hpp`, `cpp/libcudfudf_streaming/src/partition_utils.cpp`	Adds cost estimation, optional packed sizes, caller-provided reservations, and `cuda::stream_ref` support for partitioning and splitting.
C++ unpack and concatenate reservation flow `cpp/libcudf_streaming/include/cudf_streaming/partition_utils.hpp`, `cpp/libcudf_streaming/src/partition_utils.cpp`	Adds unpack-cost overloads, packed-input validation, and separate reservation accounting for unspilling and concatenation.
Python API bindings and public declarations `python/cudf_streaming/cudf_streaming/partition_utils.*`, `python/cudf_streaming/cudf_streaming/__init__.py`	Exposes cost functions and optional `MemoryReservation` parameters through Cython bindings, stubs, declarations, stream conversions, and package exports.
Reservation and cost-estimation validation `cpp/libcudf_streaming/tests/test_partition.cpp`, `python/cudf_streaming/cudf_streaming/tests/test_integration_partition.py`	Tests reserved round trips, invalid reservations, empty tables, spill and unspill accounting, reusable cost estimation, generators, and consumed partitions.

Estimated code review effort: 4 (Complex) | ~45 minutes

Merge Risk: 🔴 Critical · up to 7e352

The change adds caller-managed memory reservations, but the current implementation still has binding declarations inconsistent with the C++ APIs, which can block builds, and failed reservations may consume caller-owned partition data so the operation cannot be retried safely. These issues should be fixed before merging.

Suggested reviewers: quasiben, nirandaperera, vyasr, bdice

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 12.90% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 31 functions across 6 files. (1 skipped: …	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title accurately identifies the central change: callers can provide reservations for partitioning memory. It does not mention splitting, unpacking, or cost-estimation APIs, but those are supportin…
Description check	✅ Passed	The description clearly explains caller-provided memory reservations, cost-estimation APIs, Python bindings, compatibility, and test coverage. It is directly related to the changeset.

Full details: Docstring Coverage

Explanation

Docstring coverage is 12.90% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 31 functions across 6 files. (1 skipped: 1 unsupported.)

Full details: Title check

Explanation

The title accurately identifies the central change: callers can provide reservations for partitioning memory. It does not mention splitting, unpacking, or cost-estimation APIs, but those are supporting aspects of the changeset.

Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

…e-partitioning-memory

wence-

Very small suggestions

…e-partitioning-memory

nirandaperera

some suggestions

…e-partitioning-memory

nirandaperera · 2026-08-27T18:00:55Z

+std::unordered_map<rapidsmpf::shuffler::PartID, rapidsmpf::PackedData> split_and_pack_impl(
+  cudf::table_view const& table,
+  std::vector<cudf::size_type> const& splits,
+  rmm::cuda_stream_view stream,
+  rapidsmpf::BufferResource* br,
+  rapidsmpf::AllowOverbooking allow_overbooking,
+  rapidsmpf::MemoryReservation* reservation);


I think we can request a mandatory reservation in every *_impl methods.

Suggested change

std::unordered_map<rapidsmpf::shuffler::PartID, rapidsmpf::PackedData> split_and_pack_impl(

cudf::table_view const& table,

std::vector<cudf::size_type> const& splits,

rmm::cuda_stream_view stream,

rapidsmpf::BufferResource* br,

rapidsmpf::AllowOverbooking allow_overbooking,

rapidsmpf::MemoryReservation* reservation);

std::unordered_map<rapidsmpf::shuffler::PartID, rapidsmpf::PackedData> split_and_pack_impl(

cudf::table_view const& table,

std::vector<cudf::size_type> const& splits,

rmm::cuda_stream_view stream,

rapidsmpf::MemoryReservation& reservation);

The signature is nicer, br and allow_overbooking both fall out since MemoryReservation carries br(). But it moves the reserving up into the public non-reservation overload, and that costs the callers who use it today, which is everything: the benchmarks, the tests and cudf-polars. They would pay an extra packed_size() sync, and hold twice the peak claim, since partition_and_pack() currently reserves the reorder and the pack in sequence rather than together.

No I meant in the, *_impl methods only. We have,

auto partition_and_pack_impl(..., res); // anon ns auto partition_and_pack(..., br, ...){ // original API // create res // call impl } auto partition_and_pack(..., res, ...) { // new API // call impl }

Maybe that way, we dont really need an impl as well?

Ah, I see what you mean. I was hesitant because it would worsen the repeated cudf::packed_size() calls @wence- mentions. But both are handled now: it does what you suggest, and an optional packed_bytes argument lets the size be computed once and passed down rather than recomputed in each phase.

`MemoryReservation::split(size)` reduces a reservation by `size` and returns a new reservation of that size on the same buffer resource and memory type. The total reserved by the buffer resource is unchanged, the bytes only move between the two reservations. This lets a caller scope part of a reservation to the allocation it covers. The returned reservation releases its bytes when it goes out of scope, so the reservation is never counted on top of memory that has already been allocated, and there is no size argument at the release point to disagree with what was allocated. The alternative is `BufferResource::release(reservation, size)`, which works but leaves the caller to pick the release point and to repeat the size. `split()` also composes with a caller-provided reservation: a function can carve off what it needs without caring whether the caller supplied the reservation or it made one itself. Adopted by cudf's `partition_and_pack()` and friends, which take a caller-provided reservation and split off one sub-reservation per allocation phase: NVIDIA/cudf#23833 Authors: - Mads R. B. Kristensen (https://github.com/madsbk) Approvers: - Niranda Perera (https://github.com/nirandaperera) URL: #1177

…e-partitioning-memory

nirandaperera

@madsbk Last change requests! Promise :inno

nirandaperera · 2026-08-28T16:27:33Z

 #include <cudf/partitioning.hpp>
 #include <cudf/table/table.hpp>

+#include <cuda/stream_ref>


stream_ref header is deprecated 😞

Suggested change

#include <cuda/stream_ref>

#include <cuda/stream>

nirandaperera · 2026-08-28T16:33:37Z

+  if (!packed_bytes.has_value()) {
+    packed_bytes = table.num_rows() == 0 ? 0 : cudf::packed_size(table, stream, temp_mr);
+  }


nit. empty tables are handled by packed_size.
https://github.com/NVIDIA/cudf/blob/main/cpp/src/copying/contiguous_split.cu#L2216

Suggested change

if (!packed_bytes.has_value()) {

packed_bytes = table.num_rows() == 0 ? 0 : cudf::packed_size(table, stream, temp_mr);

}

if (!packed_bytes.has_value()) {

packed_bytes = cudf::packed_size(table, stream, temp_mr);

}

nirandaperera · 2026-08-28T16:42:20Z

+  rapidsmpf::BufferResource* br,
+  rapidsmpf::MemoryReservation& reservation,
+  std::optional<std::size_t> packed_bytes)
+{
+  RAPIDSMPF_EXPECTS(reservation.mem_type() == rapidsmpf::MemoryType::DEVICE,
+                    "reservation must be for device memory",
+                    std::invalid_argument);
+  RAPIDSMPF_EXPECTS(reservation.br() == br,
+                    "reservation must belong to the given buffer resource",
+                    std::invalid_argument);


Suggested change

rapidsmpf::BufferResource* br,

rapidsmpf::MemoryReservation& reservation,

std::optional<std::size_t> packed_bytes)

{

RAPIDSMPF_EXPECTS(reservation.mem_type() == rapidsmpf::MemoryType::DEVICE,

"reservation must be for device memory",

std::invalid_argument);

RAPIDSMPF_EXPECTS(reservation.br() == br,

"reservation must belong to the given buffer resource",

std::invalid_argument);

rapidsmpf::MemoryReservation& reservation,

std::optional<std::size_t> packed_bytes)

{

auto br = reservation.br();

RAPIDSMPF_EXPECTS(reservation.mem_type() == rapidsmpf::MemoryType::DEVICE,

"reservation must be for device memory",

std::invalid_argument);

nirandaperera · 2026-08-28T16:45:00Z

+  rapidsmpf::BufferResource* br,
+  rapidsmpf::MemoryReservation& reservation,
+  std::optional<std::size_t> packed_bytes)
+{
+  RAPIDSMPF_EXPECTS(reservation.mem_type() == rapidsmpf::MemoryType::DEVICE,
+                    "reservation must be for device memory",
+                    std::invalid_argument);
+  RAPIDSMPF_EXPECTS(reservation.br() == br,
+                    "reservation must belong to the given buffer resource",
+                    std::invalid_argument);


Suggested change

rapidsmpf::BufferResource* br,

rapidsmpf::MemoryReservation& reservation,

std::optional<std::size_t> packed_bytes)

{

RAPIDSMPF_EXPECTS(reservation.mem_type() == rapidsmpf::MemoryType::DEVICE,

"reservation must be for device memory",

std::invalid_argument);

RAPIDSMPF_EXPECTS(reservation.br() == br,

"reservation must belong to the given buffer resource",

std::invalid_argument);

rapidsmpf::MemoryReservation& reservation,

std::optional<std::size_t> packed_bytes)

{

auto br = reservation.br();

RAPIDSMPF_EXPECTS(reservation.mem_type() == rapidsmpf::MemoryType::DEVICE,

"reservation must be for device memory",

std::invalid_argument);

nirandaperera · 2026-08-28T16:47:27Z

+                                               rapidsmpf::BufferResource* br,
+                                               rapidsmpf::MemoryReservation& reservation)


nirandaperera · 2026-08-28T16:52:19Z

+  // Checked up front so an undersized reservation is caught before the first split
+  // mutates it, leaving the caller free to reserve more and retry.
+  check_reservation(reservation, total_size + non_device_size);
+
+  // The unspill consumes its reservation as it moves each partition, the concatenation
+  // only needs its bytes accounted for until it has allocated them.
+  auto unspill_res = reservation.split(non_device_size);
+  auto concat_res  = reservation.split(total_size);


I think these are redundant. Splitting the reservation for total_size + non_device_size should cover all these

Suggested change

// Checked up front so an undersized reservation is caught before the first split

// mutates it, leaving the caller free to reserve more and retry.

check_reservation(reservation, total_size + non_device_size);

// The unspill consumes its reservation as it moves each partition, the concatenation

// only needs its bytes accounted for until it has allocated them.

auto unspill_res = reservation.split(non_device_size);

auto concat_res = reservation.split(total_size);

// Checked up front so an undersized reservation is caught before the first split

// mutates it, leaving the caller free to reserve more and retry.

auto scoped_res = reservation.split(total_size + non_device_size);

nirandaperera · 2026-08-28T16:59:47Z

+ * Computed here otherwise, which syncs the stream.
+ * @return A pair of the packed size and the total cost, both in bytes.
+ */
+[[nodiscard]] std::pair<std::size_t, std::size_t> packed_and_total_size(


Nit. wondering if we really need this 😇

madsbk added 2 commits August 26, 2026 14:40

Let callers reserve partitioning memory up front

83b8db9

Merge branch 'main' of github.com:NVIDIA/cudf into let-callers-reserv…

6003041

…e-partitioning-memory

madsbk self-assigned this Aug 26, 2026

madsbk added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Aug 26, 2026

github-actions Bot added libcudf Affects libcudf (C++/CUDA) code. Python Affects Python cuDF API. labels Aug 26, 2026

github-project-automation Bot added this to cuDF Python Aug 26, 2026

github-project-automation Bot moved this to Todo in cuDF Python Aug 26, 2026

madsbk mentioned this pull request Aug 26, 2026

Apply memory backpressure to the shuffle #23834

Draft

madsbk marked this pull request as ready for review August 26, 2026 14:50

madsbk requested review from a team as code owners August 26, 2026 14:50

madsbk requested a review from quasiben August 26, 2026 14:50

NVIDIA deleted a comment from copy-pr-bot Bot Aug 26, 2026

This comment was marked as resolved.

Sign in to view

madsbk added 2 commits August 26, 2026 17:10

tests

10a9938

Merge branch 'main' of github.com:NVIDIA/cudf into let-callers-reserv…

9abf166

…e-partitioning-memory

wence- approved these changes Aug 26, 2026

View reviewed changes

Comment thread cpp/libcudf_streaming/include/cudf_streaming/partition_utils.hpp Outdated

Comment thread cpp/libcudf_streaming/src/partition_utils.cpp Outdated

Comment thread cpp/libcudf_streaming/src/partition_utils.cpp Outdated

madsbk added 2 commits August 26, 2026 22:05

docs

39da76a

Merge branch 'main' of github.com:NVIDIA/cudf into let-callers-reserv…

d8182da

…e-partitioning-memory

This comment was marked as resolved.

Sign in to view

nirandaperera requested changes Aug 27, 2026

View reviewed changes

madsbk added 3 commits August 27, 2026 07:44

cleanup

3542b80

Merge branch 'main' of github.com:NVIDIA/cudf into let-callers-reserv…

4b2739f

…e-partitioning-memory

based on review my @nirandaperera

cf4dbbc

madsbk requested a review from nirandaperera August 27, 2026 06:53

This comment was marked as resolved.

Sign in to view

more tests

62c94de

wence- reviewed Aug 27, 2026

View reviewed changes

madsbk added 2 commits August 27, 2026 14:19

ReservationHandle

5ea9c5c

Merge branch 'main' into let-callers-reserve-partitioning-memory

09eff2e

nirandaperera reviewed Aug 27, 2026

View reviewed changes

madsbk mentioned this pull request Aug 27, 2026

Add MemoryReservation::split() rapidsai/rapidsmpf#1177

Merged

Use rapidsai/rapidsmpf#1177

3ae6b75

This comment was marked as resolved.

Sign in to view

nirandaperera reviewed Aug 27, 2026

View reviewed changes

Comment thread cpp/libcudf_streaming/src/partition_utils.cpp Outdated

madsbk added 4 commits August 28, 2026 08:38

no _impl

4b180ca

std::optional<std::size_t> packed_bytes

ca9ddf7

Merge branch 'main' of github.com:NVIDIA/cudf into let-callers-reserv…

795cabc

…e-partitioning-memory

cleanup

49f7e3e

madsbk requested a review from nirandaperera August 28, 2026 07:11

check_reservation

b7c8602

This comment was marked as resolved.

Sign in to view

madsbk added 2 commits August 28, 2026 14:43

Merge branch 'main' of github.com:NVIDIA/cudf into let-callers-reserv…

a3aa119

…e-partitioning-memory

Merge branch 'main' of github.com:NVIDIA/cudf into let-callers-reserv…

7e352b0

…e-partitioning-memory

This comment was marked as resolved.

Sign in to view

madsbk added 2 commits August 28, 2026 15:44

stream_ref

3de547a

docs

3654630

nirandaperera reviewed Aug 28, 2026

View reviewed changes

		rapidsmpf::BufferResource* br,
		rapidsmpf::MemoryReservation& reservation)

Uh oh!

Conversation

madsbk commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Notes

Uh oh!

coderabbitai Bot commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Summary by CodeRabbit

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

This comment was marked as resolved.

Uh oh!

wence- left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

nirandaperera left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nirandaperera Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

This comment was marked as resolved.

Uh oh!

nirandaperera left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

madsbk commented Aug 26, 2026 •

edited

Loading

coderabbitai Bot commented Aug 26, 2026 •

edited

Loading

nirandaperera Aug 27, 2026 •

edited

Loading