Avoid blocking memory reservations in actor coroutines - #23892
Conversation
madsbk
added
improvement
github-actions
Bot
added
Python
140c1e7 to
5826e82
Compare
📝 WalkthroughSummary by CodeRabbit
WalkthroughThe streaming actor graph now uses asynchronous chunk availability and explicit memory reservations. AllGather, shuffle, ordering, join, fanout, sink, origin-stamp, and union paths use the updated coordination flow. ChangesStreaming memory coordination
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: 🟡 Moderate · up to The PR replaces blocking memory reservations with asynchronous backpressure across streaming operations, but expanded hash-key expressions can exceed the reserved memory, and some collective paths consume data before a reservation wait that may fail or be cancelled. Under memory pressure this could cause failures or strand collective state, so the change needs targeted fixes or explicit owner acceptance before merge. Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
TomAugspurger
left a comment
TomAugspurger
left a comment
There was a problem hiding this comment.
Breaking change
Two functions become coroutines
These aren't part of the public cudf-polars API, so if you were basing the "breaking" label off those two I think we can remove it.
| # Representation change: the packed input is consumed as the | ||
| # unpacked table is produced, at roughly the same size. | ||
| reservation = await reserve_memory( | ||
| self.context, | ||
| unpack_and_concat_cost(partitions), | ||
| net_memory_delta=0, | ||
| ) |
There was a problem hiding this comment.
Can you explain this reservation? IIUC, we're making a change to ensure that we have room for both partitions (already in memory) and the unpack_and_concat form.
But the comment says that "the packed input is consumed as the unpacked table is produces". So are we freeing from partitions as the unpacked form is produced or not?
There was a problem hiding this comment.
Updated the comment, the same size was refering to the net_memory_delta argument. Fixed in: ccac165
It was, removed |
madsbk
added
non-breaking
afdcf5f to
a3aeb96
Compare
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
| reserve_extra=0, | ||
| net_memory_delta=0, | ||
| ) | ||
| self._manager.allgather.insert( |
There was a problem hiding this comment.
I cannot mark the whole block, because GitHub, but could we reserve and consume the packing allocation here as well? make_table_chunks_available_or_wait(..., reserve_extra=0) consumes its reservation just to make the chunk available, but chunk.into_packed_data() calls cudf::pack() for non-PackedData chunks. The local table_chunk API explicitly documents that this allocation is not reservation-tracked, so it can still synchronously spill/block the shared actor event loop. I think this should route packing through a reservation-aware path and cover an unpacked chunk under memory pressure.
There was a problem hiding this comment.
Done in #23894, kept separate because it is a breaking C++ and Cython change.
This comment was marked as resolved.
This comment was marked as resolved.
This comment was marked as resolved.
This comment was marked as resolved.
|
Thanks for the quick reviews! |
|
/merge |
3 participants
Every actor on a rank shares one event loop, so an actor that reserves device memory synchronously stalls all the others while it spills. This moves every such reservation that runs inside an actor coroutine off the blocking path. They take their memory from
reserve_memory()now, so a request that cannot be satisfied queues alongside the other actors' and is served by priority instead of spilling on the spot. That gives these sites memory backpressure as well as an unblocked loop, since an actor waiting on a reservation lets the ones that can release memory run first.Unspilling table chunks
TableChunk.make_available_and_spill()spills synchronously. Twelve call sites move to the awaitablemake_table_chunks_available_or_wait(), which suspends until a reservation is granted so the other actors can run and release memory meanwhile.AllGatherManager.Inserter.insertand_unpack_remote_partitionbecome coroutines, updating eight call sites acrossjoin.py,repartition.py,sort.py,utils.py,ordering.pyand the AllGather tests._unpack_remote_partitionalso takes the context rather than a buffer resource. Neither is public API.Behavior change
The unspill calls previously passed
allow_overbooking=True. They now fall through to theallow_overbooking_by_defaultconfiguration option, matching every existingmake_table_chunks_available_or_wait()call site. That option ships astrue, so under the default the only difference is that these sites wait for memory before overbooking rather than overbooking immediately. With the option set tofalsethey can raise where previously they could not.