Avoid blocking memory reservations in actor coroutines - #23892
Conversation
madsbk
added
improvement
github-actions
Bot
added
Python
140c1e7 to
5826e82
Compare
|
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 (2)
🚧 Files skipped from review as they are similar to previous changes (2)
Included review availability: Your plan provides up to 12 included reviews per hour; 10 remain after this review. 📝 WalkthroughSummary by CodeRabbit
WalkthroughThe streaming actor graph now uses asynchronous chunk availability and explicit memory reservations. AllGather packing and remote unpacking reserve temporary memory before processing. Collective, 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: ⚪ Minimal · up to The change moves actor memory reservations onto an awaitable path without any supplied actionable merge-blocking risk; it is merge-ready after normal checks and review. 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
…ervations-in-actor-coroutines
2 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.