Let callers reserve partitioning memory - #23833
Conversation
madsbk
added
improvement
github-actions
Bot
added
libcudf
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe 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. ChangesPartition memory APIs
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: 🔴 Critical · up to 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: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation 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 checkExplanation 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. ✨ 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.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
`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
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.
nirandaperera
left a comment
nirandaperera
left a comment
There was a problem hiding this comment.
@madsbk Last change requests! Promise :inno
|
Thanks all |
|
/merge |
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.