From 94348cbb73fbf8881b5f32c79f7f9f12dbc7da83 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 13 Mar 2026 13:00:33 +0000 Subject: [PATCH 1/8] Remove wait_on and wait_any interfaces from FinishCounter Since all partitions now complete simultaneously, it does not make sense to offer an interface to wait for an individual partition. --- .../rapidsmpf/shuffler/finish_counter.hpp | 50 +++---------- cpp/include/rapidsmpf/shuffler/shuffler.hpp | 27 +++++-- cpp/src/shuffler/finish_counter.cpp | 26 +------ cpp/src/shuffler/shuffler.cpp | 19 ++++- cpp/tests/test_shuffler.cpp | 72 +++---------------- 5 files changed, 58 insertions(+), 136 deletions(-) diff --git a/cpp/include/rapidsmpf/shuffler/finish_counter.hpp b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp index ad257682c..bf69ba49d 100644 --- a/cpp/include/rapidsmpf/shuffler/finish_counter.hpp +++ b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp @@ -10,7 +10,6 @@ #include #include #include -#include #include #include @@ -103,51 +102,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. + * @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 * - * @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 timeout = {}); - - /** - * @brief Wait for a specific partition to be finished (blocking). Optionally a - * timeout (in ms) can be provided. - * - * 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 timeout = {}); + void wait(std::optional timeout = {}); /** * @brief Returns a description of this instance. @@ -166,9 +137,6 @@ class FinishCounter { ChunkID total_finished_chunks_{0}; ///< global finished chunk counter std::vector rank_reported_; ///< indexed by rank, prevents double-reporting std::span 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 pending_pids_; /// Set to true exactly once when all chunks have arrived. Ensures callback only fires /// once for each partition. bool all_done_{false}; diff --git a/cpp/include/rapidsmpf/shuffler/shuffler.hpp b/cpp/include/rapidsmpf/shuffler/shuffler.hpp index ff9bdecab..6cf9193fe 100644 --- a/cpp/include/rapidsmpf/shuffler/shuffler.hpp +++ b/cpp/include/rapidsmpf/shuffler/shuffler.hpp @@ -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,22 +207,38 @@ 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. + */ + void wait(std::optional timeout = {}); + + /** + * @brief Wait for any partition to finish. + * + * @deprecated Use `wait()` followed by iterating `local_partitions()` instead. + * + * All local partitions complete simultaneously, so this just calls `wait()` and + * returns an arbitrary local partition ID. * + * @param timeout Optional timeout (ms) to wait. + * @return The partition ID of a local partition. * @throws std::runtime_error if the timeout is reached. + * @throws std::out_of_range if called more times than there are local partitions. */ PartID wait_any(std::optional timeout = {}); /** * @brief Wait for a specific partition to finish (blocking). * - * @param pid The desired partition ID. - * @param timeout Optional timeout (ms) to wait. + * @deprecated Use `wait()` instead. + * + * All local partitions complete simultaneously, so this just calls `wait()`. * + * @param pid The desired partition ID (unused, retained for API compatibility). + * @param timeout Optional timeout (ms) to wait. * @throws std::runtime_error if the timeout is reached. */ void wait_on(PartID pid, std::optional timeout = {}); @@ -342,6 +358,7 @@ class Shuffler { std::vector const local_partitions_; detail::FinishCounter finish_counter_; + std::size_t wait_any_idx_{0}; ///< next index into local_partitions_ for wait_any() std::vector outbound_chunk_counter_; ///< indexed by Rank mutable std::mutex outbound_chunk_counter_mutex_; diff --git a/cpp/src/shuffler/finish_counter.cpp b/cpp/src/shuffler/finish_counter.cpp index 7ce102b6f..c99b603bd 100644 --- a/cpp/src/shuffler/finish_counter.cpp +++ b/cpp/src/shuffler/finish_counter.cpp @@ -55,7 +55,6 @@ FinishCounter::FinishCounter( n_unfinished_partitions_{safe_cast(local_partitions.size())}, rank_reported_(safe_cast(nranks), false), local_partitions_(local_partitions), - pending_pids_(local_partitions.begin(), local_partitions.end()), finished_callback_{std::forward(finished_callback)} {} bool FinishCounter::all_finished() const { @@ -101,32 +100,9 @@ void FinishCounter::add_finished_chunk() { } } -PartID FinishCounter::wait_any(std::optional timeout) { - std::unique_lock lock(mutex_); - wait_for_if_timeout_else_wait(lock, wait_cv_, timeout, [&] { - return all_done_ || pending_pids_.empty(); - }); - - RAPIDSMPF_EXPECTS( - !pending_pids_.empty(), "no more partitions to wait on", std::out_of_range - ); - - auto it = pending_pids_.begin(); - PartID pid = *it; - pending_pids_.erase(it); - return pid; -} - -void FinishCounter::wait_on( - PartID pid, std::optional timeout -) { +void FinishCounter::wait(std::optional timeout) { std::unique_lock lock(mutex_); wait_for_if_timeout_else_wait(lock, wait_cv_, timeout, [&] { return all_done_; }); - RAPIDSMPF_EXPECTS( - pending_pids_.erase(pid) > 0, - "PartID has already been extracted", - std::out_of_range - ); } std::string detail::FinishCounter::str() const { diff --git a/cpp/src/shuffler/shuffler.cpp b/cpp/src/shuffler/shuffler.cpp index 936a86ee5..760dc1bf5 100644 --- a/cpp/src/shuffler/shuffler.cpp +++ b/cpp/src/shuffler/shuffler.cpp @@ -490,14 +490,27 @@ bool Shuffler::finished() const { return finish_counter_.all_finished() && ready_postbox_.empty(); } +void Shuffler::wait(std::optional timeout) { + RAPIDSMPF_NVTX_FUNC_RANGE(); + finish_counter_.wait(std::move(timeout)); +} + PartID Shuffler::wait_any(std::optional timeout) { RAPIDSMPF_NVTX_FUNC_RANGE(); - return finish_counter_.wait_any(std::move(timeout)); + finish_counter_.wait(std::move(timeout)); + RAPIDSMPF_EXPECTS( + wait_any_idx_ < local_partitions_.size(), + "no more partitions to wait on", + std::out_of_range + ); + return local_partitions_[wait_any_idx_++]; } -void Shuffler::wait_on(PartID pid, std::optional timeout) { +void Shuffler::wait_on( + [[maybe_unused]] PartID pid, std::optional timeout +) { RAPIDSMPF_NVTX_FUNC_RANGE(); - finish_counter_.wait_on(pid, std::move(timeout)); + finish_counter_.wait(std::move(timeout)); } std::size_t Shuffler::spill(std::optional amount) { diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index 9c887a2aa..d9a1e2da3 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -466,18 +466,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { EXPECT_EQ(mr.get_main_record().num_current_allocs(), 0); } -/** - * @brief A test util that runs the wait test by first calling wait_fn lambda with no - * partitions finished, and then with one partition finished. Former case, should timeout, - * while the latter should pass. - * - * @tparam WaitFn a lambda that takes FinishCounter and PartID as arguments and returns - * the result of the wait function. - * - * @param wait_fn wait lambda - */ -template -void run_wait_test(WaitFn&& wait_fn) { +TEST(FinishCounterTests, wait_with_timeout) { auto comm = GlobalEnvironment->comm_; if (comm->rank() != 0) { @@ -485,8 +474,7 @@ void run_wait_test(WaitFn&& wait_fn) { } // Use nranks partitions so each rank owns exactly 1 partition (round robin). - rapidsmpf::shuffler::PartID out_nparts = - rapidsmpf::safe_cast(comm->nranks()); + auto out_nparts = rapidsmpf::safe_cast(comm->nranks()); auto local_partitions = rapidsmpf::shuffler::Shuffler::local_partitions( comm, out_nparts, &rapidsmpf::shuffler::Shuffler::round_robin @@ -497,11 +485,8 @@ void run_wait_test(WaitFn&& wait_fn) { comm->nranks(), local_partitions ); - // pick the single local partition to test - auto p_id = local_partitions[0]; - - // none of the partitions are finished now. So, wait_fn should timeout - EXPECT_THROW(wait_fn(finish_counter, p_id), std::runtime_error); + // none of the partitions are finished now. So, wait should timeout + EXPECT_THROW(finish_counter.wait(std::chrono::milliseconds(10)), std::runtime_error); // For nranks ranks, each rank sends 1 data chunk + 1 control, so // move_goalpost(rank, 2) per rank. @@ -515,27 +500,9 @@ void run_wait_test(WaitFn&& wait_fn) { finish_counter.add_finished_chunk(); // control chunk } - // pass the wait_fn result to extract_pid_fn. It should return p_id - EXPECT_EQ(p_id, wait_fn(finish_counter, p_id)); -} - -TEST(FinishCounterTests, wait_with_timeout) { - ASSERT_NO_FATAL_FAILURE( - run_wait_test([](rapidsmpf::shuffler::detail::FinishCounter& finish_counter, - rapidsmpf::shuffler::PartID const& /* exp_pid */) { - return finish_counter.wait_any(std::chrono::milliseconds(10)); - }) - ); -} - -TEST(FinishCounterTests, wait_on_with_timeout) { - ASSERT_NO_FATAL_FAILURE( - run_wait_test([&](rapidsmpf::shuffler::detail::FinishCounter& finish_counter, - rapidsmpf::shuffler::PartID const& exp_pid) { - finish_counter.wait_on(exp_pid, std::chrono::milliseconds(10)); - return exp_pid; // return expected PID as wait_on return void - }) - ); + // After completion, wait should return immediately + EXPECT_NO_THROW(finish_counter.wait(std::chrono::milliseconds(10))); + EXPECT_TRUE(finish_counter.all_finished()); } class FinishCounterMultithreadingTest @@ -649,31 +616,12 @@ TEST_P(FinishCounterMultithreadingTest, produce_then_consume) { EXPECT_TRUE(finish_counter->all_finished()); } -TEST_P(FinishCounterMultithreadingTest, wait_any) { +TEST_P(FinishCounterMultithreadingTest, wait) { produce_data(); std::atomic n_wait_calls{0}; auto futures = create_consumer_threads_with_wait([&](auto /* pid */) { - finish_counter->wait_any(timeout); - n_wait_calls.fetch_add(1, std::memory_order_relaxed); - }); - - EXPECT_NO_THROW(std::ranges::for_each(futures, [](auto& f) { f.get(); })); - - EXPECT_EQ(npartitions, n_wait_calls); - EXPECT_TRUE(finish_counter->all_finished()); - - // callbacks should still receive all finished partitions, even after the wait_any - auto cb_futures = create_consumer_threads_with_cb(); - EXPECT_NO_THROW(std::ranges::for_each(cb_futures, [](auto& f) { f.get(); })); -} - -TEST_P(FinishCounterMultithreadingTest, wait_on) { - produce_data(); - - std::atomic n_wait_calls{0}; - auto futures = create_consumer_threads_with_wait([&](auto pid) { - finish_counter->wait_on(pid, timeout); + finish_counter->wait(timeout); n_wait_calls.fetch_add(1, std::memory_order_relaxed); }); @@ -682,7 +630,7 @@ TEST_P(FinishCounterMultithreadingTest, wait_on) { EXPECT_EQ(npartitions, n_wait_calls); EXPECT_TRUE(finish_counter->all_finished()); - // callbacks should still receive all finished partitions, even after the wait_on + // callbacks should still receive all finished partitions, even after the wait auto cb_futures = create_consumer_threads_with_cb(); EXPECT_NO_THROW(std::ranges::for_each(cb_futures, [](auto& f) { f.get(); })); } From 7dee3f820020b67d3b25f343b07d5cb8daeb60e9 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 13 Mar 2026 13:54:13 +0000 Subject: [PATCH 2/8] Remove deprecated wait_any and wait_on Since these now just defer to wait and wait for all partitions to be ready, they are no longer necessary. --- cpp/benchmarks/bench_shuffle.cpp | 4 +-- cpp/examples/example_shuffle.cpp | 8 ++--- cpp/include/rapidsmpf/shuffler/shuffler.hpp | 29 ----------------- cpp/src/shuffler/shuffler.cpp | 18 ----------- cpp/tests/test_shuffler.cpp | 13 +++++--- cpp/tests/test_shuffler_many_streams.cpp | 6 ++-- .../benchmarks/streaming_benchmark.py | 4 +-- .../rapidsmpf/examples/bulk_mpi_shuffle.py | 6 ++-- python/rapidsmpf/rapidsmpf/examples/dask.py | 2 +- .../examples/ray/bulk_ray_shuffle.py | 6 ++-- .../examples/ray/ray_shuffle_example.py | 4 +-- .../rapidsmpf/integrations/dask/shuffler.py | 5 +-- python/rapidsmpf/rapidsmpf/shuffler.pxd | 3 +- python/rapidsmpf/rapidsmpf/shuffler.pyi | 3 +- python/rapidsmpf/rapidsmpf/shuffler.pyx | 31 +++---------------- .../rapidsmpf/tests/test_shuffler.py | 17 +++------- 16 files changed, 42 insertions(+), 117 deletions(-) diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp index ac95fc5a8..722c1fd48 100644 --- a/cpp/benchmarks/bench_shuffle.cpp +++ b/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( diff --git a/cpp/examples/example_shuffle.cpp b/cpp/examples/example_shuffle.cpp index d2a53b184..a6b95dd89 100644 --- a/cpp/examples/example_shuffle.cpp +++ b/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> 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); diff --git a/cpp/include/rapidsmpf/shuffler/shuffler.hpp b/cpp/include/rapidsmpf/shuffler/shuffler.hpp index 6cf9193fe..da6a53291 100644 --- a/cpp/include/rapidsmpf/shuffler/shuffler.hpp +++ b/cpp/include/rapidsmpf/shuffler/shuffler.hpp @@ -215,34 +215,6 @@ class Shuffler { */ void wait(std::optional timeout = {}); - /** - * @brief Wait for any partition to finish. - * - * @deprecated Use `wait()` followed by iterating `local_partitions()` instead. - * - * All local partitions complete simultaneously, so this just calls `wait()` and - * returns an arbitrary local partition ID. - * - * @param timeout Optional timeout (ms) to wait. - * @return The partition ID of a local partition. - * @throws std::runtime_error if the timeout is reached. - * @throws std::out_of_range if called more times than there are local partitions. - */ - PartID wait_any(std::optional timeout = {}); - - /** - * @brief Wait for a specific partition to finish (blocking). - * - * @deprecated Use `wait()` instead. - * - * All local partitions complete simultaneously, so this just calls `wait()`. - * - * @param pid The desired partition ID (unused, retained for API compatibility). - * @param timeout Optional timeout (ms) to wait. - * @throws std::runtime_error if the timeout is reached. - */ - void wait_on(PartID pid, std::optional timeout = {}); - /** * @brief Spills data to device if necessary. * @@ -358,7 +330,6 @@ class Shuffler { std::vector const local_partitions_; detail::FinishCounter finish_counter_; - std::size_t wait_any_idx_{0}; ///< next index into local_partitions_ for wait_any() std::vector outbound_chunk_counter_; ///< indexed by Rank mutable std::mutex outbound_chunk_counter_mutex_; diff --git a/cpp/src/shuffler/shuffler.cpp b/cpp/src/shuffler/shuffler.cpp index 760dc1bf5..3064f19f2 100644 --- a/cpp/src/shuffler/shuffler.cpp +++ b/cpp/src/shuffler/shuffler.cpp @@ -495,24 +495,6 @@ void Shuffler::wait(std::optional timeout) { finish_counter_.wait(std::move(timeout)); } -PartID Shuffler::wait_any(std::optional timeout) { - RAPIDSMPF_NVTX_FUNC_RANGE(); - finish_counter_.wait(std::move(timeout)); - RAPIDSMPF_EXPECTS( - wait_any_idx_ < local_partitions_.size(), - "no more partitions to wait on", - std::out_of_range - ); - return local_partitions_[wait_any_idx_++]; -} - -void Shuffler::wait_on( - [[maybe_unused]] PartID pid, std::optional timeout -) { - RAPIDSMPF_NVTX_FUNC_RANGE(); - finish_counter_.wait(std::move(timeout)); -} - std::size_t Shuffler::spill(std::optional amount) { RAPIDSMPF_NVTX_FUNC_RANGE(); std::size_t spill_need{0}; diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index d9a1e2da3..541dbda0c 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -161,8 +161,8 @@ void test_shuffler( // Tell the shuffler that we have no more input partitions. insert_finished_fn(); - while (!shuffler.finished()) { - auto finished_partition = shuffler.wait_any(wait_timeout); + shuffler.wait(wait_timeout); + for (auto finished_partition : shuffler.local_partitions()) { auto packed_chunks = shuffler.extract(finished_partition); auto result = rapidsmpf::unpack_and_concat( rapidsmpf::unspill_partitions( @@ -864,8 +864,8 @@ class ExtractEmptyPartitionsTest : public cudf::test::BaseFixture { } void verify_extracted_chunks(auto expected_empty_fn) { - while (!shuffler->finished()) { - auto pid = shuffler->wait_any(wait_timeout); + shuffler->wait(wait_timeout); + for (auto pid : shuffler->local_partitions()) { SCOPED_TRACE("pid: " + std::to_string(pid)); std::vector chunks; EXPECT_NO_THROW({ chunks = shuffler->extract(pid); }); @@ -948,7 +948,10 @@ TEST(ShufflerTest, multiple_shutdowns) { std::make_unique(comm, 0, comm->nranks(), &br); shuffler->insert_finished(); - std::ignore = shuffler->extract(shuffler->wait_any()); + shuffler->wait(); + for (auto pid : shuffler->local_partitions()) { + std::ignore = shuffler->extract(pid); + } constexpr int n_threads = 10; std::vector> futures; diff --git a/cpp/tests/test_shuffler_many_streams.cpp b/cpp/tests/test_shuffler_many_streams.cpp index 45f2c8c8e..a42fea215 100644 --- a/cpp/tests/test_shuffler_many_streams.cpp +++ b/cpp/tests/test_shuffler_many_streams.cpp @@ -80,13 +80,11 @@ TEST(ShufflerManyStreams, Test) { ); } - // Insert all partitions. shuffler.insert(std::move(partitions)); shuffler.insert_finished(); - // Extract and validate the partitions as they finishes. - while (!shuffler.finished()) { - auto pid = shuffler.wait_any(wait_timeout); + shuffler.wait(wait_timeout); + for (auto pid : shuffler.local_partitions()) { std::vector partition_chunks = shuffler.extract(pid); for (PackedData& chunk : partition_chunks) { auto stream = chunk.data->stream(); diff --git a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py b/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py index ff133b9af..330650c22 100644 --- a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py +++ b/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py @@ -74,8 +74,8 @@ def consume_finished_partitions( The shuffler to use. """ finished = set() - while not shuffler.finished(): - partition_id = shuffler.wait_any() + shuffler.wait() + for partition_id in shuffler.local_partitions(): assert partition_id % comm.nranks == comm.rank # discard the extracted partition splits diff --git a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py index ee434adec..688f4d2eb 100644 --- a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py @@ -220,9 +220,9 @@ def bulk_mpi_shuffle( # Tell the shuffler we are done adding local data shuffler.insert_finished() - # Write shuffled partitions to disk as they finish - while not shuffler.finished(): - partition_id = shuffler.wait_any() + # Write shuffled partitions to disk + shuffler.wait() + for partition_id in shuffler.local_partitions(): table = unpack_and_concat( unspill_partitions( shuffler.extract(partition_id), diff --git a/python/rapidsmpf/rapidsmpf/examples/dask.py b/python/rapidsmpf/rapidsmpf/examples/dask.py index 641f9139a..70759ffcb 100644 --- a/python/rapidsmpf/rapidsmpf/examples/dask.py +++ b/python/rapidsmpf/rapidsmpf/examples/dask.py @@ -141,7 +141,7 @@ def extract_partition( assert ctx.br is not None column_names = options["column_names"] - shuffler.wait_on(partition_id) + shuffler.wait() table = unpack_and_concat( unspill_partitions( shuffler.extract(partition_id), diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py index 7dde5d171..7cc146f7b 100644 --- a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py @@ -219,7 +219,7 @@ def insert_finished(self) -> None: def extract(self) -> Iterator[tuple[int, plc.Table]]: """ - Extract shuffled partitions as they become ready. + Extract shuffled partitions. Returns ------- @@ -227,8 +227,8 @@ def extract(self) -> Iterator[tuple[int, plc.Table]]: """ from rmm.pylibrmm.stream import DEFAULT_STREAM - while not self.shuffler.finished(): - partition_id = self.shuffler.wait_any() + self.shuffler.wait() + for partition_id in self.shuffler.local_partitions(): packed_chunks = self.shuffler.extract(partition_id) partition = unpack_and_concat( unspill_partitions( diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py b/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py index b00b2c436..6b99c2ef2 100644 --- a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py +++ b/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py @@ -137,8 +137,8 @@ def run(self) -> None: shuffler.insert_finished() # Extract and check shuffled partitions - while not shuffler.finished(): - partition_id = shuffler.wait_any() + shuffler.wait() + for partition_id in shuffler.local_partitions(): packed_chunks = shuffler.extract(partition_id) partition = unpack_and_concat( unspill_partitions(packed_chunks, br=br, allow_overbooking=True), diff --git a/python/rapidsmpf/rapidsmpf/integrations/dask/shuffler.py b/python/rapidsmpf/rapidsmpf/integrations/dask/shuffler.py index 186e7d710..312aae6cb 100644 --- a/python/rapidsmpf/rapidsmpf/integrations/dask/shuffler.py +++ b/python/rapidsmpf/rapidsmpf/integrations/dask/shuffler.py @@ -309,8 +309,9 @@ def rapidsmpf_shuffle_graph( **Extraction phase** Each output partition is extracted from the local - :class:`rapidsmpf.shuffler.Shuffler` object on the worker (using `rapidsmpf.shuffler.Shuffler.wait_on` - and `rapidsmpf.integrations.cudf.partition.unpack_and_concat`). + :class:`rapidsmpf.shuffler.Shuffler` object on the worker (using + `rapidsmpf.shuffler.Shuffler.wait` and + `rapidsmpf.integrations.cudf.partition.unpack_and_concat`). The extraction phase will include a single task for each of the ``partition_count_out`` partitions in the shuffled output diff --git a/python/rapidsmpf/rapidsmpf/shuffler.pxd b/python/rapidsmpf/rapidsmpf/shuffler.pxd index d654a39b2..b182f8d04 100644 --- a/python/rapidsmpf/rapidsmpf/shuffler.pxd +++ b/python/rapidsmpf/rapidsmpf/shuffler.pxd @@ -48,8 +48,7 @@ cdef extern from "" nogil: void insert_finished() except +ex_handler vector[cpp_PackedData] extract(uint32_t pid) except +ex_handler bool finished() except +ex_handler - uint32_t wait_any() except +ex_handler - void wait_on(uint32_t pid) except +ex_handler + void wait() except +ex_handler span[const uint32_t] local_partitions() except +ex_handler string str() except +ex_handler diff --git a/python/rapidsmpf/rapidsmpf/shuffler.pyi b/python/rapidsmpf/rapidsmpf/shuffler.pyi index 5ae0beb86..0f499083b 100644 --- a/python/rapidsmpf/rapidsmpf/shuffler.pyi +++ b/python/rapidsmpf/rapidsmpf/shuffler.pyi @@ -31,6 +31,5 @@ class Shuffler: def insert_finished(self) -> None: ... def extract(self, pid: int) -> list[PackedData]: ... def finished(self) -> bool: ... - def wait_any(self) -> int: ... - def wait_on(self, pid: int) -> None: ... + def wait(self) -> None: ... def local_partitions(self) -> list[int]: ... diff --git a/python/rapidsmpf/rapidsmpf/shuffler.pyx b/python/rapidsmpf/rapidsmpf/shuffler.pyx index d946204da..499a5104c 100644 --- a/python/rapidsmpf/rapidsmpf/shuffler.pyx +++ b/python/rapidsmpf/rapidsmpf/shuffler.pyx @@ -189,36 +189,15 @@ cdef class Shuffler: ret = deref(self._handle).finished() return ret - def wait_any(self): + def wait(self): """ - Wait for any partition to finish. + Wait for all partitions to finish (blocking). - This method blocks until at least one partition is marked as finished. - It is useful for processing partitions as they are completed. - - Returns - ------- - The partition ID of the next finished partition. - """ - cdef uint32_t ret - with nogil: - ret = deref(self._handle).wait_any() - return ret - - def wait_on(self, uint32_t pid): - """ - Wait for a specific partition to finish. - - This method blocks until the desired partition - is ready for processing. - - Parameters - ---------- - pid - The desired partition ID. + This method blocks until all partitions are finished and ready + to be extracted. """ with nogil: - deref(self._handle).wait_on(pid) + deref(self._handle).wait() def local_partitions(self): """ diff --git a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py index eed7d8026..cda3a72ed 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py @@ -32,13 +32,11 @@ from rapidsmpf.communicator.communicator import Communicator -@pytest.mark.parametrize("wait_on", [False, True]) @pytest.mark.parametrize("total_num_partitions", [1, 2, 3, 10]) def test_shuffler_single_nonempty_partition( comm: Communicator, device_mr: rmm.mr.CudaMemoryResource, total_num_partitions: int, - wait_on: bool, # noqa: FBT001 ) -> None: br = BufferResource(device_mr) @@ -60,17 +58,12 @@ def test_shuffler_single_nonempty_partition( shuffler.insert_chunks(packed_inputs) shuffler.insert_finished() - my_partitions = shuffler.local_partitions() - expected_partitions = set(my_partitions) + expected_partitions = set(shuffler.local_partitions()) local_outputs = [] extracted_partitions = set() - while not shuffler.finished(): - if wait_on: - partition_id = my_partitions.pop() - shuffler.wait_on(partition_id) - else: - partition_id = shuffler.wait_any() + shuffler.wait() + for partition_id in shuffler.local_partitions(): extracted_partitions.add(partition_id) packed_chunks = shuffler.extract(partition_id) partition = unpack_and_concat( @@ -162,8 +155,8 @@ def test_shuffler_uniform( expected_partitions = set(shuffler.local_partitions()) extracted_partitions = set() - while not shuffler.finished(): - partition_id = shuffler.wait_any() + shuffler.wait() + for partition_id in shuffler.local_partitions(): extracted_partitions.add(partition_id) packed_chunks = shuffler.extract(partition_id) partition = unpack_and_concat( From 3d78fd18a6a3664fd3c476287b9354e5bfb03239 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 13 Mar 2026 14:19:48 +0000 Subject: [PATCH 3/8] Shuffler finished callback fires exactly once Now that all partitions complete together, we only need the finish callback to fire once. --- .../rapidsmpf/shuffler/finish_counter.hpp | 16 +++----- cpp/src/shuffler/finish_counter.cpp | 21 +++++----- cpp/src/shuffler/shuffler.cpp | 6 ++- cpp/src/streaming/coll/shuffler.cpp | 18 +++++---- cpp/tests/test_shuffler.cpp | 40 +++++++++++++++++-- 5 files changed, 69 insertions(+), 32 deletions(-) diff --git a/cpp/include/rapidsmpf/shuffler/finish_counter.hpp b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp index bf69ba49d..a7973507a 100644 --- a/cpp/include/rapidsmpf/shuffler/finish_counter.hpp +++ b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp @@ -9,7 +9,6 @@ #include #include #include -#include #include #include @@ -43,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; + using FinishedCallback = std::function; /** * @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 local_partitions, + PartID n_local_partitions, FinishedCallback&& finished_callback = nullptr ); @@ -136,9 +133,8 @@ class FinishCounter { ChunkID total_chunk_goal_{0}; ///< sum of all rank chunk goals ChunkID total_finished_chunks_{0}; ///< global finished chunk counter std::vector rank_reported_; ///< indexed by rank, prevents double-reporting - std::span local_partitions_; ///< for firing callbacks /// 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? diff --git a/cpp/src/shuffler/finish_counter.cpp b/cpp/src/shuffler/finish_counter.cpp index c99b603bd..d34a44b68 100644 --- a/cpp/src/shuffler/finish_counter.cpp +++ b/cpp/src/shuffler/finish_counter.cpp @@ -47,15 +47,18 @@ void wait_for_if_timeout_else_wait( } // namespace FinishCounter::FinishCounter( - Rank nranks, - std::span local_partitions, - FinishedCallback&& finished_callback + Rank nranks, PartID n_local_partitions, FinishedCallback&& finished_callback ) : nranks_{nranks}, - n_unfinished_partitions_{safe_cast(local_partitions.size())}, + n_unfinished_partitions_{n_local_partitions}, rank_reported_(safe_cast(nranks), false), - local_partitions_(local_partitions), - finished_callback_{std::forward(finished_callback)} {} + // If we own no partitions we will immediately be ready. + all_done_{n_local_partitions == 0}, + finished_callback_{std::forward(finished_callback)} { + if (all_done_ && finished_callback_) { + finished_callback_(); + } +} bool FinishCounter::all_finished() const { std::unique_lock lock(mutex_); @@ -92,10 +95,8 @@ void FinishCounter::add_finished_chunk() { wait_cv_.notify_all(); // notify any waiting threads - if (finished_callback_) { // notify the callback for each partition - for (auto pid : local_partitions_) { - finished_callback_(pid); - } + if (finished_callback_) { + finished_callback_(); } } } diff --git a/cpp/src/shuffler/shuffler.cpp b/cpp/src/shuffler/shuffler.cpp index 3064f19f2..05567d2e5 100644 --- a/cpp/src/shuffler/shuffler.cpp +++ b/cpp/src/shuffler/shuffler.cpp @@ -322,7 +322,11 @@ Shuffler::Shuffler( comm_{std::move(comm)}, op_id_{op_id}, local_partitions_{local_partitions(comm_, total_num_partitions, partition_owner)}, - finish_counter_{comm_->nranks(), local_partitions_, std::move(finished_callback)}, + finish_counter_{ + comm_->nranks(), + safe_cast(local_partitions_.size()), + std::move(finished_callback) + }, outbound_chunk_counter_(safe_cast(comm_->nranks()), 0), statistics_{br_->statistics()} { RAPIDSMPF_EXPECTS( diff --git a/cpp/src/streaming/coll/shuffler.cpp b/cpp/src/streaming/coll/shuffler.cpp index 3f38cb90b..6d35ff6e7 100644 --- a/cpp/src/streaming/coll/shuffler.cpp +++ b/cpp/src/streaming/coll/shuffler.cpp @@ -101,19 +101,21 @@ ShufflerAsync::ShufflerAsync( op_id, total_num_partitions, ctx_->br().get(), - [this](shuffler::PartID pid) -> void { + [this]() -> void { shuffler_.comm()->logger()->trace( - "notifying waiters that ", pid, " is ready" + "notifying waiters that all pids are ready" ); // Libcoro may resume suspended coroutines during cv notification, using the // caller thread. Submitting a detached task ensures that the progress // thread is not used to resume the coroutines. - RAPIDSMPF_EXPECTS( - notifications_.start( - insert_and_notify(mtx_, semaphore_, latch_, ready_pids_, pid) - ), - "failed to start task to notify waiters that the partition is ready" - ); + for (auto pid : shuffler_.local_partitions()) { + RAPIDSMPF_EXPECTS( + notifications_.start( + insert_and_notify(mtx_, semaphore_, latch_, ready_pids_, pid) + ), + "failed to start task to notify waiters that the partition is ready" + ); + } }, std::move(partition_owner) ) {} diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index 541dbda0c..e6158d5a5 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -466,6 +466,38 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { EXPECT_EQ(mr.get_main_record().num_current_allocs(), 0); } +TEST(FinishCounterTests, zero_local_partitions_fires_callback) { + bool callback_fired = false; + rapidsmpf::shuffler::detail::FinishCounter finish_counter( + /*nranks=*/2, /*n_local_partitions=*/0, [&]() { callback_fired = true; } + ); + + EXPECT_TRUE(callback_fired); + EXPECT_TRUE(finish_counter.all_finished()); + EXPECT_NO_THROW(finish_counter.wait(std::chrono::milliseconds(10))); +} + +TEST(FinishCounterTests, nonzero_local_partitions_fires_callback) { + bool callback_fired = false; + rapidsmpf::shuffler::detail::FinishCounter finish_counter( + /*nranks=*/1, /*n_local_partitions=*/2, [&]() { callback_fired = true; } + ); + + EXPECT_FALSE(callback_fired); + EXPECT_FALSE(finish_counter.all_finished()); + + // One rank sends 3 chunks total. + finish_counter.move_goalpost(0, 3); + finish_counter.add_finished_chunk(); + finish_counter.add_finished_chunk(); + EXPECT_FALSE(callback_fired); + + finish_counter.add_finished_chunk(); + EXPECT_TRUE(callback_fired); + EXPECT_TRUE(finish_counter.all_finished()); + EXPECT_NO_THROW(finish_counter.wait(std::chrono::milliseconds(10))); +} + TEST(FinishCounterTests, wait_with_timeout) { auto comm = GlobalEnvironment->comm_; @@ -482,7 +514,7 @@ TEST(FinishCounterTests, wait_with_timeout) { ASSERT_EQ(local_partitions.size(), 1); rapidsmpf::shuffler::detail::FinishCounter finish_counter( - comm->nranks(), local_partitions + comm->nranks(), local_partitions.size() ); // none of the partitions are finished now. So, wait should timeout @@ -529,10 +561,12 @@ class FinishCounterMultithreadingTest n_finished_pids = 0; finish_counter = std::make_unique( - nranks, local_partitions, [&](rapidsmpf::shuffler::PartID pid) { + nranks, npartitions, [&]() { { std::lock_guard lock(mtx); - finished_pids.push_back(pid); + for (auto pid : local_partitions) { + finished_pids.push_back(pid); + } } cv.notify_all(); } From 30206a60a0cb63841b5091c288eccb7bf67163c7 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 13 Mar 2026 16:55:37 +0000 Subject: [PATCH 4/8] Simplify async shuffle notification mechanism Now that the shuffle callback fires only once, we can use the same suspend and wake mechanism that the async allgather uses. The async shuffle now holds an event that is set by the finish callback when all partitions are complete. We must now `co_await insert_finished()` and can then extract all of our local partitions without blocking. --- cpp/benchmarks/streaming/ndsh/join.cpp | 7 +- .../rapidsmpf/streaming/coll/shuffler.hpp | 127 +++-------- cpp/src/streaming/coll/shuffler.cpp | 201 ++---------------- cpp/tests/streaming/test_shuffler.cpp | 161 ++------------ .../rapidsmpf/streaming/coll/shuffler.pxd | 1 + .../rapidsmpf/streaming/coll/shuffler.pyi | 7 +- .../rapidsmpf/streaming/coll/shuffler.pyx | 152 +------------ .../tests/streaming/test_shuffler.py | 42 +--- 8 files changed, 78 insertions(+), 620 deletions(-) diff --git a/cpp/benchmarks/streaming/ndsh/join.cpp b/cpp/benchmarks/streaming/ndsh/join.cpp index 843c17316..4654d7091 100644 --- a/cpp/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/benchmarks/streaming/ndsh/join.cpp @@ -582,8 +582,7 @@ 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( @@ -591,9 +590,7 @@ streaming::Actor shuffle( std::make_unique( 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() diff --git a/cpp/include/rapidsmpf/streaming/coll/shuffler.hpp b/cpp/include/rapidsmpf/streaming/coll/shuffler.hpp index a1aabc264..b8ef5d0f2 100644 --- a/cpp/include/rapidsmpf/streaming/coll/shuffler.hpp +++ b/cpp/include/rapidsmpf/streaming/coll/shuffler.hpp @@ -5,7 +5,7 @@ #pragma once -#include +#include #include #include @@ -16,19 +16,16 @@ 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} @@ -36,16 +33,12 @@ namespace rapidsmpf::streaming { * 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>> 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>; - - /** - * @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> extract_any_async(); + [[nodiscard]] std::vector 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 ctx_; - coro::task_group - 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 ready_pids_; - std::unordered_set extracted_pids_; }; namespace actor { diff --git a/cpp/src/streaming/coll/shuffler.cpp b/cpp/src/streaming/coll/shuffler.cpp index 6d35ff6e7..196a78f7a 100644 --- a/cpp/src/streaming/coll/shuffler.cpp +++ b/cpp/src/streaming/coll/shuffler.cpp @@ -3,7 +3,6 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include #include #include @@ -16,73 +15,6 @@ namespace rapidsmpf::streaming { -namespace { - -// Async notification mechanism: -// The underlying shuffle object has a callback that fires every time a partition -// arrives and, if the shuffle has N local partitions, is guaranteed to fire exactly N -// times. To wrap an async interface around this we use the following scheme: -// We attach a callback to the shuffle object that spawns a background coroutine task -// that we track in a `coro::task_group`. This task inserts the id of the ready -// partition into a set of `ready_pids_`. Consumers move received ids from `ready_pids_` -// to `extracted_pids_` and extract the partition. Insertion to, and extraction from, the -// ready and extracted sets is protected by a std::mutex (the manipulation of these -// objects does not cross coroutine suspension points). -// -// To keep track of when all notifications have been received we use a `coro::latch`. -// This must be awaited before the async shuffle goes out of scope to ensure that all -// notifications have arrived, after which we yield until the task container is empty -// ensuring that all notifications have been processed. -// -// Extraction waiters wait on acquisition of a semaphore that is released by the -// notification task and then inspect the ready_pid set. If it contains anything, that -// partition is removed from the ready set and moved into the extracted set. The -// extraction then completes and extracts a partition. -// -// Note, we do not use `coro::condition_variable` for this signalling because it -// currently has race conditions between notification of sleeping waiters and new -// waiters arriving. See https://github.com/jbaldwin/libcoro/issues/398 for details. -// -// So, in sum, the latch is required so that we do not have dangling references to the -// shuffle in the notification callback, the task_group allows us to wait until all -// notifications have truly finished firing, and a semaphore is used to release the -// "resource" of arrived partitions as they appear. - -/** - * @brief Inserts a partition ID into a ready set and notifies all waiting tasks. - * - * @param mtx The mutex to use for synchronization. - * @param semaphore Semaphore releasing resources to consumers. - * @param latch Counting notifications so we can shutdown when all notifications have been - * received. - * @param ready_pids The ready set to insert the ready partition ID into. - * @param pid The partition ID to insert. - * @return A coroutine task that completes when the partition ID is inserted into the set. - */ -coro::task insert_and_notify( - std::mutex& mtx, - coro::semaphore::max()>& semaphore, - coro::latch& latch, - std::unordered_set& ready_pids, - shuffler::PartID pid -) { - // Note: this coroutine does not need to be scheduled, because it is offloaded to the - // thread pool using a task_group. - { - std::unique_lock lock(mtx); - RAPIDSMPF_EXPECTS( - ready_pids.insert(pid).second, - "something went wrong, pid is already in the ready set!" - ); - } - // keeping track of how many notifications we've received. - latch.count_down(); - // Let a consumer know a pid is ready. - co_await semaphore.release(); -} - -} // namespace - ShufflerAsync::ShufflerAsync( std::shared_ptr ctx, std::shared_ptr comm, @@ -91,47 +23,26 @@ ShufflerAsync::ShufflerAsync( shuffler::Shuffler::PartitionOwner partition_owner ) : ctx_(std::move(ctx)), - notifications_(ctx_->executor()->get()), - latch_{static_cast(shuffler::Shuffler::local_partitions( - comm, total_num_partitions, partition_owner - ) - .size())}, shuffler_( std::move(comm), op_id, total_num_partitions, ctx_->br().get(), - [this]() -> void { - shuffler_.comm()->logger()->trace( - "notifying waiters that all pids are ready" - ); - // Libcoro may resume suspended coroutines during cv notification, using the - // caller thread. Submitting a detached task ensures that the progress - // thread is not used to resume the coroutines. - for (auto pid : shuffler_.local_partitions()) { - RAPIDSMPF_EXPECTS( - notifications_.start( - insert_and_notify(mtx_, semaphore_, latch_, ready_pids_, pid) - ), - "failed to start task to notify waiters that the partition is ready" - ); - } + [this]() { + // Schedule waiters to resume on the executor. + // This doesn't resume the frame immediately so we don't have to track + // completion of this callback with a task_group. + event_.set(ctx_->executor()->get()); }, std::move(partition_owner) ) {} ShufflerAsync::~ShufflerAsync() noexcept { RAPIDSMPF_EXPECTS_FATAL( - notifications_.empty(), - "~ShufflerAsync: not all notification tasks complete, remember to await the " + event_.is_set(), + "~ShufflerAsync: shuffle not complete, remember to await the " "finish token from this->insert_finished()" ); - if (!ready_pids_.empty()) { - comm()->logger()->warn("~ShufflerAsync: still ready partitions"); - } - if (extracted_pids_.size() != shuffler_.local_partitions().size()) { - comm()->logger()->warn("~ShufflerAsync: not all partitions have been extracted"); - } } std::span ShufflerAsync::local_partitions() const { @@ -144,90 +55,11 @@ void ShufflerAsync::insert(std::unordered_map&& ch Actor ShufflerAsync::insert_finished() { shuffler_.insert_finished(); - return finished_drain(); + co_await event_; } -coro::task>> ShufflerAsync::extract_async( - shuffler::PartID pid -) { - // Ensure that `pid` is owned by this rank. - RAPIDSMPF_EXPECTS( - shuffler_.partition_owner(comm(), pid, shuffler_.total_num_partitions) - == comm()->rank(), - "the pid isn't owned by this rank, see ShufflerAsync::partition_owner()", - std::out_of_range - ); - - while (true) { - // Note this loop might be unfair, but does not suffer from starvation because - // eventually every pid will either be in ready_pids_ or extracted_pids_. - // We don't care if the semaphore is shut down here, our pid might still be in the - // ready set to extract. - std::ignore = co_await semaphore_.acquire(); - // Note: we cannot rely on RAII for unlocking in the "success" case because the - // resumption of another coroutine might send us directly to a frame that needs - // the lock. - std::unique_lock lock(mtx_); - if (extracted_pids_.contains(pid)) { - lock.unlock(); - // Someone else got our partition. - co_return std::nullopt; - } - if (ready_pids_.erase(pid) > 0) { - RAPIDSMPF_EXPECTS( - extracted_pids_.emplace(pid).second, - "something went wrong, pid was both in the ready and the extracted set!" - ); - lock.unlock(); - co_return shuffler_.extract(pid); - } - // The pid we are waiting on is not yet available, release so that someone else - // can try and extract the pid that was available. - lock.unlock(); - co_await semaphore_.release(); - } -} - -coro::task> -ShufflerAsync::extract_any_async() { - // We don't care if the semaphore is shut down here, there might still be pids in the - // ready set to extract. - std::ignore = co_await semaphore_.acquire(); - // Note: we cannot rely on RAII for unlocking in the "success" case because the - // resumption of another coroutine might send us directly to a frame that needs - // the lock. - std::unique_lock lock(mtx_); - // Did we wake up because a partition is ready?. - if (!ready_pids_.empty()) { - // Move a pid from the ready to the extracted set. - auto pid = ready_pids_.extract(ready_pids_.begin()).value(); - RAPIDSMPF_EXPECTS( - extracted_pids_.emplace(pid).second, - "something went wrong, pid is already in the extracted set!" - ); - lock.unlock(); - co_return std::make_pair(pid, shuffler_.extract(pid)); - } - // If not, we were released because all partitions have been extracted. - lock.unlock(); - co_return std::nullopt; -} - -Actor ShufflerAsync::finished_drain() { - // Wait for all notifications to have fired. - co_await latch_; - - // Now wait for them to complete, otherwise coroutine frame unwinding can reach the - // shuffler's destructor while the notification callback still references members. - // If there are no local partitions, these coroutines were never started, and - // awaiting them would deadlock. - if (!local_partitions().empty()) { - co_await notifications_; - } - - // And wake up any pending extraction tasks. - co_await ctx_->executor()->yield(); - co_await semaphore_.shutdown(); +std::vector ShufflerAsync::extract(shuffler::PartID pid) { + return shuffler_.extract(pid); } namespace actor { @@ -256,20 +88,17 @@ Actor shuffler( shuffler_async.insert(std::move(partition_map.data)); } - auto finish_token = shuffler_async.insert_finished(); - - for ([[maybe_unused]] auto& _ : shuffler_async.local_partitions()) { - auto finished = co_await shuffler_async.extract_any_async(); - RAPIDSMPF_EXPECTS(finished.has_value(), "extract_any_async returned null"); + co_await shuffler_async.insert_finished(); + for (auto pid : shuffler_async.local_partitions()) { + auto chunks = shuffler_async.extract(pid); co_await ch_out->send(to_message( - finished->first, + pid, std::make_unique( - PartitionVectorChunk{.data = std::move(finished->second)} + PartitionVectorChunk{.data = std::move(chunks)} ) )); } - co_await finish_token; co_await ch_out->drain(ctx->executor()); } diff --git a/cpp/tests/streaming/test_shuffler.cpp b/cpp/tests/streaming/test_shuffler.cpp index 410022d9e..454fe49a0 100644 --- a/cpp/tests/streaming/test_shuffler.cpp +++ b/cpp/tests/streaming/test_shuffler.cpp @@ -3,8 +3,6 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include - #include #include @@ -158,20 +156,19 @@ TEST_P(StreamingShuffler, basic_shuffler) { })); } -class ShufflerAsyncTest : public BaseStreamingShuffle, - public ::testing::WithParamInterface< - std::tuple> { +class ShufflerAsyncTest + : public BaseStreamingShuffle, + public ::testing::WithParamInterface> { protected: int n_threads; std::size_t n_inserts; std::uint32_t n_partitions; - int n_consumers; static constexpr OpID op_id = 0; static constexpr std::size_t n_elements = 100; void SetUp() override { - std::tie(n_threads, n_inserts, n_partitions, n_consumers) = GetParam(); + std::tie(n_threads, n_inserts, n_partitions) = GetParam(); BaseStreamingShuffle::SetUpWithThreads(n_threads); GlobalEnvironment->barrier(); // prevent accidental mixup between shufflers @@ -189,42 +186,18 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Combine( ::testing::Values(1, 2, 4), // number of streaming threads ::testing::Values(1, 10), // number of inserts - ::testing::Values(1, 10, 100), // number of partitions - ::testing::Values(1, 4) // number of consumers + ::testing::Values(1, 10, 100) // number of partitions ), [](const testing::TestParamInfo& info) { return "nthreads_" + std::to_string(std::get<0>(info.param)) + "_ninserts_" + std::to_string(std::get<1>(info.param)) + "_nparts_" - + std::to_string(std::get<2>(info.param)) + "_nconsumers_" - + std::to_string(std::get<3>(info.param)); + + std::to_string(std::get<2>(info.param)); } ); -TEST_P(ShufflerAsyncTest, multi_consumer_extract) { +TEST_P(ShufflerAsyncTest, insert_wait_extract) { auto comm = GlobalEnvironment->comm_; auto shuffler = std::make_unique(ctx, comm, op_id, n_partitions); - // extract data (executed by thread pool) - auto extract_task = [](int tid, - auto* shuffler, - auto* ctx, - std::mutex& mtx, - std::vector& finished_pids, - std::size_t& n_chunks_received) -> Actor { - co_await ctx->executor()->schedule(); - ctx->logger()->debug(tid, " extract task started"); - - while (true) { - auto result = co_await shuffler->extract_any_async(); - if (!result.has_value()) { - break; - } - auto lock = std::unique_lock(mtx); - auto& [pid, chunks] = *result; - n_chunks_received += chunks.size(); - finished_pids.push_back(pid); - } - ctx->logger()->debug(tid, " extract task finished"); - }; for (std::size_t i = 0; i < n_inserts; ++i) { std::unordered_map data; @@ -235,122 +208,20 @@ TEST_P(ShufflerAsyncTest, multi_consumer_extract) { shuffler->insert(std::move(data)); } - auto finish_token = shuffler->insert_finished(); - - std::mutex mtx; - std::vector finished_pids; - std::size_t n_chunks_received = 0; - std::vector tasks; - for (int i = 0; i < n_consumers; ++i) { - tasks.emplace_back(extract_task( - i, shuffler.get(), ctx.get(), mtx, finished_pids, n_chunks_received - )); - } - tasks.push_back(ctx->executor()->schedule(std::move(finish_token))); - run_actor_network(std::move(tasks)); + coro::sync_wait(shuffler->insert_finished()); auto local_pids = shuffler::Shuffler::local_partitions( comm, n_partitions, &shuffler::Shuffler::round_robin ); - EXPECT_EQ(n_inserts * local_pids.size() * comm->nranks(), n_chunks_received); - - std::ranges::sort(finished_pids); - EXPECT_EQ(local_pids, finished_pids); -} - -TEST_F(BaseStreamingShuffle, extract_any_before_extract) { - GlobalEnvironment->barrier(); // prevent accidental mixup between shufflers - static constexpr OpID op_id = 0; - static constexpr std::size_t n_partitions = 10; - { - auto comm = GlobalEnvironment->comm_; - auto shuffler = std::make_unique(ctx, comm, op_id, n_partitions); - - // all empty partitions - auto finish_token = shuffler->insert_finished(); - - auto local_pids = shuffler::Shuffler::local_partitions( - comm, n_partitions, &shuffler::Shuffler::round_robin - ); - - std::size_t parts_extracted = 0; - // For this test we need to await the shuffler being finished and drained, i.e. - // ensure all insertion notifications have been received before extracting. This - // is only because we sync_wait each individual extract_any_async. - coro::sync_wait(finish_token); - while (true) { // extract all partitions - if (!coro::sync_wait(shuffler->extract_any_async()).has_value()) { - break; - } - parts_extracted++; - } - EXPECT_EQ(local_pids.size(), parts_extracted); - // now extract should return std::nullopt. - for (auto pid : local_pids) { - EXPECT_EQ(coro::sync_wait(shuffler->extract_async(pid)), std::nullopt); - } - } - GlobalEnvironment->barrier(); // prevent accidental mixup between shufflers -} -class CompetingShufflerAsyncTest : public BaseStreamingShuffle { - public: - void SetUp() override { - BaseStreamingShuffle::SetUp(); - GlobalEnvironment->barrier(); - } - - void TearDown() override { - GlobalEnvironment->barrier(); - BaseStreamingShuffle::TearDown(); - } - - protected: - // produce_results_fn is a function that produces the results of the extract_any_async - // and extract_async coroutines. - void run_test(auto produce_results_fn) { - static constexpr OpID op_id = 0; - auto comm = GlobalEnvironment->comm_; - shuffler::PartID const n_partitions = comm->nranks(); - shuffler::PartID const this_pid = comm->rank(); - - auto shuffler = std::make_unique(ctx, comm, op_id, 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); - - // if extract_any_result is valid, then extract_result should return nullopt - if (extract_any_result.return_value().has_value()) { - EXPECT_EQ(extract_any_result.return_value()->first, this_pid); - EXPECT_EQ(extract_result.return_value(), std::nullopt); - } else { - // else extract_result should be valid and an empty vector - EXPECT_TRUE(extract_result.return_value().has_value()); - EXPECT_EQ(extract_result.return_value()->size(), 0); - } + std::vector finished_pids; + std::size_t n_chunks_received = 0; + for (auto pid : local_pids) { + auto chunks = shuffler->extract(pid); + n_chunks_received += chunks.size(); + finished_pids.push_back(pid); } -}; -TEST_F(CompetingShufflerAsyncTest, extract_any_then_extract) { - EXPECT_NO_FATAL_FAILURE(run_test([&](auto shuffler, auto this_pid) { - return coro::sync_wait( - coro::when_all( - shuffler->extract_any_async(), shuffler->extract_async(this_pid) - ) - ); - })); -} - -TEST_F(CompetingShufflerAsyncTest, extract_then_extract_any) { - EXPECT_NO_FATAL_FAILURE(run_test([&](auto shuffler, auto this_pid) { - auto [extract_result, extract_any_result] = coro::sync_wait( - coro::when_all( - shuffler->extract_async(this_pid), shuffler->extract_any_async() - ) - ); - // rotate the results to match the order of the coroutines - return std::make_tuple(std::move(extract_any_result), std::move(extract_result)); - })); + EXPECT_EQ(n_inserts * local_pids.size() * comm->nranks(), n_chunks_received); + EXPECT_EQ(local_pids, finished_pids); } diff --git a/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pxd b/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pxd index e367f7c6b..2fc7ccc66 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pxd +++ b/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pxd @@ -39,6 +39,7 @@ cdef extern from "" nogil: const shared_ptr[cpp_Communicator]& comm() except +ex_handler void insert(unordered_map[uint32_t, cpp_PackedData] chunks) except +ex_handler span[const uint32_t] local_partitions() except +ex_handler + vector[cpp_PackedData] extract(uint32_t pid) except +ex_handler cdef class ShufflerAsync: diff --git a/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pyi b/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pyi index 7641e2b13..4ce252d00 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pyi +++ b/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pyi @@ -34,10 +34,5 @@ class ShufflerAsync: def comm(self) -> Communicator: ... def insert(self, chunks: Mapping[int, PackedData]) -> None: ... async def insert_finished(self, ctx: Context) -> None: ... + def extract(self, pid: int) -> list[PackedData]: ... def local_partitions(self) -> list[int]: ... - async def extract_async( - self, ctx: Context, pid: int - ) -> list[PackedData] | None: ... - async def extract_any_async( - self, ctx: Context - ) -> tuple[int, list[PackedData]] | None: ... diff --git a/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pyx b/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pyx index 79f98d18c..d2405b9cb 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pyx +++ b/python/rapidsmpf/rapidsmpf/streaming/coll/shuffler.pyx @@ -6,10 +6,9 @@ from cpython.ref cimport Py_INCREF from cython.operator cimport dereference as deref from libc.stdint cimport int32_t, uint32_t from libcpp.memory cimport make_unique, shared_ptr -from libcpp.optional cimport optional from libcpp.span cimport span from libcpp.unordered_map cimport unordered_map -from libcpp.utility cimport move, pair +from libcpp.utility cimport move from libcpp.vector cimport vector from rapidsmpf._detail.exception_handling cimport ex_handler @@ -33,73 +32,6 @@ import asyncio cdef extern from * nogil: """ namespace { - coro::task extract_async_task( - rapidsmpf::streaming::ShufflerAsync *shuffle, - std::uint32_t pid, - std::shared_ptr>> output - ) { - *output = co_await shuffle->extract_async(pid); - } - - std::shared_ptr>> - cpp_extract_async( - std::shared_ptr ctx, - rapidsmpf::streaming::ShufflerAsync *shuffle, - std::uint32_t pid, - void (*cpp_set_py_future)(void*, const char *), - rapidsmpf::OwningWrapper py_future - ) { - auto output = std::make_shared< - std::optional> - >(); - RAPIDSMPF_EXPECTS( - ctx->executor()->spawn_detached( - cython_libcoro_task_wrapper( - cpp_set_py_future, - std::move(py_future), - extract_async_task(shuffle, pid, output) - ) - ), - "libcoro's spawn_detached() failed to spawn task" - ); - return output; - } - - coro::task extract_any_async_task( - rapidsmpf::streaming::ShufflerAsync *shuffle, - std::shared_ptr< - std::optional>> - > output - ) { - *output = co_await shuffle->extract_any_async(); - } - - std::shared_ptr< - std::optional>> - > cpp_extract_any_async( - std::shared_ptr ctx, - rapidsmpf::streaming::ShufflerAsync *shuffle, - void (*cpp_set_py_future)(void*, const char *), - rapidsmpf::OwningWrapper py_future - ) { - auto output = std::make_shared< - std::optional>> - >(); - RAPIDSMPF_EXPECTS( - ctx->executor()->spawn_detached( - cython_libcoro_task_wrapper( - cpp_set_py_future, - std::move(py_future), - extract_any_async_task( - shuffle, output - ) - ) - ), - "libcoro's spawn_detached() failed to spawn task" - ); - return output; - } - coro::task insert_finished_task( rapidsmpf::streaming::ShufflerAsync *shuffle ) { @@ -125,22 +57,6 @@ cdef extern from * nogil: } } // namespace """ - shared_ptr[optional[vector[cpp_PackedData]]] cpp_extract_async( - shared_ptr[cpp_Context] ctx, - cpp_ShufflerAsync *shuffle, - uint32_t pid, - void (*cpp_set_py_future)(void*, const char *), - cpp_OwningWrapper py_future - ) except +ex_handler - - shared_ptr[optional[pair[uint32_t, vector[cpp_PackedData]]]] \ - cpp_extract_any_async( - shared_ptr[cpp_Context] ctx, - cpp_ShufflerAsync *shuffle, - void (*cpp_set_py_future)(void*, const char *), - cpp_OwningWrapper py_future - ) except +ex_handler - void cpp_insert_finished( shared_ptr[cpp_Context] ctx, cpp_ShufflerAsync *shuffle, @@ -303,76 +219,26 @@ cdef class ShufflerAsync: ) await ret - async def extract_async(self, Context ctx not None, uint32_t pid): + def extract(self, uint32_t pid): """ - Suspend and extract a partition from the shuffle. + Extract all chunks belonging to the specified partition. + + Must only be called after awaiting :meth:`insert_finished`. Parameters ---------- - ctx - Streaming context. pid The partition to extract. Returns ------- list[PackedData] - The PackedData representing the extracted partition. - None - If the partition has already been extracted. - """ - ret = asyncio.get_running_loop().create_future() - Py_INCREF(ret) - cdef shared_ptr[optional[vector[cpp_PackedData]]] c_ret - with nogil: - c_ret = cpp_extract_async( - ctx._handle, - self._handle.get(), - pid, - cpp_set_py_future, - move(cpp_OwningWrapper(ret, py_deleter)) - ) - await ret - if deref(c_ret).has_value(): - return packed_data_vector_to_list(move(deref(deref(c_ret)))) - else: - return None - - async def extract_any_async(self, Context ctx not None): - """ - Suspend and extract any partition from the shuffle. - - Parameters - ---------- - ctx - Streaming context. - - Returns - ------- - tuple[int, list[PackedData]] - The identifier for the extracted partition and the PackedData - of the partition. - None - If there are no more partitions to extract. + The PackedData chunks associated with the partition. """ - ret = asyncio.get_running_loop().create_future() - Py_INCREF(ret) - cdef shared_ptr[optional[pair[uint32_t, vector[cpp_PackedData]]]] c_ret + cdef vector[cpp_PackedData] c_ret with nogil: - c_ret = cpp_extract_any_async( - ctx._handle, - self._handle.get(), - cpp_set_py_future, - move(cpp_OwningWrapper(ret, py_deleter)) - ) - await ret - if deref(c_ret).has_value(): - return ( - deref(c_ret).value().first, - packed_data_vector_to_list(move(deref(c_ret).value().second)) - ) - else: - return None + c_ret = deref(self._handle).extract(pid) + return packed_data_vector_to_list(move(c_ret)) def local_partitions(self): """ diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py index cce9fc0e6..12ad8e50e 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py @@ -155,7 +155,6 @@ async def do_shuffle( op_id: int, num_partitions: int, *, - use_extract_any: bool, partition_assignment: PartitionAssignment = PartitionAssignment.ROUND_ROBIN, ) -> None: shuffle = ShufflerAsync( @@ -170,38 +169,24 @@ async def do_shuffle( split_and_pack(chunk.table_view(), splits, chunk.stream, context.br()) ) await shuffle.insert_finished(context) - if use_extract_any: - while (out := await shuffle.extract_any_async(context)) is not None: - pid, data = out - stream = context.get_stream_from_pool() - unpacked = TableChunk.from_pylibcudf_table( - unpack_and_concat(data, stream, context.br()), - stream, - exclusive_view=True, - ) - await ch_out.send(context, Message(pid, unpacked)) - else: - for pid in shuffle.local_partitions(): - pd = await shuffle.extract_async(context, pid) - assert pd is not None - stream = context.get_stream_from_pool() - unpacked = TableChunk.from_pylibcudf_table( - unpack_and_concat(pd, stream, context.br()), - stream, - exclusive_view=True, - ) - await ch_out.send(context, Message(pid, unpacked)) + for pid in shuffle.local_partitions(): + data = shuffle.extract(pid) + stream = context.get_stream_from_pool() + unpacked = TableChunk.from_pylibcudf_table( + unpack_and_concat(data, stream, context.br()), + stream, + exclusive_view=True, + ) + await ch_out.send(context, Message(pid, unpacked)) await ch_out.drain(context) @pytest.mark.parametrize("num_partitions", [4, 8]) -@pytest.mark.parametrize("use_extract_any", [False, True]) def test_shuffler_runtime_obeys_contiguous_assignment( context: Context, comm: Communicator, py_executor: ThreadPoolExecutor, num_partitions: int, - use_extract_any: bool, # noqa: FBT001 ) -> None: actors: list[CppActor | PyActor] = [] @@ -219,7 +204,6 @@ def test_shuffler_runtime_obeys_contiguous_assignment( ch_shuffled, op_id, num_partitions, - use_extract_any=use_extract_any, partition_assignment=PartitionAssignment.CONTIGUOUS, ) ) @@ -242,12 +226,10 @@ def test_shuffler_runtime_obeys_contiguous_assignment( assert len(received_pids) == len(expected_local) -@pytest.mark.parametrize("use_extract_any", [False, True]) def test_shuffler_object_interface( context: Context, comm: Communicator, py_executor: ThreadPoolExecutor, - use_extract_any: bool, # noqa: FBT001 ) -> None: actors: list[CppActor | PyActor] = [] @@ -266,7 +248,6 @@ def test_shuffler_object_interface( ch_shuffled, op_id, num_partitions, - use_extract_any=use_extract_any, ) ) actor, deferred = pull_from_channel(context, ch_shuffled) @@ -276,10 +257,7 @@ def test_shuffler_object_interface( messages = deferred.release() # TODO: single rank only assertions assert len(messages) == 5 - if use_extract_any: - assert {msg.sequence_number for msg in messages} == set(range(num_partitions)) - else: - assert [msg.sequence_number for msg in messages] == list(range(num_partitions)) + assert [msg.sequence_number for msg in messages] == list(range(num_partitions)) chunks = [(msg.sequence_number, TableChunk.from_message(msg)) for msg in messages] full_column = np.arange(num_rows * num_chunks, dtype=np.int32) From 3c08cbeadc4b40059438d1cb338eeaaa20fc3138 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 13 Mar 2026 17:18:15 +0000 Subject: [PATCH 5/8] Add test that ranks with no owned partitions are woken --- cpp/tests/streaming/test_shuffler.cpp | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/cpp/tests/streaming/test_shuffler.cpp b/cpp/tests/streaming/test_shuffler.cpp index 454fe49a0..b612db768 100644 --- a/cpp/tests/streaming/test_shuffler.cpp +++ b/cpp/tests/streaming/test_shuffler.cpp @@ -29,6 +29,29 @@ namespace actor = rapidsmpf::streaming::actor; class BaseStreamingShuffle : public BaseStreamingFixture {}; +TEST_F(BaseStreamingShuffle, zero_owned_partitions_completes) { + auto comm = GlobalEnvironment->comm_; + if (comm->nranks() < 2) { + GTEST_SKIP() << "Need at least 2 ranks so that some rank owns 0 partitions"; + } + constexpr Rank owner = 0; + auto collapse = [](std::shared_ptr const&, + shuffler::PartID, + shuffler::PartID) -> Rank { return owner; }; + constexpr OpID op_id = 0; + constexpr shuffler::PartID total = 4; + auto shuffler = std::make_unique(ctx, comm, op_id, total, collapse); + + coro::sync_wait(shuffler->insert_finished()); + + auto local_pids = shuffler->local_partitions(); + if (comm->rank() == owner) { + EXPECT_EQ(local_pids.size(), total); + } else { + EXPECT_TRUE(local_pids.empty()); + } +} + class StreamingShuffler : public BaseStreamingShuffle, public ::testing::WithParamInterface { public: From baac65dac7fdaeb47352e66e6bdb761ccf449254 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Fri, 13 Mar 2026 17:27:19 +0000 Subject: [PATCH 6/8] Simplify and only test with four streaming threads We don't have competing coroutines now, so no need for parameterising over the thread pool size. --- cpp/tests/streaming/test_shuffler.cpp | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/cpp/tests/streaming/test_shuffler.cpp b/cpp/tests/streaming/test_shuffler.cpp index b612db768..38355bbfa 100644 --- a/cpp/tests/streaming/test_shuffler.cpp +++ b/cpp/tests/streaming/test_shuffler.cpp @@ -181,9 +181,8 @@ TEST_P(StreamingShuffler, basic_shuffler) { class ShufflerAsyncTest : public BaseStreamingShuffle, - public ::testing::WithParamInterface> { + public ::testing::WithParamInterface> { protected: - int n_threads; std::size_t n_inserts; std::uint32_t n_partitions; @@ -191,9 +190,9 @@ class ShufflerAsyncTest static constexpr std::size_t n_elements = 100; void SetUp() override { - std::tie(n_threads, n_inserts, n_partitions) = GetParam(); + std::tie(n_inserts, n_partitions) = GetParam(); - BaseStreamingShuffle::SetUpWithThreads(n_threads); + BaseStreamingShuffle::SetUpWithThreads(4); GlobalEnvironment->barrier(); // prevent accidental mixup between shufflers } @@ -207,14 +206,12 @@ INSTANTIATE_TEST_SUITE_P( StreamingShuffler, ShufflerAsyncTest, ::testing::Combine( - ::testing::Values(1, 2, 4), // number of streaming threads ::testing::Values(1, 10), // number of inserts ::testing::Values(1, 10, 100) // number of partitions ), [](const testing::TestParamInfo& info) { - return "nthreads_" + std::to_string(std::get<0>(info.param)) + "_ninserts_" - + std::to_string(std::get<1>(info.param)) + "_nparts_" - + std::to_string(std::get<2>(info.param)); + return "ninserts_" + std::to_string(std::get<0>(info.param)) + "_nparts_" + + std::to_string(std::get<1>(info.param)); } ); From fb839e7778b6e28250e690d02dcb9b95e85155d1 Mon Sep 17 00:00:00 2001 From: "Mads R. B. Kristensen" Date: Mon, 16 Mar 2026 08:51:57 +0100 Subject: [PATCH 7/8] fix stale comments --- cpp/include/rapidsmpf/shuffler/finish_counter.hpp | 5 ++--- cpp/include/rapidsmpf/shuffler/shuffler.hpp | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/cpp/include/rapidsmpf/shuffler/finish_counter.hpp b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp index a7973507a..89f1eee4c 100644 --- a/cpp/include/rapidsmpf/shuffler/finish_counter.hpp +++ b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp @@ -125,9 +125,8 @@ 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 diff --git a/cpp/include/rapidsmpf/shuffler/shuffler.hpp b/cpp/include/rapidsmpf/shuffler/shuffler.hpp index da6a53291..305a10300 100644 --- a/cpp/include/rapidsmpf/shuffler/shuffler.hpp +++ b/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 From 4e3b35067a3e3becf4438b3ffec7aa227dda6301 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Mon, 16 Mar 2026 11:36:16 +0000 Subject: [PATCH 8/8] Remove timing-dependent spill check test Sleeping the main thread and letting the periodic spill thread in the background provides no guarantee on how the OS time-slices the execution, so we can't really check anything here. --- cpp/tests/test_spill_manager.cpp | 27 --------------------------- 1 file changed, 27 deletions(-) diff --git a/cpp/tests/test_spill_manager.cpp b/cpp/tests/test_spill_manager.cpp index b83fc7b92..298a9234b 100644 --- a/cpp/tests/test_spill_manager.cpp +++ b/cpp/tests/test_spill_manager.cpp @@ -75,30 +75,3 @@ TEST(SpillManager, SpillFunction) { EXPECT_EQ(br.spill_manager().spill_to_make_headroom(-100_KiB), 0); EXPECT_EQ(br.memory_available(MemoryType::DEVICE)(), 100_KiB); } - -TEST(SpillManager, PeriodicSpillCheck) { - // Create a buffer resource that always trigger spilling (always reports - // negative available memory). - std::chrono::milliseconds period{1}; - BufferResource br{ - cudf::get_current_device_resource_ref(), - PinnedMemoryResource::Disabled, - {{MemoryType::DEVICE, []() -> std::int64_t { return -100_KiB; }}}, - period, - }; - - // Spill function that increases `mem` for each call. - std::int64_t num_calls = 0; - SpillManager::SpillFunction func = - [&num_calls](std::size_t /* amount */) -> std::size_t { return ++num_calls; }; - br.spill_manager().add_spill_function(func, 0); - - std::this_thread::sleep_for(period * 100); - // With no overhead, we should see 100 spill calls but we allow wiggle room. - if (!is_running_under_valgrind()) { - EXPECT_THAT(num_calls, testing::AllOf(testing::Gt(10), testing::Lt(200))); - } else { - // In valgrind, we cannot expect it to run more than once. - EXPECT_THAT(num_calls, testing::AllOf(testing::Gt(1), testing::Lt(200))); - } -}