From 88208180a8870a09cfe2cccd2070ff4c446ed0fc Mon Sep 17 00:00:00 2001 From: niranda perera Date: Mon, 8 Jun 2026 14:05:56 -0700 Subject: [PATCH 01/14] rebased Signed-off-by: niranda perera --- cpp/include/rapidsmpf/utils/misc.hpp | 45 +++ cpp/tests/CMakeLists.txt | 13 +- cpp/tests/streaming/test_allgather.cpp | 2 - cpp/tests/test_shuffler.cpp | 537 +++++++++---------------- cpp/tests/utils.hpp | 51 ++- 5 files changed, 279 insertions(+), 369 deletions(-) diff --git a/cpp/include/rapidsmpf/utils/misc.hpp b/cpp/include/rapidsmpf/utils/misc.hpp index 96ddfe185..0407731d7 100644 --- a/cpp/include/rapidsmpf/utils/misc.hpp +++ b/cpp/include/rapidsmpf/utils/misc.hpp @@ -4,8 +4,10 @@ */ #pragma once +#include #include #include +#include #include #include #include @@ -193,6 +195,49 @@ constexpr T safe_div(T x, T y) { return (y == 0) ? 0 : x / y; } +/** + * @brief Computes the ceiling of the division of two integers. + * + * Returns the smallest integer not less than `x / y`. Both operands must be + * non-negative and the denominator must be non-zero. + * + * @tparam T An integral type. + * @param x The numerator (must be non-negative). + * @param y The denominator (must be positive). + * @return T The ceiling of `x / y`. + */ +template +constexpr T ceil_div(T x, T y) { + return (x + y - 1) / y; +} + +/** + * @brief Splits the index range `[0, count)` into exactly `num_chunks` contiguous chunks. + * + * Each chunk is a half-open `[begin, end)` index pair. Chunks are front-loaded with + * size `ceil(count / num_chunks)`; the last non-empty chunk may be smaller and, when + * `count < num_chunks`, the trailing chunks are empty (`begin == end`). The chunks + * exactly tile `[0, count)`, so their sizes sum to `count`. + * + * Unlike `std::ranges::chunk_view` (C++23), this always yields exactly `num_chunks` + * chunks, injecting empty trailing chunks as needed. + * + * @param count The number of elements to split. + * @param num_chunks The number of chunks to produce (must be positive). + * @return A lazy view of `num_chunks` `std::pair` `[begin, + * end)` index pairs. + */ +[[nodiscard]] inline auto chunk_indices(std::size_t count, std::size_t num_chunks) { + std::size_t const chunk_size = ceil_div(count, num_chunks); + return std::views::iota(std::size_t{0}, num_chunks) + | std::views::transform([count, chunk_size](std::size_t k) { + return std::pair{ + std::min(k * chunk_size, count), + std::min((k + 1) * chunk_size, count) + }; + }); +} + // Macro to concatenate two tokens x and y. #define RAPIDSMPF_CONCAT_DETAIL_(x, y) x##y #define RAPIDSMPF_CONCAT(x, y) RAPIDSMPF_CONCAT_DETAIL_(x, y) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 12421123a..b90449df2 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -86,6 +86,7 @@ target_sources( test_pausable_thread_loop.cpp test_progress_thread.cpp test_rmm_resource_adaptor.cpp + test_shuffler.cpp test_sparse_alltoall.cpp test_spill_manager.cpp test_spilling.cpp @@ -99,7 +100,8 @@ target_sources( if(RAPIDSMPF_HAVE_STREAMING) target_sources( test_sources - PRIVATE streaming/test_allreduce.cpp + PRIVATE streaming/test_allgather.cpp + streaming/test_allreduce.cpp streaming/test_channel.cpp streaming/test_error_handling.cpp streaming/test_fanout.cpp @@ -113,19 +115,14 @@ endif() # cudf-dependent test sources, gated behind BUILD_CUDF_TESTS if(BUILD_CUDF_TESTS) - target_sources( - test_sources PRIVATE test_partition.cpp test_shuffler.cpp test_shuffler_many_streams.cpp - ) + target_sources(test_sources PRIVATE test_partition.cpp test_shuffler_many_streams.cpp) target_link_libraries( test_sources PRIVATE cudf_streaming::cudf_streaming cudf::cudftestutil cudf::cudftestutil_impl ) target_compile_definitions(test_sources PRIVATE RAPIDSMPF_HAVE_CUDF) if(RAPIDSMPF_HAVE_STREAMING) - target_sources( - test_sources PRIVATE streaming/test_allgather.cpp streaming/test_leaf_actor.cpp - streaming/test_shuffler.cpp - ) + target_sources(test_sources PRIVATE streaming/test_leaf_actor.cpp streaming/test_shuffler.cpp) endif() endif() diff --git a/cpp/tests/streaming/test_allgather.cpp b/cpp/tests/streaming/test_allgather.cpp index 6805f3100..c94541baa 100644 --- a/cpp/tests/streaming/test_allgather.cpp +++ b/cpp/tests/streaming/test_allgather.cpp @@ -10,8 +10,6 @@ #include #include -#include - #include #include diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index 9c8786929..4c5da8cdc 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -6,17 +6,15 @@ #include #include #include +#include #include #include #include #include -#include -#include -#include -#include -#include +#include +#include #include #include @@ -32,7 +30,7 @@ extern Environment* GlobalEnvironment; TEST(ReceivedChunks, spill_skips_control_messages) { - auto mr = cudf::get_current_device_resource_ref(); + auto mr = rmm::mr::get_current_device_resource_ref(); auto br = rapidsmpf::BufferResource::create(mr); rapidsmpf::shuffler::detail::ReceivedChunks received; @@ -49,9 +47,9 @@ TEST(ReceivedChunks, spill_skips_control_messages) { } TEST(ReceivedChunks, spill_respects_amount) { - auto mr = cudf::get_current_device_resource_ref(); + auto mr = rmm::mr::get_current_device_resource_ref(); auto br = rapidsmpf::BufferResource::create(mr); - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_default; rapidsmpf::shuffler::detail::ReceivedChunks received; constexpr std::size_t chunk_size = 100; @@ -74,8 +72,8 @@ TEST(ReceivedChunks, spill_respects_amount) { } TEST(MetadataMessage, round_trip) { - auto stream = cudf::get_default_stream(); - auto mr = cudf::get_current_device_resource_ref(); + auto stream = rmm::cuda_stream_default; + auto mr = rmm::mr::get_current_device_resource_ref(); auto br = rapidsmpf::BufferResource::create(mr); auto metadata = iota_vector(100); @@ -108,6 +106,8 @@ TEST(MetadataMessage, round_trip) { EXPECT_EQ(metadata, *result.release_metadata_buffer()); } +namespace { + using MemoryLimitsMap = std::unordered_map; // Help function to get the `memory_limits` argument for a `BufferResource` @@ -130,20 +130,95 @@ MemoryLimitsMap get_memory_limits_map(rapidsmpf::MemoryType priorities) { return ret; } -/// @tparam InsertFn: lambda that inserts the packed chunks into the shuffler. -/// Signature: void(std::vector&& packed_chunks) -/// @tparam InsertFinishedFn: lambda that inserts the finished flag into the shuffler. -/// Signature: void() -template +// Conservation-preserving data model shared by the shuffler round-trip tests. +// +// We split the index range [0, total_num_rows) into total_num_partitions^2 contiguous +// sub-regions via chunk_indices (front-loaded, so when N < P*P the trailing sub-regions +// are empty). Sub-region (local_pidx, split_idx) is piece k = local_pidx*P + split_idx +// and is routed to destination partition split_idx; input region local_pidx is the union +// of its P sub-regions. The pieces exactly tile [0,N), so the total shuffled data == N +// regardless of rank/partition counts (conservation). A per-shuffle `base` offset is +// added to every value so distinct shuffles carry distinct data. + +// Produces the non-empty sub-regions of one owned input region `local_pidx`, keyed by +// destination partition. Since local_partitions() across ranks partition [0,P), every +// input region is produced exactly once, so rows are not replicated. +std::unordered_map +make_partition_data( + rapidsmpf::shuffler::PartID total_num_partitions, + std::size_t total_num_rows, + rapidsmpf::shuffler::PartID local_pidx, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource& br, + std::int64_t base = 0 +) { + auto const P = static_cast(total_num_partitions); + auto const pieces = rapidsmpf::chunk_indices(total_num_rows, P * P); + + std::unordered_map chunks; + for (rapidsmpf::shuffler::PartID split_idx = 0; split_idx < total_num_partitions; + ++split_idx) + { + auto [start, end] = pieces[static_cast(local_pidx) * P + split_idx]; + if (end > start) { + chunks.emplace( + split_idx, + generate_packed_data( + end - start, base + static_cast(start), stream, br + ) + ); + } + } + return chunks; +} + +// Verifies that the `received` chunks for partition `j` match the non-empty sub-regions +// expected for it. +void validate_partition_data( + std::vector received, + rapidsmpf::shuffler::PartID total_num_partitions, + std::size_t total_num_rows, + rapidsmpf::shuffler::PartID j, + rapidsmpf::BufferResource& br, + std::int64_t base = 0 +) { + auto const P = static_cast(total_num_partitions); + auto const pieces = rapidsmpf::chunk_indices(total_num_rows, P * P); + + // Locally recompute the non-empty (offset, count) sub-regions expected for partition + // j, in increasing input-region-index (== increasing offset) order. + std::vector> expected; + for (rapidsmpf::shuffler::PartID i = 0; i < total_num_partitions; ++i) { + auto [start, end] = pieces[static_cast(i) * P + j]; + if (end > start) { + expected.emplace_back(base + static_cast(start), end - start); + } + } + + EXPECT_EQ(received.size(), expected.size()); + + // Sort received chunks by their first metadata int64 (== offset) so they align 1:1 + // with the expected list, which is already in offset order. + std::sort(received.begin(), received.end(), [](auto const& a, auto const& b) { + std::int64_t oa{}, ob{}; + std::memcpy(&oa, a.metadata->data(), sizeof(std::int64_t)); + std::memcpy(&ob, b.metadata->data(), sizeof(std::int64_t)); + return oa < ob; + }); + + for (std::size_t k = 0; k < received.size() && k < expected.size(); ++k) { + auto const [off, cnt] = expected[k]; + auto const cs = received[k].stream(); + EXPECT_NO_FATAL_FAILURE( + validate_packed_data(std::move(received[k]), cnt, off, cs, br) + ); + } +} + void test_shuffler( - std::shared_ptr const& comm, rapidsmpf::shuffler::Shuffler& shuffler, rapidsmpf::shuffler::PartID total_num_partitions, - InsertFn&& insert_fn, - InsertFinishedFn&& insert_finished_fn, std::size_t total_num_rows, - std::int64_t seed, - cudf::hash_id hash_fn, rmm::cuda_stream_view stream, rapidsmpf::BufferResource* br ) { @@ -151,95 +226,35 @@ void test_shuffler( // shuffle shouldn't get near 30s. std::chrono::seconds const wait_timeout(30); - // Every rank creates the full input table and all the expected partitions (also - // partitions this rank might not get after the shuffle). - cudf::table full_input_table = random_table_with_index(seed, total_num_rows, 0, 10); - auto [expect_partitions, owner] = cudf_streaming::integrations::partition_and_split( - full_input_table, - {1}, - static_cast(total_num_partitions), - hash_fn, - seed, - stream, - br, - rapidsmpf::AllowOverbooking::YES - ); - - cudf::size_type row_offset = 0; - cudf::size_type partiton_size = - full_input_table.num_rows() / static_cast(total_num_partitions); - for (rapidsmpf::shuffler::PartID i = 0; i < total_num_partitions; ++i) { - // To simulate that `full_input_table` is distributed between multiple ranks, - // we divided them into `total_num_partitions` number of partitions and pick - // the partitions this rank should use as input. We pick using round robin but - // any distribution would work (as long as no rows are picked by multiple - // ranks). - // TODO: we should test different distributions of the input partitions. - if (rapidsmpf::shuffler::Shuffler::round_robin(comm, i, total_num_partitions) - == comm->rank()) - { - cudf::size_type row_end = row_offset + partiton_size; - if (i == total_num_partitions - 1) { - // Include the reminder of rows in the very last partition. - row_end = full_input_table.num_rows(); - } - // Select the partition from the full input table. - auto slice = cudf::slice(full_input_table, {row_offset, row_end}).at(0); - // Hash the `slice` into chunks and pack (serialize) them. - auto packed_chunks = cudf_streaming::integrations::partition_and_pack( - slice, - {1}, - static_cast(total_num_partitions), - hash_fn, - seed, - stream, - br, - rapidsmpf::AllowOverbooking::YES - ); - // Add the chunks to the shuffle - insert_fn(std::move(packed_chunks)); - } - row_offset += partiton_size; + for (rapidsmpf::shuffler::PartID local_pidx : shuffler.local_partitions()) { + shuffler.insert(make_partition_data( + total_num_partitions, total_num_rows, local_pidx, stream, *br + )); } - // Tell the shuffler that we have no more input partitions. - insert_finished_fn(); + shuffler.insert_finished(); EXPECT_NO_THROW(shuffler.wait(wait_timeout)); - for (auto finished_partition : shuffler.local_partitions()) { - auto packed_chunks = shuffler.extract(finished_partition); - auto result = cudf_streaming::integrations::unpack_and_concat( - rapidsmpf::unspill_partitions( - std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES - ), - stream, - br, - rapidsmpf::AllowOverbooking::YES - ); - // We should only receive the partitions assigned to this rank. - EXPECT_EQ( - shuffler.partition_owner(comm, finished_partition, total_num_partitions), - comm->rank() - ); - - // Check the result while ignoring the row order. - CUDF_TEST_EXPECT_TABLES_EQUIVALENT( - sort_table(result), sort_table(expect_partitions[finished_partition]) + for (auto j : shuffler.local_partitions()) { + validate_partition_data( + shuffler.extract(j), total_num_partitions, total_num_rows, j, *br ); } } +} // namespace + class MemoryLimits_NumPartition - : public cudf::test::BaseFixtureWithParam< + : public ::testing::TestWithParam< std::tuple> { public: void SetUp() override { - stream = cudf::get_default_stream(); - memory_limits = std::get<0>(GetParam()); - total_num_partitions = std::get<1>(GetParam()); - total_num_rows = std::get<2>(GetParam()); + stream = rmm::cuda_stream_default; + std::tie(memory_limits, total_num_partitions, total_num_rows) = GetParam(); br = rapidsmpf::BufferResource::create( - mr(), rapidsmpf::PinnedMemoryResource::Disabled, memory_limits + rmm::mr::get_current_device_resource_ref(), + rapidsmpf::PinnedMemoryResource::Disabled, + memory_limits ); shuffler = std::make_unique( @@ -258,8 +273,6 @@ class MemoryLimits_NumPartition MemoryLimitsMap memory_limits; rapidsmpf::shuffler::PartID total_num_partitions; std::size_t total_num_rows; - std::int64_t seed = 42; - cudf::hash_id hash_fn = cudf::hash_id::HASH_MURMUR3; rmm::cuda_stream_view stream; std::shared_ptr br; std::unique_ptr shuffler; @@ -285,82 +298,54 @@ INSTANTIATE_TEST_SUITE_P( ); TEST_P(MemoryLimits_NumPartition, round_trip) { - EXPECT_NO_FATAL_FAILURE(test_shuffler( - GlobalEnvironment->comm_, - *shuffler, - total_num_partitions, - [&](auto&& packed_chunks) { shuffler->insert(std::move(packed_chunks)); }, - [&]() { shuffler->insert_finished(); }, - total_num_rows, - seed, - hash_fn, - stream, - br.get() - )); + EXPECT_NO_FATAL_FAILURE( + test_shuffler(*shuffler, total_num_partitions, total_num_rows, stream, br.get()) + ); } // Test that the same communicator can be used concurrently by multiple shufflers in // separate threads -class ConcurrentShuffleTest - : public cudf::test::BaseFixtureWithParam> { +class ConcurrentShuffleTest : public ::testing::TestWithParam< + std::tuple> { public: void SetUp() override { - num_shufflers = std::get<0>(GetParam()); - total_num_partitions = - static_cast(std::get<1>(GetParam())); + std::tie(num_shufflers, total_num_partitions) = GetParam(); // these resources will be used by multiple threads to instantiate shufflers - br = rapidsmpf::BufferResource::create(mr()); - stream = cudf::get_default_stream(); + br = + rapidsmpf::BufferResource::create(rmm::mr::get_current_device_resource_ref()); + stream = rmm::cuda_stream_default; } void TearDown() override {} // test run for each thread. The test follows the same logic as // `MemoryLimits_NumPartition` test, but without any memory limitations - template - void RunTest(int t_id, InsertFn&& insert_fn, InsertFinishedFn&& insert_finished_fn) { + void RunTest(std::size_t t_id) { rapidsmpf::shuffler::Shuffler shuffler( GlobalEnvironment->comm_, - t_id, // op_id, use t_id as a proxy + static_cast(t_id), // op_id, use t_id as a proxy total_num_partitions, br.get() ); EXPECT_NO_FATAL_FAILURE(test_shuffler( - GlobalEnvironment->comm_, shuffler, total_num_partitions, - [&](auto&& packed_chunks) { insert_fn(shuffler, std::move(packed_chunks)); }, - [&]() { insert_finished_fn(shuffler); }, 100'000, // total_num_rows - t_id, // seed - cudf::hash_id::HASH_MURMUR3, stream, br.get() )); } - template - void RunTestTemplate(InsertFn insert_fn, InsertFinishedFn insert_finished_fn) { + void RunTestTemplate() { std::vector> futures; - futures.reserve(static_cast(num_shufflers)); - - for (int t_id = 0; t_id < num_shufflers; t_id++) { - // pass a copy of the insert_fn and insert_finished_fn to each thread - futures.push_back( - std::async( - std::launch::async, - [this, - t_id, - insert_fn1 = insert_fn, - insert_finished_fn1 = insert_finished_fn] { - ASSERT_NO_FATAL_FAILURE(this->RunTest( - t_id, std::move(insert_fn1), std::move(insert_finished_fn1) - )); - } - ) - ); + futures.reserve(num_shufflers); + + for (std::size_t t_id = 0; t_id < num_shufflers; t_id++) { + futures.push_back(std::async(std::launch::async, [this, t_id] { + ASSERT_NO_FATAL_FAILURE(this->RunTest(t_id)); + })); } for (auto& f : futures) { @@ -368,7 +353,7 @@ class ConcurrentShuffleTest } } - int num_shufflers; + std::size_t num_shufflers; rapidsmpf::shuffler::PartID total_num_partitions; rmm::cuda_stream_view stream; @@ -376,12 +361,7 @@ class ConcurrentShuffleTest }; TEST_P(ConcurrentShuffleTest, round_trip) { - ASSERT_NO_FATAL_FAILURE(RunTestTemplate( - [&](auto& shuffler, auto&& packed_chunks) { - shuffler.insert(std::move(packed_chunks)); - }, - [&](auto& shuffler) { shuffler.insert_finished(); } - )); + ASSERT_NO_FATAL_FAILURE(RunTestTemplate()); } // test different `num_shufflers` and `total_num_partitions`. @@ -389,8 +369,12 @@ INSTANTIATE_TEST_SUITE_P( ConcurrentShuffle, ConcurrentShuffleTest, testing::Combine( - testing::ValuesIn({1, 2, 4}), // num_shufflers - testing::ValuesIn({1, 10, 100}) // total_num_partitions + testing::Values(std::size_t{1}, std::size_t{2}, std::size_t{4}), // num_shufflers + testing::Values( // total_num_partitions + rapidsmpf::shuffler::PartID{1}, + rapidsmpf::shuffler::PartID{10}, + rapidsmpf::shuffler::PartID{100} + ) ), [](const testing::TestParamInfo& info) { return "num_shufflers_" + std::to_string(std::get<0>(info.param)) @@ -400,13 +384,11 @@ INSTANTIATE_TEST_SUITE_P( TEST(Shuffler, SpillOnInsertAndExtraction) { rapidsmpf::shuffler::PartID const total_num_partitions = 2; - std::int64_t const seed = 42; - cudf::hash_id const hash_fn = cudf::hash_id::HASH_MURMUR3; - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_default; // Use RapidsMPF's memory resource adaptor so the test can observe per-rank // allocation counts via `get_main_record().num_current_allocs()`. - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; // Control spilling by adjusting the DEVICE memory limit at runtime. // `memory_available(DEVICE)` is computed as `limit - current_allocated()`, so a @@ -433,17 +415,12 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { total_num_partitions, br.get() ); - cudf::table input_table = random_table_with_index(seed, 1000, 0, 10); - auto input_chunks = cudf_streaming::integrations::partition_and_pack( - input_table, - {1}, - total_num_partitions, - hash_fn, - seed, - stream, - br.get(), - rapidsmpf::AllowOverbooking::YES - ); // with overbooking + // Create one non-empty chunk per partition. Each chunk owns a single device + // buffer, so we start with exactly `total_num_partitions` device allocations. + std::unordered_map input_chunks; + for (rapidsmpf::shuffler::PartID pid = 0; pid < total_num_partitions; ++pid) { + input_chunks.emplace(pid, generate_packed_data(1000, 0, stream, *br)); + } // Insert spills does nothing when device memory is available, we start // with 2 device allocations. @@ -675,7 +652,7 @@ TEST_P(ContiguousPartitionAssignmentTest, contiguous) { TEST(Shuffler, ShutdownWhilePaused) { auto progress_thread = GlobalEnvironment->comm_->progress_thread(); - auto mr = cudf::get_current_device_resource_ref(); + auto mr = rmm::mr::get_current_device_resource_ref(); auto br = rapidsmpf::BufferResource::create(mr); @@ -692,27 +669,15 @@ TEST(Shuffler, ShutdownWhilePaused) { EXPECT_NO_THROW(shuffler.shutdown()); } -// check cudf pack conditions for empty table -TEST(EmptyPartitions, cudf_pack) { - auto stream = cudf::get_default_stream(); - cudf::table tbl = random_table_with_index(0, 0, 0, 0); - EXPECT_EQ(0, tbl.num_rows()); - - // following conditions should be met for an empty cudf table - auto packed = cudf::pack(tbl, stream); - EXPECT_TRUE(packed.metadata); - EXPECT_TRUE(packed.gpu_data); - EXPECT_EQ(0, packed.gpu_data->size()); -} - -class ExtractEmptyPartitionsTest : public cudf::test::BaseFixture { +class ExtractEmptyPartitionsTest : public ::testing::Test { public: static constexpr rapidsmpf::shuffler::PartID nparts = 10; static constexpr auto wait_timeout = std::chrono::seconds(30); void SetUp() override { - stream = cudf::get_default_stream(); - br = rapidsmpf::BufferResource::create(mr()); + stream = rmm::cuda_stream_default; + br = + rapidsmpf::BufferResource::create(rmm::mr::get_current_device_resource_ref()); shuffler = std::make_unique( GlobalEnvironment->comm_, 0, nparts, br.get() @@ -811,7 +776,8 @@ TEST_F(ExtractEmptyPartitionsTest, SomeEmptyAndNonEmptyInsertions) { TEST(ShufflerTest, multiple_shutdowns) { auto& comm = GlobalEnvironment->comm_; - auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = + rapidsmpf::BufferResource::create(rmm::mr::get_current_device_resource_ref()); auto shuffler = std::make_unique( comm, 0, comm->nranks(), br.get() ); @@ -835,83 +801,44 @@ TEST(ShufflerTest, multiple_shutdowns) { // Test that multiple threads can call wait() concurrently. TEST(Shuffler, concurrent_wait) { auto const& comm = GlobalEnvironment->comm_; - auto stream = cudf::get_default_stream(); - auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = + rapidsmpf::BufferResource::create(rmm::mr::get_current_device_resource_ref()); // Use more partitions than ranks so each rank owns multiple partitions, ensuring // multiple threads call wait() concurrently on the same shuffler. auto const total_num_partitions = rapidsmpf::safe_cast(comm->nranks()) * 8; constexpr std::size_t total_num_rows = 1000; - constexpr cudf::hash_id hash_fn = cudf::hash_id::HASH_MURMUR3; - constexpr std::int64_t seed = 42; constexpr auto wait_timeout = std::chrono::seconds{30}; rapidsmpf::shuffler::Shuffler shuffler(comm, 0, total_num_partitions, br.get()); - cudf::table full_input = random_table_with_index(seed, total_num_rows, 0, 10); - auto [expected, owner] = cudf_streaming::integrations::partition_and_split( - full_input, - {1}, - static_cast(total_num_partitions), - hash_fn, - seed, - stream, - br.get(), - rapidsmpf::AllowOverbooking::YES - ); - + // Insert each owned input region concurrently, each thread using its own pool stream. { std::vector> insert_futures; - cudf::size_type row_offset = 0; - cudf::size_type part_size = - full_input.num_rows() / static_cast(total_num_partitions); - for (rapidsmpf::shuffler::PartID i = 0; i < total_num_partitions; ++i) { - if (rapidsmpf::shuffler::Shuffler::round_robin(comm, i, total_num_partitions) - == comm->rank()) - { - cudf::size_type row_end = row_offset + part_size; - if (i == total_num_partitions - 1) { - row_end = full_input.num_rows(); - } - auto slice = cudf::slice(full_input, {row_offset, row_end}).at(0); - insert_futures.push_back(std::async(std::launch::async, [&, slice] { - shuffler.insert( - cudf_streaming::integrations::partition_and_pack( - slice, - {1}, - static_cast(total_num_partitions), - hash_fn, - seed, - br->stream_pool().get_stream(), - br.get(), - rapidsmpf::AllowOverbooking::YES - ) - ); - })); - } - row_offset += part_size; + for (rapidsmpf::shuffler::PartID local_pidx : shuffler.local_partitions()) { + insert_futures.push_back(std::async(std::launch::async, [&, local_pidx] { + shuffler.insert(make_partition_data( + total_num_partitions, + total_num_rows, + local_pidx, + br->stream_pool().get_stream(), + *br + )); + })); } std::ranges::for_each(insert_futures, [](auto& f) { f.get(); }); shuffler.insert_finished(); } - auto local_pids = shuffler.local_partitions(); + // Wait + extract + validate each local partition concurrently, so multiple threads + // call wait() on the same shuffler at once. std::vector> futures; - for (auto pid : local_pids) { - futures.push_back(std::async(std::launch::async, [&, pid] { + for (auto j : shuffler.local_partitions()) { + futures.push_back(std::async(std::launch::async, [&, j] { EXPECT_NO_THROW(shuffler.wait(wait_timeout)); - auto chunks = shuffler.extract(pid); - auto result = cudf_streaming::integrations::unpack_and_concat( - rapidsmpf::unspill_partitions( - std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES - ), - stream, - br.get(), - rapidsmpf::AllowOverbooking::YES - ); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT( - sort_table(result), sort_table(expected[pid]) + validate_partition_data( + shuffler.extract(j), total_num_partitions, total_num_rows, j, *br ); })); } @@ -931,11 +858,10 @@ TEST(Shuffler, opid_reuse) { GTEST_SKIP() << "OpID reuse test requires multiple ranks"; } - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_default; auto const total_num_partitions = rapidsmpf::safe_cast(comm->nranks()); constexpr std::size_t total_num_rows = 1000; - constexpr cudf::hash_id hash_fn = cudf::hash_id::HASH_MURMUR3; constexpr rapidsmpf::OpID op_id = 0; constexpr auto wait_timeout = std::chrono::seconds{30}; @@ -953,61 +879,21 @@ TEST(Shuffler, opid_reuse) { shuffler_br = delayed_br.get(); } - auto insert_data = [&](rapidsmpf::shuffler::Shuffler& shuffler, std::int64_t seed) { - cudf::table full_input = random_table_with_index(seed, total_num_rows, 0, 10); - cudf::size_type row_offset = 0; - cudf::size_type part_size = - full_input.num_rows() / static_cast(total_num_partitions); - for (rapidsmpf::shuffler::PartID i = 0; i < total_num_partitions; ++i) { - if (rapidsmpf::shuffler::Shuffler::round_robin(comm, i, total_num_partitions) - == comm->rank()) - { - cudf::size_type row_end = row_offset + part_size; - if (i == total_num_partitions - 1) { - row_end = full_input.num_rows(); - } - auto slice = cudf::slice(full_input, {row_offset, row_end}).at(0); - auto packed = cudf_streaming::integrations::partition_and_pack( - slice, - {1}, - static_cast(total_num_partitions), - hash_fn, - seed, - stream, - br.get(), - rapidsmpf::AllowOverbooking::YES - ); - shuffler.insert(std::move(packed)); - } - row_offset += part_size; + // Each shuffle uses a distinct base offset (in place of a seed) so the two shuffles + // carry different data; a cross-matched message would therefore fail validation. + auto insert_data = [&](rapidsmpf::shuffler::Shuffler& shuffler, std::int64_t base) { + for (rapidsmpf::shuffler::PartID local_pidx : shuffler.local_partitions()) { + shuffler.insert(make_partition_data( + total_num_partitions, total_num_rows, local_pidx, stream, *br, base + )); } }; auto validate_results = [&](rapidsmpf::shuffler::Shuffler& shuffler, - std::int64_t seed) { - cudf::table full_input = random_table_with_index(seed, total_num_rows, 0, 10); - auto [expected, owner] = cudf_streaming::integrations::partition_and_split( - full_input, - {1}, - static_cast(total_num_partitions), - hash_fn, - seed, - stream, - br.get(), - rapidsmpf::AllowOverbooking::YES - ); - for (auto pid : shuffler.local_partitions()) { - auto chunks = shuffler.extract(pid); - auto result = cudf_streaming::integrations::unpack_and_concat( - rapidsmpf::unspill_partitions( - std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES - ), - stream, - br.get(), - rapidsmpf::AllowOverbooking::YES - ); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT( - sort_table(result), sort_table(expected[pid]) + std::int64_t base) { + for (auto j : shuffler.local_partitions()) { + validate_partition_data( + shuffler.extract(j), total_num_partitions, total_num_rows, j, *br, base ); } }; @@ -1039,10 +925,9 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { GTEST_SKIP() << "OpID reuse test requires multiple ranks"; } - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_default; constexpr rapidsmpf::shuffler::PartID total_num_partitions = 1; constexpr std::size_t total_num_rows = 1000; - constexpr cudf::hash_id hash_fn = cudf::hash_id::HASH_MURMUR3; constexpr rapidsmpf::OpID op_id = 0; constexpr auto wait_timeout = std::chrono::seconds{30}; @@ -1060,51 +945,23 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { shuffler_br = delayed_br.get(); } - auto insert_data = [&](rapidsmpf::shuffler::Shuffler& shuffler, std::int64_t seed) { - cudf::table full_input = random_table_with_index(seed, total_num_rows, 0, 10); - // With total_num_partitions=1, only rank 0 owns the single partition. - if (rapidsmpf::shuffler::Shuffler::round_robin(comm, 0, total_num_partitions) - == comm->rank()) - { - auto packed = cudf_streaming::integrations::partition_and_pack( - full_input, - {1}, - static_cast(total_num_partitions), - hash_fn, - seed, - stream, - br.get(), - rapidsmpf::AllowOverbooking::YES - ); - shuffler.insert(std::move(packed)); + // Each shuffle uses a distinct base offset (in place of a seed) so the two shuffles + // carry different data; a cross-matched message would therefore fail validation. + // With total_num_partitions=1, only rank 0 owns the single partition; all other ranks + // have empty local_partitions(), so they insert/validate nothing. + auto insert_data = [&](rapidsmpf::shuffler::Shuffler& shuffler, std::int64_t base) { + for (rapidsmpf::shuffler::PartID local_pidx : shuffler.local_partitions()) { + shuffler.insert(make_partition_data( + total_num_partitions, total_num_rows, local_pidx, stream, *br, base + )); } }; auto validate_results = [&](rapidsmpf::shuffler::Shuffler& shuffler, - std::int64_t seed) { - cudf::table full_input = random_table_with_index(seed, total_num_rows, 0, 10); - auto [expected, owner] = cudf_streaming::integrations::partition_and_split( - full_input, - {1}, - static_cast(total_num_partitions), - hash_fn, - seed, - stream, - br.get(), - rapidsmpf::AllowOverbooking::YES - ); - for (auto pid : shuffler.local_partitions()) { - auto chunks = shuffler.extract(pid); - auto result = cudf_streaming::integrations::unpack_and_concat( - rapidsmpf::unspill_partitions( - std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES - ), - stream, - br.get(), - rapidsmpf::AllowOverbooking::YES - ); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT( - sort_table(result), sort_table(expected[pid]) + std::int64_t base) { + for (auto j : shuffler.local_partitions()) { + validate_partition_data( + shuffler.extract(j), total_num_partitions, total_num_rows, j, *br, base ); } }; diff --git a/cpp/tests/utils.hpp b/cpp/tests/utils.hpp index bc51e460b..94eace6ac 100644 --- a/cpp/tests/utils.hpp +++ b/cpp/tests/utils.hpp @@ -36,6 +36,7 @@ #include #include +#include #include /** @@ -192,58 +193,70 @@ template /** * @brief Generate a packed data object with the given number of elements and offset. * - * Both metadata and GPU data contain the same integer sequence. + * Both metadata and GPU data contain the same integer sequence of type T. * + * @tparam T Element type stored in the buffer (default: int). * @param n_elements Number of elements in the sequence. * @param offset Starting value of the sequence. * @param stream CUDA stream for device allocation. * @param br Buffer resource used for allocations. * @return A packed data object containing metadata and GPU data. */ +template [[nodiscard]] inline rapidsmpf::PackedData generate_packed_data( - int n_elements, - int offset, + std::size_t n_elements, + T offset, rmm::cuda_stream_view stream, - rapidsmpf::BufferResource& br + rapidsmpf::BufferResource& br, + rapidsmpf::AllowOverbooking allow_overbooking = rapidsmpf::AllowOverbooking::YES ) { - auto values = iota_vector(n_elements, offset); - - auto metadata = std::make_unique>(n_elements * sizeof(int)); - std::memcpy(metadata->data(), values.data(), n_elements * sizeof(int)); + auto const values = iota_vector(n_elements, offset); + auto const* bytes = reinterpret_cast(values.data()); - auto data = std::make_unique( - values.data(), n_elements * sizeof(int), stream, br.device_mr() + auto metadata = std::make_unique>( + bytes, bytes + values.size() * sizeof(T) ); + auto [reservation, _] = + br.reserve(rapidsmpf::MemoryType::DEVICE, metadata->size(), allow_overbooking); + auto data = br.make_buffer(stream, std::move(reservation)); + + data->write_access([d_ptr = metadata->data(), m_size = metadata->size()]( + std::byte* ptr, rmm::cuda_stream_view op_stream + ) { + RAPIDSMPF_CUDA_TRY(rapidsmpf::cuda_memcpy_async(ptr, d_ptr, m_size, op_stream)); + }); - return {std::move(metadata), br.move(std::move(data), stream)}; + return {std::move(metadata), std::move(data)}; } /** * @brief Validate a packed data object by checking metadata and GPU data contents. * + * @tparam T Element type stored in the buffer (default: int). * @param packed_data Packed data object to validate. * @param n_elements Expected number of elements. * @param offset Expected starting value of the sequence. * @param stream CUDA stream used for device-host transfers. * @param br Buffer resource used for host allocation. */ +template inline void validate_packed_data( rapidsmpf::PackedData&& packed_data, - int n_elements, - int offset, + std::size_t n_elements, + T offset, rmm::cuda_stream_view stream, rapidsmpf::BufferResource& br ) { auto const& metadata = *packed_data.metadata; - EXPECT_EQ(n_elements * sizeof(int), metadata.size()); + EXPECT_EQ(n_elements * sizeof(T), metadata.size()); - for (int i = 0; i < n_elements; i++) { - int val; - std::memcpy(&val, metadata.data() + i * sizeof(int), sizeof(int)); - EXPECT_EQ(offset + i, val); + for (std::size_t i = 0; i < n_elements; i++) { + T val; + std::memcpy(&val, metadata.data() + i * sizeof(T), sizeof(T)); + EXPECT_EQ(offset + static_cast(i), val); } - EXPECT_EQ(n_elements * sizeof(int), packed_data.data->size); + EXPECT_EQ(n_elements * sizeof(T), packed_data.data->size); auto res = br.reserve_or_fail(packed_data.data->size, rapidsmpf::MemoryType::HOST); auto data_on_host = br.move_to_host_buffer(std::move(packed_data.data), res); From 85e34437b6dc2101c56a686b79d7afdda267d213 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Tue, 9 Jun 2026 13:38:23 -0700 Subject: [PATCH 02/14] cudf-free streaming shuffler Signed-off-by: niranda perera --- cpp/tests/CMakeLists.txt | 3 +- cpp/tests/streaming/test_shuffler.cpp | 224 +++++++++++++------------- cpp/tests/test_shuffler.cpp | 8 +- 3 files changed, 121 insertions(+), 114 deletions(-) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index b90449df2..571d529cb 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -110,6 +110,7 @@ if(RAPIDSMPF_HAVE_STREAMING) streaming/test_message.cpp streaming/test_sparse_alltoall.cpp streaming/test_spillable_messages.cpp + streaming/test_shuffler.cpp ) endif() @@ -122,7 +123,7 @@ if(BUILD_CUDF_TESTS) target_compile_definitions(test_sources PRIVATE RAPIDSMPF_HAVE_CUDF) if(RAPIDSMPF_HAVE_STREAMING) - target_sources(test_sources PRIVATE streaming/test_leaf_actor.cpp streaming/test_shuffler.cpp) + target_sources(test_sources PRIVATE streaming/test_leaf_actor.cpp) endif() endif() diff --git a/cpp/tests/streaming/test_shuffler.cpp b/cpp/tests/streaming/test_shuffler.cpp index 7bfa893a5..5dab73b1b 100644 --- a/cpp/tests/streaming/test_shuffler.cpp +++ b/cpp/tests/streaming/test_shuffler.cpp @@ -3,18 +3,16 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include +#include + #include #include -#include -#include -#include -#include -#include - #include #include #include +#include #include #include #include @@ -55,13 +53,12 @@ TEST_F(BaseStreamingShuffle, zero_owned_partitions_completes) { class StreamingShuffler : public BaseStreamingShuffle, public ::testing::WithParamInterface { public: - const unsigned int num_partitions = 10; - const unsigned int num_rows = 1000; - const unsigned int num_chunks = 5; - const unsigned int chunk_size = num_rows / num_chunks; - const std::int64_t seed = 42; - const cudf::hash_id hash_function = cudf::hash_id::HASH_MURMUR3; - const OpID op_id = 0; + static constexpr size_t num_partitions = 10; + static constexpr size_t num_rows = 1000; + static constexpr size_t num_chunks = 5; + static constexpr OpID op_id = 0; + + shuffler::Shuffler::PartitionOwner partition_owner = shuffler::Shuffler::round_robin; void SetUp() override { BaseStreamingShuffle::SetUpWithThreads(GetParam()); @@ -70,100 +67,6 @@ class StreamingShuffler : public BaseStreamingShuffle, void TearDown() override { BaseStreamingShuffle::TearDown(); } - - void run_test(auto make_shuffler_actor_fn) { - // Create the full input table and slice it into chunks. - cudf::table full_input_table = random_table_with_index(seed, num_rows, 0, 10); - std::vector input_chunks; - for (unsigned int i = 0; i < num_chunks; ++i) { - input_chunks.emplace_back( - cudf_streaming::streaming::to_message( - i, - std::make_unique( - std::make_unique( - cudf::slice( - full_input_table, - {static_cast(i * chunk_size), - static_cast((i + 1) * chunk_size)}, - stream - ) - .at(0), - stream, - ctx->br()->device_mr() - ), - stream - ) - ) - ); - } - - // Create and run the streaming pipeline. - std::vector output_chunks; - { - std::vector actors; - auto ch1 = ctx->create_channel(); - actors.push_back(actor::push_to_channel(ctx, ch1, std::move(input_chunks))); - - auto ch2 = ctx->create_channel(); - actors.push_back( - cudf_streaming::streaming::actor::partition_and_pack( - ctx, ch1, ch2, {1}, num_partitions, hash_function, seed - ) - ); - - auto ch3 = ctx->create_channel(); - actors.emplace_back(make_shuffler_actor_fn(ch2, ch3)); - - auto ch4 = ctx->create_channel(); - actors.push_back( - cudf_streaming::streaming::actor::unpack_and_concat(ctx, ch3, ch4) - ); - - actors.push_back(actor::pull_from_channel(ctx, ch4, output_chunks)); - - run_actor_network(std::move(actors)); - } - - auto comm = GlobalEnvironment->comm_; - std::unique_ptr expected_table; - if (comm->nranks() == 1) { // full_input table is expected - expected_table = std::make_unique(std::move(full_input_table)); - } else { // full_input table is replicated on all ranks - // local partitions - auto [table, offsets] = cudf::hash_partition( - full_input_table.view(), {1}, num_partitions, hash_function, seed - ); - - auto local_pids = shuffler::Shuffler::local_partitions( - comm, num_partitions, shuffler::Shuffler::round_robin - ); - - // every partition is replicated on all ranks - std::vector expected_tables; - for (auto pid : local_pids) { - auto t_view = - cudf::slice(table->view(), {offsets[pid], offsets[pid + 1]}).at(0); - // this will be replicated on all ranks - for (rapidsmpf::Rank rank = 0; rank < comm->nranks(); ++rank) { - expected_tables.push_back(t_view); - } - } - expected_table = cudf::concatenate(expected_tables); - } - - // Concat all output chunks to a single table. - std::vector output_chunks_as_views; - for (auto& chunk : output_chunks) { - output_chunks_as_views.push_back( - chunk.get().table_view() - ); - } - auto result_table = cudf::concatenate(output_chunks_as_views); - - CUDF_TEST_EXPECT_TABLES_EQUIVALENT( - sort_table(result_table->view()), sort_table(expected_table->view()) - ); - } }; INSTANTIATE_TEST_SUITE_P( @@ -175,12 +78,111 @@ INSTANTIATE_TEST_SUITE_P( } ); +// Verifies end-to-end correctness of the streaming shuffler actor. +// +// Each rank sends num_chunks messages, where each message is a PartitionMapChunk covering +// all num_partitions partitions. The data for partition j in chunk c_idx on rank r is a +// contiguous integer sequence whose values encode both the rank and the row range, making +// them globally unique and independently verifiable. +// +// After shuffling, each local partition should have received exactly num_chunks * nranks +// packed-data items (one per (rank, chunk) pair). The items are sorted by their first +// element — which equals rank * num_rows + row_start — to reconstruct rank-major, +// chunk-minor order, and then validated against the expected row range from `pieces`. TEST_P(StreamingShuffler, basic_shuffler) { - EXPECT_NO_FATAL_FAILURE(run_test([&](auto ch_in, auto ch_out) -> Actor { - return actor::shuffler( - ctx, GlobalEnvironment->comm_, ch_in, ch_out, op_id, num_partitions + auto comm = GlobalEnvironment->comm_; + // split a span [0, num_rows) into num_partitions * num_chunks contiguous pieces. + auto const pieces = rapidsmpf::chunk_indices(num_rows, num_partitions * num_chunks); + + // each rank contributes num_chunks messages, each covering num_partitions pieces. + const int64_t base = static_cast(comm->rank()) * num_rows; + std::vector input_chunks; // Message contains a PartitionMapChunk + for (size_t chunk_idx = 0; chunk_idx < num_chunks; ++chunk_idx) { + ContentDescription cd{}; + std::unordered_map chunks; + chunks.reserve(num_partitions); + for (size_t j = 0; j < num_partitions; ++j) { + auto [start, end] = pieces[chunk_idx * num_partitions + j]; + // end > start is guaranteed. + auto [it, _] = chunks.emplace( + static_cast(j), + generate_packed_data( + end - start, base + static_cast(start), stream, *br + ) + ); + cd.content_size(it->second.data->mem_type()) += it->second.data->size; + } + input_chunks.emplace_back(Message( + chunk_idx, + std::make_unique(std::move(chunks)), + std::move(cd) + )); + } + EXPECT_EQ(input_chunks.size(), num_chunks); + + // Create and run the streaming pipeline. + std::vector output_chunks; + { + std::vector actors; + auto ch1 = ctx->create_channel(); + actors.push_back(actor::push_to_channel(ctx, ch1, std::move(input_chunks))); + + auto ch2 = ctx->create_channel(); + actors.emplace_back( + actor::shuffler(ctx, comm, ch1, ch2, op_id, num_partitions, partition_owner) ); - })); + + actors.push_back(actor::pull_from_channel(ctx, ch2, output_chunks)); + + run_actor_network(std::move(actors)); + } + + auto local_pids = + shuffler::Shuffler::local_partitions(comm, num_partitions, partition_owner); + + // Since all partitions are non-empty, each local partition ID should a corresponding + // output chunk. + EXPECT_EQ(local_pids.size(), output_chunks.size()); + const size_t n_ranks = static_cast(comm->nranks()); + for (auto& chunk : output_chunks) { + auto pid = chunk.sequence_number(); + std::erase_if(local_pids, [pid](auto& p) { return p == pid; }); + + auto p_vec = chunk.release(); + // for each local pid, it should receive num_chunks * nranks chunks. + EXPECT_EQ(p_vec.data.size(), num_chunks * n_ranks); + + // since values are offset by rank, if we sort packed data by their first element, + // then it will be in rank & chunk-index order. + std::ranges::sort(p_vec.data, [](auto& a, auto& b) { + std::int64_t oa{}, ob{}; + std::memcpy(&oa, a.metadata->data(), sizeof(std::int64_t)); + std::memcpy(&ob, b.metadata->data(), sizeof(std::int64_t)); + return oa < ob; + }); + + // p_vec.data is sorted by first element, so entries are ordered + // (rank=0,chunk=0), (rank=0,chunk=1), ..., (rank=1,chunk=0), ... + // i.e. flat index r*num_chunks + c_idx. + for (size_t r = 0; r < n_ranks; ++r) { + auto r_base = static_cast(r) * num_rows; // rank-base offset + for (size_t c_idx = 0; c_idx < num_chunks; ++c_idx) { + const auto [start, end] = pieces[c_idx * num_partitions + pid]; + SCOPED_TRACE( + "pid=" + std::to_string(pid) + ", rank=" + std::to_string(r) + + ", chunk_idx=" + std::to_string(c_idx) + ); + validate_packed_data( + std::move(p_vec.data[r * num_chunks + c_idx]), + end - start, + r_base + static_cast(start), + stream, + *br + ); + } + } + } + EXPECT_TRUE(local_pids.empty()); } class ShufflerAsyncTest diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index 4c5da8cdc..36e2c5fc1 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -235,9 +235,13 @@ void test_shuffler( EXPECT_NO_THROW(shuffler.wait(wait_timeout)); - for (auto j : shuffler.local_partitions()) { + for (auto local_pidx : shuffler.local_partitions()) { validate_partition_data( - shuffler.extract(j), total_num_partitions, total_num_rows, j, *br + shuffler.extract(local_pidx), + total_num_partitions, + total_num_rows, + local_pidx, + *br ); } } From 5f066f78c7aee718c5c14e1375d0d17b93a08b3d Mon Sep 17 00:00:00 2001 From: niranda perera Date: Tue, 9 Jun 2026 16:23:03 -0700 Subject: [PATCH 03/14] python tests Signed-off-by: niranda perera --- .../rapidsmpf/memory/packed_data.pyi | 11 + .../rapidsmpf/memory/packed_data.pyx | 70 ++++ python/rapidsmpf/rapidsmpf/testing.py | 185 ++++++++++ .../tests/streaming/test_shuffler.py | 329 ++++-------------- .../rapidsmpf/tests/test_shuffler.py | 175 ++-------- 5 files changed, 353 insertions(+), 417 deletions(-) diff --git a/python/rapidsmpf/rapidsmpf/memory/packed_data.pyi b/python/rapidsmpf/rapidsmpf/memory/packed_data.pyi index f044f4d48..d1beddfd2 100644 --- a/python/rapidsmpf/rapidsmpf/memory/packed_data.pyi +++ b/python/rapidsmpf/rapidsmpf/memory/packed_data.pyi @@ -1,5 +1,8 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 +from rmm.pylibrmm.device_buffer import DeviceBuffer +from rmm.pylibrmm.stream import Stream + from rapidsmpf.memory.buffer_resource import BufferResource class PackedData: @@ -8,4 +11,12 @@ class PackedData: def from_host_bytes( cls, data: bytes | bytearray, br: BufferResource ) -> PackedData: ... + @classmethod + def from_device_buffer( + cls, + gpu_data: DeviceBuffer, + metadata: bytes | bytearray, + stream: Stream, + br: BufferResource, + ) -> PackedData: ... def to_host_bytes(self) -> bytes: ... diff --git a/python/rapidsmpf/rapidsmpf/memory/packed_data.pyx b/python/rapidsmpf/rapidsmpf/memory/packed_data.pyx index 124faa089..d4142d765 100644 --- a/python/rapidsmpf/rapidsmpf/memory/packed_data.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/packed_data.pyx @@ -7,6 +7,8 @@ from libcpp.utility cimport move from libcpp.vector cimport vector from rmm.librmm.cuda_stream_view cimport cuda_stream_view from rmm.librmm.device_buffer cimport device_buffer +from rmm.pylibrmm.device_buffer cimport DeviceBuffer +from rmm.pylibrmm.stream cimport Stream from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.memory.buffer_resource cimport (BufferResource, @@ -57,6 +59,21 @@ cdef extern from *: ); } + std::unique_ptr cpp_packed_data_from_device_buffer( + const std::uint8_t* metadata, + std::size_t metadata_size, + std::unique_ptr gpu_data, + rmm::cuda_stream_view stream, + rapidsmpf::BufferResource* br + ) { + auto meta = std::make_unique>( + metadata, metadata + metadata_size + ); + return std::make_unique( + std::move(meta), br->move(std::move(gpu_data), stream) + ); + } + std::vector cpp_packed_data_to_host_bytes( rapidsmpf::PackedData* pd ) { @@ -86,6 +103,14 @@ cdef extern from *: cpp_BufferResource* br, ) except + nogil + unique_ptr[cpp_PackedData] cpp_packed_data_from_device_buffer( + const uint8_t* metadata, + size_t metadata_size, + unique_ptr[device_buffer] gpu_data, + cuda_stream_view stream, + cpp_BufferResource* br, + ) except +ex_handler nogil + vector[uint8_t] cpp_packed_data_to_host_bytes( cpp_PackedData* pd, ) except + nogil @@ -141,6 +166,51 @@ cdef class PackedData: ret._br = br return ret + @classmethod + def from_device_buffer( + cls, + DeviceBuffer gpu_data not None, + const uint8_t[::1] metadata not None, + Stream stream not None, + BufferResource br not None, + ): + """ + Construct a PackedData from an rmm device buffer and host metadata. + + Takes ownership of ``gpu_data``; the input buffer is left empty after this + call. The metadata bytes are copied into a host-side metadata buffer. + + Parameters + ---------- + gpu_data + Device buffer holding the data payload. Consumed by this call. + metadata + Contiguous buffer of host bytes (bytes, bytearray, or buffer-protocol + object). Must be non-empty. + stream + CUDA stream used to take ownership of the device buffer. + br + Buffer resource for memory management. + + Returns + ------- + A new PackedData instance owning the device buffer. + """ + cdef cpp_BufferResource* _br = br.ptr() + cdef size_t meta_size = len(metadata) + cdef const uint8_t* meta_ptr = NULL + if meta_size > 0: + meta_ptr = &metadata[0] + cdef cuda_stream_view sv = stream.view() + cdef unique_ptr[device_buffer] gpu = move(gpu_data.c_obj) + cdef PackedData ret = cls.__new__(cls) + with nogil: + ret.c_obj = cpp_packed_data_from_device_buffer( + meta_ptr, meta_size, move(gpu), sv, _br + ) + ret._br = br + return ret + def to_host_bytes(self) -> bytes: """ Extract the host bytes from this PackedData. diff --git a/python/rapidsmpf/rapidsmpf/testing.py b/python/rapidsmpf/rapidsmpf/testing.py index 60207f286..a4dbb145f 100644 --- a/python/rapidsmpf/rapidsmpf/testing.py +++ b/python/rapidsmpf/rapidsmpf/testing.py @@ -6,13 +6,23 @@ from typing import TYPE_CHECKING +import numpy as np import pylibcudf +import rmm from rmm.pylibrmm.stream import DEFAULT_STREAM +from rapidsmpf.memory.packed_data import PackedData + if TYPE_CHECKING: from rmm.pylibrmm.stream import Stream + from rapidsmpf.memory.buffer_resource import BufferResource + +# Element type stored in the synthetic packed payloads. Values are wide enough +# that any per-shuffle ``base`` offsets used by callers never collide. +_DTYPE = np.int64 + def assert_eq( left: pylibcudf.Table, @@ -63,3 +73,178 @@ def assert_eq( ) if not pylibcudf.table_equality.tables_equal(left, right, stream=stream): raise AssertionError(f"Table are not equal with {sort_rows=}") + + +def chunk_indices(count: int, num_chunks: int) -> list[tuple[int, int]]: + """ + Split ``[0, count)`` into ``num_chunks`` contiguous, front-loaded pieces. + + Mirrors ``rapidsmpf::chunk_indices``: when ``count < num_chunks`` the trailing + pieces are empty (``start == end``). The returned pieces exactly tile + ``[0, count)``. + + Parameters + ---------- + count + Size of the range to split. + num_chunks + Number of pieces to split the range into. + + Returns + ------- + A list of ``(start, end)`` pairs, one per chunk. + """ + chunk_size = -(-count // num_chunks) # ceil division + return [ + (min(k * chunk_size, count), min((k + 1) * chunk_size, count)) + for k in range(num_chunks) + ] + + +def generate_packed_data( + n_elements: int, offset: int, stream: Stream, br: BufferResource +) -> PackedData: + """ + Build a ``PackedData`` holding the int sequence ``[offset, offset + n_elements)``. + + The sequence is stored as a device buffer payload (and mirrored in the metadata) + so it survives a shuffle round-trip and can be validated by + :func:`validate_packed_data`. + + Parameters + ---------- + n_elements + Number of elements in the sequence. + offset + Starting value of the sequence. + stream + CUDA stream used for the device allocation. + br + Buffer resource used for memory management. + + Returns + ------- + A ``PackedData`` containing the integer sequence. + """ + data = np.arange(offset, offset + n_elements, dtype=_DTYPE).tobytes() + gpu_data = rmm.DeviceBuffer.to_device(data, stream=stream) + return PackedData.from_device_buffer(gpu_data, data, stream, br) + + +def validate_packed_data(packed_data: PackedData, n_elements: int, offset: int) -> None: + """ + Check that ``packed_data`` holds the sequence ``[offset, offset + n_elements)``. + + Parameters + ---------- + packed_data + Packed data to validate. + n_elements + Expected number of elements. + offset + Expected starting value of the sequence. + """ + values = np.frombuffer(packed_data.to_host_bytes(), dtype=_DTYPE) + np.testing.assert_array_equal( + values, np.arange(offset, offset + n_elements, dtype=_DTYPE) + ) + + +def make_partition_data( + total_num_partitions: int, + total_num_rows: int, + local_pid: int, + stream: Stream, + br: BufferResource, + base: int = 0, +) -> dict[int, PackedData]: + """ + Produce the non-empty sub-regions of one owned input region ``local_pid``. + + The index range ``[0, total_num_rows)`` is split into ``P * P`` contiguous + sub-regions (``P == total_num_partitions``). Sub-region ``(local_pid, + split_idx)`` is piece ``local_pid * P + split_idx`` and is routed to + destination partition ``split_idx``. Since ``local_partitions()`` across ranks + partition ``[0, P)``, every input region is produced exactly once, so rows are + not replicated and the total shuffled data equals ``total_num_rows``. + + Parameters + ---------- + total_num_partitions + Total number of partitions in the shuffle. + total_num_rows + Total number of rows tiled across all input regions. + local_pid + Index of the owned input region to produce. + stream + CUDA stream used for device allocations. + br + Buffer resource used for memory management. + base + Offset added to every value so distinct shuffles carry distinct data. + + Returns + ------- + A map of destination partition ID to its packed sub-region. + """ + P = total_num_partitions + pieces = chunk_indices(total_num_rows, P * P) + + chunks: dict[int, PackedData] = {} + for split_idx in range(P): + start, end = pieces[local_pid * P + split_idx] + if end > start: + chunks[split_idx] = generate_packed_data( + end - start, base + start, stream, br + ) + return chunks + + +def validate_partition_data( + received: list[PackedData], + total_num_partitions: int, + total_num_rows: int, + local_pid: int, + base: int = 0, +) -> None: + """ + Verify that the ``received`` chunks for partition ``local_pid`` are as expected. + + Checks the chunks against the non-empty sub-regions expected for partition + ``local_pid`` under the conserved, front-loaded data model of + :func:`make_partition_data`. + + Parameters + ---------- + received + The packed data chunks extracted for partition ``local_pid``. + total_num_partitions + Total number of partitions in the shuffle. + total_num_rows + Total number of rows tiled across all input regions. + local_pid + Partition ID being validated. + base + Offset that was added to every value when the data was generated. + """ + P = total_num_partitions + pieces = chunk_indices(total_num_rows, P * P) + + # Recompute the non-empty (offset, count) sub-regions expected for partition + # local_pid, in increasing input-region-index (== increasing offset) order. + expected: list[tuple[int, int]] = [] + for i in range(P): + start, end = pieces[i * P + local_pid] + if end > start: + expected.append((base + start, end - start)) + + assert len(received) == len(expected) + + # Decode each received chunk and sort by its first element (== offset) so they + # align 1:1 with the expected list, which is already in offset order. + decoded = sorted( + (np.frombuffer(pd.to_host_bytes(), dtype=_DTYPE) for pd in received), + key=lambda arr: int(arr[0]), + ) + for arr, (off, cnt) in zip(decoded, expected, strict=True): + np.testing.assert_array_equal(arr, np.arange(off, off + cnt, dtype=_DTYPE)) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py index 6f6fd6a9b..5f4e364c6 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py @@ -3,290 +3,95 @@ from __future__ import annotations +import asyncio from typing import TYPE_CHECKING -import cupy as cp -import numpy as np -import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming") -from cudf_streaming.integrations.partition import split_and_pack, unpack_and_concat -from cudf_streaming.streaming.partition import ( - partition_and_pack, - unpack_and_concat as streaming_unpack_and_concat, +from rapidsmpf.streaming.coll.shuffler import ShufflerAsync +from rapidsmpf.testing import ( + generate_packed_data, + make_partition_data, + validate_partition_data, ) -from cudf_streaming.streaming.table_chunk import TableChunk - -from rapidsmpf.shuffler import PartitionAssignment -from rapidsmpf.streaming.coll.shuffler import ( - ShufflerAsync, - shuffler, -) -from rapidsmpf.streaming.core.actor import define_actor, run_actor_network -from rapidsmpf.streaming.core.leaf_actor import pull_from_channel, push_to_channel -from rapidsmpf.streaming.core.message import Message -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") if TYPE_CHECKING: - from collections.abc import Awaitable - from rmm.pylibrmm.stream import Stream from rapidsmpf.communicator.communicator import Communicator - from rapidsmpf.streaming.chunks.partition import ( - PartitionMapChunk, - PartitionVectorChunk, - ) - from rapidsmpf.streaming.core.actor import CppActor - from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context -@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) -def test_single_rank_shuffler( - context: Context, comm: Communicator, stream: Stream, num_partitions: int -) -> None: - if comm.nranks != 1: - pytest.skip("Only support single-rank runs") - - num_rows = 1000 - num_chunks = 5 - chunk_size = num_rows // num_chunks - op_id = 0 - # We start a full dataframe. - df = plc.Table( - [ - plc.Column.from_array(cp.arange(num_rows, dtype=cp.int32)), - plc.Column.from_array( - cp.random.randint(0, 10, size=num_rows, dtype=cp.int32) - ), - ] - ) - - # That we slice into chunks and wrap as TableChunk (sequence_number=i). - input_chunks: list[Message[TableChunk]] = [] - for i in range(num_chunks): - lo = i * chunk_size - hi = (i + 1) * chunk_size - df_chunk = plc.copying.slice(df, [lo, hi])[0] - chunk = TableChunk.from_pylibcudf_table( - table=df_chunk, - stream=stream, - exclusive_view=False, - br=context.br(), - ) - input_chunks.append(Message(i, chunk)) - - # Build the streaming pipeline: - # push -> partition/pack -> shuffle -> unpack/concat -> pull. - actors = [] - - ch1: Channel[TableChunk] = context.create_channel() - actors.append(push_to_channel(context, ch1, input_chunks)) - - ch2: Channel[PartitionMapChunk] = context.create_channel() - actors.append( - partition_and_pack( - context, - ch_in=ch1, - ch_out=ch2, - columns_to_hash=(1,), - num_partitions=num_partitions, - ) - ) - - ch3: Channel[PartitionVectorChunk] = context.create_channel() - actors.append( - shuffler( - context, - comm, - ch_in=ch2, - ch_out=ch3, - op_id=op_id, - total_num_partitions=num_partitions, - ) - ) - - ch4: Channel[TableChunk] = context.create_channel() - actors.append(streaming_unpack_and_concat(context, ch_in=ch3, ch_out=ch4)) - - pull_actor, out_messages = pull_from_channel(context, ch_in=ch4) - actors.append(pull_actor) - - run_actor_network(context, actors=actors) - - output_chunks = [ - TableChunk.from_message(msg, br=context.br()) for msg in out_messages.release() - ] - - result = plc.concatenate.concatenate( - [chunk.table_view() for chunk in output_chunks] - ) - assert_eq(result, df, sort_rows=0) - - -@define_actor() -async def generate_inputs( - context: Context, ch: Channel[TableChunk], num_rows: int, num_chunks: int -) -> None: - for i in range(num_chunks): - stream = context.get_stream_from_pool() - table = plc.Table( - [ - plc.Column.from_array( - np.arange(num_rows, dtype=np.int32) + i * num_rows, stream=stream - ) - ] - ) - msg = Message( - i, - TableChunk.from_pylibcudf_table( - table, stream, exclusive_view=True, br=context.br() - ), - ) - await ch.send(context, msg) - await ch.drain(context) - - -@define_actor() -async def do_shuffle( - context: Context, - comm: Communicator, - ch_in: Channel[TableChunk], - ch_out: Channel[TableChunk], - op_id: int, - num_partitions: int, - *, - partition_assignment: PartitionAssignment = PartitionAssignment.ROUND_ROBIN, -) -> None: - shuffle = ShufflerAsync( - context, comm, op_id, num_partitions, partition_assignment=partition_assignment - ) - while (msg := await ch_in.recv(context)) is not None: - chunk = TableChunk.from_message(msg, br=context.br()) - num_rows = chunk.table_view().num_rows() - part_size = num_rows // num_partitions + (num_rows % num_partitions) - splits = range(part_size, num_rows, part_size) - shuffle.insert( - split_and_pack(chunk.table_view(), splits, chunk.stream, context.br()) - ) - await shuffle.insert_finished(context) - for pid in shuffle.local_partitions(): - data = shuffle.extract(pid) - stream = context.get_stream_from_pool() - unpacked = TableChunk.from_pylibcudf_table( - unpack_and_concat(data, stream, context.br()), - stream, - exclusive_view=True, - br=context.br(), - ) - await ch_out.send(context, Message(pid, unpacked)) - await ch_out.drain(context) - - -@pytest.mark.parametrize("num_partitions", [4, 8]) -def test_shuffler_runtime_obeys_contiguous_assignment( +@pytest.mark.parametrize("total_num_partitions", [1, 2, 5, 10]) +@pytest.mark.parametrize("total_num_rows", [1, 100, 1000]) +def test_shuffler_round_trip( context: Context, comm: Communicator, - num_partitions: int, + stream: Stream, + total_num_partitions: int, + total_num_rows: int, ) -> None: - if comm.nranks != 1: - pytest.skip("Only support single-rank runs") - - actors: list[CppActor | Awaitable[None]] = [] - - num_rows = 200 - num_chunks = 3 - op_id = 0 - ch_in: Channel[TableChunk] = context.create_channel() - actors.append(generate_inputs(context, ch_in, num_rows, num_chunks)) - ch_shuffled: Channel[TableChunk] = context.create_channel() - actors.append( - do_shuffle( - context, - comm, - ch_in, - ch_shuffled, - op_id, - num_partitions, - partition_assignment=PartitionAssignment.CONTIGUOUS, + """ + End-to-end correctness of the async streaming shuffler. + + Each rank inserts the input regions it owns and, after shuffling, every local + partition is validated against the conserved data model. + """ + br = context.br() + shuffler = ShufflerAsync(context, comm, 0, total_num_partitions) + + for local_pidx in shuffler.local_partitions(): + chunks = make_partition_data( + total_num_partitions, total_num_rows, local_pidx, stream, br ) - ) - actor, deferred = pull_from_channel(context, ch_shuffled) - actors.append(actor) + if chunks: + shuffler.insert(chunks) - run_actor_network(context, actors=actors) - messages = deferred.release() - received_pids = [msg.sequence_number for msg in messages] + asyncio.run(shuffler.insert_finished(context)) - nranks = comm.nranks - rank = comm.rank - expected_local = list( - range( - rank * num_partitions // nranks, - (rank + 1) * num_partitions // nranks, + for local_pidx in shuffler.local_partitions(): + validate_partition_data( + shuffler.extract(local_pidx), + total_num_partitions, + total_num_rows, + local_pidx, ) - ) - assert set(received_pids) == set(expected_local) - assert len(received_pids) == len(expected_local) -def test_shuffler_object_interface( +@pytest.mark.parametrize("n_inserts", [1, 10]) +@pytest.mark.parametrize("n_partitions", [1, 10, 100]) +def test_shuffler_insert_wait_extract( context: Context, comm: Communicator, + stream: Stream, + n_inserts: int, + n_partitions: int, ) -> None: - if comm.nranks != 1: - pytest.skip("Only support single-rank runs") - actors: list[CppActor | Awaitable[None]] = [] - - num_partitions = 5 - num_rows = 100 - num_chunks = 4 - op_id = 0 - ch_in: Channel[TableChunk] = context.create_channel() - actors.append(generate_inputs(context, ch_in, num_rows, num_chunks)) - ch_shuffled: Channel[TableChunk] = context.create_channel() - actors.append( - do_shuffle( - context, - comm, - ch_in, - ch_shuffled, - op_id, - num_partitions, - ) - ) - actor, deferred = pull_from_channel(context, ch_shuffled) - actors.append(actor) - - run_actor_network(context, actors=actors) - messages = deferred.release() - # TODO: single rank only assertions - assert len(messages) == 5 - assert [msg.sequence_number for msg in messages] == list(range(num_partitions)) - chunks = [ - (msg.sequence_number, TableChunk.from_message(msg, br=context.br())) - for msg in messages - ] - - full_column = np.arange(num_rows * num_chunks, dtype=np.int32) - part_size = num_rows // num_partitions + (num_rows % num_partitions) - splits = [*range(0, num_rows, part_size), num_rows] - for pid, table in chunks: - expect = plc.Column.from_array( - np.concat( - [ - full_column[i * num_rows : (i + 1) * num_rows][ - splits[pid] : splits[pid + 1] - ] - for i in range(num_chunks) - ] - ), - stream=table.stream, - ) - got = table.table_view() - table.stream.synchronize() - assert_eq(plc.Table([expect]), got, sort_rows=0) + """ + Each rank inserts ``n_inserts`` full partition maps; after shuffling each local + partition must receive exactly ``n_inserts * nranks`` chunks. + """ + n_elements = 100 + br = context.br() + shuffler = ShufflerAsync(context, comm, 0, n_partitions) + + for _ in range(n_inserts): + data = { + pid: generate_packed_data(n_elements, 0, stream, br) + for pid in range(n_partitions) + } + shuffler.insert(data) + + asyncio.run(shuffler.insert_finished(context)) + + local_pids = shuffler.local_partitions() + + finished_pids = [] + n_chunks_received = 0 + for pid in local_pids: + chunks = shuffler.extract(pid) + n_chunks_received += len(chunks) + finished_pids.append(pid) + + assert n_chunks_received == n_inserts * len(local_pids) * comm.nranks + assert finished_pids == local_pids diff --git a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py index e4cdb6cb8..ae7e3dea9 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py @@ -2,27 +2,13 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -import math from typing import TYPE_CHECKING -import numpy as np -import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming") -from cudf_streaming.integrations.partition import ( - partition_and_pack, - unpack_and_concat, -) - from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.memory.spill import unspill_partitions -from rapidsmpf.shuffler import ( - Shuffler, -) -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") +from rapidsmpf.shuffler import Shuffler +from rapidsmpf.testing import make_partition_data, validate_partition_data if TYPE_CHECKING: import rmm.mr @@ -31,113 +17,17 @@ from rapidsmpf.communicator.communicator import Communicator -@pytest.mark.parametrize("total_num_partitions", [1, 2, 3, 10]) -def test_shuffler_single_nonempty_partition( - comm: Communicator, - device_mr: rmm.mr.CudaMemoryResource, - stream: Stream, - total_num_partitions: int, -) -> None: - br = BufferResource(device_mr) - - shuffler = Shuffler( - comm, - op_id=0, - total_num_partitions=total_num_partitions, - br=br, - ) - - df = plc.Table( - [ - plc.Column.from_iterable_of_py( - [1, 2, 3], plc.DataType(plc.TypeId.INT64), stream=stream - ), - plc.Column.from_iterable_of_py( - [42, 42, 42], plc.DataType(plc.TypeId.INT64), stream=stream - ), - ] - ) - packed_inputs = partition_and_pack( - df, - columns_to_hash=(1,), - num_partitions=total_num_partitions, - br=br, - stream=stream, - ) - shuffler.insert_chunks(packed_inputs) - shuffler.insert_finished() - - expected_partitions = set(shuffler.local_partitions()) - - local_outputs = [] - extracted_partitions = set() - shuffler.wait() - for partition_id in shuffler.local_partitions(): - extracted_partitions.add(partition_id) - packed_chunks = shuffler.extract(partition_id) - partition = unpack_and_concat( - unspill_partitions(packed_chunks, br=br, allow_overbooking=True), - br=br, - stream=stream, - ) - local_outputs.append(partition) - shuffler.shutdown() - assert extracted_partitions == expected_partitions - # Everything should go to a single rank thus we should get the whole dataframe or nothing. - if len(local_outputs) == 0: - return - res = plc.concatenate.concatenate(local_outputs) - # Each rank has `df` thus each rank contribute to the rows of `df` to the expected result. - expect = plc.concatenate.concatenate([df] * comm.nranks, stream=stream) - if res.num_rows() > 0: - assert_eq(res, expect, sort_rows=0) - - -@pytest.mark.parametrize("batch_size", [None, 10]) -@pytest.mark.parametrize("total_num_partitions", [1, 2, 3, 10]) -def test_shuffler_uniform( +@pytest.mark.parametrize("total_num_partitions", [1, 2, 5, 10]) +@pytest.mark.parametrize("total_num_rows", [1, 9, 100, 100_000]) +def test_shuffler_round_trip( comm: Communicator, device_mr: rmm.mr.CudaMemoryResource, stream: Stream, - batch_size: int | None, total_num_partitions: int, + total_num_rows: int, ) -> None: + """End-to-end shuffle of a conserved, front-loaded data model.""" br = BufferResource(device_mr) - - # Every rank creates the full input dataframe and all the expected partitions - # (also partitions this rank might not get after the shuffle). - num_rows = 100 - np.random.seed(42) # Make sure all ranks create the same input dataframe. - df = plc.Table( - [ - plc.Column.from_iterable_of_py( - range(num_rows), plc.DataType(plc.TypeId.INT64), stream=stream - ), - plc.Column.from_array(np.random.randint(0, 1000, num_rows), stream=stream), - plc.Column.from_iterable_of_py( - ["cat", "dog"] * (num_rows // 2), - plc.DataType(plc.TypeId.STRING), - stream=stream, - ), - ] - ) - columns_to_hash = (1,) - - expected = { - partition_id: unpack_and_concat( - [packed], - br=br, - stream=stream, - ) - for partition_id, packed in partition_and_pack( - df, - columns_to_hash=columns_to_hash, - num_partitions=total_num_partitions, - br=br, - stream=stream, - ).items() - } - shuffler = Shuffler( comm, op_id=0, @@ -145,47 +35,22 @@ def test_shuffler_uniform( br=br, ) - # Slice df and submit local slices to shuffler - stride = math.ceil(num_rows / comm.nranks) - local_df = plc.copying.slice( - df, - [comm.rank * stride, min((comm.rank + 1) * stride, num_rows)], - stream=stream, - )[0] - num_rows_local = local_df.num_rows() - batch_size = batch_size or num_rows_local - for i in range(0, num_rows_local, batch_size): - batch = plc.copying.slice( - local_df, [i, min(i + batch_size, num_rows_local)], stream=stream - )[0] - packed_inputs = partition_and_pack( - batch, - columns_to_hash=columns_to_hash, - num_partitions=total_num_partitions, - br=br, - stream=stream, + # Insert every owned input region, wait, then extract and validate. + for local_pidx in shuffler.local_partitions(): + chunks = make_partition_data( + total_num_partitions, total_num_rows, local_pidx, stream, br ) - shuffler.insert_chunks(packed_inputs) - - # Tell shuffler we are done adding data + if chunks: + shuffler.insert_chunks(chunks) shuffler.insert_finished() - - expected_partitions = set(shuffler.local_partitions()) - extracted_partitions = set() shuffler.wait() - for partition_id in shuffler.local_partitions(): - extracted_partitions.add(partition_id) - packed_chunks = shuffler.extract(partition_id) - partition = unpack_and_concat( - unspill_partitions(packed_chunks, br=br, allow_overbooking=True), - br=br, - stream=stream, - ) - assert_eq( - partition, - expected[partition_id], - sort_rows=0, + + for local_pidx in shuffler.local_partitions(): + validate_partition_data( + shuffler.extract(local_pidx), + total_num_partitions, + total_num_rows, + local_pidx, ) shuffler.shutdown() - assert extracted_partitions == expected_partitions From 1aad864e7cd3771a559a6f1076da7bf74348d32b Mon Sep 17 00:00:00 2001 From: niranda perera Date: Tue, 9 Jun 2026 16:36:34 -0700 Subject: [PATCH 04/14] addressign comments Signed-off-by: niranda perera --- cpp/include/rapidsmpf/utils/misc.hpp | 5 ++-- cpp/tests/streaming/test_shuffler.cpp | 15 +++++------ cpp/tests/test_misc.cpp | 38 +++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 11 deletions(-) diff --git a/cpp/include/rapidsmpf/utils/misc.hpp b/cpp/include/rapidsmpf/utils/misc.hpp index 0407731d7..26412fb9d 100644 --- a/cpp/include/rapidsmpf/utils/misc.hpp +++ b/cpp/include/rapidsmpf/utils/misc.hpp @@ -199,7 +199,8 @@ constexpr T safe_div(T x, T y) { * @brief Computes the ceiling of the division of two integers. * * Returns the smallest integer not less than `x / y`. Both operands must be - * non-negative and the denominator must be non-zero. + * non-negative and the denominator must be non-zero. Computed as `x / y + (x % y != 0)` + * to avoid the overflow (and signed UB) * * @tparam T An integral type. * @param x The numerator (must be non-negative). @@ -208,7 +209,7 @@ constexpr T safe_div(T x, T y) { */ template constexpr T ceil_div(T x, T y) { - return (x + y - 1) / y; + return x / y + (x % y != 0 ? T{1} : T{0}); } /** diff --git a/cpp/tests/streaming/test_shuffler.cpp b/cpp/tests/streaming/test_shuffler.cpp index 5dab73b1b..ab92a94d5 100644 --- a/cpp/tests/streaming/test_shuffler.cpp +++ b/cpp/tests/streaming/test_shuffler.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -98,25 +99,21 @@ TEST_P(StreamingShuffler, basic_shuffler) { const int64_t base = static_cast(comm->rank()) * num_rows; std::vector input_chunks; // Message contains a PartitionMapChunk for (size_t chunk_idx = 0; chunk_idx < num_chunks; ++chunk_idx) { - ContentDescription cd{}; std::unordered_map chunks; chunks.reserve(num_partitions); for (size_t j = 0; j < num_partitions; ++j) { auto [start, end] = pieces[chunk_idx * num_partitions + j]; // end > start is guaranteed. - auto [it, _] = chunks.emplace( + chunks.emplace( static_cast(j), generate_packed_data( end - start, base + static_cast(start), stream, *br ) ); - cd.content_size(it->second.data->mem_type()) += it->second.data->size; } - input_chunks.emplace_back(Message( - chunk_idx, - std::make_unique(std::move(chunks)), - std::move(cd) - )); + input_chunks.emplace_back( + to_message(chunk_idx, std::make_unique(std::move(chunks))) + ); } EXPECT_EQ(input_chunks.size(), num_chunks); @@ -150,7 +147,7 @@ TEST_P(StreamingShuffler, basic_shuffler) { auto p_vec = chunk.release(); // for each local pid, it should receive num_chunks * nranks chunks. - EXPECT_EQ(p_vec.data.size(), num_chunks * n_ranks); + ASSERT_EQ(p_vec.data.size(), num_chunks * n_ranks); // since values are offset by rank, if we sort packed data by their first element, // then it will be in rank & chunk-index order. diff --git a/cpp/tests/test_misc.cpp b/cpp/tests/test_misc.cpp index d77facb91..08e62475e 100644 --- a/cpp/tests/test_misc.cpp +++ b/cpp/tests/test_misc.cpp @@ -233,6 +233,44 @@ TEST(MiscTest, SafeCastZeroAndBoundaries) { EXPECT_THROW(safe_cast(std::uint16_t{256}), std::overflow_error); } +// Test ceil_div edge cases. `ceil_div` is constexpr, so these are all +// constant-evaluated and verified at compile time. +TEST(MiscTest, CeilDiv) { + // Basic behavior. + static_assert(ceil_div(0, 1) == 0); + static_assert(ceil_div(1, 1) == 1); + static_assert(ceil_div(10, 5) == 2); // exact division + static_assert(ceil_div(11, 5) == 3); // remainder rounds up + static_assert(ceil_div(9, 5) == 2); + static_assert(ceil_div(1, 5) == 1); // numerator smaller than denominator + static_assert(ceil_div(5, 1) == 5); // denominator of one + + // A zero numerator always yields zero, regardless of denominator. + static_assert(ceil_div(0u, 1u) == 0u); + static_assert(ceil_div(0u, 7u) == 0u); + static_assert( + ceil_div(std::size_t{0}, std::numeric_limits::max()) + == std::size_t{0} + ); + + // Large unsigned values that would overflow the naive `(x + y - 1) / y`. + constexpr auto umax = std::numeric_limits::max(); + static_assert(ceil_div(umax, umax) == std::uint64_t{1}); // `x + y - 1` wraps to 0 + static_assert(ceil_div(umax - 1, umax) == std::uint64_t{1}); + static_assert(ceil_div(umax, std::uint64_t{1}) == umax); + // `umax` is odd, so ceil(umax / 2) rounds up to 2^63. + static_assert(ceil_div(umax, std::uint64_t{2}) == (umax / 2) + 1); + + // Large signed values near the maximum must not trigger signed-overflow UB. + constexpr auto imax = std::numeric_limits::max(); + static_assert(ceil_div(imax, imax) == std::int64_t{1}); + static_assert(ceil_div(imax - 1, imax) == std::int64_t{1}); + static_assert(ceil_div(imax, std::int64_t{1}) == imax); + static_assert(ceil_div(imax, std::int64_t{2}) == (imax / 2) + 1); + + SUCCEED(); +} + // Test mixed signed/unsigned of different sizes TEST(MiscTest, SafeCastMixedSignednessAndSize) { // std::int64_t to std::uint32_t From bcee4f27381c13895fffa5eb6a4a5c0b770a1ef1 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 10 Jun 2026 14:33:07 -0700 Subject: [PATCH 05/14] add docs Signed-off-by: niranda perera --- cpp/tests/test_shuffler.cpp | 100 ++++++++++++++++++++++++++++-------- 1 file changed, 79 insertions(+), 21 deletions(-) diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index 36e2c5fc1..54cbd021b 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -110,8 +110,19 @@ namespace { using MemoryLimitsMap = std::unordered_map; -// Help function to get the `memory_limits` argument for a `BufferResource` -// that prioritizes the specified memory type. +/** + * @brief Build a `memory_limits` map for a `BufferResource` that prioritizes one memory + * type. + * + * All memory types are initialised to unlimited. If @p priorities is not + * `MemoryType::DEVICE`, the device-memory limit is then set to zero, forcing + * the `BufferResource` to allocate exclusively in host memory. Host memory is + * never zeroed because it backs metadata and control-message allocations that + * must always succeed. + * + * @param priorities The memory type to keep unlimited (all others are zeroed). + * @return A map from each `MemoryType` to its byte limit (`std::int64_t`). + */ MemoryLimitsMap get_memory_limits_map(rapidsmpf::MemoryType priorities) { using namespace rapidsmpf; @@ -130,19 +141,35 @@ MemoryLimitsMap get_memory_limits_map(rapidsmpf::MemoryType priorities) { return ret; } -// Conservation-preserving data model shared by the shuffler round-trip tests. -// -// We split the index range [0, total_num_rows) into total_num_partitions^2 contiguous -// sub-regions via chunk_indices (front-loaded, so when N < P*P the trailing sub-regions -// are empty). Sub-region (local_pidx, split_idx) is piece k = local_pidx*P + split_idx -// and is routed to destination partition split_idx; input region local_pidx is the union -// of its P sub-regions. The pieces exactly tile [0,N), so the total shuffled data == N -// regardless of rank/partition counts (conservation). A per-shuffle `base` offset is -// added to every value so distinct shuffles carry distinct data. - -// Produces the non-empty sub-regions of one owned input region `local_pidx`, keyed by -// destination partition. Since local_partitions() across ranks partition [0,P), every -// input region is produced exactly once, so rows are not replicated. +/// Conservation-preserving data model shared by the shuffler round-trip tests. +/// +/// The index range `[0, total_num_rows)` is split into `total_num_partitions^2` +/// contiguous sub-regions via `chunk_indices` (front-loaded, so trailing sub-regions +/// are empty when `N < P*P`). Sub-region `(local_pidx, split_idx)` is piece +/// `k = local_pidx * P + split_idx` and is routed to destination partition +/// `split_idx`. The pieces exactly tile `[0, N)`, so the total shuffled row +/// count equals `N` regardless of rank or partition counts (conservation). A +/// per-shuffle `base` offset is added to every value so distinct shuffles carry +/// distinct data. + + +/** + * @brief Build the input data for one owned partition region ready for insertion. + * + * Produces all non-empty sub-regions of the input region `local_pidx`, keyed by + * their destination partition. Because `local_partitions()` across all ranks + * partitions `[0, P)`, every input region is produced exactly once and rows are + * never replicated. + * + * @param total_num_partitions Total number of shuffle partitions `P`. + * @param total_num_rows Total row count `N` tiled across all sub-regions. + * @param local_pidx Index of the locally-owned input region to generate. + * @param stream CUDA stream used for device allocations. + * @param br Buffer resource used to allocate packed data. + * @param base Offset added to every generated value (default 0). + * @return Map from destination `PartID` to the corresponding `PackedData` chunk; + * empty sub-regions are omitted. + */ std::unordered_map make_partition_data( rapidsmpf::shuffler::PartID total_num_partitions, @@ -172,8 +199,22 @@ make_partition_data( return chunks; } -// Verifies that the `received` chunks for partition `j` match the non-empty sub-regions -// expected for it. +/** + * @brief Verify that received chunks for a partition match the expected sub-regions. + * + * Recomputes the non-empty `(offset, count)` sub-regions expected for partition + * `j` from the same conservation model used by `make_partition_data`, then + * checks that @p received contains exactly those chunks (in any order). Chunks + * are sorted by their embedded offset before comparison so the validation is + * order-independent. + * + * @param received Chunks extracted from the shuffler for partition `j`. + * @param total_num_partitions Total number of shuffle partitions `P`. + * @param total_num_rows Total row count `N` used to tile sub-regions. + * @param j Destination partition index being validated. + * @param br Buffer resource used for unpacking received data. + * @param base Offset that was added to every generated value (default 0). + */ void validate_partition_data( std::vector received, rapidsmpf::shuffler::PartID total_num_partitions, @@ -215,12 +256,27 @@ void validate_partition_data( } } +/** + * @brief Execute a full shuffler round-trip and validate every local partition. + * + * For each locally-owned partition, generates input data with `make_partition_data`, + * inserts it into the shuffler, signals insertion completion, then waits (with a + * 30-second timeout to catch deadlocks) and validates every received partition + * with `validate_partition_data`. + * + * @param shuffler The shuffler instance under test. + * @param total_num_partitions Total number of shuffle partitions `P`. + * @param total_num_rows Total row count `N` distributed across all sub-regions. + * @param stream CUDA stream used for device allocations. + * @param br Buffer resource used to allocate and validate data. + */ void test_shuffler( rapidsmpf::shuffler::Shuffler& shuffler, rapidsmpf::shuffler::PartID total_num_partitions, std::size_t total_num_rows, rmm::cuda_stream_view stream, - rapidsmpf::BufferResource* br + rapidsmpf::BufferResource* br, + std::int64_t base = 0 ) { // To expose unexpected deadlocks, we use a 30s timeout. In a normal run, the // shuffle shouldn't get near 30s. @@ -228,7 +284,7 @@ void test_shuffler( for (rapidsmpf::shuffler::PartID local_pidx : shuffler.local_partitions()) { shuffler.insert(make_partition_data( - total_num_partitions, total_num_rows, local_pidx, stream, *br + total_num_partitions, total_num_rows, local_pidx, stream, *br, base )); } shuffler.insert_finished(); @@ -241,7 +297,8 @@ void test_shuffler( total_num_partitions, total_num_rows, local_pidx, - *br + *br, + base ); } } @@ -338,7 +395,8 @@ class ConcurrentShuffleTest : public ::testing::TestWithParam< total_num_partitions, 100'000, // total_num_rows stream, - br.get() + br.get(), + static_cast(t_id) )); } From 30953424db49cb17d970457e328f676cfc0f8926 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Wed, 10 Jun 2026 16:58:13 -0700 Subject: [PATCH 06/14] addressing comments Signed-off-by: niranda perera --- .../rapidsmpf/streaming/chunks/partition.pyi | 12 ++ .../rapidsmpf/streaming/chunks/partition.pyx | 178 +++++++++++++++++- python/rapidsmpf/rapidsmpf/testing.py | 3 +- .../tests/streaming/test_shuffler.py | 174 +++++++++++++++++ 4 files changed, 365 insertions(+), 2 deletions(-) diff --git a/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyi b/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyi index f94280f0e..c4fb73400 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyi +++ b/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyi @@ -2,21 +2,33 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations +from collections.abc import Mapping, Sequence from typing import Self from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.core.message import Message class PartitionMapChunk: + @classmethod + def from_packed_data_map( + cls: type[Self], data: Mapping[int, PackedData], br: BufferResource + ) -> Self: ... @classmethod def from_message( cls: type[Self], message: Message[Self], br: BufferResource ) -> Self: ... + def to_packed_data_map(self) -> dict[int, PackedData]: ... def into_message(self, sequence_number: int, message: Message[Self]) -> None: ... class PartitionVectorChunk: + @classmethod + def from_packed_data_list( + cls: type[Self], data: Sequence[PackedData], br: BufferResource + ) -> Self: ... @classmethod def from_message( cls: type[Self], message: Message[Self], br: BufferResource ) -> Self: ... + def to_packed_data_list(self) -> list[PackedData]: ... def into_message(self, sequence_number: int, message: Message[Self]) -> None: ... diff --git a/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyx b/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyx index d3197ea29..2d88866a4 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyx +++ b/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyx @@ -1,12 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -from libc.stdint cimport uint64_t +from libc.stdint cimport uint32_t, uint64_t from libcpp.memory cimport make_unique, unique_ptr from libcpp.utility cimport move +from libcpp.vector cimport vector from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.memory.buffer_resource cimport BufferResource +from rapidsmpf.memory.packed_data cimport PackedData, cpp_PackedData from rapidsmpf.streaming.core.message cimport Message, cpp_Message @@ -19,6 +21,72 @@ cdef extern from "" nogil: except +ex_handler +# Move PackedData into a chunk's container. We implement these in C++ because +# PackedData doesn't have a default ctor. +cdef extern from * nogil: + """ + namespace { + void cpp_insert_into_partition_map( + rapidsmpf::streaming::PartitionMapChunk* chunk, + std::uint32_t pid, + std::unique_ptr packed_data + ) { + chunk->data.emplace(pid, std::move(*packed_data)); + } + + void cpp_append_to_partition_vector( + rapidsmpf::streaming::PartitionVectorChunk* chunk, + std::unique_ptr packed_data + ) { + chunk->data.push_back(std::move(*packed_data)); + } + + void cpp_drain_partition_map( + rapidsmpf::streaming::PartitionMapChunk* chunk, + std::vector& keys, + std::vector>& values + ) { + keys.reserve(chunk->data.size()); + values.reserve(chunk->data.size()); + for (auto& [pid, pd] : chunk->data) { + keys.push_back(pid); + values.push_back(std::make_unique(std::move(pd))); + } + chunk->data.clear(); + } + + void cpp_drain_partition_vector( + rapidsmpf::streaming::PartitionVectorChunk* chunk, + std::vector>& values + ) { + values.reserve(chunk->data.size()); + for (auto& pd : chunk->data) { + values.push_back(std::make_unique(std::move(pd))); + } + chunk->data.clear(); + } + } // namespace + """ + void cpp_insert_into_partition_map( + cpp_PartitionMapChunk* chunk, + uint32_t pid, + unique_ptr[cpp_PackedData] packed_data, + ) except +ex_handler + void cpp_append_to_partition_vector( + cpp_PartitionVectorChunk* chunk, + unique_ptr[cpp_PackedData] packed_data, + ) except +ex_handler + void cpp_drain_partition_map( + cpp_PartitionMapChunk* chunk, + vector[uint32_t]& keys, + vector[unique_ptr[cpp_PackedData]]& values, + ) except +ex_handler + void cpp_drain_partition_vector( + cpp_PartitionVectorChunk* chunk, + vector[unique_ptr[cpp_PackedData]]& values, + ) except +ex_handler + + cdef class PartitionMapChunk: def __init__(self): raise ValueError("use the `from_*` factory functions") @@ -27,6 +95,62 @@ cdef class PartitionMapChunk: with nogil: self._handle.reset() + @staticmethod + def from_packed_data_map(dict data not None, BufferResource br not None): + """ + Construct a PartitionMapChunk from a mapping of partition ID to PackedData. + + Parameters + ---------- + data + Mapping of partition ID to the + :class:`~rapidsmpf.memory.packed_data.PackedData` it holds. Each PackedData + is consumed and left empty after this call. + br + Buffer resource kept alive for the lifetime of the chunk. + + Returns + ------- + A new PartitionMapChunk owning the given packed data. + + Raises + ------ + ValueError + If any of the provided PackedData objects is empty. + """ + cdef unique_ptr[cpp_PartitionMapChunk] handle = make_unique[ + cpp_PartitionMapChunk + ]() + cdef uint32_t pid + cdef PackedData pd + for key, value in data.items(): + pid = key + pd = value + if not pd.c_obj: + raise ValueError("PackedData was empty") + cpp_insert_into_partition_map(handle.get(), pid, move(pd.c_obj)) + return PartitionMapChunk.from_handle(move(handle), br) + + def to_packed_data_map(self): + """ + Extract the partition data as a mapping of partition ID to PackedData. + + The chunk is drained and left empty after this call. + + Returns + ------- + A dict mapping partition ID to the + :class:`~rapidsmpf.memory.packed_data.PackedData` it holds. + """ + cdef vector[uint32_t] keys + cdef vector[unique_ptr[cpp_PackedData]] values + cpp_drain_partition_map(self._handle.get(), keys, values) + cdef dict ret = {} + cdef size_t i + for i in range(values.size()): + ret[keys[i]] = PackedData.from_librapidsmpf(move(values[i]), self._br) + return ret + @staticmethod cdef PartitionMapChunk from_handle( unique_ptr[cpp_PartitionMapChunk] handle, BufferResource br @@ -147,6 +271,58 @@ cdef class PartitionVectorChunk: with nogil: self._handle.reset() + @staticmethod + def from_packed_data_list(list data not None, BufferResource br not None): + """ + Construct a PartitionVectorChunk from a sequence of PackedData. + + Parameters + ---------- + data + Sequence of :class:`~rapidsmpf.memory.packed_data.PackedData` objects, + stored in order. Each PackedData is consumed and left empty after this + call. + br + Buffer resource kept alive for the lifetime of the chunk. + + Returns + ------- + A new PartitionVectorChunk owning the given packed data. + + Raises + ------ + ValueError + If any of the provided PackedData objects is empty. + """ + cdef unique_ptr[cpp_PartitionVectorChunk] handle = make_unique[ + cpp_PartitionVectorChunk + ]() + cdef PackedData pd + for value in data: + pd = value + if not pd.c_obj: + raise ValueError("PackedData was empty") + cpp_append_to_partition_vector(handle.get(), move(pd.c_obj)) + return PartitionVectorChunk.from_handle(move(handle), br) + + def to_packed_data_list(self): + """ + Extract the partition data as a list of PackedData. + + The chunk is drained and left empty after this call. + + Returns + ------- + A list of :class:`~rapidsmpf.memory.packed_data.PackedData`, in order. + """ + cdef vector[unique_ptr[cpp_PackedData]] values + cpp_drain_partition_vector(self._handle.get(), values) + cdef list ret = [] + cdef size_t i + for i in range(values.size()): + ret.append(PackedData.from_librapidsmpf(move(values[i]), self._br)) + return ret + @staticmethod cdef PartitionVectorChunk from_handle( unique_ptr[cpp_PartitionVectorChunk] handle, BufferResource br diff --git a/python/rapidsmpf/rapidsmpf/testing.py b/python/rapidsmpf/rapidsmpf/testing.py index a4dbb145f..a38f4b88a 100644 --- a/python/rapidsmpf/rapidsmpf/testing.py +++ b/python/rapidsmpf/rapidsmpf/testing.py @@ -127,7 +127,8 @@ def generate_packed_data( A ``PackedData`` containing the integer sequence. """ data = np.arange(offset, offset + n_elements, dtype=_DTYPE).tobytes() - gpu_data = rmm.DeviceBuffer.to_device(data, stream=stream) + gpu_data = rmm.DeviceBuffer(size=len(data), stream=stream, mr=br.device_mr) + gpu_data.copy_from_host(data, stream=stream) return PackedData.from_device_buffer(gpu_data, data, stream, br) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py index 5f4e364c6..88a735c08 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py @@ -6,19 +6,33 @@ import asyncio from typing import TYPE_CHECKING +import numpy as np import pytest +from rapidsmpf.shuffler import PartitionAssignment +from rapidsmpf.streaming.chunks.partition import ( + PartitionMapChunk, + PartitionVectorChunk, +) from rapidsmpf.streaming.coll.shuffler import ShufflerAsync +from rapidsmpf.streaming.core.actor import define_actor, run_actor_network +from rapidsmpf.streaming.core.leaf_actor import pull_from_channel +from rapidsmpf.streaming.core.message import Message from rapidsmpf.testing import ( generate_packed_data, make_partition_data, + validate_packed_data, validate_partition_data, ) if TYPE_CHECKING: + from collections.abc import Awaitable + from rmm.pylibrmm.stream import Stream from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.streaming.core.actor import CppActor + from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context @@ -95,3 +109,163 @@ def test_shuffler_insert_wait_extract( assert n_chunks_received == n_inserts * len(local_pids) * comm.nranks assert finished_pids == local_pids + + +@define_actor() +async def generate_inputs( + context: Context, + ch: Channel[PartitionMapChunk], + num_rows: int, + num_chunks: int, + num_partitions: int, +) -> None: + br = context.br() + for i in range(num_chunks): + stream = context.get_stream_from_pool() + data = { + pid: generate_packed_data( + num_rows, (i * num_partitions + pid) * num_rows, stream, br + ) + for pid in range(num_partitions) + } + msg = Message(i, PartitionMapChunk.from_packed_data_map(data, br)) + await ch.send(context, msg) + await ch.drain(context) + + +@define_actor() +async def do_shuffle( + context: Context, + comm: Communicator, + ch_in: Channel[PartitionMapChunk], + ch_out: Channel[PartitionVectorChunk], + op_id: int, + num_partitions: int, + *, + partition_assignment: PartitionAssignment = PartitionAssignment.ROUND_ROBIN, +) -> None: + shuffle = ShufflerAsync( + context, comm, op_id, num_partitions, partition_assignment=partition_assignment + ) + while (msg := await ch_in.recv(context)) is not None: + chunk = PartitionMapChunk.from_message(msg, br=context.br()) + shuffle.insert(chunk.to_packed_data_map()) + await shuffle.insert_finished(context) + for pid in shuffle.local_partitions(): + data = shuffle.extract(pid) + out_chunk = PartitionVectorChunk.from_packed_data_list(data, context.br()) + await ch_out.send(context, Message(pid, out_chunk)) + await ch_out.drain(context) + + +@pytest.mark.parametrize("num_partitions", [4, 8]) +def test_shuffler_runtime_obeys_contiguous_assignment( + context: Context, + comm: Communicator, + num_partitions: int, +) -> None: + if comm.nranks != 1: + pytest.skip("Only support single-rank runs") + + actors: list[CppActor | Awaitable[None]] = [] + + num_rows = 200 + num_chunks = 3 + op_id = 0 + ch_in: Channel[PartitionMapChunk] = context.create_channel() + actors.append(generate_inputs(context, ch_in, num_rows, num_chunks, num_partitions)) + ch_shuffled: Channel[PartitionVectorChunk] = context.create_channel() + actors.append( + do_shuffle( + context, + comm, + ch_in, + ch_shuffled, + op_id, + num_partitions, + partition_assignment=PartitionAssignment.CONTIGUOUS, + ) + ) + actor, deferred = pull_from_channel(context, ch_shuffled) + actors.append(actor) + + run_actor_network(context, actors=actors) + messages = deferred.release() + received_pids = [msg.sequence_number for msg in messages] + + # Single rank, so every partition is local to this rank. + assert set(received_pids) == set(range(num_partitions)) + + # Validate the data routed to each local partition. Across the ``num_chunks`` + # inputs, partition ``pid`` receives the packed sequence ``generate_inputs`` + # produced for ``(chunk i, pid)``, which starts at value + # ``(i * num_partitions + pid) * num_rows``. The shuffler makes no ordering + # guarantee, so match each received chunk to its expected input by start value. + for msg in messages: + pid = msg.sequence_number + packed = PartitionVectorChunk.from_message( + msg, br=context.br() + ).to_packed_data_list() + assert len(packed) == num_chunks + by_offset = { + int(np.frombuffer(pd.to_host_bytes(), dtype=np.int64)[0]): pd + for pd in packed + } + for i in range(num_chunks): + offset = (i * num_partitions + pid) * num_rows + validate_packed_data(by_offset[offset], num_rows, offset) + + +def test_shuffler_object_interface( + context: Context, + comm: Communicator, +) -> None: + if comm.nranks != 1: + pytest.skip("Only support single-rank runs") + actors: list[CppActor | Awaitable[None]] = [] + + num_partitions = 5 + num_rows = 100 + num_chunks = 4 + op_id = 0 + ch_in: Channel[PartitionMapChunk] = context.create_channel() + actors.append(generate_inputs(context, ch_in, num_rows, num_chunks, num_partitions)) + ch_shuffled: Channel[PartitionVectorChunk] = context.create_channel() + actors.append( + do_shuffle( + context, + comm, + ch_in, + ch_shuffled, + op_id, + num_partitions, + ) + ) + actor, deferred = pull_from_channel(context, ch_shuffled) + actors.append(actor) + + run_actor_network(context, actors=actors) + messages = deferred.release() + # TODO: single rank only assertions + assert len(messages) == num_partitions + assert [msg.sequence_number for msg in messages] == list(range(num_partitions)) + chunks = [ + (msg.sequence_number, PartitionVectorChunk.from_message(msg, br=context.br())) + for msg in messages + ] + + # Each destination partition ``pid`` receives, across the ``num_chunks`` inputs, + # the packed sequence generated by ``generate_inputs`` for ``(chunk i, pid)``, + # which starts at value ``(i * num_partitions + pid) * num_rows``. The shuffler + # makes no ordering guarantee, so match each received chunk to its expected + # input chunk by starting value. + for pid, vec_chunk in chunks: + packed = vec_chunk.to_packed_data_list() + assert len(packed) == num_chunks + by_offset = { + int(np.frombuffer(pd.to_host_bytes(), dtype=np.int64)[0]): pd + for pd in packed + } + for i in range(num_chunks): + offset = (i * num_partitions + pid) * num_rows + validate_packed_data(by_offset[offset], num_rows, offset) From 394d9a44148ba22cf2866489a0756a16e4f8b549 Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 11 Jun 2026 08:17:09 -0700 Subject: [PATCH 07/14] using shuffler actor Signed-off-by: niranda perera --- .../tests/streaming/test_shuffler.py | 31 ++----------------- 1 file changed, 3 insertions(+), 28 deletions(-) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py index 88a735c08..49dcaea04 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py @@ -14,7 +14,7 @@ PartitionMapChunk, PartitionVectorChunk, ) -from rapidsmpf.streaming.coll.shuffler import ShufflerAsync +from rapidsmpf.streaming.coll.shuffler import ShufflerAsync, shuffler from rapidsmpf.streaming.core.actor import define_actor, run_actor_network from rapidsmpf.streaming.core.leaf_actor import pull_from_channel from rapidsmpf.streaming.core.message import Message @@ -133,31 +133,6 @@ async def generate_inputs( await ch.drain(context) -@define_actor() -async def do_shuffle( - context: Context, - comm: Communicator, - ch_in: Channel[PartitionMapChunk], - ch_out: Channel[PartitionVectorChunk], - op_id: int, - num_partitions: int, - *, - partition_assignment: PartitionAssignment = PartitionAssignment.ROUND_ROBIN, -) -> None: - shuffle = ShufflerAsync( - context, comm, op_id, num_partitions, partition_assignment=partition_assignment - ) - while (msg := await ch_in.recv(context)) is not None: - chunk = PartitionMapChunk.from_message(msg, br=context.br()) - shuffle.insert(chunk.to_packed_data_map()) - await shuffle.insert_finished(context) - for pid in shuffle.local_partitions(): - data = shuffle.extract(pid) - out_chunk = PartitionVectorChunk.from_packed_data_list(data, context.br()) - await ch_out.send(context, Message(pid, out_chunk)) - await ch_out.drain(context) - - @pytest.mark.parametrize("num_partitions", [4, 8]) def test_shuffler_runtime_obeys_contiguous_assignment( context: Context, @@ -176,7 +151,7 @@ def test_shuffler_runtime_obeys_contiguous_assignment( actors.append(generate_inputs(context, ch_in, num_rows, num_chunks, num_partitions)) ch_shuffled: Channel[PartitionVectorChunk] = context.create_channel() actors.append( - do_shuffle( + shuffler( context, comm, ch_in, @@ -232,7 +207,7 @@ def test_shuffler_object_interface( actors.append(generate_inputs(context, ch_in, num_rows, num_chunks, num_partitions)) ch_shuffled: Channel[PartitionVectorChunk] = context.create_channel() actors.append( - do_shuffle( + shuffler( context, comm, ch_in, From 328cf4c37fe033b6c5b5cf4516f106f4eba191cc Mon Sep 17 00:00:00 2001 From: niranda perera Date: Thu, 11 Jun 2026 08:29:21 -0700 Subject: [PATCH 08/14] relax cython type Signed-off-by: niranda perera --- python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyx b/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyx index 2d88866a4..e2a3de782 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyx +++ b/python/rapidsmpf/rapidsmpf/streaming/chunks/partition.pyx @@ -96,7 +96,7 @@ cdef class PartitionMapChunk: self._handle.reset() @staticmethod - def from_packed_data_map(dict data not None, BufferResource br not None): + def from_packed_data_map(data not None, BufferResource br not None): """ Construct a PartitionMapChunk from a mapping of partition ID to PackedData. @@ -272,7 +272,7 @@ cdef class PartitionVectorChunk: self._handle.reset() @staticmethod - def from_packed_data_list(list data not None, BufferResource br not None): + def from_packed_data_list(data not None, BufferResource br not None): """ Construct a PartitionVectorChunk from a sequence of PackedData. From 15e85626d3fb4452b0af3ed651bd635fee3d598c Mon Sep 17 00:00:00 2001 From: Niranda Perera Date: Thu, 11 Jun 2026 09:07:13 -0700 Subject: [PATCH 09/14] Update python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py Co-authored-by: Peter Andreas Entschev --- python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py index 49dcaea04..cb67c2780 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py @@ -170,6 +170,7 @@ def test_shuffler_runtime_obeys_contiguous_assignment( # Single rank, so every partition is local to this rank. assert set(received_pids) == set(range(num_partitions)) + assert len(received_pids) == num_partitions # Validate the data routed to each local partition. Across the ``num_chunks`` # inputs, partition ``pid`` receives the packed sequence ``generate_inputs`` From 52b2909d934b96aa203ed9230fd555ad16052eaf Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 11 Jun 2026 05:50:41 -0700 Subject: [PATCH 10/14] test: restore cudf-free core coverage --- cpp/tests/CMakeLists.txt | 3 +- cpp/tests/test_shuffler_many_streams.cpp | 11 +- .../rapidsmpf/tests/test_allgather.py | 120 +++--------------- .../rapidsmpf/tests/test_sparse_alltoall.py | 61 ++------- 4 files changed, 35 insertions(+), 160 deletions(-) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 571d529cb..8306a7854 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -87,6 +87,7 @@ target_sources( test_progress_thread.cpp test_rmm_resource_adaptor.cpp test_shuffler.cpp + test_shuffler_many_streams.cpp test_sparse_alltoall.cpp test_spill_manager.cpp test_spilling.cpp @@ -116,7 +117,7 @@ endif() # cudf-dependent test sources, gated behind BUILD_CUDF_TESTS if(BUILD_CUDF_TESTS) - target_sources(test_sources PRIVATE test_partition.cpp test_shuffler_many_streams.cpp) + target_sources(test_sources PRIVATE test_partition.cpp) target_link_libraries( test_sources PRIVATE cudf_streaming::cudf_streaming cudf::cudftestutil cudf::cudftestutil_impl ) diff --git a/cpp/tests/test_shuffler_many_streams.cpp b/cpp/tests/test_shuffler_many_streams.cpp index 1043611c1..361499a04 100644 --- a/cpp/tests/test_shuffler_many_streams.cpp +++ b/cpp/tests/test_shuffler_many_streams.cpp @@ -3,16 +3,15 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include +#include #include +#include #include #include -#include -#include -#include -#include -#include +#include #include #include @@ -50,7 +49,7 @@ TEST(ShufflerManyStreams, Test) { std::mt19937 random_generator{42}; constexpr std::size_t chunksize = 1 << 20; constexpr int num_partitions = 100; - auto br = BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); // Create a CUDA stream for each partition. // To stress-test stream handling, assign random priorities so streams are more diff --git a/python/rapidsmpf/rapidsmpf/tests/test_allgather.py b/python/rapidsmpf/rapidsmpf/tests/test_allgather.py index 244bb6c1b..6c7a1e78e 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_allgather.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_allgather.py @@ -8,107 +8,18 @@ from typing import TYPE_CHECKING import numpy as np -import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming") -from cudf_streaming.integrations.partition import unpack_and_concat -from pylibcudf.contiguous_split import pack - from rapidsmpf.coll import AllGather from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.memory.packed_data import PackedData -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") +from rapidsmpf.testing import generate_packed_data, validate_packed_data if TYPE_CHECKING: import rmm.mr from rmm.pylibrmm.stream import Stream from rapidsmpf.communicator.communicator import Communicator - - -def generate_packed_data( - n_elements: int, offset: int, stream: Stream, br: BufferResource -) -> PackedData: - """ - Generate a packed data object with the given number of elements and offset. - - Both metadata and gpu_data contain the same sequential integer data starting - from the specified offset. - - Parameters - ---------- - n_elements - Number of integer elements to generate - offset - Starting value for the sequence (offset, offset+1, offset+2, ...) - stream - CUDA stream for operations - br - Buffer resource for memory allocation - - Returns - ------- - Packed data containing the generated sequence - """ - # Generate sequential integers starting from offset - values = np.arange(offset, offset + n_elements, dtype=np.int32) - table = plc.Table([plc.Column.from_array(values, stream=stream)]) - packed_columns = pack(table, stream=stream) - return PackedData.from_cudf_packed_columns(packed_columns, stream, br) # type: ignore[attr-defined, no-any-return] - - -def validate_packed_data( - packed_data: PackedData, - n_elements: int, - offset: int, - stream: Stream, - br: BufferResource, -) -> None: - """ - Validate a packed data object by checking its contents. - - For now, this is a simplified validation that just checks we can - convert the packed data back to a pylibcudf table and that it has - the expected number of rows. - - Parameters - ---------- - packed_data - The packed data to validate - n_elements - Expected number of elements - offset - Expected starting offset value (currently not fully validated) - stream - CUDA stream for operations - - Raises - ------ - AssertionError - If the data doesn't match expectations - """ - # unpack_and_concat expects a list of PackedData - result_table = unpack_and_concat([packed_data], stream, br) - - # Verify the row count matches expected - assert result_table.num_rows() == n_elements - - if n_elements > 0: - # Basic validation - check that we have the expected structure - assert result_table.num_columns() == 1 - - expected_table = plc.Table( - [ - plc.Column.from_array( - np.arange(offset, offset + n_elements, dtype=np.int32), - stream=stream, - ) - ] - ) - assert_eq(result_table, expected_table) + from rapidsmpf.memory.packed_data import PackedData def gen_offset(i: int, r: int) -> int: @@ -181,18 +92,25 @@ def test_basic_allgather( result_idx = r * n_inserts + i expected_offset = gen_offset(i, r) validate_packed_data( - results[result_idx], n_elements, expected_offset, stream, br + results[result_idx], n_elements, expected_offset ) else: - # For unordered results, just verify all expected offsets are present - - # For unordered results, we can't easily determine the exact order, - # so we just validate that each result is valid and has the right size - for result in results: - # Use our validation function with dummy offset (we can't easily extract the real offset) - # This will at least verify the structure and size - result_table = unpack_and_concat([result], stream, br) - assert result_table.num_rows() == n_elements + if n_elements == 0: + for result in results: + validate_packed_data(result, 0, 0) + else: + expected_offsets = { + gen_offset(i, r) for r in range(n_ranks) for i in range(n_inserts) + } + actual_by_offset: dict[int, PackedData] = { + int( + np.frombuffer(result.to_host_bytes(), dtype=np.int64)[0] + ): result + for result in results + } + assert set(actual_by_offset) == expected_offsets + for offset, result in actual_by_offset.items(): + validate_packed_data(result, n_elements, offset) def test_insert_finished_raises_in_context( diff --git a/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py b/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py index 1cd77c47b..58a190582 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py @@ -7,18 +7,11 @@ from typing import TYPE_CHECKING import numpy as np -import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming") -from cudf_streaming.integrations.partition import unpack_and_concat - from rapidsmpf.coll.sparse_alltoall import SparseAlltoall from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.memory.packed_data import PackedData -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") +from rapidsmpf.testing import generate_packed_data, validate_packed_data if TYPE_CHECKING: import rmm.mr @@ -27,25 +20,6 @@ from rapidsmpf.communicator.communicator import Communicator -def generate_packed_data( - n_elements: int, offset: int, stream: Stream, br: BufferResource -) -> PackedData: - """Generate packed integer data with a predictable payload.""" - values = np.arange(offset, offset + n_elements, dtype=np.int32) - table = plc.Table([plc.Column.from_array(values, stream=stream)]) - packed_columns = plc.contiguous_split.pack(table, stream=stream) - return PackedData.from_cudf_packed_columns(packed_columns, stream, br) # type: ignore[attr-defined, no-any-return] - - -def unpack_table( - packed_data: PackedData, - stream: Stream, - br: BufferResource, -) -> plc.Table: - """Unpack a PackedData payload into a single-column table.""" - return unpack_and_concat([packed_data], stream, br) - - def expected_peers(comm: Communicator) -> tuple[list[int], list[int]]: """Return the immediate non-self neighbors in the communicator.""" peers = [] @@ -103,19 +77,7 @@ def test_basic( results = sparse_alltoall.extract(src) assert len(results) == n_inserts for ordinal, result in enumerate(results): - expected = plc.Table( - [ - plc.Column.from_array( - np.arange( - make_offset(src, comm.rank, ordinal), - make_offset(src, comm.rank, ordinal) + 4, - dtype=np.int32, - ), - stream=stream, - ) - ] - ) - assert_eq(unpack_table(result, stream, br), expected) + validate_packed_data(result, 4, make_offset(src, comm.rank, ordinal)) def test_non_participating_ranks( @@ -156,18 +118,13 @@ def test_non_participating_ranks( if comm.rank == 1: results = sparse_alltoall.extract(0) assert len(results) == 2 - assert_eq( - unpack_table(results[0], stream, br), - plc.Table( - [plc.Column.from_array(np.array([11], dtype=np.int32), stream=stream)] - ), - ) - assert_eq( - unpack_table(results[1], stream, br), - plc.Table( - [plc.Column.from_array(np.array([29], dtype=np.int32), stream=stream)] - ), - ) + actual_offsets = [ + int(np.frombuffer(result.to_host_bytes(), dtype=np.int64)[0]) + for result in results + ] + assert actual_offsets == [11, 29] + validate_packed_data(results[0], 1, 11) + validate_packed_data(results[1], 1, 29) def test_invalid_peers_raise( From 20f574f6e9d2bf4633539a49b9d312abc452b1ec Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 11 Jun 2026 06:07:00 -0700 Subject: [PATCH 11/14] test: make spill roundtrip cudf-free --- .../rapidsmpf/rapidsmpf/tests/test_spill.py | 60 ++++++------------- 1 file changed, 17 insertions(+), 43 deletions(-) diff --git a/python/rapidsmpf/rapidsmpf/tests/test_spill.py b/python/rapidsmpf/rapidsmpf/tests/test_spill.py index 7652ebf99..80f5b8339 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_spill.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_spill.py @@ -4,60 +4,34 @@ from typing import TYPE_CHECKING -import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming") -from cudf_streaming.integrations.partition import ( - partition_and_pack, - unpack_and_concat, -) - -from rmm.pylibrmm.stream import DEFAULT_STREAM - from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.memory.spill import spill_partitions, unspill_partitions -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") +from rapidsmpf.testing import generate_packed_data, validate_packed_data if TYPE_CHECKING: import rmm.mr + from rmm.pylibrmm.stream import Stream -def _make_table(cols: list[list[int]]) -> plc.Table: - # Assigns empty column inputs as int64 - return plc.Table( - [ - plc.Column.from_iterable_of_py(col, plc.DataType(plc.TypeId.INT64)) - for col in cols - ] - ) - - -@pytest.mark.parametrize("cols", [[[1, 2, 3], [2, 2, 1]], [[], []]]) -@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) +@pytest.mark.parametrize("num_elements", [0, 1, 10]) +@pytest.mark.parametrize("num_partitions", [1, 2, 10]) def test_spill_unspill_roundtrip( - device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int + device_mr: rmm.mr.CudaMemoryResource, + stream: Stream, + num_elements: int, + num_partitions: int, ) -> None: br = BufferResource(device_mr) - expect = _make_table(cols) - partitions = partition_and_pack( - expect, - columns_to_hash=(1,), - num_partitions=num_partitions, - br=br, - stream=DEFAULT_STREAM, - ) - - # Spill roundtrip - spilled = spill_partitions(partitions.values(), br=br) + partitions = [ + generate_packed_data(num_elements, partition_id * 100, stream, br) + for partition_id in range(num_partitions) + ] + + spilled = spill_partitions(partitions, br=br) unspilled = unspill_partitions(spilled, br=br, allow_overbooking=False) - got = unpack_and_concat( - unspilled, - br=br, - stream=DEFAULT_STREAM, - ) - # Since the row order isn't preserved, we sort the rows by the first column. - assert_eq(expect, got, sort_rows=0) + assert len(unspilled) == num_partitions + for partition_id, partition in enumerate(unspilled): + validate_packed_data(partition, num_elements, partition_id * 100) From 1e0330b8a10979dfd88b44a4729fca74b45fcb20 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 11 Jun 2026 05:50:54 -0700 Subject: [PATCH 12/14] test: restore cudf-free streaming coverage --- cpp/tests/CMakeLists.txt | 4 +- cpp/tests/streaming/test_leaf_actor.cpp | 29 +--- .../tests/streaming/test_allgather.py | 126 ++++++------------ .../tests/streaming/test_define_actor.py | 85 ++++-------- .../rapidsmpf/tests/streaming/test_fanout.py | 114 ++++++---------- .../tests/streaming/test_leaf_actor.py | 42 ++---- .../tests/streaming/test_sparse_alltoall.py | 44 ++---- 7 files changed, 130 insertions(+), 314 deletions(-) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 8306a7854..9c007de4e 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -106,6 +106,7 @@ if(RAPIDSMPF_HAVE_STREAMING) streaming/test_channel.cpp streaming/test_error_handling.cpp streaming/test_fanout.cpp + streaming/test_leaf_actor.cpp streaming/test_lineariser.cpp streaming/test_memory_reserve_or_wait.cpp streaming/test_message.cpp @@ -123,9 +124,6 @@ if(BUILD_CUDF_TESTS) ) target_compile_definitions(test_sources PRIVATE RAPIDSMPF_HAVE_CUDF) - if(RAPIDSMPF_HAVE_STREAMING) - target_sources(test_sources PRIVATE streaming/test_leaf_actor.cpp) - endif() endif() if(RAPIDSMPF_HAVE_MPI) diff --git a/cpp/tests/streaming/test_leaf_actor.cpp b/cpp/tests/streaming/test_leaf_actor.cpp index cd5e90290..481644b68 100644 --- a/cpp/tests/streaming/test_leaf_actor.cpp +++ b/cpp/tests/streaming/test_leaf_actor.cpp @@ -4,17 +4,15 @@ */ #include +#include #include +#include +#include #include #include #include -#include -#include - -#include -#include #include #include #include @@ -23,7 +21,6 @@ #include #include -#include "../utils.hpp" #include "base_streaming_fixture.hpp" using namespace rapidsmpf; @@ -33,12 +30,11 @@ namespace actor = rapidsmpf::streaming::actor; using StreamingLeafTasks = BaseStreamingFixture; TEST_F(StreamingLeafTasks, PushAndPullChunks) { - constexpr int num_rows = 100; constexpr int num_chunks = 10; - std::vector expects; + std::vector expects; for (int i = 0; i < num_chunks; ++i) { - expects.emplace_back(random_table_with_index(i, num_rows, 0, 10)); + expects.push_back(i * 10); } std::vector actors; @@ -49,15 +45,7 @@ TEST_F(StreamingLeafTasks, PushAndPullChunks) { std::vector inputs; for (int i = 0; i < num_chunks; ++i) { inputs.emplace_back( - cudf_streaming::streaming::to_message( - i, - std::make_unique( - std::make_unique( - expects[i], stream, ctx->br()->device_mr() - ), - stream - ) - ) + i, std::make_unique(expects[i]), ContentDescription{} ); } @@ -72,10 +60,7 @@ TEST_F(StreamingLeafTasks, PushAndPullChunks) { EXPECT_EQ(expects.size(), outputs.size()); for (std::size_t i = 0; i < expects.size(); ++i) { EXPECT_EQ(outputs[i].sequence_number(), i); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT( - outputs[i].get().table_view(), - expects[i].view() - ); + EXPECT_EQ(outputs[i].release(), expects[i]); } } diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_allgather.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_allgather.py index d8d6ed019..447e5dcc7 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_allgather.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_allgather.py @@ -6,23 +6,14 @@ from contextlib import nullcontext from typing import TYPE_CHECKING -import numpy as np -import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming") -from cudf_streaming.integrations.partition import unpack_and_concat -from cudf_streaming.streaming.table_chunk import TableChunk - -from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.chunks.packed_data import PackedDataChunk from rapidsmpf.streaming.coll.allgather import AllGather, allgather from rapidsmpf.streaming.core.actor import define_actor, run_actor_network from rapidsmpf.streaming.core.leaf_actor import pull_from_channel, push_to_channel from rapidsmpf.streaming.core.message import Message -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") +from rapidsmpf.testing import generate_packed_data, validate_packed_data if TYPE_CHECKING: from collections.abc import Awaitable @@ -33,34 +24,28 @@ from rapidsmpf.streaming.core.context import Context +def _make_chunk(context: Context, num_rows: int, offset: int) -> PackedDataChunk: + stream = context.get_stream_from_pool() + return PackedDataChunk.from_packed_data( + generate_packed_data(num_rows, offset, stream, context.br()), + br=context.br(), + ) + + +def _validate_message( + context: Context, msg: Message[PackedDataChunk], num_rows: int, offset: int +) -> None: + packed = PackedDataChunk.from_message(msg, br=context.br()).to_packed_data() + validate_packed_data(packed, num_rows, offset) + + def test_allgather_actor(context: Context, comm: Communicator) -> None: if comm.nranks != 1: pytest.skip("Only support single-rank runs") num_rows = 1000 op_id = 0 - stream = context.get_stream_from_pool() - input_tables = [ - plc.Table( - [ - plc.Column.from_array( - np.arange(num_rows, dtype=np.int32) + i * num_rows, stream=stream - ) - ] - ) - for i in range(3) - ] - inputs = [ - PackedDataChunk.from_packed_data( - PackedData.from_cudf_packed_columns( # type: ignore[attr-defined] - plc.contiguous_split.pack(table, stream=stream), - stream, - context.br(), - ), - br=context.br(), - ) - for table in input_tables - ] + inputs = [_make_chunk(context, num_rows, i * num_rows) for i in range(3)] actors = [] ch1: Channel[PackedDataChunk] = context.create_channel() @@ -77,18 +62,11 @@ def test_allgather_actor(context: Context, comm: Communicator) -> None: actors.append(actor) run_actor_network(context, actors=actors) - result = unpack_and_concat( - ( - PackedDataChunk.from_message(msg, br=context.br()).to_packed_data() - for msg in deferred.release() - ), - stream, - context.br(), - ) - - expect = plc.concatenate.concatenate(input_tables, stream=stream) - stream.synchronize() - assert_eq(result, expect) + results = deferred.release() + assert len(results) == len(inputs) + for i, msg in enumerate(results): + assert msg.sequence_number == i + _validate_message(context, msg, num_rows, i * num_rows) @define_actor() @@ -96,35 +74,16 @@ async def generate_inputs( context: Context, ch: Channel[PackedDataChunk], num_rows: int, num_chunks: int ) -> None: for i in range(num_chunks): - stream = context.get_stream_from_pool() - table = plc.Table( - [ - plc.Column.from_array( - np.arange(num_rows, dtype=np.int32) + i * num_rows, stream=stream - ) - ] - ) - msg = Message( - i, - PackedDataChunk.from_packed_data( - PackedData.from_cudf_packed_columns( # type: ignore[attr-defined] - plc.contiguous_split.pack(table, stream=stream), - stream, - context.br(), - ), - br=context.br(), - ), - ) - await ch.send(context, msg) + await ch.send(context, Message(i, _make_chunk(context, num_rows, i * num_rows))) await ch.drain(context) @define_actor() -async def allgather_and_concat( +async def allgather_and_forward( context: Context, comm: Communicator, ch_in: Channel[PackedDataChunk], - ch_out: Channel[TableChunk], + ch_out: Channel[PackedDataChunk], op_id: int, use_context_manager: bool, # noqa: FBT001 ) -> None: @@ -137,12 +96,14 @@ async def allgather_and_concat( if not use_context_manager: gather.insert_finished() gathered = await gather.extract_all(context, ordered=True) - stream = context.get_stream_from_pool() - table = unpack_and_concat(gathered, stream, context.br()) - to_send = TableChunk.from_pylibcudf_table( - table, stream, exclusive_view=True, br=context.br() - ) - await ch_out.send(context, Message(0, to_send)) + for sequence, packed in enumerate(gathered): + await ch_out.send( + context, + Message( + sequence, + PackedDataChunk.from_packed_data(packed, br=context.br()), + ), + ) await ch_out.drain(context) @@ -158,29 +119,22 @@ def test_allgather_object_interface( pytest.skip("Only support single-rank runs") ch_in: Channel[PackedDataChunk] = context.create_channel() - ch_out: Channel[TableChunk] = context.create_channel() + ch_out: Channel[PackedDataChunk] = context.create_channel() actors: list[CppActor | Awaitable[None]] = [] num_rows = 100 num_chunks = 10 op_id = 0 actors.append(generate_inputs(context, ch_in, num_rows, num_chunks)) actors.append( - allgather_and_concat(context, comm, ch_in, ch_out, op_id, use_context_manager) + allgather_and_forward(context, comm, ch_in, ch_out, op_id, use_context_manager) ) actor, deferred = pull_from_channel(context, ch_out) actors.append(actor) run_actor_network(context, actors=actors) - (result_msg,) = deferred.release() - result = TableChunk.from_message(result_msg, br=context.br()) - expect = plc.Table( - [ - plc.Column.from_array( - np.arange(num_rows * num_chunks, dtype=np.int32), stream=result.stream - ) - ] - ) - got = result.table_view() - result.stream.synchronize() - assert_eq(expect, got) + results = deferred.release() + assert len(results) == num_chunks + for i, msg in enumerate(results): + assert msg.sequence_number == i + _validate_message(context, msg, num_rows, i * num_rows) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_define_actor.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_define_actor.py index 955fc0fb8..55e67a7d7 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_define_actor.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_define_actor.py @@ -5,64 +5,35 @@ from typing import TYPE_CHECKING -import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming") -from cudf_streaming.streaming.table_chunk import TableChunk - from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk from rapidsmpf.streaming.core.actor import define_actor, run_actor_network from rapidsmpf.streaming.core.leaf_actor import pull_from_channel, push_to_channel from rapidsmpf.streaming.core.message import Message -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") @pytest.fixture -def expects() -> list[plc.Table]: - return [ - plc.Table( - [ - plc.Column.from_iterable_of_py( - [1 * seq, 2 * seq, 3 * seq], plc.DataType(plc.TypeId.INT64) - ) - ] - ) - for seq in range(10) - ] +def expects() -> list[tuple[int, int, int]]: + return [(seq, seq * 2, seq * 3) for seq in range(10)] if TYPE_CHECKING: - from rmm.pylibrmm.stream import Stream - from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context -def test_send_table_chunks( - context: Context, stream: Stream, expects: list[plc.Table] +def test_send_arbitrary_chunks( + context: Context, expects: list[tuple[int, int, int]] ) -> None: - ch1: Channel[TableChunk] = context.create_channel() + ch1: Channel[ArbitraryChunk[tuple[int, int, int]]] = context.create_channel() - # The actor access `ch1` both through the `ch_out` parameter and the closure. + # The actor accesses `ch1` both through the `ch_out` parameter and the closure. @define_actor(extra_channels=(ch1,)) async def actor1(ctx: Context, /, ch_out: Channel) -> None: for seq, chunk in enumerate(expects): - await ch1.send( - context, - Message( - seq, - TableChunk.from_pylibcudf_table( - table=chunk, - stream=stream, - exclusive_view=False, - br=context.br(), - ), - ), - ) - await ch_out.drain(context) + await ch1.send(ctx, Message(seq, ArbitraryChunk(chunk))) + await ch_out.drain(ctx) actor2, output = pull_from_channel(context, ch_in=ch1) @@ -77,18 +48,17 @@ async def actor1(ctx: Context, /, ch_out: Channel) -> None: results = output.release() for seq, (result, expect) in enumerate(zip(results, expects, strict=True)): assert result.sequence_number == seq - tbl = TableChunk.from_message(result, br=context.br()) - assert_eq(tbl.table_view(), expect) + assert ArbitraryChunk.from_message(result).release() == expect def test_shutdown(context: Context) -> None: @define_actor() - async def actor1(ctx: Context, ch_out: Channel[TableChunk]) -> None: + async def actor1(ctx: Context, ch_out: Channel[ArbitraryChunk[int]]) -> None: await ch_out.shutdown(ctx) # Calling shutdown multiple times is allowed. await ch_out.shutdown(ctx) - ch1: Channel[TableChunk] = context.create_channel() + ch1: Channel[ArbitraryChunk[int]] = context.create_channel() actor2, output = pull_from_channel(context, ch_in=ch1) run_actor_network( @@ -104,10 +74,10 @@ async def actor1(ctx: Context, ch_out: Channel[TableChunk]) -> None: def test_send_error(context: Context) -> None: @define_actor() - async def actor1(ctx: Context, ch_out: Channel[TableChunk]) -> None: + async def actor1(ctx: Context, ch_out: Channel[ArbitraryChunk[int]]) -> None: raise RuntimeError("MyError") - ch1: Channel[TableChunk] = context.create_channel() + ch1: Channel[ArbitraryChunk[int]] = context.create_channel() actor2, output = pull_from_channel(context, ch_in=ch1) with pytest.RaisesGroup( @@ -127,43 +97,38 @@ async def actor1(ctx: Context, ch_out: Channel[TableChunk]) -> None: assert output.release() == [] -def test_recv_table_chunks( - context: Context, stream: Stream, expects: list[plc.Table] +def test_recv_arbitrary_chunks( + context: Context, expects: list[tuple[int, int, int]] ) -> None: - table_chunks = [ - Message( - seq, - TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=False, br=context.br() - ), - ) - for seq, expect in enumerate(expects) + chunks = [ + Message(seq, ArbitraryChunk(expect)) for seq, expect in enumerate(expects) ] - results: list[Message[TableChunk]] = [] + results: list[Message[ArbitraryChunk[tuple[int, int, int]]]] = [] @define_actor() - async def actor1(ctx: Context, ch_in: Channel[TableChunk]) -> None: + async def actor1( + ctx: Context, ch_in: Channel[ArbitraryChunk[tuple[int, int, int]]] + ) -> None: while True: - chunk = await ch_in.recv(context) + chunk = await ch_in.recv(ctx) if chunk is None: break results.append(chunk) - ch1: Channel[TableChunk] = context.create_channel() + ch1: Channel[ArbitraryChunk[tuple[int, int, int]]] = context.create_channel() run_actor_network( context, actors=[ - push_to_channel(context, ch_out=ch1, messages=table_chunks), + push_to_channel(context, ch_out=ch1, messages=chunks), actor1(context, ch_in=ch1), ], ) for seq, (result, expect) in enumerate(zip(results, expects, strict=True)): assert result.sequence_number == seq - tbl = TableChunk.from_message(result, br=context.br()) - assert_eq(tbl.table_view(), expect) + assert ArbitraryChunk.from_message(result).release() == expect @pytest.mark.filterwarnings("error") diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_fanout.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_fanout.py index c41b641f8..8253589be 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_fanout.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_fanout.py @@ -7,101 +7,85 @@ from typing import TYPE_CHECKING -import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming") -from cudf_streaming.streaming.table_chunk import TableChunk - +from rapidsmpf.streaming.chunks.packed_data import PackedDataChunk from rapidsmpf.streaming.core.actor import run_actor_network from rapidsmpf.streaming.core.fanout import FanoutPolicy, fanout from rapidsmpf.streaming.core.leaf_actor import pull_from_channel, push_to_channel from rapidsmpf.streaming.core.message import Message -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") +from rapidsmpf.testing import generate_packed_data, validate_packed_data -_INT64 = plc.DataType(plc.TypeId.INT64) +if TYPE_CHECKING: + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context -def _ab_table(i: int) -> plc.Table: - return plc.Table( - [ - plc.Column.from_iterable_of_py( - [i, i + 1, i + 2], plc.DataType(plc.TypeId.INT64) - ), - plc.Column.from_iterable_of_py( - [i * 10, i * 10 + 1, i * 10 + 2], plc.DataType(plc.TypeId.INT64) - ), - ] +def _message( + context: Context, sequence_number: int, n_elements: int = 3 +) -> Message[PackedDataChunk]: + stream = context.get_stream_from_pool() + chunk = PackedDataChunk.from_packed_data( + generate_packed_data( + n_elements, + sequence_number * 10, + stream, + context.br(), + ), + br=context.br(), ) + return Message(sequence_number, chunk) -if TYPE_CHECKING: - from rmm.pylibrmm.stream import Stream - - from rapidsmpf.streaming.core.channel import Channel - from rapidsmpf.streaming.core.context import Context +def _validate( + context: Context, + msg: Message[PackedDataChunk], + sequence_number: int, + n_elements: int = 3, +) -> None: + assert msg.sequence_number == sequence_number + packed = PackedDataChunk.from_message(msg, br=context.br()).to_packed_data() + validate_packed_data(packed, n_elements, sequence_number * 10) @pytest.mark.parametrize("policy", [FanoutPolicy.BOUNDED, FanoutPolicy.UNBOUNDED]) -def test_fanout_basic(context: Context, stream: Stream, policy: FanoutPolicy) -> None: +def test_fanout_basic(context: Context, policy: FanoutPolicy) -> None: """Test basic fanout functionality with multiple output channels.""" - # Create channels - ch_in: Channel[TableChunk] = context.create_channel() - ch_out1: Channel[TableChunk] = context.create_channel() - ch_out2: Channel[TableChunk] = context.create_channel() + ch_in: Channel[PackedDataChunk] = context.create_channel() + ch_out1: Channel[PackedDataChunk] = context.create_channel() + ch_out2: Channel[PackedDataChunk] = context.create_channel() - # Create test messages - messages = [] - for i in range(5): - chunk = TableChunk.from_pylibcudf_table( - _ab_table(i), stream, exclusive_view=False, br=context.br() - ) - messages.append(Message(i, chunk)) + messages = [_message(context, i) for i in range(5)] - # Create actors push_actor = push_to_channel(context, ch_in, messages) fanout_actor = fanout(context, ch_in, [ch_out1, ch_out2], policy) pull_actor1, output1 = pull_from_channel(context, ch_out1) pull_actor2, output2 = pull_from_channel(context, ch_out2) - # Run pipeline run_actor_network( context, actors=[push_actor, fanout_actor, pull_actor1, pull_actor2], ) - # Verify results results1 = output1.release() results2 = output2.release() assert len(results1) == 5, f"Expected 5 messages in output1, got {len(results1)}" assert len(results2) == 5, f"Expected 5 messages in output2, got {len(results2)}" - # Check that both outputs received the same sequence numbers and data for i in range(5): - assert results1[i].sequence_number == i - assert results2[i].sequence_number == i - - chunk1 = TableChunk.from_message(results1[i], br=context.br()) - chunk2 = TableChunk.from_message(results2[i], br=context.br()) - - # Verify data is correct - expected_table = _ab_table(i) - assert_eq(chunk1.table_view(), expected_table) - assert_eq(chunk2.table_view(), expected_table) + _validate(context, results1[i], i) + _validate(context, results2[i], i) @pytest.mark.parametrize("num_outputs", [1, 3, 5]) @pytest.mark.parametrize("policy", [FanoutPolicy.BOUNDED, FanoutPolicy.UNBOUNDED]) def test_fanout_multiple_outputs( - context: Context, stream: Stream, num_outputs: int, policy: FanoutPolicy + context: Context, num_outputs: int, policy: FanoutPolicy ) -> None: """Test fanout with varying numbers of output channels.""" - # Create channels - ch_in: Channel[TableChunk] = context.create_channel() - chs_out: list[Channel[TableChunk]] = [ + ch_in: Channel[PackedDataChunk] = context.create_channel() + chs_out: list[Channel[PackedDataChunk]] = [ context.create_channel() for _ in range(num_outputs) ] @@ -110,22 +94,8 @@ def test_fanout_multiple_outputs( fanout(context, ch_in, chs_out, policy) return - # Create test messages - messages = [] - for i in range(3): - table = plc.Table( - [ - plc.Column.from_iterable_of_py( - [i * 10, i * 10 + 1], plc.DataType(plc.TypeId.INT64) - ), - ] - ) - chunk = TableChunk.from_pylibcudf_table( - table, stream, exclusive_view=False, br=context.br() - ) - messages.append(Message(i, chunk)) + messages = [_message(context, i, n_elements=2) for i in range(3)] - # Create actors push_actor = push_to_channel(context, ch_in, messages) fanout_actor = fanout(context, ch_in, chs_out, policy) pull_actors = [] @@ -135,25 +105,23 @@ def test_fanout_multiple_outputs( pull_actors.append(pull_actor) outputs.append(output) - # Run pipeline run_actor_network( context, actors=[push_actor, fanout_actor, *pull_actors], ) - # Verify all outputs received the messages for output_idx, output in enumerate(outputs): results = output.release() assert len(results) == 3, ( f"Output {output_idx}: Expected 3 messages, got {len(results)}" ) for i in range(3): - assert results[i].sequence_number == i + _validate(context, results[i], i, n_elements=2) -def test_fanout_empty_outputs(context: Context, stream: Stream) -> None: +def test_fanout_empty_outputs(context: Context) -> None: """Test fanout with empty output list raises value error.""" - ch_in: Channel[TableChunk] = context.create_channel() + ch_in: Channel[PackedDataChunk] = context.create_channel() with pytest.raises(ValueError): fanout(context, ch_in, [], FanoutPolicy.BOUNDED) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_leaf_actor.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_leaf_actor.py index af5d36ee7..86d7fbb75 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_leaf_actor.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_leaf_actor.py @@ -5,53 +5,27 @@ from typing import TYPE_CHECKING -import pylibcudf as plc -import pytest - -pytest.importorskip("cudf_streaming") -from cudf_streaming.streaming.table_chunk import TableChunk - +from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk from rapidsmpf.streaming.core.actor import run_actor_network from rapidsmpf.streaming.core.leaf_actor import pull_from_channel, push_to_channel from rapidsmpf.streaming.core.message import Message -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") if TYPE_CHECKING: - from rmm.pylibrmm.stream import Stream - from rapidsmpf.streaming.core.channel import Channel from rapidsmpf.streaming.core.context import Context -def test_roundtrip(context: Context, stream: Stream) -> None: - expects = [ - plc.Table( - [ - plc.Column.from_iterable_of_py( - [1 * seq, 2 * seq, 3 * seq], plc.DataType(plc.TypeId.INT64) - ) - ] - ) - for seq in range(10) - ] - table_chunks = [ - Message( - seq, - TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=False, br=context.br() - ), - ) - for seq, expect in enumerate(expects) +def test_roundtrip(context: Context) -> None: + expects = [(seq, seq * 2, seq * 3) for seq in range(10)] + chunks = [ + Message(seq, ArbitraryChunk(expect)) for seq, expect in enumerate(expects) ] - ch1: Channel[TableChunk] = context.create_channel() - actor1 = push_to_channel(context, ch_out=ch1, messages=table_chunks) + ch1: Channel[ArbitraryChunk[tuple[int, int, int]]] = context.create_channel() + actor1 = push_to_channel(context, ch_out=ch1, messages=chunks) actor2, output = pull_from_channel(context, ch_in=ch1) run_actor_network(context, actors=(actor1, actor2)) results = output.release() for seq, (result, expect) in enumerate(zip(results, expects, strict=True)): assert result.sequence_number == seq - tbl = TableChunk.from_message(result, br=context.br()) - assert_eq(tbl.table_view(), expect) + assert ArbitraryChunk.from_message(result).release() == expect diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_sparse_alltoall.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_sparse_alltoall.py index db2bbe566..b1b6eb68c 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_sparse_alltoall.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_sparse_alltoall.py @@ -6,37 +6,20 @@ import asyncio from typing import TYPE_CHECKING -import numpy as np -import pylibcudf as plc import pytest -pytest.importorskip("cudf_streaming") -from cudf_streaming.integrations.partition import unpack_and_concat - -from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.coll.sparse_alltoall import SparseAlltoall -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") +from rapidsmpf.testing import generate_packed_data, validate_packed_data if TYPE_CHECKING: from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.core.context import Context -def make_packed_data(context: Context, values: np.ndarray) -> PackedData: - stream = context.get_stream_from_pool() - table = plc.Table([plc.Column.from_array(values, stream=stream)]) - return PackedData.from_cudf_packed_columns( # type: ignore[attr-defined, no-any-return] - plc.contiguous_split.pack(table, stream=stream), - stream, - context.br(), - ) - - -def unpack_table(context: Context, packed_data: PackedData) -> plc.Table: +def make_packed_data(context: Context, value: int) -> PackedData: stream = context.get_stream_from_pool() - return unpack_and_concat([packed_data], stream, context.br()) + return generate_packed_data(1, value, stream, context.br()) def test_sparse_alltoall_non_participating_ranks( @@ -64,24 +47,13 @@ def test_sparse_alltoall_non_participating_ranks( ) if comm.rank == 0: - exchange.insert(1, make_packed_data(context, np.array([11], dtype=np.int32))) - exchange.insert(1, make_packed_data(context, np.array([29], dtype=np.int32))) + exchange.insert(1, make_packed_data(context, 11)) + exchange.insert(1, make_packed_data(context, 29)) asyncio.run(exchange.insert_finished(context)) if comm.rank == 1: results = exchange.extract(0) assert len(results) == 2 - stream = context.get_stream_from_pool() - assert_eq( - unpack_table(context, results[0]), - plc.Table( - [plc.Column.from_array(np.array([11], dtype=np.int32), stream=stream)] - ), - ) - assert_eq( - unpack_table(context, results[1]), - plc.Table( - [plc.Column.from_array(np.array([29], dtype=np.int32), stream=stream)] - ), - ) + validate_packed_data(results[0], 1, 11) + validate_packed_data(results[1], 1, 29) From 8f00daaae4ba695fb160151995e476dfb3eb32bd Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 11 Jun 2026 09:52:48 -0700 Subject: [PATCH 13/14] Remove cuDF dependencies --- ci/checks/doxygen.sh | 8 +- ci/run_cpp_benchmark_smoketests.sh | 19 - ci/run_cpp_example_smoketests.sh | 11 +- ci/run_cpp_tools_smoketests.sh | 4 +- ci/test_cpp.sh | 16 - cmake/Modules/ConfigureCUDA.cmake | 2 +- cmake/thirdparty/get_cudf_streaming.cmake | 44 - .../all_cuda-129_arch-aarch64.yaml | 7 - .../all_cuda-129_arch-x86_64.yaml | 7 - .../all_cuda-132_arch-aarch64.yaml | 7 - .../all_cuda-132_arch-x86_64.yaml | 7 - conda/recipes/librapidsmpf/recipe.yaml | 10 +- cpp/CMakeLists.txt | 10 - cpp/benchmarks/CMakeLists.txt | 100 -- cpp/benchmarks/bench_pack.cpp | 324 ------ cpp/benchmarks/bench_partition.cpp | 175 --- cpp/benchmarks/bench_shuffle.cpp | 720 ------------ cpp/benchmarks/streaming/CMakeLists.txt | 38 - .../streaming/bench_streaming_shuffle.cpp | 479 -------- cpp/benchmarks/streaming/data_generator.hpp | 77 -- cpp/benchmarks/streaming/ndsh/CMakeLists.txt | 76 -- cpp/benchmarks/streaming/ndsh/bench_read.cpp | 438 ------- cpp/benchmarks/streaming/ndsh/concatenate.cpp | 88 -- cpp/benchmarks/streaming/ndsh/concatenate.hpp | 38 - cpp/benchmarks/streaming/ndsh/groupby.cpp | 78 -- cpp/benchmarks/streaming/ndsh/groupby.hpp | 51 - cpp/benchmarks/streaming/ndsh/join.cpp | 620 ---------- cpp/benchmarks/streaming/ndsh/join.hpp | 219 ---- .../streaming/ndsh/parquet_writer.cpp | 70 -- .../streaming/ndsh/parquet_writer.hpp | 34 - cpp/benchmarks/streaming/ndsh/q01.cpp | 462 -------- cpp/benchmarks/streaming/ndsh/q03.cpp | 731 ------------ cpp/benchmarks/streaming/ndsh/q04.cpp | 551 --------- cpp/benchmarks/streaming/ndsh/q09.cpp | 678 ----------- cpp/benchmarks/streaming/ndsh/q21.cpp | 1029 ----------------- cpp/benchmarks/streaming/ndsh/sort.cpp | 70 -- cpp/benchmarks/streaming/ndsh/sort.hpp | 39 - cpp/benchmarks/streaming/ndsh/sql/q01.sql | 21 - cpp/benchmarks/streaming/ndsh/sql/q03.sql | 23 - cpp/benchmarks/streaming/ndsh/sql/q04.sql | 21 - cpp/benchmarks/streaming/ndsh/sql/q09.sql | 32 - cpp/benchmarks/streaming/ndsh/sql/q17.sql | 17 - cpp/benchmarks/streaming/ndsh/sql/q18.sql | 33 - cpp/benchmarks/streaming/ndsh/sql/q21.sql | 40 - cpp/benchmarks/streaming/ndsh/utils.cpp | 427 ------- cpp/benchmarks/streaming/ndsh/utils.hpp | 323 ------ cpp/benchmarks/utils/random_data.cu | 112 -- cpp/benchmarks/utils/random_data.hpp | 117 -- cpp/compute-sanitizer-suppressions.xml | 113 -- cpp/examples/CMakeLists.txt | 28 - cpp/examples/example_shuffle.cpp | 139 --- .../rapidsmpf/bootstrap/slurm_backend.hpp | 4 +- .../rapidsmpf/memory/buffer_resource.hpp | 17 +- cpp/include/rapidsmpf/owning_wrapper.hpp | 16 +- .../rapidsmpf/shuffler/finish_counter.hpp | 2 +- cpp/scripts/ndsh.py | 879 -------------- cpp/src/memory/host_memory_resource.cpp | 1 - cpp/tests/CMakeLists.txt | 10 - cpp/tests/test_partition.cpp | 103 -- cpp/tests/utils.hpp | 59 - cpp/tools/rrun.cpp | 4 +- dependencies.yaml | 49 - .../source/background/shuffle-architecture.md | 2 +- docs/source/cpp/index.md | 11 +- docs/source/getting-started.md | 4 +- docs/source/python/index.md | 2 +- docs/source/python/quickstart.md | 12 +- python/rapidsmpf/pyproject.toml | 1 - .../rapidsmpf/benchmarks/__init__.py | 3 - .../benchmarks/streaming_benchmark.py | 388 ------- .../rapidsmpf/examples/bulk_mpi_shuffle.py | 524 --------- .../rapidsmpf/examples/ray/__init__.py | 3 - .../examples/ray/bulk_ray_shuffle.py | 451 -------- .../examples/ray/ray_shuffle_example.py | 195 ---- .../rapidsmpf/examples/streaming/__init__.py | 3 - .../examples/streaming/basic_example.py | 149 --- .../rapidsmpf/memory/memory_reservation.pyx | 4 +- python/rapidsmpf/rapidsmpf/shuffler.pyx | 4 +- .../rapidsmpf/streaming/core/actor.pyx | 10 +- python/rapidsmpf/rapidsmpf/testing.py | 53 - .../tests/streaming/test_examples.py | 13 - .../rapidsmpf/tests/test_examples.py | 119 -- .../rapidsmpf/tests/test_partition.py | 107 -- python/rapidsmpf/rapidsmpf/tests/test_ray.py | 41 - 84 files changed, 38 insertions(+), 11988 deletions(-) delete mode 100644 cmake/thirdparty/get_cudf_streaming.cmake delete mode 100644 cpp/benchmarks/bench_pack.cpp delete mode 100644 cpp/benchmarks/bench_partition.cpp delete mode 100644 cpp/benchmarks/bench_shuffle.cpp delete mode 100644 cpp/benchmarks/streaming/CMakeLists.txt delete mode 100644 cpp/benchmarks/streaming/bench_streaming_shuffle.cpp delete mode 100644 cpp/benchmarks/streaming/data_generator.hpp delete mode 100644 cpp/benchmarks/streaming/ndsh/CMakeLists.txt delete mode 100644 cpp/benchmarks/streaming/ndsh/bench_read.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/concatenate.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/concatenate.hpp delete mode 100644 cpp/benchmarks/streaming/ndsh/groupby.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/groupby.hpp delete mode 100644 cpp/benchmarks/streaming/ndsh/join.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/join.hpp delete mode 100644 cpp/benchmarks/streaming/ndsh/parquet_writer.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/parquet_writer.hpp delete mode 100644 cpp/benchmarks/streaming/ndsh/q01.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/q03.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/q04.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/q09.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/q21.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/sort.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/sort.hpp delete mode 100644 cpp/benchmarks/streaming/ndsh/sql/q01.sql delete mode 100644 cpp/benchmarks/streaming/ndsh/sql/q03.sql delete mode 100644 cpp/benchmarks/streaming/ndsh/sql/q04.sql delete mode 100644 cpp/benchmarks/streaming/ndsh/sql/q09.sql delete mode 100644 cpp/benchmarks/streaming/ndsh/sql/q17.sql delete mode 100644 cpp/benchmarks/streaming/ndsh/sql/q18.sql delete mode 100644 cpp/benchmarks/streaming/ndsh/sql/q21.sql delete mode 100644 cpp/benchmarks/streaming/ndsh/utils.cpp delete mode 100644 cpp/benchmarks/streaming/ndsh/utils.hpp delete mode 100644 cpp/benchmarks/utils/random_data.cu delete mode 100644 cpp/benchmarks/utils/random_data.hpp delete mode 100644 cpp/examples/example_shuffle.cpp delete mode 100755 cpp/scripts/ndsh.py delete mode 100644 cpp/tests/test_partition.cpp delete mode 100644 python/rapidsmpf/rapidsmpf/benchmarks/__init__.py delete mode 100644 python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py delete mode 100644 python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py delete mode 100644 python/rapidsmpf/rapidsmpf/examples/ray/__init__.py delete mode 100644 python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py delete mode 100644 python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py delete mode 100644 python/rapidsmpf/rapidsmpf/examples/streaming/__init__.py delete mode 100644 python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py delete mode 100644 python/rapidsmpf/rapidsmpf/tests/streaming/test_examples.py delete mode 100644 python/rapidsmpf/rapidsmpf/tests/test_examples.py delete mode 100644 python/rapidsmpf/rapidsmpf/tests/test_partition.py diff --git a/ci/checks/doxygen.sh b/ci/checks/doxygen.sh index 4b48c65ce..f70b59692 100755 --- a/ci/checks/doxygen.sh +++ b/ci/checks/doxygen.sh @@ -1,9 +1,9 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -############################### -# cuDF doxygen warnings check # -############################### +####################################### +# rapidsmpf doxygen warnings check # +####################################### # skip if doxygen is not installed if ! [ -x "$(command -v doxygen)" ]; then diff --git a/ci/run_cpp_benchmark_smoketests.sh b/ci/run_cpp_benchmark_smoketests.sh index 1e12e5e5a..651807cc3 100755 --- a/ci/run_cpp_benchmark_smoketests.sh +++ b/ci/run_cpp_benchmark_smoketests.sh @@ -15,28 +15,12 @@ export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 export OMPI_MCA_opal_cuda_support=1 # enable CUDA support in OpenMPI # Ensure that benchmarks are runnable -python "${TIMEOUT_TOOL_PATH}" 30 \ - mpirun --map-by node --bind-to none -np 3 ./bench_shuffle -m cuda - python "${TIMEOUT_TOOL_PATH}" 30 \ mpirun --map-by node --bind-to none -np 3 ./bench_comm -m cuda RAPIDSMPF_SMOKE_TEST_MODE="ON" \ python "${TIMEOUT_TOOL_PATH}" 30 ./bench_memory_resources -python "${TIMEOUT_TOOL_PATH}" 30 \ - ./bench_streaming_shuffle -m cuda - -# Ensure that shuffle benchmark with CUPTI monitor is runnable and creates the expected csv files -python "${TIMEOUT_TOOL_PATH}" 30 \ - mpirun --map-by node --bind-to none -np 3 ./bench_shuffle -m cuda -M cupti_shuffle -for i in {0..2}; do - if [[ ! -f cupti_shuffle${i}.csv ]]; then - echo "Error: cupti_shuffle${i}.csv was not created!" - exit 1 - fi -done - # Ensure that comm benchmark with CUPTI monitor is runnable and creates the expected csv files python "${TIMEOUT_TOOL_PATH}" 30 \ mpirun --map-by node --bind-to none -np 3 ./bench_comm -m cuda -M cupti_comm @@ -46,6 +30,3 @@ for i in {0..2}; do exit 1 fi done - -# bench pack smoketest (only run 1MB buffer benchmarks) -python "${TIMEOUT_TOOL_PATH}" 30 ./bench_pack --benchmark_filter="/1/" --benchmark_min_time=0s diff --git a/ci/run_cpp_example_smoketests.sh b/ci/run_cpp_example_smoketests.sh index cc196fbb2..c6bdc2b7f 100755 --- a/ci/run_cpp_example_smoketests.sh +++ b/ci/run_cpp_example_smoketests.sh @@ -1,5 +1,5 @@ #!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 set -xeuo pipefail @@ -9,15 +9,6 @@ TIMEOUT_TOOL_PATH="$(dirname "$(realpath "${BASH_SOURCE[0]}")")"/timeout_with_st # Support customizing the ctests' install location cd "${INSTALL_PREFIX:-${CONDA_PREFIX:-/usr}}/bin/examples/librapidsmpf/" -# OpenMPI specific options -export OMPI_ALLOW_RUN_AS_ROOT=1 # CI runs as root -export OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1 -export OMPI_MCA_opal_cuda_support=1 # enable CUDA support in OpenMPI - -# Ensure that shuffle example is runnable -python "${TIMEOUT_TOOL_PATH}" 30 \ - mpirun --map-by node --bind-to none -np 2 ./example_shuffle - # Ensure that cupti monitor example is runnable and creates the expected csv file python "${TIMEOUT_TOOL_PATH}" 30 ./example_cupti_monitor if [[ ! -f cupti_monitor_example.csv ]]; then diff --git a/ci/run_cpp_tools_smoketests.sh b/ci/run_cpp_tools_smoketests.sh index f6437b02b..852c00f49 100755 --- a/ci/run_cpp_tools_smoketests.sh +++ b/ci/run_cpp_tools_smoketests.sh @@ -29,8 +29,6 @@ python "${TIMEOUT_TOOL_PATH}" 30 \ python "${TIMEOUT_TOOL_PATH}" 30 \ rrun --tag-output -n 3 -g 0,0,0 ./bench_comm -m cuda -C ucxx python "${TIMEOUT_TOOL_PATH}" 30 \ - rrun --tag-output -n 3 -g 0,0,0 ./bench_shuffle -m cuda -C ucxx -python "${TIMEOUT_TOOL_PATH}" 30 \ - rrun --tag-output -n 1 -g 0 ./bench_streaming_shuffle -m cuda -C ucxx + rrun --tag-output -n 1 -g 0 -x RAPIDSMPF_SMOKE_TEST_MODE=ON ./bench_memory_resources topology_discovery | python "${VALIDATE_TOPOLOGY_PATH}" - diff --git a/ci/test_cpp.sh b/ci/test_cpp.sh index 8c5046591..326ac4095 100755 --- a/ci/test_cpp.sh +++ b/ci/test_cpp.sh @@ -67,21 +67,5 @@ rapids-logger "Run tools smoketests" rapids-logger "Run rrun gtests" ./run_rrun_tests.sh -BENCHMARKS_DIR=$CONDA_PREFIX/bin/benchmarks/librapidsmpf - -rapids-logger "Run NDSH benchmarks" -python ../cpp/scripts/ndsh.py run \ - --input-dir scale-1/ \ - --output-dir validation/ \ - --generate-data \ - --benchmark-dir "${BENCHMARKS_DIR}" \ - --benchmark-args='--no-pinned-host-memory' - -rapids-logger "Validate NDSH benchmarks" -python ../cpp/scripts/ndsh.py validate \ - --results-path validation/output \ - --expected-path validation/expected \ - --ignore-timezone - rapids-logger "Test script exiting with exit code: $EXITCODE" exit ${EXITCODE} diff --git a/cmake/Modules/ConfigureCUDA.cmake b/cmake/Modules/ConfigureCUDA.cmake index eb8b5b92f..a160fbd15 100644 --- a/cmake/Modules/ConfigureCUDA.cmake +++ b/cmake/Modules/ConfigureCUDA.cmake @@ -53,6 +53,6 @@ endif() # Debug options if(CMAKE_BUILD_TYPE MATCHES Debug) - message(VERBOSE "CUDF: Building with debugging flags") + message(VERBOSE "RAPIDSMPF: Building with debugging flags") list(APPEND RAPIDSMPF_CUDA_FLAGS -Xcompiler=-rdynamic) endif() diff --git a/cmake/thirdparty/get_cudf_streaming.cmake b/cmake/thirdparty/get_cudf_streaming.cmake deleted file mode 100644 index 0ae1c676e..000000000 --- a/cmake/thirdparty/get_cudf_streaming.cmake +++ /dev/null @@ -1,44 +0,0 @@ -# ================================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ================================================================================= - -# This function finds cudf_streaming for test/benchmark use only. It does NOT add cudf_streaming to -# any rapidsmpf export set. -# -# When BUILD_TESTS is enabled (in addition to BUILD_CUDF_TESTS, which gates inclusion of this file), -# the cudf `testing` component is also imported so that tests can link against `cudf::cudftestutil` -# and `cudf::cudftestutil_impl`. The default `find_dependency(cudf)` triggered transitively by -# cudf_streaming does not request the `testing` component, so we request it explicitly here. -# Benchmarks and examples only need `cudf_streaming::cudf_streaming` and do not pay for the testing -# component. This is a temporary measure until BUILD_CUDF_TESTS is removed entirely. -function(find_and_configure_cudf_streaming) - - set(oneValueArgs VERSION GIT_REPO GIT_TAG) - cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - if(NOT TARGET cudf_streaming::cudf_streaming) - rapids_cpm_find( - cudf_streaming ${PKG_VERSION} - GLOBAL_TARGETS cudf_streaming::cudf_streaming - CPM_ARGS - GIT_REPOSITORY ${PKG_GIT_REPO} - GIT_TAG ${PKG_GIT_TAG} - GIT_SHALLOW TRUE SOURCE_SUBDIR cpp/libcudf_streaming - OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_EXAMPLES OFF" - ) - endif() - - # Only the cudf-dependent tests link against cudf::cudftestutil{,_impl}; skip this when tests are - # disabled (e.g. benchmarks/examples-only builds). - if(BUILD_TESTS AND NOT TARGET cudf::cudftestutil) - find_package(cudf ${PKG_VERSION} REQUIRED COMPONENTS testing) - endif() -endfunction() - -find_and_configure_cudf_streaming( - VERSION ${RAPIDS_VERSION} GIT_REPO https://github.com/rapidsai/cudf.git GIT_TAG - "${RAPIDS_BRANCH}" -) diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 5184df872..31a729fb3 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -24,11 +24,9 @@ dependencies: - cxx-compiler - cython>=3.2.2 - doxygen=1.9.1 -- duckdb - gcc_linux-aarch64=14.* - gdb - ipython -- libcudf-streaming==26.8.*,>=0.0.0a0 - libnuma - libpmix-devel >=5.0,<6.0 - librmm==26.8.*,>=0.0.0a0 @@ -38,14 +36,11 @@ dependencies: - myst-nb - myst-parser - ninja -- numpy - numpy >=1.23,<3.0 - numpydoc - openmpi >=5.0 -- pip - pre-commit - psutil -- pyarrow - pydata-sphinx-theme>=0.15.4 - pytest - python>=3.11 @@ -59,6 +54,4 @@ dependencies: - sysroot_linux-aarch64=2.28 - ucxx==0.51.*,>=0.0.0a0 - valgrind -- pip: - - tpchgen-cli name: all_cuda-129_arch-aarch64 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 12271fda4..55aa23ac0 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -24,11 +24,9 @@ dependencies: - cxx-compiler - cython>=3.2.2 - doxygen=1.9.1 -- duckdb - gcc_linux-64=14.* - gdb - ipython -- libcudf-streaming==26.8.*,>=0.0.0a0 - libnuma - libpmix-devel >=5.0,<6.0 - librmm==26.8.*,>=0.0.0a0 @@ -38,14 +36,11 @@ dependencies: - myst-nb - myst-parser - ninja -- numpy - numpy >=1.23,<3.0 - numpydoc - openmpi >=5.0 -- pip - pre-commit - psutil -- pyarrow - pydata-sphinx-theme>=0.15.4 - pytest - python>=3.11 @@ -59,6 +54,4 @@ dependencies: - sysroot_linux-64=2.28 - ucxx==0.51.*,>=0.0.0a0 - valgrind -- pip: - - tpchgen-cli name: all_cuda-129_arch-x86_64 diff --git a/conda/environments/all_cuda-132_arch-aarch64.yaml b/conda/environments/all_cuda-132_arch-aarch64.yaml index ea4435c41..d95a1d818 100644 --- a/conda/environments/all_cuda-132_arch-aarch64.yaml +++ b/conda/environments/all_cuda-132_arch-aarch64.yaml @@ -24,11 +24,9 @@ dependencies: - cxx-compiler - cython>=3.2.2 - doxygen=1.9.1 -- duckdb - gcc_linux-aarch64=14.* - gdb - ipython -- libcudf-streaming==26.8.*,>=0.0.0a0 - libnuma - libpmix-devel >=5.0,<6.0 - librmm==26.8.*,>=0.0.0a0 @@ -38,14 +36,11 @@ dependencies: - myst-nb - myst-parser - ninja -- numpy - numpy >=1.23,<3.0 - numpydoc - openmpi >=5.0 -- pip - pre-commit - psutil -- pyarrow - pydata-sphinx-theme>=0.15.4 - pytest - python>=3.11 @@ -59,6 +54,4 @@ dependencies: - sysroot_linux-aarch64=2.28 - ucxx==0.51.*,>=0.0.0a0 - valgrind -- pip: - - tpchgen-cli name: all_cuda-132_arch-aarch64 diff --git a/conda/environments/all_cuda-132_arch-x86_64.yaml b/conda/environments/all_cuda-132_arch-x86_64.yaml index f5cde57c7..f6c791511 100644 --- a/conda/environments/all_cuda-132_arch-x86_64.yaml +++ b/conda/environments/all_cuda-132_arch-x86_64.yaml @@ -24,11 +24,9 @@ dependencies: - cxx-compiler - cython>=3.2.2 - doxygen=1.9.1 -- duckdb - gcc_linux-64=14.* - gdb - ipython -- libcudf-streaming==26.8.*,>=0.0.0a0 - libnuma - libpmix-devel >=5.0,<6.0 - librmm==26.8.*,>=0.0.0a0 @@ -38,14 +36,11 @@ dependencies: - myst-nb - myst-parser - ninja -- numpy - numpy >=1.23,<3.0 - numpydoc - openmpi >=5.0 -- pip - pre-commit - psutil -- pyarrow - pydata-sphinx-theme>=0.15.4 - pytest - python>=3.11 @@ -59,6 +54,4 @@ dependencies: - sysroot_linux-64=2.28 - ucxx==0.51.*,>=0.0.0a0 - valgrind -- pip: - - tpchgen-cli name: all_cuda-132_arch-x86_64 diff --git a/conda/recipes/librapidsmpf/recipe.yaml b/conda/recipes/librapidsmpf/recipe.yaml index 111af0c7e..4c5c44bd6 100644 --- a/conda/recipes/librapidsmpf/recipe.yaml +++ b/conda/recipes/librapidsmpf/recipe.yaml @@ -29,7 +29,7 @@ cache: export CXXFLAGS=$(echo $CXXFLAGS | sed -E 's@\-fdebug\-prefix\-map[^ ]*@@g') set +x - ./build.sh -v -n --cmake-args="\"-DBUILD_CUPTI_SUPPORT=ON -DBUILD_CUDF_TESTS=ON\"" --no-clang-tidy librapidsmpf + ./build.sh -v -n --cmake-args="\"-DBUILD_CUPTI_SUPPORT=ON\"" --no-clang-tidy librapidsmpf secrets: - AWS_ACCESS_KEY_ID - AWS_SECRET_ACCESS_KEY @@ -78,8 +78,6 @@ cache: - cuda-cupti-dev - cuda-nvml-dev - libnuma - - libcudf =${{ minor_version }} - - libcudf-streaming =${{ minor_version }} - libpmix-devel >=5.0,<6.0 - librmm =${{ minor_version }} - libucxx ${{ ucxx_version }} @@ -132,8 +130,6 @@ outputs: - ${{ stdlib("c") }} by_name: - cuda-cupti - - libcudf - - libcudf-streaming - librmm - openmpi about: @@ -160,8 +156,6 @@ outputs: - cuda-cupti-dev - cuda-nvml-dev - cuda-version =${{ cuda_version }} - - libcudf =${{ minor_version }} - - libcudf-streaming =${{ minor_version }} - librmm =${{ minor_version }} - openmpi >=5.0 - libnuma @@ -169,8 +163,6 @@ outputs: run: - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} - cuda-cupti - - libcudf =${{ minor_version }} - - libcudf-streaming =${{ minor_version }} - librmm =${{ minor_version }} - openmpi >=5.0 # See - libucxx ${{ ucxx_version }} diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index d9b7cc8e4..131c75bb8 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -54,7 +54,6 @@ option(BUILD_SHARED_LIBS "Build RapidsMPF shared library" ON) option(RAPIDSMPF_CLANG_TIDY "Enable clang-tidy during compilation" OFF) option(RAPIDSMPF_ASAN "Enable AddressSanitizer" OFF) option(RAPIDSMPF_VERBOSE_INFO "Enable detail mode" OFF) -option(BUILD_CUDF_TESTS "Build tests/benchmarks that require cudf" OFF) message(STATUS "librapidsmpf build options:") message(STATUS " BUILD_MPI_SUPPORT : ${BUILD_MPI_SUPPORT}") @@ -70,7 +69,6 @@ message(STATUS " BUILD_SHARED_LIBS : ${BUILD_SHARED_LIBS}") message(STATUS " RAPIDSMPF_CLANG_TIDY : ${RAPIDSMPF_CLANG_TIDY}") message(STATUS " RAPIDSMPF_ASAN : ${RAPIDSMPF_ASAN}") message(STATUS " RAPIDSMPF_VERBOSE_INFO : ${RAPIDSMPF_VERBOSE_INFO}") -message(STATUS " BUILD_CUDF_TESTS : ${BUILD_CUDF_TESTS}") # Copy options to our prefix to prevent upstream projects from modifying them. set(RAPIDSMPF_HAVE_MPI ${BUILD_MPI_SUPPORT}) @@ -336,18 +334,10 @@ target_compile_definitions( rapids_cuda_set_runtime(rapidsmpf USE_STATIC ON) -# Guard: when BUILD_CUDF_TESTS is enabled, cudf_streaming is fetched via CPM and its own -# cmake/thirdparty/get_rapidsmpf.cmake calls rapids_cpm_find(rapidsmpf ...) which creates this alias -# target. Without the guard we get "cannot create ALIAS target because another target with the same -# name already exists" due to the circular dependency (rapidsmpf -> cudf_streaming -> rapidsmpf). if(NOT TARGET rapidsmpf::rapidsmpf) add_library(rapidsmpf::rapidsmpf ALIAS rapidsmpf) endif() -if(BUILD_CUDF_TESTS) - include(../cmake/thirdparty/get_cudf_streaming.cmake) -endif() - # ################################################################################################## # * linter configuration --------------------------------------------------------------------------- if(RAPIDSMPF_CLANG_TIDY) diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt index 5b4d715e8..25c4084a4 100644 --- a/cpp/benchmarks/CMakeLists.txt +++ b/cpp/benchmarks/CMakeLists.txt @@ -60,103 +60,3 @@ install( DESTINATION bin/benchmarks/librapidsmpf EXCLUDE_FROM_ALL ) - -if(BUILD_CUDF_TESTS) - add_library(bench_utils INTERFACE) - target_sources(bench_utils INTERFACE utils/random_data.cu) - - add_executable(bench_shuffle "bench_shuffle.cpp") - set_target_properties( - bench_shuffle - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - CXX_EXTENSIONS ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON - ) - target_compile_options( - bench_shuffle PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" - ) - target_link_libraries( - bench_shuffle - PRIVATE rapidsmpf::rapidsmpf - ucxx::ucxx - cudf_streaming::cudf_streaming - $ - $ - maybe_asan - bench_utils - ) - install( - TARGETS bench_shuffle - COMPONENT benchmarking - DESTINATION bin/benchmarks/librapidsmpf - EXCLUDE_FROM_ALL - ) - - add_executable(bench_partition "bench_partition.cpp") - set_target_properties( - bench_partition - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - CXX_EXTENSIONS ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON - ) - target_compile_options( - bench_partition PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" - ) - target_link_libraries( - bench_partition - PRIVATE rapidsmpf::rapidsmpf - ucxx::ucxx - cudf_streaming::cudf_streaming - benchmark::benchmark - benchmark::benchmark_main - $ - $ - maybe_asan - bench_utils - ) - install( - TARGETS bench_partition - COMPONENT benchmarking - DESTINATION bin/benchmarks/librapidsmpf - EXCLUDE_FROM_ALL - ) - - add_executable(bench_pack "bench_pack.cpp") - set_target_properties( - bench_pack - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - CXX_EXTENSIONS ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON - LINK_FLAGS "-Wl,--allow-shlib-undefined" - ) - target_compile_options( - bench_pack PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" - ) - target_link_libraries( - bench_pack - PRIVATE rapidsmpf::rapidsmpf cudf_streaming::cudf_streaming benchmark::benchmark - benchmark::benchmark_main $ maybe_asan bench_utils - ) - install( - TARGETS bench_pack - COMPONENT benchmarking - DESTINATION bin/benchmarks/librapidsmpf - EXCLUDE_FROM_ALL - ) -endif() - -if(BUILD_CUDF_TESTS AND RAPIDSMPF_HAVE_STREAMING) - add_subdirectory(streaming) -endif() diff --git a/cpp/benchmarks/bench_pack.cpp b/cpp/benchmarks/bench_pack.cpp deleted file mode 100644 index 0494044d0..000000000 --- a/cpp/benchmarks/bench_pack.cpp +++ /dev/null @@ -1,324 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include "utils/random_data.hpp" - -constexpr std::size_t MB = 1024 * 1024; - -/** - * @brief Runs the cudf::pack benchmark - * @param state The benchmark state - * @param table_size_mb The size of the table in MB - * @param table_mr The memory resource for the table - * @param pack_mr The memory resource for the packed data - * @param stream The CUDA stream to use - */ -void run_pack( - benchmark::State& state, - std::size_t table_size_mb, - rmm::device_async_resource_ref table_mr, - rmm::device_async_resource_ref pack_mr, - rmm::cuda_stream_view stream -) { - auto const table_size_bytes = table_size_mb * MB; - - // Calculate number of rows for a single-column table of the desired size - auto const nrows = - rapidsmpf::safe_cast(table_size_bytes / sizeof(random_data_t)); - auto table = random_table(1, nrows, 0, 1000, stream, table_mr); - - // Warm up - auto warm_up = cudf::pack(table.view(), stream, pack_mr); - stream.synchronize(); - - for (auto _ : state) { - auto packed = cudf::pack(table.view(), stream, pack_mr); - benchmark::DoNotOptimize(packed); - stream.synchronize(); - } - - state.SetBytesProcessed( - static_cast(state.iterations()) - * static_cast(table_size_bytes) - ); - state.counters["table_size_mb"] = static_cast(table_size_mb); - state.counters["num_rows"] = nrows; -} - -/** - * @brief Benchmark for cudf::pack with device memory - */ -static void BM_Pack_device(benchmark::State& state) { - auto const table_size_mb = static_cast(state.range(0)); - - rmm::cuda_stream_view stream = rmm::cuda_stream_default; - - // Create memory resources - rmm::mr::pool_memory_resource pool_mr{ - rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) - }; - run_pack(state, table_size_mb, pool_mr, pool_mr, stream); -} - -/** - * @brief Benchmark for cudf::pack with pinned memory - */ -static void BM_Pack_pinned(benchmark::State& state) { - state.SkipWithMessage("Skipping until cudf#20886 is fixed"); - /* if (!rapidsmpf::is_pinned_memory_resources_supported()) { - state.SkipWithMessage("Pinned memory resources are not supported"); - return; - } - - auto const table_size_mb = static_cast(state.range(0)); - - rmm::cuda_stream_view stream = rmm::cuda_stream_default; - - // Create memory resources - rmm::mr::pool_memory_resource pool_mr{ - rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) - }; - - run_pack(state, table_size_mb, pool_mr, pinned_mr, stream); */ -} - -/** - * @brief Runs the cudf::chunked_pack benchmark - * @param state The benchmark state - * @param bounce_buffer_size The size of the bounce buffer in bytes - * @param table_size The size of the table in bytes - * @param table_mr The memory resource for the table - * @param pack_mr The memory resource for the packed data - * @param stream The CUDA stream to use - */ -void run_chunked_pack( - benchmark::State& state, - std::size_t bounce_buffer_size, - std::size_t table_size, - rmm::device_async_resource_ref table_mr, - rmm::device_async_resource_ref pack_mr, - rmm::cuda_stream_view stream -) { - // Calculate number of rows for a single-column table of the desired size - auto const nrows = - rapidsmpf::safe_cast(table_size / sizeof(random_data_t)); - auto table = random_table(1, nrows, 0, 1000, stream, table_mr); - - // Create the chunked_pack instance to get total output size - std::size_t total_size; - { - cudf::chunked_pack packer(table.view(), bounce_buffer_size, stream, table_mr); - total_size = packer.get_total_contiguous_size(); - } - - // Allocate bounce buffer and destination buffer using the pack_mr - rmm::device_buffer bounce_buffer(bounce_buffer_size, stream, pack_mr); - rmm::device_buffer destination(total_size, stream, pack_mr); - - auto run_packer = [&] { - cudf::chunked_pack packer(table.view(), bounce_buffer_size, stream, pack_mr); - - std::size_t offset = 0; - while (packer.has_next()) { - auto const bytes_copied = packer.next( - cudf::device_span( - static_cast(bounce_buffer.data()), bounce_buffer_size - ) - ); - RAPIDSMPF_CUDA_TRY( - rapidsmpf::cuda_memcpy_async( - static_cast(destination.data()) + offset, - bounce_buffer.data(), - bytes_copied, - stream - ) - ); - offset += bytes_copied; - } - }; - - { - run_packer(); - stream.synchronize(); - } - - for (auto _ : state) { - run_packer(); - benchmark::DoNotOptimize(destination); - stream.synchronize(); - } - - state.SetBytesProcessed( - static_cast(state.iterations()) - * static_cast(table_size) - ); - state.counters["table_size_mb"] = - static_cast(table_size) / static_cast(MB); - state.counters["num_rows"] = nrows; - state.counters["bounce_buffer_mb"] = - static_cast(bounce_buffer_size) / static_cast(MB); -} - -/** - * @brief Benchmark for cudf::chunked_pack with device memory - */ -static void BM_ChunkedPack_device(benchmark::State& state) { - auto const table_size_mb = static_cast(state.range(0)); - auto const table_size_bytes = table_size_mb * MB; - - // Bounce buffer size: max(1MB, table_size / 10) - auto const bounce_buffer_size = std::max(MB, table_size_bytes / 10); - - rmm::cuda_stream_view stream = rmm::cuda_stream_default; - - rmm::mr::pool_memory_resource pool_mr{ - rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) - }; - - run_chunked_pack( - state, bounce_buffer_size, table_size_bytes, pool_mr, pool_mr, stream - ); -} - -/** - * @brief Benchmark for cudf::chunked_pack pinned memory - */ -static void BM_ChunkedPack_pinned(benchmark::State& state) { - state.SkipWithMessage("Skipping until cudf#20886 is fixed"); - /* if (!rapidsmpf::is_pinned_memory_resources_supported()) { - state.SkipWithMessage("Pinned memory resources are not supported"); - return; - } - - auto const table_size_mb = static_cast(state.range(0)); - auto const table_size_bytes = table_size_mb * MB; - - // Bounce buffer size: max(1MB, table_size / 10) - auto const bounce_buffer_size = std::max(MB, table_size_bytes / 10); - - rmm::cuda_stream_view stream = rmm::cuda_stream_default; - - rmm::mr::pool_memory_resource pool_mr{ - rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) - }; - rapidsmpf::PinnedMemoryResource pinned_mr; - - run_chunked_pack( - state, bounce_buffer_size, table_size_bytes, pool_mr, pinned_mr, stream - ); */ -} - -// Custom argument generator for the benchmark -void PackArguments(benchmark::Benchmark* b) { - // Test different table sizes in MB (minimum 1MB as requested) - for (auto size_mb : {1, 10, 100, 500, 1000, 2000, 4000}) { - b->Args({size_mb}); - } -} - -// Register the benchmarks -BENCHMARK(BM_Pack_device) - ->Apply(PackArguments) - ->UseRealTime() - ->Unit(benchmark::kMillisecond); -BENCHMARK(BM_Pack_pinned) - ->Apply(PackArguments) - ->UseRealTime() - ->Unit(benchmark::kMillisecond); - -BENCHMARK(BM_ChunkedPack_device) - ->Apply(PackArguments) - ->UseRealTime() - ->Unit(benchmark::kMillisecond); - -BENCHMARK(BM_ChunkedPack_pinned) - ->Apply(PackArguments) - ->UseRealTime() - ->Unit(benchmark::kMillisecond); - -/** - * @brief Benchmark for cudf::chunked_pack in device memory varying the bounce buffer size - * and keeping table size fixed at 1GB - */ -static void BM_ChunkedPack_fixed_table_device(benchmark::State& state) { - auto const bounce_buffer_size = static_cast(state.range(0)) * MB; - constexpr std::size_t table_size_bytes = 1024 * MB; - - rmm::cuda_stream_view stream = rmm::cuda_stream_default; - - // Create memory resources - rmm::mr::pool_memory_resource pool_mr{ - rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) - }; - - run_chunked_pack( - state, bounce_buffer_size, table_size_bytes, pool_mr, pool_mr, stream - ); -} - -/** - * @brief Benchmark for cudf::chunked_pack in pinned memory varying the bounce buffer size - * and keeping table size fixed at 1GB - */ -static void BM_ChunkedPack_fixed_table_pinned(benchmark::State& state) { - state.SkipWithMessage("Skipping until cudf#20886 is fixed"); - /* if (!rapidsmpf::is_pinned_memory_resources_supported()) { - state.SkipWithMessage("Pinned memory resources are not supported"); - return; - } - - auto const bounce_buffer_size = static_cast(state.range(0)) * MB; - constexpr std::size_t table_size_bytes = 1024 * MB; - - rmm::cuda_stream_view stream = rmm::cuda_stream_default; - - rmm::mr::pool_memory_resource pool_mr{ - rmm::mr::cuda_async_memory_resource{}, rmm::percent_of_free_device_memory(40) - }; - rapidsmpf::PinnedMemoryResource pinned_mr; - - run_chunked_pack( - state, bounce_buffer_size, table_size_bytes, pool_mr, pinned_mr, stream - ); */ -} - -// Custom argument generator for the benchmark -void ChunkedPackArguments(benchmark::Benchmark* b) { - // Test different table sizes in MB (minimum 1MB as requested) - for (auto bounce_buf_sz_mb : {1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024}) { - b->Args({bounce_buf_sz_mb}); - } -} - -BENCHMARK(BM_ChunkedPack_fixed_table_device) - ->Apply(ChunkedPackArguments) - ->UseRealTime() - ->Unit(benchmark::kMillisecond); - -BENCHMARK(BM_ChunkedPack_fixed_table_pinned) - ->Apply(ChunkedPackArguments) - ->UseRealTime() - ->Unit(benchmark::kMillisecond); - -BENCHMARK_MAIN(); diff --git a/cpp/benchmarks/bench_partition.cpp b/cpp/benchmarks/bench_partition.cpp deleted file mode 100644 index 779a91274..000000000 --- a/cpp/benchmarks/bench_partition.cpp +++ /dev/null @@ -1,175 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -// Helper function to create a table with a single int column -std::unique_ptr create_int_table( - cudf::size_type num_rows, rmm::cuda_stream_view stream -) { - auto data = rmm::device_buffer( - rapidsmpf::safe_cast(num_rows) * sizeof(std::int32_t), stream - ); - auto validity = rmm::device_buffer(0, stream); // No nulls - - auto column = std::make_unique( - cudf::data_type{cudf::type_id::INT32}, - num_rows, - std::move(data), - std::move(validity), - 0 - ); - - std::vector> columns; - columns.push_back(std::move(column)); - return std::make_unique(std::move(columns)); -} - -static void BM_PartitionAndPack(benchmark::State& state) { - const std::int64_t local_size = std::int64_t{state.range(1)} * 1000000; - int num_rows = int(local_size / std::int64_t{sizeof(std::int32_t)}); - - const int num_partitions = state.range(1); - - rmm::cuda_stream_view stream = rmm::cuda_stream_default; - - // Get total GPU memory - cudaDeviceProp prop; - cudaGetDeviceProperties(&prop, 0); - std::size_t total_memory = prop.totalGlobalMem; - - // Calculate 50% of GPU memory - auto pool_size = static_cast(total_memory * 0.5); - - // Create a pool memory resource with 50% of GPU memory - rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_memory_resource{}, pool_size}; - auto br = rapidsmpf::BufferResource::create(pool_mr); - - // Create input table - auto table = create_int_table(num_rows, stream); - - // Columns to hash (just the first column) - std::vector columns_to_hash{0}; - - for (auto _ : state) { - auto pack_partitions = cudf_streaming::integrations::partition_and_pack( - *table, - columns_to_hash, - num_partitions, - cudf::hash_id::HASH_MURMUR3, - cudf::DEFAULT_HASH_SEED, - stream, - br.get() - ); - benchmark::DoNotOptimize(pack_partitions); - cudaStreamSynchronize(stream); - } - - // Set metrics - state.SetBytesProcessed(state.iterations() * local_size); - state.counters["num_rows"] = num_rows; - state.counters["total_nparts"] = num_partitions; - state.counters["splits"] = num_partitions; -} - -static void BM_PartitionAndPackCurrentImpl(benchmark::State& state) { - const int nranks = state.range(0); - const std::int64_t local_size = std::int64_t{state.range(1)} * 1000000; - const int num_partitions = state.range(2); - - int total_npartitions = nranks * num_partitions; - int num_rows = - int(local_size / std::int64_t{sizeof(std::int32_t)} - / std::int64_t{num_partitions}); - - rmm::cuda_stream_view stream = rmm::cuda_stream_default; - - // Get total GPU memory - cudaDeviceProp prop; - cudaGetDeviceProperties(&prop, 0); - std::size_t total_memory = prop.totalGlobalMem; - - // Calculate 50% of GPU memory - auto pool_size = static_cast(total_memory * 0.5); - - // Create a pool memory resource with 50% of GPU memory - rmm::mr::pool_memory_resource pool_mr{rmm::mr::cuda_memory_resource{}, pool_size}; - auto br = rapidsmpf::BufferResource::create(pool_mr); - - // Create input table - auto table = create_int_table(num_rows, stream); - - // Columns to hash (just the first column) - std::vector columns_to_hash{0}; - - for (auto _ : state) { - for (int i = 0; i < num_partitions; i++) { - auto pack_partitions = cudf_streaming::integrations::partition_and_pack( - *table, - columns_to_hash, - total_npartitions, - cudf::hash_id::HASH_MURMUR3, - cudf::DEFAULT_HASH_SEED, - stream, - br.get() - ); - benchmark::DoNotOptimize(pack_partitions); - } - cudaStreamSynchronize(stream); - } - - // Set metrics - state.SetBytesProcessed(state.iterations() * local_size); - state.counters["num_rows"] = num_rows; - state.counters["total_nparts"] = total_npartitions; - state.counters["splits"] = total_npartitions * num_partitions; -} - -// Custom argument generator for the benchmark -void CustomArguments(benchmark::Benchmark* b) { - // Test different combinations of table sizes and partitions - for (auto nranks : {4}) { - for (int size_mb : {4000}) { - for (auto partitions : {2, 8, 32, 128, 512, 1024}) { - b->Args({nranks, size_mb, partitions}); - } - } - } -} - -// Register the benchmark with custom arguments -BENCHMARK(BM_PartitionAndPackCurrentImpl) - ->Apply(CustomArguments) - ->UseRealTime() - ->Unit(benchmark::kMillisecond); - - -// Register the benchmark with custom arguments -BENCHMARK(BM_PartitionAndPack) - ->Args({4000, 2}) - ->Args({4000, 8}) - ->Args({4000, 32}) - ->Args({4000, 128}) - ->Args({4000, 512}) - ->Args({4000, 1024}) - ->UseRealTime() - ->Unit(benchmark::kMillisecond); - -BENCHMARK_MAIN(); diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp deleted file mode 100644 index 25164fcfd..000000000 --- a/cpp/benchmarks/bench_shuffle.cpp +++ /dev/null @@ -1,720 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include - -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#ifdef RAPIDSMPF_HAVE_CUPTI -#include -#endif - -#include "utils/misc.hpp" -#include "utils/random_data.hpp" -#include "utils/rmm_utils.hpp" - -class ArgumentParser { - public: - ArgumentParser(int argc, char* const* argv, bool use_mpi = true) { - int rank = 0; - int nranks = 1; - - if (use_mpi) { - RAPIDSMPF_EXPECTS( - rapidsmpf::mpi::is_initialized() == true, "MPI is not initialized" - ); - - RAPIDSMPF_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &rank)); - RAPIDSMPF_MPI(MPI_Comm_size(MPI_COMM_WORLD, &nranks)); - } else { - // When not using MPI, expect to be using bootstrap mode (rrun) - nranks = rapidsmpf::bootstrap::get_nranks(); - } - try { - int option; - while ((option = getopt(argc, argv, "C:r:w:c:n:p:o:m:l:LigsbxhM:")) != -1) { - switch (option) { - case 'h': - { - std::stringstream ss; - ss << "Usage: " << argv[0] << " [options]\n" - << "Options:\n" - << " -C Communicator {mpi, ucxx} (default: mpi)\n" - << " -r Number of runs (default: 1)\n" - << " -w Number of warmup runs (default: 0)\n" - << " -c Number of columns in the input tables " - "(default: 1)\n" - << " -n Number of rows per rank (default: 1M)\n" - << " -p Number of partitions (input tables) per " - "rank (default: 1)\n" - << " -o Number of output partitions per rank " - "(default: 1)\n" - << " -m RMM memory resource {cuda, pool, async, " - "managed} (default: pool)\n" - << " -l Device memory limit in MiB (default:-1, " - "unlimited)\n" - << " -L Disable Pinned host memory (default: " - " unlimited)\n" - << " -g Use pre-partitioned (hash) input tables " - "(default: unset, hash partition during insertion)\n" - << " -s Discard output chunks to simulate streaming " - "(default: disabled)\n" - << " -b Disallow memory overbooking when generating " - "input data (default: allow memory overbooking)\n" - << " -x Enable memory profiler (default: disabled)\n" -#ifdef RAPIDSMPF_HAVE_CUPTI - << " -M Enable CUPTI memory monitoring and save CSV " - "files with given path prefix. For example, /tmp/test will " - "write files to /tmp/test_.csv (default: disabled)\n" -#endif - << " -h Display this help message\n"; - if (rank == 0) { - std::cerr << ss.str(); - } - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, 0)); - } else { - std::exit(0); - } - } - break; - case 'C': - comm_type = std::string{optarg}; - if (!(comm_type == "mpi" || comm_type == "ucxx")) { - if (rank == 0) { - std::cerr << "-C (Communicator) must be one of {mpi, ucxx}" - << std::endl; - } - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1)); - } else { - std::exit(-1); - } - } - break; - case 'r': - parse_integer(num_runs, optarg); - break; - case 'w': - parse_integer(num_warmups, optarg); - break; - case 'c': - parse_integer(num_columns, optarg); - break; - case 'n': - parse_integer(num_local_rows, optarg); - break; - case 'p': - parse_integer(num_local_partitions, optarg); - break; - case 'o': - parse_integer(num_output_partitions, optarg); - break; - case 'm': - rmm_mr = std::string{optarg}; - if (!(rmm_mr == "cuda" || rmm_mr == "pool" || rmm_mr == "async" - || rmm_mr == "managed")) - { - if (rank == 0) { - std::cerr << "-m (RMM memory resource) must be one of " - "{cuda, pool, async, managed}" - << std::endl; - } - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1)); - } else { - std::exit(-1); - } - } - break; - case 'l': - parse_integer(device_mem_limit_mb, optarg); - break; - case 'L': - pinned_mem_disable = true; - break; - case 'g': - hash_partition_with_datagen = true; - break; - case 's': - enable_output_discard = true; - break; - case 'b': - input_data_allow_overbooking = rapidsmpf::AllowOverbooking::NO; - break; - case 'x': - enable_memory_profiler = true; - break; -#ifdef RAPIDSMPF_HAVE_CUPTI - case 'M': - cupti_csv_prefix = std::string{optarg}; - enable_cupti_monitoring = true; - break; -#endif - case '?': - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1)); - } else { - std::exit(-1); - } - break; - default: - RAPIDSMPF_FAIL("unknown option", std::invalid_argument); - } - } - if (optind < argc) { - RAPIDSMPF_FAIL("unknown option", std::invalid_argument); - } - } catch (std::exception const& e) { - if (rank == 0) { - std::cerr << "Error parsing arguments: " << e.what() << std::endl; - } - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1)); - } else { - std::exit(-1); - } - } - - local_nbytes = - num_columns * num_local_rows * num_local_partitions * sizeof(std::int32_t); - total_nbytes = local_nbytes * static_cast(nranks); - if (rmm_mr == "cuda") { - if (rank == 0) { - std::cout << "WARNING: using the default cuda memory resource " - "(-m cuda) might leak memory! A limitation in UCX " - "means that device memory send through IPC can " - "never be freed." - << std::endl; - } - } - } - - void pprint(rapidsmpf::Communicator& comm) const { - if (comm.rank() > 0) { - return; - } - std::stringstream ss; - ss << "Arguments:\n"; - ss << " -c " << comm_type << " (communicator)\n"; - ss << " -r " << num_runs << " (number of runs)\n"; - ss << " -w " << num_warmups << " (number of warmup runs)\n"; - ss << " -c " << num_columns << " (number of columns)\n"; - ss << " -n " << num_local_rows << " (number of rows per rank)\n"; - ss << " -p " << num_local_partitions - << " (number of input partitions per rank)\n"; - ss << " -o " << num_output_partitions - << " (number of output partitions per rank)\n"; - ss << " -m " << rmm_mr << " (RMM memory resource)\n"; - if (device_mem_limit_mb >= 0) { - ss << " -l " << device_mem_limit_mb << " (device memory limit in MiB)\n"; - } - if (pinned_mem_disable) { - ss << " -L (disable pinned host memory)\n"; - } - if (enable_output_discard) { - ss << " -s (enable output discard to simulate streaming)\n"; - } - if (input_data_allow_overbooking == rapidsmpf::AllowOverbooking::NO) { - ss << " -b (disallow memory overbooking when generating input data)\n"; - } - if (enable_memory_profiler) { - ss << " -x (enable memory profiling)\n"; - } - if (hash_partition_with_datagen) { - ss << " -g (use pre-partitioned input tables)\n"; - } - if (enable_cupti_monitoring) { - ss << " -M " << cupti_csv_prefix << " (CUPTI memory monitoring enabled)\n"; - } - ss << "Local size: " << rapidsmpf::format_nbytes(local_nbytes) << "\n"; - ss << "Total size: " << rapidsmpf::format_nbytes(total_nbytes) << "\n"; - comm.logger()->print(ss.str()); - } - - std::uint64_t num_runs{1}; - std::uint64_t num_warmups{0}; - std::uint32_t num_columns{1}; - std::uint64_t num_local_rows{1 << 20}; - rapidsmpf::shuffler::PartID num_local_partitions{1}; - rapidsmpf::shuffler::PartID num_output_partitions{1}; - std::string rmm_mr{"pool"}; - std::string comm_type{"mpi"}; - std::uint64_t local_nbytes; - std::uint64_t total_nbytes; - bool enable_output_discard{false}; - rapidsmpf::AllowOverbooking input_data_allow_overbooking{ - rapidsmpf::AllowOverbooking::YES - }; - bool enable_memory_profiler{false}; - bool hash_partition_with_datagen{false}; - std::int64_t device_mem_limit_mb{-1}; - bool pinned_mem_disable{false}; - bool enable_cupti_monitoring{false}; - std::string cupti_csv_prefix; -}; - -void barrier(std::shared_ptr& comm) { - bool use_bootstrap = rapidsmpf::bootstrap::is_running_with_rrun(); - if (!use_bootstrap) { - RAPIDSMPF_MPI(MPI_Barrier(MPI_COMM_WORLD)); - } else { - std::dynamic_pointer_cast(comm)->barrier(); - } -} - -rapidsmpf::Duration do_run( - rapidsmpf::shuffler::PartID const total_num_partitions, - std::shared_ptr& comm, - ArgumentParser const& args, - rmm::cuda_stream_view stream, - rapidsmpf::BufferResource* br, - std::shared_ptr statistics, - auto&& shuffle_insert_fn -) { - std::vector> output_partitions; - output_partitions.reserve(total_num_partitions); - - barrier(comm); - - auto const t0_elapsed = rapidsmpf::Clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Shuffling", total_num_partitions); - if (args.enable_memory_profiler) { - RAPIDSMPF_MEMORY_PROFILE(statistics, br->device_mr(), "shuffling"); - } - rapidsmpf::shuffler::Shuffler shuffler( - comm, - 0, // op_id - total_num_partitions, - br, - rapidsmpf::shuffler::Shuffler::round_robin - ); - - // insert partitions into the shuffler - shuffle_insert_fn(shuffler); - - shuffler.wait(); - for (auto finished_partition : shuffler.local_partitions()) { - auto packed_chunks = shuffler.extract(finished_partition); - auto output_partition = cudf_streaming::integrations::unpack_and_concat( - rapidsmpf::unspill_partitions( - std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES - ), - stream, - br - ); - if (!args.enable_output_discard) { - output_partitions.emplace_back(std::move(output_partition)); - } - } - stream.synchronize(); - } - - auto const elapsed = rapidsmpf::Clock::now() - t0_elapsed; - - // Check the shuffle result (this test only works for non-empty partitions - // thus we only check large shuffles). - if (args.num_local_rows >= 1000000) { - for (const auto& output_partition : output_partitions) { - auto [parts, owner] = cudf_streaming::integrations::partition_and_split( - output_partition->view(), - {0}, - static_cast(total_num_partitions), - cudf::hash_id::HASH_MURMUR3, - cudf::DEFAULT_HASH_SEED, - stream, - br - ); - RAPIDSMPF_EXPECTS( - std::count_if( - parts.begin(), - parts.end(), - [](auto const& table) { return table.num_rows() > 0; } - ) == 1, - "all rows in an output partition should hash to the same" - ); - } - } - - barrier(comm); - - return elapsed; -} - -// generate input partitions by applying a transform function to each table -template < - typename TransformFn, - typename InputPartitionsT = - std::remove_reference_t>> -std::vector generate_input_partitions( - ArgumentParser const& args, - rmm::cuda_stream_view stream, - rapidsmpf::BufferResource* br, - TransformFn&& transform_fn -) { - auto const num_columns = rapidsmpf::safe_cast(args.num_columns); - auto const num_local_rows = - rapidsmpf::safe_cast(args.num_local_rows); - std::int32_t const min_val = 0; - std::int32_t const max_val = num_local_rows; - - std::vector input_partitions; - input_partitions.reserve(args.num_local_partitions); - for (rapidsmpf::shuffler::PartID i = 0; i < args.num_local_partitions; ++i) { - std::size_t size_lb = random_table_size_lower_bound(num_columns, num_local_rows); - - // reserve at least size_lb and spill if necessary. - auto res = br->reserve_device_memory_and_spill( - size_lb, args.input_data_allow_overbooking - ); - cudf::table table = random_table( - num_columns, num_local_rows, min_val, max_val, stream, br->device_mr() - ); - input_partitions.emplace_back(transform_fn(std::move(table))); - } - stream.synchronize(); - return input_partitions; -} - -/** - * Helper function to iterate over input partitions and insert them into the shuffler. - * - * @param shuffler Shuffler to insert the partitions into. - * @param input_partitions This is either a vector or - * vector>. Former will be forwarded to to - * partition_and_pack to generate a unordered_map for each table. - * @param make_chunk_fn Function to make a chunk from a partition. - */ -void do_insert( - rapidsmpf::shuffler::Shuffler& shuffler, auto&& input_partitions, auto&& make_chunk_fn -) { - // Convert a partition into chunks and insert into the shuffler. - for (auto&& partition : input_partitions) { - shuffler.insert(std::move(make_chunk_fn(partition))); - } - - // Tell the shuffler that we have no more data. - shuffler.insert_finished(); -} - -/** - * @brief Runs shuffle by partitioning the input tables and inserting them into the - * shuffler. - * - * This function generates random input tables and partitions them into the number of - * input partitions specified by the user. It then inserts the partitions into the - * shuffler and runs the shuffle. Each input partition will be partitioned into - * `num_output_partitions * nranks`, resulting in, `num_local_partitions * - * num_output_partitions * nranks` chunks being inserted into the shuffler. Each chunk - * size will be `~num_local_rows/(num_output_partitions * nranks)` rows. - * - * @param comm Communicator for the shuffler - * @param args Command line arguments - * @param stream CUDA stream for the shuffler - * @param br Buffer resource for the shuffler - * @param statistics Statistics for the shuffler - * @return Duration of the run - */ -rapidsmpf::Duration run_hash_partition_inline( - std::shared_ptr& comm, - ArgumentParser const& args, - rmm::cuda_stream_view stream, - rapidsmpf::BufferResource* br, - std::shared_ptr statistics -) { - rapidsmpf::shuffler::PartID const total_num_partitions = - args.num_output_partitions - * static_cast(comm->nranks()); - - std::vector input_partitions = - generate_input_partitions(args, stream, br, std::identity{}); - - auto make_chunk_fn = [&](cudf::table const& partition) { - return cudf_streaming::integrations::partition_and_pack( - partition, - {0}, - static_cast(total_num_partitions), - cudf::hash_id::HASH_MURMUR3, - cudf::DEFAULT_HASH_SEED, - stream, - br - ); - }; - - return do_run( - total_num_partitions, comm, args, stream, br, statistics, [&](auto& shuffler) { - do_insert(shuffler, std::move(input_partitions), std::move(make_chunk_fn)); - } - ); -} - -/** - * @brief Runs shuffle by using pre-partitioned input tables. - * - * This is similar to the hash partitioning, but the input tables are already - * partitioned before being inserted into the shuffler. - * - * @param comm Communicator for the shuffler - * @param args Command line arguments - * @param stream CUDA stream for the shuffler - * @param br Buffer resource for the shuffler - * @param statistics Statistics for the shuffler - * @return Duration of the run - */ -rapidsmpf::Duration run_hash_partition_with_datagen( - std::shared_ptr& comm, - ArgumentParser const& args, - rmm::cuda_stream_view stream, - rapidsmpf::BufferResource* br, - std::shared_ptr statistics -) { - rapidsmpf::shuffler::PartID const total_num_partitions = - args.num_output_partitions - * static_cast(comm->nranks()); - - std::vector> - input_partitions = - generate_input_partitions(args, stream, br, [&](cudf::table&& table) { - return cudf_streaming::integrations::partition_and_pack( - table, - {0}, - static_cast(total_num_partitions), - cudf::hash_id::HASH_MURMUR3, - cudf::DEFAULT_HASH_SEED, - stream, - br - ); - }); - - return do_run( - total_num_partitions, comm, args, stream, br, statistics, [&](auto& shuffler) { - do_insert(shuffler, std::move(input_partitions), std::identity{}); - } - ); -} - -int main(int argc, char** argv) { - bool use_bootstrap = rapidsmpf::bootstrap::is_running_with_rrun(); - - // Explicitly initialize MPI with thread support, as this is needed for both mpi - // and ucxx communicators when not using bootstrap mode. - int provided = 0; - if (!use_bootstrap) { - RAPIDSMPF_MPI(MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &provided)); - - RAPIDSMPF_EXPECTS( - provided == MPI_THREAD_MULTIPLE, - "didn't get the requested thread level support: MPI_THREAD_MULTIPLE" - ); - } - - ArgumentParser args{argc, argv, !use_bootstrap}; - - // Initialize configuration options from environment variables. - rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()}; - - set_current_rmm_resource(args.rmm_mr); - rapidsmpf::RmmResourceAdaptor stat_enabled_mr = set_device_mem_resource_with_stats(); - - std::unordered_map memory_limits{}; - if (args.device_mem_limit_mb >= 0) { - memory_limits[rapidsmpf::MemoryType::DEVICE] = args.device_mem_limit_mb << 20; - } - - auto stats = rapidsmpf::Statistics::create(); - - // We're only going to measure the last run, so disable initially. - stats->disable(); - auto br = rapidsmpf::BufferResource::create( - stat_enabled_mr, - args.pinned_mem_disable ? rapidsmpf::PinnedMemoryResource::Disabled - : rapidsmpf::PinnedMemoryResource::make_if_available(), - std::move(memory_limits), - std::chrono::milliseconds{1}, - std::make_shared( - 16, rmm::cuda_stream::flags::non_blocking - ), - stats - ); - - std::shared_ptr comm; - auto progress_thread = std::make_shared(stats); - if (args.comm_type == "mpi") { - if (use_bootstrap) { - std::cerr - << "Error: MPI communicator requires MPI initialization. Don't use with " - "rrun or unset RRUN_RANK." - << std::endl; - return 1; - } - rapidsmpf::mpi::init(&argc, &argv); - comm = std::make_shared(MPI_COMM_WORLD, options, progress_thread); - } else if (args.comm_type == "ucxx") { - if (use_bootstrap) { - // Launched with rrun - use bootstrap backend - comm = rapidsmpf::bootstrap::create_ucxx_comm( - progress_thread, rapidsmpf::bootstrap::BackendType::AUTO, options - ); - } else { - // Launched with mpirun - use MPI bootstrap - comm = - rapidsmpf::ucxx::init_using_mpi(MPI_COMM_WORLD, options, progress_thread); - } - } else { - std::cerr << "Error: Unknown communicator type: " << args.comm_type << std::endl; - return 1; - } - - args.pprint(*comm); - - auto& log = comm->logger(); - rmm::cuda_stream_view stream = cudf::get_default_stream(); - - // Print benchmark/hardware info. - { - std::stringstream ss; - auto const cur_dev = rmm::get_current_cuda_device().value(); - std::string pci_bus_id(16, '\0'); // Preallocate space for the PCI bus ID - RAPIDSMPF_CUDA_TRY( - cudaDeviceGetPCIBusId(pci_bus_id.data(), pci_bus_id.size(), cur_dev) - ); - cudaDeviceProp properties; - RAPIDSMPF_CUDA_TRY(cudaGetDeviceProperties(&properties, 0)); - ss << "Hardware setup: \n"; - ss << " GPU (" << properties.name << "): \n"; - ss << " Device number: " << cur_dev << "\n"; - ss << " PCI Bus ID: " << pci_bus_id.substr(0, pci_bus_id.find('\0')) << "\n"; - ss << " Total Memory: " - << rapidsmpf::format_nbytes(properties.totalGlobalMem, 0) << "\n"; - ss << " Comm: " << *comm << "\n"; - log->print(ss.str()); - } - -#ifdef RAPIDSMPF_HAVE_CUPTI - // Create CUPTI monitor if enabled - std::unique_ptr cupti_monitor; - if (args.enable_cupti_monitoring) { - cupti_monitor = std::make_unique(); - cupti_monitor->start_monitoring(); - log->print("CUPTI memory monitoring enabled"); - } -#endif - - std::vector elapsed_vec; - std::uint64_t const total_num_runs = args.num_warmups + args.num_runs; - for (std::uint64_t i = 0; i < total_num_runs; ++i) { - // Enable statistics before the last run so only last-run data is reported. - if (i == total_num_runs - 1) { - stats->enable(); - } - double elapsed; - if (args.hash_partition_with_datagen) { - elapsed = run_hash_partition_with_datagen(comm, args, stream, br.get(), stats) - .count(); - } else { - elapsed = - run_hash_partition_inline(comm, args, stream, br.get(), stats).count(); - } - std::stringstream ss; - ss << "elapsed: " << rapidsmpf::format_duration(elapsed) - << " | local throughput: " - << rapidsmpf::format_nbytes(args.local_nbytes / elapsed) - << "/s | global throughput: " - << rapidsmpf::format_nbytes(args.total_nbytes / elapsed) << "/s"; - if (i < args.num_warmups) { - ss << " (warmup run)"; - } - log->print(ss.str()); - if (i >= args.num_warmups) { - elapsed_vec.push_back(elapsed); - } - } - - { - auto const elapsed_mean = harmonic_mean(elapsed_vec); - std::stringstream ss; - ss << "means: " << rapidsmpf::format_duration(elapsed_mean) - << " | local throughput: " - << rapidsmpf::format_nbytes(args.local_nbytes / elapsed_mean) - << "/s | global throughput: " - << rapidsmpf::format_nbytes(args.total_nbytes / elapsed_mean) << "/s" - << " | in_parts: " << args.num_local_partitions - << " | out_parts: " << args.num_output_partitions - << " | nranks: " << comm->nranks(); - if (args.enable_memory_profiler) { - auto record = stat_enabled_mr.get_main_record(); - ss << " | device memory peak: " << rapidsmpf::format_nbytes(record.peak()) - << " | device memory total: " - << rapidsmpf::format_nbytes( - record.total() / static_cast(total_num_runs) - ) - << " (avg)"; - } - log->print(ss.str()); - } - - if (args.enable_memory_profiler) { - log->print(stats->report( - {.mr = stat_enabled_mr, .header = "Statistics (of the last run):"} - )); - } else { - log->print(stats->report({.header = "Statistics (of the last run):"})); - } - -#ifdef RAPIDSMPF_HAVE_CUPTI - // Save CUPTI monitoring results to CSV file - if (args.enable_cupti_monitoring && cupti_monitor) { - cupti_monitor->stop_monitoring(); - - std::string csv_filename = - args.cupti_csv_prefix + std::to_string(comm->rank()) + ".csv"; - try { - cupti_monitor->write_csv(csv_filename); - log->print( - "CUPTI memory data written to " + csv_filename + " (" - + std::to_string(cupti_monitor->get_sample_count()) + " samples, " - + std::to_string(cupti_monitor->get_total_callback_count()) - + " callbacks)" - ); - - // Print callback summary for rank 0 - if (comm->rank() == 0) { - log->print( - "CUPTI Callback Summary:\n" + cupti_monitor->get_callback_summary() - ); - } - } catch (std::exception const& e) { - log->print("Failed to write CUPTI CSV file: " + std::string(e.what())); - } - } -#endif - - if (!use_bootstrap) { - RAPIDSMPF_MPI(MPI_Finalize()); - } - return 0; -} diff --git a/cpp/benchmarks/streaming/CMakeLists.txt b/cpp/benchmarks/streaming/CMakeLists.txt deleted file mode 100644 index e8c6542ae..000000000 --- a/cpp/benchmarks/streaming/CMakeLists.txt +++ /dev/null @@ -1,38 +0,0 @@ -# ================================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ================================================================================= - -add_executable(bench_streaming_shuffle "bench_streaming_shuffle.cpp") -set_target_properties( - bench_streaming_shuffle - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON -) -target_compile_options( - bench_streaming_shuffle PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" -) -target_link_libraries( - bench_streaming_shuffle - PRIVATE rapidsmpf::rapidsmpf - ucxx::ucxx - cudf_streaming::cudf_streaming - $ - $ - maybe_asan - bench_utils -) -install( - TARGETS bench_streaming_shuffle - COMPONENT benchmarking - DESTINATION bin/benchmarks/librapidsmpf - EXCLUDE_FROM_ALL -) - -add_subdirectory(ndsh) diff --git a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp deleted file mode 100644 index 9806b49ec..000000000 --- a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp +++ /dev/null @@ -1,479 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include - -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "../utils/misc.hpp" -#include "../utils/rmm_utils.hpp" -#include "data_generator.hpp" - -class ArgumentParser { - public: - ArgumentParser(int argc, char* const* argv, bool use_mpi = true) { - int rank = 0; - int nranks = 1; - - if (use_mpi) { - RAPIDSMPF_EXPECTS( - rapidsmpf::mpi::is_initialized() == true, "MPI is not initialized" - ); - - RAPIDSMPF_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &rank)); - RAPIDSMPF_MPI(MPI_Comm_size(MPI_COMM_WORLD, &nranks)); - } else { - // When not using MPI, expect to be using bootstrap mode (rrun) - nranks = rapidsmpf::bootstrap::get_nranks(); - } - try { - int option; - while ((option = getopt(argc, argv, "C:r:w:c:n:p:o:m:l:Lxh")) != -1) { - switch (option) { - case 'h': - { - std::stringstream ss; - ss << "Usage: " << argv[0] << " [options]\n" - << "Options:\n" - << " -C Communicator {mpi, ucxx} (default: mpi)\n" - << " -r Number of runs (default: 1)\n" - << " -w Number of warmup runs (default: 0)\n" - << " -c Number of columns in the input tables " - "(default: 1)\n" - << " -n Number of rows per rank (default: 1M)\n" - << " -p Number of partitions (input tables) per " - "rank (default: 1)\n" - << " -o Number of output partitions per rank " - "(default: 1)\n" - << " -m RMM memory resource {cuda, pool, async, " - "managed} (default: pool)\n" - << " -l Device memory limit in MiB (default:-1, " - "unlimited)\n" - << " -L Disable Pinned host memory (default: " - " unlimited)\n" - << " -x Enable memory profiler (default: disabled)\n" - << " -h Display this help message\n"; - if (rank == 0) { - std::cerr << ss.str(); - } - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, 0)); - } else { - std::exit(0); - } - } - break; - case 'C': - comm_type = std::string{optarg}; - if (!(comm_type == "mpi" || comm_type == "ucxx")) { - if (rank == 0) { - std::cerr << "-C (Communicator) must be one of {mpi, ucxx}" - << std::endl; - } - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1)); - } else { - std::exit(-1); - } - } - break; - case 'r': - parse_integer(num_runs, optarg); - break; - case 'w': - parse_integer(num_warmups, optarg); - break; - case 'c': - parse_integer(num_columns, optarg); - break; - case 'n': - parse_integer(num_local_rows, optarg); - break; - case 'p': - parse_integer(num_local_partitions, optarg); - break; - case 'o': - parse_integer(num_output_partitions, optarg); - break; - case 'm': - rmm_mr = std::string{optarg}; - if (!(rmm_mr == "cuda" || rmm_mr == "pool" || rmm_mr == "async" - || rmm_mr == "managed")) - { - if (rank == 0) { - std::cerr << "-m (RMM memory resource) must be one of " - "{cuda, pool, async, managed}" - << std::endl; - } - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1)); - } else { - std::exit(-1); - } - } - break; - case 'l': - parse_integer(device_mem_limit_mb, optarg); - break; - case 'L': - pinned_mem_disable = true; - break; - case 'x': - enable_memory_profiler = true; - break; - case '?': - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1)); - } else { - std::exit(-1); - } - break; - default: - RAPIDSMPF_FAIL("unknown option", std::invalid_argument); - } - } - if (optind < argc) { - RAPIDSMPF_FAIL("unknown option", std::invalid_argument); - } - } catch (std::exception const& e) { - if (rank == 0) { - std::cerr << "Error parsing arguments: " << e.what() << std::endl; - } - if (use_mpi) { - RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1)); - } else { - std::exit(-1); - } - } - - local_nbytes = - num_columns * num_local_rows * num_local_partitions * sizeof(std::int32_t); - total_nbytes = local_nbytes * static_cast(nranks); - if (rmm_mr == "cuda") { - if (rank == 0) { - std::cout << "WARNING: using the default cuda memory resource " - "(-m cuda) might leak memory! A limitation in UCX " - "means that device memory send through IPC can " - "never be freed." - << std::endl; - } - } - } - - void pprint(rapidsmpf::Communicator& comm) const { - if (comm.rank() > 0) { - return; - } - std::stringstream ss; - ss << "Arguments:\n"; - ss << " -c " << comm_type << " (communicator)\n"; - ss << " -r " << num_runs << " (number of runs)\n"; - ss << " -w " << num_warmups << " (number of warmup runs)\n"; - ss << " -c " << num_columns << " (number of columns)\n"; - ss << " -n " << num_local_rows << " (number of rows per rank)\n"; - ss << " -p " << num_local_partitions - << " (number of input partitions per rank)\n"; - ss << " -o " << num_output_partitions - << " (number of output partitions per rank)\n"; - ss << " -m " << rmm_mr << " (RMM memory resource)\n"; - if (device_mem_limit_mb >= 0) { - ss << " -l " << device_mem_limit_mb << " (device memory limit in MiB)\n"; - } - if (pinned_mem_disable) { - ss << " -L (disable pinned host memory)\n"; - } - if (enable_memory_profiler) { - ss << " -x (enable memory profiling)\n"; - } - ss << "Local size: " << rapidsmpf::format_nbytes(local_nbytes) << "\n"; - ss << "Total size: " << rapidsmpf::format_nbytes(total_nbytes) << "\n"; - comm.logger()->print(ss.str()); - } - - std::uint64_t num_runs{1}; - std::uint64_t num_warmups{0}; - std::uint32_t num_columns{1}; - std::uint64_t num_local_rows{1 << 20}; - rapidsmpf::shuffler::PartID num_local_partitions{1}; - rapidsmpf::shuffler::PartID num_output_partitions{1}; - std::string rmm_mr{"pool"}; - std::string comm_type{"mpi"}; - std::uint64_t local_nbytes; - std::uint64_t total_nbytes; - bool enable_memory_profiler{false}; - std::int64_t device_mem_limit_mb{-1}; - bool pinned_mem_disable{false}; -}; - -rapidsmpf::streaming::Actor consumer( - std::shared_ptr ctx, - std::shared_ptr ch_in -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in}; - co_await ctx->executor()->schedule(); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - } -} - -rapidsmpf::Duration run( - std::shared_ptr ctx, - std::shared_ptr comm, - ArgumentParser const& args, - rmm::cuda_stream_view stream -) { - constexpr std::int32_t min_val = 0; - constexpr std::int32_t max_val = 10; - constexpr cudf::hash_id hash_function = cudf::hash_id::HASH_MURMUR3; - constexpr std::uint32_t seed = cudf::DEFAULT_HASH_SEED; - rapidsmpf::shuffler::PartID const total_num_partitions = - args.num_output_partitions - * static_cast(comm->nranks()); - constexpr rapidsmpf::OpID op_id = 0; - - // Create streaming pipeline. - std::vector actors; - { - auto ch1 = ctx->create_channel(); - auto const num_columns = rapidsmpf::safe_cast(args.num_columns); - auto const num_local_rows = - rapidsmpf::safe_cast(args.num_local_rows); - actors.push_back( - rapidsmpf::streaming::actor::random_table_generator( - ctx, - stream, - ch1, - args.num_local_partitions, - num_columns, - num_local_rows, - min_val, - max_val - ) - ); - auto ch2 = ctx->create_channel(); - actors.push_back( - cudf_streaming::streaming::actor::partition_and_pack( - ctx, - ch1, - ch2, - {0}, - static_cast(total_num_partitions), - hash_function, - seed - ) - ); - auto ch3 = ctx->create_channel(); - actors.push_back( - rapidsmpf::streaming::actor::shuffler( - ctx, comm, ch2, ch3, op_id, total_num_partitions - ) - ); - auto ch4 = ctx->create_channel(); - actors.push_back( - cudf_streaming::streaming::actor::unpack_and_concat(ctx, ch3, ch4) - ); - actors.push_back(consumer(ctx, ch4)); - } - auto const t0_elapsed = rapidsmpf::Clock::now(); - rapidsmpf::streaming::run_actor_network(std::move(actors)); - return rapidsmpf::Clock::now() - t0_elapsed; -} - -int main(int argc, char** argv) { - bool use_bootstrap = rapidsmpf::bootstrap::is_running_with_rrun(); - - // Explicitly initialize MPI with thread support, as this is needed for both mpi - // and ucxx communicators when not using bootstrap mode. - int provided = 0; - if (!use_bootstrap) { - RAPIDSMPF_MPI(MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &provided)); - - RAPIDSMPF_EXPECTS( - provided == MPI_THREAD_MULTIPLE, - "didn't get the requested thread level support: MPI_THREAD_MULTIPLE" - ); - } - ArgumentParser args{argc, argv, !use_bootstrap}; - - // Initialize configuration options from environment variables. - rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()}; - auto progress_thread = std::make_shared(); - - std::shared_ptr comm; - if (args.comm_type == "mpi") { - if (use_bootstrap) { - std::cerr - << "Error: MPI communicator requires MPI initialization. Don't use with " - "rrun or unset RRUN_RANK." - << std::endl; - return 1; - } - rapidsmpf::mpi::init(&argc, &argv); - comm = std::make_shared(MPI_COMM_WORLD, options, progress_thread); - } else if (args.comm_type == "ucxx") { - if (use_bootstrap) { - // Launched with rrun - use bootstrap backend - comm = rapidsmpf::bootstrap::create_ucxx_comm( - progress_thread, rapidsmpf::bootstrap::BackendType::AUTO, options - ); - } else { - // Launched with mpirun - use MPI bootstrap - comm = - rapidsmpf::ucxx::init_using_mpi(MPI_COMM_WORLD, options, progress_thread); - } - } else { - std::cerr << "Error: Unknown communicator type: " << args.comm_type << std::endl; - return 1; - } - - args.pprint(*comm); - - RAPIDSMPF_EXPECTS(comm->nranks() == 1, "only single-rank runs are supported"); - - set_current_rmm_resource(args.rmm_mr); - auto stat_enabled_mr = set_device_mem_resource_with_stats(); - std::unordered_map memory_limits{}; - if (args.device_mem_limit_mb >= 0) { - memory_limits[rapidsmpf::MemoryType::DEVICE] = args.device_mem_limit_mb << 20; - } - - auto stats = rapidsmpf::Statistics::create(); - - auto pinned_mr = args.pinned_mem_disable - ? rapidsmpf::PinnedMemoryResource::Disabled - : rapidsmpf::PinnedMemoryResource::make_if_available(); - auto br = rapidsmpf::BufferResource::create( - stat_enabled_mr, - pinned_mr, - std::move(memory_limits), - std::nullopt, - std::make_shared( - 16, rmm::cuda_stream::flags::non_blocking - ), - stats - ); - - auto& log = *comm->logger(); - rmm::cuda_stream_view stream = cudf::get_default_stream(); - - // Print benchmark/hardware info. - { - std::stringstream ss; - auto const cur_dev = rmm::get_current_cuda_device().value(); - std::string pci_bus_id(16, '\0'); // Preallocate space for the PCI bus ID - RAPIDSMPF_CUDA_TRY( - cudaDeviceGetPCIBusId(pci_bus_id.data(), pci_bus_id.size(), cur_dev) - ); - cudaDeviceProp properties; - RAPIDSMPF_CUDA_TRY(cudaGetDeviceProperties(&properties, 0)); - ss << "Hardware setup: \n"; - ss << " GPU (" << properties.name << "): \n"; - ss << " Device number: " << cur_dev << "\n"; - ss << " PCI Bus ID: " << pci_bus_id.substr(0, pci_bus_id.find('\0')) << "\n"; - ss << " Total Memory: " - << rapidsmpf::format_nbytes(properties.totalGlobalMem, 0) << "\n"; - ss << " Comm: " << *comm << "\n"; - log.print(ss.str()); - } - - auto ctx = - std::make_shared(options, comm->logger(), br); - - std::vector elapsed_vec; - std::uint64_t const total_num_runs = args.num_warmups + args.num_runs; - for (std::uint64_t i = 0; i < total_num_runs; ++i) { - // Clear statistics before the last run so only the final run is reported. - if (i == total_num_runs - 1) { - ctx->statistics()->clear(); - } - double const elapsed = run(ctx, comm, args, stream).count(); - std::stringstream ss; - ss << "elapsed: " << rapidsmpf::format_duration(elapsed) - << " | local throughput: " - << rapidsmpf::format_nbytes(args.local_nbytes / elapsed) - << "/s | global throughput: " - << rapidsmpf::format_nbytes(args.total_nbytes / elapsed) << "/s"; - if (i < args.num_warmups) { - ss << " (warmup run)"; - } - log.print(ss.str()); - if (i >= args.num_warmups) { - elapsed_vec.push_back(elapsed); - } - } - - if (!use_bootstrap) { - RAPIDSMPF_MPI(MPI_Barrier(MPI_COMM_WORLD)); - } else { - std::dynamic_pointer_cast(comm)->barrier(); - } - - { - auto const elapsed_mean = harmonic_mean(elapsed_vec); - std::stringstream ss; - ss << "means: " << rapidsmpf::format_duration(elapsed_mean) - << " | local throughput: " - << rapidsmpf::format_nbytes(args.local_nbytes / elapsed_mean) - << "/s | global throughput: " - << rapidsmpf::format_nbytes(args.total_nbytes / elapsed_mean) << "/s" - << " | in_parts: " << args.num_local_partitions - << " | out_parts: " << args.num_output_partitions - << " | nranks: " << comm->nranks(); - if (args.enable_memory_profiler) { - auto record = stat_enabled_mr.get_main_record(); - ss << " | device memory peak: " << rapidsmpf::format_nbytes(record.peak()) - << " | device memory total: " - << rapidsmpf::format_nbytes( - record.total() / static_cast(total_num_runs) - ) - << " (avg)"; - } - log.print(ss.str()); - } - - auto statistics = ctx->statistics(); - if (args.enable_memory_profiler) { - log.print(statistics->report({ - .mr = stat_enabled_mr, - .pinned_mr = pinned_mr, - .header = "Statistics (of the last run):", - })); - } else { - log.print(statistics->report({.header = "Statistics (of the last run):"})); - } - - if (!use_bootstrap) { - RAPIDSMPF_MPI(MPI_Finalize()); - } - return 0; -} diff --git a/cpp/benchmarks/streaming/data_generator.hpp b/cpp/benchmarks/streaming/data_generator.hpp deleted file mode 100644 index d72e3ca0c..000000000 --- a/cpp/benchmarks/streaming/data_generator.hpp +++ /dev/null @@ -1,77 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include -#include -#include - -#include -#include -#include -#include - -#include "../utils/random_data.hpp" - -namespace rapidsmpf::streaming::actor { - -/** - * @brief Asynchronously generates and sends a sequence of random numeric tables. - * - * This is a streaming version of `rapidsmpf::random_table_generator` that operates on - * table chunks using channels. - * - * It creates a specified number of cuDF tables with random `std::int32_t` values, each - * consisting of `ncolumns` columns and `nrows` rows. The values are uniformly - * distributed in the range [`min_val`, `max_val`]. Each generated table is wrapped - * in a `TableChunk` and sent to the provided output channel in streaming fashion. - * - * @param ctx The actor context to use. - * @param stream The CUDA stream on which to create the random tables. TODO: use a pool - * of CUDA streams. - * @param ch_out Output channel to which generated `TableChunk` objects are sent. - * @param num_blocks Number of tables (chunks) to generate and send. - * @param ncolumns Number of columns per generated table. - * @param nrows Number of rows per column in each table. - * @param min_val Minimum inclusive value for the generated random integers. - * @param max_val Maximum inclusive value for the generated random integers. - * @return A streaming actor that completes once all random tables have been generated - * and sent, and the channel has been drained. - */ -inline Actor random_table_generator( - std::shared_ptr ctx, - rmm::cuda_stream_view stream, - std::shared_ptr ch_out, - std::uint64_t num_blocks, - cudf::size_type ncolumns, - cudf::size_type nrows, - std::int32_t min_val, - std::int32_t max_val -) { - ShutdownAtExit c{ch_out}; - co_await ctx->executor()->schedule(); - auto nbytes = rapidsmpf::safe_cast(ncolumns) - * rapidsmpf::safe_cast(nrows) * sizeof(std::int32_t); - for (std::uint64_t seq = 0; seq < num_blocks; ++seq) { - auto res = - ctx->br()->reserve_device_memory_and_spill(nbytes, AllowOverbooking::NO); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - seq, - std::make_unique( - std::make_unique(random_table( - ncolumns, nrows, min_val, max_val, stream, ctx->br()->device_mr() - )), - stream - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} - - -} // namespace rapidsmpf::streaming::actor diff --git a/cpp/benchmarks/streaming/ndsh/CMakeLists.txt b/cpp/benchmarks/streaming/ndsh/CMakeLists.txt deleted file mode 100644 index 45e0c96a0..000000000 --- a/cpp/benchmarks/streaming/ndsh/CMakeLists.txt +++ /dev/null @@ -1,76 +0,0 @@ -# ================================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ================================================================================= - -if(NOT RAPIDSMPF_HAVE_MPI) - message(FATAL_ERROR "Streaming NDSH benchmarks require MPI support") -endif() - -if(NOT RAPIDSMPF_HAVE_STREAMING) - message(FATAL_ERROR "Streaming NDSH benchmarks require streaming support") -endif() - -add_library( - rapidsmpfndsh concatenate.cpp groupby.cpp join.cpp parquet_writer.cpp sort.cpp utils.cpp -) - -set_target_properties( - rapidsmpfndsh - PROPERTIES BUILD_RPATH "\$ORIGIN" - INSTALL_RPATH "\$ORIGIN" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON - POSITION_INDEPENDENT_CODE ON - INTERFACE_POSITION_INDEPENDENT_CODE ON -) - -target_compile_options( - rapidsmpfndsh PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" -) -target_link_libraries( - rapidsmpfndsh - PRIVATE rapidsmpf::rapidsmpf cudf_streaming::cudf_streaming cuco::cuco - $ $ maybe_asan -) - -set(RAPIDSMPFNDSH_QUERIES q01 q03 q04 q09 q21 bench_read) - -foreach(query IN ITEMS ${RAPIDSMPFNDSH_QUERIES}) - add_executable(${query} "${query}.cpp") - set_target_properties( - ${query} - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks/ndsh" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON - ) - target_compile_options( - ${query} PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" - ) - target_link_libraries( - ${query} - PRIVATE rapidsmpfndsh rapidsmpf::rapidsmpf cudf_streaming::cudf_streaming - $ $ maybe_asan - ) -endforeach() - -install( - TARGETS rapidsmpfndsh - COMPONENT benchmarking - DESTINATION ${lib_dir} - EXCLUDE_FROM_ALL -) -install( - TARGETS ${RAPIDSMPFNDSH_QUERIES} - COMPONENT benchmarking - DESTINATION bin/benchmarks/librapidsmpf - EXCLUDE_FROM_ALL -) diff --git a/cpp/benchmarks/streaming/ndsh/bench_read.cpp b/cpp/benchmarks/streaming/ndsh/bench_read.cpp deleted file mode 100644 index 5e1125f9b..000000000 --- a/cpp/benchmarks/streaming/ndsh/bench_read.cpp +++ /dev/null @@ -1,438 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include - -#include "utils.hpp" - -namespace { - -rapidsmpf::streaming::Actor read_parquet( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::optional> columns, - std::string const& input_directory, - std::string const& input_file -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, input_file) - ); - auto options = - cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)).build(); - if (columns.has_value()) { - options.set_column_names(*columns); - } - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor consume_channel_parallel( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::size_t num_consumers -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in}; - std::atomic estimated_total_bytes{0}; - auto task = [&]() -> rapidsmpf::streaming::Actor { - co_await ctx->executor()->schedule(); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - if (msg.holds()) { - auto chunk = co_await msg.release() - .make_available(ctx); - ctx->logger()->print( - "Consumed chunk with ", - chunk.table_view().num_rows(), - " rows and ", - chunk.table_view().num_columns(), - " columns" - ); - estimated_total_bytes.fetch_add( - chunk.data_alloc_size(rapidsmpf::MemoryType::DEVICE) - ); - } - } - }; - std::vector tasks; - for (std::size_t i = 0; i < num_consumers; i++) { - tasks.push_back(task()); - } - rapidsmpf::streaming::coro_results(co_await coro::when_all(std::move(tasks))); - ctx->logger()->print( - "Table was around ", rmm::detail::format_bytes(estimated_total_bytes.load()) - ); -} - -///< @brief Configuration options for the benchmark -struct ProgramOptions { - int num_streaming_threads{1}; ///< Number of streaming threads to use - int num_iterations{2}; ///< Number of iterations of query to run - int num_streams{16}; ///< Number of streams in stream pool - rapidsmpf::ndsh::CommType comm_type{ - rapidsmpf::ndsh::CommType::UCXX - }; ///< Type of communicator to create - cudf::size_type num_rows_per_chunk{ - 100'000'000 - }; ///< Number of rows to produce per chunk read - std::size_t num_producers{ - 1 - }; ///< Number of simultaneous read_parquet chunk producers. - std::size_t num_consumers{1}; ///< Number of simultaneous chunk consumers. - std::string input_directory; ///< Directory containing input files. - std::string input_file; ///< Basename of input file to read. - std::optional> columns{std::nullopt}; ///< Columns to read. -}; - -ProgramOptions parse_arguments(int argc, char** argv) { - ProgramOptions options; - - static constexpr std:: - array(rapidsmpf::ndsh::CommType::MAX)> - comm_names{"single", "mpi", "ucxx"}; - - auto print_usage = [&argv, &options]() { - std::cerr - << "Usage: " << argv[0] << " [options]\n" - << "Options:\n" - << " --num-streaming-threads Number of streaming threads (default: " - << options.num_streaming_threads << ")\n" - << " --num-iterations Number of iterations (default: " - << options.num_iterations << ")\n" - << " --num-streams Number of streams in stream pool " - "(default: " - << options.num_streams << ")\n" - << " --num-rows-per-chunk Number of rows per chunk (default: " - << options.num_rows_per_chunk << ")\n" - << " --num-producers Number of concurrent read_parquet " - "producers (default: " - << options.num_producers << ")\n" - << " --num-consumers Number of concurrent consumers (default: " - << options.num_consumers << ")\n" - << " --comm-type Communicator type: single, mpi, ucxx " - "(default: " - << comm_names[static_cast(options.comm_type)] << ")\n" - << " --input-directory Input directory path (required)\n" - << " --input-file Input file basename relative to input " - "directory (required)\n" - << " --columns Comma-separated column names to read " - "(optional, default all columns)\n" - << " --help Show this help message\n"; - }; - - // NOLINTBEGIN(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays,modernize-use-designated-initializers) - static struct option long_options[] = { - {"num-streaming-threads", required_argument, nullptr, 1}, - {"num-rows-per-chunk", required_argument, nullptr, 2}, - {"num-producers", required_argument, nullptr, 3}, - {"num-consumers", required_argument, nullptr, 4}, - {"input-directory", required_argument, nullptr, 5}, - {"input-file", required_argument, nullptr, 6}, - {"help", no_argument, nullptr, 7}, - {"num-iterations", required_argument, nullptr, 8}, - {"num-streams", required_argument, nullptr, 9}, - {"comm-type", required_argument, nullptr, 10}, - {"columns", required_argument, nullptr, 11}, - {nullptr, 0, nullptr, 0} - }; - // NOLINTEND(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays,modernize-use-designated-initializers) - - int opt; - int option_index = 0; - - bool saw_input_directory = false; - bool saw_input_file = false; - - auto parse_i64 = [](char const* s, char const* opt_name) -> long long { - if (s == nullptr || *s == '\0') { - std::cerr << "Error: " << opt_name << " requires a value\n"; - std::exit(1); - } - errno = 0; - char* end = nullptr; - auto const v = std::strtoll(s, &end, 10); - if (errno != 0 || end == s || *end != '\0') { - std::cerr << "Error: invalid integer for " << opt_name << ": '" << s << "'\n"; - std::exit(1); - } - return v; - }; - - auto parse_u64 = [](char const* s, char const* opt_name) -> unsigned long long { - if (s == nullptr || *s == '\0') { - std::cerr << "Error: " << opt_name << " requires a value\n"; - std::exit(1); - } - errno = 0; - char* end = nullptr; - auto const v = std::strtoull(s, &end, 10); - if (errno != 0 || end == s || *end != '\0') { - std::cerr << "Error: invalid non-negative integer for " << opt_name << ": '" - << s << "'\n"; - std::exit(1); - } - return v; - }; - - auto require_positive_i32 = [&](char const* s, char const* opt_name) -> int { - auto const v = parse_i64(s, opt_name); - if (v <= 0 || v > std::numeric_limits::max()) { - std::cerr << "Error: " << opt_name << " must be in [1, " - << std::numeric_limits::max() << "], got '" << s << "'\n"; - std::exit(1); - } - return static_cast(v); - }; - - auto require_positive_size_t = [&](char const* s, - char const* opt_name) -> std::size_t { - auto const v = parse_u64(s, opt_name); - if (v == 0 || v > std::numeric_limits::max()) { - std::cerr << "Error: " << opt_name << " must be in [1, " - << std::numeric_limits::max() << "], got '" << s - << "'\n"; - std::exit(1); - } - return static_cast(v); - }; - - auto parse_columns = [](char const* s) -> std::optional> { - if (s == nullptr) { - return std::nullopt; - } - std::string str{s}; - if (str.empty()) { - return std::nullopt; - } - std::vector cols; - std::size_t start = 0; - while (start <= str.size()) { - auto const comma = str.find(',', start); - auto const end = (comma == std::string::npos) ? str.size() : comma; - auto const token = str.substr(start, end - start); - if (token.empty()) { - std::cerr << "Error: --columns contains an empty column name\n"; - std::exit(1); - } - cols.push_back(token); - if (comma == std::string::npos) { - break; - } - start = comma + 1; - } - return cols; - }; - - while ((opt = getopt_long(argc, argv, "", long_options, &option_index)) != -1) { - switch (opt) { - case 1: // --num-streaming-threads - options.num_streaming_threads = - require_positive_i32(optarg, "--num-streaming-threads"); - break; - case 2: // --num-rows-per-chunk - options.num_rows_per_chunk = - require_positive_i32(optarg, "--num-rows-per-chunk"); - break; - case 3: // --num-producers - options.num_producers = require_positive_size_t(optarg, "--num-producers"); - break; - case 4: // --num-consumers - options.num_consumers = require_positive_size_t(optarg, "--num-consumers"); - break; - case 5: // --input-directory - if (optarg == nullptr || *optarg == '\0') { - std::cerr << "Error: --input-directory requires a non-empty value\n"; - std::exit(1); - } - options.input_directory = optarg; - saw_input_directory = true; - break; - case 6: // --input-file - if (optarg == nullptr || *optarg == '\0') { - std::cerr << "Error: --input-file requires a non-empty value\n"; - std::exit(1); - } - options.input_file = optarg; - saw_input_file = true; - break; - case 7: // --help - print_usage(); - std::exit(0); - case 8: // --num-iterations - options.num_iterations = require_positive_i32(optarg, "--num-iterations"); - break; - case 9: // --num-streams - options.num_streams = require_positive_i32(optarg, "--num-streams"); - break; - case 10: - { // --comm-type - if (optarg == nullptr || *optarg == '\0') { - std::cerr << "Error: --comm-type requires a value\n"; - std::exit(1); - } - std::string_view const s{optarg}; - auto parsed = std::optional{}; - for (std::size_t i = 0; i < comm_names.size(); ++i) { - if (s == comm_names[i]) { - parsed = static_cast(i); - break; - } - } - if (!parsed.has_value()) { - std::cerr << "Error: invalid --comm-type '" << s - << "' (expected: single, mpi, ucxx)\n"; - std::exit(1); - } - options.comm_type = *parsed; - break; - } - case 11: // --columns - options.columns = parse_columns(optarg); - break; - case '?': - if (optopt == 0 && optind > 1) { - std::cerr << "Error: Unknown option '" << argv[optind - 1] << "'\n\n"; - } - print_usage(); - std::exit(1); - default: - print_usage(); - std::exit(1); - } - } - - // Check if required options were provided - if (!saw_input_directory || !saw_input_file) { - if (!saw_input_directory) { - std::cerr << "Error: --input-directory is required\n"; - } - if (!saw_input_file) { - std::cerr << "Error: --input-file is required\n"; - } - std::cerr << std::endl; - print_usage(); - std::exit(1); - } - - return options; -} - -} // namespace - -/** - * @brief Run a simple benchmark reading a table from parquet files. - */ -int main(int argc, char** argv) { - rapidsmpf::ndsh::FinalizeMPI finalize{}; - cudaFree(nullptr); - // work around https://github.com/rapidsai/cudf/issues/20849 - cudf::initialize(); - auto mr = rmm::mr::cuda_async_memory_resource{}; - auto arguments = parse_arguments(argc, argv); - rapidsmpf::ndsh::ProgramOptions ctx_arguments{ - .num_streaming_threads = arguments.num_streaming_threads, - .num_iterations = arguments.num_iterations, - .num_streams = arguments.num_streams, - .comm_type = arguments.comm_type, - .num_rows_per_chunk = arguments.num_rows_per_chunk, - .output_file = "", - .input_directory = arguments.input_directory - }; - - auto [ctx, comm] = rapidsmpf::ndsh::create_context(ctx_arguments, std::move(mr)); - std::vector timings; - for (int i = 0; i < arguments.num_iterations; i++) { - std::vector actors; - auto start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Constructing read_parquet pipeline"); - - // Input data channels - auto ch_out = ctx->create_channel(); - actors.push_back(read_parquet( - ctx, - comm, - ch_out, - arguments.num_producers, - arguments.num_rows_per_chunk, - arguments.columns, - arguments.input_directory, - arguments.input_file - )); - actors.push_back( - consume_channel_parallel(ctx, ch_out, arguments.num_consumers) - ); - } - auto end = std::chrono::steady_clock::now(); - std::chrono::duration pipeline = end - start; - start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("read_parquet iteration"); - rapidsmpf::streaming::run_actor_network(std::move(actors)); - } - end = std::chrono::steady_clock::now(); - std::chrono::duration compute = end - start; - timings.push_back(pipeline.count()); - timings.push_back(compute.count()); - auto statistics = ctx->statistics(); - comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} - )); - statistics->clear(); - } - - if (comm->rank() == 0) { - for (int i = 0; i < arguments.num_iterations; i++) { - comm->logger()->print( - "Iteration ", - i, - " pipeline construction time [s]: ", - timings[rapidsmpf::safe_cast(2 * i)] - ); - comm->logger()->print( - "Iteration ", - i, - " compute time [s]: ", - timings[rapidsmpf::safe_cast(2 * i + 1)] - ); - } - } - return 0; -} diff --git a/cpp/benchmarks/streaming/ndsh/concatenate.cpp b/cpp/benchmarks/streaming/ndsh/concatenate.cpp deleted file mode 100644 index e80799800..000000000 --- a/cpp/benchmarks/streaming/ndsh/concatenate.cpp +++ /dev/null @@ -1,88 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "concatenate.hpp" - -#include -#include - -#include -#include -#include -#include - -#include -#include -#include -#include -#include - -namespace rapidsmpf::ndsh { - - -streaming::Actor concatenate( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - ConcatOrder order -) { - streaming::ShutdownAtExit c{ch_in, ch_out}; - CudaEvent event; - std::vector messages; - ctx->logger()->print("Concatenate"); - auto concat_stream = ctx->br()->stream_pool().get_stream(); - while (!ch_out->is_shutdown()) { - co_await ctx->executor()->schedule(); - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - messages.push_back(std::move(msg)); - } - if (messages.size() == 0) { - co_await ch_out->send( - cudf_streaming::streaming::to_message( - 0, - std::make_unique( - std::make_unique(), concat_stream - ) - ) - ); - } else if (messages.size() == 1) { - co_await ch_out->send(std::move(messages[0])); - } else { - std::vector chunks; - std::vector views; - if (order == ConcatOrder::LINEARIZE) { - std::ranges::sort(messages, std::less{}, [](auto&& msg) { - return msg.sequence_number(); - }); - } - chunks.reserve(messages.size()); - views.reserve(messages.size()); - for (auto&& msg : messages) { - auto chunk = co_await msg.release() - .make_available(ctx); - cuda_stream_join(concat_stream, chunk.stream(), &event); - views.push_back(chunk.table_view()); - chunks.push_back(std::move(chunk)); - } - auto result = std::make_unique( - cudf::concatenate(views, concat_stream, ctx->br()->device_mr()), concat_stream - ); - cuda_stream_join( - chunks | std::views::transform([](auto&& chunk) { return chunk.stream(); }), - std::ranges::single_view(concat_stream), - &event - ); - chunks.clear(); - co_await ch_out->send( - cudf_streaming::streaming::to_message(0, std::move(result)) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/concatenate.hpp b/cpp/benchmarks/streaming/ndsh/concatenate.hpp deleted file mode 100644 index 47450eb12..000000000 --- a/cpp/benchmarks/streaming/ndsh/concatenate.hpp +++ /dev/null @@ -1,38 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include - -#include -#include -#include - -namespace rapidsmpf::ndsh { - -///< @brief Should the concatenation respect input ordering? -enum class ConcatOrder : bool { - DONT_CARE, ///< No, we don't need ordering - LINEARIZE, ///< Yes, maintain input ordering -}; - -/** - * @brief Concatenate all table chunks from an input channel. - * - * @param ctx Streaming context. - * @param ch_in Input channel of `TableChunk`s. - * @param ch_out Output channel of concatenated chunks, contains at most one message. - * @param order Do we care about maintaining the input ordering? - * - * @return Coroutine representing the concatenation. - */ -streaming::Actor concatenate( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - ConcatOrder order = ConcatOrder::DONT_CARE -); - -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/groupby.cpp b/cpp/benchmarks/streaming/ndsh/groupby.cpp deleted file mode 100644 index 6c6337218..000000000 --- a/cpp/benchmarks/streaming/ndsh/groupby.cpp +++ /dev/null @@ -1,78 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "groupby.hpp" - -#include -#include -#include - -#include -#include -#include -#include - -#include -#include -#include - -namespace rapidsmpf::ndsh { - -streaming::Actor chunkwise_group_by( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys, - std::vector requests, - cudf::null_policy null_policy -) { - streaming::ShutdownAtExit c{ch_in, ch_out}; - co_await ctx->executor()->schedule(); - while (!ch_out->is_shutdown()) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - auto stream = chunk.stream(); - auto table = chunk.table_view(); - auto agg_requests = std::vector(); - agg_requests.reserve(requests.size()); - std::ranges::transform( - requests, std::back_inserter(agg_requests), [&table](auto&& req) { - std::vector> reqs; - for (auto&& x : req.requests) { - reqs.push_back(x()); - } - return cudf::groupby::aggregation_request{ - table.column(req.column_idx), std::move(reqs) - }; - } - ); - auto grouper = - cudf::groupby::groupby(table.select(keys), null_policy, cudf::sorted::NO); - - auto [keys, aggregated] = - grouper.aggregate(agg_requests, stream, ctx->br()->device_mr()); - std::ignore = std::move(chunk); - auto result = keys->release(); - for (auto&& a : aggregated) { - std::ranges::move(a.results, std::back_inserter(result)); - } - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::make_unique(std::move(result)), stream - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/groupby.hpp b/cpp/benchmarks/streaming/ndsh/groupby.hpp deleted file mode 100644 index 000bdf322..000000000 --- a/cpp/benchmarks/streaming/ndsh/groupby.hpp +++ /dev/null @@ -1,51 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include -#include -#include - -#include -#include - -#include -#include -#include - -namespace rapidsmpf::ndsh { - -///< @brief Description of aggregation requests on a given column -struct groupby_request { - cudf::size_type column_idx; ///< Index of column in input table to aggregate - std::vector()>> - requests; ///< Functions to generate aggregations to perform on the column -}; - -/** - * @brief Perform a chunkwise grouped aggregation. - * - * @note Grouped chunks are not further grouped together. - * - * @param ctx Streaming context. - * @param ch_in `TableChunk`s to aggregate - * @param ch_out Output channel of grouped `TableChunk`s - * @param keys Column indices of the key columns in the input channel. - * @param requests Vector of aggregation requests referencing columns in the input - * channel. - * @param null_policy How nulls in the key columns are treated. - * - * @return Coroutine representing the completion of the aggregation. - */ -streaming::Actor chunkwise_group_by( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys, - std::vector requests, - cudf::null_policy null_policy - -); -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/join.cpp b/cpp/benchmarks/streaming/ndsh/join.cpp deleted file mode 100644 index ab7c46519..000000000 --- a/cpp/benchmarks/streaming/ndsh/join.cpp +++ /dev/null @@ -1,620 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "join.hpp" - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace rapidsmpf::ndsh { - -coro::task broadcast( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - OpID tag, - streaming::AllGather::Ordered ordered -) { - streaming::ShutdownAtExit c{ch_in}; - co_await ctx->executor()->schedule(); - CudaEvent event; - comm->logger()->print("Broadcast ", static_cast(tag)); - if (comm->nranks() == 1) { - std::vector chunks; - std::vector views; - auto gather_stream = ctx->br()->stream_pool().get_stream(); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = co_await msg.release() - .make_available(ctx); - cuda_stream_join(gather_stream, chunk.stream(), &event); - views.push_back(chunk.table_view()); - chunks.push_back(std::move(chunk)); - } - if (chunks.size() == 1) { - co_return cudf_streaming::streaming::to_message( - 0, - std::make_unique( - std::move(chunks[0]) - ) - ); - } else { - RAPIDSMPF_EXPECTS(chunks.size() > 0, "No chunks in broadcast"); - auto result = cudf::concatenate(views, gather_stream, ctx->br()->device_mr()); - // So that deallocation of the consitutent tables is stream-ordered wrt the - // concatenation. - cuda_stream_join( - chunks - | std::views::transform([](auto&& chunk) { return chunk.stream(); }), - std::ranges::single_view(gather_stream), - &event - ); - co_return cudf_streaming::streaming::to_message( - 0, - std::make_unique( - std::move(result), gather_stream - ) - ); - } - } else { - streaming::AllGather gatherer{ctx, comm, tag}; - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - // TODO: If this chunk is already in pack form, this is unnecessary. - auto chunk = co_await msg.release() - .make_available(ctx); - auto pack = - cudf::pack(chunk.table_view(), chunk.stream(), ctx->br()->device_mr()); - auto packed_data = PackedData( - std::move(pack.metadata), - ctx->br()->move(std::move(pack.gpu_data), chunk.stream()) - ); - gatherer.insert(msg.sequence_number(), {std::move(packed_data)}); - } - gatherer.insert_finished(); - auto result = co_await gatherer.extract_all(ordered); - if (result.size() == 1) { - co_return cudf_streaming::streaming::to_message( - 0, - std::make_unique( - std::make_unique(std::move(result[0])) - ) - ); - } else { - auto stream = ctx->br()->stream_pool().get_stream(); - co_return cudf_streaming::streaming::to_message( - 0, - std::make_unique( - cudf_streaming::integrations::unpack_and_concat( - rapidsmpf::unspill_partitions( - std::move(result), ctx->br().get(), AllowOverbooking::YES - ), - stream, - ctx->br().get() - ), - stream - ) - ); - } - } -} - -streaming::Actor broadcast( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - OpID tag, - streaming::AllGather::Ordered ordered -) { - streaming::ShutdownAtExit c{ch_in, ch_out}; - co_await ctx->executor()->schedule(); - co_await ch_out->send(co_await broadcast(ctx, comm, ch_in, tag, ordered)); - co_await ch_out->drain(ctx->executor()); -} - -/** - * @brief Join a table chunk against a build hash table returning a message of the result. - * - * @param ctx Streaming context - * @param left_chunk Chunk to join. Used as the probe table in a filtered join. - * @param right_chunk Chunk to join. Used as the build table in a filtered join. - * @param left_carrier Columns from `left_chunk` to include in the output. - * @param left_on Key column indices in `left_chunk`. - * @param right_on Key column indices in `right_chunk`. - * @param sequence Sequence number of the output - * @param left_event Event recording the availability of `left_chunk`. - * - * @return Message of `TableChunk` containing the result of the semi join. - */ -streaming::Message semi_join_chunk( - std::shared_ptr ctx, - cudf_streaming::streaming::TableChunk const& left_chunk, - cudf_streaming::streaming::TableChunk&& right_chunk, - cudf::table_view left_carrier, - std::vector left_on, - std::vector right_on, - std::uint64_t sequence, - CudaEvent* left_event -) { - auto chunk_stream = right_chunk.stream(); - - left_event->stream_wait(chunk_stream); - - // At this point, both left_chunk and right_chunk are valid on - // either stream. We'll do everything from here out on the - // right_chunk.stream(), so that we don't introduce false dependencies - // between the different chunks. - - auto joiner = cudf::filtered_join( - right_chunk.table_view().select(right_on), - cudf::null_equality::UNEQUAL, - chunk_stream - ); - - auto match = joiner.semi_join( - left_chunk.table_view().select(left_on), chunk_stream, ctx->br()->device_mr() - ); - - ctx->logger()->debug( - "semi_join_chunk: left.num_rows()=", left_chunk.table_view().num_rows() - ); - ctx->logger()->debug("semi_join_chunk: match.size()=", match->size()); - - cudf::column_view indices = cudf::device_span(*match); - auto result_columns = cudf::gather( - left_carrier, - indices, - cudf::out_of_bounds_policy::DONT_CHECK, - chunk_stream, - ctx->br()->device_mr() - ) - ->release(); - - auto result_table = std::make_unique(std::move(result_columns)); - // Deallocation of the join indices will happen on chunk_stream, so add stream dep - cuda_stream_join(left_chunk.stream(), chunk_stream); - - return cudf_streaming::streaming::to_message( - sequence, - std::make_unique( - std::move(result_table), chunk_stream - ) - ); -} - -/** - * @brief Join a table chunk against a build hash table returning a message of the result. - * - * @param ctx Streaming context. - * @param right_chunk Chunk to join. Must be on device e.g. use make_available() on the - * chunk. - * @param sequence Sequence number of the output - * @param joiner hash_join object, representing the build table. - * @param build_carrier Columns from the build-side table to be included in the output. - * @param right_on Key column indiecs in `right_chunk`. - * @param build_stream Stream the `joiner` will be deallocated on. - * @param build_event Event recording the creation of the `joiner`. - * @param tmp_event Preallocated event used for internal stream ordering. - * - * @return Message of `TableChunk` containing the result of the inner join. - */ -streaming::Message inner_join_chunk( - std::shared_ptr ctx, - cudf_streaming::streaming::TableChunk&& right_chunk, - std::uint64_t sequence, - cudf::hash_join& joiner, - cudf::table_view build_carrier, - std::vector right_on, - rmm::cuda_stream_view build_stream, - CudaEvent* build_event, - CudaEvent* tmp_event - -) { - auto chunk_stream = right_chunk.stream(); - build_event->stream_wait(chunk_stream); - auto probe_table = right_chunk.table_view(); - auto probe_keys = probe_table.select(right_on); - auto [probe_match, build_match] = - joiner.inner_join(probe_keys, std::nullopt, chunk_stream, ctx->br()->device_mr()); - - cudf::column_view build_indices = - cudf::device_span(*build_match); - cudf::column_view probe_indices = - cudf::device_span(*probe_match); - // build_carrier is valid on build_stream, but chunk_stream is - // waiting for build_stream work to be done, so running this on - // chunk_stream is fine. - auto result_columns = cudf::gather( - build_carrier, - build_indices, - cudf::out_of_bounds_policy::DONT_CHECK, - chunk_stream, - ctx->br()->device_mr() - ) - ->release(); - // drop key columns from probe table. - std::vector to_keep; - std::ranges::copy_if( - std::ranges::iota_view(0, probe_table.num_columns()), - std::back_inserter(to_keep), - [&](auto i) { return std::ranges::find(right_on, i) == right_on.end(); } - ); - std::ranges::move( - cudf::gather( - probe_table.select(to_keep), - probe_indices, - cudf::out_of_bounds_policy::DONT_CHECK, - chunk_stream, - ctx->br()->device_mr() - ) - ->release(), - std::back_inserter(result_columns) - ); - // Deallocation of the join indices will happen on build_stream, so add stream dep - // This also ensure deallocation of the hash_join object waits for completion. - cuda_stream_join(build_stream, chunk_stream, tmp_event); - return cudf_streaming::streaming::to_message( - sequence, - std::make_unique( - std::make_unique(std::move(result_columns)), chunk_stream - ) - ); -} - -streaming::Actor inner_join_broadcast( - std::shared_ptr ctx, - std::shared_ptr comm, - // We will always choose left as build table and do "broadcast" joins - std::shared_ptr left, - std::shared_ptr right, - std::shared_ptr ch_out, - std::vector left_on, - std::vector right_on, - OpID tag, - KeepKeys keep_keys -) { - streaming::ShutdownAtExit c{left, right, ch_out}; - co_await ctx->executor()->schedule(); - comm->logger()->print("Inner broadcast join ", static_cast(tag)); - auto build_table = co_await ( - (co_await broadcast(ctx, comm, left, tag, streaming::AllGather::Ordered::NO)) - .release() - .make_available(ctx) - ); - comm->logger()->print( - "Build table has ", build_table.table_view().num_rows(), " rows" - ); - - auto joiner = cudf::hash_join( - build_table.table_view().select(left_on), - cudf::null_equality::UNEQUAL, - build_table.stream() - ); - CudaEvent build_event; - build_event.record(build_table.stream()); - CudaEvent tmp_event; - cudf::table_view build_carrier; - if (keep_keys == KeepKeys::YES) { - build_carrier = build_table.table_view(); - } else { - std::vector to_keep; - std::ranges::copy_if( - std::ranges::iota_view(0, build_table.table_view().num_columns()), - std::back_inserter(to_keep), - [&](auto i) { return std::ranges::find(left_on, i) == left_on.end(); } - ); - build_carrier = build_table.table_view().select(to_keep); - } - while (!ch_out->is_shutdown()) { - auto right_msg = co_await right->receive(); - if (right_msg.empty()) { - break; - } - co_await ch_out->send(inner_join_chunk( - ctx, - right_msg.release(), - right_msg.sequence_number(), - joiner, - build_carrier, - right_on, - build_table.stream(), - &build_event, - &tmp_event - )); - } - - co_await ch_out->drain(ctx->executor()); -} - -streaming::Actor inner_join_shuffle( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr left, - std::shared_ptr right, - std::shared_ptr ch_out, - std::vector left_on, - std::vector right_on, - KeepKeys keep_keys -) { - streaming::ShutdownAtExit c{left, right, ch_out}; - comm->logger()->print("Inner shuffle join"); - co_await ctx->executor()->schedule(); - CudaEvent build_event; - CudaEvent tmp_event; - while (!ch_out->is_shutdown()) { - // Requirement: two shuffles kick out partitions in the same order - auto left_msg = co_await left->receive(); - auto right_msg = co_await right->receive(); - if (left_msg.empty()) { - RAPIDSMPF_EXPECTS( - right_msg.empty(), "Left does not have same number of partitions as right" - ); - break; - } - RAPIDSMPF_EXPECTS( - left_msg.sequence_number() == right_msg.sequence_number(), - "Mismatching sequence numbers" - ); - // TODO: currently always using left as build table. - auto build_chunk = - co_await left_msg.release() - .make_available(ctx); - auto build_stream = build_chunk.stream(); - auto joiner = cudf::hash_join( - build_chunk.table_view().select(left_on), - cudf::null_equality::UNEQUAL, - build_stream - ); - build_event.record(build_stream); - cudf::table_view build_carrier; - if (keep_keys == KeepKeys::YES) { - build_carrier = build_chunk.table_view(); - } else { - std::vector to_keep; - std::ranges::copy_if( - std::ranges::iota_view(0, build_chunk.table_view().num_columns()), - std::back_inserter(to_keep), - [&](auto i) { return std::ranges::find(left_on, i) == left_on.end(); } - ); - build_carrier = build_chunk.table_view().select(to_keep); - } - co_await ch_out->send(inner_join_chunk( - ctx, - right_msg.release(), - left_msg.sequence_number(), - joiner, - build_carrier, - right_on, - build_stream, - &build_event, - &tmp_event - )); - } - co_await ch_out->drain(ctx->executor()); -} - -streaming::Actor left_semi_join_broadcast_left( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr left, - std::shared_ptr right, - std::shared_ptr ch_out, - std::vector left_on, - std::vector right_on, - OpID tag, - KeepKeys keep_keys -) { - streaming::ShutdownAtExit c{left, right, ch_out}; - co_await ctx->executor()->schedule(); - comm->logger()->print("Left semi broadcast join ", static_cast(tag)); - auto left_table = co_await (co_await broadcast(ctx, comm, left, tag)) - .release() - .make_available(ctx); - comm->logger()->print( - "Left (probe) table has ", left_table.table_view().num_rows(), " rows" - ); - CudaEvent left_event; - left_event.record(left_table.stream()); - - cudf::table_view left_carrier; - if (keep_keys == KeepKeys::YES) { - left_carrier = left_table.table_view(); - } else { - std::vector to_keep; - std::ranges::copy_if( - std::ranges::iota_view(0, left_table.table_view().num_columns()), - std::back_inserter(to_keep), - [&](auto i) { return std::ranges::find(left_on, i) == left_on.end(); } - ); - left_carrier = left_table.table_view().select(to_keep); - } - - while (!ch_out->is_shutdown()) { - auto right_msg = co_await right->receive(); - if (right_msg.empty()) { - break; - } - // The ``right`` table has been hash-partitioned (via a shuffle) on - // the join key. Thanks to the hash-partitioning, we don't need to worry - // about deduplicating matches across partitions. Anything that matches - // in the semi-join belongs in the output. - auto right_chunk = - co_await right_msg.release() - .make_available(ctx); - co_await ch_out->send(semi_join_chunk( - ctx, - left_table, - std::move(right_chunk), - left_carrier, - left_on, - right_on, - right_msg.sequence_number(), - &left_event - )); - } - - co_await ch_out->drain(ctx->executor()); -} - -streaming::Actor left_semi_join_shuffle( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr left, - std::shared_ptr right, - std::shared_ptr ch_out, - std::vector left_on, - std::vector right_on, - KeepKeys keep_keys -) { - streaming::ShutdownAtExit c{left, right, ch_out}; - comm->logger()->print("Shuffle left semi join"); - - co_await ctx->executor()->schedule(); - CudaEvent left_event; - - while (!ch_out->is_shutdown()) { - // Requirement: two shuffles kick out partitions in the same order - auto left_msg = co_await left->receive(); - auto right_msg = co_await right->receive(); - - if (left_msg.empty()) { - RAPIDSMPF_EXPECTS( - right_msg.empty(), "Left does not have same number of partitions as right" - ); - break; - } - RAPIDSMPF_EXPECTS( - left_msg.sequence_number() == right_msg.sequence_number(), - "Mismatching sequence numbers" - ); - - auto left_chunk = - co_await left_msg.release() - .make_available(ctx); - auto right_chunk = - co_await right_msg.release() - .make_available(ctx); - - left_event.record(left_chunk.stream()); - - cudf::table_view left_carrier; - if (keep_keys == KeepKeys::YES) { - left_carrier = left_chunk.table_view(); - } else { - std::vector to_keep; - std::ranges::copy_if( - std::ranges::iota_view(0, left_chunk.table_view().num_columns()), - std::back_inserter(to_keep), - [&](auto i) { return std::ranges::find(left_on, i) == left_on.end(); } - ); - left_carrier = left_chunk.table_view().select(to_keep); - } - - co_await ch_out->send(semi_join_chunk( - ctx, - left_chunk, - std::move(right_chunk), - left_carrier, - left_on, - right_on, - left_msg.sequence_number(), - &left_event - )); - } -} - -streaming::Actor shuffle( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys, - std::uint32_t num_partitions, - OpID tag -) { - streaming::ShutdownAtExit c{ch_in, ch_out}; - co_await ctx->executor()->schedule(); - comm->logger()->print("Shuffle ", static_cast(tag)); - streaming::ShufflerAsync shuffler(ctx, comm, tag, num_partitions); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - comm->logger()->debug("Shuffle: no more input"); - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - auto packed = cudf_streaming::integrations::partition_and_pack( - chunk.table_view(), - keys, - static_cast(num_partitions), - cudf::hash_id::HASH_MURMUR3, - 0, - chunk.stream(), - ctx->br().get() - ); - shuffler.insert(std::move(packed)); - } - co_await shuffler.insert_finished(); - for (auto pid : shuffler.local_partitions()) { - auto packed_data = shuffler.extract(pid); - auto stream = ctx->br()->stream_pool().get_stream(); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - pid, - std::make_unique( - cudf_streaming::integrations::unpack_and_concat( - rapidsmpf::unspill_partitions( - std::move(packed_data), ctx->br().get(), AllowOverbooking::YES - ), - stream, - ctx->br().get() - ), - stream - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/join.hpp b/cpp/benchmarks/streaming/ndsh/join.hpp deleted file mode 100644 index 50339ee52..000000000 --- a/cpp/benchmarks/streaming/ndsh/join.hpp +++ /dev/null @@ -1,219 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include -#include -#include - -#include - -#include -#include -#include -#include - -namespace rapidsmpf::ndsh { -///< @brief Treatment of keys in the result of a join -enum class KeepKeys : bool { - NO, ///< Key columns do not appear in the output - YES, ///< Key columns do appear in the output -}; - -/** - * @brief Broadcast the concatenation of all input messages to all ranks. - * - * @note Receives all input chunks, gathers from all ranks, and then provides concatenated - * output. - * - * @param ctx Streaming context - * @param comm Communicator for the collective operation. - * @param ch_in Input channel of `TableChunk`s - * @param tag Disambiguating tag for allgather - * @param ordered Should the concatenated output be ordered - * - * @return Message containing the concatenation of all the input table chunks. - */ -[[nodiscard]] coro::task broadcast( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - OpID tag, - streaming::AllGather::Ordered ordered = streaming::AllGather::Ordered::YES -); - -/** - * @brief Broadcast the concatenation of all input messages to all ranks. - * - * @note Receives all input chunks, gathers from all ranks, and then provides concatenated - * output. - * - * @param ctx Streaming context - * @param comm Communicator for the collective operation. - * @param ch_in Input channel of `TableChunk`s - * @param ch_out Input channel of a single `TableChunk` - * @param tag Disambiguating tag for allgather - * @param ordered Should the concatenated output be ordered - * - * @return Coroutine representing the broadcast - */ -[[nodiscard]] streaming::Actor broadcast( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - OpID tag, - streaming::AllGather::Ordered ordered = streaming::AllGather::Ordered::YES -); - -/** - * @brief Perform a streaming inner join between two tables. - * - * @note This performs a broadcast join, broadcasting the table represented by the `left` - * channel to all ranks, and then streaming through the chunks of the `right` channel. - * - * @param ctx Streaming context. - * @param comm Communicator for the collective operation. - * @param left Channel of `TableChunk`s used as the broadcasted build side. - * @param right Channel of `TableChunk`s joined in turn against the build side. - * @param ch_out Output channel of `TableChunk`s. - * @param left_on Column indices of the keys in the left table. - * @param right_on Column indices of the keys in the right table. - * @param tag Disambiguating tag for the broadcast of the left table. - * @param keep_keys Does the result contain the key columns, or only "carrier" value - * columns - * - * @return Coroutine representing the completion of the join. - */ -[[nodiscard]] streaming::Actor inner_join_broadcast( - std::shared_ptr ctx, - std::shared_ptr comm, - // We will always choose left as build table and do "broadcast" joins - std::shared_ptr left, - std::shared_ptr right, - std::shared_ptr ch_out, - std::vector left_on, - std::vector right_on, - OpID tag, - KeepKeys keep_keys = KeepKeys::YES -); -/** - * @brief Perform a streaming inner join between two tables. - * - * @note This performs a shuffle join, the left and right channels are required to provide - * hash-partitioned data in-order. - * - * @param ctx Streaming context. - * @param comm Communicator for the collective operation. - * @param left Channel of `TableChunk`s in hash-partitioned order. - * @param right Channel of `TableChunk`s in matching hash-partitioned order. - * @param ch_out Output channel of `TableChunk`s. - * @param left_on Column indices of the keys in the left table. - * @param right_on Column indices of the keys in the right table. - * @param keep_keys Does the result contain the key columns, or only "carrier" value - * columns - * - * @return Coroutine representing the completion of the join. - */ -[[nodiscard]] streaming::Actor inner_join_shuffle( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr left, - std::shared_ptr right, - std::shared_ptr ch_out, - std::vector left_on, - std::vector right_on, - KeepKeys keep_keys = KeepKeys::YES -); - -/** - * @brief Perform a streaming left semi join between two tables. - * - * @note This performs a broadcast join, broadcasting the table represented by the `left` - * channel to all ranks, and then streaming through the chunks of the `right` channel. - * The `right` channel is required to provide hash-partitioned data in-order. - * All of the chunks from the `left` channel must fit in memory at once. - * - * @param ctx Streaming context. - * @param comm Communicator for the collective operation. - * @param left Channel of `TableChunk`s. - * @param right Channel of `TableChunk`s in hash-partitioned order (shuffled). - * @param ch_out Output channel of `TableChunk`s. - * @param left_on Column indices of the keys in the left table. - * @param right_on Column indices of the keys in the right table. - * @param tag Disambiguating tag for the broadcast of the left table. - * @param keep_keys Does the result contain the key columns, or only "carrier" value - * columns - * - * @return Coroutine representing the completion of the join. - */ -streaming::Actor left_semi_join_broadcast_left( - std::shared_ptr ctx, - std::shared_ptr comm, - // We will always choose left as build table and do "broadcast" joins - std::shared_ptr left, - std::shared_ptr right, - std::shared_ptr ch_out, - std::vector left_on, - std::vector right_on, - OpID tag, - KeepKeys keep_keys -); - -/** - * @brief Perform a streaming left semi join between two tables. - * - * @note This performs a shuffle join, the left and right channels are required to provide - * hash-partitioned data in-order. - * - * @param ctx Streaming context. - * @param comm Communicator for the collective operation. - * @param left Channel of `TableChunk`s in hash-partitioned order. - * @param right Channel of `TableChunk`s in matching hash-partitioned order. - * @param ch_out Output channel of `TableChunk`s. - * @param left_on Column indices of the keys in the left table. - * @param right_on Column indices of the keys in the right table. - * @param tag Disambiguating tag for the broadcast of the left table. - * @param keep_keys Does the result contain the key columns, or only "carrier" value - * columns - * - * @return Coroutine representing the completion of the join. - */ - -streaming::Actor left_semi_join_shuffle( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr left, - std::shared_ptr right, - std::shared_ptr ch_out, - std::vector left_on, - std::vector right_on, - KeepKeys keep_keys = KeepKeys::YES -); - -/** - * @brief Shuffle the input channel by hash-partitioning on given key columns. - * - * @param ctx Streaming context. - * @param comm Communicator for the collective operation. - * @param ch_in Channel of `TableChunk`s to shuffle. - * @param ch_out Channel of shuffled `TableChunk`s. - * @param keys Indices of key columns to shuffle on. - * @param num_partitions Number of output partitions of the shuffle. - * @param tag Disambiguating tag for the shuffle. - * - * @return Coroutine representing the completion of the shuffle. - */ -[[nodiscard]] streaming::Actor shuffle( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys, - std::uint32_t num_partitions, - OpID tag -); - -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/parquet_writer.cpp b/cpp/benchmarks/streaming/ndsh/parquet_writer.cpp deleted file mode 100644 index 2ffe49909..000000000 --- a/cpp/benchmarks/streaming/ndsh/parquet_writer.cpp +++ /dev/null @@ -1,70 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "parquet_writer.hpp" - -#include -#include -#include - -#include -#include -#include - -#include -#include -#include - -namespace rapidsmpf::ndsh { - -rapidsmpf::streaming::Actor write_parquet( - std::shared_ptr ctx, - std::shared_ptr ch_in, - cudf::io::sink_info sink, - std::vector column_names -) { - streaming::ShutdownAtExit c{ch_in}; - co_await ctx->executor()->schedule(); - auto builder = cudf::io::chunked_parquet_writer_options::builder(sink); - auto msg = co_await ch_in->receive(); - RAPIDSMPF_EXPECTS(!msg.empty(), "Writing from empty channel not supported"); - auto chunk = - co_await msg.release().make_available(ctx); - auto table = chunk.table_view(); - auto metadata = cudf::io::table_input_metadata(table); - CudaEvent event; - auto write_stream = chunk.stream(); - RAPIDSMPF_EXPECTS( - column_names.size() == metadata.column_metadata.size(), - "Mismatching number of column names and chunk columns" - ); - for (std::size_t i = 0; i < column_names.size(); i++) { - metadata.column_metadata[i].set_name(column_names[i]); - } - builder = builder.metadata(metadata); - auto options = builder.build(); - auto writer = cudf::io::chunked_parquet_writer(options, write_stream); - writer.write(table); - while (true) { - msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - chunk = - co_await msg.release().make_available( - ctx - ); - table = chunk.table_view(); - RAPIDSMPF_EXPECTS( - static_cast(table.num_columns()) == column_names.size(), - "Mismatching number of column names and chunk columns" - ); - cuda_stream_join(write_stream, chunk.stream(), &event); - writer.write(table); - cuda_stream_join(chunk.stream(), write_stream, &event); - } - writer.close(); -} -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/parquet_writer.hpp b/cpp/benchmarks/streaming/ndsh/parquet_writer.hpp deleted file mode 100644 index 324404069..000000000 --- a/cpp/benchmarks/streaming/ndsh/parquet_writer.hpp +++ /dev/null @@ -1,34 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include -#include - -#include -#include - -#include -#include - -namespace rapidsmpf::ndsh { - -/** - * @brief Write chunks in a channel to an output sink - * - * @param ctx Streaming context - * @param ch_in Input channel of `TableChunk`s - * @param sink Sink to write into - * @param column_names Names of the columns to add to the parquet metadata - * - * @return Coroutine representing the write - */ -[[nodiscard]] rapidsmpf::streaming::Actor write_parquet( - std::shared_ptr ctx, - std::shared_ptr ch_in, - cudf::io::sink_info sink, - std::vector column_names -); -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/q01.cpp b/cpp/benchmarks/streaming/ndsh/q01.cpp deleted file mode 100644 index a84fdb5c9..000000000 --- a/cpp/benchmarks/streaming/ndsh/q01.cpp +++ /dev/null @@ -1,462 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "concatenate.hpp" -#include "groupby.hpp" -#include "join.hpp" -#include "parquet_writer.hpp" -#include "sort.hpp" -#include "utils.hpp" - -namespace { - -rapidsmpf::streaming::Actor read_lineitem( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory, - bool use_date32 -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "lineitem") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({ - "l_returnflag", // 0 - "l_linestatus", // 1 - "l_quantity", // 2 - "l_extendedprice", // 3 - "l_discount", // 4 - "l_tax" // 5 - }) - .build(); - auto stream = ctx->br()->stream_pool().get_stream(); - // l_shipdate <= DATE '1998-09-02' - constexpr auto date = cuda::std::chrono::year_month_day( - cuda::std::chrono::year(1998), - cuda::std::chrono::month(9), - cuda::std::chrono::day(2) - ); - auto filter_expr = - use_date32 ? rapidsmpf::ndsh::make_date_filter( - stream, date, "l_shipdate", cudf::ast::ast_operator::LESS_EQUAL - ) - : rapidsmpf::ndsh::make_date_filter( - stream, date, "l_shipdate", cudf::ast::ast_operator::LESS_EQUAL - ); - return cudf_streaming::streaming::actor::read_parquet( - ctx, - comm, - ch_out, - num_producers, - options, - num_rows_per_chunk, - std::move(filter_expr) - ); -} - -std::vector chunkwise_groupby_requests() { - auto requests = std::vector(); - std::vector()>> aggs; - // sum(l_quantity), sum(l_extendedprice), sum(disc_price), sum(charge), - // sum(l_discount) - for (cudf::size_type idx = 2; idx < 7; idx++) { - aggs.emplace_back(cudf::make_sum_aggregation); - requests.emplace_back(idx, std::move(aggs)); - } - // count(*) - aggs.emplace_back([]() { - return cudf::make_count_aggregation( - cudf::null_policy::INCLUDE - ); - }); - requests.emplace_back(0, std::move(aggs)); - return requests; -} - -std::vector final_groupby_requests() { - auto requests = std::vector(); - std::vector()>> aggs; - // sum(l_quantity), sum(l_extendedprice), sum(disc_price), sum(charge), - // sum(l_discount), sum(count(*)) - for (cudf::size_type idx = 2; idx < 8; idx++) { - aggs.emplace_back(cudf::make_sum_aggregation); - requests.emplace_back(idx, std::move(aggs)); - } - return requests; -} - -rapidsmpf::streaming::Actor postprocess_group_by( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - co_await ctx->executor()->schedule(); - auto msg = co_await ch_in->receive(); - RAPIDSMPF_EXPECTS( - (co_await ch_in->receive()).empty(), "Expecting concatenated input at this point" - ); - auto chunk = - co_await msg.release().make_available(ctx); - auto stream = chunk.stream(); - auto columns = - cudf::table{chunk.table_view(), stream, ctx->br()->device_mr()}.release(); - std::ignore = std::move(chunk); - auto count = std::move(columns.back()); - columns.pop_back(); - auto discount = std::move(columns.back()); - columns.pop_back(); - for (std::size_t i = 2; i < 4; i++) { - columns.push_back( - cudf::binary_operation( - columns[i]->view(), - count->view(), - cudf::binary_operator::TRUE_DIV, - cudf::data_type(cudf::type_id::FLOAT64), - stream, - ctx->br()->device_mr() - ) - ); - } - columns.push_back( - cudf::binary_operation( - discount->view(), - count->view(), - cudf::binary_operator::TRUE_DIV, - cudf::data_type(cudf::type_id::FLOAT64), - stream, - ctx->br()->device_mr() - ) - ); - columns.push_back(std::move(count)); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::make_unique(std::move(columns)), stream - ) - ) - ); - co_await ch_out->drain(ctx->executor()); -} - -// In: l_returnflag, l_linestatus, l_quantity, l_extendedprice, -// l_discount, l_tax -// Out: l_returnflag, l_linestatus, l_quantity, l_extendedprice, -// disc_price = (l_extendedprice * (1 - l_discount)), -// charge = (l_extendedprice * (1 - l_discount) * (1 + l_tax)), -// l_discount -rapidsmpf::streaming::Actor select_columns_for_groupby( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - - co_await ctx->executor()->schedule(); - while (!ch_out->is_shutdown()) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - auto chunk_stream = chunk.stream(); - auto sequence_number = msg.sequence_number(); - auto table = chunk.table_view(); - // l_returnflag, l_linestatus, l_quantity, l_extendedprice - auto result = - cudf::table(table.select({0, 1, 2, 3}), chunk_stream, ctx->br()->device_mr()) - .release(); - result.reserve(7); - auto extendedprice = table.column(3); - auto discount = table.column(4); - auto tax = table.column(5); - std::string udf_disc_price = - R"***( -static __device__ void calculate_disc_price(double *disc_price, double extprice, double discount) { - *disc_price = extprice * (1 - discount); -} - )***"; - std::string udf_charge = - R"***( -static __device__ void calculate_charge(double *charge, double discprice, double tax) { - *charge = discprice * (1 + tax); -} - )***"; - - // disc_price - result.push_back( - cudf::transform_extended( - std::vector{extendedprice, discount}, - udf_disc_price, - cudf::data_type(cudf::type_id::FLOAT64), - cudf::udf_source_type::CUDA, - std::nullopt, - cudf::null_aware::NO, - std::nullopt, - cudf::output_nullability::PRESERVE, - chunk_stream, - ctx->br()->device_mr() - ) - ); - // charge - result.push_back( - cudf::transform_extended( - std::vector{result.back()->view(), tax}, - udf_charge, - cudf::data_type(cudf::type_id::FLOAT64), - cudf::udf_source_type::CUDA, - std::nullopt, - cudf::null_aware::NO, - std::nullopt, - cudf::output_nullability::PRESERVE, - chunk_stream, - ctx->br()->device_mr() - ) - ); - // l_discount - result.push_back( - std::make_unique(discount, chunk_stream, ctx->br()->device_mr()) - ); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - sequence_number, - std::make_unique( - std::make_unique(std::move(result)), chunk_stream - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} -} // namespace - -/** - * @brief Run a derived version of TPC-H query 1. - * - * The SQL form of the query is: - * @code{.sql} - * select - * l_returnflag, - * l_linestatus, - * sum(l_quantity) as sum_qty, - * sum(l_extendedprice) as sum_base_price, - * sum(l_extendedprice * (1 - l_discount)) as sum_disc_price, - * sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge, - * avg(l_quantity) as avg_qty, - * avg(l_extendedprice) as avg_price, - * avg(l_discount) as avg_disc, - * count(*) as count_order - * from - * lineitem - * where - * l_shipdate <= DATE '1998-09-02' - * group by - * l_returnflag, - * l_linestatus - * order by - * l_returnflag, - * l_linestatus - * @endcode{} - */ -int main(int argc, char** argv) { - rapidsmpf::ndsh::FinalizeMPI finalize{}; - cudaFree(nullptr); - // work around https://github.com/rapidsai/cudf/issues/20849 - cudf::initialize(); - auto mr = rmm::mr::cuda_async_memory_resource{}; - auto arguments = rapidsmpf::ndsh::parse_arguments(argc, argv); - auto [ctx, comm] = rapidsmpf::ndsh::create_context(arguments, std::move(mr)); - std::string output_path = arguments.output_file; - - // Detect date column type from parquet metadata before timed section - auto const column_types = - rapidsmpf::ndsh::detail::get_column_types(arguments.input_directory, "lineitem"); - bool const use_date32 = - column_types.at("l_shipdate").id() == cudf::type_id::TIMESTAMP_DAYS; - - std::vector timings; - for (int i = 0; i < arguments.num_iterations; i++) { - int op_id = 0; - std::vector actors; - auto start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Constructing Q1 pipeline"); - - // Input data channels - auto lineitem = ctx->create_channel(); - // Out: l_returnflag, l_linestatus, l_quantity, l_extendedprice, - // l_discount, l_tax - actors.push_back(read_lineitem( - ctx, - comm, - lineitem, - /* num_tickets */ 4, - arguments.num_rows_per_chunk, - arguments.input_directory, - use_date32 - )); - - auto chunkwise_groupby_input = ctx->create_channel(); - // Out: l_returnflag, l_linestatus, l_quantity, l_extendedprice, - // disc_price = (l_extendedprice * (1 - l_discount)), - // charge = (l_extendedprice * (1 - l_discount) * (1 + l_tax)) - // l_discount - actors.push_back( - select_columns_for_groupby(ctx, lineitem, chunkwise_groupby_input) - ); - auto chunkwise_groupby_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - chunkwise_groupby_input, - chunkwise_groupby_output, - {0, 1}, - chunkwise_groupby_requests(), - cudf::null_policy::INCLUDE - ) - ); - auto final_groupby_input = ctx->create_channel(); - if (comm->nranks() > 1) { - actors.push_back( - rapidsmpf::ndsh::broadcast( - ctx, - comm, - chunkwise_groupby_output, - final_groupby_input, - static_cast(10 * i + op_id++), - rapidsmpf::streaming::AllGather::Ordered::NO - ) - ); - } else { - actors.push_back( - rapidsmpf::ndsh::concatenate( - ctx, chunkwise_groupby_output, final_groupby_input - ) - ); - } - if (comm->rank() == 0) { - auto final_groupby_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - final_groupby_input, - final_groupby_output, - {0, 1}, - final_groupby_requests(), - cudf::null_policy::INCLUDE - ) - ); - auto sorted_input = ctx->create_channel(); - actors.push_back( - postprocess_group_by(ctx, final_groupby_output, sorted_input) - ); - auto sorted_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_sort_by( - ctx, - sorted_input, - sorted_output, - {0, 1}, - {0, 1, 2, 3, 4, 5, 6, 7, 8, 9}, - {cudf::order::ASCENDING, cudf::order::ASCENDING}, - {cudf::null_order::BEFORE, cudf::null_order::BEFORE} - ) - ); - actors.push_back( - rapidsmpf::ndsh::write_parquet( - ctx, - sorted_output, - cudf::io::sink_info(output_path), - {"l_returnflag", - "l_linestatus", - "sum_qty", - "sum_base_price", - "sum_disc_price", - "sum_charge", - "avg_qty", - "avg_price", - "avg_disc", - "count_order"} - - ) - ); - } else { - actors.push_back(rapidsmpf::ndsh::sink_channel(ctx, final_groupby_input)); - } - } - auto end = std::chrono::steady_clock::now(); - std::chrono::duration pipeline = end - start; - start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Q1 Iteration"); - rapidsmpf::streaming::run_actor_network(std::move(actors)); - } - end = std::chrono::steady_clock::now(); - std::chrono::duration compute = end - start; - timings.push_back(pipeline.count()); - timings.push_back(compute.count()); - auto statistics = ctx->statistics(); - comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} - )); - statistics->clear(); - } - - if (comm->rank() == 0) { - for (int i = 0; i < arguments.num_iterations; i++) { - comm->logger()->print( - "Iteration ", - i, - " pipeline construction time [s]: ", - timings[rapidsmpf::safe_cast(2 * i)] - ); - comm->logger()->print( - "Iteration ", - i, - " compute time [s]: ", - timings[rapidsmpf::safe_cast(2 * i + 1)] - ); - } - } - return 0; -} diff --git a/cpp/benchmarks/streaming/ndsh/q03.cpp b/cpp/benchmarks/streaming/ndsh/q03.cpp deleted file mode 100644 index 2a2e46f3f..000000000 --- a/cpp/benchmarks/streaming/ndsh/q03.cpp +++ /dev/null @@ -1,731 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "concatenate.hpp" -#include "groupby.hpp" -#include "join.hpp" -#include "parquet_writer.hpp" -#include "utils.hpp" - -namespace { - -rapidsmpf::streaming::Actor read_customer( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "customer") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({"c_custkey"}) // 0 - .build(); - auto filter_expr = [&]() -> std::unique_ptr { - auto stream = ctx->br()->stream_pool().get_stream(); - auto owner = new std::vector; - owner->push_back(std::make_shared("BUILDING", true, stream)); - owner->push_back( - std::make_shared( - *std::any_cast>(owner->at(0)) - ) - ); - owner->push_back( - std::make_shared("c_mktsegment") - ); - owner->push_back( - std::make_shared( - cudf::ast::ast_operator::EQUAL, - *std::any_cast>( - owner->at(2) - ), - *std::any_cast>(owner->at(1)) - ) - ); - return std::make_unique( - stream, - *std::any_cast>(owner->back()), - rapidsmpf::OwningWrapper(static_cast(owner), [](void* p) { - delete static_cast*>(p); - }) - ); - }(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, - comm, - ch_out, - num_producers, - options, - num_rows_per_chunk, - std::move(filter_expr) - ); -} - -rapidsmpf::streaming::Actor read_lineitem( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory, - bool use_date32 -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "lineitem") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({ - "l_orderkey", // 0 - "l_extendedprice", // 1 - "l_discount", // 2 - }) - .build(); - auto stream = ctx->br()->stream_pool().get_stream(); - // l_shipdate > DATE '1995-03-15' - constexpr auto date = cuda::std::chrono::year_month_day( - cuda::std::chrono::year(1995), - cuda::std::chrono::month(3), - cuda::std::chrono::day(15) - ); - auto filter_expr = - use_date32 ? rapidsmpf::ndsh::make_date_filter( - stream, date, "l_shipdate", cudf::ast::ast_operator::GREATER - ) - : rapidsmpf::ndsh::make_date_filter( - stream, date, "l_shipdate", cudf::ast::ast_operator::GREATER - ); - return cudf_streaming::streaming::actor::read_parquet( - ctx, - comm, - ch_out, - num_producers, - options, - num_rows_per_chunk, - std::move(filter_expr) - ); -} - -rapidsmpf::streaming::Actor read_orders( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory, - bool use_date32 -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "orders") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({ - "o_orderkey", // 0 - "o_orderdate", // 1 - "o_shippriority", // 2 - "o_custkey" // 3 - }) - .build(); - auto stream = ctx->br()->stream_pool().get_stream(); - // o_orderdate < DATE '1995-03-15' - constexpr auto date = cuda::std::chrono::year_month_day( - cuda::std::chrono::year(1995), - cuda::std::chrono::month(3), - cuda::std::chrono::day(15) - ); - auto filter_expr = - use_date32 ? rapidsmpf::ndsh::make_date_filter( - stream, date, "o_orderdate", cudf::ast::ast_operator::LESS - ) - : rapidsmpf::ndsh::make_date_filter( - stream, date, "o_orderdate", cudf::ast::ast_operator::LESS - ); - return cudf_streaming::streaming::actor::read_parquet( - ctx, - comm, - ch_out, - num_producers, - options, - num_rows_per_chunk, - std::move(filter_expr) - ); -} - -std::vector chunkwise_groupby_requests() { - auto requests = std::vector(); - std::vector()>> aggs; - // sum(revenue) - aggs.emplace_back(cudf::make_sum_aggregation); - requests.emplace_back(3, std::move(aggs)); - return requests; -} - -// In: o_orderkey, o_orderdate, o_shippriority, l_extendedprice, l_discount -// Out: o_orderkey, o_orderdate, o_shippriority, revenue = (l_extendedprice - (1 - -// l_discount)) -rapidsmpf::streaming::Actor select_columns_for_groupby( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - - co_await ctx->executor()->schedule(); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - auto chunk_stream = chunk.stream(); - auto sequence_number = msg.sequence_number(); - auto table = chunk.table_view(); - std::vector> result; - result.reserve(4); - - // o_orderkey - result.push_back( - std::make_unique( - table.column(0), chunk_stream, ctx->br()->device_mr() - ) - ); - // o_orderdate - result.push_back( - std::make_unique( - table.column(1), chunk_stream, ctx->br()->device_mr() - ) - ); - // o_shippriority - result.push_back( - std::make_unique( - table.column(2), chunk_stream, ctx->br()->device_mr() - ) - ); - auto extendedprice = table.column(3); - auto discount = table.column(4); - std::string udf = - R"***( -static __device__ void calculate_revenue(double *revenue, double extprice, double discount) { - *revenue = extprice * (1 - discount); -} - )***"; - - // revenue - result.push_back( - cudf::transform_extended( - std::vector{extendedprice, discount}, - udf, - cudf::data_type(cudf::type_id::FLOAT64), - cudf::udf_source_type::CUDA, - std::nullopt, - cudf::null_aware::NO, - std::nullopt, - cudf::output_nullability::PRESERVE, - chunk_stream, - ctx->br()->device_mr() - ) - ); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - sequence_number, - std::make_unique( - std::make_unique(std::move(result)), chunk_stream - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -rapidsmpf::streaming::Actor top_k_by( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys, - std::vector values, - std::vector order, - cudf::size_type k -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - - co_await ctx->executor()->schedule(); - std::vector> partials; - std::vector chunk_streams; - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - auto const indices = cudf::sorted_order( - chunk.table_view().select(keys), - order, - {}, - chunk.stream(), - ctx->br()->device_mr() - ); - partials.push_back( - cudf::gather( - chunk.table_view().select(values), - cudf::split(indices->view(), {k}, chunk.stream()).front(), - cudf::out_of_bounds_policy::DONT_CHECK, - chunk.stream(), - ctx->br()->device_mr() - ) - ); - chunk_streams.push_back(chunk.stream()); - } - - // TODO: multi-node - RAPIDSMPF_EXPECTS(chunk_streams.size() > 0, "No chunks to sort"); - auto out_stream = chunk_streams.front(); - rapidsmpf::CudaEvent event; - rapidsmpf::cuda_stream_join( - std::ranges::single_view{out_stream}, chunk_streams, &event - ); - std::vector views; - std::ranges::transform(partials, std::back_inserter(views), [](auto& t) { - return t->view(); - }); - auto merged = cudf::merge(views, keys, order, {}, out_stream, ctx->br()->device_mr()); - auto result = std::make_unique( - cudf::slice(merged->view(), {0, 10}, out_stream), - out_stream, - ctx->br()->device_mr() - ); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - 0, - std::make_unique( - std::move(result), out_stream - ) - ) - ); - co_await ch_out->drain(ctx->executor()); -} - -rapidsmpf::streaming::Actor fanout_bounded( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - std::shared_ptr ch1_out, - std::vector ch1_cols, - std::shared_ptr ch2_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch1_out, ch2_out}; - - co_await ctx->executor()->schedule(); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - // Here, we know that copying ch1_cols (a single col) is better than copying - // ch2_cols (the whole table) - std::vector> tasks; - if (!ch1_out->is_shutdown()) { - auto msg1 = cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::make_unique( - chunk.table_view().select(ch1_cols), - chunk.stream(), - ctx->br()->device_mr() - ), - chunk.stream() - ) - ); - tasks.push_back(ch1_out->send(std::move(msg1))); - } - if (!ch2_out->is_shutdown()) { - // TODO: We know here that ch2 wants the whole table. - tasks.push_back(ch2_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::move(chunk) - ) - ) - )); - } - if (!std::ranges::any_of( - rapidsmpf::streaming::coro_results( - co_await coro::when_all(std::move(tasks)) - ), - std::identity{} - )) - { - comm->logger()->print("Breaking after ", msg.sequence_number()); - break; - }; - } - - rapidsmpf::streaming::coro_results( - co_await coro::when_all( - ch1_out->drain(ctx->executor()), ch2_out->drain(ctx->executor()) - ) - ); -} -} // namespace - -/** - * @brief Run a derived version of TPC-H query 3. - * - * The SQL form of the query is: - * @code{.sql} - * select - * l_orderkey, - * sum(l_extendedprice * (1 - l_discount)) as revenue, - * o_orderdate, - * o_shippriority - * from - * customer, - * orders, - * lineitem - * where - * c_mktsegment = 'BUILDING' - * and c_custkey = o_custkey - * and l_orderkey = o_orderkey - * and o_orderdate < '1995-03-15' - * and l_shipdate > '1995-03-15' - * group by - * l_orderkey, - * o_orderdate, - * o_shippriority - * order by - * revenue desc, - * o_orderdate - * limit 10 - * @endcode{} - */ -int main(int argc, char** argv) { - rapidsmpf::ndsh::FinalizeMPI finalize{}; - cudaFree(nullptr); - // work around https://github.com/rapidsai/cudf/issues/20849 - cudf::initialize(); - auto mr = rmm::mr::cuda_async_memory_resource{}; - auto arguments = rapidsmpf::ndsh::parse_arguments(argc, argv); - auto [ctx, comm] = rapidsmpf::ndsh::create_context(arguments, std::move(mr)); - std::string output_path = arguments.output_file; - - // Detect date column types from parquet metadata before timed section - auto const lineitem_types = - rapidsmpf::ndsh::detail::get_column_types(arguments.input_directory, "lineitem"); - bool const lineitem_use_date32 = - lineitem_types.at("l_shipdate").id() == cudf::type_id::TIMESTAMP_DAYS; - auto const orders_types = - rapidsmpf::ndsh::detail::get_column_types(arguments.input_directory, "orders"); - bool const orders_use_date32 = - orders_types.at("o_orderdate").id() == cudf::type_id::TIMESTAMP_DAYS; - - std::vector timings; - int l2size; - int device; - RAPIDSMPF_CUDA_TRY(cudaGetDevice(&device)); - RAPIDSMPF_CUDA_TRY(cudaDeviceGetAttribute(&l2size, cudaDevAttrL2CacheSize, device)); - auto const num_filter_blocks = - cudf_streaming::integrations::BloomFilter::fitting_num_blocks( - static_cast(l2size) - ); - - for (int i = 0; i < arguments.num_iterations; i++) { - int op_id{0}; - std::vector actors; - auto start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Constructing Q3 pipeline"); - auto customer = ctx->create_channel(); - auto lineitem = ctx->create_channel(); - auto orders = ctx->create_channel(); - - auto customer_x_orders = ctx->create_channel(); - auto customer_x_orders_x_lineitem = ctx->create_channel(); - - // Out: "c_custkey" - actors.push_back(read_customer( - ctx, - comm, - customer, - /* num_tickets */ 2, - arguments.num_rows_per_chunk, - arguments.input_directory - )); - // Out: o_orderkey, o_orderdate, o_shippriority, o_custkey - actors.push_back(read_orders( - ctx, - comm, - orders, - 6, - arguments.num_rows_per_chunk, - arguments.input_directory, - orders_use_date32 - )); - // join c_custkey = o_custkey - // Out: o_orderkey, o_orderdate, o_shippriority - actors.push_back( - rapidsmpf::ndsh::inner_join_broadcast( - ctx, - comm, - customer, - orders, - customer_x_orders, - {0}, - {3}, - static_cast(10 * i + op_id++), - rapidsmpf::ndsh::KeepKeys::NO - ) - ); - auto bloom_filter_input = ctx->create_channel(); - auto bloom_filter_output = ctx->create_channel(); - auto customer_x_orders_input = ctx->create_channel(); - actors.push_back(fanout_bounded( - ctx, - comm, - customer_x_orders, - bloom_filter_input, - {0}, - customer_x_orders_input - )); - auto bloom_filter = cudf_streaming::streaming::BloomFilter( - ctx, comm, cudf::DEFAULT_HASH_SEED, num_filter_blocks - ); - actors.push_back(bloom_filter.build( - bloom_filter_input, - bloom_filter_output, - static_cast(10 * i + op_id++) - )); - // Out: l_orderkey, l_extendedprice, l_discount - actors.push_back(read_lineitem( - ctx, - comm, - lineitem, - /* num_tickets */ 4, - arguments.num_rows_per_chunk, - arguments.input_directory, - lineitem_use_date32 - )); - auto lineitem_output = ctx->create_channel(); - actors.push_back( - bloom_filter.apply(bloom_filter_output, lineitem, lineitem_output, {0}) - ); - // join o_orderkey = l_orderkey - // Out: o_orderkey, o_orderdate, o_shippriority, l_extendedprice, - // l_discount - if (arguments.use_shuffle_join) { - auto lineitem_shuffled = ctx->create_channel(); - auto customer_x_orders_shuffled = ctx->create_channel(); - std::uint32_t num_partitions = 16; - actors.push_back( - rapidsmpf::ndsh::shuffle( - ctx, - comm, - lineitem_output, - lineitem_shuffled, - {0}, - num_partitions, - static_cast(10 * i + op_id++) - ) - ); - actors.push_back( - rapidsmpf::ndsh::shuffle( - ctx, - comm, - customer_x_orders_input, - customer_x_orders_shuffled, - {0}, - num_partitions, - static_cast(10 * i + op_id++) - ) - ); - actors.push_back( - rapidsmpf::ndsh::inner_join_shuffle( - ctx, - comm, - customer_x_orders_shuffled, - lineitem_shuffled, - customer_x_orders_x_lineitem, - {0}, - {0}, - rapidsmpf::ndsh::KeepKeys::YES - ) - ); - } else { - actors.push_back( - rapidsmpf::ndsh::inner_join_broadcast( - ctx, - comm, - customer_x_orders_input, - lineitem_output, - customer_x_orders_x_lineitem, - {0}, - {0}, - static_cast(10 * i + op_id++), - rapidsmpf::ndsh::KeepKeys::YES - ) - ); - } - auto groupby_input = ctx->create_channel(); - // Out: o_orderkey, o_orderdate, o_shippriority, revenue - actors.push_back(select_columns_for_groupby( - ctx, customer_x_orders_x_lineitem, groupby_input - )); - auto chunkwise_groupby_output = ctx->create_channel(); - // Out: o_orderkey, o_orderdate, o_shippriority, revenue - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - groupby_input, - chunkwise_groupby_output, - {0, 1, 2}, - chunkwise_groupby_requests(), - cudf::null_policy::INCLUDE - ) - ); - auto final_groupby_input = ctx->create_channel(); - if (comm->nranks() > 1) { - actors.push_back( - rapidsmpf::ndsh::broadcast( - ctx, - comm, - chunkwise_groupby_output, - final_groupby_input, - static_cast(10 * i + op_id++), - rapidsmpf::streaming::AllGather::Ordered::NO - ) - ); - } else { - actors.push_back( - rapidsmpf::ndsh::concatenate( - ctx, chunkwise_groupby_output, final_groupby_input - ) - ); - } - if (comm->rank() == 0) { - auto final_groupby_output = ctx->create_channel(); - // Out: o_orderkey, o_orderdate, o_shippriority, revenue - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - final_groupby_input, - final_groupby_output, - {0, 1, 2}, - chunkwise_groupby_requests(), - cudf::null_policy::INCLUDE - - ) - ); - auto topk = ctx->create_channel(); - // Out: o_orderkey, revenue, o_orderdate, o_shippriority - actors.push_back(top_k_by( - ctx, - final_groupby_output, - topk, - {3, 1}, - {0, 3, 1, 2}, - {cudf::order::DESCENDING, cudf::order::ASCENDING}, - 10 - )); - actors.push_back( - rapidsmpf::ndsh::write_parquet( - ctx, - topk, - cudf::io::sink_info(output_path), - {"l_orderkey", "revenue", "o_orderdate", "o_shippriority"} - ) - ); - } else { - actors.push_back(rapidsmpf::ndsh::sink_channel(ctx, final_groupby_input)); - } - } - auto end = std::chrono::steady_clock::now(); - std::chrono::duration pipeline = end - start; - start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Q3 Iteration"); - rapidsmpf::streaming::run_actor_network(std::move(actors)); - } - end = std::chrono::steady_clock::now(); - std::chrono::duration compute = end - start; - timings.push_back(pipeline.count()); - timings.push_back(compute.count()); - auto statistics = ctx->statistics(); - comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} - )); - statistics->clear(); - } - if (comm->rank() == 0) { - for (int i = 0; i < arguments.num_iterations; i++) { - comm->logger()->print( - "Iteration ", - i, - " pipeline construction time [s]: ", - timings[rapidsmpf::safe_cast(2 * i)] - ); - comm->logger()->print( - "Iteration ", - i, - " compute time [s]: ", - timings[rapidsmpf::safe_cast(2 * i + 1)] - ); - } - } - return 0; -} diff --git a/cpp/benchmarks/streaming/ndsh/q04.cpp b/cpp/benchmarks/streaming/ndsh/q04.cpp deleted file mode 100644 index 2c9825ab4..000000000 --- a/cpp/benchmarks/streaming/ndsh/q04.cpp +++ /dev/null @@ -1,551 +0,0 @@ -/** - - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include -#include -#include - -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "concatenate.hpp" -#include "groupby.hpp" -#include "join.hpp" -#include "parquet_writer.hpp" -#include "sort.hpp" -#include "utils.hpp" - -namespace { - -std::vector chunkwise_groupby_requests() { - auto requests = std::vector(); - std::vector()>> aggs; - // count(*) - aggs.emplace_back([]() { - return cudf::make_count_aggregation( - cudf::null_policy::INCLUDE - ); - }); - requests.emplace_back(0, std::move(aggs)); - return requests; -} - -std::vector final_groupby_requests() { - auto requests = std::vector(); - std::vector()>> aggs; - // sum of partial counts - aggs.emplace_back([]() { - return cudf::make_sum_aggregation(); - }); - requests.emplace_back(1, std::move(aggs)); // column 1 is order_count - return requests; -} - -rapidsmpf::streaming::Actor read_lineitem( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "lineitem") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({ - "l_commitdate", // used in filter - "l_receiptdate", // used in filter - "l_orderkey", // used in join - }) - .build(); - - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor read_orders( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory, - bool use_date32 -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "orders") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({ - "o_orderkey", // used in join - "o_orderpriority", // used in group by - }) - .build(); - - auto stream = ctx->br()->stream_pool().get_stream(); - // 1993-07-01 <= o_orderdate < 1993-10-01 - constexpr auto start_date = cuda::std::chrono::year_month_day( - cuda::std::chrono::year(1993), - cuda::std::chrono::month(7), - cuda::std::chrono::day(1) - ); - constexpr auto end_date = cuda::std::chrono::year_month_day( - cuda::std::chrono::year(1993), - cuda::std::chrono::month(10), - cuda::std::chrono::day(1) - ); - auto filter = use_date32 - ? rapidsmpf::ndsh::make_date_range_filter( - stream, start_date, end_date, "o_orderdate" - ) - : rapidsmpf::ndsh::make_date_range_filter( - stream, start_date, end_date, "o_orderdate" - ); - - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk, std::move(filter) - ); -} - -rapidsmpf::streaming::Actor filter_lineitem( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - co_await ctx->executor()->schedule(); - auto mr = ctx->br()->device_mr(); - - while (!ch_out->is_shutdown()) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - auto chunk_stream = chunk.stream(); - auto table = chunk.table_view(); - - auto l_commitdate = table.column(0); - auto l_receiptdate = table.column(1); - auto mask = cudf::binary_operation( - l_commitdate, - l_receiptdate, - cudf::binary_operator::LESS, - cudf::data_type(cudf::type_id::BOOL8), - chunk_stream, - mr - ); - auto filtered_table = - cudf::apply_boolean_mask(table.select({2}), mask->view(), chunk_stream, mr); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::move(filtered_table), chunk_stream - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -[[maybe_unused]] -rapidsmpf::streaming::Actor fanout_bounded( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - std::shared_ptr ch1_out, - std::vector ch1_cols, - std::shared_ptr ch2_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch1_out, ch2_out}; - co_await ctx->executor()->schedule(); - - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); // Here, we know that copying ch1_cols (a single col) is better than - // copying - // ch2_cols (the whole table) - std::vector> tasks; - if (!ch1_out->is_shutdown()) { - auto msg1 = cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::make_unique( - chunk.table_view().select(ch1_cols), - chunk.stream(), - ctx->br()->device_mr() - ), - chunk.stream() - ) - ); - tasks.push_back(ch1_out->send(std::move(msg1))); - } - if (!ch2_out->is_shutdown()) { - // TODO: We know here that ch2 wants the whole table. - tasks.push_back(ch2_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::move(chunk) - ) - ) - )); - } - if (!std::ranges::any_of( - rapidsmpf::streaming::coro_results( - co_await coro::when_all(std::move(tasks)) - ), - std::identity{} - )) - { - comm->logger()->print("Breaking after ", msg.sequence_number()); - break; - }; - } - - rapidsmpf::streaming::coro_results( - co_await coro::when_all( - ch1_out->drain(ctx->executor()), ch2_out->drain(ctx->executor()) - ) - ); -} - -} // namespace - -/** - * @brief Run a derived version of TPCH-query 4. - * - * The SQL form of the query is: - * @code{.sql} - * - * SELECT - * o_orderpriority, - * count(*) as order_count - * FROM - * orders - * where - * o_orderdate >= TIMESTAMP '1993-07-01' - * and o_orderdate < TIMESTAMP '1993-07-01' + INTERVAL '3' MONTH - * and EXISTS ( - * SELECT - * * - * FROM - * lineitem - * WHERE - * l_orderkey = o_orderkey - * and l_commitdate < l_receiptdate - * ) - * GROUP BY - * o_orderpriority - * ORDER BY - * o_orderpriority - * @endcode{} - * - * The "exists" clause is translated into a left-semi join in libcudf. - */ -int main(int argc, char** argv) { - cudaFree(nullptr); - - rapidsmpf::ndsh::FinalizeMPI finalize{}; - cudaFree(nullptr); - // work around https://github.com/rapidsai/cudf/issues/20849 - cudf::initialize(); - auto mr = rmm::mr::cuda_async_memory_resource{}; - auto arguments = rapidsmpf::ndsh::parse_arguments(argc, argv); - auto [ctx, comm] = rapidsmpf::ndsh::create_context(arguments, std::move(mr)); - std::string output_path = arguments.output_file; - std::vector timings; - - // Detect date column types from parquet metadata before timed section - auto const orders_types = - rapidsmpf::ndsh::detail::get_column_types(arguments.input_directory, "orders"); - bool const orders_use_date32 = - orders_types.at("o_orderdate").id() == cudf::type_id::TIMESTAMP_DAYS; - - int l2size; - int device; - RAPIDSMPF_CUDA_TRY(cudaGetDevice(&device)); - RAPIDSMPF_CUDA_TRY(cudaDeviceGetAttribute(&l2size, cudaDevAttrL2CacheSize, device)); - auto const num_filter_blocks = - cudf_streaming::integrations::BloomFilter::fitting_num_blocks( - static_cast(l2size) - ); - - for (int i = 0; i < arguments.num_iterations; i++) { - rapidsmpf::OpID op_id{0}; - std::vector actors; - auto start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Constructing Q4 pipeline"); - // Convention for channel names: express the *output*. - /* Lineitem Table */ - // [l_commitdate, l_receiptdate, l_orderkey] - auto lineitem = ctx->create_channel(); - // [l_orderkey] - auto filtered_lineitem = ctx->create_channel(); - // [l_orderkey] - auto filtered_lineitem_shuffled = ctx->create_channel(); - - /* Orders Table */ - // [o_orderkey, o_orderpriority] - auto order = ctx->create_channel(); - - // [o_orderpriority] - auto orders_x_lineitem = ctx->create_channel(); - // [o_orderpriority, order_count] - auto grouped_chunkwise = ctx->create_channel(); - - actors.push_back(read_lineitem( - ctx, - comm, - lineitem, - 4, - arguments.num_rows_per_chunk, - arguments.input_directory - )); - actors.push_back( - filter_lineitem(ctx, lineitem, filtered_lineitem) - ); // l_orderkey - actors.push_back(read_orders( - ctx, - comm, - order, - 4, - arguments.num_rows_per_chunk, - arguments.input_directory, - orders_use_date32 - )); - - // Fanout filtered orders: one for bloom filter, one for join - auto bloom_filter_input = ctx->create_channel(); - auto orders_for_join = ctx->create_channel(); - actors.push_back( - fanout_bounded(ctx, comm, order, bloom_filter_input, {0}, orders_for_join) - ); - - // Build bloom filter from filtered orders' o_orderkey - auto bloom_filter_output = ctx->create_channel(); - auto bloom_filter = cudf_streaming::streaming::BloomFilter( - ctx, comm, cudf::DEFAULT_HASH_SEED, num_filter_blocks - ); - actors.push_back(bloom_filter.build( - bloom_filter_input, - bloom_filter_output, - static_cast(10 * i + op_id++) - )); - - // Apply bloom filter to filtered lineitem before shuffling - auto bloom_filtered_lineitem = ctx->create_channel(); - actors.push_back(bloom_filter.apply( - bloom_filter_output, filtered_lineitem, bloom_filtered_lineitem, {0} - )); - - // We unconditionally shuffle the filtered lineitem table. This is - // necessary to correctly handle duplicates in the left-semi join. - // Failing to shuffle (hash partition) the right table on the join - // key could allow a record to match multiple times from the - // multiple partitions of the right table. - - // TODO: configurable num_partitions - std::uint32_t num_partitions = 16; - actors.push_back( - rapidsmpf::ndsh::shuffle( - ctx, - comm, - bloom_filtered_lineitem, - filtered_lineitem_shuffled, - {0}, - num_partitions, - static_cast(10 * i + op_id++) - ) - ); - - if (arguments.use_shuffle_join) { - auto filtered_order_shuffled = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::shuffle( - ctx, - comm, - orders_for_join, - filtered_order_shuffled, - {0}, - num_partitions, - static_cast(10 * i + op_id++) - ) - ); - - actors.push_back( - rapidsmpf::ndsh::left_semi_join_shuffle( - ctx, - comm, - filtered_order_shuffled, - filtered_lineitem_shuffled, - orders_x_lineitem, - {0}, - {0} - ) - ); - } else { - actors.push_back( - rapidsmpf::ndsh::left_semi_join_broadcast_left( - ctx, - comm, - orders_for_join, - filtered_lineitem_shuffled, - orders_x_lineitem, - {0}, - {0}, - static_cast(10 * i + op_id++), - rapidsmpf::ndsh::KeepKeys::NO - ) - ); - } - - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - orders_x_lineitem, - grouped_chunkwise, - {0}, - chunkwise_groupby_requests(), - cudf::null_policy::INCLUDE - ) - ); - auto final_groupby_input = ctx->create_channel(); - if (comm->nranks() > 1) { - actors.push_back( - rapidsmpf::ndsh::broadcast( - ctx, - comm, - grouped_chunkwise, - final_groupby_input, - static_cast(10 * i + op_id++), - rapidsmpf::streaming::AllGather::Ordered::NO - ) - ); - } else { - actors.push_back( - rapidsmpf::ndsh::concatenate( - ctx, grouped_chunkwise, final_groupby_input - ) - ); - } - if (comm->rank() == 0) { - auto final_groupby_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - final_groupby_input, - final_groupby_output, - {0}, - final_groupby_requests(), - cudf::null_policy::INCLUDE - ) - ); - auto sorted_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_sort_by( - ctx, - final_groupby_output, - sorted_output, - {0}, - {0, 1}, - {cudf::order::ASCENDING}, - {cudf::null_order::BEFORE} - ) - ); - actors.push_back( - rapidsmpf::ndsh::write_parquet( - ctx, - sorted_output, - cudf::io::sink_info(output_path), - {"o_orderpriority", "order_count"} - ) - ); - } else { - actors.push_back(rapidsmpf::ndsh::sink_channel(ctx, final_groupby_input)); - } - } - auto end = std::chrono::steady_clock::now(); - std::chrono::duration pipeline = end - start; - start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Q4 Iteration"); - rapidsmpf::streaming::run_actor_network(std::move(actors)); - } - end = std::chrono::steady_clock::now(); - std::chrono::duration compute = end - start; - timings.push_back(pipeline.count()); - timings.push_back(compute.count()); - auto statistics = ctx->statistics(); - comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} - )); - statistics->clear(); - } - - if (comm->rank() == 0) { - for (int i = 0; i < arguments.num_iterations; i++) { - comm->logger()->print( - "Iteration ", - i, - " pipeline construction time [s]: ", - timings[rapidsmpf::safe_cast(2 * i)] - ); - comm->logger()->print( - "Iteration ", - i, - " compute time [s]: ", - timings[rapidsmpf::safe_cast(2 * i + 1)] - ); - } - } - return 0; -} diff --git a/cpp/benchmarks/streaming/ndsh/q09.cpp b/cpp/benchmarks/streaming/ndsh/q09.cpp deleted file mode 100644 index dbe038209..000000000 --- a/cpp/benchmarks/streaming/ndsh/q09.cpp +++ /dev/null @@ -1,678 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -#include "concatenate.hpp" -#include "groupby.hpp" -#include "join.hpp" -#include "parquet_writer.hpp" -#include "sort.hpp" -#include "utils.hpp" - -using rapidsmpf::safe_cast; - -namespace { - -rapidsmpf::streaming::Actor read_lineitem( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "lineitem") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names( - {"l_discount", - "l_extendedprice", - "l_orderkey", - "l_partkey", - "l_quantity", - "l_suppkey"} - ) - .build(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor read_nation( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "nation") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({"n_name", "n_nationkey"}) - .build(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor read_orders( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "orders") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({"o_orderdate", "o_orderkey"}) - .build(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor read_part( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "part") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({"p_partkey", "p_name"}) - .build(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor read_partsupp( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "partsupp") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({"ps_partkey", "ps_suppkey", "ps_supplycost"}) - .build(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor read_supplier( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "supplier") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({"s_nationkey", "s_suppkey"}) - .build(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor filter_part( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - auto mr = ctx->br()->device_mr(); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - co_await ctx->executor()->schedule(); - auto chunk = - co_await msg.release().make_available( - ctx - ); - auto chunk_stream = chunk.stream(); - auto table = chunk.table_view(); - auto p_name = table.column(1); - auto target = cudf::make_string_scalar("green", chunk_stream, mr); - auto mask = cudf::strings::contains( - p_name, *static_cast(target.get()), chunk_stream, mr - ); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - cudf::apply_boolean_mask( - table.select({0}), mask->view(), chunk_stream, mr - ), - chunk_stream - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -rapidsmpf::streaming::Actor select_columns( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - // n_name, ps_supplycost, l_discount, l_extendedprice, l_quantity, o_orderdate - - // Select n_name, year_part_of(o_orderdate), amount = (extendedprice * (1 - // - discount)) - (ps_supplycost * l_quantity) group by n_name year agg - // sum(amount).round(2) sort by n_name, o_year descending = true, false - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - co_await ctx->executor()->schedule(); - auto chunk = - co_await msg.release().make_available( - ctx - ); - auto chunk_stream = chunk.stream(); - auto sequence_number = msg.sequence_number(); - auto table = chunk.table_view(); - std::vector> result; - result.reserve(3); - // n_name - result.push_back( - std::make_unique( - table.column(0), chunk_stream, ctx->br()->device_mr() - ) - ); - result.push_back( - cudf::datetime::extract_datetime_component( - table.column(5), - cudf::datetime::datetime_component::YEAR, - chunk_stream, - ctx->br()->device_mr() - ) - ); - auto discount = table.column(2); - auto extendedprice = table.column(3); - auto supplycost = table.column(1); - auto quantity = table.column(4); - std::string udf = - R"***( -static __device__ void calculate_amount(double *amount, double discount, double extprice, double supplycost, double quantity) { - *amount = extprice * (1 - discount) - supplycost * quantity; -} - )***"; - result.push_back( - cudf::transform_extended( - std::vector{ - discount, extendedprice, supplycost, quantity - }, - udf, - cudf::data_type(cudf::type_id::FLOAT64), - cudf::udf_source_type::CUDA, - std::nullopt, - cudf::null_aware::NO, - std::nullopt, - cudf::output_nullability::PRESERVE, - chunk_stream, - ctx->br()->device_mr() - ) - ); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - sequence_number, - std::make_unique( - std::make_unique(std::move(result)), chunk_stream - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -std::vector chunkwise_groupby_requests() { - auto requests = std::vector(); - std::vector()>> aggs; - aggs.emplace_back(cudf::make_sum_aggregation); - requests.emplace_back(2, std::move(aggs)); - return requests; -} - -rapidsmpf::streaming::Actor round_sum_profit( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - co_await ctx->executor()->schedule(); - auto msg = co_await ch_in->receive(); - RAPIDSMPF_EXPECTS(!msg.empty(), "Expecting to see a single chunk"); - auto next = co_await ch_in->receive(); - RAPIDSMPF_EXPECTS(next.empty(), "Not expecting to see a second chunk"); - auto chunk = - co_await msg.release().make_available(ctx); - auto table = chunk.table_view(); -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wdeprecated-declarations" - // cudf::round_decimal does not support float types - auto rounded = cudf::round( - table.column(2), - 2, - cudf::rounding_method::HALF_EVEN, - chunk.stream(), - ctx->br()->device_mr() - ); -#pragma GCC diagnostic pop - auto result = cudf_streaming::streaming::to_message( - 0, - std::make_unique( - std::make_unique( - cudf::table_view({table.column(0), table.column(1), rounded->view()}), - chunk.stream(), - ctx->br()->device_mr() - ), - chunk.stream() - ) - ); - co_await ch_out->send(std::move(result)); - co_await ch_out->drain(ctx->executor()); -} - -} // namespace - -/** - * @brief Run a derived version of TPC-H query 9. - * - * The SQL form of the query is: - * @code{.sql} - * select - * nation, - * o_year, - * round(sum(amount), 2) as sum_profit - * from - * ( - * select - * n_name as nation, - * year(o_orderdate) as o_year, - * l_extendedprice * (1 - l_discount) - ps_supplycost * l_quantity as amount - * from - * part, - * supplier, - * lineitem, - * partsupp, - * orders, - * nation - * where - * s_suppkey = l_suppkey - * and ps_suppkey = l_suppkey - * and ps_partkey = l_partkey - * and p_partkey = l_partkey - * and o_orderkey = l_orderkey - * and s_nationkey = n_nationkey - * and p_name like '%green%' - * ) as profit - * group by - * nation, - * o_year - * order by - * nation, - * o_year desc - * @endcode{} - */ -int main(int argc, char** argv) { - rapidsmpf::ndsh::FinalizeMPI finalize{}; - cudaFree(nullptr); - // work around https://github.com/rapidsai/cudf/issues/20849 - cudf::initialize(); - auto mr = rmm::mr::cuda_async_memory_resource{}; - auto arguments = rapidsmpf::ndsh::parse_arguments(argc, argv); - auto [ctx, comm] = rapidsmpf::ndsh::create_context(arguments, std::move(mr)); - std::string output_path = arguments.output_file; - std::vector timings; - for (int i = 0; i < arguments.num_iterations; i++) { - rapidsmpf::OpID op_id{0}; - std::vector actors; - auto start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Constructing Q9 pipeline"); - auto part = ctx->create_channel(); - auto filtered_part = ctx->create_channel(); - auto partsupp = ctx->create_channel(); - auto part_x_partsupp = ctx->create_channel(); - auto supplier = ctx->create_channel(); - auto lineitem = ctx->create_channel(); - auto supplier_x_part_x_partsupp = ctx->create_channel(); - auto supplier_x_part_x_partsupp_x_lineitem = ctx->create_channel(); - actors.push_back(read_part( - ctx, - comm, - part, - /* num_tickets */ 4, - arguments.num_rows_per_chunk, - arguments.input_directory - )); // p_partkey, p_name - actors.push_back(filter_part(ctx, part, filtered_part)); // p_partkey - actors.push_back(read_partsupp( - ctx, - comm, - partsupp, - /* num_tickets */ 4, - arguments.num_rows_per_chunk, - arguments.input_directory - )); // ps_partkey, ps_suppkey, ps_supplycost - actors.push_back( - // p_partkey x ps_partkey - rapidsmpf::ndsh::inner_join_broadcast( - ctx, - comm, - filtered_part, - partsupp, - part_x_partsupp, - {0}, - {0}, - rapidsmpf::OpID{static_cast(10 * i + op_id++)} - ) // p_partkey/ps_partkey, ps_suppkey, ps_supplycost - ); - actors.push_back(read_supplier( - ctx, - comm, - supplier, - /* num_tickets */ 4, - arguments.num_rows_per_chunk, - arguments.input_directory - )); // s_nationkey, s_suppkey - actors.push_back( - // s_suppkey x ps_suppkey - rapidsmpf::ndsh::inner_join_broadcast( - ctx, - comm, - supplier, - part_x_partsupp, - supplier_x_part_x_partsupp, - {1}, - {1}, - rapidsmpf::OpID{static_cast(10 * i + op_id++)} - - ) // s_nationkey, s_suppkey/ps_suppkey, p_partkey/ps_partkey, - // ps_supplycost - ); - actors.push_back(read_lineitem( - ctx, - comm, - lineitem, - /* num_tickets */ 4, - arguments.num_rows_per_chunk, - arguments.input_directory - )); // l_discount, l_extendedprice, l_orderkey, l_partkey, l_quantity, - // l_suppkey - actors.push_back( - // [p_partkey, ps_suppkey] x [l_partkey, l_suppkey] - rapidsmpf::ndsh::inner_join_broadcast( - ctx, - comm, - supplier_x_part_x_partsupp, - lineitem, - supplier_x_part_x_partsupp_x_lineitem, - {2, 1}, - {3, 5}, - rapidsmpf::OpID{static_cast(10 * i + op_id++)}, - rapidsmpf::ndsh::KeepKeys::NO - ) // s_nationkey, ps_supplycost, - // l_discount, l_extendedprice, l_orderkey, l_quantity - ); - auto nation = ctx->create_channel(); - auto orders = ctx->create_channel(); - actors.push_back( - read_nation( - ctx, - comm, - nation, - /* num_tickets */ 4, - arguments.num_rows_per_chunk, - arguments.input_directory - ) // n_name, n_nationkey - ); - actors.push_back( - read_orders( - ctx, - comm, - orders, - /* num_tickets */ 4, - arguments.num_rows_per_chunk, - arguments.input_directory - ) // o_orderdate, o_orderkey - ); - auto all_joined = ctx->create_channel(); - auto supplier_x_part_x_partsupp_x_lineitem_x_orders = ctx->create_channel(); - if (arguments.use_shuffle_join) { - auto supplier_x_part_x_partsupp_x_lineitem_shuffled = - ctx->create_channel(); - auto orders_shuffled = ctx->create_channel(); - // TODO: customisable - std::uint32_t num_partitions = 16; - actors.push_back( - rapidsmpf::ndsh::shuffle( - ctx, - comm, - supplier_x_part_x_partsupp_x_lineitem, - supplier_x_part_x_partsupp_x_lineitem_shuffled, - {4}, - num_partitions, - rapidsmpf::OpID{static_cast(10 * i + op_id++)} - ) - ); - actors.push_back( - rapidsmpf::ndsh::shuffle( - ctx, - comm, - orders, - orders_shuffled, - {1}, - num_partitions, - rapidsmpf::OpID{static_cast(10 * i + op_id++)} - ) - ); - actors.push_back( - // l_orderkey x o_orderkey - rapidsmpf::ndsh::inner_join_shuffle( - ctx, - comm, - supplier_x_part_x_partsupp_x_lineitem_shuffled, - orders_shuffled, - supplier_x_part_x_partsupp_x_lineitem_x_orders, - {4}, - {1}, - rapidsmpf::ndsh::KeepKeys::NO - ) // s_nationkey, ps_supplycost, l_discount, l_extendedprice, - // l_quantity, o_orderdate - ); - } else { - actors.push_back( - // l_orderkey x o_orderkey - rapidsmpf::ndsh::inner_join_broadcast( - ctx, - comm, - supplier_x_part_x_partsupp_x_lineitem, - orders, - supplier_x_part_x_partsupp_x_lineitem_x_orders, - {4}, - {1}, - rapidsmpf::OpID{static_cast(10 * i + op_id++)}, - rapidsmpf::ndsh::KeepKeys::NO - ) // s_nationkey, ps_supplycost, l_discount, l_extendedprice, - // l_quantity, o_orderdate - ); - } - actors.push_back( - // n_nationkey x s_nationkey - rapidsmpf::ndsh::inner_join_broadcast( - ctx, - comm, - nation, - supplier_x_part_x_partsupp_x_lineitem_x_orders, - all_joined, - {1}, - {0}, - rapidsmpf::OpID{static_cast(10 * i + op_id++)}, - rapidsmpf::ndsh::KeepKeys::NO - ) // n_name, ps_supplycost, l_discount, l_extendedprice, - // l_quantity, o_orderdate - ); - auto chunkwise_groupby_input = ctx->create_channel(); - actors.push_back(select_columns(ctx, all_joined, chunkwise_groupby_input)); - auto chunkwise_groupby_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - chunkwise_groupby_input, - chunkwise_groupby_output, - {0, 1}, - chunkwise_groupby_requests(), - cudf::null_policy::INCLUDE - ) - ); - auto final_groupby_input = ctx->create_channel(); - if (comm->nranks() > 1) { - actors.push_back( - rapidsmpf::ndsh::broadcast( - ctx, - comm, - chunkwise_groupby_output, - final_groupby_input, - static_cast(10 * i + op_id++), - rapidsmpf::streaming::AllGather::Ordered::NO - ) - ); - } else { - actors.push_back( - rapidsmpf::ndsh::concatenate( - ctx, chunkwise_groupby_output, final_groupby_input - ) - ); - } - if (comm->rank() == 0) { - auto final_groupby_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - final_groupby_input, - final_groupby_output, - {0, 1}, - chunkwise_groupby_requests(), - cudf::null_policy::INCLUDE - ) - ); - auto sorted_input = ctx->create_channel(); - actors.push_back( - round_sum_profit(ctx, final_groupby_output, sorted_input) - ); - auto sorted_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_sort_by( - ctx, - sorted_input, - sorted_output, - {0, 1}, - {0, 1, 2}, - {cudf::order::ASCENDING, cudf::order::DESCENDING}, - {cudf::null_order::BEFORE, cudf::null_order::BEFORE} - ) - ); - actors.push_back( - rapidsmpf::ndsh::write_parquet( - ctx, - sorted_output, - cudf::io::sink_info{output_path}, - {"nation", "o_year", "sum_profit"} - ) - ); - } else { - actors.push_back(rapidsmpf::ndsh::sink_channel(ctx, final_groupby_input)); - } - } - auto end = std::chrono::steady_clock::now(); - std::chrono::duration pipeline = end - start; - start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Q9 Iteration"); - rapidsmpf::streaming::run_actor_network(std::move(actors)); - } - end = std::chrono::steady_clock::now(); - std::chrono::duration compute = end - start; - timings.push_back(pipeline.count()); - timings.push_back(compute.count()); - auto statistics = ctx->statistics(); - comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} - )); - statistics->clear(); - } - if (comm->rank() == 0) { - for (std::size_t i = 0; i < safe_cast(arguments.num_iterations); i++) - { - comm->logger()->print( - "Iteration ", i, " pipeline construction time [s]: ", timings[2 * i] - ); - comm->logger()->print( - "Iteration ", i, " compute time [s]: ", timings[2 * i + 1] - ); - } - } - return 0; -} diff --git a/cpp/benchmarks/streaming/ndsh/q21.cpp b/cpp/benchmarks/streaming/ndsh/q21.cpp deleted file mode 100644 index fb69d98ab..000000000 --- a/cpp/benchmarks/streaming/ndsh/q21.cpp +++ /dev/null @@ -1,1029 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "concatenate.hpp" -#include "groupby.hpp" -#include "join.hpp" -#include "parquet_writer.hpp" -#include "sort.hpp" -#include "utils.hpp" - -using rapidsmpf::safe_cast; - -namespace { - -rapidsmpf::streaming::Actor read_lineitem( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const input_directory, - std::vector columns, - std::shared_ptr latch = nullptr -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "lineitem") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names(columns) - .build(); - if (latch != nullptr) { - co_await *latch; - } - co_return co_await cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor read_nation( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "nation") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({"n_nationkey"}) - .build(); - // filter: "n_name" == "SAUDI ARABIA" - auto filter_expr = [&]() -> std::unique_ptr { - auto stream = ctx->br()->stream_pool().get_stream(); - auto owner = new std::vector; - constexpr auto name = "SAUDI ARABIA"; - owner->push_back( - std::make_shared( - name, /* is_valid = */ true, stream, ctx->br()->device_mr() - ) - ); - owner->push_back( - std::make_shared( - *std::any_cast>(owner->at(0)) - ) - ); - owner->push_back(std::make_shared("n_name")); - owner->push_back( - std::make_shared( - cudf::ast::ast_operator::EQUAL, - *std::any_cast>( - owner->at(2) - ), - *std::any_cast>(owner->at(1)) - ) - ); - return std::make_unique( - stream, - *std::any_cast>(owner->back()), - rapidsmpf::OwningWrapper(static_cast(owner), [](void* p) { - delete static_cast*>(p); - }) - ); - }(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, - comm, - ch_out, - num_producers, - options, - num_rows_per_chunk, - std::move(filter_expr) - ); -} - -rapidsmpf::streaming::Actor read_orders( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "orders") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({"o_orderkey"}) - .build(); - // filter: "o_orderstatus" == "F" - auto filter_expr = [&]() -> std::unique_ptr { - auto stream = ctx->br()->stream_pool().get_stream(); - auto owner = new std::vector; - constexpr auto status = "F"; - owner->push_back( - std::make_shared( - status, /* is_valid = */ true, stream, ctx->br()->device_mr() - ) - ); - owner->push_back( - std::make_shared( - *std::any_cast>(owner->at(0)) - ) - ); - owner->push_back( - std::make_shared("o_orderstatus") - ); - owner->push_back( - std::make_shared( - cudf::ast::ast_operator::EQUAL, - *std::any_cast>( - owner->at(2) - ), - *std::any_cast>(owner->at(1)) - ) - ); - return std::make_unique( - stream, - *std::any_cast>(owner->back()), - rapidsmpf::OwningWrapper(static_cast(owner), [](void* p) { - delete static_cast*>(p); - }) - ); - }(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, - comm, - ch_out, - num_producers, - options, - num_rows_per_chunk, - std::move(filter_expr) - ); -} - -rapidsmpf::streaming::Actor read_orders_with_bloom_filter( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr bloom_filter_in, - std::shared_ptr ch_out, - std::vector filter_keys, - cudf_streaming::streaming::BloomFilter bloom_filter, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const input_directory -) { - auto filter_passthrough = ctx->create_channel(); - auto orders_passthrough = ctx->create_channel(); - rapidsmpf::streaming::ShutdownAtExit c{ - bloom_filter_in, ch_out, filter_passthrough, orders_passthrough - }; - co_await ctx->executor()->schedule(); - // We want to await the bloom_filter being ready before kicking off the tasks to read - // orders and apply the filter. This way, the read won't start until the bloom filter - // is ready and we won't stack up num_producers chunks waiting for ages. - auto filter = co_await bloom_filter_in->receive(); - auto passthrough = [&]() -> coro::task { - co_await filter_passthrough->send(std::move(filter)); - co_await filter_passthrough->drain(ctx->executor()); - }; - rapidsmpf::streaming::coro_results( - co_await coro::when_all( - passthrough(), - read_orders( - ctx, - comm, - orders_passthrough, - num_producers, - num_rows_per_chunk, - input_directory - ), - bloom_filter.apply( - filter_passthrough, orders_passthrough, ch_out, filter_keys - ) - ) - ); - co_await ch_out->drain(ctx->executor()); -} - -rapidsmpf::streaming::Actor read_supplier( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::size_type num_rows_per_chunk, - std::string const& input_directory -) { - auto files = rapidsmpf::ndsh::detail::list_parquet_files( - rapidsmpf::ndsh::detail::get_table_path(input_directory, "supplier") - ); - auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) - .column_names({"s_suppkey", "s_nationkey", "s_name"}) - .build(); - return cudf_streaming::streaming::actor::read_parquet( - ctx, comm, ch_out, num_producers, options, num_rows_per_chunk - ); -} - -rapidsmpf::streaming::Actor filter_lineitem( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - - while (!ch_out->is_shutdown()) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - - auto mask = cudf::binary_operation( - chunk.table_view().column(2), - chunk.table_view().column(3), - cudf::binary_operator::GREATER, - cudf::data_type(cudf::type_id::BOOL8), - chunk.stream(), - ctx->br()->device_mr() - ); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - cudf::apply_boolean_mask( - chunk.table_view().select({0, 1}), - mask->view(), - chunk.stream(), - ctx->br()->device_mr() - ), - chunk.stream() - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -rapidsmpf::streaming::Actor filter_grouped_greater( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::shared_ptr latch -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - - while (!ch_out->is_shutdown()) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - - auto mask = cudf::binary_operation( - chunk.table_view().column(1), - cudf::numeric_scalar( - 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr() - ), - cudf::binary_operator::GREATER, - cudf::data_type(cudf::type_id::BOOL8), - chunk.stream(), - ctx->br()->device_mr() - ); - latch->count_down(); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - cudf::apply_boolean_mask( - chunk.table_view().select({0}), - mask->view(), - chunk.stream(), - ctx->br()->device_mr() - ), - chunk.stream() - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -rapidsmpf::streaming::Actor filter_grouped_equal( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - - auto mask = cudf::binary_operation( - chunk.table_view().column(1), - cudf::numeric_scalar( - 1, /* is_valid = */ true, chunk.stream(), ctx->br()->device_mr() - ), - cudf::binary_operator::EQUAL, - cudf::data_type(cudf::type_id::BOOL8), - chunk.stream(), - ctx->br()->device_mr() - ); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - cudf::apply_boolean_mask( - chunk.table_view().select({0}), - mask->view(), - chunk.stream(), - ctx->br()->device_mr() - ), - chunk.stream() - ) - ) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -rapidsmpf::streaming::Actor fanout_bounded( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - std::shared_ptr ch1_out, - std::vector ch1_cols, - std::shared_ptr ch2_out -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch1_out, ch2_out}; - - co_await ctx->executor()->schedule(); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - // Here, we know that copying ch1_cols (a single col) is better than copying - // ch2_cols (the whole table) - std::vector> tasks; - if (!ch1_out->is_shutdown()) { - auto msg1 = cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::make_unique( - chunk.table_view().select(ch1_cols), - chunk.stream(), - ctx->br()->device_mr() - ), - chunk.stream() - ) - ); - tasks.push_back(ch1_out->send(std::move(msg1))); - } - if (!ch2_out->is_shutdown()) { - // TODO: We know here that ch2 wants the whole table. - tasks.push_back(ch2_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::move(chunk) - ) - ) - )); - } - if (!std::ranges::any_of( - rapidsmpf::streaming::coro_results( - co_await coro::when_all(std::move(tasks)) - ), - std::identity{} - )) - { - comm->logger()->print("Breaking after ", msg.sequence_number()); - break; - }; - } - - rapidsmpf::streaming::coro_results( - co_await coro::when_all( - ch1_out->drain(ctx->executor()), ch2_out->drain(ctx->executor()) - ) - ); -} - -rapidsmpf::streaming::Actor slice( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::int64_t global_start, - std::int64_t global_end -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - co_await ctx->executor()->schedule(); - std::int64_t current_row = 0; - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - - if (global_start == global_end) { - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - cudf::empty_like(chunk.table_view()), chunk.stream() - ) - ) - ); - break; - } - - auto num_rows = chunk.table_view().num_rows(); - - std::int64_t chunk_start = current_row; - std::int64_t chunk_end = current_row + num_rows; - - std::int64_t slice_start = std::max(chunk_start, global_start); - std::int64_t slice_end = std::min(chunk_end, global_end); - - if (slice_start < slice_end) { - auto local_start = static_cast(slice_start - chunk_start); - auto local_end = static_cast(slice_end - chunk_start); - - if (local_start == 0 && local_end == num_rows) { - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::move(chunk) - ) - ) - ); - } else { - auto sliced_table = std::make_unique( - cudf::slice(chunk.table_view(), {local_start, local_end})[0], - chunk.stream(), - ctx->br()->device_mr() - ); - co_await ch_out->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::move(sliced_table), chunk.stream() - ) - ) - ); - } - } - current_row += num_rows; - if (current_row >= global_end) { - break; - } - } - co_await ch_out->drain(ctx->executor()); -} - -std::vector count_groupby_request() { - auto requests = std::vector(); - std::vector()>> aggs; - // count(*) - aggs.emplace_back([]() { - return cudf::make_count_aggregation( - cudf::null_policy::INCLUDE - ); - }); - requests.emplace_back(0, std::move(aggs)); - return requests; -} - -std::vector sum_groupby_request( - cudf::size_type column -) { - auto requests = std::vector(); - std::vector()>> aggs; - // count(*) - aggs.emplace_back(cudf::make_sum_aggregation); - requests.emplace_back(column, std::move(aggs)); - return requests; -} - -rapidsmpf::streaming::Actor populate_bloom_filter( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys, - rapidsmpf::OpID tag, - cudf_streaming::streaming::BloomFilter bloom_filter -) { - rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; - auto passthrough = ctx->create_channel(); - auto selector = [&]() -> rapidsmpf::streaming::Actor { - rapidsmpf::streaming::ShutdownAtExit c{passthrough}; - co_await ctx->executor()->schedule(); - while (!passthrough->is_shutdown()) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = co_await msg.release() - .make_available(ctx); - auto stream = chunk.stream(); - auto out = std::make_unique( - chunk.table_view().select(keys), stream, ctx->br()->device_mr() - ); - std::ignore = std::move(chunk); - co_await passthrough->send( - cudf_streaming::streaming::to_message( - msg.sequence_number(), - std::make_unique( - std::move(out), stream - ) - ) - ); - } - comm->logger()->print("Sent all things through filter"); - co_await passthrough->drain(ctx->executor()); - }; - rapidsmpf::streaming::coro_results( - co_await coro::when_all(selector(), bloom_filter.build(passthrough, ch_out, tag)) - ); - co_await ch_out->drain(ctx->executor()); -} -} // namespace - -/** - * @brief Run a derived version of TPC-H query 21. - * - * The SQL form of the query is: - * @code{.sql} - * select - * s_name, - * count(*) as numwait - * from - * supplier, - * lineitem l1, - * orders, - * nation - * where - * s_suppkey = l1.l_suppkey - * and o_orderkey = l1.l_orderkey - * and o_orderstatus = 'F' - * and l1.l_receiptdate > l1.l_commitdate - * and exists ( - * select - * * - * from - * lineitem l2 - * where - * l2.l_orderkey = l1.l_orderkey - * and l2.l_suppkey <> l1.l_suppkey - * ) - * and not exists ( - * select - * * - * from - * lineitem l3 - * where - * l3.l_orderkey = l1.l_orderkey - * and l3.l_suppkey <> l1.l_suppkey - * and l3.l_receiptdate > l3.l_commitdate - * ) - * and s_nationkey = n_nationkey - * and n_name = 'SAUDI ARABIA' - * group by - * s_name - * order by - * numwait desc, - * s_name - * limit 100 - * @endcode{} - */ -int main(int argc, char** argv) { - rapidsmpf::ndsh::FinalizeMPI finalize{}; - cudaFree(nullptr); - cudf::initialize(); - auto mr = rmm::mr::cuda_async_memory_resource{}; - auto arguments = rapidsmpf::ndsh::parse_arguments(argc, argv); - auto [ctx, comm] = rapidsmpf::ndsh::create_context(arguments, std::move(mr)); - std::string output_path = arguments.output_file; - std::vector timings; - int l2size; - int device; - RAPIDSMPF_CUDA_TRY(cudaGetDevice(&device)); - RAPIDSMPF_CUDA_TRY(cudaDeviceGetAttribute(&l2size, cudaDevAttrL2CacheSize, device)); - auto const num_filter_blocks = - cudf_streaming::integrations::BloomFilter::fitting_num_blocks( - static_cast(l2size) - ); - for (int i = 0; i < arguments.num_iterations; i++) { - int op_id{0}; - std::vector actors; - auto start = std::chrono::steady_clock::now(); - // TODO: configurable/adaptive - std::uint32_t num_shuffle_partitions = 16; - { - // Idea: - // We need to shuffle lineitem, but then we always join/groupby on l_orderkey - // The join against orders will be done via shuffle join. The selectivity of - // the supplier x nation join is about 1/25, so that can be a broadcast join. - // - // We commit to reading the lineitem table twice because one of the reads only - // needs a single column and we can do all of the processing of it before - // reading the second time, keeping memory under control. - - RAPIDSMPF_NVTX_SCOPED_RANGE("Constructing Q21 pipeline"); - auto lineitem_orderkey = ctx->create_channel(); - actors.push_back(read_lineitem( - ctx, - comm, - lineitem_orderkey, - /* num_tickets */ 2, - arguments.num_rows_per_chunk, - arguments.input_directory, - {"l_orderkey"} - )); // "l_orderkey" - auto lineitem_orderkey_grouped = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - lineitem_orderkey, - lineitem_orderkey_grouped, - {0}, - count_groupby_request(), - cudf::null_policy::INCLUDE - ) - ); // l_orderkey, count(*) - auto lineitem_orderkey_shuffled = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::shuffle( - ctx, - comm, - lineitem_orderkey_grouped, - lineitem_orderkey_shuffled, - {0}, - num_shuffle_partitions, - static_cast(10 * i) + op_id++ - ) - ); // l_orderkey, count(*) [shuffled on l_orderkey] - auto lineitem_orderkey_shuffled_grouped = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - lineitem_orderkey_shuffled, - lineitem_orderkey_shuffled_grouped, - {0}, - sum_groupby_request(1), - cudf::null_policy::INCLUDE - ) - ); // l_orderkey, sum(count(*)) [groupby done] - auto lineitem_orderkey_filtered = ctx->create_channel(); - auto latch = std::make_shared(1); - actors.push_back(filter_grouped_greater( - ctx, lineitem_orderkey_shuffled_grouped, lineitem_orderkey_filtered, latch - )); // l_orderkey [sum(count(*)) > 1, releases lineitem read] - auto lineitem_suppkey = ctx->create_channel(); - actors.push_back(read_lineitem( - ctx, - comm, - lineitem_suppkey, - /* num_tickets */ 2, - arguments.num_rows_per_chunk, - arguments.input_directory, - {"l_orderkey", "l_suppkey", "l_receiptdate", "l_commitdate"}, - latch - )); // l_orderkey, l_suppkey, l_receiptdate, l_commitdate - // [released once filter_grouped_greater has seen an input] - auto lineitem_suppkey_filtered = ctx->create_channel(); - actors.push_back( - filter_lineitem(ctx, lineitem_suppkey, lineitem_suppkey_filtered) - ); // l_orderkey, l_suppkey - auto lineitem_suppkey_shuffled = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::shuffle( - ctx, - comm, - lineitem_suppkey_filtered, - lineitem_suppkey_shuffled, - {0}, - num_shuffle_partitions, - static_cast(10 * i) + op_id++ - ) - ); // l_orderkey, l_suppkey [shuffled on l_orderkey] - auto lineitem_self_joined = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::inner_join_shuffle( - ctx, - comm, - lineitem_orderkey_filtered, - lineitem_suppkey_shuffled, - lineitem_self_joined, - {0}, - {0} - ) - ); // l_orderkey, l_suppkey [join complete] - - auto joined_grouped_input = ctx->create_channel(); - auto joined_input = ctx->create_channel(); - actors.push_back(fanout_bounded( - ctx, comm, lineitem_self_joined, joined_grouped_input, {0}, joined_input - )); // l_orderkey (in joined_grouped_input), - // l_orderkey l_suppkey (in joined_input) - auto joined_grouped_len = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - joined_grouped_input, - joined_grouped_len, - {0}, - count_groupby_request(), - cudf::null_policy::INCLUDE - ) - ); // l_orderkey, count(*) [complete, because partitioned on l_orderkey] - auto joined_grouped_filter = ctx->create_channel(); - actors.push_back( - filter_grouped_equal(ctx, joined_grouped_len, joined_grouped_filter) - ); // l_orderkey [count(*) == 1] - auto lineitem_joined = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::inner_join_shuffle( - ctx, - comm, - joined_grouped_filter, - joined_input, - lineitem_joined, - {0}, - {0} - ) - ); // l_orderkey, l_suppkey - auto supplier = ctx->create_channel(); - auto nation = ctx->create_channel(); - actors.push_back(read_supplier( - ctx, - comm, - supplier, - 2, - arguments.num_rows_per_chunk, - arguments.input_directory - )); // s_suppkey, s_nationkey, s_name - actors.push_back(read_nation( - ctx, - comm, - nation, - 1, - arguments.num_rows_per_chunk, - arguments.input_directory - )); // n_nationkey - auto supp_x_nation = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::inner_join_broadcast( - ctx, - comm, - nation, - supplier, - supp_x_nation, - {0}, - {1}, - static_cast(10 * i + op_id++), - rapidsmpf::ndsh::KeepKeys::NO - ) - ); // s_suppkey, s_name - auto supp_nation_lineitem = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::inner_join_broadcast( - ctx, - comm, - supp_x_nation, - lineitem_joined, - supp_nation_lineitem, - {0}, - {1}, - static_cast(10 * i + op_id++), - rapidsmpf::ndsh::KeepKeys::NO - ) - ); // s_name, l_orderkey [this join is quite selective] - // OK, we're going to pre-filter the orders table before joining using a bloom - // filter. - // This has two consequences: - // 1. We shuffle less data - // 2. This acts as a latch: the orders table is not read and shuffled until - // the other side is "ready", reducing memory pressure. - auto bloom_input = ctx->create_channel(); - auto snl_passthrough = ctx->create_channel(); - // Bloom filter needs to see all the input before we can release the orders - // read, so need unbounded fanout. - actors.push_back( - rapidsmpf::streaming::actor::fanout( - ctx, - supp_nation_lineitem, - {bloom_input, snl_passthrough}, - rapidsmpf::streaming::actor::FanoutPolicy::UNBOUNDED - ) - ); - auto bloom_output = ctx->create_channel(); - auto bloom_filter = cudf_streaming::streaming::BloomFilter( - ctx, comm, cudf::DEFAULT_HASH_SEED, num_filter_blocks - ); - // Select the relevant key column(s) and build filter. - actors.push_back(populate_bloom_filter( - ctx, - comm, - bloom_input, - bloom_output, - {1}, - static_cast(10 * i + op_id++), - bloom_filter - )); - auto orders = ctx->create_channel(); - auto shuffled_orders = ctx->create_channel(); - // OK, now we obtain the filter, and release the orders read which we apply - // the filter to before sending on to the shuffle. - actors.push_back(read_orders_with_bloom_filter( - ctx, - comm, - bloom_output, - orders, - {0}, - bloom_filter, - 2, - arguments.num_rows_per_chunk, - arguments.input_directory - )); // o_orderkey - actors.push_back( - rapidsmpf::ndsh::shuffle( - ctx, - comm, - orders, - shuffled_orders, - {0}, - num_shuffle_partitions, - static_cast(10 * i + op_id++) - ) - ); // o_orderkey [shuffled on o_orderkey] - auto all_joined = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::inner_join_shuffle( - ctx, - comm, - snl_passthrough, - shuffled_orders, - all_joined, - {1}, - {0}, - rapidsmpf::ndsh::KeepKeys::NO - ) - ); // s_name - auto chunked_groupby = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - all_joined, - chunked_groupby, - {0}, - count_groupby_request(), - cudf::null_policy::INCLUDE - ) - ); // s_name, count(*) - auto final_groupby_input = ctx->create_channel(); - if (comm->nranks() > 1) { - actors.push_back( - rapidsmpf::ndsh::broadcast( - ctx, - comm, - chunked_groupby, - final_groupby_input, - static_cast(10 * i + op_id++), - rapidsmpf::streaming::AllGather::Ordered::NO - ) - ); - } else { - actors.push_back( - rapidsmpf::ndsh::concatenate( - ctx, chunked_groupby, final_groupby_input - ) - ); - } - if (comm->rank() == 0) { - auto final_groupby_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_group_by( - ctx, - final_groupby_input, - final_groupby_output, - {0}, - sum_groupby_request(1), - cudf::null_policy::INCLUDE - ) - ); // s_name, sum(count(*)) [only a single partition now due to the - // broadcast] - auto sorted_output = ctx->create_channel(); - actors.push_back( - rapidsmpf::ndsh::chunkwise_sort_by( - ctx, - final_groupby_output, - sorted_output, - {1, 0}, - {0, 1}, - {cudf::order::DESCENDING, cudf::order::ASCENDING}, - {cudf::null_order::BEFORE, cudf::null_order::BEFORE} - ) - ); - auto sliced = ctx->create_channel(); - actors.push_back(slice(ctx, sorted_output, sliced, 0, 100)); - actors.push_back( - rapidsmpf::ndsh::write_parquet( - ctx, - sliced, - cudf::io::sink_info{arguments.output_file}, - {"s_name", "numwait"} - ) - ); - } else { - actors.push_back(rapidsmpf::ndsh::sink_channel(ctx, final_groupby_input)); - } - } - auto end = std::chrono::steady_clock::now(); - std::chrono::duration pipeline = end - start; - start = std::chrono::steady_clock::now(); - { - RAPIDSMPF_NVTX_SCOPED_RANGE("Q21 Iteration"); - rapidsmpf::streaming::run_actor_network(std::move(actors)); - } - end = std::chrono::steady_clock::now(); - std::chrono::duration compute = end - start; - timings.push_back(pipeline.count()); - timings.push_back(compute.count()); - auto statistics = ctx->statistics(); - comm->logger()->print(statistics->report( - {.mr = ctx->br()->device_mr(), .pinned_mr = ctx->br()->try_pinned_mr()} - )); - statistics->clear(); - } - - if (comm->rank() == 0) { - for (std::size_t i = 0; i < safe_cast(arguments.num_iterations); i++) - { - comm->logger()->print( - "Iteration ", i, " pipeline construction time [s]: ", timings[2 * i] - ); - comm->logger()->print( - "Iteration ", i, " compute time [s]: ", timings[2 * i + 1] - ); - } - } - return 0; -} diff --git a/cpp/benchmarks/streaming/ndsh/sort.cpp b/cpp/benchmarks/streaming/ndsh/sort.cpp deleted file mode 100644 index 78ebc8f2a..000000000 --- a/cpp/benchmarks/streaming/ndsh/sort.cpp +++ /dev/null @@ -1,70 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "sort.hpp" - -#include -#include -#include - -#include -#include -#include - -#include -#include - -namespace rapidsmpf::ndsh { - -rapidsmpf::streaming::Actor chunkwise_sort_by( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys, - std::vector values, - std::vector order, - std::vector null_order -) { - streaming::ShutdownAtExit c{ch_in, ch_out}; - co_await ctx->executor()->schedule(); - auto make_table = [&](cudf_streaming::streaming::TableChunk& chunk) { - if (std::ranges::equal(keys, values)) { - return cudf::sort( - chunk.table_view().select(keys), - order, - null_order, - chunk.stream(), - ctx->br()->device_mr() - ); - } else { - return cudf::sort_by_key( - chunk.table_view().select(values), - chunk.table_view().select(keys), - order, - null_order, - chunk.stream(), - ctx->br()->device_mr() - ); - } - }; - while (!ch_out->is_shutdown()) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = - co_await msg.release().make_available( - ctx - ); - co_await ch_out->send(to_message( - msg.sequence_number(), - std::make_unique( - make_table(chunk), chunk.stream() - ) - )); - } - co_await ch_out->drain(ctx->executor()); -} -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/sort.hpp b/cpp/benchmarks/streaming/ndsh/sort.hpp deleted file mode 100644 index e7e38607e..000000000 --- a/cpp/benchmarks/streaming/ndsh/sort.hpp +++ /dev/null @@ -1,39 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include -#include - -#include - -#include -#include - -namespace rapidsmpf::ndsh { - -/** - * @brief Sort chunks in a channel - * - * @param ctx Streaming context - * @param ch_in Input channel of `TableChunk`s - * @param ch_out Output channel of sorted `TableChunk`s - * @param keys Indices of key columns in the input channel - * @param values Indices of value columns in the input channel - * @param order Sort order for each column named in `keys` - * @param null_order Null precedence for each column named in `keys` - * - * @return Coroutine representing the sort - */ -[[nodiscard]] rapidsmpf::streaming::Actor chunkwise_sort_by( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys, - std::vector values, - std::vector order, - std::vector null_order -); -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/sql/q01.sql b/cpp/benchmarks/streaming/ndsh/sql/q01.sql deleted file mode 100644 index 8b769a130..000000000 --- a/cpp/benchmarks/streaming/ndsh/sql/q01.sql +++ /dev/null @@ -1,21 +0,0 @@ -select - l_returnflag, - l_linestatus, - sum(l_quantity) as sum_qty, - sum(l_extendedprice) as sum_base_price, - sum(l_extendedprice * (1 - l_discount)) as sum_disc_price, - sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge, - avg(l_quantity) as avg_qty, - avg(l_extendedprice) as avg_price, - avg(l_discount) as avg_disc, - count(*) as count_order -from - lineitem -where - l_shipdate <= DATE '1998-09-02' -group by - l_returnflag, - l_linestatus -order by - l_returnflag, - l_linestatus diff --git a/cpp/benchmarks/streaming/ndsh/sql/q03.sql b/cpp/benchmarks/streaming/ndsh/sql/q03.sql deleted file mode 100644 index c93865c19..000000000 --- a/cpp/benchmarks/streaming/ndsh/sql/q03.sql +++ /dev/null @@ -1,23 +0,0 @@ -select - l_orderkey, - sum(l_extendedprice * (1 - l_discount)) as revenue, - o_orderdate, - o_shippriority -from - customer, - orders, - lineitem -where - c_mktsegment = 'BUILDING' - and c_custkey = o_custkey - and l_orderkey = o_orderkey - and o_orderdate < '1995-03-15' - and l_shipdate > '1995-03-15' -group by - l_orderkey, - o_orderdate, - o_shippriority -order by - revenue desc, - o_orderdate -limit 10 diff --git a/cpp/benchmarks/streaming/ndsh/sql/q04.sql b/cpp/benchmarks/streaming/ndsh/sql/q04.sql deleted file mode 100644 index 84a9733ed..000000000 --- a/cpp/benchmarks/streaming/ndsh/sql/q04.sql +++ /dev/null @@ -1,21 +0,0 @@ -select - o_orderpriority, - count(*) as order_count -from - orders -where - o_orderdate >= timestamp '1993-07-01' - and o_orderdate < timestamp '1993-07-01' + interval '3' month - and exists ( - select - * - from - lineitem - where - l_orderkey = o_orderkey - and l_commitdate < l_receiptdate - ) -group by - o_orderpriority -order by - o_orderpriority diff --git a/cpp/benchmarks/streaming/ndsh/sql/q09.sql b/cpp/benchmarks/streaming/ndsh/sql/q09.sql deleted file mode 100644 index 7b6a0d609..000000000 --- a/cpp/benchmarks/streaming/ndsh/sql/q09.sql +++ /dev/null @@ -1,32 +0,0 @@ -select - nation, - o_year, - round(sum(amount), 2) as sum_profit -from - ( - select - n_name as nation, - year(o_orderdate)::INT16 as o_year, - l_extendedprice * (1 - l_discount) - ps_supplycost * l_quantity as amount - from - part, - supplier, - lineitem, - partsupp, - orders, - nation - where - s_suppkey = l_suppkey - and ps_suppkey = l_suppkey - and ps_partkey = l_partkey - and p_partkey = l_partkey - and o_orderkey = l_orderkey - and s_nationkey = n_nationkey - and p_name like '%green%' - ) as profit -group by - nation, - o_year -order by - nation, - o_year desc diff --git a/cpp/benchmarks/streaming/ndsh/sql/q17.sql b/cpp/benchmarks/streaming/ndsh/sql/q17.sql deleted file mode 100644 index 13a64f7bf..000000000 --- a/cpp/benchmarks/streaming/ndsh/sql/q17.sql +++ /dev/null @@ -1,17 +0,0 @@ -select - round(sum(l_extendedprice) / 7.0, 2) as avg_yearly -from - lineitem, - part -where - p_partkey = l_partkey - and p_brand = 'Brand#23' - and p_container = 'MED BOX' - and l_quantity < ( - select - 0.2 * avg(l_quantity) - from - read_parquet('/datasets/toaugspurger/tpch-rs/scale-10/lineitem/*.parquet') as lineitem - where - l_partkey = p_partkey - ) diff --git a/cpp/benchmarks/streaming/ndsh/sql/q18.sql b/cpp/benchmarks/streaming/ndsh/sql/q18.sql deleted file mode 100644 index bf61e83cd..000000000 --- a/cpp/benchmarks/streaming/ndsh/sql/q18.sql +++ /dev/null @@ -1,33 +0,0 @@ -select - c_name, - c_custkey, - o_orderkey, - o_orderdate as o_orderdat, - o_totalprice, - sum(l_quantity) as col6 -from - customer, - orders, - lineitem -where - o_orderkey in ( - select - l_orderkey - from - lineitem - group by - l_orderkey having - sum(l_quantity) > 300 - ) - and c_custkey = o_custkey - and o_orderkey = l_orderkey -group by - c_name, - c_custkey, - o_orderkey, - o_orderdate, - o_totalprice -order by - o_totalprice desc, - o_orderdate -limit 100 diff --git a/cpp/benchmarks/streaming/ndsh/sql/q21.sql b/cpp/benchmarks/streaming/ndsh/sql/q21.sql deleted file mode 100644 index eae182761..000000000 --- a/cpp/benchmarks/streaming/ndsh/sql/q21.sql +++ /dev/null @@ -1,40 +0,0 @@ -select - s_name, - count(*) as numwait -from - supplier, - lineitem l1, - orders, - nation -where - s_suppkey = l1.l_suppkey - and o_orderkey = l1.l_orderkey - and o_orderstatus = 'F' - and l1.l_receiptdate > l1.l_commitdate - and exists ( - select - * - from - lineitem l2 - where - l2.l_orderkey = l1.l_orderkey - and l2.l_suppkey <> l1.l_suppkey - ) - and not exists ( - select - * - from - lineitem l3 - where - l3.l_orderkey = l1.l_orderkey - and l3.l_suppkey <> l1.l_suppkey - and l3.l_receiptdate > l3.l_commitdate - ) - and s_nationkey = n_nationkey - and n_name = 'SAUDI ARABIA' -group by - s_name -order by - numwait desc, - s_name -limit 100 diff --git a/cpp/benchmarks/streaming/ndsh/utils.cpp b/cpp/benchmarks/streaming/ndsh/utils.cpp deleted file mode 100644 index 76237873a..000000000 --- a/cpp/benchmarks/streaming/ndsh/utils.cpp +++ /dev/null @@ -1,427 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include "utils.hpp" - -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace rapidsmpf::ndsh { -namespace detail { -std::vector list_parquet_files(std::string const root_path) { - auto root_entry = std::filesystem::directory_entry(std::filesystem::path(root_path)); - RAPIDSMPF_EXPECTS( - root_entry.exists() - && (root_entry.is_regular_file() || root_entry.is_directory()), - "Invalid file path", - std::runtime_error - ); - if (root_entry.is_regular_file()) { - RAPIDSMPF_EXPECTS( - root_path.ends_with(".parquet"), "Invalid filename", std::runtime_error - ); - return {root_path}; - } - std::vector result; - for (auto const& entry : std::filesystem::directory_iterator(root_path)) { - if (entry.is_regular_file()) { - std::string filename = entry.path().filename().string(); - if (filename.ends_with(".parquet")) { - result.push_back(entry.path()); - } - } - } - return result; -} - -std::string get_table_path( - std::string const& input_directory, std::string const& table_name -) { - auto dir = input_directory.empty() ? "." : input_directory; - auto file_path = dir + "/" + table_name + ".parquet"; - - if (std::filesystem::exists(file_path)) { - return file_path; - } - - return dir + "/" + table_name + "/"; -} - -std::map get_column_types( - std::string const& input_directory, std::string const& table_name -) { - auto files = list_parquet_files(get_table_path(input_directory, table_name)); - RAPIDSMPF_EXPECTS(!files.empty(), "No parquet files found for table " + table_name); - - auto metadata = cudf::io::read_parquet_metadata(cudf::io::source_info(files[0])); - auto const& root = metadata.schema().root(); - - std::map result; - for (std::size_t i = 0; i < root.num_children(); ++i) { - auto const& column = root.child(safe_cast(i)); - result.emplace(column.name(), column.cudf_type()); - } - return result; -} - -} // namespace detail - -streaming::Actor sink_channel( - std::shared_ptr ctx, std::shared_ptr ch -) { - co_await ctx->executor()->schedule(); - co_await ch->shutdown(); -} - -streaming::Actor consume_channel( - std::shared_ptr ctx, std::shared_ptr ch_in -) { - streaming::ShutdownAtExit c{ch_in}; - co_await ctx->executor()->schedule(); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - if (msg.holds()) { - auto chunk = co_await msg.release() - .make_available(ctx); - ctx->logger()->print( - "Consumed chunk with ", - chunk.table_view().num_rows(), - " rows and ", - chunk.table_view().num_columns(), - " columns" - ); - } - } -} - -std::pair, std::shared_ptr> -create_context( - ProgramOptions& arguments, cuda::mr::any_resource mr -) { - rmm::mr::set_current_device_resource(mr); - std::unordered_map memory_limits{}; - if (arguments.spill_device_limit.has_value()) { - auto limit_size = rmm::align_down( - (rmm::available_device_memory().second - * static_cast(arguments.spill_device_limit.value() * 100) - / 100), - rmm::CUDA_ALLOCATION_ALIGNMENT - ); - - memory_limits[MemoryType::DEVICE] = static_cast(limit_size); - } - auto statistics = Statistics::create(); - - RAPIDSMPF_EXPECTS( - arguments.no_pinned_host_memory || is_pinned_memory_resources_supported(), - "Pinned host memory is not supported on this system. " - "CUDA " RAPIDSMPF_PINNED_MEM_RES_MIN_CUDA_VERSION_STR - " is one of the requirements, but additional platform or driver constraints may " - "apply. If needed, use `--no-pinned-host-memory` to disable pinned host memory, " - "noting that this may significantly degrade spilling performance.", - std::invalid_argument - ); - - auto br = BufferResource::create( - std::move(mr), - arguments.no_pinned_host_memory ? PinnedMemoryResource::Disabled - : PinnedMemoryResource::make_if_available(), - std::move(memory_limits), - arguments.periodic_spill, - std::make_shared( - arguments.num_streams, rmm::cuda_stream::flags::non_blocking - ), - statistics - ); - auto environment = config::get_environment_variables(); - environment["NUM_STREAMING_THREADS"] = - std::to_string(arguments.num_streaming_threads); - auto options = config::Options(environment); - auto progress_thread = std::make_shared(statistics); - std::shared_ptr comm; - switch (arguments.comm_type) { - case CommType::MPI: - RAPIDSMPF_EXPECTS( - !bootstrap::is_running_with_rrun(), "Can't use MPI communicator with rrun" - ); - mpi::init(nullptr, nullptr); - - comm = std::make_shared(MPI_COMM_WORLD, options, progress_thread); - break; - case CommType::SINGLE: - comm = std::make_shared(options, progress_thread); - break; - case CommType::UCXX: - if (bootstrap::is_running_with_rrun()) { - comm = bootstrap::create_ucxx_comm( - progress_thread, bootstrap::BackendType::AUTO, options - ); - } else { - mpi::init(nullptr, nullptr); - comm = ucxx::init_using_mpi(MPI_COMM_WORLD, options, progress_thread); - } - break; - default: - RAPIDSMPF_FAIL("Unknown communicator type"); - } - auto ctx = std::make_shared(options, comm->logger(), br); - if (comm->rank() == 0) { - comm->logger()->print( - "Execution context on ", - comm->nranks(), - " ranks has ", - ctx->executor()->num_streaming_threads(), - " threads" - ); - } - return {ctx, comm}; -} - -ProgramOptions parse_arguments(int argc, char** argv) { - ProgramOptions options; - - static constexpr std::array(CommType::MAX)> - comm_names{"single", "mpi", "ucxx"}; - - auto print_usage = [&argv, &options]() { - std::cerr - << "Usage: " << argv[0] << " [options]\n" - << "Options:\n" - << " --num-streaming-threads Number of streaming threads (default: " - << options.num_streaming_threads << ")\n" - << " --num-iterations Number of iterations (default: " - << options.num_iterations << ")\n" - << " --num-streams Number of streams in stream pool " - "(default: " - << options.num_streams << ")\n" - << " --num-rows-per-chunk Number of rows per chunk (default: " - << options.num_rows_per_chunk << ")\n" - << " --spill-device-limit Fractional spill device limit as " - "fraction " - "of total device memory (default: " - << (options.spill_device_limit.has_value() - ? std::to_string(options.spill_device_limit.value()) - : "None") - << ")\n" - << " --no-pinned-host-memory Disable pinned host memory (default: " - << (options.no_pinned_host_memory ? "true" : "false") << ")\n" - << " --periodic-spill Duration in milliseconds between periodic " - "spilling checks (default: " - << (options.periodic_spill.has_value() - ? std::to_string(options.periodic_spill.value().count()) - : "None") - << ")\n" - << " --comm-type Communicator type: single, mpi, ucxx " - "(default: " - << comm_names[static_cast(options.comm_type)] << ")\n" - << " --use-shuffle-join Use shuffle join (default: " - << (options.use_shuffle_join ? "true" : "false") << ")\n" - << " --output-file Output file path (required)\n" - << " --input-directory Input directory path (required)\n" - << " --help Show this help message\n"; - }; - - // NOLINTBEGIN(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays,modernize-use-designated-initializers) - static struct option long_options[] = { - {"num-streaming-threads", required_argument, nullptr, 1}, - {"num-rows-per-chunk", required_argument, nullptr, 2}, - {"use-shuffle-join", no_argument, nullptr, 3}, - {"output-file", required_argument, nullptr, 4}, - {"input-directory", required_argument, nullptr, 5}, - {"help", no_argument, nullptr, 6}, - {"spill-device-limit", required_argument, nullptr, 7}, - {"num-iterations", required_argument, nullptr, 8}, - {"num-streams", required_argument, nullptr, 9}, - {"comm-type", required_argument, nullptr, 10}, - {"periodic-spill", required_argument, nullptr, 11}, - {"no-pinned-host-memory", no_argument, nullptr, 12}, - {nullptr, 0, nullptr, 0} - }; - // NOLINTEND(modernize-avoid-c-arrays,cppcoreguidelines-avoid-c-arrays,modernize-use-designated-initializers) - - int opt; - int option_index = 0; - - bool saw_output_file = false; - bool saw_input_directory = false; - - while ((opt = getopt_long(argc, argv, "", long_options, &option_index)) != -1) { - switch (opt) { - case 1: - { - char* endptr; - long val = std::strtol(optarg, &endptr, 10); - if (*endptr != '\0' || val <= 0) { - std::cerr << "Error: Invalid value for --num-streaming-threads: " - << optarg << "\n\n"; - print_usage(); - std::exit(1); - } - options.num_streaming_threads = static_cast(val); - break; - } - case 2: - { - char* endptr; - long val = std::strtol(optarg, &endptr, 10); - if (*endptr != '\0' || val <= 0) { - std::cerr << "Error: Invalid value for --num-rows-per-chunk: " - << optarg << "\n\n"; - print_usage(); - std::exit(1); - } - options.num_rows_per_chunk = static_cast(val); - break; - } - case 3: - options.use_shuffle_join = true; - break; - case 4: - options.output_file = optarg; - saw_output_file = true; - break; - case 5: - options.input_directory = optarg; - saw_input_directory = true; - break; - case 6: - print_usage(); - std::exit(0); - case 7: - { - char* endptr; - double val = std::strtod(optarg, &endptr); - if (*endptr != '\0' || val < 0.0 || val > 1.0) { - std::cerr << "Error: Invalid value for --spill-device-limit: " - << optarg << " (must be between 0.0 and 1.0)\n\n"; - print_usage(); - std::exit(1); - } - options.spill_device_limit = val; - break; - } - case 8: - { - char* endptr; - long val = std::strtol(optarg, &endptr, 10); - if (*endptr != '\0' || val <= 0) { - std::cerr << "Error: Invalid value for --num-iterations: " << optarg - << "\n\n"; - print_usage(); - std::exit(1); - } - options.num_iterations = static_cast(val); - break; - } - case 9: - { - char* endptr; - long val = std::strtol(optarg, &endptr, 10); - if (*endptr != '\0' || val <= 0) { - std::cerr << "Error: Invalid value for --num-streams: " << optarg - << "\n\n"; - print_usage(); - std::exit(1); - } - options.num_streams = static_cast(val); - break; - } - case 10: - { - std::string comm_type = optarg; - if (comm_type == "mpi") { - options.comm_type = CommType::MPI; - } else if (comm_type == "single") { - options.comm_type = CommType::SINGLE; - } else if (comm_type == "ucxx") { - options.comm_type = CommType::UCXX; - } else { - std::cerr << "Error: Invalid value for --comm-type: " << optarg - << " (must be one of " << comm_names[0]; - for (std::size_t i = 1; i < comm_names.size(); ++i) { - std::cerr << ", " << comm_names[i]; - } - std::cerr << ")\n\n"; - print_usage(); - std::exit(1); - } - break; - } - case 11: - { - char* endptr; - long val = std::strtol(optarg, &endptr, 10); - if (*endptr != '\0' || val <= 0) { - std::cerr << "Error: Invalid value for --periodic-spill: " << optarg - << "\n\n"; - print_usage(); - std::exit(1); - } - options.periodic_spill = std::chrono::milliseconds(val); - break; - } - case 12: - options.no_pinned_host_memory = true; - break; - case '?': - if (optopt == 0 && optind > 1) { - std::cerr << "Error: Unknown option '" << argv[optind - 1] << "'\n\n"; - } - print_usage(); - std::exit(1); - default: - print_usage(); - std::exit(1); - } - } - - // Check if required options were provided - if (!saw_output_file || !saw_input_directory) { - if (!saw_output_file) { - std::cerr << "Error: --output-file is required\n"; - } - if (!saw_input_directory) { - std::cerr << "Error: --input-directory is required\n"; - } - std::cerr << std::endl; - print_usage(); - std::exit(1); - } - - return options; -} -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/streaming/ndsh/utils.hpp b/cpp/benchmarks/streaming/ndsh/utils.hpp deleted file mode 100644 index bd63c45b8..000000000 --- a/cpp/benchmarks/streaming/ndsh/utils.hpp +++ /dev/null @@ -1,323 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include -#include -#include -#include -#include -#include -#include - -#include - -#include - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -namespace rapidsmpf::ndsh { -namespace detail { - -/** - * @brief List all parquet files in a given path. - * - * @param root_path The path to look in. - * - * @return If `root_path` names a regular file that ends with `.parquet` then a singleton - * vector of just that file. If `root_path` is a directory, then a vector containing all - * regular files in that directory whose name ends with `.parquet`, in the order they are - * listed. - * - * @throws std::runtime_error if the `root_path` doesn't name a regular file or a - * directory. Or if it does name a regular file, but that file doesn't end in `.parquet`. - */ -[[nodiscard]] std::vector list_parquet_files(std::string const root_path); - -/** - * @brief Get the path to a given table - * - * @param input_directory Input directory - * @param table_name Name of table to find. - * - * @return Path to given table. - */ -[[nodiscard]] std::string get_table_path( - std::string const& input_directory, std::string const& table_name -); - -/** - * @brief Get cudf data types for all columns from parquet metadata. - * - * Reads parquet metadata to determine the cudf data type for each column. - * The data types are inferred from the first file found for the given table. - * - * @param input_directory Directory containing input parquet files - * @param table_name Name of the table (e.g., "lineitem") - * @return Map from column name to cudf data type - */ -[[nodiscard]] std::map get_column_types( - std::string const& input_directory, std::string const& table_name -); - -} // namespace detail - -/** - * @brief Create a date comparison filter expression. - * - * Creates a filter that compares a date column against a literal date value. - * The operation will be equivalent to - * " DATE '--'". - * - * @tparam timestamp_type The timestamp type to use for the filter scalar - * (e.g., cudf::timestamp_D or cudf::timestamp_ms) - * @param stream CUDA stream to use - * @param date The date to compare against - * @param column_name The name of the column to compare - * @param op The comparison operator (e.g., LESS, LESS_EQUAL, GREATER) - * @return Filter expression with proper lifetime management - */ -template -std::unique_ptr make_date_filter( - rmm::cuda_stream_view stream, - cuda::std::chrono::year_month_day date, - std::string const& column_name, - cudf::ast::ast_operator op -) { - auto owner = new std::vector; - auto sys_days = cuda::std::chrono::sys_days(date); - owner->push_back( - std::make_shared>( - sys_days.time_since_epoch(), true, stream - ) - ); - owner->push_back( - std::make_shared( - *std::any_cast>>( - owner->at(0) - ) - ) - ); - owner->push_back(std::make_shared(column_name)); - owner->push_back( - std::make_shared( - op, - *std::any_cast>( - owner->at(2) - ), - *std::any_cast>(owner->at(1)) - ) - ); - return std::make_unique( - stream, - *std::any_cast>(owner->back()), - OwningWrapper(static_cast(owner), [](void* p) { - delete static_cast*>(p); - }) - ); -} - -/** - * @brief Create a date range filter expression. - * - * Creates a filter that checks if a date column falls within a half-open range. - * The operation will be equivalent to - * " >= DATE '' AND < DATE ''". - * - * @tparam timestamp_type The timestamp type to use for the filter scalars - * (e.g., cudf::timestamp_D or cudf::timestamp_ms) - * @param stream CUDA stream to use - * @param start_date The start date (inclusive) of the range - * @param end_date The end date (exclusive) of the range - * @param column_name The name of the column to compare - * @return Filter expression with proper lifetime management - */ -template -std::unique_ptr make_date_range_filter( - rmm::cuda_stream_view stream, - cuda::std::chrono::year_month_day start_date, - cuda::std::chrono::year_month_day end_date, - std::string const& column_name -) { - auto owner = new std::vector; - - // 0: column_reference - owner->push_back(std::make_shared(column_name)); - - // 1, 2: Scalars for start and end dates - owner->push_back( - std::make_shared>( - cuda::std::chrono::sys_days(start_date).time_since_epoch(), true, stream - ) - ); - owner->push_back( - std::make_shared>( - cuda::std::chrono::sys_days(end_date).time_since_epoch(), true, stream - ) - ); - - // 3, 4: Literals for start and end dates - owner->push_back( - std::make_shared( - *std::any_cast>>( - owner->at(1) - ) - ) - ); - owner->push_back( - std::make_shared( - *std::any_cast>>( - owner->at(2) - ) - ) - ); - - // 5: (GE, column, literal) - owner->push_back( - std::make_shared( - cudf::ast::ast_operator::GREATER_EQUAL, - *std::any_cast>( - owner->at(0) - ), - *std::any_cast>(owner->at(3)) - ) - ); - - // 6: (LT, column, literal) - owner->push_back( - std::make_shared( - cudf::ast::ast_operator::LESS, - *std::any_cast>( - owner->at(0) - ), - *std::any_cast>(owner->at(4)) - ) - ); - - // 7: (AND, GE, LT) - owner->push_back( - std::make_shared( - cudf::ast::ast_operator::LOGICAL_AND, - *std::any_cast>(owner->at(5)), - *std::any_cast>(owner->at(6)) - ) - ); - - return std::make_unique( - stream, - *std::any_cast>(owner->back()), - OwningWrapper(static_cast(owner), [](void* p) { - delete static_cast*>(p); - }) - ); -} - -/** - * @brief Sink messages into a channel and discard them. - * - * @param ctx Streaming context - * @param ch Channel to discard messages from. - * - * @return Coroutine representing the shutdown and discard of the channel. - */ -[[nodiscard]] streaming::Actor sink_channel( - std::shared_ptr ctx, std::shared_ptr ch -); - -/** - * @brief Consume messages from a channel and discard them. - * - * @param ctx Streaming context - * @param ch Channel to consume messages from. - * - * @note If the channel contains `TableChunk`s, moves them to device and prints small - * amount of detail about them (row and column count). - * - * @return Coroutine representing consuming and discarding messages in channel. - */ -[[nodiscard]] streaming::Actor consume_channel( - std::shared_ptr ctx, std::shared_ptr ch_in -); - -///< @brief Communicator type to use -enum class CommType : std::uint8_t { - SINGLE, ///< Single process communicator - MPI, ///< MPI backed communicator - UCXX, ///< UCXX backed communicator - MAX, ///< Max value -}; - -///< @brief Configuration options for the query -struct ProgramOptions { - int num_streaming_threads{1}; ///< Number of streaming threads to use - int num_iterations{2}; ///< Number of iterations of query to run - int num_streams{16}; ///< Number of streams in stream pool - CommType comm_type{CommType::UCXX}; ///< Type of communicator to create - std::optional - periodic_spill; ///< Duration between background periodic spilling checks - cudf::size_type num_rows_per_chunk{ - 100'000'000 - }; ///< Number of rows to produce per chunk read - std::optional spill_device_limit{ - std::nullopt - }; ///< Optional fractional spill limit - bool no_pinned_host_memory{false}; ///< Disable pinned host memory? - bool use_shuffle_join = false; ///< Use shuffle join for "big" joins? - std::string output_file; ///< File to write output to - std::string input_directory; ///< Directory containing input files. -}; - -/** - * @brief Parse commandline arguments - * - * @param argc Number of arguments - * @param argv Arguments - * - * @return `ProgramOptions` struct with parsed arguments. - */ -ProgramOptions parse_arguments(int argc, char** argv); - -/** - * @brief Create a streaming execution context and communicator for a query. - * - * @param arguments Arguments to configure the context - * @param mr The device memory resource to use for all allocations. - * - * @return Pair of shared pointer to new streaming context and communicator. - */ -std::pair, std::shared_ptr> -create_context( - ProgramOptions& arguments, cuda::mr::any_resource mr -); - -/** - * @brief Finalize MPI when going out of scope. - */ -struct FinalizeMPI { - ~FinalizeMPI() noexcept { - if (rapidsmpf::mpi::is_initialized()) { - int flag; - RAPIDSMPF_MPI(MPI_Finalized(&flag)); - if (!flag) { - RAPIDSMPF_MPI(MPI_Finalize()); - } - } - } -}; -} // namespace rapidsmpf::ndsh diff --git a/cpp/benchmarks/utils/random_data.cu b/cpp/benchmarks/utils/random_data.cu deleted file mode 100644 index 021c384ad..000000000 --- a/cpp/benchmarks/utils/random_data.cu +++ /dev/null @@ -1,112 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include - -#include -#include - -#include -#include -#include -#include -#include - -#include -#include - -#include "random_data.hpp" - -rmm::device_uvector random_device_vector( - std::size_t nelem, - std::int32_t min_val, - std::int32_t max_val, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -) { - // Fill vector with random data. - using index_t = std::int64_t; - auto const end_index = rapidsmpf::safe_cast(nelem); - rmm::device_uvector vec(nelem, stream, mr); - thrust::counting_iterator const begin(0); - thrust::counting_iterator const end(end_index); - thrust::transform( - rmm::exec_policy(stream), - begin, - end, - vec.begin(), - [min_val, max_val] __device__(index_t index) { - thrust::default_random_engine engine( - static_cast(index) - ); - thrust::uniform_int_distribution dist(min_val, max_val); - return dist(engine); - } - ); - return vec; -} - -std::unique_ptr random_column( - cudf::size_type nrows, - std::int32_t min_val, - std::int32_t max_val, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -) { - auto vec = random_device_vector( - rapidsmpf::safe_cast(nrows), min_val, max_val, stream, mr - ); - return std::make_unique( - std::move(vec), rmm::device_buffer{0, stream, mr}, 0 - ); -} - -cudf::table random_table( - cudf::size_type ncolumns, - cudf::size_type nrows, - std::int32_t min_val, - std::int32_t max_val, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -) { - std::vector> cols; - for (auto i = 0; i < ncolumns; ++i) { - cols.push_back(random_column(nrows, min_val, max_val, stream, mr)); - } - return cudf::table(std::move(cols)); -} - -void random_fill(rapidsmpf::Buffer& buffer, rmm::device_async_resource_ref mr) { - switch (buffer.mem_type()) { - case rapidsmpf::MemoryType::DEVICE: - { - auto const num_elements = std::max( - std::size_t{1}, - buffer.size / sizeof(random_data_t) - + (buffer.size % sizeof(random_data_t) != 0) - ); - auto vec = random_device_vector( - num_elements, - std::numeric_limits::min(), - std::numeric_limits::max(), - buffer.stream(), - mr - ); - buffer.write_access([&](std::byte* buffer_data, - rmm::cuda_stream_view stream) { - RAPIDSMPF_CUDA_TRY( - rapidsmpf::cuda_memcpy_async( - buffer_data, vec.data(), buffer.size, stream - ) - ); - }); - break; - } - default: - RAPIDSMPF_FAIL("unsupported memory type", std::invalid_argument); - } -} diff --git a/cpp/benchmarks/utils/random_data.hpp b/cpp/benchmarks/utils/random_data.hpp deleted file mode 100644 index 983f086a1..000000000 --- a/cpp/benchmarks/utils/random_data.hpp +++ /dev/null @@ -1,117 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include -#include - -#include -#include -#include - -#include -#include - - -/** - * @brief The type of random data to generate. - */ -using random_data_t = std::int32_t; - -/** - * @brief The lower bound of the size of a random table. - * - * @param ncolumns The number of columns in the table. - * @param nrows The number of rows in the table. - */ -std::size_t constexpr random_table_size_lower_bound( - cudf::size_type ncolumns, cudf::size_type nrows -) { - return rapidsmpf::safe_cast(ncolumns) - * rapidsmpf::safe_cast(nrows) * sizeof(random_data_t); -} - -/** - * @brief Generates a random numeric device vector (std::int32_t). - * - * Creates a device vector with random integer values uniformly distributed in - * the range `[min_val, max_val]`. - * - * @param nelem Number of elements in the generated vector. - * @param min_val Minimum value (inclusive) for the random data. - * @param max_val Maximum value (inclusive) for the random data. - * @param stream CUDA stream to use for memory and kernel operations. - * @param mr Device memory resource for allocating the device vector. - * @return A unique pointer to the generated device vector. - * - * @note The function uses the specified CUDA stream for asynchronous operations. - */ -rmm::device_uvector random_device_vector( - std::size_t nelem, - std::int32_t min_val, - std::int32_t max_val, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -); - -/** - * @brief Generates a random numeric column (std::int32_t). - * - * Creates a cuDF column with random integer values uniformly distributed in the range - * `[min_val, max_val]`. - * - * @param nrows Number of rows in the generated column. - * @param min_val Minimum value (inclusive) for the random data. - * @param max_val Maximum value (inclusive) for the random data. - * @param stream CUDA stream to use for memory and kernel operations. - * @param mr Device memory resource for allocating the column. - * @return A unique pointer to the generated cuDF column. - * - * @note The function uses the specified CUDA stream for asynchronous operations. - */ -std::unique_ptr random_column( - cudf::size_type nrows, - std::int32_t min_val, - std::int32_t max_val, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -); - -/** - * @brief Generates a random numeric table (std::int32_t). - * - * Creates a cuDF table consisting of multiple columns with random integer values, each - * uniformly distributed in the range `[min_val, max_val]`. - * - * @param ncolumns Number of columns in the generated table. - * @param nrows Number of rows in each column of the table. - * @param min_val Minimum value (inclusive) for the random data. - * @param max_val Maximum value (inclusive) for the random data. - * @param stream CUDA stream to use for memory and kernel operations. - * @param mr Device memory resource for allocating the table. - * @return A cuDF table containing the generated random columns. - * - * @note Each column in the table will have the same number of rows and data distribution. - */ -cudf::table random_table( - cudf::size_type ncolumns, - cudf::size_type nrows, - std::int32_t min_val, - std::int32_t max_val, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -); - -/** - * @brief Fill a rapidsmpf buffer with random data (std::int32_t). - * - * Using buffer's CUDA stream. - * - * @param buffer The buffer to fill. - * @param mr Device memory resource for allocating temporary random data. - * - * @throws std::invalid_argument if the memory type of `buffer` isn't supported. - */ -void random_fill(rapidsmpf::Buffer& buffer, rmm::device_async_resource_ref mr); diff --git a/cpp/compute-sanitizer-suppressions.xml b/cpp/compute-sanitizer-suppressions.xml index 5e5cd718f..632c74ffe 100644 --- a/cpp/compute-sanitizer-suppressions.xml +++ b/cpp/compute-sanitizer-suppressions.xml @@ -51,117 +51,4 @@ - - InitcheckApiError - Error - - Host API uninitialized memory access - 126881900793856 - 1664 - 126881900794656 - - - error - - .*/libcuda.so.* - - - .*/libcudart.so.* - - - .*/libcudart.so.12 - - - cudaMemcpyAsync - .*/libcudart.so.12 - - - rapidsmpf::buffer_copy const - .*/memory/buffer.cpp - .*/librapidsmpf.so - - - - - - - rapidsmpf::buffer_copy - .*/memory/buffer.cpp - /.*/librapidsmpf.so - - - rapidsmpf::BufferResource::move - .*/memory/buffer_resource.cpp - /.*/librapidsmpf.so - - - .*/tests/streaming/test_table_chunk.cpp - /.*/gtests/mpi_tests - - - - - - InitcheckApiError - Error - - Host API uninitialized memory access - - - error - - .*/libcuda.so.* - - - .*/libcudart.so.* - - - .*/libcudart.so.* - - - cudaMemcpyAsync - .*/libcudart.so.* - - - rapidsmpf::buffer_copy const - .*/memory/buffer.cpp - .*/librapidsmpf.so - - - .*rapidsmpf::buffer_copy - .*/c++/bits/invoke.h - .*/librapidsmpf.so - - - .*rapidsmpf::buffer_copy - .*/c++/bits/invoke.h - .*/librapidsmpf.so - - - .*rapidsmpf::buffer_copy - .*/c++/functional - .*/librapidsmpf.so - - - std::invoke_result<rapidsmpf::buffer_copy - */memory/buffer.hpp - .*/librapidsmpf.so - - - rapidsmpf::buffer_copy - .*/memory/buffer.cpp - .*/librapidsmpf.so - - - rapidsmpf::BufferResource::move - .*/memory/buffer_resource.cpp - .*/librapidsmpf.so - - - cudf_streaming::streaming::TableChunk::spill_to_host - .*/streaming/cudf/table_chunk.cpp - .*/librapidsmpf.so - - - diff --git a/cpp/examples/CMakeLists.txt b/cpp/examples/CMakeLists.txt index 8f107fe32..f8474e660 100644 --- a/cpp/examples/CMakeLists.txt +++ b/cpp/examples/CMakeLists.txt @@ -5,34 +5,6 @@ # cmake-format: on # ================================================================================= -if(BUILD_CUDF_TESTS) - add_executable(example_shuffle "example_shuffle.cpp" "../benchmarks/utils/random_data.cu") - set_target_properties( - example_shuffle - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/examples" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - CXX_EXTENSIONS ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON - ) - target_compile_options( - example_shuffle PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" - ) - target_link_libraries( - example_shuffle - PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx cudf_streaming::cudf_streaming - $ $ maybe_asan - ) - install( - TARGETS example_shuffle - COMPONENT testing - DESTINATION bin/examples/librapidsmpf - EXCLUDE_FROM_ALL - ) -endif() - if(RAPIDSMPF_HAVE_CUPTI) add_executable(example_cupti_monitor "example_cupti_monitor.cpp") set_target_properties( diff --git a/cpp/examples/example_shuffle.cpp b/cpp/examples/example_shuffle.cpp deleted file mode 100644 index 98bcc6868..000000000 --- a/cpp/examples/example_shuffle.cpp +++ /dev/null @@ -1,139 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include -#include - -#include - -#include -#include -#include -#include -#include -#include - -#include "../benchmarks/utils/random_data.hpp" - -// An example of how to use the shuffler. -int main(int argc, char** argv) { - // In this example we use the MPI backed. For convenience, rapidsmpf provides an - // optional MPI-init function that initialize MPI with thread support. - rapidsmpf::mpi::init(&argc, &argv); - - // Initialize configuration options from environment variables. - rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()}; - - // Create a statistics instance for the shuffler that tracks useful information. - auto stats = rapidsmpf::Statistics::create(); - - // The communicator has a progress thread where the shuffler event loop executes. A - // single progress thread may be used by multiple shufflers simultaneously. - auto progress_thread = std::make_shared(stats); - - // Now we have to create a Communicator, which we will use throughout the - // example. Multiple concurrent shuffles are possible on the same communicator by - // providing differentiating "OpID" arguments. - std::shared_ptr comm = - std::make_shared(MPI_COMM_WORLD, options, progress_thread); - - - // The Communicator provides a logger. - auto& log = comm->logger(); - - // We will use the same stream, memory, and buffer resource throughout the example. - rmm::cuda_stream_view stream = cudf::get_default_stream(); - rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); - auto br = rapidsmpf::BufferResource::create(mr); - - // As input data, we use a helper function from the benchmark suite. It creates a - // random cudf table with 2 columns and 100 rows. In this example, each MPI rank - // creates its own local input and we only have one input per rank but each rank - // could take any number of inputs. - cudf::table local_input = random_table(2, 100, 0, 10, stream, mr); - - // The total number of inputs equals the number of ranks, in this case. - auto const total_num_partitions = - static_cast(comm->nranks()); - - // We create a new shuffler instance, which represents a single shuffle. It takes - // a Communicator, the total number of partitions, and a "owner function", which - // map partitions to their destination ranks. All ranks must use the same owner - // function, in this example we use the included round-robin owner function. - rapidsmpf::shuffler::Shuffler shuffler( - comm, - 0, // op_id - total_num_partitions, - br.get(), - rapidsmpf::shuffler::Shuffler::round_robin // partition owner - ); - - // It is our own responsibility to partition and pack (serialize) the input for - // the shuffle. The shuffler only handles raw host and device buffers. However, it - // does provide a convenience function that hash partitions a cudf table and packs - // each partition. The result is a mapping of `PartID`, globally unique partition - // identifiers, to their packed partitions. - std::unordered_map packed_inputs = - cudf_streaming::integrations::partition_and_pack( - local_input, - {0}, // columns_to_hash - static_cast(total_num_partitions), - cudf::hash_id::HASH_MURMUR3, - cudf::DEFAULT_HASH_SEED, - stream, - br.get() - ); - - // Now, we can insert the packed partitions into the shuffler. This operation is - // non-blocking and we can continue inserting new input partitions. E.g., a pipeline - // could read, hash-partition, pack, and insert, one parquet-file at a time while the - // distributed shuffle is being processed underneath. - shuffler.insert(std::move(packed_inputs)); - - // When we are finished inserting data, we tell the shuffler. This sends one control - // message per target rank, informing each that this rank has finished inserting data. - shuffler.insert_finished(); - - // Vector to hold the local results of the shuffle operation. - std::vector> local_outputs; - - // Wait for all partitions to finish. - shuffler.wait(); - - // Process the shuffle results for each partition. - for (auto finished_partition : shuffler.local_partitions()) { - // Extract the finished partition's data from the Shuffler. - auto packed_chunks = shuffler.extract(finished_partition); - - // Unpack (deserialize) and concatenate the chunks into a single table using a - // convenience function. - local_outputs.push_back( - cudf_streaming::integrations::unpack_and_concat( - rapidsmpf::unspill_partitions( - std::move(packed_chunks), br.get(), rapidsmpf::AllowOverbooking::YES - ), - stream, - br.get() - ) - ); - } - // At this point, `local_outputs` contains the local result of the shuffle. - // Let's log the result. - log->print( - "Finished shuffle with ", local_outputs.size(), " local output partitions" - ); - - // Log the statistics report. - log->print(stats->report()); - - // Shutdown the Shuffler explicitly or let it go out of scope for cleanup. - shuffler.shutdown(); - - // Finalize the execution, `RAPIDSMPF_MPI` is a convenience macro that - // checks for MPI errors. - RAPIDSMPF_MPI(MPI_Finalize()); -} diff --git a/cpp/include/rapidsmpf/bootstrap/slurm_backend.hpp b/cpp/include/rapidsmpf/bootstrap/slurm_backend.hpp index ebeec10c6..38c9a9f28 100644 --- a/cpp/include/rapidsmpf/bootstrap/slurm_backend.hpp +++ b/cpp/include/rapidsmpf/bootstrap/slurm_backend.hpp @@ -37,7 +37,7 @@ namespace rapidsmpf::bootstrap::detail { * --cpus-per-task=36 \ * --gpus-per-task=1 \ * --gres=gpu:4 \ - * rrun ./benchmarks/bench_shuffle -C ucxx + * rrun ./benchmarks/bench_comm -C ucxx * * # Hybrid mode: one task per node, 4 GPUs per task, two nodes. * srun \ @@ -47,7 +47,7 @@ namespace rapidsmpf::bootstrap::detail { * --cpus-per-task=144 \ * --gpus-per-task=4 \ * --gres=gpu:4 \ - * rrun -n 4 ./benchmarks/bench_shuffle -C ucxx + * rrun -n 4 ./benchmarks/bench_comm -C ucxx * ``` */ class SlurmBackend : public Backend { diff --git a/cpp/include/rapidsmpf/memory/buffer_resource.hpp b/cpp/include/rapidsmpf/memory/buffer_resource.hpp index d477f903a..6307eeb94 100644 --- a/cpp/include/rapidsmpf/memory/buffer_resource.hpp +++ b/cpp/include/rapidsmpf/memory/buffer_resource.hpp @@ -182,26 +182,13 @@ class BufferResource : public std::enable_shared_from_this { * mr.allocate(...); // safe * @endcode * - * In the common case, no explicit promotion is needed because RMM and cuDF containers - * that store a memory resource do this internally: + * In the common case, no explicit promotion is needed because RMM containers that + * store a memory resource do this internally: * @code * auto br = BufferResource::create(...); * rmm::device_buffer buf{1024, stream, br->device_mr()}; * br.reset(); // safe: `buf` keeps the BufferResource alive internally * @endcode - * - * Returned objects from cuDF APIs typically behave the same way: - * @code - * auto br = BufferResource::create(...); - * auto col = cudf::make_numeric_column( - * cudf::data_type{cudf::type_id::INT32}, - * 1000, - * cudf::mask_state::UNALLOCATED, - * stream, - * br->device_mr() - * ); - * br.reset(); // safe: `col` keeps the BufferResource alive internally - * @endcode */ [[nodiscard]] rmm::device_async_resource_ref device_mr() noexcept; diff --git a/cpp/include/rapidsmpf/owning_wrapper.hpp b/cpp/include/rapidsmpf/owning_wrapper.hpp index f7560b06e..dca037a66 100644 --- a/cpp/include/rapidsmpf/owning_wrapper.hpp +++ b/cpp/include/rapidsmpf/owning_wrapper.hpp @@ -1,5 +1,5 @@ /** - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -16,14 +16,12 @@ namespace rapidsmpf { * 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. + * For example, a C++ message payload may hold non-owning views into a Python-owned + * object. 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 keeping that owner 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 diff --git a/cpp/include/rapidsmpf/shuffler/finish_counter.hpp b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp index 4468b658f..6414d4bba 100644 --- a/cpp/include/rapidsmpf/shuffler/finish_counter.hpp +++ b/cpp/include/rapidsmpf/shuffler/finish_counter.hpp @@ -14,7 +14,7 @@ * @namespace rapidsmpf::shuffler * @brief Shuffler interfaces. * - * A shuffle service for cuDF tables. Use `Shuffler` to perform a single shuffle. + * A shuffle service for partitioned payloads. Use `Shuffler` to perform a single shuffle. */ namespace rapidsmpf::shuffler { diff --git a/cpp/scripts/ndsh.py b/cpp/scripts/ndsh.py deleted file mode 100755 index a268334af..000000000 --- a/cpp/scripts/ndsh.py +++ /dev/null @@ -1,879 +0,0 @@ -#!/usr/bin/env python3 -# /// script -# requires-python = ">=3.11" -# dependencies = [ -# "duckdb", -# "numpy", -# "pyarrow", -# "tpchgen-cli", -# ] -# /// - -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -""" -Validation script for NDSH benchmarks. - -This script validates the correctness of NDSH benchmark outputs by: -1. Running SQL queries via DuckDB to generate expected results -2. Running the C++ benchmark binaries -3. Comparing the benchmark output against the DuckDB result - -Usage: - # Run benchmarks and generate expected results - python validate_ndsh.py run \\ - --benchmark-dir /path/to/build/benchmarks/ndsh \\ - --sql-dir /path/to/sql/queries \\ - --input-dir /raid/rapidsmpf/data/tpch/scale-1.0 \\ - --output-dir /tmp/validation - - # Validate results against expected - python validate_ndsh.py validate \\ - --results-path /tmp/validation/output \\ - --expected-path /tmp/validation/expected -""" - -from __future__ import annotations - -import argparse -import re -import subprocess -import hashlib -import sys -from pathlib import Path - -import duckdb -import numpy as np -import pyarrow as pa -import pyarrow.parquet as pq - -TPCH_TABLES = [ - "customer", - "lineitem", - "nation", - "orders", - "part", - "partsupp", - "region", - "supplier", -] - -# simple precaution to ensure that the SQL hasn't changed -# from what we expect. -QUERY_HASHES = { - "q01": "cccf0c3d9302ee4b56a2bac2f2aa05ab", - "q03": "1fdb8b3d8044f7d72c7bb021d025ea70", - "q04": "ed586bdb2b1495d3b1b19d1217dd6750", - "q09": "de7c61303a841c512857ee50412b54b9", - "q17": "7be3e180995f841ece86ffb2de9cf2b0", - "q18": "38b13cbbaeea09c224cc53b532507f76", - "q21": "61bcc0a1239c0feefc54b683027df014", -} - - -def discover_benchmarks( - benchmark_dir: Path, sql_dir: Path -) -> list[tuple[str, Path, Path]]: - """ - Discover benchmark binaries and their corresponding SQL files. - - Parameters - ---------- - benchmark_dir - Directory containing benchmark binaries (q04, q09, etc.) - sql_dir - Directory containing SQL query files (q04.sql, q09.sql, etc.) - - Returns - ------- - List of tuples with the following elements: - - - query_name - - binary_path - - sql_path - - These are sorted by query_name. - """ - benchmarks = [] - pattern = re.compile(r"^q(\d+)$") - - for binary in benchmark_dir.iterdir(): - if not binary.is_file(): - continue - match = pattern.match(binary.name) - if not match: - continue - - query_name = binary.name - sql_path = sql_dir / f"{query_name}.sql" - - if sql_path.exists(): - benchmarks.append((query_name, binary, sql_path)) - else: - print(f"Warning: No SQL file found for {query_name} at {sql_path}") - - return sorted(benchmarks) - - -def discover_parquet_files(directory: Path) -> dict[str, Path]: - """ - Discover parquet files matching the qDD.parquet pattern. - - Parameters - ---------- - directory - Directory containing parquet files (q03.parquet, q09.parquet, etc.) - - Returns - ------- - Dictionary mapping query_name (e.g., 'q03') to file path. - """ - pattern = re.compile(r"^q_?(\d+)\.parquet$") - files = {} - - for file in directory.iterdir(): - if not file.is_file(): - continue - match = pattern.match(file.name) - if match: - query_name = f"q{match.group(1)}" - files[query_name] = file - - return files - - -def generate_expected(sql_path: Path, input_dir: Path, output_path: Path) -> None: - """ - Generate expected results by running a SQL query via DuckDB. - - Parameters - ---------- - sql_path - Path to the SQL query file - input_dir: - Directory containing TPC-H parquet files - output_path - Path to write the expected parquet result - """ - con = duckdb.connect() - - # Register TPC-H tables as views from parquet files - for table in TPCH_TABLES: - # Try both single file and directory patterns - single_file = input_dir / f"{table}.parquet" - directory = input_dir / table - - if single_file.exists(): - parquet_path = single_file - elif directory.exists() and directory.is_dir(): - parquet_path = directory / "*.parquet" - else: - raise FileNotFoundError(f"Table {table} not found in {input_dir}") - - con.execute( - f"CREATE VIEW {table} AS SELECT * FROM read_parquet('{parquet_path}')" - ) - - # Read and execute the query - query = sql_path.read_text() - query_hash = hashlib.md5(query.encode()).hexdigest() - query_id = sql_path.stem - - if query_id not in QUERY_HASHES: - raise ValueError(f"Query {query_id} from file {sql_path} not found in QUERY_HASHES. Please update scripts/ndsh.py with the new hash.") - if query_hash != QUERY_HASHES[query_id]: - raise ValueError(f"Query {query_id} from file {sql_path} has changed. Please update scripts/ndsh.py with the new hash using 'md5sum {sql_path}'.") - - result = con.sql(query).arrow().read_all() - - # Write result to parquet - pq.write_table(result, output_path) - print(f" Generated expected: {output_path} ({result.num_rows} rows)") - - -def generate_data(input_dir: Path) -> None: - """ - Generate data for the benchmarks. - - This uses tpchgen-cli to generate the data and casts some columns - to the types expected by the benchmarks. - """ - print(f"Generating data for {input_dir}...") - subprocess.check_output( - [ - "tpchgen-cli", - "--scale-factor", - "1", - "--format", - "parquet", - "--output-dir", - str(input_dir), - ] - ) - - # Some of our queries are written expecting float (Double) - casts = { - ("customer", "c_nationkey"): pa.int32(), - ("customer", "c_acctbal"): pa.float64(), - ("lineitem", "l_linenumber"): pa.int64(), - ("lineitem", "l_quantity"): pa.float64(), - ("lineitem", "l_extendedprice"): pa.float64(), - ("lineitem", "l_discount"): pa.float64(), - ("lineitem", "l_tax"): pa.float64(), - ("lineitem", "l_shipdate"): pa.date32(), - ("lineitem", "l_commitdate"): pa.date32(), - ("lineitem", "l_receiptdate"): pa.date32(), - ("nation", "n_nationkey"): pa.int32(), - ("nation", "n_regionkey"): pa.int32(), - ("orders", "o_totalprice"): pa.float64(), - ("orders", "o_orderdate"): pa.date32(), - ("part", "p_retailprice"): pa.float64(), - ("partsupp", "ps_availqty"): pa.int64(), - ("partsupp", "ps_supplycost"): pa.float64(), - ("region", "r_regionkey"): pa.int32(), - ("supplier", "s_nationkey"): pa.int32(), - ("supplier", "s_acctbal"): pa.float64(), - } - - for table_name in TPCH_TABLES: - file = (input_dir / table_name).with_suffix(".parquet") - table = pq.read_table(file) - schema = table.schema - for i, field in enumerate(schema): - if cast := casts.get((table_name, field.name)): - schema = schema.set(i, field.with_type(cast)) - - pq.write_table(table.cast(schema), file) - - -def run_benchmark( - binary_path: Path, - input_dir: Path, - output_path: Path, - extra_args: list[str] | None = None, -) -> subprocess.CompletedProcess: - """ - Run a benchmark binary. - - Parameters - ---------- - binary_path - Path to the benchmark binary - input_dir - Directory containing TPC-H parquet files - output_path - Path for benchmark output - extra_args - Additional arguments to pass to the benchmark - - Returns - ------- - CompletedProcess result - """ - cmd = [ - "mpirun", - "-np", - "1", - "--allow-run-as-root", - str(binary_path), - "--input-directory", - str(input_dir), - "--output-file", - str(output_path), - ] - - if extra_args: - cmd.extend(extra_args) - - print(f" Running: {' '.join(cmd)}") - - return subprocess.run( - cmd, - check=False, - capture_output=True, - text=True, - ) - - -def _types_compatible( - o_type: pa.DataType, - e_type: pa.DataType, - *, - ignore_timezone: bool = False, - ignore_string_type: bool = False, - ignore_integer_sign: bool = False, - ignore_integer_size: bool = False, - ignore_decimal_int: bool = False, -) -> bool: - """ - Check if two Arrow types are compatible given the ignore flags. - - Returns True if the types should be considered equal. - """ - if o_type.equals(e_type): - return True - - # Ignore differences in timezone and precision for timestamps - if ( - ignore_timezone - and pa.types.is_timestamp(o_type) - and pa.types.is_timestamp(e_type) - ): - return True - - # Ignore large_string vs string differences - if ignore_string_type: - string_types = {pa.string(), pa.large_string()} - if o_type in string_types and e_type in string_types: - return True - - # Check integer compatibility - if pa.types.is_integer(o_type) and pa.types.is_integer(e_type): - o_signed = pa.types.is_signed_integer(o_type) - e_signed = pa.types.is_signed_integer(e_type) - o_width = o_type.bit_width - e_width = e_type.bit_width - - sign_matches = o_signed == e_signed or ignore_integer_sign - size_matches = o_width == e_width or ignore_integer_size - - if sign_matches and size_matches: - return True - - # Ignore decimal vs integer differences - if ignore_decimal_int: - o_is_numeric = pa.types.is_integer(o_type) or pa.types.is_decimal(o_type) - e_is_numeric = pa.types.is_integer(e_type) or pa.types.is_decimal(e_type) - if o_is_numeric and e_is_numeric: - return True - - return False - - -def compare_parquet( - output_path: Path, - expected_path: Path, - decimal: int = 2, - *, - ignore_timezone: bool = False, - ignore_string_type: bool = False, - ignore_integer_sign: bool = False, - ignore_integer_size: bool = False, - ignore_decimal_int: bool = False, -) -> tuple[bool, str | None]: - """ - Compare two parquet files for exact equality. - - Parameters - ---------- - output_path - Path to the benchmark output parquet - expected_path - Path to the expected parquet - decimal - Number of decimal places to compare for floating point values - ignore_timezone - Ignore differences in timezone and precision for timestamp types - ignore_string_type - Ignore differences between string and large_string types. - Note that the values will still be compared. - ignore_integer_sign - Ignore differences between signed and unsigned integer types - Note that the values will still be compared. - ignore_integer_size - Ignore differences in integer bit width (e.g., int32 vs int64) - Note that the values will still be compared. - ignore_decimal_int - Ignore differences between decimal and integer types - Note that the values will still be compared. - - Returns - ------- - Tuple of boolean indicating success and list of error messages. A non-empty list indicates failure. - """ - try: - output = pq.read_table(output_path) - expected = pq.read_table(expected_path) - except Exception as e: - return False, f"Failed to read parquet files: {e}" - - # Check the schema and data by validating... - # 1. names... - if output.schema.names != expected.schema.names: - return ( - False, - f"Schema name mismatch: {output.schema.names} != {expected.schema.names}", - ) - - # 2. types... - errors = [] - for name in output.schema.names: - o_field = output.schema.field(name) - e_field = expected.schema.field(name) - # We only care about the type, not the metadata or nullability - if not _types_compatible( - o_field.type, - e_field.type, - ignore_timezone=ignore_timezone, - ignore_string_type=ignore_string_type, - ignore_integer_sign=ignore_integer_sign, - ignore_integer_size=ignore_integer_size, - ignore_decimal_int=ignore_decimal_int, - ): - errors.append(f"\t{name}: {o_field.type} != {e_field.type}") - if errors: - return False, "\n".join(["Field type mismatch (output != expected)", *errors]) - - # 3. row count... - if output.num_rows != expected.num_rows: - return False, ( - f"Row count mismatch: output={output.num_rows}, expected={expected.num_rows}" - ) - - # 4. and values. - for name, out_col, expected_col in zip( - output.column_names, output.columns, expected.columns, strict=False - ): - if pa.types.is_floating(out_col.type): - # We don't promise exact equality - try: - np.testing.assert_array_almost_equal( - out_col.to_numpy(), expected_col.to_numpy(), decimal=decimal - ) - except AssertionError as e: - errors.append(f"{name} differs. {e}") - else: - try: - np.testing.assert_array_equal( - out_col.to_numpy(), expected_col.to_numpy() - ) - except AssertionError as e: - errors.append(f"{name} differs. {e}") - - if errors: - return False, "\n".join(errors) - - return True, None - - -def run_single_benchmark( - query_name: str, - binary_path: Path, - sql_path: Path, - input_dir: Path, - output_dir: Path, - expected_dir: Path, - extra_args: list[str] | None = None, - *, - reuse_expected: bool = False, - reuse_output: bool = False, -) -> bool: - """ - Run a single benchmark and generate expected results. - - Parameters - ---------- - query_name - Name of the query to run (e.g., 'q03') - binary_path - Path to the benchmark binary - sql_path - Path to the SQL query file - input_dir - Directory containing TPC-H parquet files - output_dir - Directory for benchmark output - expected_dir - Directory for expected results - extra_args - Additional arguments to pass to the benchmark - reuse_expected - Skip generating expected results if the expected file already exists - reuse_output - Skip running the benchmark if the output file already exists - - Returns - ------- - True if both operations succeed, False otherwise. - """ - print(f"\nRunning {query_name}...") - - expected_path = expected_dir / f"{query_name}.parquet" - benchmark_output = output_dir / f"{query_name}.parquet" - - # Generate expected - if reuse_expected and expected_path.exists(): - print(f" Reusing existing expected: {expected_path}") - else: - print(" Generating expected via DuckDB...") - try: - generate_expected(sql_path, input_dir, expected_path) - except Exception as e: - print(f" FAILED: Expected generation error: {e}") - return False - - # Run benchmark - if reuse_output and benchmark_output.exists(): - print(f" Reusing existing output: {benchmark_output}") - else: - result = run_benchmark(binary_path, input_dir, benchmark_output, extra_args) - - if result.returncode != 0: - print(f" FAILED: Benchmark exited with code {result.returncode}") - print(f" stdout: {result.stdout[:1000] if result.stdout else '(empty)'}") - print(f" stderr: {result.stderr[:1000] if result.stderr else '(empty)'}") - return False - - if not benchmark_output.exists(): - print(f" FAILED: Benchmark did not produce output file: {benchmark_output}") - return False - - print(" SUCCESS") - return True - - -def cmd_run(args: argparse.Namespace) -> int: - """Execute the 'run' subcommand.""" - # Validate paths - if not args.benchmark_dir.exists(): - print(f"Error: Benchmark directory does not exist: {args.benchmark_dir}") - return 1 - - if not args.sql_dir.exists(): - print(f"Error: SQL directory does not exist: {args.sql_dir}") - return 1 - - if args.generate_data: - generate_data(args.input_dir) - - if not args.input_dir.exists(): - print(f"Error: Input directory does not exist: {args.input_dir}") - return 1 - - output_dir = args.output_dir - output_dir.mkdir(parents=True, exist_ok=True) - - # Create subdirectories for output and expected - benchmark_output_dir = output_dir / "output" - expected_output_dir = output_dir / "expected" - benchmark_output_dir.mkdir(parents=True, exist_ok=True) - expected_output_dir.mkdir(parents=True, exist_ok=True) - - # Parse extra benchmark args - extra_args = args.benchmark_args.split() if args.benchmark_args else None - - # Discover benchmarks - benchmarks = discover_benchmarks(args.benchmark_dir, args.sql_dir) - - if not benchmarks: - print("No benchmarks found!") - return 1 - - # Filter to specific queries if requested - if args.queries: - benchmarks = [ - (name, binary, sql) - for name, binary, sql in benchmarks - if int(name.lstrip("q")) in args.queries - ] - if not benchmarks: - print(f"No matching benchmarks found for queries: {args.queries}") - return 1 - - print(f"Found {len(benchmarks)} benchmark(s) to run:") - for name, binary, sql in benchmarks: - print(f" {name}: {binary} + {sql}") - - # Run benchmarks - results = {} - for query_name, binary_path, sql_path in benchmarks: - passed = run_single_benchmark( - query_name, - binary_path, - sql_path, - args.input_dir, - benchmark_output_dir, - expected_output_dir, - extra_args, - reuse_expected=args.reuse_expected, - reuse_output=args.reuse_output, - ) - results[query_name] = passed - - # Summary - print("\n" + "=" * 60) - print("RUN SUMMARY") - print("=" * 60) - - passed = sum(results.values()) - failed = len(results) - passed - - print(f"Total: {passed} succeeded, {failed} failed") - print(f"\nOutput directory: {output_dir}") - print(f" Results: {benchmark_output_dir}") - print(f" Expected: {expected_output_dir}") - - return int(failed > 0) - - -def cmd_run_and_validate(args: argparse.Namespace) -> int: - """Execute the 'run-and-validate' subcommand.""" - # First run the benchmarks - run_result = cmd_run(args) - if run_result != 0: - print("\nRun phase failed, skipping validation.") - return run_result - - # Set up paths for validation based on run output - args.results_path = args.output_dir / "output" - args.expected_path = args.output_dir / "expected" - - return cmd_validate(args) - - -def cmd_validate(args: argparse.Namespace) -> int: - """Execute the 'validate' subcommand.""" - if not args.results_path.exists(): - print(f"Error: Results directory does not exist: {args.results_path}") - return 1 - - if not args.expected_path.exists(): - print(f"Error: Expected directory does not exist: {args.expected_path}") - return 1 - - # Discover parquet files in both directories - # But we treat *results* as the source of truth. If we have a result - # but not an expected we error; if we have an expected but not a result, - # that's fine. - results_files = discover_parquet_files(args.results_path) - - if not results_files: - print(f"No qDD.parquet files found in results directory: {args.results_path}") - return 1 - - # Filter to specific queries if requested - if args.queries: - results_files = { - name: path - for name, path in results_files.items() - if int(name.lstrip("q")) in args.queries - } - if not results_files: - print(f"No matching result files found for queries: {args.queries}") - return 1 - - print(f"\nValidating {len(results_files)} query(ies):") - - # Validate each matching pair - results = {} - for query_name in results_files.keys(): - print(f"\nValidating {query_name}...") - result_path = results_files[query_name] - expected_path = args.expected_path / f"{query_name}.parquet" - - if not expected_path.exists(): - print(f" FAILED: Expected file does not exist: {expected_path}") - results[query_name] = False - - is_equal, message = compare_parquet( - result_path, - expected_path, - decimal=args.decimal, - ignore_timezone=args.ignore_timezone, - ignore_string_type=args.ignore_string_type, - ignore_integer_sign=args.ignore_integer_sign, - ignore_integer_size=args.ignore_integer_size, - ignore_decimal_int=args.ignore_decimal_int, - ) - - if is_equal: - print(" PASSED") - results[query_name] = True - else: - print(f" FAILED:\n{message}") - results[query_name] = False - - # Summary - print("\n" + "=" * 60) - print("VALIDATION SUMMARY") - print("=" * 60) - - passed = sum(results.values()) - failed = len(results) - passed - - for query_name, result in sorted(results.items()): - status = "PASSED" if result else "FAILED" - print(f" {query_name}: {status}") - - print("-" * 60) - print(f"Total: {passed} passed, {failed} failed") - - return int(failed > 0) - - -def query_type(query: str) -> list[int]: - if query == "all": - return list(range(1, 23)) - else: - return [int(q) for q in query.split(",")] - - -def main(): - """Run the NDSH validation tool.""" - parser = argparse.ArgumentParser( - description="NDSH benchmark runner and validator", - formatter_class=argparse.RawDescriptionHelpFormatter, - epilog=__doc__, - ) - - subparsers = parser.add_subparsers(dest="command", required=True) - - # Parent parser for run-related arguments - run_parent = argparse.ArgumentParser(add_help=False) - run_parent.add_argument( - "--benchmark-dir", - type=Path, - help="Directory containing benchmark binaries (q04, q09, etc.)", - default=Path(__file__).parent.parent.parent.joinpath( - "cpp/build/benchmarks/ndsh" - ), - ) - run_parent.add_argument( - "--sql-dir", - type=Path, - help="Directory containing SQL query files (q04.sql, q09.sql, etc.)", - default=Path(__file__).parent.parent.parent.joinpath( - "cpp/benchmarks/streaming/ndsh/sql" - ), - ) - run_parent.add_argument( - "--input-dir", - type=Path, - required=True, - help="Directory containing TPC-H input parquet files", - ) - run_parent.add_argument( - "--output-dir", - type=Path, - required=True, - help="Directory for output files", - ) - run_parent.add_argument( - "-q", - "--queries", - help="Comma-separated list of SQL query numbers to run or the string 'all'", - type=query_type, - default="all", - ) - run_parent.add_argument( - "--benchmark-args", - type=str, - default="", - help="Additional arguments to pass to benchmark binaries (space-separated)", - ) - run_parent.add_argument( - "--reuse-expected", - action="store_true", - help="Skip generating expected results if the expected file already exists", - ) - run_parent.add_argument( - "--reuse-output", - action="store_true", - help="Skip running the benchmark if the output file already exists", - ) - run_parent.add_argument( - "--generate-data", - action="store_true", - help="Generate data for the benchmarks", - ) - - # Parent parser for validation comparison options (not the paths) - validate_options_parent = argparse.ArgumentParser(add_help=False) - validate_options_parent.add_argument( - "-d", - "--decimal", - type=int, - default=2, - help="Number of decimal places to compare for floating point values (default: 2)", - ) - validate_options_parent.add_argument( - "--ignore-timezone", - action="store_true", - help="Ignore differences in timezone and precision for timestamp types", - ) - validate_options_parent.add_argument( - "--ignore-string-type", - action="store_true", - help="Ignore differences between string and large_string types", - ) - validate_options_parent.add_argument( - "--ignore-integer-sign", - action="store_true", - help="Ignore differences between signed and unsigned integer types", - ) - validate_options_parent.add_argument( - "--ignore-integer-size", - action="store_true", - help="Ignore differences in integer bit width (e.g., int32 vs int64)", - ) - validate_options_parent.add_argument( - "--ignore-decimal-int", - action="store_true", - help="Ignore differences between decimal and integer types", - ) - - # 'run' subcommand - inherits from run_parent - subparsers.add_parser( - "run", - parents=[run_parent], - help="Run benchmarks and generate expected results", - description="Run C++ benchmark binaries and generate expected results via DuckDB.", - ) - - # 'validate' subcommand - inherits comparison options, adds its own paths - validate_parser = subparsers.add_parser( - "validate", - parents=[validate_options_parent], - help="Compare results against expected", - description="Validate benchmark results by comparing parquet files against expected results.", - ) - validate_parser.add_argument( - "--results-path", - type=Path, - required=True, - help="Directory containing benchmark result parquet files (qDD.parquet)", - ) - validate_parser.add_argument( - "--expected-path", - type=Path, - required=True, - help="Directory containing expected parquet files (qDD.parquet)", - ) - validate_parser.add_argument( - "-q", - "--queries", - help="Comma-separated list of SQL query numbers to validate or the string 'all'", - type=query_type, - default="all", - ) - - # 'run-and-validate' subcommand - inherits from BOTH parents - subparsers.add_parser( - "run-and-validate", - parents=[run_parent, validate_options_parent], - help="Run benchmarks and validate results in one step", - description="Run C++ benchmark binaries, generate expected results via DuckDB, and validate.", - ) - - args = parser.parse_args() - - if args.command == "run": - sys.exit(cmd_run(args)) - elif args.command == "validate": - sys.exit(cmd_validate(args)) - elif args.command == "run-and-validate": - sys.exit(cmd_run_and_validate(args)) - - -if __name__ == "__main__": - main() diff --git a/cpp/src/memory/host_memory_resource.cpp b/cpp/src/memory/host_memory_resource.cpp index a7b024be2..8b285cbba 100644 --- a/cpp/src/memory/host_memory_resource.cpp +++ b/cpp/src/memory/host_memory_resource.cpp @@ -20,7 +20,6 @@ namespace { * Attempts to mark the specified memory region as eligible for Transparent Huge Pages * (THP) using `madvise(MADV_HUGEPAGE)`. This is a best-effort optimization that can * improve device to host memory transfer performance for sufficiently large buffers. - * See . * * @param ptr Pointer to the start of the memory region. * @param size Size of the region in bytes. diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 9c007de4e..93cd78bd6 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -116,16 +116,6 @@ if(RAPIDSMPF_HAVE_STREAMING) ) endif() -# cudf-dependent test sources, gated behind BUILD_CUDF_TESTS -if(BUILD_CUDF_TESTS) - target_sources(test_sources PRIVATE test_partition.cpp) - target_link_libraries( - test_sources PRIVATE cudf_streaming::cudf_streaming cudf::cudftestutil cudf::cudftestutil_impl - ) - target_compile_definitions(test_sources PRIVATE RAPIDSMPF_HAVE_CUDF) - -endif() - if(RAPIDSMPF_HAVE_MPI) add_executable(mpi_tests main/mpi.cpp) set_target_properties( diff --git a/cpp/tests/test_partition.cpp b/cpp/tests/test_partition.cpp deleted file mode 100644 index 8ba164e9a..000000000 --- a/cpp/tests/test_partition.cpp +++ /dev/null @@ -1,103 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include - -#include "utils.hpp" - -using namespace rapidsmpf; -using namespace cudf_streaming::integrations; - -class NumOfPartitions : public cudf::test::BaseFixtureWithParam> {}; - -// test different `num_partitions` and `num_rows`. -INSTANTIATE_TEST_SUITE_P( - Partitions, - NumOfPartitions, - testing::Combine( - testing::Range(1, 10), // num_partitions - testing::Range(1, 100, 9) // num_rows - ) -); - -TEST_P(NumOfPartitions, partition_and_pack) { - int const num_partitions = std::get<0>(GetParam()); - int const num_rows = std::get<1>(GetParam()); - std::int64_t const seed = 42; - cudf::hash_id const hash_fn = cudf::hash_id::HASH_MURMUR3; - auto stream = cudf::get_default_stream(); - auto br = rapidsmpf::BufferResource::create(mr()); - - cudf::table expect = - random_table_with_index(seed, static_cast(num_rows), 0, 10); - - auto chunks = cudf_streaming::integrations::partition_and_pack( - expect, {1}, num_partitions, hash_fn, seed, stream, br.get() - ); - - // Convert to a vector - std::vector chunks_vector; - for (auto& [_, chunk] : chunks) { - chunks_vector.push_back(std::move(chunk)); - } - EXPECT_EQ(chunks_vector.size(), num_partitions); - - auto result = cudf_streaming::integrations::unpack_and_concat( - std::move(chunks_vector), stream, br.get() - ); - - // Compare the input table with the result. We ignore the row order by - // sorting by their index (first column). - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(sort_table(expect), sort_table(result)); -} - -TEST_P(NumOfPartitions, split_and_pack) { - int const num_partitions = std::get<0>(GetParam()); - int const num_rows = std::get<1>(GetParam()); - std::int64_t const seed = 42; - auto stream = cudf::get_default_stream(); - auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); - - cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); - - std::vector splits; - for (int i = 1; i < num_partitions; ++i) { - splits.emplace_back(i * num_rows / num_partitions); - } - - auto chunks = - cudf_streaming::integrations::split_and_pack(expect, splits, stream, br.get()); - - // Convert to a vector (restoring the original order). - std::vector chunks_vector; - for (int i = 0; i < num_partitions; ++i) { - chunks_vector.emplace_back(std::move(chunks.at(i))); - } - EXPECT_EQ(chunks_vector.size(), num_partitions); - - auto result = cudf_streaming::integrations::unpack_and_concat( - std::move(chunks_vector), stream, br.get() - ); - - // Compare the input table with the result. - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, *result); -} diff --git a/cpp/tests/utils.hpp b/cpp/tests/utils.hpp index 94eace6ac..3042a7b4b 100644 --- a/cpp/tests/utils.hpp +++ b/cpp/tests/utils.hpp @@ -23,14 +23,6 @@ #include -#ifdef RAPIDSMPF_HAVE_CUDF -#include -#include -#include -#include -#include -#endif // RAPIDSMPF_HAVE_CUDF - #include #include @@ -119,57 +111,6 @@ template return ret; } -#ifdef RAPIDSMPF_HAVE_CUDF -template -[[nodiscard]] inline std::unique_ptr iota_column( - std::size_t nrows, T start = 0 -) { - std::vector vec = iota_vector(nrows, start); - cudf::test::fixed_width_column_wrapper ret(vec.begin(), vec.end()); - return ret.release(); -} - -[[nodiscard]] inline std::unique_ptr random_column( - std::int64_t seed, - std::size_t nrows, - std::int64_t min = std::numeric_limits::min(), - std::int64_t max = std::numeric_limits::max() -) { - std::vector vec = random_vector(seed, nrows, min, max); - cudf::test::fixed_width_column_wrapper ret(vec.begin(), vec.end()); - return ret.release(); -} - -[[nodiscard]] inline cudf::table random_table_with_index( - std::int64_t seed, - std::size_t nrows, - std::int64_t min = std::numeric_limits::min(), - std::int64_t max = std::numeric_limits::max() -) { - std::vector> cols; - cols.push_back(iota_column(nrows)); - cols.push_back(random_column(seed, nrows, min, max)); - return cudf::table(std::move(cols)); -} - -[[nodiscard]] inline cudf::table sort_table( - cudf::table_view const& table, - std::vector const& /* column_indices */ = {0} -) { - if (table.num_columns() == 0) { - return cudf::table(table); - } - return cudf::gather(table, cudf::sorted_order(table.select({0}))->view())->release(); -} - -[[nodiscard]] inline cudf::table sort_table( - std::unique_ptr const& table, - std::vector const& column_indices = {0} -) { - return sort_table(table->view(), column_indices); -} -#endif // RAPIDSMPF_HAVE_CUDF - /// @brief Create a PackedData object from a host buffer [[nodiscard]] inline rapidsmpf::PackedData create_packed_data( std::span metadata, diff --git a/cpp/tools/rrun.cpp b/cpp/tools/rrun.cpp index 119c144cf..76198c854 100644 --- a/cpp/tools/rrun.cpp +++ b/cpp/tools/rrun.cpp @@ -92,11 +92,11 @@ void print_usage(std::string_view prog_name) { << " # Passthrough: multiple (4) tasks per node, one task per GPU, two nodes.\n" << " srun --mpi=pmix --nodes=2 --ntasks-per-node=4 --cpus-per-task=36 \\\n" << " --gpus-per-task=1 --gres=gpu:4 \\\n" - << " rrun ./benchmarks/bench_shuffle -C ucxx\n\n" + << " rrun ./benchmarks/bench_comm -C ucxx\n\n" << " # Hybrid mode: one task per node, 4 GPUs per task, two nodes.\n" << " srun --mpi=pmix --nodes=2 --ntasks-per-node=1 --cpus-per-task=144 \\\n" << " --gpus-per-task=4 --gres=gpu:4 \\\n" - << " rrun -n 4 ./benchmarks/bench_shuffle -C ucxx\n\n" + << " rrun -n 4 ./benchmarks/bench_comm -C ucxx\n\n" << std::endl; } diff --git a/dependencies.yaml b/dependencies.yaml index b3a35bbd3..4fedb06ae 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -27,7 +27,6 @@ files: - py_version - rapids_build_skbuild - test_cpp - - test_cpp_ndsh - test_python - docs devcontainers: @@ -51,7 +50,6 @@ files: - py_version - rapids_build_skbuild - test_cpp - - test_cpp_ndsh - test_python - docs test_cpp: @@ -62,7 +60,6 @@ files: - depends_on_librapidsmpf_tests - py_version - test_cpp - - test_cpp_ndsh test_python: output: none includes: @@ -75,8 +72,6 @@ files: - py_version - run_rapidsmpf - test_python - - depends_on_pylibcudf - - depends_on_cudf_streaming checks: output: none includes: @@ -170,8 +165,6 @@ files: includes: - depends_on_ray - test_python - - depends_on_pylibcudf - - depends_on_cudf_streaming channels: - rapidsai-nightly - rapidsai @@ -390,18 +383,6 @@ dependencies: - libnuma - openmpi >=5.0 # See - valgrind - - libcudf-streaming==26.8.*,>=0.0.0a0 - test_cpp_ndsh: - common: - - output_types: conda - packages: - # NDSH validation dependencies - - duckdb - - numpy - - pyarrow - - pip - - pip: - - tpchgen-cli test_python: common: - output_types: conda @@ -620,33 +601,3 @@ dependencies: - output_types: [requirements, pyproject] packages: - ray>=2.55.1 - depends_on_pylibcudf: - common: - - output_types: conda - packages: - - &pylibcudf_unsuffixed pylibcudf==26.8.*,>=0.0.0a0 - - output_types: requirements - packages: - # pip recognizes the index as a global option for the requirements.txt file - - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple - specific: - - output_types: [requirements, pyproject] - matrices: - - matrix: - cuda: "12.*" - cuda_suffixed: "true" - packages: - - pylibcudf-cu12==26.8.*,>=0.0.0a0 - - matrix: - cuda: "13.*" - cuda_suffixed: "true" - packages: - - pylibcudf-cu13==26.8.*,>=0.0.0a0 - - matrix: - packages: - - *pylibcudf_unsuffixed - depends_on_cudf_streaming: - common: - - output_types: conda - packages: - - cudf-streaming==26.8.*,>=0.0.0a0 diff --git a/docs/source/background/shuffle-architecture.md b/docs/source/background/shuffle-architecture.md index 0c837f826..ee1400ef4 100644 --- a/docs/source/background/shuffle-architecture.md +++ b/docs/source/background/shuffle-architecture.md @@ -140,7 +140,7 @@ it under nsys to capture a report: ``` $ nsys profile -o spill --trace cuda,nvtx \ - python -m rapidsmpf.benchmarks.streaming_benchmark --spill-device '1MiB' --out-nparts 4 --part-size 1MiB --local-size 24MiB + python spill.py ``` The `rapidsmpf.report` command line interface can analyze the rapidsmpf diff --git a/docs/source/cpp/index.md b/docs/source/cpp/index.md index ec8b29fe2..687ad688f 100644 --- a/docs/source/cpp/index.md +++ b/docs/source/cpp/index.md @@ -12,23 +12,16 @@ The C++ API reference is available at The C++ API provides access to all core RapidsMPF subsystems: - **Communicator** — MPI and UCXX backends for inter-process communication. -- **Shuffler** — Out-of-core, distributed table shuffle service. +- **Shuffler** — Out-of-core, distributed payload shuffle service. - **Streaming Engine** — Asynchronous multi-GPU pipeline with Channels, Actors, and Messages. - **Memory** — BufferResource, spilling, pinned memory, and packed data utilities. - **Config** — Configuration options and environment-variable parsing. -## Table Shuffle Service +## Shuffle Service See {doc}`../background/shuffle-architecture` for an in-depth explanation of the shuffle design. -The following is a complete MPI program that uses the RapidsMPF shuffler: - -```{literalinclude} ../../../cpp/examples/example_shuffle.cpp -:language: cpp -:lines: 7- -``` - ## rrun — Distributed Launcher RapidsMPF includes `rrun`, a lightweight launcher that eliminates the MPI dependency diff --git a/docs/source/getting-started.md b/docs/source/getting-started.md index 7e363eccb..c58493276 100644 --- a/docs/source/getting-started.md +++ b/docs/source/getting-started.md @@ -88,14 +88,14 @@ mpirun -np 2 cpp/build/gtests/mpi_tests cd cpp/build && ctest -R mpi_tests_2 ``` -We can also run the shuffle benchmark. To assign each MPI rank its own GPU, we use a +We can also run the communication benchmark. To assign each MPI rank its own GPU, we use a [binder script](https://github.com/LStuber/binding/blob/master/binder.sh): ```bash # The binder script requires numactl: mamba install numactl wget https://raw.githubusercontent.com/LStuber/binding/refs/heads/master/binder.sh chmod a+x binder.sh -mpirun -np 2 ./binder.sh cpp/build/benchmarks/bench_shuffle +mpirun -np 2 ./binder.sh cpp/build/benchmarks/bench_comm -C mpi ``` ## UCX diff --git a/docs/source/python/index.md b/docs/source/python/index.md index 7a2fef201..865fdff30 100644 --- a/docs/source/python/index.md +++ b/docs/source/python/index.md @@ -6,7 +6,7 @@ frameworks. ## Quickstart -- {doc}`quickstart` — Streaming Engine example +- {doc}`quickstart` — Streaming Engine overview ## API Reference diff --git a/docs/source/python/quickstart.md b/docs/source/python/quickstart.md index b1076c7b1..cb89292a5 100644 --- a/docs/source/python/quickstart.md +++ b/docs/source/python/quickstart.md @@ -9,12 +9,6 @@ some larger runtime. ## Streaming Engine -Basic streaming pipeline example in Python. In this example we have 3 {term}`Actor`s -in the {term}`Network`: push_to_channel->count_num_rows->pull_from_channel. - -*note: push_to_channel/pull_from_channel are convenience functions which simulate scans/writes* - -```{literalinclude} ../../../python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py -:language: python -:lines: 34- -``` +The Python streaming API exposes {term}`Actor`, {term}`Channel`, and message primitives +for downstream libraries that need to assemble their own pipelines. See +{doc}`api` for the available classes and functions. diff --git a/python/rapidsmpf/pyproject.toml b/python/rapidsmpf/pyproject.toml index b9e8a6d59..d03357d20 100644 --- a/python/rapidsmpf/pyproject.toml +++ b/python/rapidsmpf/pyproject.toml @@ -41,7 +41,6 @@ test = [ "cupy-cuda13x>=13.6.0,!=14.0.0,!=14.1.0", "numpy >=1.23,<3.0", "psutil", - "pylibcudf==26.8.*,>=0.0.0a0", "pytest", "ray>=2.55.1", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. diff --git a/python/rapidsmpf/rapidsmpf/benchmarks/__init__.py b/python/rapidsmpf/rapidsmpf/benchmarks/__init__.py deleted file mode 100644 index 4e1455b8e..000000000 --- a/python/rapidsmpf/rapidsmpf/benchmarks/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""RapidsMPF Python benchmarks.""" diff --git a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py b/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py deleted file mode 100644 index 7e2dc9201..000000000 --- a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py +++ /dev/null @@ -1,388 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -"""Example performing a streaming shuffle.""" - -from __future__ import annotations - -import argparse -import threading -import time -from typing import TYPE_CHECKING - -import cupy as cp -import pylibcudf as plc -from mpi4py import MPI -from pylibcudf.contiguous_split import pack - -import rmm.mr -from rmm.pylibrmm.stream import DEFAULT_STREAM - -import rapidsmpf.bootstrap -import rapidsmpf.communicator.mpi -from rapidsmpf.config import Options, get_environment_variables -from rapidsmpf.memory.buffer import MemoryType -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.memory.packed_data import PackedData -from rapidsmpf.progress_thread import ProgressThread -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor -from rapidsmpf.shuffler import Shuffler -from rapidsmpf.statistics import Statistics -from rapidsmpf.utils.string import format_bytes, parse_bytes - -if TYPE_CHECKING: - from rapidsmpf.communicator.communicator import Communicator - - -def generate_partition(size_bytes: int) -> plc.Table: - """ - Generate a random partition of data. - - Parameters - ---------- - size_bytes - size of the table in bytes - - Returns - ------- - plc.Table - """ - n_rows = size_bytes // 8 # each row is 8 bytes - return plc.Table( - [ - plc.Column.from_array(cp.arange(0, n_rows, dtype=cp.int32)), - plc.Column.from_array(cp.arange(0, n_rows, dtype=cp.float32)), - ] - ) - - -def consume_finished_partitions( - total_partitions: int, - comm: Communicator, - shuffler: Shuffler, -) -> None: - """ - Consume the finished partitions. - - Parameters - ---------- - total_partitions - The total number of partitions. - comm - The communicator to use. - shuffler - The shuffler to use. - """ - finished = set() - shuffler.wait() - for partition_id in shuffler.local_partitions(): - assert partition_id % comm.nranks == comm.rank - - # discard the extracted partition splits - shuffler.extract(partition_id) - - finished.add(partition_id) - - # all gather len(finished) to determine if all partitions have finished - comm.logger.print(f"Received parts: {len(finished)}") - finished_parts: int = MPI.COMM_WORLD.allreduce(len(finished), op=MPI.SUM) - assert finished_parts == total_partitions, "all partitions have not finished" - - -def streaming_shuffle( - comm: Communicator, - br: BufferResource, - output_nparts: int, - local_size: int, - part_size: int, - insert_delay_ms: float, - wait_timeout: int | None, -) -> None: - """ - Run shuffle operation in a streaming fashion. - - Main thread will produce local partitions and stream them through the shuffler. A separate - consumer thread will consume the finished partitions, and discard them. - - Parameters - ---------- - comm - The communicator to use. - br - The buffer resource to use. - output_nparts - The total number of output partitions. - local_size - The size of the local partition. - part_size - The size of each partition. - insert_delay_ms - A delay (ms) before inserting a partition to the shuffler. - wait_timeout - Timeout to wait for completion - """ - assert local_size >= part_size, "local_size must be >= part_size" - assert local_size % part_size == 0, "local_size must be divisible by part_size" - assert part_size >= 8 * output_nparts, "part_size must be >= 8 * output_nparts" - assert part_size % output_nparts == 0, ( - "part_size must be divisible by output_nparts" - ) - - # create a shuffler instance - shuffler = Shuffler( - comm, - op_id=0, - total_num_partitions=output_nparts, - br=br, - ) - - # create a thread to consume the finished partitions. It is a daemon thread, so it - # will not block the main thread from exiting in case of an error. - consumer_thread = threading.Thread( - target=consume_finished_partitions, - args=(output_nparts, comm, shuffler), - daemon=True, - ) - - # start the consumer thread. This will wait for any finished partition. - consumer_thread.start() - - n_parts_local = local_size // part_size - - # simulate a hash partition by splitting a partition into total_num_partitions - split_size = part_size // output_nparts - dummy_table = generate_partition(split_size) - - comm.logger.print(f"num local partitions:{n_parts_local} split size:{split_size}") - for p in range(n_parts_local): - # generate chunks for a single local partition by deep copying the dummy table - # as packed columns - # NOTE: This will require part_size amount of GPU memory. - chunks: dict[int, PackedData] = {} - for i in range(output_nparts): - chunks[i] = PackedData.from_cudf_packed_columns( # type: ignore[attr-defined] - pack(dummy_table), DEFAULT_STREAM, br - ) - - if p > 0 and insert_delay_ms > 0: - time.sleep(insert_delay_ms / 1000) - - shuffler.insert_chunks(chunks) - # finish inserting all partitions - shuffler.insert_finished() - - # wait for the consumer thread to finish. - consumer_thread.join(timeout=wait_timeout) - - -def ucxx_mpi_setup(options: Options, progress_thread: ProgressThread) -> Communicator: - """ - Bootstrap UCXX cluster using MPI. - - Parameters - ---------- - options - Configuration options. - progress_thread - Progress thread for the communicator. - - Returns - ------- - Communicator - A new ucxx communicator. - """ - import ucxx._lib.libucxx as ucx_api - - from rapidsmpf.communicator.ucxx import ( - barrier, - get_root_ucxx_address, - new_communicator, - ) - - if MPI.COMM_WORLD.Get_rank() == 0: - comm = new_communicator( - MPI.COMM_WORLD.size, None, None, options, progress_thread - ) - root_address_str = get_root_ucxx_address(comm) - else: - root_address_str = None - - root_address_str = MPI.COMM_WORLD.bcast(root_address_str, root=0) - - if MPI.COMM_WORLD.Get_rank() != 0: - root_address = ucx_api.UCXAddress.create_from_buffer(root_address_str) - comm = new_communicator( - MPI.COMM_WORLD.size, None, root_address, options, progress_thread - ) - - assert comm.nranks == MPI.COMM_WORLD.size - barrier(comm) - return comm - - -def setup_and_run(args: argparse.Namespace) -> None: - """ - Setup the args. - - Parameters - ---------- - args - The arguments to parse. - """ - options = Options(get_environment_variables()) - - # Create a RMM stack with both a device pool and statistics. - mr = RmmResourceAdaptor( - rmm.mr.PoolMemoryResource( - rmm.mr.CudaMemoryResource(), - initial_pool_size=args.rmm_pool_size, - maximum_pool_size=args.rmm_pool_size, - ) - ) - rmm.mr.set_current_device_resource(mr) - - stats = Statistics(enable=args.statistics) - progress_thread = ProgressThread(stats) - if args.comm == "mpi": - comm = rapidsmpf.communicator.mpi.new_communicator( - MPI.COMM_WORLD, options, progress_thread - ) - elif args.comm == "ucxx": - if rapidsmpf.bootstrap.is_running_with_rrun(): - raise ValueError( - "UCXX communicator is not supported with rrun yet, due to missing allreduce support" - ) - else: - comm = ucxx_mpi_setup(options, progress_thread) - - # Create a buffer resource that limits device memory if `--spill-device` - # is not None. - memory_limits = ( - None if args.spill_device is None else {MemoryType.DEVICE: args.spill_device} - ) - br = BufferResource(mr, memory_limits=memory_limits, statistics=stats) - - args.out_nparts = args.out_nparts if args.out_nparts is not None else comm.nranks - args.part_size = args.part_size if args.part_size is not None else args.local_size - - if comm.rank == 0: - comm.logger.print(str(vars(args))) - - MPI.COMM_WORLD.barrier() - start_time = MPI.Wtime() - streaming_shuffle( - comm, - br, - args.out_nparts, - args.local_size, - args.part_size, - args.insert_delay_ms, - args.wait_timeout, - ) - elapsed_time = MPI.Wtime() - start_time - MPI.COMM_WORLD.barrier() - - mem_peak = format_bytes(mr.get_main_record().peak()) - comm.logger.print( - f"elapsed: {elapsed_time:.2f} sec | rmm device memory peak: {mem_peak}" - ) - - if args.statistics: - comm.logger.print(stats.report(mr=mr)) - - -def parse_args( - args: list[str] | None = None, -) -> argparse.Namespace: # numpydoc ignore=PR01,RT01 - """Parse command line arguments.""" - parser = argparse.ArgumentParser( - prog="streaming shuffle", description="Streaming shuffle example" - ) - - parser.add_argument( - "--out-nparts", - type=int, - help="Number of output partitions. Default: n_ranks in the cluster", - default=None, - ) - - parser.add_argument( - "--local-size", - type=parse_bytes, - default="1MiB", - help="Local data size. Default: 1MiB", - ) - - parser.add_argument( - "--part-size", - type=parse_bytes, - default=None, - help="Partition size. Local size will be split into partitions of this size. Default: local_sz.", - ) - parser.add_argument( - "--comm", - type=str, - default="mpi", - help="Communicator type", - choices={"mpi", "ucxx"}, - ) - - parser.add_argument( - "--rmm-pool-size", - type=parse_bytes, - default=(int(parse_bytes(rmm.mr.available_device_memory()[1]) * 0.8) // 256) - * 256, - help=( - "The size of the RMM pool as a string with unit such as '2MiB' and '4KiB'. " - "Default to 80%% of the total device memory, which is %(default)s." - ), - ) - - parser.add_argument( - "--spill-device", - type=lambda x: None if x is None else parse_bytes(x), - default=None, - help=( - "Spilling device-to-host threshold as a string with unit such as '2MiB' " - "and '4KiB'. Default is no spilling" - ), - ) - - parser.add_argument( - "--report", - action=argparse.BooleanOptionalAction, - default=True, - help="Print the statistics report", - ) - - parser.add_argument( - "--statistics", - default=False, - action="store_true", - help="Enable statistics.", - ) - - parser.add_argument( - "--insert-delay-ms", - type=float, - help="A delay (ms) before inserting a partition to the shuffler. Default: 0", - default=0, - ) - - parser.add_argument( - "--wait-timeout", - type=int, - default=None, - help="Wait timeout to finish. Default, wait indefinitely", - ) - - return parser.parse_args(args) - - -def main(args: list[str] | None = None) -> None: # numpydoc ignore=PR01 - """Streaming shuffle.""" - parsed = parse_args(args) - setup_and_run(parsed) - - -if __name__ == "__main__": - main() diff --git a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py deleted file mode 100644 index 6be1a0d29..000000000 --- a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py +++ /dev/null @@ -1,524 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Bulk-synchronous MPI shuffle.""" - -from __future__ import annotations - -import argparse -import math -import uuid -from pathlib import Path -from typing import TYPE_CHECKING - -import pylibcudf as plc -from cudf_streaming.integrations.partition import ( - partition_and_pack, - unpack_and_concat, -) -from mpi4py import MPI - -import rmm.mr -from rmm.pylibrmm.stream import DEFAULT_STREAM - -import rapidsmpf.bootstrap -import rapidsmpf.communicator.mpi -from rapidsmpf.config import Options, get_environment_variables -from rapidsmpf.memory.buffer import MemoryType -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.memory.spill import unspill_partitions -from rapidsmpf.progress_thread import ProgressThread -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor -from rapidsmpf.shuffler import Shuffler -from rapidsmpf.statistics import Statistics -from rapidsmpf.utils.string import format_bytes, parse_bytes - -try: - from rapidsmpf.cupti import CuptiMonitor - - CUPTI_AVAILABLE = True -except ImportError: - CUPTI_AVAILABLE = False - -if TYPE_CHECKING: - from collections.abc import Callable - - from rapidsmpf.communicator.communicator import Communicator - - -def barrier(comm: Communicator) -> None: - """ - Blocks until all processes in the communicator have reached this point. - - Parameters - ---------- - comm - The communicator to barrier. - """ - if rapidsmpf.bootstrap.is_running_with_rrun(): - from rapidsmpf.communicator.ucxx import barrier as ucxx_barrier - - ucxx_barrier(comm) - else: - MPI.COMM_WORLD.barrier() - - -def read_batch(paths: list[str]) -> tuple[plc.Table, list[str]]: - """ - Read a single batch of Parquet files. - - Parameters - ---------- - paths - List of file paths to the Parquet files. - - Returns - ------- - plc.Table - The table containing the data read from the Parquet files. - list of str - Column names from the Parquet files, excluding nested children. - """ - options = plc.io.parquet.ParquetReaderOptions.builder( - plc.io.SourceInfo(paths) - ).build() - tbl_w_meta = plc.io.parquet.read_parquet(options) - return (tbl_w_meta.tbl, tbl_w_meta.column_names(include_children=False)) - - -def write_table( - table: plc.Table, output_path: str, id: int | str, column_names: list[str] | None -) -> None: - """ - Write a pylibcudf Table to a Parquet file. - - Parameters - ---------- - table - The table to be written to the Parquet file. - output_path : str - Directory where the Parquet file will be written. - id - Unique identifier used to generate the filename using `part.{id}.parque`. - column_names - List of column names. - """ - path = f"{output_path}/part.{id}.parquet" - builder = plc.io.parquet.ParquetWriterOptions.builder( - plc.io.SinkInfo([path]), table - ) - if column_names is not None: - metadata = plc.io.types.TableInputMetadata(table) - for col_meta, name in zip(metadata.column_metadata, column_names, strict=True): - col_meta.set_name(name) - builder = builder.metadata(metadata) - plc.io.parquet.write_parquet(builder.build()) - - -def bulk_mpi_shuffle( - paths: list[str], - shuffle_on: list[str], - output_path: str, - comm: Communicator, - br: BufferResource, - *, - num_output_files: int | None = None, - batchsize: int = 1, - read_func: Callable = read_batch, - write_func: Callable = write_table, - baseline: bool = False, - statistics: Statistics | None = None, -) -> None: - """ - Perform a bulk-synchronous dataset shuffle. - - Parameters - ---------- - paths - List of file paths to shuffle. This list contains all files in the - dataset (not just the files that will be processed by the local rank). - shuffle_on - List of column names to shuffle on. - output_path - Path of the output directory where the data will be written. This - directory does not need to be on a shared filesystem. - comm - The communicator to use. - br - Buffer resource to use. - num_output_files - Number of output files to produce. Default will preserve the - input file count. - batchsize - Number of files to read at once on each rank. - read_func - Call-back function to read the input data. This function must accept a - list of file paths, and return a pylibcudf Table and the list of column - names in the table. Default logic will use `pylibcudf.read_parquet`. - write_func - Call-back function to write shuffled data to disk. This function must - accept `table`, `output_path`, `id`, and `column_names` arguments. - Default logic will write the pylibcudf table to a parquet file - (e.g. `f"{output_path}/part.{id}.parquet"`). - baseline - Whether to skip the shuffle and run a simple IO baseline. - statistics - The statistics instance to use. If None, statistics is disabled. - - Notes - ----- - This function is executed on each rank of the MPI communicator in a - bulk-synchronous fashion. This means all ranks are expected to call - this same function with the same arguments. - """ - # Create output directory if necessary - Path(output_path).mkdir(exist_ok=True) - - # Determine which files to process on this rank - num_input_files = len(paths) - num_output_files = num_output_files or num_input_files - total_num_partitions = num_output_files - files_per_rank = math.ceil(num_input_files / comm.nranks) - start = files_per_rank * comm.rank - finish = start + files_per_rank - local_files = paths[start:finish] - num_local_files = len(local_files) - num_batches = math.ceil(num_local_files / batchsize) - - if baseline: - # Skip the shuffle - Run IO baseline - for batch_id in range(num_batches): - batch = local_files[batch_id * batchsize : (batch_id + 1) * batchsize] - table, columns = read_func(batch) - write_func( - table, - output_path, - str(uuid.uuid4()), - columns, - ) - else: - br = BufferResource(rmm.mr.get_current_device_resource()) - shuffler = Shuffler( - comm, - op_id=0, - total_num_partitions=total_num_partitions, - br=br, - ) - - # Read batches and submit them to the shuffler - column_names = None - for batch_id in range(num_batches): - batch = local_files[batch_id * batchsize : (batch_id + 1) * batchsize] - table, columns = read_func(batch) - if column_names is None: - column_names = columns - columns_to_hash = tuple(columns.index(val) for val in shuffle_on) - packed_inputs = partition_and_pack( - table, - columns_to_hash=columns_to_hash, - num_partitions=total_num_partitions, - br=br, - stream=DEFAULT_STREAM, - ) - shuffler.insert_chunks(packed_inputs) - - # Tell the shuffler we are done adding local data - shuffler.insert_finished() - - # Write shuffled partitions to disk - shuffler.wait() - for partition_id in shuffler.local_partitions(): - table = unpack_and_concat( - unspill_partitions( - shuffler.extract(partition_id), - br=br, - allow_overbooking=True, - ), - br=br, - stream=DEFAULT_STREAM, - ) - write_func( - table, - output_path, - partition_id, - column_names, - ) - shuffler.shutdown() - - -def ucxx_mpi_setup(options: Options, progress_thread: ProgressThread) -> Communicator: - """ - Bootstrap UCXX cluster using MPI. - - Parameters - ---------- - options - Configuration options. - progress_thread - Progress thread for the initialized communicator. - - Returns - ------- - Communicator - A new ucxx communicator. - """ - import ucxx._lib.libucxx as ucx_api - - from rapidsmpf.communicator.ucxx import ( - barrier, - get_root_ucxx_address, - new_communicator, - ) - - if MPI.COMM_WORLD.Get_rank() == 0: - comm = new_communicator( - MPI.COMM_WORLD.size, None, None, options, progress_thread - ) - root_address_bytes = get_root_ucxx_address(comm) - else: - root_address_bytes = None - - root_address_bytes = MPI.COMM_WORLD.bcast(root_address_bytes, root=0) - - if MPI.COMM_WORLD.Get_rank() != 0: - root_address = ucx_api.UCXAddress.create_from_buffer(root_address_bytes) - comm = new_communicator( - MPI.COMM_WORLD.size, None, root_address, options, progress_thread - ) - - assert comm.nranks == MPI.COMM_WORLD.size - barrier(comm) - return comm - - -def setup_and_run(args: argparse.Namespace) -> None: - """ - Set up the environment and run the shuffle example. - - Parameters - ---------- - args - Command-line arguments containing the configuration for the shuffle example. - """ - options = Options(get_environment_variables()) - - # Create a RMM stack with both a device pool and statistics. - mr = RmmResourceAdaptor( - rmm.mr.PoolMemoryResource( - rmm.mr.CudaMemoryResource(), - initial_pool_size=args.rmm_pool_size, - maximum_pool_size=args.rmm_pool_size, - ) - ) - rmm.mr.set_current_device_resource(mr) - - # Create a buffer resource that limits device memory if `--spill-device` - # is not None. - memory_limits = ( - None if args.spill_device is None else {MemoryType.DEVICE: args.spill_device} - ) - br = BufferResource(mr, memory_limits=memory_limits) - - stats = Statistics(enable=args.statistics) - - progress_thread = ProgressThread(stats) - if args.cluster_type == "mpi": - comm = rapidsmpf.communicator.mpi.new_communicator( - MPI.COMM_WORLD, options, progress_thread - ) - elif args.cluster_type == "ucxx": - if rapidsmpf.bootstrap.is_running_with_rrun(): - comm = rapidsmpf.bootstrap.create_ucxx_comm( - progress_thread, - type=rapidsmpf.bootstrap.BackendType.AUTO, - options=options, - ) - else: - comm = ucxx_mpi_setup(options, progress_thread) - cupti_monitor = None - if args.monitor_memory is not None: - if not CUPTI_AVAILABLE: - if comm.rank == 0: - comm.logger.print( - "WARNING: --memory-monitor specified but CUPTI support not available. " - "CUPTI monitoring disabled." - ) - else: - cupti_monitor = CuptiMonitor(enable_periodic_sampling=False) - if comm.rank == 0: - comm.logger.print("CUPTI memory monitoring enabled") - - if comm.rank == 0: - spill_device = ( - "disabled" if args.spill_device is None else format_bytes(args.spill_device) - ) - comm.logger.print( - f"""\ -Shuffle: - input: {args.input} - output: {args.output} - on: {args.on} - --cluster-type: {args.cluster_type} - --n-output-files: {args.n_output_files} - --batchsize: {args.batchsize} - --baseline: {args.baseline} - --rmm-pool-size: {format_bytes(args.rmm_pool_size)} - --spill-device: {spill_device}""" - ) - - barrier(comm) - - if cupti_monitor is not None: - cupti_monitor.start_monitoring() - - start_time = MPI.Wtime() - bulk_mpi_shuffle( - paths=sorted(map(str, args.input.glob("**/*"))), - shuffle_on=args.on.split(","), - output_path=args.output, - comm=comm, - br=br, - num_output_files=args.n_output_files, - batchsize=args.batchsize, - baseline=args.baseline, - statistics=stats, - ) - elapsed_time = MPI.Wtime() - start_time - barrier(comm) - - if cupti_monitor is not None: - cupti_monitor.stop_monitoring() - - csv_filename = f"{args.monitor_memory}_{comm.rank}.csv" - try: - # Write CSV files - cupti_monitor.write_csv(csv_filename) - comm.logger.print( - f"CUPTI memory data written to {csv_filename} " - f"({cupti_monitor.get_sample_count()} samples, " - f"{cupti_monitor.get_total_callback_count()} callbacks)" - ) - - # Print callback summary for rank 0 - if comm.rank == 0: - comm.logger.print( - f"CUPTI Callback Summary:\n{cupti_monitor.get_callback_summary()}" - ) - except Exception as e: - comm.logger.print(f"Failed to write CUPTI CSV file: {e}") - - mem_peak = format_bytes(mr.get_main_record().peak()) - comm.logger.print( - f"elapsed: {elapsed_time:.2f} sec | rmm device memory peak: {mem_peak}" - ) - if stats.enabled: - comm.logger.print(stats.report(mr=mr)) - - -def dir_path(path: str) -> Path: - """ - Validate that the given path is a directory and return a Path object. - - Parameters - ---------- - path - The path to check. - - Returns - ------- - Path - A Path object representing the directory. - """ - ret = Path(path) - if not ret.is_dir(): - raise ValueError() - return ret - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - prog="Bulk-synchronous MPI shuffle", - description="Shuffle a dataset at rest (on disk) on both ends.", - ) - parser.add_argument( - "input", - type=dir_path, - metavar="INPUT_DIR_PATH", - help="Input directory path.", - ) - parser.add_argument( - "output", - metavar="OUTPUT_DIR_PATH", - type=Path, - help="Output directory path.", - ) - parser.add_argument( - "on", - metavar="COLUMN_LIST", - type=str, - help="Comma-separated list of column names to shuffle on.", - ) - parser.add_argument( - "--n-output-files", - type=int, - default=None, - help="Number of output files. Default preserves input file count.", - ) - parser.add_argument( - "--batchsize", - type=int, - default=1, - help="Number of files to read on each MPI rank at once.", - ) - parser.add_argument( - "--baseline", - default=False, - action="store_true", - help="Run an IO baseline without any shuffling.", - ) - parser.add_argument( - "--rmm-pool-size", - type=parse_bytes, - default=format_bytes(int(rmm.mr.available_device_memory()[1] * 0.8)), - help=( - "The size of the RMM pool as a string with unit such as '2MiB' and '4KiB'. " - "Default to 80%% of the total device memory, which is %(default)s." - ), - ) - parser.add_argument( - "--spill-device", - type=lambda x: None if x is None else parse_bytes(x), - default=None, - help=( - "Spilling device-to-host threshold as a string with unit such as '2MiB' " - "and '4KiB'. Default is no spilling" - ), - ) - parser.add_argument( - "--statistics", - default=False, - action="store_true", - help="Enable statistics.", - ) - parser.add_argument( - "--cluster-type", - type=str, - default="mpi", - choices=("mpi", "ucxx"), - help=( - "Cluster type to setup. Regardless of the cluster type selected it must " - "be launched with 'mpirun'." - ), - ) - parser.add_argument( - "--monitor-memory", - type=str, - default=None, - help=( - "Enable memory monitoring with CUPTI and save CSV files with given path " - "prefix. For example, /tmp/test will write files to /tmp/test_.csv. " - "Requires CUPTI support to be compiled in." - ), - ) - args = parser.parse_args() - args.rmm_pool_size = (args.rmm_pool_size // 256) * 256 # Align to 256 bytes - setup_and_run(args) diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/__init__.py b/python/rapidsmpf/rapidsmpf/examples/ray/__init__.py deleted file mode 100644 index ec8cda507..000000000 --- a/python/rapidsmpf/rapidsmpf/examples/ray/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Submodule for Ray examples.""" diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py deleted file mode 100644 index 8508e0a5a..000000000 --- a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py +++ /dev/null @@ -1,451 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -"""Example running a Bulk RapidsMPF Shuffle operation using Ray and UCXX communication.""" - -from __future__ import annotations - -import argparse -import math -import os -import time -from pathlib import Path -from typing import TYPE_CHECKING - -import pylibcudf as plc -import ray -from cudf_streaming.integrations.partition import ( - partition_and_pack, - unpack_and_concat, -) - -import rmm.mr - -from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster -from rapidsmpf.memory.buffer import MemoryType -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.memory.spill import unspill_partitions -from rapidsmpf.rmm_resource_adaptor import RmmResourceAdaptor -from rapidsmpf.shuffler import Shuffler -from rapidsmpf.statistics import Statistics -from rapidsmpf.utils.cudf import pylibcudf_to_cudf_dataframe -from rapidsmpf.utils.string import format_bytes, parse_bytes - -if TYPE_CHECKING: - from collections.abc import Iterator - - -@ray.remote(num_gpus=1, num_cpus=4) -class BulkRayShufflerActor(RapidsMPFActor): - """ - Actor that performs a bulk shuffle operation using Ray. - - Parameters - ---------- - nranks - Number of ranks in the communication group. - total_nparts - Total number of output partitions. - shuffle_on - List of column names to shuffle on. - batchsize - Number of files to process in a batch. - output_path - Path to write output files. - rmm_pool_size - Size of the RMM memory pool in bytes. - spill_device - Device memory limit for spilling to host in bytes. - enable_statistics - Whether to collect statistics. - """ - - def __init__( - self, - nranks: int, - total_nparts: int, - shuffle_on: list[str], - batchsize: int = 1, - output_path: str = "./", - rmm_pool_size: int = 1024 * 1024 * 1024, - spill_device: int | None = None, - *, - enable_statistics: bool = False, - ): - self.batchsize = batchsize - self.shuffle_on = shuffle_on - self.output_path = output_path - self.total_nparts = total_nparts - self.rmm_pool_size = rmm_pool_size - self.spill_device = spill_device - - # Initialize actor-local resources (statistics, memory resource) - self.mr = RmmResourceAdaptor( - rmm.mr.PoolMemoryResource( - rmm.mr.CudaMemoryResource(), - initial_pool_size=self.rmm_pool_size, - maximum_pool_size=self.rmm_pool_size, - ) - ) - rmm.mr.set_current_device_resource(self.mr) - # Create a buffer resource that limits device memory if `--spill-device` - memory_limits = ( - None - if self.spill_device is None - else {MemoryType.DEVICE: self.spill_device} - ) - br = BufferResource(self.mr, memory_limits=memory_limits) - self.br = br - super().__init__(nranks, Statistics(enable=enable_statistics)) - - def setup_worker(self, root_address_bytes: bytes) -> None: - """ - Setup the UCXX communication and a shuffle operation. - - Parameters - ---------- - root_address_bytes - Address of the root worker for UCXX initialization. - """ - super().setup_worker(root_address_bytes) - self.shuffler: Shuffler = Shuffler( - self.comm, - 0, - total_num_partitions=self.total_nparts, - br=self.br, - ) - - def cleanup(self) -> None: - """Cleanup the UCXX communication and the shuffle operation.""" - self.comm.logger.info(self.statistics.report()) - if self.shuffler is not None: - self.shuffler.shutdown() - - def read_batch(self, paths: list[str]) -> tuple[plc.Table, list[str]]: - """ - Read a single batch of Parquet files. - - Parameters - ---------- - paths - List of file paths to the Parquet files. - - Returns - ------- - A tuple containing the read in table and the column names. - """ - options = plc.io.parquet.ParquetReaderOptions.builder( - plc.io.SourceInfo(paths) - ).build() - tbl_w_meta = plc.io.parquet.read_parquet(options) - return (tbl_w_meta.tbl, tbl_w_meta.column_names(include_children=False)) - - def write_table( - self, - table: plc.Table, - output_path: str, - id: int | str, - column_names: list[str], - ) -> None: - """ - Write a pylibcudf Table to a Parquet file. - - Parameters - ---------- - table - The table to write. - output_path - The path to write the table to. - id - Partition id used for naming the output file. - column_names - The column names of the table. - """ - path = f"{output_path}/part.{id}.parquet" - pylibcudf_to_cudf_dataframe( - table, - column_names=column_names, - ).to_parquet(path) - - def insert_chunk(self, table: plc.Table, column_names: list[str]) -> None: - """ - Insert a pylibcudf Table into the shuffler. - - Parameters - ---------- - table - The table to insert. - column_names - The column names of the table. - """ - from rmm.pylibrmm.stream import DEFAULT_STREAM - - columns_to_hash = tuple(column_names.index(val) for val in self.shuffle_on) - packed_inputs = partition_and_pack( - table, - columns_to_hash=columns_to_hash, - num_partitions=self.total_nparts, - br=self.br, - stream=DEFAULT_STREAM, - ) - self.shuffler.insert_chunks(packed_inputs) - - def read_and_insert(self, paths: list[str]) -> list[str]: - """ - Read the list of parquet files every batchsize and insert the partitions into the shuffler. - - Parameters - ---------- - paths - List of file paths to the Parquet files. - - Returns - ------- - The column names of the table. - """ - for i in range(0, len(paths), self.batchsize): - tbl, column_names = self.read_batch(paths[i : i + self.batchsize]) - self.insert_chunk(tbl, column_names) - self.insert_finished() - return column_names - - def insert_finished(self) -> None: - """Tell the shuffler that we are done inserting data.""" - self.shuffler.insert_finished() - self.comm.logger.info("Insert finished") - - def extract(self) -> Iterator[tuple[int, plc.Table]]: - """ - Extract shuffled partitions. - - Returns - ------- - An iterator over the shuffled partitions. - """ - from rmm.pylibrmm.stream import DEFAULT_STREAM - - self.shuffler.wait() - for partition_id in self.shuffler.local_partitions(): - packed_chunks = self.shuffler.extract(partition_id) - partition = unpack_and_concat( - unspill_partitions( - packed_chunks, - br=self.br, - allow_overbooking=True, - ), - br=self.br, - stream=DEFAULT_STREAM, - ) - yield partition_id, partition - - def extract_and_write(self, column_names: list[str]) -> None: - """ - Extract and write shuffled partitions. - - Parameters - ---------- - column_names - The column names of the table. - """ - for partition_id, partition in self.extract(): - self.write_table(partition, self.output_path, partition_id, column_names) - - -def bulk_ray_shuffle( - paths: list[str], - shuffle_on: list[str], - output_path: str, - num_workers: int = 2, - batchsize: int = 1, - num_output_files: int | None = None, - rmm_pool_size: int = 1024 * 1024 * 1024, - spill_device: int | None = None, - *, - enable_statistics: bool = False, -) -> None: - """ - Perform a bulk shuffle operation using Ray and UCXX communication. - - Parameters - ---------- - paths - The list of paths to the input files. - shuffle_on - The list of column names to shuffle on. - output_path - The directory to write the shuffled data. - num_workers - The number of workers to use. - batchsize - The number of files to read on each rank at once. - num_output_files - The number of output files to write. - rmm_pool_size - The size of the RMM pool. - spill_device - Device memory limit for spilling to host. - enable_statistics - Whether to collect statistics. - """ - # Initialize the UCXX cluster - num_input_files = len(paths) - num_output_files = num_output_files or num_input_files - total_num_partitions = num_output_files - files_per_rank = math.ceil(num_input_files / num_workers) - - actors = setup_ray_ucxx_cluster( - BulkRayShufflerActor, - num_workers=num_workers, - total_nparts=total_num_partitions, - shuffle_on=shuffle_on, - batchsize=batchsize, - output_path=output_path, - enable_statistics=enable_statistics, - rmm_pool_size=rmm_pool_size, - spill_device=spill_device, - ) - start_time = time.time() - insert_tasks = [] - for i, actor in enumerate(actors): - # Calculate the start and end indices for this actor's files - start = i * files_per_rank - # Use min to ensure we don't go beyond the end of the paths list - end = min(start + files_per_rank, num_input_files) - insert_tasks.append(actor.read_and_insert.remote(paths[start:end])) - column_names = ray.get(insert_tasks) - ray.get( - [ - actor.extract_and_write.remote(column_name) - for actor, column_name in zip(actors, column_names, strict=False) - ] - ) - end_time = time.time() - print(f"Time taken: {end_time - start_time} seconds") - ray.get([actor.cleanup.remote() for actor in actors]) - - -def dir_path(path: str) -> Path: - """ - Validate that the given path is a directory and return a Path object. - - Parameters - ---------- - path - The path to check. - - Returns - ------- - Path - A Path object representing the directory. - - Raises - ------ - ValueError - If the path is not a directory. - """ - ret = Path(path) - if not ret.is_dir(): - raise ValueError(f"{path} path is not a directory") - return ret - - -def setup_and_run(args: argparse.Namespace) -> None: - """ - Setup and run the bulk shuffle operation. - - Parameters - ---------- - args : argparse.Namespace - The parsed command line arguments. - """ - if args.ray_address or os.environ.get("RAY_ADDRESS") is not None: - ray.init(address="auto") # connect to existing cluster - else: - ray.init(num_gpus=args.num_workers, dashboard_host="0.0.0.0") - - bulk_ray_shuffle( - paths=sorted(map(str, args.input.glob("**/*"))), - shuffle_on=args.on.split(","), - output_path=args.output, - num_workers=args.num_workers, - batchsize=args.batchsize, - num_output_files=args.n_output_files, - enable_statistics=args.statistics, - rmm_pool_size=args.rmm_pool_size, - spill_device=args.spill_device, - ) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - prog="Bulk-synchronous Ray shuffle", - description="Shuffle a dataset at rest (on disk) on both ends.", - ) - parser.add_argument( - "--num-workers", - type=int, - default=2, - help="Number of workers to use.", - ) - parser.add_argument( - "input", - type=dir_path, - metavar="INPUT_DIR_PATH", - help="Input directory path.", - ) - parser.add_argument( - "output", - type=dir_path, - metavar="OUTPUT_DIR_PATH", - help="Output directory path.", - ) - parser.add_argument( - "on", - metavar="COLUMN_LIST", - type=str, - help="Comma-separated list of column names to shuffle on.", - ) - parser.add_argument( - "--n-output-files", - type=int, - default=None, - help="Number of output files. Default preserves input file count.", - ) - parser.add_argument( - "--batchsize", - type=int, - default=1, - help="Number of files to read on each MPI rank at once.", - ) - parser.add_argument( - "--rmm-pool-size", - type=parse_bytes, - default=format_bytes(int(rmm.mr.available_device_memory()[1] * 0.8)), - help=( - "The size of the RMM pool as a string with unit such as '2MiB' and '4KiB'. " - "Default to 80%% of the total device memory, which is %(default)s." - ), - ) - parser.add_argument( - "--spill-device", - type=lambda x: None if x is None else parse_bytes(x), - default=None, - help=( - "Spilling device-to-host threshold as a string with unit such as '2MiB' " - "and '4KiB'. Default is no spilling" - ), - ) - parser.add_argument( - "--statistics", - default=False, - action="store_true", - help="Enable statistics.", - ) - parser.add_argument( - "--ray-address", - type=str, - default=None, - help="Connect to an existing Ray cluster.", - ) - args = parser.parse_args() - args.rmm_pool_size = (args.rmm_pool_size // 256) * 256 # Align to 256 bytes - setup_and_run(args) diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py b/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py deleted file mode 100644 index 743038b55..000000000 --- a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py +++ /dev/null @@ -1,195 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -"""Example running a RapidsMPF Shuffle operation using Ray and UCXX communication.""" - -from __future__ import annotations - -import argparse -import math - -import numpy as np -import pylibcudf as plc -import ray -from cudf_streaming.integrations.partition import ( - partition_and_pack, - unpack_and_concat, -) - -import rmm - -from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.memory.spill import unspill_partitions -from rapidsmpf.shuffler import Shuffler -from rapidsmpf.testing import assert_eq - - -class ShufflingActor(RapidsMPFActor): - """ - An example of a Ray actor that performs a shuffle operation. - - Parameters - ---------- - nranks - Number of ranks. - num_rows - Number of rows in the input dataframe. - batch_size - Batch size (rows) of the input. The input dataframe will be split into batches of this size. - total_nparts - Total number of partitions into which the input dataframe will be partitioned. - """ - - def __init__( - self, - nranks: int, - num_rows: int = 100, - batch_size: int = -1, - total_nparts: int = -1, - ): - super().__init__(nranks, statistics=None) - self._num_rows: int = num_rows - self._batch_size: int = batch_size - self._total_nparts: int = total_nparts if total_nparts > 0 else nranks - - def _gen_table(self) -> plc.Table: - """ - Generate a dummy table with three columns ("a", "b", "c"). - - Returns - ------- - plc.Table - The input table. - """ - # Every rank creates the full input table and all the expected partitions - # (also partitions this rank might not get after the shuffle). - - np.random.seed(42) # Make sure all ranks create the same input table. - - return plc.Table( - [ - plc.Column.from_iterable_of_py( - range(self._num_rows), plc.DataType(plc.TypeId.INT64) - ), - plc.Column.from_array(np.random.randint(0, 1000, self._num_rows)), - plc.Column.from_iterable_of_py( - ["cat", "dog"] * (self._num_rows // 2), - plc.DataType(plc.TypeId.STRING), - ), - ] - ) - - def run(self) -> None: - """Run the shuffle operation, and this will be called remotely from the client.""" - # If DEFAULT_STREAM was imported outside of this context, it will be pickled, - # and it is not serializable. Therefore, we need to import it here. - from rmm.pylibrmm.stream import DEFAULT_STREAM - - df = self._gen_table() - columns_to_hash = (1,) - - mr = rmm.mr.get_current_device_resource() - br = BufferResource(mr) - stream = DEFAULT_STREAM # use the default stream - - # Calculate the expected output partitions on all ranks - expected = { - partition_id: unpack_and_concat( - [packed], - br=br, - stream=stream, - ) - for partition_id, packed in partition_and_pack( - df, - columns_to_hash=columns_to_hash, - num_partitions=self._total_nparts, - br=br, - stream=stream, - ).items() - } - - shuffler = Shuffler( - self.comm, - 0, - total_num_partitions=self._total_nparts, - br=br, - ) - - # Slice df and submit local slices to shuffler - stride = math.ceil(self._num_rows / self.comm.nranks) - local_df = plc.copying.slice( - df, - [ - self.comm.rank * stride, - min((self.comm.rank + 1) * stride, self._num_rows), - ], - )[0] - num_rows_local = local_df.num_rows() - self._batch_size = num_rows_local if self._batch_size < 0 else self._batch_size - for i in range(0, num_rows_local, self._batch_size): - batch = plc.copying.slice( - local_df, [i, min(i + self._batch_size, num_rows_local)] - )[0] - packed_inputs = partition_and_pack( - batch, - columns_to_hash=columns_to_hash, - num_partitions=self._total_nparts, - br=br, - stream=stream, - ) - shuffler.insert_chunks(packed_inputs) - - # Tell shuffler we are done adding data - shuffler.insert_finished() - - # Extract and check shuffled partitions - shuffler.wait() - for partition_id in shuffler.local_partitions(): - packed_chunks = shuffler.extract(partition_id) - partition = unpack_and_concat( - unspill_partitions(packed_chunks, br=br, allow_overbooking=True), - br=br, - stream=stream, - ) - assert_eq( - partition, - expected[partition_id], - sort_rows=0, - ) - - shuffler.shutdown() - - -@ray.remote(num_gpus=1) -class GpuShufflingActor(ShufflingActor): - """Shuffle example class with 1 GPU resource.""" - - -if __name__ == "__main__": - parser = argparse.ArgumentParser( - description="RapidsMPF Ray Shuffling Actor example ", - ) - parser.add_argument("--nranks", type=int, default=1) - parser.add_argument("--num_rows", type=int, default=100) - parser.add_argument("--batch_size", type=int, default=-1) - parser.add_argument("--total_nparts", type=int, default=-1) - args = parser.parse_args() - - ray.init() # init ray with all resources - - # Create shufflling actors - gpu_actors = setup_ray_ucxx_cluster( - GpuShufflingActor, - args.nranks, - args.num_rows, - args.batch_size, - args.total_nparts, - ) - - try: - # run the ShufflingActor.run method remotely - ray.get([actor.run.remote() for actor in gpu_actors]) # type: ignore - - finally: - for actor in gpu_actors: - ray.kill(actor) diff --git a/python/rapidsmpf/rapidsmpf/examples/streaming/__init__.py b/python/rapidsmpf/rapidsmpf/examples/streaming/__init__.py deleted file mode 100644 index abab35f90..000000000 --- a/python/rapidsmpf/rapidsmpf/examples/streaming/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Submodule for streaming examples.""" diff --git a/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py b/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py deleted file mode 100644 index 483a43c73..000000000 --- a/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py +++ /dev/null @@ -1,149 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Basic streaming example.""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import pylibcudf -from cudf_streaming.streaming.table_chunk import TableChunk - -import rmm.mr -from rmm.pylibrmm.stream import DEFAULT_STREAM - -from rapidsmpf.communicator.single import ( - new_communicator as single_process_comm, -) -from rapidsmpf.config import Options, get_environment_variables -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.progress_thread import ProgressThread -from rapidsmpf.streaming.core.actor import ( - define_actor, - run_actor_network, -) -from rapidsmpf.streaming.core.context import Context -from rapidsmpf.streaming.core.leaf_actor import pull_from_channel, push_to_channel -from rapidsmpf.streaming.core.message import Message - -if TYPE_CHECKING: - from collections.abc import Awaitable - - from rapidsmpf.streaming.core.actor import CppActor - from rapidsmpf.streaming.core.channel import Channel - - -def main() -> int: - """Basic example of a streaming graph.""" - # Initialize configuration options from environment variables. - options = Options(get_environment_variables()) - - # Create a communicator and context that will be used by all streaming actors. - comm = single_process_comm(options, ProgressThread()) - ctx = Context( - logger=comm.logger, - br=BufferResource(rmm.mr.get_current_device_resource()), - options=options, - ) - - # Create some pylibcudf tables as input to the streaming graph. - tables = [ - pylibcudf.Table( - [ - pylibcudf.Column.from_iterable_of_py( - [1 * seq, 2 * seq, 3 * seq], - pylibcudf.DataType(pylibcudf.TypeId.INT64), - ) - ] - ) - for seq in range(10) - ] - - # Wrap tables in TableChunk objects before sending them into the graph. - # A TableChunk contains a pylibcudf table, a sequence number, and a CUDA stream. - table_chunks = [ - Message( - seq, - TableChunk.from_pylibcudf_table( - expect, DEFAULT_STREAM, exclusive_view=False, br=ctx.br() - ), - ) - for seq, expect in enumerate(tables) - ] - - # Create input and output channels for table chunks. - ch1: Channel[TableChunk] = ctx.create_channel() - ch2: Channel[TableChunk] = ctx.create_channel() - - # Actor 1: producer that pushes messages into the graph. - # This is a native C++ actor that runs as a coroutine with minimal Python overhead. - actor1: CppActor = push_to_channel(ctx, ch_out=ch1, messages=table_chunks) - - # Actor 2: Python actor that counts the total number of rows. - # Runs as a Python coroutine (asyncio), which comes with overhead, - # but releases the GIL on `await` and when calling into C++ APIs. - @define_actor() - async def count_num_rows( - ctx: Context, ch_in: Channel, ch_out: Channel, total_num_rows: list[int] - ) -> None: - assert len(total_num_rows) == 1, "should be a scalar" - msg: Message[TableChunk] | None - while (msg := await ch_in.recv(ctx)) is not None: - # Convert the message back into a table chunk (releases the message). - table = TableChunk.from_message(msg, br=ctx.br()) - - # Accumulate the number of rows. - total_num_rows[0] += table.table_view().num_rows() - - # The message is now empty since it was released. - assert msg.empty() - - # Wrap the table chunk in a new message. - msg = Message(msg.sequence_number, table) - - # Forward the message to the output channel. - await ch_out.send(ctx, msg) - - # `msg == None` indicates the channel is closed, i.e. we are done. - # Before exiting, drain the output channel to close it gracefully. - await ch_out.drain(ctx) - - # Actors return None, so if we want an "output" value we can use either a closure - # or an output parameter like `total_num_rows`. - total_num_rows = [0] # Wrap scalar in a list to make it mutable in-place. - actor2: Awaitable[None] = count_num_rows( - ctx, ch_in=ch1, ch_out=ch2, total_num_rows=total_num_rows - ) - - # Actor 3: consumer that pulls messages from the graph. - # Like push_to_channel(), it returns a CppActor. It also returns a placeholder - # object that will be populated with the pulled messages after execution. - actor3, out_messages = pull_from_channel(ctx, ch_in=ch2) - - # Run all actors. This blocks until every actor has completed. - run_actor_network( - ctx, - actors=( - actor1, - actor2, - actor3, - ), - ) - - # Collect and verify results. - expect = 0 - for msg in out_messages.release(): - table = TableChunk.from_message(msg, br=ctx.br()).table_view() - expect += table.num_rows() - assert total_num_rows[0] == expect - - # Shut down the context explicitly to ensure it happens on the same thread that - # created it. Alternatively, use `with Context(...) as ctx:` to shut it down - # automatically. - ctx.shutdown() - - return total_num_rows[0] - - -if __name__ == "__main__": - print(f"total_num_rows: {main()}") diff --git a/python/rapidsmpf/rapidsmpf/memory/memory_reservation.pyx b/python/rapidsmpf/rapidsmpf/memory/memory_reservation.pyx index 78c4534fc..cb75e7d9a 100644 --- a/python/rapidsmpf/rapidsmpf/memory/memory_reservation.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/memory_reservation.pyx @@ -105,8 +105,8 @@ def opaque_memory_usage(MemoryReservation reservation not None): This context manager is intended for code paths that use memory outside of RapidsMPF's memory reservation system, for example internal allocations in - libcudf or other third-party libraries. The memory may be of any type covered - by a :class:`MemoryReservation`, most commonly device memory. + third-party libraries. The memory may be of any type covered by a + :class:`MemoryReservation`, most commonly device memory. While the context is active, the provided memory reservation is considered consumed by the enclosed code block. On exit, the reservation is cleared, diff --git a/python/rapidsmpf/rapidsmpf/shuffler.pyx b/python/rapidsmpf/rapidsmpf/shuffler.pyx index 6a9f33b6c..194c37d2c 100644 --- a/python/rapidsmpf/rapidsmpf/shuffler.pyx +++ b/python/rapidsmpf/rapidsmpf/shuffler.pyx @@ -49,8 +49,8 @@ cdef class Shuffler: Notes ----- This class is designed to handle distributed operations by partitioning data - and redistributing it across ranks in a cluster. It is typically used in - distributed data processing workflows involving cuDF tables. + and redistributing it across ranks in a cluster. It operates on caller-provided + packed payloads and is independent of any particular dataframe implementation. The caller promises that inserted buffers are stream-ordered with respect to their own stream, and extracted buffers are likewise guaranteed to be stream- diff --git a/python/rapidsmpf/rapidsmpf/streaming/core/actor.pyx b/python/rapidsmpf/rapidsmpf/streaming/core/actor.pyx index f125fe67c..f97aea2c8 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/core/actor.pyx +++ b/python/rapidsmpf/rapidsmpf/streaming/core/actor.pyx @@ -315,16 +315,13 @@ def run_actor_network(Context ctx not None, *, actors): Examples -------- - >>> ch: Channel[TableChunk] = context.create_channel() + >>> ch: Channel = context.create_channel() >>> cpp_actor, output = pull_from_channel(context, ch_in=ch) ... >>> @define_actor() ... async def python_actor(ctx: Context, ch_out: Channel) -> None: ... # Send one message and close. - ... await ch_out.send( - ... context, - ... Message(42, TableChunk.from_pylibcudf_table(...)) - ... ) + ... await ch_out.send(context, Message(42, payload)) ... await ch_out.drain(context) ... >>> run_actor_network( @@ -332,8 +329,7 @@ def run_actor_network(Context ctx not None, *, actors): ... actors=[cpp_actor, python_actor(context, ch_out=ch)] ... ) >>> results = output.release() - >>> tbl = TableChunk.from_message(results[0]) - >>> tbl.sequence_number + >>> results[0].sequence_number 42 """ diff --git a/python/rapidsmpf/rapidsmpf/testing.py b/python/rapidsmpf/rapidsmpf/testing.py index a38f4b88a..c5a747bee 100644 --- a/python/rapidsmpf/rapidsmpf/testing.py +++ b/python/rapidsmpf/rapidsmpf/testing.py @@ -7,10 +7,8 @@ from typing import TYPE_CHECKING import numpy as np -import pylibcudf import rmm -from rmm.pylibrmm.stream import DEFAULT_STREAM from rapidsmpf.memory.packed_data import PackedData @@ -24,57 +22,6 @@ _DTYPE = np.int64 -def assert_eq( - left: pylibcudf.Table, - right: pylibcudf.Table, - *, - sort_rows: int | None = None, - stream: Stream | None = None, -) -> None: - """ - Assert that two tables are equivalent using pylibcudf. - - Parameters - ---------- - left - plc.Table to compare. - right - plc.Table to compare. - sort_rows - If not None, sort both tables by this column before comparing. - An ``int`` is treated as a column index. - stream - CUDA stream to use for the comparison. - - Raises - ------ - AssertionError - If the two tables do not compare equal. - """ - if stream is None: - stream = DEFAULT_STREAM - - if sort_rows is not None: - column_order = [pylibcudf.types.Order.ASCENDING] - null_precedence = [pylibcudf.types.NullOrder.BEFORE] - left = pylibcudf.sorting.stable_sort_by_key( - left, - pylibcudf.Table([left.columns()[sort_rows]]), - column_order, - null_precedence, - stream=stream, - ) - right = pylibcudf.sorting.stable_sort_by_key( - right, - pylibcudf.Table([right.columns()[sort_rows]]), - column_order, - null_precedence, - stream=stream, - ) - if not pylibcudf.table_equality.tables_equal(left, right, stream=stream): - raise AssertionError(f"Table are not equal with {sort_rows=}") - - def chunk_indices(count: int, num_chunks: int) -> list[tuple[int, int]]: """ Split ``[0, count)`` into ``num_chunks`` contiguous, front-loaded pieces. diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_examples.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_examples.py deleted file mode 100644 index d67a73bce..000000000 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_examples.py +++ /dev/null @@ -1,13 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -import pytest - -cudf = pytest.importorskip("cudf") - -from rapidsmpf.examples.streaming import basic_example # noqa: E402 - - -def test_basic_streaming_example() -> None: - basic_example.main() diff --git a/python/rapidsmpf/rapidsmpf/tests/test_examples.py b/python/rapidsmpf/rapidsmpf/tests/test_examples.py deleted file mode 100644 index f46be682e..000000000 --- a/python/rapidsmpf/rapidsmpf/tests/test_examples.py +++ /dev/null @@ -1,119 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from pathlib import Path -from typing import TYPE_CHECKING - -import numpy as np -import pylibcudf as plc -import pytest - -pytest.importorskip("cudf_streaming") - -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") - -MPI = pytest.importorskip("mpi4py.MPI") -from rapidsmpf.examples.bulk_mpi_shuffle import bulk_mpi_shuffle # noqa: E402 - -if TYPE_CHECKING: - import py.path - - import rmm.mr - - from rapidsmpf.communicator.communicator import Communicator - - -def _write_parquet(table: plc.Table, column_names: list[str], path: str) -> None: - metadata = plc.io.types.TableInputMetadata(table) - for col_meta, name in zip(metadata.column_metadata, column_names, strict=True): - col_meta.set_name(name) - options = ( - plc.io.parquet.ParquetWriterOptions.builder(plc.io.SinkInfo([path]), table) - .metadata(metadata) - .build() - ) - plc.io.parquet.write_parquet(options) - - -def _read_parquet(paths: list[str]) -> plc.Table: - options = plc.io.parquet.ParquetReaderOptions.builder( - plc.io.SourceInfo(paths) - ).build() - return plc.io.parquet.read_parquet(options).tbl - - -@pytest.mark.parametrize("batchsize", [1, 2, 3]) -@pytest.mark.parametrize("num_output_files", [10, 5]) -def test_bulk_shuffle( - comm: Communicator, - tmpdir: py.path.local.LocalPath, - device_mr: rmm.mr.CudaMemoryResource, - batchsize: int, - num_output_files: int, -) -> None: - # Get mpi-compatible tmpdir - mpi_comm = MPI.COMM_WORLD - rank = comm.rank - name = str(tmpdir) if rank == 0 else None - name = mpi_comm.bcast(name, root=0) - mpi_tmpdir = type(tmpdir)(name) - - # Generate input dataset - num_files = 15 - num_rows = 100 - np.random.seed(42) - dataset_dir = mpi_tmpdir.join("dataset") - if rank == 0: - mpi_tmpdir.mkdir("dataset") - for i in range(num_files): - table = plc.Table( - [ - plc.Column.from_iterable_of_py( - range(i * num_rows, (i + 1) * num_rows), - plc.DataType(plc.TypeId.INT64), - ), - plc.Column.from_array(np.random.randint(0, 1000, num_rows)), - plc.Column.from_iterable_of_py( - [i] * num_rows, plc.DataType(plc.TypeId.INT64) - ), - ] - ) - _write_parquet( - table, - ["a", "b", "c"], - str(dataset_dir.join(f"part.{i}.parquet")), - ) - mpi_tmpdir.mkdir("output") - input_paths = sorted(map(str, Path(dataset_dir).glob("**/*"))) - else: - input_paths = None - input_paths = mpi_comm.bcast(input_paths, root=0) - assert isinstance(input_paths, list) # for mypy - output_dir = str(mpi_tmpdir.join("output")) - - # Use a default buffer resource. - br = BufferResource(device_mr) - - # Perform a the shuffle - bulk_mpi_shuffle( - paths=input_paths, - shuffle_on=["b"], - output_path=output_dir, - comm=comm, - br=br, - batchsize=batchsize, - num_output_files=num_output_files, - ) - mpi_comm.barrier() - - # Check that original and shuffled data match - if rank == 0: - shuffled_paths = sorted(map(str, Path(output_dir).glob("**/*"))) - df_original = _read_parquet(input_paths) - df_shuffled = _read_parquet(shuffled_paths) - assert_eq(df_original, df_shuffled, sort_rows=0) - mpi_comm.barrier() diff --git a/python/rapidsmpf/rapidsmpf/tests/test_partition.py b/python/rapidsmpf/rapidsmpf/tests/test_partition.py deleted file mode 100644 index 38ce185bc..000000000 --- a/python/rapidsmpf/rapidsmpf/tests/test_partition.py +++ /dev/null @@ -1,107 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from typing import TYPE_CHECKING - -import numpy as np -import pylibcudf as plc -import pytest - -pytest.importorskip("cudf_streaming") -from cudf_streaming.integrations.partition import ( - partition_and_pack, - split_and_pack, - unpack_and_concat, -) - -from rmm.pylibrmm.stream import DEFAULT_STREAM - -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.testing import assert_eq - -cudf = pytest.importorskip("cudf") - -if TYPE_CHECKING: - import rmm.mr - - -def _make_table(cols: list[list[int]]) -> plc.Table: - # Assigns empty column inputs as int64 - return plc.Table( - [ - plc.Column.from_iterable_of_py(col, plc.DataType(plc.TypeId.INT64)) - for col in cols - ] - ) - - -@pytest.mark.parametrize("cols", [[[1, 2, 3], [2, 2, 1]], [[], []]]) -@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) -def test_partition_and_pack_unpack( - device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int -) -> None: - br = BufferResource(device_mr) - expect = _make_table(cols) - partitions = partition_and_pack( - expect, - columns_to_hash=(1,), - num_partitions=num_partitions, - br=br, - stream=DEFAULT_STREAM, - ) - got = unpack_and_concat( - tuple(partitions.values()), - br=br, - stream=DEFAULT_STREAM, - ) - # Since the row order isn't preserved, we sort the rows by the first column. - assert_eq(expect, got, sort_rows=0) - - -@pytest.mark.parametrize( - "cols", - [ - [[1, 2, 3], [2, 2, 1]], - [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]], - [[], []], - ], -) -@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) -def test_split_and_pack_unpack( - device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int -) -> None: - br = BufferResource(device_mr) - expect = _make_table(cols) - splits = np.linspace(0, expect.num_rows(), num_partitions, endpoint=False)[ - 1: - ].astype(int) - partitions = split_and_pack( - expect, - splits=splits, - br=br, - stream=DEFAULT_STREAM, - ) - got = unpack_and_concat( - tuple(partitions[i] for i in range(num_partitions)), - br=br, - stream=DEFAULT_STREAM, - ) - - assert_eq(expect, got) - - -@pytest.mark.parametrize("cols", [[[1, 2, 3], [2, 2, 1]], [[], []]]) -@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) -def test_split_and_pack_unpack_out_of_range( - device_mr: rmm.mr.CudaMemoryResource, cols: list[list[int]], num_partitions: int -) -> None: - br = BufferResource(device_mr) - expect = _make_table(cols) - with pytest.raises(IndexError): - split_and_pack( - expect, - splits=[100], - br=br, - stream=DEFAULT_STREAM, - ) diff --git a/python/rapidsmpf/rapidsmpf/tests/test_ray.py b/python/rapidsmpf/rapidsmpf/tests/test_ray.py index cbdcc4cee..abe25789c 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_ray.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_ray.py @@ -2,26 +2,19 @@ # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -import functools import os from typing import TYPE_CHECKING -from cuda.core import system - os.environ["RAY_DEDUP_LOGS"] = "0" os.environ["RAY_IGNORE_UNHANDLED_ERRORS"] = "1" import pytest ray = pytest.importorskip("ray") -cudf = pytest.importorskip("cudf") if TYPE_CHECKING: from collections.abc import Generator -from rapidsmpf.examples.ray.ray_shuffle_example import ( # noqa: E402 - ShufflingActor, -) from rapidsmpf.integrations.ray import ( # noqa: E402 RapidsMPFActor, setup_ray_ucxx_cluster, @@ -126,37 +119,3 @@ class NonRapidsMPFActor: ... with pytest.raises(TypeError): setup_ray_ucxx_cluster(NonRapidsMPFActor, 1) - - -@functools.cache -def get_gpu_count() -> int: - return system.get_num_devices() # type: ignore - - -@pytest.mark.parametrize("num_workers", [1, 4]) -@pytest.mark.parametrize("batch_size", [-1, 10]) -@pytest.mark.parametrize("total_num_partitions", [1, 10]) -def test_ray_shuffle_actor( - ray_cluster: None, num_workers: int, batch_size: int, total_num_partitions: int -) -> None: - gpu_count = get_gpu_count() - - # Test shuffling actor that uses 1/num_workers fractional GPUs if - # gpu_count < num_workers or 1 GPU otherwise - @ray.remote(num_gpus=(gpu_count / num_workers) if gpu_count < num_workers else 1) - class TestShufflingActor(ShufflingActor): ... - - # setup the UCXX cluster using TestShufflingActor - gpu_actors = setup_ray_ucxx_cluster( - TestShufflingActor, - num_workers, - batch_size=batch_size, - total_nparts=total_num_partitions, - ) - - try: - # call run on all actors remotely - ray.get([actor.run.remote() for actor in gpu_actors]) - finally: - for actor in gpu_actors: - ray.kill(actor) From 82556c3880d81589eebf58b00c35f511cc5b2777 Mon Sep 17 00:00:00 2001 From: Peter Andreas Entschev Date: Thu, 11 Jun 2026 14:23:55 -0700 Subject: [PATCH 14/14] Add link back --- cpp/src/memory/host_memory_resource.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/cpp/src/memory/host_memory_resource.cpp b/cpp/src/memory/host_memory_resource.cpp index 8b285cbba..a7b024be2 100644 --- a/cpp/src/memory/host_memory_resource.cpp +++ b/cpp/src/memory/host_memory_resource.cpp @@ -20,6 +20,7 @@ namespace { * Attempts to mark the specified memory region as eligible for Transparent Huge Pages * (THP) using `madvise(MADV_HUGEPAGE)`. This is a best-effort optimization that can * improve device to host memory transfer performance for sufficiently large buffers. + * See . * * @param ptr Pointer to the start of the memory region. * @param size Size of the region in bytes.