Expose async collectives to Python (#685)

This adds bindings for both the streaming node versions, and the object-based implementation.

- Closes #682

Authors:
  - Lawrence Mitchell (https://github.com/wence-)

Approvers:
  - Mads R. B. Kristensen (https://github.com/madsbk)
  - Richard (Rick) Zamora (https://github.com/rjzamora)

URL: #685

Author: wence-

cpp/include/rapidsmpf/streaming/chunks/packed_data.hpp

@@ -12,46 +12,33 @@
 namespace rapidsmpf::streaming {
 
 /**
- * @brief Chunk of `PackedData`.
- */
-struct PackedDataChunk {
-    /**
-     * @brief Packed data payload.
-     */
-    PackedData data;
-};
-
-/**
- * @brief Generate a content description for a `PackedDataChunk`.
+ * @brief Generate a content description for `PackedData`.
  *
  * @param obj The object's content to describe.
  * @return A new content description.
  */
-inline ContentDescription get_content_description(PackedDataChunk const& obj) {
+inline ContentDescription get_content_description(PackedData const& obj) {
     return ContentDescription{
-        {{obj.data.data->mem_type(), obj.data.data->size}},
-        ContentDescription::Spillable::YES
+        {{obj.data->mem_type(), obj.data->size}}, ContentDescription::Spillable::YES
     };
 }
 
 /**
- * @brief Wrap a `PackedDataChunk` into a `Message`.
+ * @brief Wrap `PackedData` into a `Message`.
  *
  * @param sequence_number Ordering identifier for the message.
  * @param chunk The chunk to wrap into a message.
  * @return A `Message` encapsulating the provided chunk as its payload.
  */
-Message to_message(
-    std::uint64_t sequence_number, std::unique_ptr<PackedDataChunk> chunk
-) {
+Message to_message(std::uint64_t sequence_number, std::unique_ptr<PackedData> chunk) {
     auto cd = get_content_description(*chunk);
     return Message{
         sequence_number,
         std::move(chunk),
         cd,
         [](Message const& msg, MemoryReservation& reservation) -> Message {
-            auto const& self = msg.get<PackedDataChunk>();
-            auto chunk = std::make_unique<PackedDataChunk>(self.data.copy(reservation));
+            auto const& self = msg.get<PackedData>();
+            auto chunk = std::make_unique<PackedData>(self.copy(reservation));
             auto cd = get_content_description(*chunk);
             return Message{msg.sequence_number(), std::move(chunk), cd, msg.copy_cb()};
         }

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

@@ -12,7 +12,6 @@
 #include <rapidsmpf/allgather/allgather.hpp>
 #include <rapidsmpf/communicator/communicator.hpp>
 #include <rapidsmpf/memory/packed_data.hpp>
-#include <rapidsmpf/streaming/chunks/packed_data.hpp>
 #include <rapidsmpf/streaming/core/channel.hpp>
 #include <rapidsmpf/streaming/core/context.hpp>
 
@@ -61,7 +60,7 @@ class AllGather {
      * @param sequence_number The sequence number for this chunk.
      * @param chunk The chunk to insert.
      */
-    void insert(std::uint64_t sequence_number, PackedDataChunk&& chunk);
+    void insert(std::uint64_t sequence_number, PackedData&& chunk);
 
     /// @copydoc rapidsmpf::allgather::AllGather::insert_finished()
     void insert_finished();
@@ -76,7 +75,7 @@ class AllGather {
      * @return Coroutine that completes when all data is available for extraction and
      * returns the data.
      */
-    coro::task<std::vector<PackedDataChunk>> extract_all(Ordered ordered = Ordered::YES);
+    coro::task<std::vector<PackedData>> extract_all(Ordered ordered = Ordered::YES);
 
   private:
     coro::event
@@ -94,8 +93,8 @@ namespace node {
  * packed data received through `Channel`s.
  *
  * @param ctx The streaming context to use.
- * @param ch_in Input channel providing `PackedDataChunk`s to be gathered.
- * @param ch_out Output channel where the gathered `PackedDataChunk`s are sent.
+ * @param ch_in Input channel providing `PackedData`s to be gathered.
+ * @param ch_out Output channel where the gathered `PackedData`s are sent.
  * @param op_id Unique identifier for the operation.
  * @param ordered If the extracted data should be sent to the output channel with sequence
  * numbers corresponding to the global total order of input chunks. If yes, then the

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

@@ -125,7 +125,7 @@ class ShufflerAsync {
      *
      * @return A coroutine that, when awaited, indicates the shuffle has completed.
      */
-    [[nodiscard]] Node insert_finished(std::vector<shuffler::PartID>&& pids);
+    [[nodiscard]] Node insert_finished();
 
     /**
      * @brief Asynchronously extracts all data for a specific partition.

cpp/src/streaming/coll/allgather.cpp

@@ -3,8 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-#include <algorithm>
-
 #include <rapidsmpf/memory/packed_data.hpp>
 #include <rapidsmpf/streaming/chunks/packed_data.hpp>
 #include <rapidsmpf/streaming/coll/allgather.hpp>
@@ -41,27 +39,19 @@ AllGather::~AllGather() {
     return ctx_;
 }
 
-void AllGather::insert(std::uint64_t sequence_number, PackedDataChunk&& packed_data) {
-    gatherer_.insert(sequence_number, std::move(packed_data.data));
+void AllGather::insert(std::uint64_t sequence_number, PackedData&& packed_data) {
+    gatherer_.insert(sequence_number, std::move(packed_data));
 }
 
 void AllGather::insert_finished() {
     gatherer_.insert_finished();
 }
 
-coro::task<std::vector<PackedDataChunk>> AllGather::extract_all(
-    AllGather::Ordered ordered
-) {
+coro::task<std::vector<PackedData>> AllGather::extract_all(AllGather::Ordered ordered) {
     // Wait until we're notified that everything is done.
     co_await event_;
     // And now this will not block.
-    auto data = gatherer_.wait_and_extract(ordered);
-    std::vector<PackedDataChunk> result;
-    result.reserve(data.size());
-    std::ranges::transform(data, std::back_inserter(result), [](auto&& pd) {
-        return PackedDataChunk{.data = std::move(pd)};
-    });
-    co_return result;
+    co_return gatherer_.wait_and_extract(ordered);
 }
 
 namespace node {
@@ -80,14 +70,14 @@ Node allgather(
         if (msg.empty()) {
             break;
         }
-        gatherer.insert(msg.sequence_number(), msg.release<PackedDataChunk>());
+        gatherer.insert(msg.sequence_number(), msg.release<PackedData>());
     }
     gatherer.insert_finished();
     auto data = co_await gatherer.extract_all(ordered);
     std::uint64_t sequence{0};
     for (auto&& chunk : data) {
         co_await ch_out->send(
-            to_message(sequence++, std::make_unique<PackedDataChunk>(std::move(chunk)))
+            to_message(sequence++, std::make_unique<PackedData>(std::move(chunk)))
         );
     }
     co_await ch_out->drain(ctx->executor());

cpp/src/streaming/coll/shuffler.cpp

@@ -153,7 +153,9 @@ void ShufflerAsync::insert(std::unordered_map<shuffler::PartID, PackedData>&& ch
     shuffler_.insert(std::move(chunks));
 }
 
-Node ShufflerAsync::insert_finished(std::vector<shuffler::PartID>&& pids) {
+Node ShufflerAsync::insert_finished() {
+    std::vector<shuffler::PartID> pids(total_num_partitions());
+    std::iota(pids.begin(), pids.end(), shuffler::PartID{0});
     shuffler_.insert_finished(std::move(pids));
     return finished_drain();
 }
@@ -258,12 +260,7 @@ Node shuffler(
         shuffler_async.insert(std::move(partition_map.data));
     }
 
-    // Tell the shuffler that we have no more input data.
-    std::vector<rapidsmpf::shuffler::PartID> finished(
-        shuffler_async.total_num_partitions()
-    );
-    std::iota(finished.begin(), finished.end(), 0);
-    auto finish_token = shuffler_async.insert_finished(std::move(finished));
+    auto finish_token = shuffler_async.insert_finished();
 
     for ([[maybe_unused]] auto& _ : shuffler_async.local_partitions()) {
         auto finished = co_await shuffler_async.extract_any_async();

cpp/tests/streaming/test_allgather.cpp

@@ -87,9 +87,7 @@ TEST_P(StreamingAllGather, basic) {
         });
         auto meta = std::make_unique<std::vector<std::uint8_t>>(sizeof(int));
         std::memcpy(meta->data(), &size, sizeof(int));
-        allgather.insert(
-            sequence, streaming::PackedDataChunk{{std::move(meta), std::move(buf)}}
-        );
+        allgather.insert(sequence, PackedData{std::move(meta), std::move(buf)});
         latch.count_down();
         co_return;
     };
@@ -104,20 +102,20 @@ TEST_P(StreamingAllGather, basic) {
         std::size_t offset{0};
         for (auto& pd : data) {
             RAPIDSMPF_EXPECTS(
-                pd.data.metadata->size() == sizeof(int), "Invalid metadata buffer size"
+                pd.metadata->size() == sizeof(int), "Invalid metadata buffer size"
             );
             int msize;
-            std::memcpy(&msize, pd.data.metadata->data(), sizeof(int));
+            std::memcpy(&msize, pd.metadata->data(), sizeof(int));
             RAPIDSMPF_EXPECTS(msize == size, "Corrupted metadata value");
             RAPIDSMPF_CUDA_TRY(cudaMemcpyAsync(
                 result.data() + offset,
-                pd.data.data->data(),
-                pd.data.data->size,
+                pd.data->data(),
+                pd.data->size,
                 cudaMemcpyDefault,
-                pd.data.data->stream()
+                pd.data->stream()
             ));
             offset += msize;
-            pd.data.data->stream().synchronize();
+            pd.data->stream().synchronize();
         }
     };
 
@@ -170,9 +168,7 @@ TEST_P(StreamingAllGather, streaming_node) {
         input_messages.emplace_back(
             streaming::to_message(
                 insertion_id,
-                std::make_unique<streaming::PackedDataChunk>(streaming::PackedDataChunk{
-                    PackedData{std::move(meta), std::move(buf)}
-                })
+                std::make_unique<PackedData>(std::move(meta), std::move(buf))
             )
         );
     }
@@ -191,8 +187,7 @@ TEST_P(StreamingAllGather, streaming_node) {
     std::vector<int> actual(size * size * n_inserts);
     std::size_t offset{0};
     for (auto& msg : output_messages) {
-        auto chunk = msg.release<streaming::PackedDataChunk>();
-        auto& pd = chunk.data;
+        auto pd = msg.release<PackedData>();
         RAPIDSMPF_EXPECTS(
             pd.metadata->size() == sizeof(int), "Invalid metadata buffer size"
         );

cpp/tests/streaming/test_shuffler.cpp

@@ -232,8 +232,7 @@ TEST_P(ShufflerAsyncTest, multi_consumer_extract) {
         shuffler->insert(std::move(data));
     }
 
-    auto finish_token =
-        shuffler->insert_finished(iota_vector<shuffler::PartID>(n_partitions));
+    auto finish_token = shuffler->insert_finished();
 
     std::mutex mtx;
     std::vector<shuffler::PartID> finished_pids;
@@ -264,8 +263,7 @@ TEST_F(BaseStreamingShuffle, extract_any_before_extract) {
         auto shuffler = std::make_unique<ShufflerAsync>(ctx, op_id, n_partitions);
 
         // all empty partitions
-        auto finish_token =
-            shuffler->insert_finished(iota_vector<shuffler::PartID>(n_partitions));
+        auto finish_token = shuffler->insert_finished();
 
         auto local_pids = shuffler::Shuffler::local_partitions(
             ctx->comm(), n_partitions, shuffler::Shuffler::round_robin
@@ -313,8 +311,7 @@ class CompetingShufflerAsyncTest : public BaseStreamingShuffle {
 
         auto shuffler = std::make_unique<ShufflerAsync>(ctx, op_id, n_partitions);
 
-        auto finish_token =
-            shuffler->insert_finished(iota_vector<shuffler::PartID>(n_partitions));
+        auto finish_token = shuffler->insert_finished();
         coro::sync_wait(finish_token);
         auto [extract_any_result, extract_result] =
             produce_results_fn(shuffler.get(), this_pid);

python/rapidsmpf/rapidsmpf/shuffler.pxd

@@ -37,6 +37,24 @@ cdef extern from "<rapidsmpf/shuffler/shuffler.hpp>" nogil:
         uint32_t wait_any() except +
         void wait_on(uint32_t pid) except +
         string str() except +
+# Insert PackedData into a partition map. We implement this in C++ because
+# PackedData doesn't have a default ctor.
+cdef extern from *:
+    """
+    void cpp_insert_chunk_into_partition_map(
+        std::unordered_map<std::uint32_t, rapidsmpf::PackedData> &partition_map,
+        std::uint32_t pid,
+        std::unique_ptr<rapidsmpf::PackedData> packed_data
+    ) {
+        partition_map.insert({pid, std::move(*packed_data)});
+    }
+    """
+    void cpp_insert_chunk_into_partition_map(
+        unordered_map[uint32_t, cpp_PackedData] &partition_map,
+        uint32_t pid,
+        unique_ptr[cpp_PackedData] packed_data,
+    ) except + nogil
+
 
 cdef class Shuffler:
     cdef unique_ptr[cpp_Shuffler] _handle

python/rapidsmpf/rapidsmpf/shuffler.pyx

@@ -15,25 +15,6 @@ from rapidsmpf.progress_thread cimport ProgressThread
 from rapidsmpf.statistics cimport Statistics
 
 
-# Insert PackedData into a partition map. We implement this in C++ because
-# PackedData doesn't have a default ctor.
-cdef extern from *:
-    """
-    void cpp_insert_chunk_into_partition_map(
-        std::unordered_map<std::uint32_t, rapidsmpf::PackedData> &partition_map,
-        std::uint32_t pid,
-        std::unique_ptr<rapidsmpf::PackedData> packed_data
-    ) {
-        partition_map.insert({pid, std::move(*packed_data)});
-    }
-    """
-    void cpp_insert_chunk_into_partition_map(
-        unordered_map[uint32_t, cpp_PackedData] &partition_map,
-        uint32_t pid,
-        unique_ptr[cpp_PackedData] packed_data,
-    ) except + nogil
-
-
 cdef class Shuffler:
     """
     Shuffle service for partitioned data.

python/rapidsmpf/rapidsmpf/streaming/chunks/CMakeLists.txt

@@ -5,7 +5,7 @@
 # cmake-format: on
 # =================================================================================
 
-set(cython_modules arbitrary.pyx partition.pyx utils.pyx)
+set(cython_modules arbitrary.pyx packed_data.pyx partition.pyx utils.pyx)
 
 rapids_cython_create_modules(
   CXX