diff --git a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp index b119ebc3b2e5..3fcc8f5ddaf1 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/detail/device_bloom_filter.hpp +++ b/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,12 +32,8 @@ 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. @@ -41,14 +41,12 @@ struct device_bloom_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 diff --git a/cpp/libcudf_streaming/src/bloom_filter.cpp b/cpp/libcudf_streaming/src/bloom_filter.cpp index 54dbfcf52cea..78acf8431f05 100644 --- a/cpp/libcudf_streaming/src/bloom_filter.cpp +++ b/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(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(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()) { diff --git a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu index 155cc522d320..0525f139babf 100644 --- a/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu +++ b/cpp/libcudf_streaming/src/detail/device_bloom_filter.cu @@ -30,7 +30,6 @@ #include #include -#include #include @@ -72,15 +71,9 @@ 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(storage_) % std::alignment_of_v == 0, "Allocation for bloom filter is not aligned."); @@ -88,19 +81,19 @@ device_bloom_filter::device_bloom_filter(std::size_t num_blocks, 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(storage), stream); + return device_bloom_filter(num_blocks, seed, const_cast(storage)); } std::unique_ptr device_bloom_filter::storage(std::size_t num_blocks, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr) { - return std::make_unique(num_blocks * sizeof(StorageType), stream, mr); + return std::make_unique( + num_blocks * sizeof(StorageType), std::alignment_of_v, 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(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(), "Hash values do not have correct type"); @@ -136,9 +128,8 @@ rmm::device_uvector device_bloom_filter::contains(cudf::table_view const& RAPIDSMPF_NVTX_FUNC_RANGE(); auto filter_ref = BloomFilterRefType{ static_cast(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 result{static_cast(view.size()), stream, mr}; filter_ref.contains_async(view.begin(), view.end(), 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_; } diff --git a/python/cudf_streaming/cudf_streaming/bloom_filter.pyx b/python/cudf_streaming/cudf_streaming/bloom_filter.pyx index c0554d5ddf21..5116a8aa44e2 100644 --- a/python/cudf_streaming/cudf_streaming/bloom_filter.pyx +++ b/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(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(ret, py_deleter)), ) - await ret + await await_cpp_future( + ret, on_cancel=lambda: shutdown_channels(ctx, bloom_filter, ch_in, ch_out) + ) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py b/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py index 2f4c055bc308..d00e7926e50e 100644 --- a/python/cudf_streaming/cudf_streaming/tests/test_bloom_filter.py +++ b/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