Use new rapidsmpf utilities to handle cancellation in bloom filter awaitables

cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp

@@ -19,6 +19,10 @@ namespace cudf_streaming::detail {
 
 /**
  * @brief A bloom filter, used for approximate set membership queries.
+ *
+ * @note All methods of this class launch work on the streams provided. It is the caller's
+ * responsibility to ensure that data are valid to read/write on that stream as
+ * appropriate.
  */
 struct device_bloom_filter {
   /**
@@ -28,27 +32,21 @@ struct device_bloom_filter {
    * @param seed Seed used for hashing each value.
    * @param storage Storage to view as a bloom filter, must be appropriately
    * initialized.
-   * @param stream CUDA stream for device operations.
    */
-  device_bloom_filter(std::size_t num_blocks,
-                      std::uint64_t seed,
-                      void* storage,
-                      rmm::cuda_stream_view stream);
+  device_bloom_filter(std::size_t num_blocks, std::uint64_t seed, void* storage);
 
   /**
    * @brief Create a read-only filter.
    *
    * @param num_blocks Number of blocks in the filter.
    * @param seed Seed used for hashing each value.
    * @param storage View of storage, must be appropriately initialized.
-   * @param stream CUDA stream for device operations.
    *
    * @return A const-qualified bloom filter viewing the underlying storage.
    */
   static device_bloom_filter const view(std::size_t num_blocks,
                                         std::uint64_t seed,
-                                        void const* storage,
-                                        rmm::cuda_stream_view stream);
+                                        void const* storage);
 
   /**
    * @brief Create uninitialized storage for a filter.
@@ -98,11 +96,6 @@ struct device_bloom_filter {
                                                    rmm::cuda_stream_view stream,
                                                    rmm::device_async_resource_ref mr) const;
 
-  /**
-   * @brief @return The stream the underlying storage is valid on.
-   */
-  [[nodiscard]] rmm::cuda_stream_view stream() const noexcept;
-
   /**
    * @brief @return Pointer to the underlying storage.
    */
@@ -127,10 +120,9 @@ struct device_bloom_filter {
   [[nodiscard]] static std::size_t fitting_num_blocks(std::size_t l2size) noexcept;
 
  private:
-  std::size_t num_blocks_;        ///< Number of blocks used in the filter.
-  std::uint64_t seed_;            ///< Seed used when hashing values.
-  void* storage_;                 ///< Backing storage.
-  rmm::cuda_stream_view stream_;  ///< Stream storage is valid on.
+  std::size_t num_blocks_;  ///< Number of blocks used in the filter.
+  std::uint64_t seed_;      ///< Seed used when hashing values.
+  void* storage_;           ///< Backing storage.
 };
 
 }  // namespace cudf_streaming::detail

cpp/libcudf_streaming/src/bloom_filter.cpp

