Moving spill/unspill_partitions from cudf_streaming

cpp/CMakeLists.txt

@@ -168,10 +168,6 @@ if(BUILD_SLURM_SUPPORT)
   endif()
 endif()
 
-if(BUILD_CUDF_TESTS)
-  include(../cmake/thirdparty/get_cudf_streaming.cmake)
-endif()
-
 # ##################################################################################################
 # * library targets --------------------------------------------------------------------------------
 
@@ -206,6 +202,7 @@ add_library(
   src/memory/memory_type.cpp
   src/memory/pinned_memory_resource.cpp
   src/memory/scoped_memory_record.cpp
+  src/memory/spill.cpp
   src/memory/spill_manager.cpp
   src/pausable_thread_loop.cpp
   src/progress_thread.cpp
@@ -347,6 +344,10 @@ if(NOT TARGET rapidsmpf::rapidsmpf)
   add_library(rapidsmpf::rapidsmpf ALIAS rapidsmpf)
 endif()
 
+if(BUILD_CUDF_TESTS)
+  include(../cmake/thirdparty/get_cudf_streaming.cmake)
+endif()
+
 # ##################################################################################################
 # * linter configuration ---------------------------------------------------------------------------
 if(RAPIDSMPF_CLANG_TIDY)

cpp/benchmarks/bench_shuffle.cpp

@@ -20,6 +20,7 @@
 #include <rapidsmpf/communicator/ucxx.hpp>
 #include <rapidsmpf/communicator/ucxx_utils.hpp>
 #include <rapidsmpf/error.hpp>
+#include <rapidsmpf/memory/spill.hpp>
 #include <rapidsmpf/nvtx.hpp>
 #include <rapidsmpf/progress_thread.hpp>
 #include <rapidsmpf/shuffler/shuffler.hpp>
@@ -320,7 +321,7 @@ rapidsmpf::Duration do_run(
         for (auto finished_partition : shuffler.local_partitions()) {
             auto packed_chunks = shuffler.extract(finished_partition);
             auto output_partition = cudf_streaming::integrations::unpack_and_concat(
-                cudf_streaming::integrations::unspill_partitions(
+                rapidsmpf::unspill_partitions(
                     std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES
                 ),
                 stream,

cpp/benchmarks/streaming/ndsh/join.cpp

@@ -28,6 +28,7 @@
 #include <rapidsmpf/cuda_stream.hpp>
 #include <rapidsmpf/error.hpp>
 #include <rapidsmpf/memory/packed_data.hpp>
+#include <rapidsmpf/memory/spill.hpp>
 #include <rapidsmpf/shuffler/chunk.hpp>
 #include <rapidsmpf/streaming/coll/allgather.hpp>
 #include <rapidsmpf/streaming/coll/shuffler.hpp>
@@ -121,7 +122,7 @@ coro::task<streaming::Message> broadcast(
                 0,
                 std::make_unique<cudf_streaming::streaming::TableChunk>(
                     cudf_streaming::integrations::unpack_and_concat(
-                        cudf_streaming::integrations::unspill_partitions(
+                        rapidsmpf::unspill_partitions(
                             std::move(result), ctx->br().get(), AllowOverbooking::YES
                         ),
                         stream,
@@ -602,7 +603,7 @@ streaming::Actor shuffle(
                 pid,
                 std::make_unique<cudf_streaming::streaming::TableChunk>(
                     cudf_streaming::integrations::unpack_and_concat(
-                        cudf_streaming::integrations::unspill_partitions(
+                        rapidsmpf::unspill_partitions(
                             std::move(packed_data), ctx->br().get(), AllowOverbooking::YES
                         ),
                         stream,

cpp/examples/example_shuffle.cpp

@@ -13,6 +13,7 @@
 #include <rapidsmpf/communicator/mpi.hpp>
 #include <rapidsmpf/error.hpp>
 #include <rapidsmpf/memory/packed_data.hpp>
+#include <rapidsmpf/memory/spill.hpp>
 #include <rapidsmpf/shuffler/shuffler.hpp>
 #include <rapidsmpf/statistics.hpp>
 
@@ -112,7 +113,7 @@ int main(int argc, char** argv) {
         // convenience function.
         local_outputs.push_back(
             cudf_streaming::integrations::unpack_and_concat(
-                cudf_streaming::integrations::unspill_partitions(
+                rapidsmpf::unspill_partitions(
                     std::move(packed_chunks), br.get(), rapidsmpf::AllowOverbooking::YES
                 ),
                 stream,

cpp/include/rapidsmpf/memory/spill.hpp

@@ -0,0 +1,64 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <vector>
+
+#include <rapidsmpf/memory/buffer_resource.hpp>
+#include <rapidsmpf/memory/packed_data.hpp>
+
+namespace rapidsmpf {
+
+/**
+ * @brief Spill partitions from device memory to host memory.
+ *
+ * Moves the buffer of each `PackedData` from device memory to host memory using
+ * the provided buffer resource and the buffer's CUDA stream. Partitions that are
+ * already in host memory are passed through unchanged.
+ *
+ * For device-resident partitions, a host memory reservation is made before moving
+ * the buffer. If the reservation fails due to insufficient host memory, an exception
+ * is thrown. Overbooking is not allowed.
+ *
+ * @param partitions The partitions to spill.
+ * @param br Buffer resource used to reserve host memory and perform the move.
+ *
+ * @return A vector of `PackedData`, where each buffer resides in host memory.
+ *
+ * @throws rapidsmpf::reservation_error If host memory reservation fails.
+ */
+std::vector<PackedData> spill_partitions(
+    std::vector<PackedData>&& partitions, BufferResource* br
+);
+
+/**
+ * @brief Move spilled partitions (i.e., packed tables in host memory) back to device
+ * memory.
+ *
+ * Each partition is inspected to determine whether its buffer resides in device memory.
+ * Buffers already in device memory are left untouched. Host-resident buffers are moved
+ * to device memory using the provided buffer resource and the buffer's CUDA stream.
+ *
+ * If insufficient device memory is available, the buffer resource's spill manager is
+ * invoked to free memory. If overbooking occurs and spilling fails to reclaim enough
+ * memory, behavior depends on the `allow_overbooking` flag.
+ *
+ * @param partitions The partitions to unspill, potentially containing host-resident data.
+ * @param br Buffer resource responsible for memory reservation and spills.
+ * @param allow_overbooking If false, ensures enough memory is freed to satisfy the
+ * reservation; otherwise, allows overbooking even if spilling was insufficient.
+ *
+ * @return A vector of `PackedData`, each with a buffer in device memory.
+ *
+ * @throws rapidsmpf::reservation_error If overbooking exceeds the amount spilled and
+ *         `allow_overbooking` is false.
+ */
+std::vector<PackedData> unspill_partitions(
+    std::vector<PackedData>&& partitions,
+    BufferResource* br,
+    AllowOverbooking allow_overbooking
+);
+
+}  // namespace rapidsmpf

cpp/src/memory/spill.cpp

@@ -0,0 +1,62 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <utility>
+
+#include <rapidsmpf/memory/buffer.hpp>
+#include <rapidsmpf/memory/buffer_resource.hpp>
+#include <rapidsmpf/memory/memory_type.hpp>
+#include <rapidsmpf/memory/packed_data.hpp>
+#include <rapidsmpf/memory/spill.hpp>
+
+namespace rapidsmpf {
+
+std::vector<PackedData> spill_partitions(
+    std::vector<PackedData>&& partitions, BufferResource* br
+) {
+    // Sum the total size of all packed data in device memory.
+    std::size_t device_size{0};
+    for (auto& [_, data] : partitions) {
+        if (data->mem_type() == MemoryType::DEVICE) {
+            device_size += data->size;
+        }
+    }
+    // Spill each partition to host memory.
+    auto reservation = br->reserve_or_fail(device_size, SPILL_TARGET_MEMORY_TYPES);
+    std::vector<PackedData> ret;
+    ret.reserve(partitions.size());
+    for (auto& [metadata, data] : partitions) {
+        ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation));
+    }
+    return ret;
+}
+
+std::vector<PackedData> unspill_partitions(
+    std::vector<PackedData>&& partitions,
+    BufferResource* br,
+    AllowOverbooking allow_overbooking
+) {
+    auto statistics = br->statistics();
+    // Sum the total size of all packed data not in device memory already.
+    std::size_t non_device_size{0};
+    for (auto& [_, data] : partitions) {
+        if (data->mem_type() != MemoryType::DEVICE) {
+            non_device_size += data->size;
+        }
+    }
+
+    // Unspill each partition.
+    auto reservation =
+        br->reserve_device_memory_and_spill(non_device_size, allow_overbooking);
+    std::vector<PackedData> ret;
+    ret.reserve(partitions.size());
+    for (auto& [metadata, data] : partitions) {
+        ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation));
+    }
+
+    return ret;
+}
+
+}  // namespace rapidsmpf

cpp/tests/CMakeLists.txt

@@ -88,6 +88,7 @@ target_sources(
           test_rmm_resource_adaptor.cpp
           test_sparse_alltoall.cpp
           test_spill_manager.cpp
+          test_spilling.cpp
           test_statistics.cpp
           test_stream_ordered_timing.cpp
           test_system_info.cpp

cpp/tests/test_partition.cpp

@@ -101,43 +101,3 @@ TEST_P(NumOfPartitions, split_and_pack) {
     // Compare the input table with the result.
     CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, *result);
 }
-
-class SpillingTest : public ::testing::Test {
-  protected:
-    void SetUp() override {
-        br = BufferResource::create(cudf::get_current_device_resource_ref());
-        stream = cudf::get_default_stream();
-    }
-
-    std::shared_ptr<BufferResource> br;
-    rmm::cuda_stream_view stream;
-};
-
-TEST_F(SpillingTest, SpillUnspillRoundtripPreservesDataAndMetadata) {
-    std::vector<std::uint8_t> metadata{42, 99};
-    std::vector<std::uint8_t> payload{10, 20, 30};
-
-    // Create device input.
-    std::vector<rapidsmpf::PackedData> input;
-    input.push_back(create_packed_data(metadata, payload, stream, br.get()));
-
-    // Device -> Device (moves data)
-    auto on_gpu = cudf_streaming::integrations::unspill_partitions(
-        std::move(input), br.get(), AllowOverbooking::YES
-    );
-    ASSERT_EQ(on_gpu.size(), 1);
-    EXPECT_EQ(on_gpu[0].data->mem_type(), rapidsmpf::MemoryType::DEVICE);
-    EXPECT_EQ(*on_gpu[0].metadata, metadata);
-
-    // Device -> Host
-    auto back_on_host =
-        cudf_streaming::integrations::spill_partitions(std::move(on_gpu), br.get());
-    ASSERT_EQ(back_on_host.size(), 1);
-    EXPECT_EQ(back_on_host[0].data->mem_type(), rapidsmpf::MemoryType::HOST);
-    EXPECT_EQ(*back_on_host[0].metadata, metadata);
-
-    // Check that contents match original
-    auto res = br->reserve_or_fail(back_on_host[0].data->size, MemoryType::HOST);
-    auto actual = br->move_to_host_buffer(std::move(back_on_host[0].data), res);
-    EXPECT_EQ(actual->copy_to_uint8_vector(), payload);
-}

cpp/tests/test_shuffler.cpp

@@ -21,6 +21,7 @@
 #include <rapidsmpf/memory/buffer.hpp>
 #include <rapidsmpf/memory/buffer_resource.hpp>
 #include <rapidsmpf/memory/packed_data.hpp>
+#include <rapidsmpf/memory/spill.hpp>
 #include <rapidsmpf/shuffler/finish_counter.hpp>
 #include <rapidsmpf/shuffler/shuffler.hpp>
 #include <rapidsmpf/utils/misc.hpp>
@@ -207,7 +208,7 @@ void test_shuffler(
     for (auto finished_partition : shuffler.local_partitions()) {
         auto packed_chunks = shuffler.extract(finished_partition);
         auto result = cudf_streaming::integrations::unpack_and_concat(
-            cudf_streaming::integrations::unspill_partitions(
+            rapidsmpf::unspill_partitions(
                 std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES
             ),
             stream,
@@ -456,10 +457,9 @@ TEST(Shuffler, SpillOnInsertAndExtraction) {
 
     {
         // Now extract triggers spilling of the partition not being extracted.
-        std::vector<rapidsmpf::PackedData> output_chunks =
-            cudf_streaming::integrations::unspill_partitions(
-                shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES
-            );
+        std::vector<rapidsmpf::PackedData> output_chunks = rapidsmpf::unspill_partitions(
+            shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES
+        );
         EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1);
 
         // And insert also triggers spilling. We end up with zero device allocations.
@@ -470,15 +470,13 @@ TEST(Shuffler, SpillOnInsertAndExtraction) {
     }
 
     // Extract and unspill both partitions.
-    std::vector<rapidsmpf::PackedData> out0 =
-        cudf_streaming::integrations::unspill_partitions(
-            shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES
-        );
+    std::vector<rapidsmpf::PackedData> out0 = rapidsmpf::unspill_partitions(
+        shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES
+    );
     EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1);
-    std::vector<rapidsmpf::PackedData> out1 =
-        cudf_streaming::integrations::unspill_partitions(
-            shuffler.extract(1), br.get(), rapidsmpf::AllowOverbooking::YES
-        );
+    std::vector<rapidsmpf::PackedData> out1 = rapidsmpf::unspill_partitions(
+        shuffler.extract(1), br.get(), rapidsmpf::AllowOverbooking::YES
+    );
     EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2);
 
     // Disable spilling and insert the first partition.
@@ -905,7 +903,7 @@ TEST(Shuffler, concurrent_wait) {
             EXPECT_NO_THROW(shuffler.wait(wait_timeout));
             auto chunks = shuffler.extract(pid);
             auto result = cudf_streaming::integrations::unpack_and_concat(
-                cudf_streaming::integrations::unspill_partitions(
+                rapidsmpf::unspill_partitions(
                     std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES
                 ),
                 stream,
@@ -1001,7 +999,7 @@ TEST(Shuffler, opid_reuse) {
         for (auto pid : shuffler.local_partitions()) {
             auto chunks = shuffler.extract(pid);
             auto result = cudf_streaming::integrations::unpack_and_concat(
-                cudf_streaming::integrations::unspill_partitions(
+                rapidsmpf::unspill_partitions(
                     std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES
                 ),
                 stream,
@@ -1098,7 +1096,7 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) {
         for (auto pid : shuffler.local_partitions()) {
             auto chunks = shuffler.extract(pid);
             auto result = cudf_streaming::integrations::unpack_and_concat(
-                cudf_streaming::integrations::unspill_partitions(
+                rapidsmpf::unspill_partitions(
                     std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES
                 ),
                 stream,