Merge pull request #259 from madsbk/branch-25.08-merge-25.06

Branch 25.08 merge 25.06

Author: raydouglass

.gitignore

@@ -6,3 +6,4 @@ build
 .mypy_cache
 .hypothesis
 __pycache__
+compile_commands.json

conda/environments/all_cuda-128_arch-aarch64.yaml

@@ -38,7 +38,6 @@ dependencies:
 - pytest
 - python>=3.10,<3.13
 - rapids-build-backend>=0.3.0,<0.4.0.dev0
-- ray-default==2.42.*,>=0.0.0a0
 - rmm==25.8.*,>=0.0.0a0
 - scikit-build-core>=0.10.0
 - spdlog

cpp/include/rapidsmpf/progress_thread.hpp

@@ -132,10 +132,15 @@ class ProgressThread {
      *
      * @param logger The logger instance to use.
      * @param statistics The statistics instance to use (disabled by default).
+     * @param sleep The duration to sleep between each progress loop iteration.
+     * If 0, the thread yields execution instead of sleeping. Anecdotally, a 1 us
+     * sleep time (the default) is sufficient to avoid starvation and get smooth
+     * progress.
      */
     ProgressThread(
         Communicator::Logger& logger,
-        std::shared_ptr<Statistics> statistics = std::make_shared<Statistics>(false)
+        std::shared_ptr<Statistics> statistics = std::make_shared<Statistics>(false),
+        Duration sleep = std::chrono::microseconds{1}
     );
 
     ~ProgressThread();

cpp/include/rapidsmpf/shuffler/shuffler.hpp

@@ -215,7 +215,7 @@ class Shuffler {
      *
      * @param chunk The chunk to insert.
      */
-    void insert_into_outbox(detail::Chunk&& chunk);
+    void insert_into_ready_postbox(detail::Chunk&& chunk);
 
     /// @brief Get an new unique chunk ID.
     [[nodiscard]] detail::ChunkID get_new_cid();
@@ -255,10 +255,10 @@ class Shuffler {
     rmm::cuda_stream_view stream_;
     BufferResource* br_;
     bool active_{true};
-    detail::PostBox<Rank>
-        outgoing_chunks_;  ///< Outgoing chunks, that are ready to be sent to other ranks.
-    detail::PostBox<PartID> received_chunks_;  ///< Received chunks, that are ready to be
-                                               ///< extracted by the user.
+    detail::PostBox<Rank> outgoing_postbox_;  ///< Postbox for outgoing chunks, that are
+                                              ///< ready to be sent to other ranks.
+    detail::PostBox<PartID> ready_postbox_;  ///< Postbox for received chunks, that are
+                                             ///< ready to be extracted by the user.
 
     std::shared_ptr<Communicator> comm_;
     std::shared_ptr<ProgressThread> progress_thread_;

cpp/src/progress_thread.cpp

@@ -37,17 +37,20 @@ void ProgressThread::FunctionState::operator()() {
 }
 
 ProgressThread::ProgressThread(
-    Communicator::Logger& logger, std::shared_ptr<Statistics> statistics
+    Communicator::Logger& logger, std::shared_ptr<Statistics> statistics, Duration sleep
 )
-    : thread_([this]() {
-          if (!is_thread_initialized_) {
-              // This thread needs to have a cuda context associated with it.
-              // For now, do so by calling cudaFree to initialise the driver.
-              RAPIDSMPF_CUDA_TRY(cudaFree(nullptr));
-              is_thread_initialized_ = true;
-          }
-          return event_loop();
-      }),
+    : thread_(
+        [this]() {
+            if (!is_thread_initialized_) {
+                // This thread needs to have a cuda context associated with it.
+                // For now, do so by calling cudaFree to initialise the driver.
+                RAPIDSMPF_CUDA_TRY(cudaFree(nullptr));
+                is_thread_initialized_ = true;
+            }
+            return event_loop();
+        },
+        sleep
+    ),
       logger_(logger),
       statistics_(std::move(statistics)) {
     RAPIDSMPF_EXPECTS(statistics_ != nullptr, "the statistics pointer cannot be NULL");

cpp/src/shuffler/shuffler.cpp

@@ -175,7 +175,7 @@ class Shuffler::Progress {
 
         // Check for new chunks in the inbox and send off their metadata.
         auto const t0_send_metadata = Clock::now();
-        for (auto&& chunk : shuffler_.outgoing_chunks_.extract_all_ready()) {
+        for (auto&& chunk : shuffler_.outgoing_postbox_.extract_all_ready()) {
             auto dst = shuffler_.partition_owner(shuffler_.comm_, chunk.pid);
             log.trace("send metadata to ", dst, ": ", chunk);
             RAPIDSMPF_EXPECTS(
@@ -282,7 +282,7 @@ class Shuffler::Progress {
                     chunk.gpu_data =
                         std::move(allocate_buffer(0, shuffler_.stream_, shuffler_.br_));
                 }
-                shuffler_.insert_into_outbox(std::move(chunk));
+                shuffler_.insert_into_ready_postbox(std::move(chunk));
             }
         }
 
@@ -325,7 +325,7 @@ class Shuffler::Progress {
                 auto chunk = extract_value(in_transit_chunks_, cid);
                 auto future = extract_value(in_transit_futures_, cid);
                 chunk.gpu_data = shuffler_.comm_->get_gpu_data(std::move(future));
-                shuffler_.insert_into_outbox(std::move(chunk));
+                shuffler_.insert_into_ready_postbox(std::move(chunk));
             }
         }
 
@@ -356,7 +356,7 @@ class Shuffler::Progress {
                 || !(
                     fire_and_forget_.empty() && incoming_chunks_.empty()
                     && outgoing_chunks_.empty() && in_transit_chunks_.empty()
-                    && in_transit_futures_.empty() && shuffler_.outgoing_chunks_.empty()
+                    && in_transit_futures_.empty() && shuffler_.outgoing_postbox_.empty()
                 ))
                    ? ProgressThread::ProgressState::InProgress
                    : ProgressThread::ProgressState::Done;
@@ -404,13 +404,13 @@ Shuffler::Shuffler(
       partition_owner{partition_owner},
       stream_{stream},
       br_{br},
-      outgoing_chunks_{
+      outgoing_postbox_{
           [this](PartID pid) -> Rank {
               return this->partition_owner(this->comm_, pid);
           },  // extract Rank from pid
           static_cast<std::size_t>(comm->nranks())
       },
-      received_chunks_{
+      ready_postbox_{
           [](PartID pid) -> PartID { return pid; },  // identity mapping
           static_cast<std::size_t>(total_num_partitions),
       },
@@ -459,14 +459,14 @@ void Shuffler::shutdown() {
     }
 }
 
-void Shuffler::insert_into_outbox(detail::Chunk&& chunk) {
+void Shuffler::insert_into_ready_postbox(detail::Chunk&& chunk) {
     auto& log = comm_->logger();
     log.trace("insert_into_outbox: ", chunk);
     auto pid = chunk.pid;
     if (chunk.expected_num_chunks) {
         finish_counter_.move_goalpost(chunk.pid, chunk.expected_num_chunks);
     } else {
-        received_chunks_.insert(std::move(chunk));
+        ready_postbox_.insert(std::move(chunk));
     }
     finish_counter_.add_finished_chunk(pid);
 }
@@ -481,9 +481,9 @@ void Shuffler::insert(detail::Chunk&& chunk) {
             statistics_->add_bytes_stat("shuffle-payload-send", chunk.gpu_data->size);
             statistics_->add_bytes_stat("shuffle-payload-recv", chunk.gpu_data->size);
         }
-        insert_into_outbox(std::move(chunk));
+        insert_into_ready_postbox(std::move(chunk));
     } else {
-        outgoing_chunks_.insert(std::move(chunk));
+        outgoing_postbox_.insert(std::move(chunk));
     }
 }
 
@@ -555,7 +555,7 @@ std::vector<PackedData> Shuffler::extract(PartID pid) {
     // Protect the chunk extraction to make sure we don't get a chunk
     // `Shuffler::spill` is in the process of spilling.
     std::unique_lock<std::mutex> lock(outbox_spilling_mutex_);
-    auto chunks = received_chunks_.extract(pid);
+    auto chunks = ready_postbox_.extract(pid);
     lock.unlock();
     std::vector<PackedData> ret;
     ret.reserve(chunks.size());
@@ -610,7 +610,7 @@ std::size_t Shuffler::spill(std::optional<std::size_t> amount) {
     if (spill_need > 0) {
         std::lock_guard<std::mutex> lock(outbox_spilling_mutex_);
         spilled =
-            postbox_spilling(br_, comm_->logger(), stream_, received_chunks_, spill_need);
+            postbox_spilling(br_, comm_->logger(), stream_, ready_postbox_, spill_need);
     }
     return spilled;
 }
@@ -625,7 +625,7 @@ detail::ChunkID Shuffler::get_new_cid() {
 
 std::string Shuffler::str() const {
     std::stringstream ss;
-    ss << "Shuffler(outgoing=" << outgoing_chunks_ << ", received=" << received_chunks_
+    ss << "Shuffler(outgoing=" << outgoing_postbox_ << ", received=" << ready_postbox_
        << ", " << finish_counter_;
     return ss.str();
 }

dependencies.yaml

@@ -49,7 +49,7 @@ channels:
 dependencies:
   build-universal:
     common:
-      - output_types: [conda, pyproject]
+      - output_types: [conda, pyproject, requirements]
         packages:
           - &cmake_ver cmake>=3.26.4,!=3.30.0
           - ninja
@@ -87,18 +87,18 @@ dependencies:
               - cuda-nvcc
   rapids_build_skbuild:
     common:
-      - output_types: [conda, pyproject]
+      - output_types: [conda, pyproject, requirements]
         packages:
           - &rapids_build_backend rapids-build-backend>=0.3.0,<0.4.0.dev0
       - output_types: conda
         packages:
           - scikit-build-core>=0.10.0
-      - output_types: [pyproject]
+      - output_types: [pyproject, requirements]
         packages:
           - scikit-build-core[pyproject]>=0.10.0
   build-python:
     common:
-      - output_types: [conda, pyproject]
+      - output_types: [conda, pyproject, requirements]
         packages:
           - cython>=3.0.3
           - *rmm_unsuffixed
@@ -176,12 +176,14 @@ dependencies:
           - rapidsmpf==25.8.*,>=0.0.0a0
   test_python:
     common:
-      - output_types: [conda, pyproject]
+      - output_types: conda
+        packages:
+          - gdb
+      - output_types: [conda, pyproject, requirements]
         packages:
           - cudf==25.8.*,>=0.0.0a0
           - dask-cuda==25.8.*,>=0.0.0a0
           - dask-cudf==25.8.*,>=0.0.0a0
-          - gdb
           - psutil
           - pytest
           - ucxx==0.45.*,>=0.0.0a0
@@ -207,8 +209,17 @@ dependencies:
           - myst-parser
           - numpydoc
           - pydata-sphinx-theme
-          - ray-default==2.42.*,>=0.0.0a0
           - sphinx
           - sphinx-autobuild
           - sphinx-copybutton
           - ucxx==0.45.*,>=0.0.0a0
+    specific:
+      - output_types: conda
+        matrices:
+          - matrix:
+              arch: x86_64
+            packages:
+              - ray-default==2.42.*,>=0.0.0a0
+          - matrix:
+              arch: aarch64
+            packages:

python/rapidsmpf/rapidsmpf/integrations/dask/shuffler.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import threading
 from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
 
 from distributed import get_worker
@@ -22,24 +23,62 @@
 if TYPE_CHECKING:
     from collections.abc import Callable, MutableMapping, Sequence
 
-    from distributed import Worker
+    from distributed import Client, Worker
 
 
-_shuffle_counter: int = 0
+# Set of available shuffle IDs
+_shuffle_id_vacancy: set[int] = set(range(Shuffler.max_concurrent_shuffles))
+_shuffle_id_vacancy_lock: threading.Lock = threading.Lock()
 
 
-def get_shuffle_id() -> int:
+def _get_new_shuffle_id(client: Client) -> int:
     """
-    Return the unique id for a new shuffle.
+    Get a new available shuffle ID.
+
+    Since RapidsMPF only supports a limited number of shuffler instances at
+    any given time, this function maintains a shared pool of shuffle IDs.
+
+    If no IDs are available locally, it queries all workers for IDs in use,
+    updates the vacancy set accordingly, and retries. If all IDs are in use
+    across the cluster, an error is raised.
+
+    Parameters
+    ----------
+    client
+        A Dask distributed client used to query workers for active shuffle IDs.
 
     Returns
     -------
-    The enumerated integer id for the current shuffle.
+    A unique shuffle ID not currently in use.
+
+    Raises
+    ------
+    ValueError
+        If all shuffle IDs are currently in use across the cluster.
     """
-    global _shuffle_counter  # noqa: PLW0603
+    global _shuffle_id_vacancy  # noqa: PLW0603
 
-    _shuffle_counter += 1
-    return _shuffle_counter
+    with _shuffle_id_vacancy_lock:
+        if not _shuffle_id_vacancy:
+
+            def get_occupied_ids(dask_worker: Worker) -> set[int]:
+                ctx = get_worker_context(dask_worker)
+                with ctx.lock:
+                    return set(ctx.shufflers.keys())
+
+            # We start with setting all IDs as vacant and then subtract all
+            # IDs occupied on any one worker.
+            _shuffle_id_vacancy = set(range(Shuffler.max_concurrent_shuffles))
+            _shuffle_id_vacancy.difference_update(
+                *client.run(get_occupied_ids).values()
+            )
+            if not _shuffle_id_vacancy:
+                raise ValueError(
+                    f"Cannot shuffle more than {Shuffler.max_concurrent_shuffles} "
+                    "times in a single Dask compute."
+                )
+
+        return _shuffle_id_vacancy.pop()
 
 
 def get_shuffler(
@@ -280,13 +319,17 @@ def _extract_partition(
     -------
     Extracted DataFrame partition.
     """
-    if callback is None:
-        raise ValueError("Missing callback in _extract_partition.")
-    return callback(
-        partition_id,
-        column_names,
-        get_shuffler(shuffle_id),
-    )
+    shuffler = get_shuffler(shuffle_id)
+    try:
+        return callback(
+            partition_id,
+            column_names,
+            shuffler,
+        )
+    finally:
+        if shuffler.finished():
+            ctx = get_worker_context()
+            del ctx.shufflers[shuffle_id]
 
 
 def rapidsmpf_shuffle_graph(
@@ -383,7 +426,7 @@ def rapidsmpf_shuffle_graph(
     """
     # Get the shuffle id
     client = get_dask_client()
-    shuffle_id = get_shuffle_id()
+    shuffle_id = _get_new_shuffle_id(client)
 
     # Check integration argument
     if not isinstance(integration, DaskIntegration):

python/rapidsmpf/rapidsmpf/shuffler.pxd

@@ -1,7 +1,7 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
 # SPDX-License-Identifier: Apache-2.0
 
-from libc.stdint cimport uint16_t, uint32_t
+from libc.stdint cimport uint8_t, uint32_t
 from libcpp cimport bool
 from libcpp.memory cimport shared_ptr, unique_ptr
 from libcpp.string cimport string
@@ -38,7 +38,7 @@ cdef extern from "<rapidsmpf/shuffler/shuffler.hpp>" nogil:
         cpp_Shuffler(
             shared_ptr[cpp_Communicator] comm,
             shared_ptr[cpp_ProgressThread] comm,
-            uint16_t op_id,
+            uint8_t op_id,
             uint32_t total_num_partitions,
             cuda_stream_view stream,
             cpp_BufferResource *br,

python/rapidsmpf/rapidsmpf/shuffler.pyi

@@ -34,6 +34,7 @@ def unpack_and_concat(
 ) -> Table: ...
 
 class Shuffler:
+    max_concurrent_shuffles: int
     def __init__(
         self,
         comm: Communicator,