finer locking without spilling

Signed-off-by: niranda perera <[email protected]>

Author: nirandaperera

cpp/include/rapidsmpf/shuffler/chunk.hpp

@@ -469,3 +469,44 @@ inline std::ostream& operator<<(std::ostream& os, ReadyForDataMessage const& obj
 
 }  // namespace detail
 }  // namespace rapidsmpf::shuffler
+
+// Custom hash function for Chunk that uses chunk ID
+namespace std {
+template <>
+/**
+ * @brief Hash function for Chunk.
+ */
+struct hash<rapidsmpf::shuffler::detail::Chunk> {
+    /**
+     * @brief Hash function for Chunk that uses chunk ID.
+     *
+     * @param chunk The chunk to hash.
+     * @return The hash of the chunk.
+     */
+    std::size_t operator()(
+        rapidsmpf::shuffler::detail::Chunk const& chunk
+    ) const noexcept {
+        return std::hash<rapidsmpf::shuffler::detail::ChunkID>{}(chunk.chunk_id());
+    }
+};
+
+template <>
+/**
+ * @brief Equality operator for Chunk.
+ */
+struct equal_to<rapidsmpf::shuffler::detail::Chunk> {
+    /**
+     * @brief Equality operator for Chunk that uses chunk ID.
+     *
+     * @param lhs The left chunk.
+     * @param rhs The right chunk.
+     * @return True if the chunks are equal, false otherwise.
+     */
+    bool operator()(
+        rapidsmpf::shuffler::detail::Chunk const& lhs,
+        rapidsmpf::shuffler::detail::Chunk const& rhs
+    ) const noexcept {
+        return lhs.chunk_id() == rhs.chunk_id();
+    }
+};
+}  // namespace std

cpp/include/rapidsmpf/shuffler/postbox.hpp

@@ -4,13 +4,17 @@
  */
 #pragma once
 
+#include <atomic>
 #include <functional>
 #include <mutex>
+#include <ranges>
 #include <string>
 #include <unordered_map>
+#include <unordered_set>
 #include <vector>
 
 #include <rapidsmpf/error.hpp>
