Simplify async shuffle completion notification - #916
Merged
rapids-bot[bot] merged 8 commits intoMar 16, 2026
Conversation
Since all partitions now complete simultaneously, it does not make sense to offer an interface to wait for an individual partition.
Since these now just defer to wait and wait for all partitions to be ready, they are no longer necessary.
Now that all partitions complete together, we only need the finish callback to fire once.
Now that the shuffle callback fires only once, we can use the same suspend and wake mechanism that the async allgather uses. The async shuffle now holds an event that is set by the finish callback when all partitions are complete. We must now `co_await insert_finished()` and can then extract all of our local partitions without blocking.
We don't have competing coroutines now, so no need for parameterising over the thread pool size.
wence-
force-pushed
the
wence/fea/shuffle-single-notify
branch
from
March 13, 2026 17:28
074b201 to
baac65d
Compare
Member
|
Looks like we have to update/remove the |
madsbk
reviewed
Mar 16, 2026
Contributor
Author
I mean, that code is not touched in this PR at all, but yes that test looks mostly nonsense |
Sleeping the main thread and letting the periodic spill thread in the background provides no guarantee on how the OS time-slices the execution, so we can't really check anything here.
Contributor
Author
Removed that test |
Contributor
Author
|
/merge |
3 tasks
rapids-bot Bot
pushed a commit
to NVIDIA/cudf
that referenced
this pull request
Mar 16, 2026
Since rapidsai/rapidsmpf#916, extraction is no longer a coroutine, one must just `await insert_finished()` that suspends until all partitions are ready. Thereafter we just `extract_chunk` without suspending. Authors: - Lawrence Mitchell (https://github.com/wence-) Approvers: - Mads R. B. Kristensen (https://github.com/madsbk) - Tom Augspurger (https://github.com/TomAugspurger) URL: #21787
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Now that all shuffle partitions finish simultaneously, we can change the completion callback mechanism so that the completion callback fires exactly once, when all partitions are ready.
This makes tracking completion in the async shuffle much simpler, we can use the same setup as the async allgather: awaiting an event returned by
insert_finished(). Once complete, extraction of any individual partition no longer blocks.extracttherefore just allows extraction by partition ID and is no longer a coroutine, since it doesn't make sense as a suspension point any longer.