add spilling

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

Author: nirandaperera

cpp/include/rapidsmpf/shuffler/chunk.hpp

@@ -469,44 +469,3 @@ 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

@@ -10,7 +10,6 @@
 #include <ranges>
 #include <string>
 #include <unordered_map>
-#include <unordered_set>
 #include <vector>
 
 #include <rapidsmpf/error.hpp>
@@ -115,20 +114,19 @@ class PostBox {
      * @brief Spills the specified amount of data from the PostBox.
      *
      * @param br Buffer resource to use for spilling.
+     * @param log Logger to use for logging.
      * @param amount The amount of data to spill.
      * @return The amount of data spilled.
      */
-    size_t spill(BufferResource* br, size_t amount);
+    size_t spill(BufferResource* br, Communicator::Logger& log, size_t amount);
 
   private:
     /**
      * @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::vector<Chunk> chunks;  ///< Vector of chunks for the key
     };
 
     std::function<key_type(PartID)>

cpp/src/shuffler/postbox.cpp

@@ -3,9 +3,12 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include <algorithm>
+#include <ranges>
 #include <sstream>
 
 #include <rapidsmpf/communicator/communicator.hpp>
+#include <rapidsmpf/nvtx.hpp>
 #include <rapidsmpf/shuffler/chunk.hpp>
 #include <rapidsmpf/shuffler/postbox.hpp>
 #include <rapidsmpf/utils.hpp>
@@ -32,10 +35,7 @@ void PostBox<KeyType>::insert(Chunk&& chunk) {
             "PostBox.insert(): n_non_empty_keys_ is already at the maximum"
         );
     }
-    RAPIDSMPF_EXPECTS(
-        map_value.chunks.emplace(std::move(chunk)).second,
-        "PostBox.insert(): chunk already exist"
-    );
+    map_value.chunks.push_back(std::move(chunk));
 }
 
 template <typename KeyType>
@@ -53,15 +53,11 @@ std::vector<Chunk> PostBox<KeyType>::extract(PartID pid) {
 template <typename KeyType>
 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());
 
-    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()));
-    }
+    std::vector<Chunk> ret = std::move(map_value.chunks);
+    map_value.chunks.clear();
 
     RAPIDSMPF_EXPECTS(
         n_non_empty_keys_.fetch_sub(1, std::memory_order_relaxed) > 0,
@@ -77,24 +73,33 @@ std::vector<Chunk> PostBox<KeyType>::extract_all_ready() {
     // Iterate through the outer map
     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;
-            }
-        }
 
-        // if the chunks were available and are now empty, its fully extracted
-        if (chunks_available && map_value.chunks.empty()) {
+        // Partition: non-ready chunks first, ready chunks at the end
+        auto partition_point =
+            std::ranges::partition(map_value.chunks, [](const Chunk& c) {
+                return !c.is_ready();
+            }).begin();
+
+        // if the chunks are available and all are ready, then all chunks will be
+        // extracted
+        if (map_value.chunks.begin() == partition_point
+            && partition_point != map_value.chunks.end())
+        {
             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"
             );
         }
+
+        // Move ready chunks to result
+        ret.insert(
+            ret.end(),
+            std::make_move_iterator(partition_point),
+            std::make_move_iterator(map_value.chunks.end())
+        );
+
+        // Remove ready chunks from the vector
+        map_value.chunks.erase(partition_point, map_value.chunks.end());
     }
     return ret;
 }
@@ -105,9 +110,44 @@ bool PostBox<KeyType>::empty() const {
 }
 
 template <typename KeyType>
-size_t PostBox<KeyType>::spill(BufferResource* /* br */, size_t /* amount */) {
-    // TODO: implement spill
-    return 0;
+size_t PostBox<KeyType>::spill(
+    BufferResource* br, Communicator::Logger& log, size_t amount
+) {
+    RAPIDSMPF_NVTX_FUNC_RANGE();
+
+    // individually lock each key and spill the chunks in it. If we are unable to lock the
+    // key, then it will be skipped.
+    size_t total_spilled = 0;
+    for (auto& [key, map_value] : pigeonhole_) {
+        std::unique_lock lock(map_value.mutex, std::try_to_lock);
+        if (lock) {  // now all chunks in this key are locked
+            for (auto& chunk : map_value.chunks) {
+                if (chunk.is_data_buffer_set()
+                    && chunk.data_memory_type() == MemoryType::DEVICE)
+                {
+                    size_t size = chunk.concat_data_size();
+                    auto [host_reservation, host_overbooking] =
+                        br->reserve(MemoryType::HOST, size, true);
+                    if (host_overbooking > 0) {
+                        log.warn(
+                            "Cannot spill to host because of host memory overbooking: ",
+                            format_nbytes(host_overbooking)
+                        );
+                        continue;
+                    }
+                    chunk.set_data_buffer(
+                        br->move(chunk.release_data_buffer(), host_reservation)
+                    );
+                    total_spilled += size;
+                    if (total_spilled >= amount) {
+                        break;
+                    }
+                }
+            }
+        }
+    }
+
+    return total_spilled;
 }
 
 template <typename KeyType>

cpp/src/shuffler/shuffler.cpp

@@ -753,7 +753,7 @@ std::size_t Shuffler::spill(std::optional<std::size_t> amount) {
     }
     std::size_t spilled{0};
     if (spill_need > 0) {
-        spilled = ready_postbox_.spill(br_, spill_need);
+        spilled = ready_postbox_.spill(br_, comm_->logger(), spill_need);
     }
     return spilled;
 }