Enable table chunk disk copies and fix empty spill tracking

Author: nirandaperera

cpp/libcudf_streaming/src/table_chunk.cpp

@@ -150,8 +150,8 @@ table_chunk table_chunk::copy(rapidsmpf::MemoryReservation& reservation) const
   //    into the reservation-specified memory type using libcudf:
   //    a. DEVICE       - cudf-copy table_view() into device memory.
   //    b. PINNED_HOST  - cudf::pack table_view() directly into pinned memory.
-  //    c. HOST         - cudf::pack table_view() into intermediate device
-  //                      memory and then copy to host memory.
+  //    c. HOST / DISK  - cudf::pack table_view() into intermediate device
+  //                      memory and then move to host or disk memory.
   //
   // 2. The chunk data is already packed (packed_data_ != nullptr).
   //    Use buffer_copy() to copy the packed data into the reservation-
@@ -163,6 +163,17 @@ table_chunk table_chunk::copy(rapidsmpf::MemoryReservation& reservation) const
   // copy the table in device memory, or pack it to pinned/ host memory. Else, fall
   // through to case 2 (ie. use buffer_copy).
   if (is_available() && packed_data_ == nullptr) {
+    if (data_alloc_size(rapidsmpf::MemoryType::DEVICE) == 0) {
+      // Packing still produces the metadata needed to reconstruct the table.
+      // Move the empty data buffer to the requested memory type without
+      // recording a copy that transferred no bytes.
+      auto packed_columns = cudf::pack(table_view(), stream(), br->device_mr());
+      auto data           = br->move(std::move(packed_columns.gpu_data), stream());
+      data                = br->move(std::move(data), reservation);
+      return table_chunk(std::make_unique<rapidsmpf::PackedData>(std::move(packed_columns.metadata),
+                                                                 std::move(data)));
+    }
+
     switch (reservation.mem_type()) {
       case rapidsmpf::MemoryType::DEVICE:  // Case 1a.
       {
@@ -193,20 +204,18 @@ table_chunk table_chunk::copy(rapidsmpf::MemoryReservation& reservation) const
         br->release(reservation, nbytes);
         // The data leaves device memory here rather than through `BufferResource`, so
         // the spill is opened by hand and the token handed to the buffer.
-        auto host_buffer =
-          br->move(std::move(packed_pinned.gpu_data),
-                   stream(),
-                   // An empty table packs to a default-constructed `device_buffer`,
-                   // which ignores the resource and frees nothing, so it gets no token.
-                   nbytes > 0 ? std::make_shared<rapidsmpf::SpillTrackToken>() : nullptr);
+        auto host_buffer = br->move(std::move(packed_pinned.gpu_data),
+                                    stream(),
+                                    std::make_shared<rapidsmpf::SpillTrackToken>());
         return table_chunk(std::make_unique<rapidsmpf::PackedData>(
           std::move(packed_pinned.metadata), std::move(host_buffer)));
       }
       case rapidsmpf::MemoryType::HOST:  // Case 1c.
+      case rapidsmpf::MemoryType::DISK:  // Case 1c.
       {
         // We use libcudf's pack() to serialize `table_view()` into a
         // packed_columns and then we move the packed_columns' gpu_data to a
-        // new host buffer.
+        // new host/disk buffer.
         // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently,
         // `cudf::pack()` allocates device memory we haven't reserved.
         auto packed_columns = cudf::pack(table_view(), stream(), br->device_mr());

cpp/libcudf_streaming/tests/streaming/test_table_chunk.cpp

@@ -23,6 +23,7 @@
 #include <rapidsmpf/memory/buffer_resource.hpp>
 #include <rapidsmpf/owning_wrapper.hpp>
 #include <rapidsmpf/streaming/core/channel.hpp>
+#include <rapidsmpf/utils/string.hpp>
 
 #include <cstdint>
 #include <memory>
@@ -50,12 +51,13 @@ class StreamingTableChunk : public BaseStreamingFixture,
       memory_limits,                      // memory_limits
       std::chrono::milliseconds{1},       // periodic_spill_check
       stream_pool,                        // stream_pool
-      rapidsmpf::Statistics::disabled()   // statistics
-    );
+      rapidsmpf::Statistics::disabled(),  // statistics
+      spill_dir.path());
     ctx = std::make_shared<rapidsmpf::streaming::Context>(
       options, GlobalEnvironment->comm_->logger(), br);
   }
 
