Merge branch 'branch-25.10' of github.com:rapidsai/rapidsmpf into shuffler-callbacks

Author: nirandaperera

.github/workflows/build.yaml

@@ -34,7 +34,7 @@ concurrency:
 jobs:
   conda-cpp-build:
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -44,7 +44,7 @@ jobs:
   conda-python-build:
     needs: [conda-cpp-build]
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -54,7 +54,7 @@ jobs:
   upload-conda:
     needs: [conda-cpp-build, conda-python-build]
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/conda-upload-packages.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -64,7 +64,7 @@ jobs:
     if: github.ref_type == 'branch'
     needs: conda-python-build
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10
     with:
       arch: "amd64"
       branch: ${{ inputs.branch }}
@@ -76,7 +76,7 @@ jobs:
       sha: ${{ inputs.sha }}
   wheel-build-librapidsmpf:
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -90,7 +90,7 @@ jobs:
   wheel-publish-librapidsmpf:
     needs: wheel-build-librapidsmpf
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -101,7 +101,7 @@ jobs:
   wheel-build-rapidsmpf:
     needs: wheel-build-librapidsmpf
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -113,7 +113,7 @@ jobs:
   wheel-publish-rapidsmpf:
     needs: wheel-build-rapidsmpf
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-publish.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}
@@ -124,7 +124,7 @@ jobs:
   wheel-build-rapidsmpf-singlecomm:
     needs: wheel-build-rapidsmpf
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type || 'branch' }}
       branch: ${{ inputs.branch }}

.github/workflows/pr.yaml

@@ -26,13 +26,13 @@ jobs:
       - wheel-test
       - devcontainer
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-25.10
     if: always()
     with:
       needs: ${{ toJSON(needs) }}
   changed-files:
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/changed-files.yaml@branch-25.10
     with:
       files_yaml: |
         test_cpp:
@@ -58,11 +58,11 @@ jobs:
           - '!docs/**'
   checks:
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/checks.yaml@branch-25.10
   wheel-build-librapidsmpf:
     needs: checks
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10
     with:
       build_type: pull-request
       script: ci/build_wheel_librapidsmpf.sh
@@ -73,7 +73,7 @@ jobs:
   wheel-build-rapidsmpf:
     needs: [checks, wheel-build-librapidsmpf]
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10
     with:
       build_type: pull-request
       script: ci/build_wheel_rapidsmpf.sh
@@ -82,7 +82,7 @@ jobs:
   wheel-build-rapidsmpf-singlecomm:
     needs: [checks, wheel-build-librapidsmpf]
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-build.yaml@branch-25.10
     with:
       build_type: pull-request
       script: ci/build_wheel_singlecomm.sh
@@ -92,22 +92,22 @@ jobs:
     if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python
     needs: [changed-files, wheel-build-rapidsmpf]
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.10
     with:
       build_type: pull-request
       container-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g --ulimit=nofile=1000000:1000000"
       script: ci/test_wheel.sh
   conda-cpp-build:
     needs: checks
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-build.yaml@branch-25.10
     with:
       build_type: pull-request
       script: ci/build_cpp.sh
   conda-cpp-linters:
     secrets: inherit
     needs: checks
-    uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10
     with:
       build_type: pull-request
       script: "ci/cpp_linters.sh"
@@ -116,23 +116,23 @@ jobs:
     if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp
     needs: [changed-files, conda-cpp-build]
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.10
     with:
       build_type: pull-request
       container-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g --ulimit=nofile=1000000:1000000"
       script: ci/test_cpp.sh
   conda-python-build:
     needs: conda-cpp-build
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/conda-python-build.yaml@branch-25.10
     with:
       build_type: pull-request
       script: ci/build_python.sh
   conda-python-tests:
     if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python
     needs: [changed-files, conda-python-build]
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.10
     with:
       build_type: pull-request
       container-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g --ulimit=nofile=1000000:1000000"
@@ -141,7 +141,7 @@ jobs:
   docs-build:
     needs: conda-python-build
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10
     with:
       build_type: pull-request
       node_type: "cpu8"
@@ -150,7 +150,7 @@ jobs:
       script: "ci/build_docs.sh"
   devcontainer:
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/build-in-devcontainer.yaml@branch-25.10
     with:
       arch: '["amd64", "arm64"]'
       cuda: '["13.0"]'

.github/workflows/test.yaml

@@ -25,7 +25,7 @@ on:
 jobs:
   conda-cpp-tests:
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type }}
       branch: ${{ inputs.branch }}
@@ -35,7 +35,7 @@ jobs:
       sha: ${{ inputs.sha }}
   conda-python-tests:
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type }}
       branch: ${{ inputs.branch }}
@@ -46,7 +46,7 @@ jobs:
       sha: ${{ inputs.sha }}
   wheel-test:
     secrets: inherit
-    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@cuda13.0
+    uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@branch-25.10
     with:
       build_type: ${{ inputs.build_type }}
       branch: ${{ inputs.branch }}

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.

cpp/src/shuffler/shuffler.cpp

@@ -105,6 +105,10 @@ std::unique_ptr<Buffer> allocate_buffer(
  * @param amount The maximum amount of data (in bytes) to be spilled.
  *
  * @return The actual amount of data successfully spilled from the postbox.
+ *
+ * @warning This may temporarily empty the postbox, causing emptiness checks to return
+ * true even though the postbox is not actually empty. As a result, in the current
+ * implementation `postbox_spilling()` must not be used to spill `outgoing_postbox_`.
  */
 template <typename KeyType>
 std::size_t postbox_spilling(

python/rapidsmpf/rapidsmpf/CMakeLists.txt

@@ -3,8 +3,8 @@
 # SPDX-License-Identifier: Apache-2.0
 # =================================================================================
 
-set(cython_modules config.pyx progress_thread.pyx rmm_resource_adaptor.pyx shuffler.pyx
-                   statistics.pyx
+set(cython_modules config.pyx cuda_stream.pyx progress_thread.pyx rmm_resource_adaptor.pyx
+                   shuffler.pyx statistics.pyx
 )
 
 rapids_cython_create_modules(

python/rapidsmpf/rapidsmpf/allgather/allgather.pyx

@@ -59,13 +59,11 @@ cdef class AllGather:
         Communicator comm,
         ProgressThread progress_thread,
         uint8_t op_id,
-        stream,
+        Stream stream,
         BufferResource br,
         Statistics statistics = None,
     ):
-        if stream is None:
-            raise ValueError("stream cannot be None")
-        self._stream = Stream(stream)
+        self._stream = stream
         self._comm = comm
         self._br = br
         cdef cpp_BufferResource* br_ = br.ptr()