Add optional owner slot to TableChunks (rapidsai#536)

When a `TableChunk` is constructed from a `cudf::table_view`, we are on the hook to keep the backing data alive. When we move `TableChunk`s into `Message`s and out again from Python we cannot do this in a sane way externally, so we must stash the owning object in the `TableChunk` itself.

The `PyObject` is stored in a `unique_ptr` with a custom deleter that, when called, acquires the GIL and decrefs the object. This way, even if we consume the `TableChunk` in a C++ node (rather than a Python node) its backing storage will be deallocated.

Authors:
  - Lawrence Mitchell (https://github.com/wence-)

Approvers:
  - Mads R. B. Kristensen (https://github.com/madsbk)
  - Gil Forsyth (https://github.com/gforsyth)

URL: rapidsai#536

Author: wence-

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

@@ -0,0 +1,62 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <memory>
+
+namespace rapidsmpf::streaming {
+
+/**
+ * @brief Utility class to store an arbitrary type-erased object while another object is
+ * alive.
+ *
+ * When sending messages through `Channel`s from Python, we typically need to keep various
+ * Python objects alive since the matching C++ objects only hold views.
+ *
+ * For example, when constructing a `TableChunk` from a pylibcudf `Table`, the
+ * `TableChunk` has a non-owning `cudf::table_view` of the `Table` and someone must be
+ * responsible for keeping the `Table` alive for the lifetime of the `TableChunk`. If we
+ * want to allow creation of such objects in Python with the ability to sink them on the
+ * C++ side we cannot rely on the Python side of things keeping the `Table` alive (the
+ * reference disappears!). Similarly when we send a message through a `Channel` the sender
+ * will, once pushed into the channel, drop the reference to the message payload and so,
+ * again, we need some way of keeping the payload alive.
+ *
+ * To square this circle, such C++ objects have an `OwningWrapper` slot that stores a
+ * type-erased pointer with, as far as we are concerned, unique ownership semantics. When
+ * this object is destroyed, the custom deleter runs and can do whatever deallocation is
+ * necessary.
+ *
+ * @warning Behaviour is undefined if the unique ownership semantic is not respected. The
+ * deleter may be called from any thread at any time, the implementer of the deleter is
+ * responsible for correct synchronisation with (for example) the Python GIL. Furthermore,
+ * the deleter may not throw: if an error occurs, the only safe thing to do is
+ * `std::terminate`.
+ *
+ * @warning When using this `OwningWrapper` inside a C++ object, make sure it is
+ * constructed first and destructed last.
+ */
+class OwningWrapper {
+  public:
+    /// @brief Callback used to delete the owned object.
+    using deleter_type = void (*)(void*);
+
+    OwningWrapper() = default;
+
+    /**
+     * @brief Take ownership and responsibility for the destruction of an object.
+     *
+     * @param obj Type-erased object to own.
+     * @param deleter Function called to destruct the object.
+     */
+    explicit OwningWrapper(void* obj, deleter_type deleter)
+        : obj_{owning_type(obj, deleter)} {}
+
+  private:
+    using owning_type = std::unique_ptr<void, deleter_type>;
+    owning_type obj_{nullptr, [](void*) {}};
+};
+}  // namespace rapidsmpf::streaming

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

@@ -19,6 +19,7 @@
 #include <rapidsmpf/streaming/core/channel.hpp>
 #include <rapidsmpf/streaming/core/context.hpp>
 #include <rapidsmpf/streaming/core/node.hpp>
+#include <rapidsmpf/streaming/cudf/owning_wrapper.hpp>
 
 namespace rapidsmpf::streaming {
 
@@ -58,12 +59,17 @@ class TableChunk {
      * @param table_view Device-resident table view.
      * @param device_alloc_size The number of bytes in device memory.
      * @param stream The CUDA stream on which the table was created.
+     * @param owner Optional object owning the memory backing the @p table_view. If it
+     * exists this object will be destructed last when the @p TableChunk is destroyed.
+     * This is typically used when constructing a @p TableChunk from python and we need to
+     * keep the owning python object alive.
      */
     TableChunk(
         std::uint64_t sequence_number,
         cudf::table_view table_view,
         std::size_t device_alloc_size,
-        rmm::cuda_stream_view stream
+        rmm::cuda_stream_view stream,
+        OwningWrapper&& owner = {}
     );
 
     /**
@@ -183,6 +189,9 @@ class TableChunk {
     [[nodiscard]] TableChunk spill_to_host(BufferResource* br);
 
   private:
+    ///< @brief Optional owning object if the TableChunk was constructed from a
+    ///< table_view.
+    OwningWrapper owner_{};
     std::uint64_t sequence_number_;
 
     // At most, one of the following unique pointers is non-null. If all of them are null,

cpp/src/shuffler/shuffler.cpp

@@ -4,7 +4,6 @@
  */
 
 #include <algorithm>
-#include <concepts>
 #include <functional>
 #include <memory>
 #include <numeric>

cpp/src/streaming/cudf/shuffler.cpp

@@ -122,7 +122,7 @@ coro::task<std::vector<PackedData>> ShufflerAsync::extract_async(shuffler::PartI
     auto all_extracted = all_extracted_unsafe();
 
     auto chunks = shuffler_.extract(pid);
-    lock.unlock();  // no longer need the lock
+    lock.unlock();
 
     // if all partitions have been extracted, notify all waiting tasks.
     if (all_extracted) {

cpp/src/streaming/cudf/table_chunk.cpp

@@ -25,9 +25,13 @@ TableChunk::TableChunk(
     std::uint64_t sequence_number,
     cudf::table_view table_view,
     std::size_t device_alloc_size,
-    rmm::cuda_stream_view stream
+    rmm::cuda_stream_view stream,
+    OwningWrapper&& owner
 )
-    : sequence_number_{sequence_number}, table_view_{table_view}, stream_{stream} {
+    : owner_{std::move(owner)},
+      sequence_number_{sequence_number},
+      table_view_{table_view},
+      stream_{stream} {
     data_alloc_size_[static_cast<std::size_t>(MemoryType::DEVICE)] = device_alloc_size;
     make_available_cost_ = 0;
 }

cpp/tests/streaming/test_shuffler.cpp

@@ -3,6 +3,8 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
+#include <algorithm>
+
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 

cpp/tests/streaming/test_table_chunk.cpp

@@ -17,6 +17,8 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
 
+#include <rapidsmpf/streaming/core/channel.hpp>
+#include <rapidsmpf/streaming/cudf/owning_wrapper.hpp>
 #include <rapidsmpf/streaming/cudf/table_chunk.hpp>
 
 #include "../utils.hpp"
@@ -43,6 +45,51 @@ TEST_F(StreamingTableChunk, FromTable) {
     CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk.table_view(), expect);
 }
 
+TEST_F(StreamingTableChunk, TableChunkOwner) {
+    constexpr unsigned int num_rows = 100;
+    constexpr std::int64_t seed = 1337;
+    constexpr std::uint64_t seq = 42;
+
+    cudf::table expect = random_table_with_index(seed, num_rows, 0, 10);
+    // Static because the deleter function is a void(*)(void*) which precludes the use of
+    // a lambda with captures.
+    static std::size_t num_deletions{0};
+    auto deleter = [](void* p) {
+        num_deletions++;
+        delete static_cast<int*>(p);
+    };
+    auto make_chunk = [&]() {
+        return TableChunk{
+            seq, expect, expect.alloc_size(), stream, OwningWrapper(new int, deleter)
+        };
+    };
+    auto check_chunk = [&](TableChunk const& chunk) {
+        EXPECT_EQ(chunk.sequence_number(), seq);
+        EXPECT_EQ(chunk.stream().value(), stream.value());
+        EXPECT_TRUE(chunk.is_available());
+        EXPECT_EQ(chunk.make_available_cost(), 0);
+        CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk.table_view(), expect);
+    };
+    {
+        auto chunk = make_chunk();
+        check_chunk(chunk);
+        EXPECT_EQ(num_deletions, 0);
+    }
+    EXPECT_EQ(num_deletions, 1);
+    {
+        auto msg = Message(std::make_unique<TableChunk>(make_chunk()));
+        EXPECT_EQ(num_deletions, 1);
+    }
+    EXPECT_EQ(num_deletions, 2);
+    {
+        auto msg = Message(std::make_unique<TableChunk>(make_chunk()));
+        auto chunk = msg.release<TableChunk>();
+        check_chunk(chunk);
+        EXPECT_EQ(num_deletions, 2);
+    }
+    EXPECT_EQ(num_deletions, 3);
+}
+
 TEST_F(StreamingTableChunk, FromTableView) {
     constexpr unsigned int num_rows = 100;
     constexpr std::int64_t seed = 1337;

dependencies.yaml

@@ -281,6 +281,7 @@ dependencies:
           - gdb
           - openmpi >=5.0  # See <https://github.com/rapidsai/rapidsmpf/issues/17>
           - valgrind
+          - gdb
   test_python:
     common:
       - output_types: conda

python/rapidsmpf/rapidsmpf/streaming/cudf/partition_chunk.pxd

@@ -21,23 +21,21 @@ cdef extern from "<rapidsmpf/streaming/cudf/partition.hpp>" nogil:
 
 cdef class PartitionMapChunk:
     cdef unique_ptr[cpp_PartitionMapChunk] _handle
-    cdef object _owner
 
     @staticmethod
     cdef PartitionMapChunk from_handle(
-        unique_ptr[cpp_PartitionMapChunk] handle, object owner
+        unique_ptr[cpp_PartitionMapChunk] handle
     )
     cdef const cpp_PartitionMapChunk* handle_ptr(self)
     cdef unique_ptr[cpp_PartitionMapChunk] release_handle(self)
 
 
 cdef class PartitionVectorChunk:
     cdef unique_ptr[cpp_PartitionVectorChunk] _handle
-    cdef object _owner
 
     @staticmethod
     cdef PartitionVectorChunk from_handle(
-        unique_ptr[cpp_PartitionVectorChunk] handle, object owner
+        unique_ptr[cpp_PartitionVectorChunk] handle
     )
     cdef const cpp_PartitionVectorChunk* handle_ptr(self)
     cdef unique_ptr[cpp_PartitionVectorChunk] release_handle(self)

python/rapidsmpf/rapidsmpf/streaming/cudf/partition_chunk.pyx

@@ -18,7 +18,7 @@ cdef class PartitionMapChunk:
 
     @staticmethod
     cdef PartitionMapChunk from_handle(
-        unique_ptr[cpp_PartitionMapChunk] handle, object owner
+        unique_ptr[cpp_PartitionMapChunk] handle
     ):
         """
         Construct a PartitionMapChunk from an existing C++ handle.
@@ -27,17 +27,14 @@ cdef class PartitionMapChunk:
         ----------
         handle
             A unique pointer to a C++ PartitionMapChunk.
-        owner
-            An optional Python object to keep alive for as long as this
-            PartitionMapChunk exists (e.g., to maintain resource lifetime).
 
         Returns
         -------
         A new PartitionMapChunk wrapping the given handle.
         """
+
         cdef PartitionMapChunk ret = PartitionMapChunk.__new__(PartitionMapChunk)
         ret._handle = move(handle)
-        ret._owner = owner
         return ret
 
     @staticmethod
@@ -58,8 +55,7 @@ cdef class PartitionMapChunk:
         return PartitionMapChunk.from_handle(
             make_unique[cpp_PartitionMapChunk](
                 message._handle.release[cpp_PartitionMapChunk]()
-            ),
-            owner = None,
+            )
         )
 
     def into_message(self, Message message not None):
@@ -84,6 +80,8 @@ cdef class PartitionMapChunk:
         --------
         The PartitionMapChunk is released and must not be used after this call.
         """
+        if not message.empty():
+            raise ValueError("cannot move into a non-empty message")
         message._handle = cpp_Message(self.release_handle())
 
     cdef const cpp_PartitionMapChunk* handle_ptr(self):
@@ -145,7 +143,7 @@ cdef class PartitionVectorChunk:
 
     @staticmethod
     cdef PartitionVectorChunk from_handle(
-        unique_ptr[cpp_PartitionVectorChunk] handle, object owner
+        unique_ptr[cpp_PartitionVectorChunk] handle
     ):
         """
         Construct a PartitionVectorChunk from an existing C++ handle.
@@ -154,9 +152,6 @@ cdef class PartitionVectorChunk:
         ----------
         handle
             A unique pointer to a C++ PartitionVectorChunk.
-        owner
-            An optional Python object to keep alive for as long as this
-            PartitionVectorChunk exists (e.g., to maintain resource lifetime).
 
         Returns
         -------
@@ -166,7 +161,6 @@ cdef class PartitionVectorChunk:
             PartitionVectorChunk
         )
         ret._handle = move(handle)
-        ret._owner = owner
         return ret
 
     @staticmethod
@@ -187,8 +181,7 @@ cdef class PartitionVectorChunk:
         return PartitionVectorChunk.from_handle(
             make_unique[cpp_PartitionVectorChunk](
                 message._handle.release[cpp_PartitionVectorChunk]()
-            ),
-            owner = None,
+            )
         )
 
     def into_message(self, Message message not None):
@@ -213,6 +206,8 @@ cdef class PartitionVectorChunk:
         --------
         The PartitionVectorChunk is released and must not be used after this call.
         """
+        if not message.empty():
+            raise ValueError("cannot move into a non-empty message")
         message._handle = cpp_Message(self.release_handle())
 
     cdef const cpp_PartitionVectorChunk* handle_ptr(self):