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/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/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/finish_counter.hpp b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp index ad257682c..89f1eee4c 100644 --- a/cpp/include/rapidsmpf/shuffler/finish_counter.hpp +++ b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp @@ -9,8 +9,6 @@ #include #include #include -#include -#include #include #include @@ -44,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 ); @@ -103,51 +99,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. - * - * @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. + * @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 * - * 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. @@ -157,20 +125,15 @@ 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 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. + /// once. bool all_done_{false}; mutable std::mutex mutex_; // TODO: use a shared_mutex lock? diff --git a/cpp/include/rapidsmpf/shuffler/shuffler.hpp b/cpp/include/rapidsmpf/shuffler/shuffler.hpp index ff9bdecab..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 @@ -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,25 +207,13 @@ 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. - */ - 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. - * * @throws std::runtime_error if the timeout is reached. */ - void wait_on(PartID pid, std::optional timeout = {}); + void wait(std::optional timeout = {}); /** * @brief Spills data to device if necessary. 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/shuffler/finish_counter.cpp b/cpp/src/shuffler/finish_counter.cpp index 7ce102b6f..d34a44b68 100644 --- a/cpp/src/shuffler/finish_counter.cpp +++ b/cpp/src/shuffler/finish_counter.cpp @@ -47,16 +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), - pending_pids_(local_partitions.begin(), local_partitions.end()), - 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_); @@ -93,40 +95,15 @@ 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_(); } } } -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..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( @@ -490,14 +494,9 @@ bool Shuffler::finished() const { return finish_counter_.all_finished() && ready_postbox_.empty(); } -PartID Shuffler::wait_any(std::optional timeout) { - RAPIDSMPF_NVTX_FUNC_RANGE(); - return finish_counter_.wait_any(std::move(timeout)); -} - -void Shuffler::wait_on(PartID pid, std::optional timeout) { +void Shuffler::wait(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/src/streaming/coll/shuffler.cpp b/cpp/src/streaming/coll/shuffler.cpp index 3f38cb90b..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,45 +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](shuffler::PartID pid) -> void { - shuffler_.comm()->logger()->trace( - "notifying waiters that ", pid, " is 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" - ); + [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 { @@ -142,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 { @@ -254,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..38355bbfa 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 @@ -31,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: @@ -158,22 +179,20 @@ 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_inserts, n_partitions) = GetParam(); - BaseStreamingShuffle::SetUpWithThreads(n_threads); + BaseStreamingShuffle::SetUpWithThreads(4); GlobalEnvironment->barrier(); // prevent accidental mixup between shufflers } @@ -187,44 +206,18 @@ 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 - ::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)); + return "ninserts_" + std::to_string(std::get<0>(info.param)) + "_nparts_" + + std::to_string(std::get<1>(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 +228,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/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index 9c887a2aa..e6158d5a5 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( @@ -466,18 +466,39 @@ 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, 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_; if (comm->rank() != 0) { @@ -485,8 +506,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 @@ -494,14 +514,11 @@ void run_wait_test(WaitFn&& wait_fn) { ASSERT_EQ(local_partitions.size(), 1); rapidsmpf::shuffler::detail::FinishCounter finish_counter( - comm->nranks(), local_partitions + comm->nranks(), local_partitions.size() ); - // 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 +532,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 @@ -562,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(); } @@ -649,31 +650,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 +664,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(); })); } @@ -916,8 +898,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); }); @@ -1000,7 +982,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/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))); - } -} 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/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) 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(