Simplify async shuffle completion notification

cpp/benchmarks/bench_shuffle.cpp

@@ -312,8 +312,8 @@ rapidsmpf::Duration do_run(
         // insert partitions into the shuffler
         shuffle_insert_fn(shuffler);
 
-        while (!shuffler.finished()) {
-            auto finished_partition = shuffler.wait_any();
+        shuffler.wait();
+        for (auto finished_partition : shuffler.local_partitions()) {
             auto packed_chunks = shuffler.extract(finished_partition);
             auto output_partition = rapidsmpf::unpack_and_concat(
                 rapidsmpf::unspill_partitions(

cpp/benchmarks/streaming/ndsh/join.cpp

@@ -582,18 +582,15 @@ streaming::Actor shuffle(
     }
     co_await shuffler.insert_finished();
     for (auto pid : shuffler.local_partitions()) {
-        auto packed_data = co_await shuffler.extract_async(pid);
-        RAPIDSMPF_EXPECTS(packed_data.has_value(), "Partition already extracted");
+        auto packed_data = shuffler.extract(pid);
         auto stream = ctx->br()->stream_pool().get_stream();
         co_await ch_out->send(
             streaming::to_message(
                 pid,
                 std::make_unique<streaming::TableChunk>(
                     unpack_and_concat(
                         unspill_partitions(
-                            std::move(*packed_data),
-                            ctx->br().get(),
-                            AllowOverbooking::YES
+                            std::move(packed_data), ctx->br().get(), AllowOverbooking::YES
                         ),
                         stream,
                         ctx->br().get()

cpp/examples/example_shuffle.cpp

@@ -99,11 +99,11 @@ int main(int argc, char** argv) {
     // Vector to hold the local results of the shuffle operation.
     std::vector<std::unique_ptr<cudf::table>> local_outputs;
 
-    // Wait for and process the shuffle results for each partition.
-    while (!shuffler.finished()) {
-        // Block until a partition is ready and retrieve its partition ID.
-        rapidsmpf::shuffler::PartID finished_partition = shuffler.wait_any();
+    // Wait for all partitions to finish.
+    shuffler.wait();
 
+    // Process the shuffle results for each partition.
+    for (auto finished_partition : shuffler.local_partitions()) {
         // Extract the finished partition's data from the Shuffler.
         auto packed_chunks = shuffler.extract(finished_partition);
 

cpp/include/rapidsmpf/shuffler/finish_counter.hpp

@@ -9,8 +9,6 @@
 #include <functional>
 #include <mutex>
 #include <optional>
-#include <span>
-#include <unordered_set>
 #include <vector>
 
 #include <rapidsmpf/communicator/communicator.hpp>
@@ -44,27 +42,25 @@ namespace detail {
 class FinishCounter {
   public:
     /**
-     * @brief Callback function type called when a partition is finished.
-     *
-     * The callback receives the partition ID of the finished partition.
+     * @brief Callback function type called when all partitions are finished.
      *
      * @warning A callback must be fast and non-blocking and should not call any of the
      * `wait*` methods. And be very careful if acquiring locks. Ideally it should be used
      * to signal a separate thread to do the actual processing.
      */
-    using FinishedCallback = std::function<void(PartID)>;
+    using FinishedCallback = std::function<void()>;
 
     /**
      * @brief Construct a finish counter.
      *
      * @param nranks The total number of ranks participating in the shuffle.
-     * @param local_partitions The partition IDs local to the current rank.
-     * @param finished_callback The callback to notify when a partition is finished
+     * @param n_local_partitions The number of local partitions owned by this rank.
+     * @param finished_callback The callback to notify when all partitions are finished
      * (optional).
      */
     FinishCounter(
         Rank nranks,
-        std::span<PartID const> local_partitions,
+        PartID n_local_partitions,
         FinishedCallback&& finished_callback = nullptr
     );
 
@@ -103,51 +99,23 @@ class FinishCounter {
     [[nodiscard]] bool all_finished() const;
 
     /**
-     * @brief Returns the partition ID of a finished partition that hasn't been waited on
-     * (blocking). Optionally a timeout (in ms) can be provided.
+     * @brief Wait for all partitions to be finished (blocking). Optionally a timeout
+     * (in ms) can be provided.
      *
      * This function blocks until all partitions are finished and ready to be processed.
-     * If the timeout is set and a partition is not available within the specified
+     * If the timeout is set and the partitions are not finished within the specified
      * timeout, a std::runtime_error will be thrown.
      *
-     * @param timeout Optional timeout (ms) to wait.
-     *
-     * @note Due to the completion mechanism once `wait_any` returns any partition, all
-     * local partitions will be available for extraction. We previously supported
-     * per-partition completion mechanisms but since the usual usecase for a shuffle is a
-     * dense all to all this did not actually provide any additional concurrency. See also
-     * https://github.com/rapidsai/rapidsmpf/pull/914
-     *
-     * @return The partition ID of a finished partition.
-     *
-     * @throws std::out_of_range If all partitions have already been waited on.
-     * @throws std::runtime_error If timeout was set and no partitions have been finished
-     * by the expiration.
-     */
-    PartID wait_any(std::optional<std::chrono::milliseconds> timeout = {});
-
-    /**
-     * @brief Wait for a specific partition to be finished (blocking). Optionally a
-     * timeout (in ms) can be provided.
+     * @note We previously supported per-partition completion mechanisms but since the
+     * usual usecase for a shuffle is a dense all to all this did not actually provide any
+     * additional concurrency. See also https://github.com/rapidsai/rapidsmpf/pull/914
      *
-     * This function blocks until all partitions are finished and the desired partition
-     * is ready to be processed. If the timeout is set and the requested partition is not
-     * available within the specified timeout, a std::runtime_error will be thrown.
-     *
-     * @param pid The desired partition ID.
      * @param timeout Optional timeout (ms) to wait.
      *
-     * @note Due to the completion mechanism once `wait_on` returns successfully, all
-     * local partitions will be available for extraction. We previously supported
-     * per-partition completion mechanisms but since the usual usecase for a shuffle is a
-     * dense all to all this did not actually provide any additional concurrency. See also
-     * https://github.com/rapidsai/rapidsmpf/pull/914
-     *
-     * @throws std::out_of_range If the desired partition is unavailable.
-     * @throws std::runtime_error If timeout was set and requested partition has been
-     * finished by the expiration.
+     * @throws std::runtime_error If timeout was set and the partitions are not all ready
+     * by the expiration time.
      */
-    void wait_on(PartID pid, std::optional<std::chrono::milliseconds> timeout = {});
+    void wait(std::optional<std::chrono::milliseconds> timeout = {});
 
     /**
      * @brief Returns a description of this instance.
@@ -157,20 +125,15 @@ class FinishCounter {
 
   private:
     Rank const nranks_;
-    PartID
-        n_unfinished_partitions_;  ///< aux counter to track the number of unfinished
-                                   ///< partitions (without using the goalposts.empty())
+    PartID n_unfinished_partitions_;  ///< aux counter to track the number of unfinished
+                                      ///< partitions; set to zero when all chunks arrive
 
     Rank n_ranks_with_goalpost_{0};  ///< how many ranks have called move_goalpost
     ChunkID total_chunk_goal_{0};  ///< sum of all rank chunk goals
     ChunkID total_finished_chunks_{0};  ///< global finished chunk counter
     std::vector<bool> rank_reported_;  ///< indexed by rank, prevents double-reporting
-    std::span<PartID const> local_partitions_;  ///< for firing callbacks
-    /// Partitions not yet consumed by wait_any/wait_on; populated at construction and
-    /// then only ever decreases in size as partitions are consumed
-    std::unordered_set<PartID> pending_pids_;
     /// Set to true exactly once when all chunks have arrived. Ensures callback only fires
-    /// once for each partition.
+    /// once.
     bool all_done_{false};
 
     mutable std::mutex mutex_;  // TODO: use a shared_mutex lock?

cpp/include/rapidsmpf/shuffler/shuffler.hpp

@@ -110,7 +110,7 @@ class Shuffler {
      * and should not be reused until all nodes has called `Shuffler::shutdown()`.
      * @param total_num_partitions Total number of partitions in the shuffle.
      * @param br Buffer resource used to allocate temporary and the shuffle result.
-     * @param finished_callback Callback to notify when a partition is finished.
+     * @param finished_callback Callback to notify when all partitions are finished.
      * @param partition_owner Function to determine partition ownership.
      *
      * @note The caller promises that inserted buffers are stream-ordered with respect
@@ -191,7 +191,7 @@ class Shuffler {
      * It is valid to extract a partition that has not yet been fully received.
      * In such cases, only the chunks received so far are returned.
      *
-     * To ensure the partition is complete, use `wait_any()`, `wait_on()`,
+     * To ensure the partition is complete, use `wait()`
      * or another appropriate synchronization mechanism beforehand.
      *
      * @param pid The ID of the partition to extract.
@@ -207,25 +207,13 @@ class Shuffler {
     [[nodiscard]] bool finished() const;
 
     /**
-     * @brief Wait for any partition to finish.
+     * @brief Wait for all partitions to finish (blocking).
      *
      * @param timeout Optional timeout (ms) to wait.
      *
-     * @return The partition ID of the next finished partition.
-     *
-     * @throws std::runtime_error if the timeout is reached.
-     */
-    PartID wait_any(std::optional<std::chrono::milliseconds> timeout = {});
-
-    /**
-     * @brief Wait for a specific partition to finish (blocking).
-     *
-     * @param pid The desired partition ID.
-     * @param timeout Optional timeout (ms) to wait.
-     *
      * @throws std::runtime_error if the timeout is reached.
      */
-    void wait_on(PartID pid, std::optional<std::chrono::milliseconds> timeout = {});
+    void wait(std::optional<std::chrono::milliseconds> timeout = {});
 
     /**
      * @brief Spills data to device if necessary.

cpp/include/rapidsmpf/streaming/coll/shuffler.hpp

@@ -5,7 +5,7 @@
 
 #pragma once
 
-#include <unordered_set>
+#include <coro/event.hpp>
 
 #include <rapidsmpf/shuffler/shuffler.hpp>
 #include <rapidsmpf/streaming/core/actor.hpp>
@@ -16,36 +16,29 @@
 namespace rapidsmpf::streaming {
 
 /**
- * @brief An asynchronous shuffler that allows concurrent insertion and extraction of
- * data.
+ * @brief An asynchronous shuffler that wraps the synchronous shuffler with a coroutine
+ * interface.
  *
  * ShufflerAsync provides an asynchronous interface to the shuffler, allowing data to be
- * inserted while previously shuffled partitions are extracted concurrently. This is
- * useful for streaming scenarios where data can be processed as soon as individual
- * partitions are ready, rather than waiting for the entire shuffle to complete.
+ * inserted and then extracted after the shuffle completes. All local partitions complete
+ * simultaneously, so extraction is non-blocking after awaiting `insert_finished()`.
  *
- * Inserting the finished flags provides a token that one must await to "finalize"
- * extractions. One can asynchronously extract partitions before awaiting this token.
- *
- * @warning The finish token _must_ be awaited otherwise the shuffle will throw in
- * destruction or deadlocks will occur.
+ * @warning The coroutine returned by `insert_finished()` _must_ be awaited before the
+ * object is destroyed, otherwise the shuffle with terminate in destruction and/or
+ * deadlocks will occur.
  *
  * Example usage:
  * @code{.cpp}
  * auto shuffle = ShufflerAsync(...);
  * while (...) {
  *   shuffle.insert(...);
  * }
- * auto finished_token = shuffle.insert_finished();
- * for (auto i = 0; i < shuffle.local_partitions().size(); i++) {
- *   auto part = co_await shuffle.extract_any_async();
+ * co_await shuffle.insert_finished();
+ * for (auto pid : shuffle.local_partitions()) {
+ *   auto chunks = shuffle.extract(pid);
+ *   // process chunks...
  * }
- * co_await finished_token;
  * @endcode{}
- *
- * @note One can launch more extraction tasks than there are partitions to extract, for
- * example if we have multiple consumers of a shuffle, the extraction will return
- * `std::nullopt` if no more partitions are available.
  */
 class ShufflerAsync {
   public:
@@ -125,101 +118,29 @@ class ShufflerAsync {
     /**
      * @copydoc rapidsmpf::shuffler::Shuffler::insert_finished()
      *
-     * @note This function itself is not a coroutine. Instead, it returns a coroutine that
-     * must be awaited to ensure the shuffler has fully completed its asynchronous
-     * operations. Awaiting this coroutine guarantees that all notifications and
-     * background tasks in the underlying shuffler have finished before destruction. The
-     * coroutine does not need to be awaited before extraction begins, but it must
-     * eventually be awaited before the shuffle object is destroyed. Any pending
-     * extractions will wake up and either extract remaining partitions or return empty
-     * results if none remain.
+     * @note This coroutine function must be awaited to ensure the shuffler has fully
+     * completed its asynchronous operations.
      *
-     * @return A coroutine that, when awaited, indicates the shuffle has completed.
+     * @return A coroutine that inserts the finish marker and suspends until the shuffle
+     * has completed. Once complete,
      */
     [[nodiscard]] Actor insert_finished();
 
     /**
-     * @brief Asynchronously extracts all data for a specific partition.
-     *
-     * This coroutine suspends until the specified partition is ready for extraction
-     * (i.e., `insert_finished` has been called for this partition and all data has been
-     * shuffled).
-     *
-     * @warning Be careful when mixing `extract_async` and `extract_any_async`.
-     * A partition intended for `extract_async` may already have been consumed by
-     * `extract_any_async`, in which case this function returns `std::nullopt`.
-     *
-     * @param pid The partition ID to extract data for.
-     * @return
-     *   - `std::nullopt` if the partition ID is not ready or has already been extracted.
-     *   - Otherwise, a vector of `PackedData` chunks belonging to the partition.
-     *
-     * @throws std::out_of_range If the partition ID isn't owned by this rank, see
-     * `partition_owner()`.
-     */
-    [[nodiscard]] coro::task<std::optional<std::vector<PackedData>>> extract_async(
-        shuffler::PartID pid
-    );
-
-    /**
-     * @brief Result type for extract_any_async operations.
-     *
-     * Contains the partition ID and associated data chunks from an extract operation.
-     */
-    using ExtractResult = std::pair<shuffler::PartID, std::vector<PackedData>>;
-
-    /**
-     * @brief Asynchronously extracts data for any ready partition.
-     *
-     * This coroutine will suspend until at least one partition is ready for extraction,
-     * then extract and return the data for one such partition. If no partitions become
-     * ready and the shuffle is finished, returns a nullopt.
+     * @brief Extract all chunks belonging to the specified partition.
      *
-     * @return `ExtractResult` containing the partition ID and data chunks, or a nullopt
-     * if all partitions has been extracted.
-     *
-     * @warning Be careful when mixing `extract_async` and `extract_any_async`.
-     * A partition intended for `extract_async` may already have been consumed by
-     * `extract_any_async`, in which case `extract_async` will later return
-     * `std::nullopt`.
+     * @param pid The ID of the partition to extract.
+     * @throws std::logic_error If the partition has already been extracted or is
+     * otherwise not available.
+     * @return A vector of PackedData chunks associated with the partition.
      */
-    [[nodiscard]] coro::task<std::optional<ExtractResult>> extract_any_async();
+    [[nodiscard]] std::vector<PackedData> extract(shuffler::PartID pid);
 
   private:
-    /**
-     * @brief Ensure that all notifications have been received and drain pending
-     * extractions.
-     *
-     * This is required to ensure that all asynchronous notifications from the underlying
-     * shuffler have completed before the shuffle destructs. Any pending extractions will
-     * wake up and extract any remaining pids (or wake up empty if no pids are remaining).
-     *
-     * @note Typically this is not called directly, the coroutine it represents is
-     * returned from `insert_finished`.
-     *
-     * @return A coroutine representing the completion of all notifications and the
-     * shutdown of the semaphore.
-     */
-    [[nodiscard]] Actor finished_drain();
-
     std::shared_ptr<Context> ctx_;
-    coro::task_group<coro::thread_pool>
-        notifications_;  ///< Container tracking the notifications that have fired.
-    Semaphore semaphore_{0};  ///< Releases resources (inserted ready pids)
-    coro::latch
-        latch_;  ///< Tracks notifications so that we know when all have been received.
-    std::mutex mtx_;  ///< Protects modification of ready_pids_ and extracted_pids_
+    coro::event
+        event_{};  ///< Event tracking whether all data has arrived and can be extracted.
     shuffler::Shuffler shuffler_;
-
-    /**
-     * @brief Tracks partition states for extraction.
-     *
-     * A received partition's ID is always in exactly one of the two sets:
-     *   - `ready_pids_`: partitions ready for extraction but not yet extracted.
-     *   - `extracted_pids_`: partitions that have already been extracted.
-     */
-    std::unordered_set<shuffler::PartID> ready_pids_;
-    std::unordered_set<shuffler::PartID> extracted_pids_;
 };
 
 namespace actor {