Use (un)spill_partitions from RapidsMPF

ci/cpp_linters.sh

@@ -24,6 +24,14 @@ set +u
 conda activate clang_tidy
 set -u
 
+# clang-tidy parses the GCC compile command with clang. Newer conda compilers add
+# this GCC-only optimization flag, which clang reports as an error.
+for flags_var in CFLAGS CXXFLAGS; do
+  if [[ -n "${!flags_var:-}" ]]; then
+    export "${flags_var}=$(printf '%s' "${!flags_var}" | sed -E 's/(^|[[:space:]])-fno-merge-constants([[:space:]]|$)/ /g; s/[[:space:]]+/ /g; s/^ //; s/ $//')"
+  fi
+done
+
 export SCCACHE_S3_PREPROCESSOR_CACHE_KEY_PREFIX="cudf-cpp-linters-preprocessor-cache"
 export SCCACHE_S3_USE_PREPROCESSOR_CACHE_MODE=true
 

cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp

@@ -13,6 +13,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>
@@ -266,7 +267,7 @@ rapidsmpf::Duration do_run(rapidsmpf::shuffler::PartID const total_num_partition
     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,
         br);

cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp

@@ -25,6 +25,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>
@@ -40,7 +41,6 @@ namespace rapidsmpf::ndsh {
 using cudf_streaming::integrations::partition_and_pack;
 using cudf_streaming::integrations::partition_and_split;
 using cudf_streaming::integrations::unpack_and_concat;
-using cudf_streaming::integrations::unspill_partitions;
 using cudf_streaming::streaming::TableChunk;
 using cudf_streaming::streaming::to_message;
 
@@ -106,10 +106,10 @@ coro::task<streaming::Message> broadcast(std::shared_ptr<streaming::Context> ctx
       co_return to_message(
         0,
         std::make_unique<cudf_streaming::streaming::TableChunk>(
-          unpack_and_concat(
-            unspill_partitions(std::move(result), ctx->br().get(), AllowOverbooking::YES),
-            stream,
-            ctx->br().get()),
+          unpack_and_concat(rapidsmpf::unspill_partitions(
+                              std::move(result), ctx->br().get(), AllowOverbooking::YES),
+                            stream,
+                            ctx->br().get()),
           stream));
     }
   }
@@ -505,10 +505,10 @@ streaming::Actor shuffle(std::shared_ptr<streaming::Context> ctx,
     co_await ch_out->send(to_message(
       pid,
       std::make_unique<cudf_streaming::streaming::TableChunk>(
-        unpack_and_concat(
-          unspill_partitions(std::move(packed_data), ctx->br().get(), AllowOverbooking::YES),
-          stream,
-          ctx->br().get()),
+        unpack_and_concat(rapidsmpf::unspill_partitions(
+                            std::move(packed_data), ctx->br().get(), AllowOverbooking::YES),
+                          stream,
+                          ctx->br().get()),
         stream)));
   }
   co_await ch_out->drain(ctx->executor());

cpp/libcudf_streaming/examples/example_shuffle.cpp

@@ -10,6 +10,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>
 #include <unistd.h>
@@ -106,7 +107,7 @@ int main(int argc, char** argv)
     // Unpack (deserialize) and concatenate the chunks into a single table using a
     // 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,
       br.get()));

cpp/libcudf_streaming/tests/test_shuffler.cpp

@@ -16,6 +16,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>
@@ -194,8 +195,7 @@ void test_shuffler(std::shared_ptr<rapidsmpf::Communicator> const& comm,
   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(
-        std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES),
+      rapidsmpf::unspill_partitions(std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES),
       stream,
       br,
       rapidsmpf::AllowOverbooking::YES);
@@ -413,9 +413,8 @@ 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.
@@ -426,11 +425,11 @@ 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.
@@ -842,7 +841,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,
         br.get(),
@@ -926,7 +925,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,
         br.get(),
@@ -1008,7 +1007,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,
         br.get(),

python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py

@@ -19,11 +19,11 @@
 from cudf_streaming.integrations.partition import (
     partition_and_pack,
     unpack_and_concat,
-    unspill_partitions,
 )
 from rapidsmpf.config import Options, get_environment_variables
 from rapidsmpf.memory.buffer import MemoryType
 from rapidsmpf.memory.buffer_resource import BufferResource
+from rapidsmpf.memory.spill import unspill_partitions
 from rapidsmpf.progress_thread import ProgressThread
 from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor
 from rapidsmpf.shuffler import Shuffler

python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py

@@ -18,11 +18,11 @@
 from cudf_streaming.integrations.partition import (
     partition_and_pack,
     unpack_and_concat,
-    unspill_partitions,
 )
 from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster
 from rapidsmpf.memory.buffer import MemoryType
 from rapidsmpf.memory.buffer_resource import BufferResource
+from rapidsmpf.memory.spill import unspill_partitions
 from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor
 from rapidsmpf.shuffler import Shuffler
 from rapidsmpf.statistics import Statistics

python/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.py

@@ -15,10 +15,10 @@
 from cudf_streaming.integrations.partition import (
     partition_and_pack,
     unpack_and_concat,
-    unspill_partitions,
 )
 from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster
 from rapidsmpf.memory.buffer_resource import BufferResource
+from rapidsmpf.memory.spill import unspill_partitions
 from rapidsmpf.shuffler import Shuffler
 from rapidsmpf.testing import assert_eq
 

python/cudf_streaming/cudf_streaming/tests/test_integration_partition.py

@@ -10,12 +10,11 @@
 
 from cudf_streaming.integrations.partition import (
     partition_and_pack,
-    spill_partitions,
     split_and_pack,
     unpack_and_concat,
-    unspill_partitions,
 )
 from rapidsmpf.memory.buffer_resource import BufferResource
+from rapidsmpf.memory.spill import spill_partitions, unspill_partitions
 from rapidsmpf.testing import assert_eq
 from rmm.pylibrmm.stream import DEFAULT_STREAM
 

python/cudf_streaming/cudf_streaming/tests/test_shuffler.py

@@ -12,9 +12,9 @@
 from cudf_streaming.integrations.partition import (
     partition_and_pack,
     unpack_and_concat,
-    unspill_partitions,
 )
 from rapidsmpf.memory.buffer_resource import BufferResource
+from rapidsmpf.memory.spill import unspill_partitions
 from rapidsmpf.shuffler import (
     Shuffler,
 )