TableChunk.copy(): allow wiggle room (#672)

Allow some wiggle room based on number of columns.

Also:
- use `estimated_memory_usage()`
- add missing python bindings for `TableChunk.copy()`

Authors:
  - Mads R. B. Kristensen (https://github.com/madsbk)

Approvers:
  - Niranda Perera (https://github.com/nirandaperera)

URL: #672

Author: madsbk

cpp/include/rapidsmpf/streaming/cudf/table_chunk.hpp

@@ -75,7 +75,6 @@ class TableChunk {
      * the TableChunk is destroyed.
      *
      * @param table_view Device-resident table view.
-     * @param device_alloc_size Number of bytes allocated in device memory.
      * @param stream CUDA stream on which the table was created.
      * @param owner Object owning the memory backing @p table_view. This object will be
      * destroyed last when the TableChunk is destroyed or spilled.
@@ -91,7 +90,6 @@ class TableChunk {
      */
     TableChunk(
         cudf::table_view table_view,
-        std::size_t device_alloc_size,
         rmm::cuda_stream_view stream,
         OwningWrapper&& owner,
         ExclusiveView exclusive_view

cpp/src/streaming/cudf/table_chunk.cpp

@@ -5,9 +5,8 @@
 
 #include <memory>
 
-#include <rmm/aligned.hpp>
-
 #include <rapidsmpf/buffer/buffer.hpp>
+#include <rapidsmpf/integrations/cudf/utils.hpp>
 #include <rapidsmpf/streaming/cudf/table_chunk.hpp>
 
 namespace rapidsmpf::streaming {
@@ -24,7 +23,6 @@ TableChunk::TableChunk(std::unique_ptr<cudf::table> table, rmm::cuda_stream_view
 
 TableChunk::TableChunk(
     cudf::table_view table_view,
-    std::size_t device_alloc_size,
     rmm::cuda_stream_view stream,
     OwningWrapper&& owner,
     ExclusiveView exclusive_view
@@ -33,7 +31,8 @@ TableChunk::TableChunk(
       table_view_{table_view},
       stream_{stream},
       is_spillable_{static_cast<bool>(exclusive_view)} {
-    data_alloc_size_[static_cast<std::size_t>(MemoryType::DEVICE)] = device_alloc_size;
+    data_alloc_size_[static_cast<std::size_t>(MemoryType::DEVICE)] =
+        estimated_memory_usage(table_view, stream_);
     make_available_cost_ = 0;
 }
 
@@ -182,11 +181,11 @@ TableChunk TableChunk::copy(MemoryReservation& reservation) const {
                     // Handle the case where `cudf::pack` allocates slightly more than the
                     // input size. This can occur because cudf uses aligned allocations,
                     // which may exceed the requested size. To accommodate this, we
-                    // slightly increase the reservation if the packed data fits within
-                    // the aligned size.
+                    // allow some wiggle room.
                     if (packed_data->data->size > reservation.size()) {
-                        auto const aligned_size = rmm::align_up(reservation.size(), 1024);
-                        if (packed_data->data->size <= aligned_size) {
+                        auto const wiggle_room =
+                            1024 * static_cast<std::size_t>(table_view().num_columns());
+                        if (packed_data->data->size <= reservation.size() + wiggle_room) {
                             reservation =
                                 br->reserve(
                                       MemoryType::HOST, packed_data->data->size, true

cpp/tests/streaming/test_table_chunk.cpp

@@ -59,11 +59,7 @@ TEST_F(StreamingTableChunk, TableChunkOwner) {
     };
     auto make_chunk = [&](TableChunk::ExclusiveView exclusive_view) {
         return TableChunk{
-            expect,
-            expect.alloc_size(),
-            stream,
-            OwningWrapper(new int, deleter),
-            exclusive_view
+            expect, stream, OwningWrapper(new int, deleter), exclusive_view
         };
     };
     auto check_chunk = [&](TableChunk const& chunk, bool is_spillable) {
@@ -360,11 +356,7 @@ TEST_F(StreamingTableChunk, ToMessageNotSpillable) {
 
     auto deleter = [](void* p) { delete static_cast<int*>(p); };
     auto chunk = std::make_unique<TableChunk>(
-        expect,
-        expect.alloc_size(),
-        stream,
-        OwningWrapper(new int, deleter),
-        TableChunk::ExclusiveView::NO
+        expect, stream, OwningWrapper(new int, deleter), TableChunk::ExclusiveView::NO
     );
 
     Message m = to_message(seq, std::move(chunk));

python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pxd

@@ -10,6 +10,7 @@ from rmm.librmm.cuda_stream_view cimport cuda_stream_view
 from rmm.pylibrmm.stream cimport Stream
 
 from rapidsmpf.buffer.buffer cimport MemoryType
+from rapidsmpf.buffer.resource cimport cpp_MemoryReservation
 
 
 cdef extern from "<rapidsmpf/streaming/cudf/table_chunk.hpp>" nogil:
@@ -20,6 +21,7 @@ cdef extern from "<rapidsmpf/streaming/cudf/table_chunk.hpp>" nogil:
         size_t make_available_cost() noexcept
         cpp_table_view table_view() except +
         bool_t is_spillable() noexcept
+        cpp_TableChunk copy(cpp_MemoryReservation& reservation) except +
 
 cdef class TableChunk:
     cdef unique_ptr[cpp_TableChunk] _handle

python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pyx

@@ -3,11 +3,9 @@
 
 from cpython.object cimport PyObject
 from cython.operator cimport dereference as deref
-from libc.stddef cimport size_t
 from libc.stdint cimport uint64_t
 from libcpp.memory cimport unique_ptr
 from libcpp.utility cimport move
-from pylibcudf.column cimport Column
 from pylibcudf.libcudf.table.table_view cimport table_view as cpp_table_view
 from pylibcudf.table cimport Table
 
@@ -38,7 +36,6 @@ cdef extern from * nogil:
 
     std::unique_ptr<rapidsmpf::streaming::TableChunk> cpp_from_table_view_with_owner(
         cudf::table_view view,
-        std::size_t device_alloc_size,
         rmm::cuda_stream_view stream,
         PyObject *owner,
         void(*py_deleter)(void *),
@@ -49,7 +46,6 @@ cdef extern from * nogil:
         Py_XINCREF(owner);
         return std::make_unique<rapidsmpf::streaming::TableChunk>(
             view,
-            device_alloc_size,
             stream,
             rapidsmpf::OwningWrapper(owner, py_deleter),
             exclusive_view ?
@@ -66,6 +62,15 @@ cdef extern from * nogil:
             table->make_available(*reservation)
         );
     }
+
+    std::unique_ptr<rapidsmpf::streaming::TableChunk> cpp_table_copy(
+        std::unique_ptr<rapidsmpf::streaming::TableChunk> const& table,
+        rapidsmpf::MemoryReservation* reservation
+    ) {
+        return std::make_unique<rapidsmpf::streaming::TableChunk>(
+            table->copy(*reservation)
+        );
+    }
     }
     """
     unique_ptr[cpp_TableChunk] cpp_release_table_chunk_from_message(
@@ -75,6 +80,9 @@ cdef extern from * nogil:
     unique_ptr[cpp_TableChunk] cpp_table_make_available(
         unique_ptr[cpp_TableChunk], cpp_MemoryReservation*
     ) except +
+    unique_ptr[cpp_TableChunk] cpp_table_copy(
+        unique_ptr[cpp_TableChunk], cpp_MemoryReservation*
+    ) except +
 
 cdef class TableChunk:
     """
@@ -161,15 +169,10 @@ cdef class TableChunk:
         ensure the stream remains valid for the lifetime of the streaming pipeline.
         """
         cdef cuda_stream_view _stream = stream.view()
-        cdef size_t device_alloc_size = 0
-        for col in table.columns():
-            device_alloc_size += (<Column?>col).device_buffer_size()
-
         cdef cpp_table_view view = table.view()
         return TableChunk.from_handle(
             cpp_from_table_view_with_owner(
                 view,
-                device_alloc_size,
                 _stream,
                 <PyObject *>table,
                 py_deleter,
@@ -429,3 +432,29 @@ cdef class TableChunk:
         True if the table chunk can be spilled, otherwise, False.
         """
         return deref(self.handle_ptr()).is_spillable()
+
+    def copy(self, MemoryReservation reservation not None):
+        """
+        Create a deep copy of this table chunk.
+
+        All buffers are allocated for the new table chunk using the provided
+        memory reservation, which also determines the target memory type of
+        the copy.
+
+        Parameters
+        ----------
+        reservation
+            Memory reservation to consume for allocating the buffers of the
+            new table chunk.
+
+        Returns
+        -------
+        TableChunk
+            A new table chunk containing a deep copy of this chunk's data and
+            metadata.
+        """
+        cdef unique_ptr[cpp_TableChunk] ret
+        cdef cpp_MemoryReservation* res = reservation._handle.get()
+        with nogil:
+            ret = cpp_table_copy(self._handle, res)
+        return TableChunk.from_handle(move(ret))

python/rapidsmpf/rapidsmpf/tests/streaming/test_table_chunk.py

@@ -106,6 +106,31 @@ def test_roundtrip(context: Context, stream: Stream, *, exclusive_view: bool) ->
     assert_eq(expect, table_chunk6.table_view())
 
 
+def test_copy_roundtrip(context: Context, stream: Stream) -> None:
+    for nrows, ncols in [(1, 1), (1000, 100), (1, 1000)]:
+        expect = cudf_to_pylibcudf_table(
+            cudf.DataFrame(
+                {
+                    f"{name}": cupy.random.random(nrows, dtype=cupy.float32)
+                    for name in range(ncols)
+                }
+            )
+        )
+
+        tbl1 = TableChunk.from_pylibcudf_table(expect, stream, exclusive_view=True)
+        res, _ = context.br().reserve(
+            MemoryType.HOST,
+            tbl1.data_alloc_size(MemoryType.DEVICE),
+            allow_overbooking=True,
+        )
+        tbl2 = tbl1.copy(res)
+        res, _ = context.br().reserve(
+            MemoryType.DEVICE, tbl2.make_available_cost(), allow_overbooking=True
+        )
+        tbl3 = tbl2.make_available(res)
+        assert_eq(expect, tbl3.table_view())
+
+
 def test_spillable_messages(context: Context, stream: Stream) -> None:
     seq = 42
     df1 = random_table(1024)