+#include <rapidsmpf/memory/buffer_resource.hpp>
 #include <rapidsmpf/shuffler/chunk.hpp>
 
 namespace rapidsmpf::shuffler::detail {
@@ -31,16 +35,21 @@ class PostBox {
      *
      * @tparam Fn The type of the function that maps a partition ID to a key.
      * @param key_map_fn A function that maps a partition ID to a key.
-     * @param num_keys_hint The number of keys to reserve space for.
+     * @param keys The keys expected to be used in the PostBox.
      *
      * @note The `key_map_fn` must be convertible to a function that takes a `PartID` and
      * returns a `KeyType`.
      */
-    template <typename Fn>
-    PostBox(Fn&& key_map_fn, size_t num_keys_hint = 0)
-        : key_map_fn_(std::move(key_map_fn)) {
-        if (num_keys_hint > 0) {
-            pigeonhole_.reserve(num_keys_hint);
+    template <typename Fn, std::ranges::input_range Range>
+        requires std::convertible_to<std::ranges::range_value_t<Range>, KeyType>
+    PostBox(Fn&& key_map_fn, Range&& keys) : key_map_fn_(std::move(key_map_fn)) {
+        pigeonhole_.reserve(std::ranges::size(keys));
+        for (const auto& key : keys) {
+            pigeonhole_.emplace(
+                std::piecewise_construct,
+                std::forward_as_tuple(key),
+                std::forward_as_tuple()
+            );
         }
     }
 
@@ -62,17 +71,6 @@ class PostBox {
      */
     bool is_empty(PartID pid) const;
 
-    /**
-     * @brief Extracts a specific chunk from the PostBox.
-     *
-     * @param pid The ID of the partition containing the chunk.
-     * @param cid The ID of the chunk to be accessed.
-     * @return The extracted chunk.
-     *
-     * @throws std::out_of_range If the chunk is not found.
-     */
-    [[nodiscard]] Chunk extract(PartID pid, ChunkID cid);
-
     /**
      * @brief Extracts all chunks associated with a specific partition.
      *
@@ -81,7 +79,7 @@ class PostBox {
      *
      * @throws std::out_of_range If the partition is not found.
      */
-    std::unordered_map<ChunkID, Chunk> extract(PartID pid);
+    std::vector<Chunk> extract(PartID pid);
 
     /**
      * @brief Extracts all chunks associated with a specific key.
@@ -91,7 +89,7 @@ class PostBox {
      *
      * @throws std::out_of_range If the key is not found.
      */
-    std::unordered_map<ChunkID, Chunk> extract_by_key(KeyType key);
+    std::vector<Chunk> extract_by_key(KeyType key);
 
     /**
      * @brief Extracts all ready chunks from the PostBox.
@@ -107,30 +105,36 @@ class PostBox {
      */
     [[nodiscard]] bool empty() const;
 
-    /**
-     * @brief Searches for chunks of the specified memory type.
-     *
-     * @param mem_type The type of memory to search within.
-     * @return A vector of tuples, where each tuple contains: PartID, ChunkID, and the
-     * size of the chunk.
-     */
-    [[nodiscard]] std::vector<std::tuple<key_type, ChunkID, std::size_t>> search(
-        MemoryType mem_type
-    ) const;
-
     /**
      * @brief Returns a description of this instance.
      * @return The description.
      */
     [[nodiscard]] std::string str() const;
 
+    /**
+     * @brief Spills the specified amount of data from the PostBox.
+     *
+     * @param br Buffer resource to use for spilling.
+     * @param amount The amount of data to spill.
+     * @return The amount of data spilled.
+     */
+    size_t spill(BufferResource* br, size_t amount);
+
   private:
-    // TODO: more fine-grained locking e.g. by locking each partition individually.
-    mutable std::mutex mutex_;
+    /**
+     * @brief Map value for the PostBox.
+     *
+     * @note The mutex is used to protect the chunks set.
+     */
+    struct MapValue {
+        mutable std::mutex mutex;  ///< Mutex to protect each key
+        std::unordered_set<Chunk> chunks;  ///< Set of chunks for the key
+    };
+
     std::function<key_type(PartID)>
         key_map_fn_;  ///< Function to map partition IDs to keys.
-    std::unordered_map<key_type, std::unordered_map<ChunkID, Chunk>>
-        pigeonhole_;  ///< Storage for chunks, organized by a key and chunk ID.
+    std::unordered_map<key_type, MapValue> pigeonhole_;  ///< Storage for chunks
+    std::atomic<size_t> n_non_empty_keys_{0};
 };
 
 /**

cpp/include/rapidsmpf/shuffler/shuffler.hpp

@@ -341,29 +341,26 @@ class Shuffler {
 
   private:
     BufferResource* br_;
+    std::shared_ptr<Communicator> comm_;
     std::atomic<bool> active_{true};
+    std::vector<PartID> const local_partitions_;
+
     detail::PostBox<Rank> outgoing_postbox_;  ///< Postbox for outgoing chunks, that are
                                               ///< ready to be sent to other ranks.
     detail::PostBox<PartID> ready_postbox_;  ///< Postbox for received chunks, that are
                                              ///< ready to be extracted by the user.
 
-    std::shared_ptr<Communicator> comm_;
     std::shared_ptr<ProgressThread> progress_thread_;
     ProgressThread::FunctionID progress_thread_function_id_;
     OpID const op_id_;
 
     SpillManager::SpillFunctionID spill_function_id_;
 
-    std::vector<PartID> const local_partitions_;
 
     detail::FinishCounter finish_counter_;
     std::unordered_map<PartID, detail::ChunkID> outbound_chunk_counter_;
     mutable std::mutex outbound_chunk_counter_mutex_;
 
-    // We protect ready_postbox extraction to avoid returning a chunk that is in the
-    // process of being spilled by `Shuffler::spill`.
-    mutable std::mutex ready_postbox_spilling_mutex_;
-
     std::atomic<detail::ChunkID> chunk_id_counter_{0};
 
     std::shared_ptr<Statistics> statistics_;

cpp/src/shuffler/postbox.cpp

@@ -22,88 +22,92 @@ void PostBox<KeyType>::insert(Chunk&& chunk) {
             "PostBox.insert(): all messages in the chunk must map to the same key"
         );
     }
-    std::lock_guard const lock(mutex_);
+    // std::lock_guard const lock(mutex_);
+    auto& map_value = pigeonhole_.at(key);
+    std::lock_guard lock(map_value.mutex);
+    if (map_value.chunks.empty()) {
+        RAPIDSMPF_EXPECTS(
+            n_non_empty_keys_.fetch_add(1, std::memory_order_relaxed) + 1
+                <= pigeonhole_.size(),
+            "PostBox.insert(): n_non_empty_keys_ is already at the maximum"
+        );
+    }
     RAPIDSMPF_EXPECTS(
-        pigeonhole_[key].emplace(chunk.chunk_id(), std::move(chunk)).second,
+        map_value.chunks.emplace(std::move(chunk)).second,
         "PostBox.insert(): chunk already exist"
     );
 }
 
 template <typename KeyType>
 bool PostBox<KeyType>::is_empty(PartID pid) const {
-    std::lock_guard const lock(mutex_);
-    return !pigeonhole_.contains(key_map_fn_(pid));
+    auto& map_value = pigeonhole_.at(key_map_fn_(pid));
+    std::lock_guard lock(map_value.mutex);
+    return map_value.chunks.empty();
 }
 
 template <typename KeyType>
-Chunk PostBox<KeyType>::extract(PartID pid, ChunkID cid) {
-    std::lock_guard const lock(mutex_);
-    return extract_item(pigeonhole_[key_map_fn_(pid)], cid).second;
+std::vector<Chunk> PostBox<KeyType>::extract(PartID pid) {
+    return extract_by_key(key_map_fn_(pid));
 }
 
 template <typename KeyType>
-std::unordered_map<ChunkID, Chunk> PostBox<KeyType>::extract(PartID pid) {
-    std::lock_guard const lock(mutex_);
-    return extract_value(pigeonhole_, key_map_fn_(pid));
-}
+std::vector<Chunk> PostBox<KeyType>::extract_by_key(KeyType key) {
+    auto& map_value = pigeonhole_.at(key);
+    std::vector<Chunk> ret;
+    std::lock_guard lock(map_value.mutex);
+    RAPIDSMPF_EXPECTS(!map_value.chunks.empty(), "PostBox.extract(): partition is empty");
+    ret.reserve(map_value.chunks.size());
 
-template <typename KeyType>
-std::unordered_map<ChunkID, Chunk> PostBox<KeyType>::extract_by_key(KeyType key) {
-    std::lock_guard const lock(mutex_);
-    return extract_value(pigeonhole_, key);
+    for (auto it = map_value.chunks.begin(); it != map_value.chunks.end();) {
+        auto node = map_value.chunks.extract(it++);
+        ret.emplace_back(std::move(node.value()));
+    }
+
+    RAPIDSMPF_EXPECTS(
+        n_non_empty_keys_.fetch_sub(1, std::memory_order_relaxed) > 0,
+        "PostBox.extract(): n_non_empty_keys_ is already 0"
+    );
+    return ret;
 }
 
 template <typename KeyType>
 std::vector<Chunk> PostBox<KeyType>::extract_all_ready() {
-    std::lock_guard const lock(mutex_);
     std::vector<Chunk> ret;
 
     // Iterate through the outer map
-    auto pid_it = pigeonhole_.begin();
-    while (pid_it != pigeonhole_.end()) {
-        // Iterate through the inner map
-        auto& chunks = pid_it->second;
-        auto chunk_it = chunks.begin();
-        while (chunk_it != chunks.end()) {
-            if (chunk_it->second.is_ready()) {
-                ret.emplace_back(std::move(chunk_it->second));
-                chunk_it = chunks.erase(chunk_it);
+    for (auto& [key, map_value] : pigeonhole_) {
+        std::lock_guard lock(map_value.mutex);
+        bool chunks_available = !map_value.chunks.empty();
+        auto chunk_it = map_value.chunks.begin();
+        while (chunk_it != map_value.chunks.end()) {
+            if (chunk_it->is_ready()) {
+                auto node = map_value.chunks.extract(chunk_it++);
+                ret.emplace_back(std::move(node.value()));
             } else {
                 ++chunk_it;
             }
         }
 
-        // Remove the pid entry if its chunks map is empty
-        if (chunks.empty()) {
-            pid_it = pigeonhole_.erase(pid_it);
-        } else {
-            ++pid_it;
+        // if the chunks were available and are now empty, its fully extracted
+        if (chunks_available && map_value.chunks.empty()) {
+            RAPIDSMPF_EXPECTS(
+                n_non_empty_keys_.fetch_sub(1, std::memory_order_relaxed) > 0,
+                "PostBox.extract_all_ready(): n_non_empty_keys_ is already 0"
+            );
         }
     }
-
     return ret;
 }
 
 template <typename KeyType>
 bool PostBox<KeyType>::empty() const {
-    std::lock_guard const lock(mutex_);
-    return pigeonhole_.empty();
+    return n_non_empty_keys_.load(std::memory_order_acquire) == 0;
 }
 
 template <typename KeyType>
-std::vector<std::tuple<KeyType, ChunkID, std::size_t>> PostBox<KeyType>::search(
-    MemoryType mem_type
-) const {
-    std::lock_guard const lock(mutex_);
-    std::vector<std::tuple<KeyType, ChunkID, std::size_t>> ret;
-    for (auto& [key, chunks] : pigeonhole_) {
-        for (auto& [cid, chunk] : chunks) {
-            if (!chunk.is_control_message(0) && chunk.data_memory_type() == mem_type) {
-                ret.emplace_back(key, cid, chunk.concat_data_size());
-            }
-        }
-    }
-    return ret;
+size_t PostBox<KeyType>::spill(BufferResource* /* br */, size_t /* amount */) {
+    // TODO: implement spill
+    return 0;
 }
 
 template <typename KeyType>
@@ -113,14 +117,14 @@ std::string PostBox<KeyType>::str() const {
     }
     std::stringstream ss;
     ss << "PostBox(";
-    for (auto const& [key, chunks] : pigeonhole_) {
+    for (auto const& [key, map_value] : pigeonhole_) {
         ss << "k=" << key << ": [";
-        for (auto const& [cid, chunk] : chunks) {
-            assert(cid == chunk.chunk_id());
+        for (auto const& chunk : map_value.chunks) {
+            // assert(cid == chunk.chunk_id());
             if (chunk.is_control_message(0)) {
                 ss << "EOP" << chunk.expected_num_chunks(0) << ", ";
             } else {
-                ss << cid << ", ";
+                ss << chunk.chunk_id() << ", ";
             }
         }
         ss << "\b\b], ";