Apply memory backpressure to the shuffle - #23886
Conversation
madsbk
added
improvement
github-actions
Bot
added
Python
e3090b3 to
a3f477a
Compare
|
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:
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)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughSummary by CodeRabbit
WalkthroughThe shuffle implementation now performs memory reservation, partitioning, packing, and unpacking asynchronously. Repartitioning and dependent operators await the new APIs. Window routing uses asynchronous indexed insertion. Shuffler tests cover asynchronous hash and index workflows. ChangesAsynchronous shuffle memory management
Estimated code review effort: 3 (Moderate) | ~25 minutes Merge Risk: ⚪ Minimal · up to The shuffle memory-reservation behavior and related async interfaces introduce no actionable merge-blocking risk at the current head. 🚥 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.
fbe65af to
4abe1d5
Compare
| # Three allocations of the chunk's packed size: the key table, the | ||
| # reorder, then the pack. | ||
| chunk_nbytes = py_split_and_pack_cost(chunk.table_view(), chunk.stream, br) | ||
| reservation = await reserve_memory( | ||
| self._manager.context, | ||
| 3 * chunk_nbytes, | ||
| # The chunk's data moves into shuffler-owned packed buffers, | ||
| # nothing lasting is added. | ||
| net_memory_delta=0, | ||
| ) |
There was a problem hiding this comment.
Could we add a brief comment explaining the 3 * chunk_nbytes assumption? Pointwise key expressions can expand their input, for example, string padding, replacement, or concatenation, so the evaluated key table is not necessarily bounded by the input chunk’s packed size. I do not think this PR needs to solve general expression-aware size estimation, but documenting the expected workload/invariant here would make the trade-off explicit and help us revisit it if this path encounters memory-pressure failures.
And could we also add a bounded-memory test for a computed shuffle key? It need not cover every expansion case, one representative expression that materializes a new key column would verify that this path waits for memory and completes correctly, rather than allocating outside the backpressure mechanism.
There was a problem hiding this comment.
Good point, added a comment explaining the 3 * chunk_nbytes assumption.
On the test, I tried a few approaches but they all ended up too intrusive and tied to internal details. A bounded-memory test does not actually distinguish the two cases here, since the buffer resource's limit is an accounting limit rather than an allocation barrier, so the unreserved version allocates and completes just the same. Asserting on the reservation itself works, but only by monkeypatching reserve_memory and pinning the exact multiplier, which locks the test to the current internals of this path.
Given the estimate is an approximation rather than a bound, I do not think a detailed test is worth it. The computed-key path itself is already covered end to end by test_join_non_col_keys_rapidsmpf.
2 participants
The shuffle now takes its device memory from
reserve_memory()on both sides, the fourShuffleManager.Insertermethods going in andextract_chunkcoming out, so a request that cannot be satisfied queues alongside the other actors' and is served by priority. Before this both let the C++ side reserve and spill internally. cudf-polars already reserves this way for scans inactor_graph/io.py.Breaking change
The insert methods and
extract_chunkare coroutines now, andLocalRepartitioner._iter_chunksis an async generator. That updates twenty-one call sites acrossgroupby.py,over.py,sort.py,shuffle.pyand the shuffler tests.Notes
Reserving inside the insert methods rather than at the call sites keeps the accounting in one place. That matters for
insert_hash_keysandinsert_index, whose Python-side reorder allocates a full table copy that nothing reserved before. They reservepartition_and_pack_cost(), which covers a reorder plus a pack, consume the reorder's share once it lands, and hand the remainder tosplit_and_pack().The cost functions call
cudf::packed_size(), which syncs the stream, and they run on the actor event loop.insert_hashandinsert_splitadd one sync per chunk,insert_hash_keysandinsert_indextwo, since they size the reorder separately from the pack. Passing apacked_byteshint through the Cython layer would remove the second one, andpartition_and_packalready takes that hint in C++._iter_chunksreserves per piece, so theTODOthere about batching pieces up totarget_partition_sizewould cut the reservation traffic as well as the unpack overhead.