+  TempDir spill_dir;
   cuda::stream_ref stream{cudaStream_t{cudaStreamDefault}};
   rmm::mr::cuda_memory_resource mr_cuda;
   std::shared_ptr<rapidsmpf::BufferResource> br;
@@ -172,7 +174,9 @@ TEST_F(StreamingTableChunk, FromPackedDataOnDevice)
 
 INSTANTIATE_TEST_SUITE_P(StreamingTableChunkWithSpillTargets,
                          StreamingTableChunk,
-                         ::testing::ValuesIn(rapidsmpf::SPILL_TARGET_MEMORY_TYPES),
+                         ::testing::ValuesIn({rapidsmpf::MemoryType::PINNED_HOST,
+                                              rapidsmpf::MemoryType::HOST,
+                                              rapidsmpf::MemoryType::DISK}),
                          [](testing::TestParamInfo<rapidsmpf::MemoryType> const& info) {
                            return std::string{rapidsmpf::to_string(info.param)};
                          });
@@ -304,9 +308,12 @@ TEST_P(StreamingTableChunk, DeviceToHostRoundTripCopy)
     }
   }
 
-  // Host to host copy.
-  auto host_res2  = br->reserve_or_fail(host_copy.data_alloc_size(spill_mem_type), spill_mem_type);
-  auto host_copy2 = host_copy.copy(host_res2);
+  // Disk-to-disk copies are unsupported; keep the disk chunk for the round trip.
+  auto host_copy2 = [&] {
+    if (spill_mem_type == rapidsmpf::MemoryType::DISK) { return std::move(host_copy); }
+    auto host_res2 = br->reserve_or_fail(host_copy.data_alloc_size(spill_mem_type), spill_mem_type);
+    return host_copy.copy(host_res2);
+  }();
   EXPECT_FALSE(host_copy2.is_available());
   EXPECT_TRUE(host_copy2.is_spillable());
   EXPECT_EQ(host_copy2.stream().get(), stream.get());
@@ -371,7 +378,8 @@ TEST_P(StreamingTableChunk, SpillTrackingOnHostCopy)
     std::unordered_map<rapidsmpf::MemoryType, std::int64_t>{},
     std::nullopt,
     std::make_shared<rapidsmpf::StreamPool>(16),
-    stats);
+    stats,
+    spill_dir.path());
 
   auto samples = [&stats] {
     return stats->has_stat("buffer-spilled-time") ? stats->get_stat("buffer-spilled-time").count()
@@ -392,10 +400,14 @@ TEST_P(StreamingTableChunk, SpillTrackingOnHostCopy)
   // no token and must not be reported as a spill.
   std::ignore = round_trip(random_table_with_index(2025, 0, 0, 5));
   EXPECT_EQ(samples(), 0UL);
+  auto const spill_mem_name = rapidsmpf::to_lower(rapidsmpf::to_string(spill_mem_type));
+  EXPECT_FALSE(stats->has_stat("copy-device-to-" + spill_mem_name + "-bytes"));
+  EXPECT_FALSE(stats->has_stat("copy-" + spill_mem_name + "-to-device-bytes"));
 
   // A non-empty one does leave the device, so the round trip is recorded once.
   std::ignore = round_trip(random_table_with_index(2025, 64, 0, 5));
   EXPECT_EQ(samples(), 1UL);
+  EXPECT_GT(stats->get_stat("copy-device-to-" + spill_mem_name + "-bytes").value(), 0);
 }
 
 TEST_F(StreamingTableChunk, ToMessageRoundTrip)