@@ -37,8 +37,8 @@ rapidsmpf::streaming::Actor bloom_filter::build(
   auto storage =
     cudf_streaming::detail::device_bloom_filter::storage(num_filter_blocks_, filter_stream, mr);
   RAPIDSMPF_CUDA_TRY(cudaMemsetAsync(storage->data(), 0, storage->size(), filter_stream));
-  auto filter = cudf_streaming::detail::device_bloom_filter(
-    num_filter_blocks_, seed_, storage->data(), filter_stream);
+  auto filter =
+    cudf_streaming::detail::device_bloom_filter(num_filter_blocks_, seed_, storage->data());
   rapidsmpf::CudaEvent build_event;
   build_event.record(filter_stream);
   while (!ch_out->is_shutdown()) {
@@ -48,6 +48,14 @@ rapidsmpf::streaming::Actor bloom_filter::build(
     chunk      = co_await chunk.make_available(
       ctx_,
       -rapidsmpf::safe_cast<std::int64_t>(chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE)));
+
+    // Reservation for the hash values in add.
+    auto res = co_await ctx_->memory(rapidsmpf::MemoryType::DEVICE)
+                 ->reserve_or_wait(rapidsmpf::safe_cast<std::size_t>(chunk.table_view().num_rows())
+                                     // TODO: no magic numbers: the hashing algorithm in
+                                     // `add` below returns an int64 column.
+                                     * sizeof(std::int64_t),
+                                   0);
     // Filter is allocated on `filter_stream`, but we run the additions on the chunk's
     // stream. The addition modifies global memory but we can safely launch two
     // kernels doing that concurrently because the updates are atomic.
@@ -67,9 +75,9 @@ rapidsmpf::streaming::Actor bloom_filter::build(
       [num_blocks = num_filter_blocks_, seed = seed_](rapidsmpf::Buffer const* left,
                                                       rapidsmpf::Buffer* right) {
         right->write_access([&](std::byte* out_bytes, rmm::cuda_stream_view stream) {
-          auto const in = cudf_streaming::detail::device_bloom_filter::view(
-            num_blocks, seed, left->data(), stream);
-          cudf_streaming::detail::device_bloom_filter(num_blocks, seed, out_bytes, stream)
+          auto const in =
+            cudf_streaming::detail::device_bloom_filter::view(num_blocks, seed, left->data());
+          cudf_streaming::detail::device_bloom_filter(num_blocks, seed, out_bytes)
             .merge(in, stream);
         });
       });
@@ -95,7 +103,7 @@ rapidsmpf::streaming::Actor bloom_filter::apply(
   auto stream = storage.stream();
   rapidsmpf::CudaEvent event;
   auto filter =
-    cudf_streaming::detail::device_bloom_filter(num_filter_blocks_, seed_, storage.data(), stream);
+    cudf_streaming::detail::device_bloom_filter(num_filter_blocks_, seed_, storage.data());
   auto meta = co_await ch_in->receive_metadata();
   if (!meta.empty()) { co_await ch_out->send_metadata(std::move(meta)); }
   while (!ch_out->is_shutdown()) {

cpp/libcudf_streaming/src/detail/device_bloom_filter.cu

@@ -30,7 +30,6 @@
 
 #include <cudf/hashing.hpp>
 #include <cudf/table/table_view.hpp>
-#include <cudf/utilities/memory_resource.hpp>
 
 #include <cudf_streaming/detail/device_bloom_filter.hpp>
 
@@ -72,35 +71,29 @@ using StorageType = BloomFilterRefType::filter_block_type;
 
 }  // namespace
 
-device_bloom_filter::device_bloom_filter(std::size_t num_blocks,
-                                         std::uint64_t seed,
-                                         void* storage,
-                                         rmm::cuda_stream_view stream)
-  : num_blocks_{num_blocks}, seed_{seed}, storage_{storage}, stream_{stream}
+device_bloom_filter::device_bloom_filter(std::size_t num_blocks, std::uint64_t seed, void* storage)
+  : num_blocks_{num_blocks}, seed_{seed}, storage_{storage}
 {
-  // TODO: use an aligned allocator adaptor to ensure this holds.
-  // Today all RMM device allocators guarantee at least 256 byte alignment, but that is
-  // an implementation detail.
   RAPIDSMPF_EXPECTS(
     reinterpret_cast<std::uintptr_t>(storage_) % std::alignment_of_v<StorageType> == 0,
     "Allocation for bloom filter is not aligned.");
 }
 
 device_bloom_filter const device_bloom_filter::view(std::size_t num_blocks,
                                                     std::uint64_t seed,
-                                                    void const* storage,
-                                                    rmm::cuda_stream_view stream)
+                                                    void const* storage)
 {
   // const-cast is safe because the returned object is also const and therefore can't
   // call methods that throw away constness.
-  return device_bloom_filter(num_blocks, seed, const_cast<void*>(storage), stream);
+  return device_bloom_filter(num_blocks, seed, const_cast<void*>(storage));
 }
 
 std::unique_ptr<rmm::device_buffer> device_bloom_filter::storage(std::size_t num_blocks,
                                                                  rmm::cuda_stream_view stream,
                                                                  rmm::device_async_resource_ref mr)
 {
-  return std::make_unique<rmm::device_buffer>(num_blocks * sizeof(StorageType), stream, mr);
+  return std::make_unique<rmm::device_buffer>(
+    num_blocks * sizeof(StorageType), std::alignment_of_v<StorageType>, stream, mr);
 }
 
 void device_bloom_filter::add(cudf::table_view const& values_to_hash,
@@ -110,8 +103,7 @@ void device_bloom_filter::add(cudf::table_view const& values_to_hash,
   RAPIDSMPF_NVTX_FUNC_RANGE();
   auto filter_ref = BloomFilterRefType{
     static_cast<StorageType*>(storage_), num_blocks_, cuco::thread_scope_device, {}};
-  auto hashes = cudf::hashing::xxhash_64(
-    values_to_hash, seed_, stream, cudf::get_current_device_resource_ref());
+  auto hashes    = cudf::hashing::xxhash_64(values_to_hash, seed_, stream, mr);
   auto hash_view = hashes->view();
   RAPIDSMPF_EXPECTS(hash_view.type().id() == cudf::type_to_id<KeyType>(),
                     "Hash values do not have correct type");
@@ -136,9 +128,8 @@ rmm::device_uvector<bool> device_bloom_filter::contains(cudf::table_view const&
   RAPIDSMPF_NVTX_FUNC_RANGE();
   auto filter_ref = BloomFilterRefType{
     static_cast<StorageType*>(storage_), num_blocks_, cuco::thread_scope_device, {}};
-  auto hashes =
-    cudf::hashing::xxhash_64(values, seed_, stream, cudf::get_current_device_resource_ref());
-  auto view = hashes->view();
+  auto hashes = cudf::hashing::xxhash_64(values, seed_, stream, mr);
+  auto view   = hashes->view();
   rmm::device_uvector<bool> result{static_cast<std::size_t>(view.size()), stream, mr};
   filter_ref.contains_async(view.begin<KeyType>(), view.end<KeyType>(), result.begin(), stream);
   return result;
@@ -149,8 +140,6 @@ std::size_t device_bloom_filter::fitting_num_blocks(std::size_t l2size) noexcept
   return (l2size * 2) / (3 * sizeof(StorageType));
 }
 
-rmm::cuda_stream_view device_bloom_filter::stream() const noexcept { return stream_; }
-
 void* device_bloom_filter::data() noexcept { return storage_; }
 
 void const* device_bloom_filter::data() const noexcept { return storage_; }

python/cudf_streaming/cudf_streaming/bloom_filter.pyx

@@ -18,6 +18,7 @@ from rapidsmpf.streaming._detail.libcoro_spawn_task cimport cpp_set_py_future
 from rapidsmpf.streaming.chunks.utils cimport py_deleter
 from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel
 from rapidsmpf.streaming.core.context cimport Context, cpp_Context
+from rapidsmpf.streaming.core.cancellation import (await_cpp_future, shutdown_channels)
 
 import asyncio
 
@@ -208,7 +209,11 @@ cdef class BloomFilter:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter)),
             )
-        await ret
+        # Note: multi-rank, if we get an exception we can't cancel the
+        # in-progress AllReduce, so this might still hang.
+        await await_cpp_future(
+            ret, on_cancel=lambda: shutdown_channels(ctx, ch_in, ch_out)
+        )
 
     async def apply(
         self,
@@ -248,4 +253,6 @@ cdef class BloomFilter:
                 cpp_set_py_future,
                 move(cpp_OwningWrapper(<void*><PyObject*>ret, py_deleter)),
             )
-        await ret
+        await await_cpp_future(
+            ret, on_cancel=lambda: shutdown_channels(ctx, bloom_filter, ch_in, ch_out)
+        )

python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py

@@ -172,3 +172,60 @@ def test_bloom_filter_empty_build_filters_all(
     )
     result.stream.synchronize()
     assert_eq(result.table_view(), expected)
+
+
+def test_bloom_filter_build_exception_no_shutdown(
+    context: Context, comm: Communicator
+) -> None:
+    if comm.nranks != 1:
+        pytest.skip("Only support single-rank runs")
+
+    stream = context.br().stream_pool.get_stream()
+    bloom = BloomFilter(
+        context,
+        comm,
+        seed=42,
+        num_filter_blocks=BloomFilter.fitting_num_blocks(1 << 20),
+    )
+    ch_in: Channel[TableChunk] = context.create_channel()
+    ch_out: Channel[BloomFilterChunk] = context.create_channel()
+    messages = [
+        Message(
+            sequence_number,
+            make_table(
+                np.arange(10, dtype=np.int32),
+                stream=stream,
+                br=context.br(),
+            ),
+        )
+        for sequence_number in range(3)
+    ]
+
+    async def recv_then_raise(
+        context: Context, ch_in: Channel[BloomFilterChunk]
+    ):
+        await ch_in.recv(context)
+        raise RuntimeError("Raising but didn't shutdown channel")
+
+    # With no consumer for ch_out, the bloom-filter build blocks while draining
+    # its output channel. pytest-timeout interrupts run_actor_network, which
+    # must cancel and drain the worker before propagating the timeout failure.
+    with pytest.RaisesGroup(
+        pytest.RaisesExc(RuntimeError, match="didn't shutdown channel")
+    ):
+        run_actor_network(
+            context,
+            actors=[
+                push_to_channel(context, ch_in, messages),
+                bloom.build(context, ch_in=ch_in, ch_out=ch_out, tag=0),
+                recv_then_raise(context, ch_out),
+            ],
+        )
+
+    async def recv_after_cancellation() -> Message | None:
+        return await asyncio.wait_for(ch_out.recv(context), timeout=1)
+
+    # Cancellation should close the output channel. Without shutting down the
+    # channels inside the bloom filter if we get a cancellation, this
+    # receive picks up the message that is still in the channel.
+    assert asyncio.run(recv_after_cancellation()) is None