Python bindings for the streaming shuffler (rapidsai#490)

Also adding bindings to `PartitionVectorChunk`

Authors:
  - Mads R. B. Kristensen (https://github.com/madsbk)

Approvers:
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)
  - Peter Andreas Entschev (https://github.com/pentschev)
  - Lawrence Mitchell (https://github.com/wence-)

URL: rapidsai#490

Author: madsbk

cpp/benchmarks/streaming/data_generator.hpp

@@ -29,7 +29,7 @@ namespace rapidsmpf::streaming::node {
  * distributed in the range [`min_val`, `max_val`]. Each generated table is wrapped
  * in a `TableChunk` and sent to the provided output channel in streaming fashion.
  *
- * @param ctx The context to use.
+ * @param ctx The node context to use.
  * @param stream The CUDA stream on which to create the random tables. TODO: use a pool
  * of CUDA streams.
  * @param ch_out Output channel to which generated `TableChunk` objects are sent.

cpp/include/rapidsmpf/streaming/core/leaf_node.hpp

@@ -18,7 +18,7 @@ namespace rapidsmpf::streaming::node {
  * Sends each message of the input vector into the channel in order,
  * marking the end of the stream once done.
  *
- * @param ctx The context to use.
+ * @param ctx The node context to use.
  * @param ch_out Output channel to which messages will be sent.
  * @param messages Input vector containing the messages to send.
  * @return Streaming node representing the asynchronous operation.
@@ -37,7 +37,7 @@ Node push_to_channel(
  * Receives messages from the channel until it is closed and appends them
  * to the provided output vector.
  *
- * @param ctx The context to use.
+ * @param ctx The node context to use.
  * @param ch_in Input channel providing messages.
  * @param out_messages Output vector to store the received messages.
  * @return Streaming node representing the asynchronous operation.

cpp/include/rapidsmpf/streaming/cudf/partition.hpp

@@ -83,7 +83,7 @@ namespace node {
  * `num_partitions` based on a hash of the selected columns, packs the resulting
  * partitions, and sends them to an output channel.
  *
- * @param ctx The context to use.
+ * @param ctx The node context to use.
  * @param ch_in Input channel providing `TableChunk`s to partition.
  * @param ch_out Output channel to which `PartitionMapChunk`s are sent.
  * @param columns_to_hash Indices of input columns to hash.
@@ -119,7 +119,7 @@ Node partition_and_pack(
  * them, and sends the resulting tables to the output channel. Empty partitions are
  * ignored.
  *
- * @param ctx The context to use.
+ * @param ctx The node context to use.
  * @param ch_in Input channel providing packed partitions as PartitionMapChunk or
  * PartitionVectorChunk.
  * @param ch_out Output channel to which unpacked and concatenated tables table are sent.

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

@@ -23,8 +23,7 @@ namespace rapidsmpf::streaming::node {
  * It consumes partitioned input data from the input channel and produces output chunks
  * grouped by `partition_owner`.
  *
- * @param ctx The streaming context providing communication, memory, stream, and execution
- * resources.
+ * @param ctx The context to use.
  * @param stream The CUDA stream on which to perform the shuffling. If chunks from the
  * input channel aren't created on `stream`, the streams are all synchronized.
  * @param ch_in Input channel providing PartitionMapChunk to be shuffled.

python/rapidsmpf/rapidsmpf/streaming/core/context.pyx

@@ -64,7 +64,7 @@ cdef class Context:
         return self._options
 
     def comm(self):
-        return self._options
+        return self._comm
 
     def br(self):
         return self._br

python/rapidsmpf/rapidsmpf/streaming/cudf/CMakeLists.txt

@@ -3,7 +3,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # =================================================================================
 
-set(cython_modules partition.pyx partition_chunk.pyx table_chunk.pyx)
+set(cython_modules partition.pyx partition_chunk.pyx shuffler.pyx table_chunk.pyx)
 
 rapids_cython_create_modules(
   CXX

python/rapidsmpf/rapidsmpf/streaming/cudf/partition.pyi

@@ -6,7 +6,10 @@ from collections.abc import Iterable
 from rapidsmpf.streaming.core.channel import Channel
 from rapidsmpf.streaming.core.context import Context
 from rapidsmpf.streaming.core.node import CppNode
-from rapidsmpf.streaming.cudf.partition_chunk import PartitionMapChunk
+from rapidsmpf.streaming.cudf.partition_chunk import (
+    PartitionMapChunk,
+    PartitionVectorChunk,
+)
 from rapidsmpf.streaming.cudf.table_chunk import TableChunk
 
 def partition_and_pack(
@@ -18,6 +21,8 @@ def partition_and_pack(
 ) -> CppNode: ...
 def unpack_and_concat(
     ctx: Context,
-    ch_in: Channel[PartitionMapChunk],
+    ch_in: Channel[PartitionMapChunk]
+    | Channel[PartitionVectorChunk]
+    | Channel[PartitionMapChunk | PartitionVectorChunk],
     ch_out: Channel[TableChunk],
 ) -> CppNode: ...

python/rapidsmpf/rapidsmpf/streaming/cudf/partition_chunk.pxd

@@ -4,6 +4,7 @@
 from libc.stdint cimport uint32_t, uint64_t
 from libcpp.memory cimport unique_ptr
 from libcpp.unordered_map cimport unordered_map
+from libcpp.vector cimport vector
 from rmm.librmm.cuda_stream_view cimport cuda_stream_view
 from rmm.pylibrmm.stream cimport Stream
 
@@ -16,6 +17,11 @@ cdef extern from "<rapidsmpf/streaming/cudf/partition.hpp>" nogil:
         unordered_map[uint32_t, cpp_PackedData] data
         cuda_stream_view stream
 
+    cdef cppclass cpp_PartitionVectorChunk "rapidsmpf::streaming::PartitionVectorChunk":
+        uint64_t sequence_number
+        vector[cpp_PackedData] data
+        cuda_stream_view stream
+
 
 cdef class PartitionMapChunk:
     cdef unique_ptr[cpp_PartitionMapChunk] _handle
@@ -28,3 +34,16 @@ cdef class PartitionMapChunk:
     )
     cdef const cpp_PartitionMapChunk* handle_ptr(self)
     cdef unique_ptr[cpp_PartitionMapChunk] release_handle(self)
+
+
+cdef class PartitionVectorChunk:
+    cdef unique_ptr[cpp_PartitionVectorChunk] _handle
+    cdef Stream _stream
+    cdef object _owner
+
+    @staticmethod
+    cdef PartitionVectorChunk from_handle(
+        unique_ptr[cpp_PartitionVectorChunk] handle, Stream stream, object owner
+    )
+    cdef const cpp_PartitionVectorChunk* handle_ptr(self)
+    cdef unique_ptr[cpp_PartitionVectorChunk] release_handle(self)

python/rapidsmpf/rapidsmpf/streaming/cudf/partition_chunk.pyi

@@ -17,7 +17,20 @@ class PartitionMapChunk:
     @property
     def stream(self) -> Stream: ...
 
+class PartitionVectorChunk:
+    @classmethod
+    def from_message(cls, message: Message[Self]) -> Self: ...
+    def into_message(self, message: Message[Self]) -> None: ...
+    @property
+    def sequence_number(self) -> int: ...
+    @property
+    def stream(self) -> Stream: ...
+
 if TYPE_CHECKING:
     # Check that PartitionMapChunk implements Payload.
     t1: PartitionMapChunk
     t2: Payload = t1
+
+    # Check that PartitionVectorChunk implements Payload.
+    t3: PartitionVectorChunk
+    t4: Payload = t3

python/rapidsmpf/rapidsmpf/streaming/cudf/partition_chunk.pyx

@@ -32,8 +32,8 @@ cdef class PartitionMapChunk:
             The CUDA stream on which this chunk was created. If `None`,
             the stream is obtained from the handle.
         owner
-            Python object that owns the underlying buffers and must
-            be kept alive for the lifetime of this PartitionMapChunk.
+            An optional Python object to keep alive for as long as this
+            PartitionMapChunk exists (e.g., to maintain resource lifetime).
 
         Returns
         -------
@@ -156,3 +156,154 @@ cdef class PartitionMapChunk:
             The CUDA stream.
         """
         return self._stream
+
+
+cdef class PartitionVectorChunk:
+    def __init__(self):
+        raise ValueError("use the `from_*` factory functions")
+
+    def __dealloc__(self):
+        with nogil:
+            self._handle.reset()
+
+    @staticmethod
+    cdef PartitionVectorChunk from_handle(
+        unique_ptr[cpp_PartitionVectorChunk] handle, Stream stream, object owner
+    ):
+        """
+        Construct a PartitionVectorChunk from an existing C++ handle.
+
+        Parameters
+        ----------
+        handle
+            A unique pointer to a C++ PartitionVectorChunk.
+        stream
+            The CUDA stream on which this chunk was created. If `None`,
+            the stream is obtained from the handle.
+        owner
+            An optional Python object to keep alive for as long as this
+            PartitionVectorChunk exists (e.g., to maintain resource lifetime).
+
+        Returns
+        -------
+        A new PartitionVectorChunk wrapping the given handle.
+        """
+
+        if stream is None:
+            stream = Stream._from_cudaStream_t(
+                deref(handle).stream.value()
+            )
+        cdef PartitionVectorChunk ret = PartitionVectorChunk.__new__(
+            PartitionVectorChunk
+        )
+        ret._handle = move(handle)
+        ret._stream = stream
+        ret._owner = owner
+        return ret
+
+    @staticmethod
+    def from_message(Message message not None):
+        """
+        Construct a PartitionVectorChunk by consuming a Message.
+
+        Parameters
+        ----------
+        message
+            Message containing a PartitionVectorChunk. The message is released
+            and is empty after this call.
+
+        Returns
+        -------
+        A new PartitionVectorChunk extracted from the given message.
+        """
+        return PartitionVectorChunk.from_handle(
+            make_unique[cpp_PartitionVectorChunk](
+                message._handle.release[cpp_PartitionVectorChunk]()
+            ),
+            stream = None,
+            owner = None,
+        )
+
+    def into_message(self, Message message not None):
+        """
+        Move this PartitionVectorChunk into a Message.
+
+        This method is not typically called directly. Instead, it is invoked by
+        `Message.__init__()` when creating a new Message with this PartitionVectorChunk
+        as its payload.
+
+        Parameters
+        ----------
+        message
+            Message object that will take ownership of this PartitionVectorChunk.
+
+        Raises
+        ------
+        ValueError
+            If the provided message is not empty.
+
+        Warnings
+        --------
+        The PartitionVectorChunk is released and must not be used after this call.
+        """
+        message._handle = cpp_Message(self.release_handle())
+
+    cdef const cpp_PartitionVectorChunk* handle_ptr(self):
+        """
+        Return a pointer to the underlying C++ PartitionVectorChunk.
+
+        Returns
+        -------
+        Raw pointer to the underlying C++ object.
+
+        Raises
+        ------
+        ValueError
+            If the PartitionVectorChunk is uninitialized.
+        """
+        if not self._handle:
+            raise ValueError("is uninitialized, has it been released?")
+        return self._handle.get()
+
+    cdef unique_ptr[cpp_PartitionVectorChunk] release_handle(self):
+        """
+        Release ownership of the underlying C++ PartitionVectorChunk.
+
+        After this call, the current object is in a moved-from state and
+        must not be accessed.
+
+        Returns
+        -------
+        Unique pointer to the underlying C++ object.
+
+        Raises
+        ------
+        ValueError
+            If the PartitionVectorChunk is uninitialized.
+        """
+        if not self._handle:
+            raise ValueError("is uninitialized, has it been released?")
+        return move(self._handle)
+
+    @property
+    def sequence_number(self):
+        """
+        Return the sequence number of this chunk.
+
+        Returns
+        -------
+        The sequence number.
+        """
+        return deref(self.handle_ptr()).sequence_number
+
+    @property
+    def stream(self):
+        """
+        Return the CUDA stream on which this chunk was created.
+
+        Returns
+        -------
+        Stream
+            The CUDA stream.
+        """
+        return self._stream