diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 1a345e4a498f..0a0b76185416 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -713,6 +713,7 @@ jobs: matrix_filter: map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))])) build_type: pull-request script: ci/test_wheel_dask_cudf.sh + # build cudf with benchmarks but disable for cudf_streaming since we can't use MPI/UCX in wheel-based builds devcontainer: permissions: actions: read @@ -734,7 +735,9 @@ jobs: SCCACHE_DIST_FALLBACK_TO_LOCAL_COMPILE=false build_command: | sccache --zero-stats; - build-all -j0 -DBUILD_BENCHMARKS=ON --verbose 2>&1 | tee telemetry-artifacts/build.log; + configure-cudf-cpp -DBUILD_BENCHMARKS=ON; + configure-cudf_streaming-cpp -DBUILD_BENCHMARKS=OFF; + build-all -j0 --verbose 2>&1 | tee telemetry-artifacts/build.log; sccache --show-adv-stats | tee telemetry-artifacts/sccache-stats.txt; unit-tests-cudf-pandas: needs: [wheel-build-cudf, changed-files] diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 7e0ad88ddb96..f138c0896bb6 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -17,6 +17,7 @@ dependencies: - cmake>=4.0 - cramjam - cuda-cudart-dev +- cuda-cupti-dev - cuda-nvcc - cuda-nvrtc-dev - cuda-nvtx-dev @@ -47,6 +48,7 @@ dependencies: - librapidsmpf==26.8.*,>=0.0.0a0 - librdkafka>=2.11.1,<2.12.0 - librmm==26.8.*,>=0.0.0a0 +- libucxx==0.51.*,>=0.0.0a0 - make - mmh3 - moto>=4.0.8 @@ -105,6 +107,7 @@ dependencies: - structlog - sysroot_linux-aarch64==2.28 - typing_extensions>=4.0.0 +- ucxx==0.51.*,>=0.0.0a0 - xlsxwriter - zlib>=1.2.13 - zstandard diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 272914ad7635..386c654557cd 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -17,6 +17,7 @@ dependencies: - cmake>=4.0 - cramjam - cuda-cudart-dev +- cuda-cupti-dev - cuda-nvcc - cuda-nvrtc-dev - cuda-nvtx-dev @@ -47,6 +48,7 @@ dependencies: - librapidsmpf==26.8.*,>=0.0.0a0 - librdkafka>=2.11.1,<2.12.0 - librmm==26.8.*,>=0.0.0a0 +- libucxx==0.51.*,>=0.0.0a0 - make - mmh3 - moto>=4.0.8 @@ -105,6 +107,7 @@ dependencies: - structlog - sysroot_linux-64==2.28 - typing_extensions>=4.0.0 +- ucxx==0.51.*,>=0.0.0a0 - xlsxwriter - zlib>=1.2.13 - zstandard diff --git a/conda/environments/all_cuda-132_arch-aarch64.yaml b/conda/environments/all_cuda-132_arch-aarch64.yaml index 9ecba408142c..f639a44bbcc9 100644 --- a/conda/environments/all_cuda-132_arch-aarch64.yaml +++ b/conda/environments/all_cuda-132_arch-aarch64.yaml @@ -17,6 +17,7 @@ dependencies: - cmake>=4.0 - cramjam - cuda-cudart-dev +- cuda-cupti-dev - cuda-nvcc - cuda-nvrtc-dev - cuda-nvtx-dev @@ -47,6 +48,7 @@ dependencies: - librapidsmpf==26.8.*,>=0.0.0a0 - librdkafka>=2.11.1,<2.12.0 - librmm==26.8.*,>=0.0.0a0 +- libucxx==0.51.*,>=0.0.0a0 - make - mmh3 - moto>=4.0.8 @@ -105,6 +107,7 @@ dependencies: - structlog - sysroot_linux-aarch64==2.28 - typing_extensions>=4.0.0 +- ucxx==0.51.*,>=0.0.0a0 - xlsxwriter - zlib>=1.2.13 - zstandard diff --git a/conda/environments/all_cuda-132_arch-x86_64.yaml b/conda/environments/all_cuda-132_arch-x86_64.yaml index ced21db00730..83ff0e3ee92c 100644 --- a/conda/environments/all_cuda-132_arch-x86_64.yaml +++ b/conda/environments/all_cuda-132_arch-x86_64.yaml @@ -17,6 +17,7 @@ dependencies: - cmake>=4.0 - cramjam - cuda-cudart-dev +- cuda-cupti-dev - cuda-nvcc - cuda-nvrtc-dev - cuda-nvtx-dev @@ -47,6 +48,7 @@ dependencies: - librapidsmpf==26.8.*,>=0.0.0a0 - librdkafka>=2.11.1,<2.12.0 - librmm==26.8.*,>=0.0.0a0 +- libucxx==0.51.*,>=0.0.0a0 - make - mmh3 - moto>=4.0.8 @@ -105,6 +107,7 @@ dependencies: - structlog - sysroot_linux-64==2.28 - typing_extensions>=4.0.0 +- ucxx==0.51.*,>=0.0.0a0 - xlsxwriter - zlib>=1.2.13 - zstandard diff --git a/cpp/libcudf_streaming/CMakeLists.txt b/cpp/libcudf_streaming/CMakeLists.txt index 385e691d6359..aedcae4e3fdf 100644 --- a/cpp/libcudf_streaming/CMakeLists.txt +++ b/cpp/libcudf_streaming/CMakeLists.txt @@ -36,11 +36,25 @@ include(${CMAKE_CURRENT_SOURCE_DIR}/../cmake/Modules/ConfigureCUDA.cmake) # * conda environment ----------------------------------------------------------------------------- rapids_cmake_support_conda_env(conda_env MODIFY_PREFIX_PATH) +# ################################################################################################## +# * find C++ libraries from RAPIDS wheels in site-packages ----------------------------------------- +find_package(Python COMPONENTS Interpreter) +if(Python_FOUND) + include(${rapids-cmake-dir}/cython-core/find_prefix_paths.cmake) + rapids_cython_find_prefix_paths("${Python_EXECUTABLE}" wheel_prefix_paths) + set_property(GLOBAL PROPERTY FIND_LIBRARY_USE_LIB64_PATHS TRUE) + list(APPEND CMAKE_PREFIX_PATH ${wheel_prefix_paths}) +endif() + # ################################################################################################## # * Build options option(BUILD_TESTS "Build tests for libcudf_streaming" OFF) +option(BUILD_BENCHMARKS "Build benchmarks for libcudf_streaming" OFF) +option(BUILD_EXAMPLES "Build examples for libcudf_streaming" OFF) message(VERBOSE "CUDF_STREAMING: Build gtests: ${BUILD_TESTS}") +message(VERBOSE "CUDF_STREAMING: Build benchmarks: ${BUILD_BENCHMARKS}") +message(VERBOSE "CUDF_STREAMING: Build examples: ${BUILD_EXAMPLES}") # ################################################################################################## # * Dependencies @@ -51,6 +65,11 @@ include(cmake/thirdparty/get_cudf.cmake) include(cmake/thirdparty/get_rapidsmpf.cmake) include(../cmake/thirdparty/get_cucollections.cmake) +if(BUILD_BENCHMARKS) + rapids_find_package(ucxx REQUIRED) + rapids_find_package(MPI REQUIRED) +endif() + # ################################################################################################## # * library target -------------------------------------------------------------------------------- add_library( @@ -115,6 +134,16 @@ if(BUILD_TESTS) add_subdirectory(tests) endif() +# Benchmarks if enabled +if(BUILD_BENCHMARKS) + add_subdirectory(benchmarks) +endif() + +# Examples if enabled +if(BUILD_EXAMPLES) + add_subdirectory(examples) +endif() + # ################################################################################################## # * cudf_streaming Install ------------------------------------------------------------------------- rapids_cmake_install_lib_dir(lib_dir) diff --git a/cpp/libcudf_streaming/benchmarks/CMakeLists.txt b/cpp/libcudf_streaming/benchmarks/CMakeLists.txt new file mode 100644 index 000000000000..3aac8d99527b --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/CMakeLists.txt @@ -0,0 +1,80 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +# Find or install GoogleBench +include(${rapids-cmake-dir}/cpm/gbench.cmake) +rapids_cpm_gbench(BUILD_STATIC) + +add_library(bench_utils INTERFACE) +target_sources(bench_utils INTERFACE utils/random_data.cu) +target_compile_options(bench_utils INTERFACE $<$:--expt-extended-lambda>) + +add_executable(bench_shuffle "bench_shuffle.cpp") +set_target_properties( + bench_shuffle + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CUDF_STREAMING_BINARY_DIR}/benchmarks" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON +) +target_link_libraries( + bench_shuffle PRIVATE cudf_streaming rapidsmpf::rapidsmpf ucxx::ucxx MPI::MPI_CXX + $ bench_utils +) +install( + TARGETS bench_shuffle + COMPONENT benchmarking + DESTINATION bin/benchmarks/libcudf_streaming + EXCLUDE_FROM_ALL +) + +add_executable(bench_partition "bench_partition.cpp") +set_target_properties( + bench_partition + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CUDF_STREAMING_BINARY_DIR}/benchmarks" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON +) +target_link_libraries( + bench_partition PRIVATE cudf_streaming rapidsmpf::rapidsmpf benchmark::benchmark + benchmark::benchmark_main $ bench_utils +) +install( + TARGETS bench_partition + COMPONENT benchmarking + DESTINATION bin/benchmarks/libcudf_streaming + EXCLUDE_FROM_ALL +) + +add_executable(bench_pack "bench_pack.cpp") +set_target_properties( + bench_pack + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CUDF_STREAMING_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_link_libraries( + bench_pack PRIVATE cudf_streaming rapidsmpf::rapidsmpf benchmark::benchmark + benchmark::benchmark_main $ bench_utils +) +install( + TARGETS bench_pack + COMPONENT benchmarking + DESTINATION bin/benchmarks/libcudf_streaming + EXCLUDE_FROM_ALL +) + +add_subdirectory(streaming) diff --git a/cpp/libcudf_streaming/benchmarks/bench_pack.cpp b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp new file mode 100644 index 000000000000..e60510269bfd --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/bench_pack.cpp @@ -0,0 +1,305 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "utils/random_data.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include + +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/libcudf_streaming/benchmarks/bench_partition.cpp b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp new file mode 100644 index 000000000000..42e4a129ac29 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/bench_partition.cpp @@ -0,0 +1,169 @@ +/** + * 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 + +// 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; + CUDF_CUDA_TRY(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); + CUDF_CUDA_TRY(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; + CUDF_CUDA_TRY(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); + } + CUDF_CUDA_TRY(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/libcudf_streaming/benchmarks/bench_shuffle.cpp b/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp new file mode 100644 index 000000000000..4a367a401194 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp @@ -0,0 +1,610 @@ +/** + * 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 + +#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( + cudf_streaming::integrations::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 >> +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/libcudf_streaming/benchmarks/streaming/CMakeLists.txt b/cpp/libcudf_streaming/benchmarks/streaming/CMakeLists.txt new file mode 100644 index 000000000000..166f835bbf2b --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/CMakeLists.txt @@ -0,0 +1,28 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# 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 "${CUDF_STREAMING_BINARY_DIR}/benchmarks" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON +) +target_link_libraries( + bench_streaming_shuffle PRIVATE cudf_streaming rapidsmpf::rapidsmpf ucxx::ucxx MPI::MPI_CXX + $ bench_utils +) +install( + TARGETS bench_streaming_shuffle + COMPONENT benchmarking + DESTINATION bin/benchmarks/libcudf_streaming + EXCLUDE_FROM_ALL +) + +add_subdirectory(ndsh) diff --git a/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp new file mode 100644 index 000000000000..0ab0bebcff98 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -0,0 +1,390 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "../utils/misc.hpp" +#include "../utils/rmm_utils.hpp" +#include "data_generator.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 + +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, 1); 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 { + auto ucxx = std::dynamic_pointer_cast(comm); + if (ucxx == nullptr) { + log.print("Expected UCXX communicator when using bootstrap mode"); + throw std::runtime_error{"Expected UCXX communicator when using bootstrap mode"}; + } + ucxx->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/libcudf_streaming/benchmarks/streaming/data_generator.hpp b/cpp/libcudf_streaming/benchmarks/streaming/data_generator.hpp new file mode 100644 index 000000000000..bde729bb931b --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/data_generator.hpp @@ -0,0 +1,71 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include "../utils/random_data.hpp" + +#include +#include + +#include +#include +#include +#include +#include + +namespace rapidsmpf::streaming::actor { + +using cudf_streaming::streaming::TableChunk; + +/** + * @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( + 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/libcudf_streaming/benchmarks/streaming/ndsh/CMakeLists.txt b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/CMakeLists.txt new file mode 100644 index 000000000000..629c86bf99b7 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/CMakeLists.txt @@ -0,0 +1,58 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +add_library( + cudf_streaming_ndsh concatenate.cpp groupby.cpp join.cpp parquet_writer.cpp sort.cpp utils.cpp +) + +set_target_properties( + cudf_streaming_ndsh + 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_link_libraries( + cudf_streaming_ndsh PRIVATE cudf_streaming rapidsmpf::rapidsmpf cuco::cuco MPI::MPI_CXX + $ +) + +set(CUDF_STREAMING_NDSH_QUERIES q01 q03 q04 q09 q21 bench_read) + +foreach(query IN ITEMS ${CUDF_STREAMING_NDSH_QUERIES}) + add_executable(${query} "${query}.cpp") + set_target_properties( + ${query} + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CUDF_STREAMING_BINARY_DIR}/benchmarks/ndsh" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + ) + target_link_libraries( + ${query} PRIVATE cudf_streaming_ndsh cudf_streaming rapidsmpf::rapidsmpf MPI::MPI_CXX + $ + ) +endforeach() + +install( + TARGETS cudf_streaming_ndsh + COMPONENT benchmarking + DESTINATION ${lib_dir} + EXCLUDE_FROM_ALL +) +install( + TARGETS ${CUDF_STREAMING_NDSH_QUERIES} + COMPONENT benchmarking + DESTINATION bin/benchmarks/libcudf_streaming + EXCLUDE_FROM_ALL +) diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cpp new file mode 100644 index 000000000000..ff2f482e0767 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/bench_read.cpp @@ -0,0 +1,390 @@ +/** + * 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 + +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{}; + CUDF_CUDA_TRY(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/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp new file mode 100644 index 000000000000..036eb81a581c --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.cpp @@ -0,0 +1,73 @@ +/** + * 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/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.hpp new file mode 100644 index 000000000000..cff4510947e6 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/concatenate.hpp @@ -0,0 +1,36 @@ +/** + * 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/libcudf_streaming/benchmarks/streaming/ndsh/groupby.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.cpp new file mode 100644 index 000000000000..71eedc24e0fb --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.cpp @@ -0,0 +1,62 @@ +/** + * 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/libcudf_streaming/benchmarks/streaming/ndsh/groupby.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.hpp new file mode 100644 index 000000000000..33c51d512e72 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/groupby.hpp @@ -0,0 +1,50 @@ +/** + * 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/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp new file mode 100644 index 000000000000..7a7acf961d05 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.cpp @@ -0,0 +1,517 @@ +/** + * 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 + +namespace rapidsmpf::ndsh { +using cudf_streaming::integrations::partition_and_pack; +using cudf_streaming::integrations::partition_and_split; +using cudf_streaming::integrations::unpack_and_concat; +using cudf_streaming::integrations::unspill_partitions; +using cudf_streaming::streaming::TableChunk; +using cudf_streaming::streaming::to_message; + +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 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 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 to_message(0, + std::make_unique( + std::make_unique(std::move(result[0])))); + } else { + auto stream = ctx->br()->stream_pool().get_stream(); + co_return to_message( + 0, + std::make_unique( + unpack_and_concat( + 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 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 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 = 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(to_message( + pid, + std::make_unique( + unpack_and_concat( + 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/libcudf_streaming/benchmarks/streaming/ndsh/join.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.hpp new file mode 100644 index 000000000000..0ed96ca9ea33 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/join.hpp @@ -0,0 +1,209 @@ +/** + * 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/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp new file mode 100644 index 000000000000..8ddf81aa20a0 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.cpp @@ -0,0 +1,59 @@ +/** + * 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/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hpp new file mode 100644 index 000000000000..abfcbec74b9a --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/parquet_writer.hpp @@ -0,0 +1,33 @@ +/** + * 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/libcudf_streaming/benchmarks/streaming/ndsh/q01.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cpp new file mode 100644 index 000000000000..3a498ff105ff --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q01.cpp @@ -0,0 +1,385 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "concatenate.hpp" +#include "groupby.hpp" +#include "join.hpp" +#include "parquet_writer.hpp" +#include "sort.hpp" +#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 + +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{}; + CUDF_CUDA_TRY(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/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp new file mode 100644 index 000000000000..474ce717940c --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q03.cpp @@ -0,0 +1,547 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "concatenate.hpp" +#include "groupby.hpp" +#include "join.hpp" +#include "parquet_writer.hpp" +#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 +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +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{}; + CUDF_CUDA_TRY(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::streaming::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/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp new file mode 100644 index 000000000000..379e1f283166 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q04.cpp @@ -0,0 +1,435 @@ +/** + + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "concatenate.hpp" +#include "groupby.hpp" +#include "join.hpp" +#include "parquet_writer.hpp" +#include "sort.hpp" +#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 +#include +#include +#include +#include +#include + +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) +{ + rapidsmpf::ndsh::FinalizeMPI finalize{}; + CUDF_CUDA_TRY(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::streaming::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/libcudf_streaming/benchmarks/streaming/ndsh/q09.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q09.cpp new file mode 100644 index 000000000000..29e341f9a6cf --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q09.cpp @@ -0,0 +1,562 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "concatenate.hpp" +#include "groupby.hpp" +#include "join.hpp" +#include "parquet_writer.hpp" +#include "sort.hpp" +#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 +#include +#include + +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{}; + CUDF_CUDA_TRY(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/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp new file mode 100644 index 000000000000..283ce5fc5291 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/q21.cpp @@ -0,0 +1,786 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "concatenate.hpp" +#include "groupby.hpp" +#include "join.hpp" +#include "parquet_writer.hpp" +#include "sort.hpp" +#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 +#include +#include +#include +#include +#include + +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}; + bool released_lineitem_read = false; + + 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()); + 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()))); + if (!released_lineitem_read) { + latch->count_down(); + released_lineitem_read = true; + } + } + if (!released_lineitem_read) { latch->count_down(); } + 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{}; + CUDF_CUDA_TRY(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::streaming::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/libcudf_streaming/benchmarks/streaming/ndsh/sort.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.cpp new file mode 100644 index 000000000000..a1b3caad8126 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.cpp @@ -0,0 +1,56 @@ +/** + * 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 { +using cudf_streaming::streaming::TableChunk; +using cudf_streaming::streaming::to_message; + +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/libcudf_streaming/benchmarks/streaming/ndsh/sort.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.hpp new file mode 100644 index 000000000000..b042199106cc --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sort.hpp @@ -0,0 +1,38 @@ +/** + * 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/libcudf_streaming/benchmarks/streaming/ndsh/sql/q01.sql b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q01.sql new file mode 100644 index 000000000000..8b769a1301f6 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q01.sql @@ -0,0 +1,21 @@ +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/libcudf_streaming/benchmarks/streaming/ndsh/sql/q03.sql b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q03.sql new file mode 100644 index 000000000000..c93865c19bab --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q03.sql @@ -0,0 +1,23 @@ +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/libcudf_streaming/benchmarks/streaming/ndsh/sql/q04.sql b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q04.sql new file mode 100644 index 000000000000..84a9733ed71c --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q04.sql @@ -0,0 +1,21 @@ +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/libcudf_streaming/benchmarks/streaming/ndsh/sql/q09.sql b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q09.sql new file mode 100644 index 000000000000..7b6a0d609734 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q09.sql @@ -0,0 +1,32 @@ +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/libcudf_streaming/benchmarks/streaming/ndsh/sql/q17.sql b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q17.sql new file mode 100644 index 000000000000..b553d52ac2d9 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q17.sql @@ -0,0 +1,17 @@ +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 + lineitem + where + l_partkey = p_partkey + ) diff --git a/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q18.sql b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q18.sql new file mode 100644 index 000000000000..bf61e83cdbfa --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q18.sql @@ -0,0 +1,33 @@ +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/libcudf_streaming/benchmarks/streaming/ndsh/sql/q21.sql b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q21.sql new file mode 100644 index 000000000000..eae182761a6e --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/sql/q21.sql @@ -0,0 +1,40 @@ +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/libcudf_streaming/benchmarks/streaming/ndsh/utils.cpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cpp new file mode 100644 index 000000000000..086e20e647e8 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.cpp @@ -0,0 +1,373 @@ +/** + * 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/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp new file mode 100644 index 000000000000..cf9148f8adc4 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/streaming/ndsh/utils.hpp @@ -0,0 +1,267 @@ +/** + * 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 +#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/libcudf_streaming/benchmarks/utils/misc.hpp b/cpp/libcudf_streaming/benchmarks/utils/misc.hpp new file mode 100644 index 000000000000..5504a12d79cb --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/utils/misc.hpp @@ -0,0 +1,71 @@ +/** + * 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 + +/** + * @brief Computes the harmonic mean of a set of values. + * + * The harmonic mean is defined as: + * \f[ + * \text{HM} = \frac{n}{\sum_{i=1}^n \frac{1}{x_i}} + * \f] + * where \f$n\f$ is the number of values, and \f$x_i\f$ are the individual values. + * + * @param values A vector of double values for which the harmonic mean is to be computed. + * All values must be non-zero, and the vector must not be empty. + * @return The harmonic mean as a `double`. + * + * @throws std::invalid_argument If the input vector is empty. + * @throws std::domain_error If any value in the input vector is zero. + */ +double harmonic_mean(std::vector const& values) +{ + if (values.empty()) { throw std::invalid_argument("The input vector is empty."); } + double sum = 0.0; + for (double value : values) { + if (value == 0.0) { throw std::domain_error("Cannot compute harmonic mean with zero values."); } + sum += 1.0 / value; + } + return static_cast(values.size()) / sum; +} + +/** + * @brief Parses a string into an integer and ensures it is within the specified range. + * + * This function converts a string representation of an integer into a value of type T. + * It performs range validation using the provided minimum and maximum values. + * + * @tparam T The target integer type. + * @param output Reference to the output variable where the parsed integer will be stored. + * @param str The input string representing the integer. + * @param min_val The minimum allowed value. + * @param max_val The maximum allowed value. + * @throws std::invalid_argument If the string is not a valid integer. + * @throws std::out_of_range If the parsed value is outside the specified range. + */ +template +void parse_integer(T& output, + std::string const& str, + std::int64_t min_val = 0, + std::int64_t max_val = std::numeric_limits::max()) +{ + long long val; + try { + val = std::stoll(str); + } catch (std::invalid_argument const&) { + RAPIDSMPF_FAIL("cannot parse \"" + str + "\"", std::invalid_argument); + } catch (std::out_of_range const&) { + RAPIDSMPF_FAIL("\"" + str + "\" is out of range", std::out_of_range); + } + RAPIDSMPF_EXPECTS(min_val <= val && val <= max_val, "\"" + str + "\" is out of range"); + output = static_cast(val); +} diff --git a/cpp/libcudf_streaming/benchmarks/utils/random_data.cu b/cpp/libcudf_streaming/benchmarks/utils/random_data.cu new file mode 100644 index 000000000000..3b79722af8d8 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/utils/random_data.cu @@ -0,0 +1,95 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "random_data.hpp" + +#include + +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include +#include +#include + +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_nosync(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/libcudf_streaming/benchmarks/utils/random_data.hpp b/cpp/libcudf_streaming/benchmarks/utils/random_data.hpp new file mode 100644 index 000000000000..b6d5df6abdc8 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/utils/random_data.hpp @@ -0,0 +1,109 @@ +/** + * 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/libcudf_streaming/benchmarks/utils/rmm_utils.hpp b/cpp/libcudf_streaming/benchmarks/utils/rmm_utils.hpp new file mode 100644 index 000000000000..ae1d1b7cd121 --- /dev/null +++ b/cpp/libcudf_streaming/benchmarks/utils/rmm_utils.hpp @@ -0,0 +1,56 @@ +/** + * 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 + +#include + +/** + * @brief Create and set a RMM memory resource as the current device resource. + * + * @param name The name of the resource: + * - `cuda`: use the default CUDA memory resource. + * - `async`: use a CUDA async memory resource. + * - `pool`: use a memory pool backed by a CUDA memory resource. + * - `managed`: use a CUDA managed memory resource. + */ +inline void set_current_rmm_resource(std::string const& name) +{ + if (name == "cuda") { + rmm::mr::set_current_device_resource(rmm::mr::cuda_memory_resource{}); + } else if (name == "async") { + rmm::mr::set_current_device_resource(rmm::mr::cuda_async_memory_resource{}); + } else if (name == "managed") { + rmm::mr::set_current_device_resource(rmm::mr::managed_memory_resource{}); + } else if (name == "pool") { + rmm::mr::set_current_device_resource( + rmm::mr::pool_memory_resource{rmm::mr::cuda_memory_resource{}, + rmm::percent_of_free_device_memory(80), + rmm::percent_of_free_device_memory(80)}); + } else { + RAPIDSMPF_FAIL("unknown RMM resource name: " + name); + } +} + +/** + * @brief Create a statistics-enabled device memory resource wrapping the current + * device resource, and set it as the current device resource. + * + * @return A RmmResourceAdaptor (shared ownership) for accessing statistics. + */ +[[nodiscard]] inline rapidsmpf::RmmResourceAdaptor set_device_mem_resource_with_stats() +{ + rapidsmpf::RmmResourceAdaptor adaptor{rmm::mr::get_current_device_resource_ref()}; + rmm::mr::set_current_device_resource(adaptor); + return adaptor; +} diff --git a/cpp/libcudf_streaming/cmake/thirdparty/get_rapidsmpf.cmake b/cpp/libcudf_streaming/cmake/thirdparty/get_rapidsmpf.cmake index 4d96e113970e..37a25b126e13 100644 --- a/cpp/libcudf_streaming/cmake/thirdparty/get_rapidsmpf.cmake +++ b/cpp/libcudf_streaming/cmake/thirdparty/get_rapidsmpf.cmake @@ -8,6 +8,10 @@ # This function finds rapidsmpf and sets any additional necessary environment variables. function(find_and_configure_rapidsmpf VERSION) rapids_cmake_parse_version(MAJOR_MINOR ${VERSION} major_minor) + set(rapidsmpf_build_comm_support OFF) + if(BUILD_BENCHMARKS) + set(rapidsmpf_build_comm_support ON) + endif() rapids_cpm_find( rapidsmpf ${VERSION} BUILD_EXPORT_SET cudf_streaming-exports @@ -16,8 +20,12 @@ function(find_and_configure_rapidsmpf VERSION) GIT_REPOSITORY https://github.com/rapidsai/rapidsmpf.git GIT_TAG "${RAPIDS_BRANCH}" GIT_SHALLOW TRUE SOURCE_SUBDIR cpp - OPTIONS "BUILD_MPI_SUPPORT OFF" "BUILD_UCXX_SUPPORT OFF" "BUILD_SLURM_SUPPORT OFF" - "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_EXAMPLES OFF" + OPTIONS "BUILD_MPI_SUPPORT ${rapidsmpf_build_comm_support}" + "BUILD_UCXX_SUPPORT ${rapidsmpf_build_comm_support}" + "BUILD_SLURM_SUPPORT OFF" + "BUILD_TESTS OFF" + "BUILD_BENCHMARKS OFF" + "BUILD_EXAMPLES OFF" ) endfunction() diff --git a/cpp/libcudf_streaming/examples/CMakeLists.txt b/cpp/libcudf_streaming/examples/CMakeLists.txt new file mode 100644 index 000000000000..5d8a2ac7b053 --- /dev/null +++ b/cpp/libcudf_streaming/examples/CMakeLists.txt @@ -0,0 +1,29 @@ +# ============================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ============================================================================= + +add_executable(example_shuffle "example_shuffle.cpp") +set_target_properties( + example_shuffle + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CUDF_STREAMING_BINARY_DIR}/examples" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON +) +target_link_libraries( + example_shuffle PRIVATE cudf_streaming rapidsmpf::rapidsmpf $ + $ +) +target_sources( + example_shuffle PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/../benchmarks/utils/random_data.cu" +) +install( + TARGETS example_shuffle + COMPONENT examples + DESTINATION bin/examples/libcudf_streaming + EXCLUDE_FROM_ALL +) diff --git a/cpp/libcudf_streaming/examples/example_shuffle.cpp b/cpp/libcudf_streaming/examples/example_shuffle.cpp new file mode 100644 index 000000000000..f053ea820f21 --- /dev/null +++ b/cpp/libcudf_streaming/examples/example_shuffle.cpp @@ -0,0 +1,127 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "../benchmarks/utils/random_data.hpp" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +// 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( + cudf_streaming::integrations::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/libcudf_streaming/include/cudf_streaming/streaming/bloom_filter.hpp b/cpp/libcudf_streaming/include/cudf_streaming/streaming/bloom_filter.hpp index e08661e83d28..f0348fb6c05f 100644 --- a/cpp/libcudf_streaming/include/cudf_streaming/streaming/bloom_filter.hpp +++ b/cpp/libcudf_streaming/include/cudf_streaming/streaming/bloom_filter.hpp @@ -99,6 +99,18 @@ struct BloomFilter { std::shared_ptr ch_out, std::vector keys); + /** + * @brief Compute number of filter blocks that fit in the given L2 cache size. + * + * @param l2size L2 cache size in bytes. + * @return Number of filter blocks that fit. + */ + [[nodiscard]] static std::size_t fitting_num_blocks(std::size_t l2size) noexcept + { + using StorageType = std::uint32_t; + return (l2size * 2) / (3 * sizeof(StorageType)); + } + private: std::shared_ptr ctx_{}; std::shared_ptr comm_{}; diff --git a/cpp/libcudf_streaming/scripts/ndsh.py b/cpp/libcudf_streaming/scripts/ndsh.py new file mode 100755 index 000000000000..8a7c69f78e22 --- /dev/null +++ b/cpp/libcudf_streaming/scripts/ndsh.py @@ -0,0 +1,910 @@ +#!/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 +r""" +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 hashlib +import re +import subprocess +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 + continue + + 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/libcudf_streaming/tests/CMakeLists.txt b/cpp/libcudf_streaming/tests/CMakeLists.txt index 600dbd51f74a..5b6cf230d699 100644 --- a/cpp/libcudf_streaming/tests/CMakeLists.txt +++ b/cpp/libcudf_streaming/tests/CMakeLists.txt @@ -45,7 +45,12 @@ target_sources( streaming/test_channel_metadata.cpp streaming/test_partition.cpp streaming/test_cudf_utils.cpp + streaming/test_shuffler.cpp + streaming/test_allgather.cpp + streaming/test_leaf_actor.cpp test_partition.cpp + test_shuffler.cpp + test_shuffler_many_streams.cpp ) add_executable(cudf_streaming_single_tests main/single.cpp) diff --git a/cpp/libcudf_streaming/tests/streaming/test_allgather.cpp b/cpp/libcudf_streaming/tests/streaming/test_allgather.cpp new file mode 100644 index 000000000000..a0be10aaced8 --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_allgather.cpp @@ -0,0 +1,171 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "base_streaming_fixture.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +using namespace rapidsmpf; + +class StreamingAllGather : public BaseStreamingFixture, + public ::testing::WithParamInterface> { + public: + void SetUp() override { BaseStreamingFixture::SetUpWithThreads(std::get<0>(GetParam())); } + + void TearDown() override { BaseStreamingFixture::TearDown(); } +}; + +INSTANTIATE_TEST_SUITE_P(StreamingAllGather, + StreamingAllGather, + ::testing::Combine(::testing::Values(1, 4), + ::testing::Values(MemoryType::HOST, + MemoryType::DEVICE)), + [](testing::TestParamInfo const& info) { + return "nthreads_" + std::to_string(std::get<0>(info.param)) + + "_mem_type_" + + (std::get<1>(info.param) == MemoryType::HOST ? "HOST" : "DEVICE"); + }); + +TEST_P(StreamingAllGather, basic) +{ + auto mem_type = std::get<1>(GetParam()); + auto comm = GlobalEnvironment->comm_; + auto allgather = streaming::AllGather(ctx, comm, OpID{0}); + + int size = comm->nranks(); + int rank = comm->rank(); + + constexpr int n_inserts{100}; + + coro::latch latch(n_inserts); + auto insert = [&](int insertion_id, std::uint64_t sequence) -> coro::task { + std::vector data(size); + for (int i = 0; i < size; i++) { + data[i] = (n_inserts * rank + insertion_id) * size + i; + } + + auto br = ctx->br(); + auto buf = br->make_buffer(br->stream_pool().get_stream(), + br->reserve_or_fail(data.size() * sizeof(int), mem_type)); + buf->write_access([&](std::byte* buf_data, rmm::cuda_stream_view stream) { + RAPIDSMPF_CUDA_TRY( + cuda_memcpy_async(buf_data, data.data(), data.size() * sizeof(int), stream)); + }); + auto meta = std::make_unique>(sizeof(int)); + std::memcpy(meta->data(), &size, sizeof(int)); + allgather.insert(sequence, PackedData{std::move(meta), std::move(buf)}); + latch.count_down(); + co_return; + }; + auto insert_finished = [&]() -> coro::task { + co_await latch; + allgather.insert_finished(); + }; + std::vector result(size * size * n_inserts); + + auto extract = [&]() -> coro::task { + auto data = co_await allgather.extract_all(streaming::AllGather::Ordered::YES); + std::size_t offset{0}; + for (auto& pd : data) { + RAPIDSMPF_EXPECTS(pd.metadata->size() == sizeof(int), "Invalid metadata buffer size"); + int msize; + std::memcpy(&msize, pd.metadata->data(), sizeof(int)); + RAPIDSMPF_EXPECTS(msize == size, "Corrupted metadata value"); + RAPIDSMPF_CUDA_TRY(cuda_memcpy_async( + result.data() + offset, pd.data->data(), pd.data->size, pd.data->stream())); + offset += msize; + pd.data->stream().synchronize(); + } + }; + + std::vector> pipeline{}; + // Insertions can all run concurrently. + for (int i = 0; i < n_inserts; i++) { + pipeline.push_back(ctx->executor()->schedule(insert(i, i))); + } + pipeline.push_back(ctx->executor()->schedule(insert_finished())); + pipeline.push_back(ctx->executor()->schedule(extract())); + streaming::run_actor_network(std::move(pipeline)); + std::vector expected(size * size * n_inserts); + std::iota(expected.begin(), expected.end(), 0); + EXPECT_EQ(expected, result); +} + +TEST_P(StreamingAllGather, streaming_actor) +{ + auto mem_type = std::get<1>(GetParam()); + + auto ch_in = ctx->create_channel(); + auto ch_out = ctx->create_channel(); + + auto comm = GlobalEnvironment->comm_; + int size = comm->nranks(); + int rank = comm->rank(); + + constexpr int n_inserts{100}; + std::vector input_messages; + for (int insertion_id = 0; insertion_id < n_inserts; insertion_id++) { + std::vector data(size); + for (int i = 0; i < size; i++) { + data[i] = (n_inserts * rank + insertion_id) * size + i; + } + + auto br = ctx->br(); + auto buf = br->make_buffer(br->stream_pool().get_stream(), + br->reserve_or_fail(data.size() * sizeof(int), mem_type)); + buf->write_access([&](std::byte* buf_data, rmm::cuda_stream_view stream) { + RAPIDSMPF_CUDA_TRY( + cuda_memcpy_async(buf_data, data.data(), data.size() * sizeof(int), stream)); + }); + auto meta = std::make_unique>(sizeof(int)); + std::memcpy(meta->data(), &size, sizeof(int)); + input_messages.emplace_back(streaming::to_message( + insertion_id, std::make_unique(std::move(meta), std::move(buf)))); + } + std::vector output_messages; + std::vector> pipeline{}; + pipeline.push_back(ctx->executor()->schedule( + streaming::actor::push_to_channel(ctx, ch_in, std::move(input_messages)))); + pipeline.push_back( + ctx->executor()->schedule(streaming::actor::allgather(ctx, comm, ch_in, ch_out, OpID{0}))); + pipeline.push_back( + ctx->executor()->schedule(streaming::actor::pull_from_channel(ctx, ch_out, output_messages))); + streaming::run_actor_network(std::move(pipeline)); + std::vector actual(size * size * n_inserts); + std::size_t offset{0}; + for (auto& msg : output_messages) { + auto pd = msg.release(); + RAPIDSMPF_EXPECTS(pd.metadata->size() == sizeof(int), "Invalid metadata buffer size"); + int msize; + std::memcpy(&msize, pd.metadata->data(), sizeof(int)); + RAPIDSMPF_EXPECTS(msize == size, "Corrupted metadata value"); + RAPIDSMPF_CUDA_TRY( + cuda_memcpy_async(actual.data() + offset, pd.data->data(), pd.data->size, pd.stream())); + offset += msize; + pd.stream().synchronize(); + } + std::vector expected(size * size * n_inserts); + std::iota(expected.begin(), expected.end(), 0); + EXPECT_EQ(expected, actual); +} diff --git a/cpp/libcudf_streaming/tests/streaming/test_leaf_actor.cpp b/cpp/libcudf_streaming/tests/streaming/test_leaf_actor.cpp new file mode 100644 index 000000000000..648fc02e3593 --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_leaf_actor.cpp @@ -0,0 +1,325 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "../utils.hpp" +#include "base_streaming_fixture.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +using namespace rapidsmpf; +using namespace rapidsmpf::streaming; +namespace actor = rapidsmpf::streaming::actor; +using namespace cudf_streaming::streaming; + +using StreamingLeafTasks = BaseStreamingFixture; + +TEST_F(StreamingLeafTasks, PushAndPullChunks) +{ + constexpr int num_rows = 100; + constexpr int num_chunks = 10; + + std::vector expects; + for (int i = 0; i < num_chunks; ++i) { + expects.emplace_back(random_table_with_index(i, num_rows, 0, 10)); + } + + std::vector actors; + auto ch1 = ctx->create_channel(); + + // Note, we use a scope to check that coroutines keeps the input alive. + { + std::vector inputs; + for (int i = 0; i < num_chunks; ++i) { + inputs.emplace_back(to_message( + i, + std::make_unique( + std::make_unique(expects[i], stream, ctx->br()->device_mr()), stream))); + } + + actors.push_back(actor::push_to_channel(ctx, ch1, std::move(inputs))); + } + + std::vector outputs; + actors.push_back(actor::pull_from_channel(ctx, ch1, outputs)); + + run_actor_network(std::move(actors)); + + 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()); + } +} + +namespace { +Actor shutdown(std::shared_ptr ctx, + std::shared_ptr ch, + std::vector&& tasks) +{ + ShutdownAtExit c{ch}; + coro_results(co_await coro::when_all(std::move(tasks))); + co_await ch->drain(ctx->executor()); +} + +Actor producer(std::shared_ptr ctx, + std::shared_ptr ch, + int val, + bool should_throw = false) +{ + co_await ctx->executor()->schedule(); + auto ticket = co_await ch->acquire(); + auto [_, receipt] = + co_await ticket.send(Message{0, std::make_unique(val), ContentDescription{}}); + if (should_throw) { throw std::runtime_error("Producer throws"); } + EXPECT_THROW(co_await ticket.send(Message{0, std::make_unique(val), ContentDescription{}}), + std::logic_error); + co_await receipt; + EXPECT_TRUE(receipt.is_ready()); +} + +Actor consumer(std::shared_ptr ctx, + std::shared_ptr ch, + std::atomic& result, + bool should_throw = false) +{ + ShutdownAtExit c{ch}; + co_await ctx->executor()->schedule(); + while (true) { + auto msg = co_await ch->receive(); + if (should_throw) { throw std::runtime_error("Consumer throws"); } + if (msg.empty()) { break; } + auto val = msg.release(); + result.fetch_add(val, std::memory_order_relaxed); + } +} +} // namespace + +TEST_F(StreamingLeafTasks, ThrottledAdaptor) +{ + auto ch = ctx->create_channel(); + auto throttle = std::make_shared(ch, 4); + std::vector producers; + std::vector consumers; + constexpr int n_producer{100}; + constexpr int n_consumer{3}; + for (int i = 0; i < n_producer; i++) { + producers.push_back(producer(ctx, throttle, i)); + } + consumers.push_back(shutdown(ctx, ch, std::move(producers))); + std::atomic result{0}; + for (int i = 0; i < n_consumer; i++) { + consumers.push_back(consumer(ctx, ch, result)); + } + run_actor_network(std::move(consumers)); + EXPECT_EQ(result, ((n_producer - 1) * n_producer) / 2); +} + +TEST_F(StreamingLeafTasks, ThrottledAdaptorThrowInProduce) +{ + auto ch = ctx->create_channel(); + auto throttle = std::make_shared(ch, 4); + std::vector producers; + std::vector consumers; + constexpr int n_producer{10}; + for (int i = 0; i < n_producer; i++) { + producers.push_back(producer(ctx, throttle, i, i == 2)); + } + consumers.push_back(shutdown(ctx, ch, std::move(producers))); + std::atomic result; + consumers.push_back(consumer(ctx, ch, result)); + EXPECT_THROW(run_actor_network(std::move(consumers)), std::runtime_error); +} + +TEST_F(StreamingLeafTasks, ThrottledAdaptorThrowInConsume) +{ + auto ch = ctx->create_channel(); + auto throttle = std::make_shared(ch, 4); + std::vector producers; + std::vector consumers; + constexpr int n_producer{100}; + constexpr int n_consumer{3}; + for (int i = 0; i < n_producer; i++) { + producers.push_back(producer(ctx, throttle, i)); + } + consumers.push_back(shutdown(ctx, ch, std::move(producers))); + std::atomic result; + for (int i = 0; i < n_consumer; i++) { + consumers.push_back(consumer(ctx, ch, result, i == 1)); + } + EXPECT_THROW(run_actor_network(std::move(consumers)), std::runtime_error); +} + +class StreamingThrottledAdaptor : public StreamingLeafTasks, + public ::testing::WithParamInterface {}; + +INSTANTIATE_TEST_SUITE_P(InvalidMaxTickets, StreamingThrottledAdaptor, ::testing::Values(-1, 0)); + +TEST_P(StreamingThrottledAdaptor, NonPositiveThrottleThrows) +{ + if (GlobalEnvironment->comm_->rank() != 0) { + // Test is independent of size of communicator. + GTEST_SKIP() << "Test only runs on rank zero"; + } + int max_tickets = GetParam(); + auto ch = ctx->create_channel(); + EXPECT_THROW(ThrottlingAdaptor(ch, max_tickets), std::logic_error); +} + +using StreamingBoundedQueue = StreamingLeafTasks; + +TEST_F(StreamingBoundedQueue, TicketUseOnce) +{ + if (GlobalEnvironment->comm_->rank() != 0) { + // Test is independent of size of communicator. + GTEST_SKIP() << "Test only runs on rank zero"; + } + + auto q = ctx->create_bounded_queue(2); + + auto producer = [](std::shared_ptr q) -> coro::task { + auto shutdown = q->raii_shutdown(); + auto ticket = co_await q->acquire(); + EXPECT_TRUE(ticket.has_value()); + auto sent = co_await ticket->send(Message{0, std::make_unique(0), ContentDescription{}}); + EXPECT_TRUE(sent); + EXPECT_THROW(co_await ticket->send(Message{1, std::make_unique(1), ContentDescription{}}), + std::logic_error); + }; + + coro::sync_wait(producer(q)); +} + +TEST_F(StreamingBoundedQueue, ShutdownStopsProducer) +{ + if (GlobalEnvironment->comm_->rank() != 0) { + // Test is independent of size of communicator. + GTEST_SKIP() << "Test only runs on rank zero"; + } + + auto q = ctx->create_bounded_queue(2); + + auto producer = [](std::shared_ptr q) -> coro::task { + auto shutdown = q->raii_shutdown(); + while (true) { + auto ticket = co_await q->acquire(); + // We're going to shutdown before the producer gets to go. + EXPECT_TRUE(!ticket.has_value()); + if (!ticket.has_value()) { break; } + } + }; + coro::sync_wait(q->shutdown()); + coro::sync_wait(producer(q)); +} + +TEST_F(StreamingBoundedQueue, ProducerThrows) +{ + if (GlobalEnvironment->comm_->rank() != 0) { + // Test is independent of size of communicator. + GTEST_SKIP() << "Test only runs on rank zero"; + } + + auto q = ctx->create_bounded_queue(2); + + auto producer = [](std::shared_ptr q) -> coro::task { + auto shutdown = q->raii_shutdown(); + throw std::runtime_error("Producer throws"); + }; + auto consumer = [](std::shared_ptr q) -> coro::task { + auto shutdown = q->raii_shutdown(); + auto [receipt, msg] = co_await q->receive(); + EXPECT_TRUE(msg.empty()); + }; + EXPECT_THROW(coro_results(coro::sync_wait(coro::when_all(consumer(q), producer(q)))), + std::runtime_error); +} + +TEST_F(StreamingBoundedQueue, ConsumerThrows) +{ + if (GlobalEnvironment->comm_->rank() != 0) { + // Test is independent of size of communicator. + GTEST_SKIP() << "Test only runs on rank zero"; + } + + auto q = ctx->create_bounded_queue(2); + + auto producer = [](std::shared_ptr q) -> coro::task { + auto shutdown = q->raii_shutdown(); + auto ticket = co_await q->acquire(); + EXPECT_TRUE(ticket.has_value()); + EXPECT_TRUE(co_await ticket->send(Message{0, std::make_unique(1), ContentDescription{}})); + while (true) { + auto ticket = co_await q->acquire(); + if (!ticket.has_value()) { break; } + } + }; + auto consumer = [](std::shared_ptr q) -> coro::task { + auto shutdown = q->raii_shutdown(); + auto [receipt, msg] = co_await q->receive(); + EXPECT_FALSE(msg.empty()); + EXPECT_EQ(msg.sequence_number(), 0); + EXPECT_EQ(msg.release(), 1); + co_await receipt; + throw std::runtime_error("Consumer throws"); + }; + EXPECT_THROW(coro_results(coro::sync_wait(coro::when_all(consumer(q), producer(q)))), + std::runtime_error); +} + +TEST_F(StreamingBoundedQueue, MultipleAcquire) +{ + if (GlobalEnvironment->comm_->rank() != 0) { + // Test is independent of size of communicator. + GTEST_SKIP() << "Test only runs on rank zero"; + } + + auto q = ctx->create_bounded_queue(2); + + auto producer = [](std::shared_ptr q) -> coro::task { + auto shutdown = q->raii_shutdown(); + int i = 0; + while (true) { + auto ticket = co_await q->acquire(); + if (!ticket.has_value()) { + EXPECT_EQ(i, 2); + break; + } + EXPECT_TRUE(ticket.has_value()); + EXPECT_TRUE(co_await ticket->send( + Message{static_cast(i), std::make_unique(i), ContentDescription{}})); + i++; + } + }; + auto consumer = [](std::shared_ptr q) -> coro::task { + auto shutdown = q->raii_shutdown(); + // Receiving, but not releasing the tickets + auto [_, msg] = co_await q->receive(); + EXPECT_FALSE(msg.empty()); + EXPECT_EQ(msg.sequence_number(), 0); + EXPECT_EQ(msg.release(), 0); + std::tie(std::ignore, msg) = co_await q->receive(); + EXPECT_FALSE(msg.empty()); + EXPECT_EQ(msg.sequence_number(), 1); + EXPECT_EQ(msg.release(), 1); + }; + coro_results(coro::sync_wait(coro::when_all(consumer(q), producer(q)))); +} diff --git a/cpp/libcudf_streaming/tests/streaming/test_shuffler.cpp b/cpp/libcudf_streaming/tests/streaming/test_shuffler.cpp new file mode 100644 index 000000000000..711ea46b1377 --- /dev/null +++ b/cpp/libcudf_streaming/tests/streaming/test_shuffler.cpp @@ -0,0 +1,223 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "../utils.hpp" +#include "base_streaming_fixture.hpp" + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace rapidsmpf; +using namespace rapidsmpf::streaming; +using cudf_streaming::streaming::TableChunk; +namespace actor = rapidsmpf::streaming::actor; +namespace cs_actor = cudf_streaming::streaming::actor; + +class BaseStreamingShuffle : public BaseStreamingFixture {}; + +TEST_F(BaseStreamingShuffle, zero_owned_partitions_completes) +{ + auto comm = GlobalEnvironment->comm_; + if (comm->nranks() < 2) { + GTEST_SKIP() << "Need at least 2 ranks so that some rank owns 0 partitions"; + } + constexpr Rank owner = 0; + auto collapse = [](std::shared_ptr const&, + shuffler::PartID, + shuffler::PartID) -> Rank { return owner; }; + constexpr OpID op_id = 0; + constexpr shuffler::PartID total = 4; + auto shuffler = std::make_unique(ctx, comm, op_id, total, collapse); + + coro::sync_wait(shuffler->insert_finished()); + + auto local_pids = shuffler->local_partitions(); + if (comm->rank() == owner) { + EXPECT_EQ(local_pids.size(), total); + } else { + EXPECT_TRUE(local_pids.empty()); + } +} + +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; + + void SetUp() override { BaseStreamingShuffle::SetUpWithThreads(GetParam()); } + + 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( + 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( + cs_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(cs_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(StreamingShuffler, + StreamingShuffler, + ::testing::Values(1, 2, 4), + [](testing::TestParamInfo const& info) { + return "nthreads_" + std::to_string(info.param); + }); + +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); + })); +} + +class ShufflerAsyncTest + : public BaseStreamingShuffle, + public ::testing::WithParamInterface> { + protected: + std::size_t n_inserts; + std::uint32_t n_partitions; + + static constexpr OpID op_id = 0; + static constexpr std::size_t n_elements = 100; + + void SetUp() override + { + std::tie(n_inserts, n_partitions) = GetParam(); + + BaseStreamingShuffle::SetUpWithThreads(4); + } + + void TearDown() override { BaseStreamingShuffle::TearDown(); } +}; + +INSTANTIATE_TEST_SUITE_P(StreamingShuffler, + ShufflerAsyncTest, + ::testing::Combine(::testing::Values(1, 10), // number of inserts + ::testing::Values(1, 10, 100) // number of partitions + ), + [](const testing::TestParamInfo& info) { + return "ninserts_" + std::to_string(std::get<0>(info.param)) + + "_nparts_" + std::to_string(std::get<1>(info.param)); + }); + +TEST_P(ShufflerAsyncTest, insert_wait_extract) +{ + auto comm = GlobalEnvironment->comm_; + auto shuffler = std::make_unique(ctx, comm, op_id, n_partitions); + + for (std::size_t i = 0; i < n_inserts; ++i) { + std::unordered_map data; + data.reserve(n_partitions); + for (shuffler::PartID pid = 0; pid < n_partitions; ++pid) { + data.emplace(pid, generate_packed_data(n_elements, 0, stream, *br)); + } + shuffler->insert(std::move(data)); + } + + coro::sync_wait(shuffler->insert_finished()); + + auto local_pids = + shuffler::Shuffler::local_partitions(comm, n_partitions, &shuffler::Shuffler::round_robin); + + std::vector finished_pids; + std::size_t n_chunks_received = 0; + for (auto pid : local_pids) { + auto chunks = shuffler->extract(pid); + n_chunks_received += chunks.size(); + finished_pids.push_back(pid); + } + + EXPECT_EQ(n_inserts * local_pids.size() * comm->nranks(), n_chunks_received); + EXPECT_EQ(local_pids, finished_pids); +} diff --git a/cpp/libcudf_streaming/tests/test_shuffler.cpp b/cpp/libcudf_streaming/tests/test_shuffler.cpp new file mode 100644 index 000000000000..b01dd3aa8042 --- /dev/null +++ b/cpp/libcudf_streaming/tests/test_shuffler.cpp @@ -0,0 +1,1032 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "environment.hpp" +#include "utils.hpp" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +extern Environment* GlobalEnvironment; +using namespace cudf_streaming::integrations; + +TEST(ReceivedChunks, spill_skips_control_messages) +{ + auto mr = cudf::get_current_device_resource_ref(); + auto br = rapidsmpf::BufferResource::create(mr); + + rapidsmpf::shuffler::detail::ReceivedChunks received; + + // Control messages have no data buffer (data_ == nullptr); spill must skip them + // rather than calling data_memory_type(), which throws if data_ is null. + received.insert(rapidsmpf::shuffler::detail::Chunk::from_finished_partition( + /*chunk_id=*/0, /*part_id=*/0, /*expected_num_chunks=*/1)); + + EXPECT_EQ(received.spill(br.get(), /*amount=*/1024), 0UL); +} + +TEST(ReceivedChunks, spill_respects_amount) +{ + auto mr = cudf::get_current_device_resource_ref(); + auto br = rapidsmpf::BufferResource::create(mr); + auto stream = cudf::get_default_stream(); + + rapidsmpf::shuffler::detail::ReceivedChunks received; + constexpr std::size_t chunk_size = 100; + + for (rapidsmpf::shuffler::PartID pid = 0; pid < 2; ++pid) { + auto metadata = std::make_unique>(std::size_t{1}, std::uint8_t{0}); + auto res = br->reserve_or_fail(chunk_size, rapidsmpf::MemoryType::DEVICE); + auto data = br->make_buffer(chunk_size, stream, res); + received.insert(rapidsmpf::shuffler::detail::Chunk::from_packed_data( + 0, pid, rapidsmpf::PackedData{std::move(metadata), std::move(data)})); + } + + // Two partitions, one 100-byte chunk each. spill() must stop after the first + // partition satisfies the request; the outer loop must not continue into partition 1. + EXPECT_EQ(received.spill(br.get(), chunk_size), chunk_size); +} + +TEST(MetadataMessage, round_trip) +{ + auto stream = cudf::get_default_stream(); + auto mr = cudf::get_current_device_resource_ref(); + auto br = rapidsmpf::BufferResource::create(mr); + + auto metadata = iota_vector(100); + + auto expect = rapidsmpf::shuffler::detail::Chunk::from_packed_data( + 1, // chunk_id + 2, // part_id + rapidsmpf::PackedData{ + std::make_unique>(metadata), // non-empty metadata + br->move(std::make_unique(), stream) // empty data + }); + + // Extract the metadata from then chunk. + auto msg = expect.serialize(); + EXPECT_FALSE(expect.is_metadata_buffer_set()); + + // Create a new chunk by deserializing the message. + auto result = rapidsmpf::shuffler::detail::Chunk::deserialize(*msg, br.get()); + + EXPECT_TRUE(expect.data_size() == 0 || result.is_data_buffer_set()); + // They should be identical. + EXPECT_EQ(expect.part_id(), result.part_id()); + EXPECT_EQ(expect.chunk_id(), result.chunk_id()); + EXPECT_EQ(expect.expected_num_chunks(), result.expected_num_chunks()); + EXPECT_EQ(expect.data_size(), result.data_size()); + EXPECT_EQ(expect.metadata_size(), result.metadata_size()); + + // The metadata should be identical to the original. + EXPECT_EQ(metadata, *result.release_metadata_buffer()); +} + +using MemoryLimitsMap = std::unordered_map; + +// Help function to get the `memory_limits` argument for a `BufferResource` +// that prioritizes the specified memory type. +MemoryLimitsMap get_memory_limits_map(rapidsmpf::MemoryType priorities) +{ + using namespace rapidsmpf; + + // We set all memory types to be unlimited. + MemoryLimitsMap ret = {{MemoryType::DEVICE, std::numeric_limits::max()}, + {MemoryType::HOST, std::numeric_limits::max()}}; + + // And then set device memory to zero if it isn't prioritized. + if (priorities != MemoryType::DEVICE) { ret.at(MemoryType::DEVICE) = 0; } + // Note, we never set host memory to zero because it is used to allocate + // stuff like metadata and control messages. + 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 +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) +{ + // To expose unexpected deadlocks, we use a 30s timeout. In a normal run, the + // 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; + } + // Tell the shuffler that we have no more input partitions. + insert_finished_fn(); + + 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( + cudf_streaming::integrations::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])); + } +} + +class MemoryLimits_NumPartition + : public cudf::test::BaseFixtureWithParam< + 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()); + br = rapidsmpf::BufferResource::create( + mr(), rapidsmpf::PinnedMemoryResource::Disabled, memory_limits); + + shuffler = std::make_unique(GlobalEnvironment->comm_, + 0, // op_id + total_num_partitions, + br.get()); + } + + void TearDown() override { shuffler.reset(); } + + protected: + 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; +}; + +// test different `memory_available` and `total_num_partitions`. +INSTANTIATE_TEST_SUITE_P( + Shuffler, + MemoryLimits_NumPartition, + testing::Combine(testing::ValuesIn({get_memory_limits_map(rapidsmpf::MemoryType::HOST), + get_memory_limits_map(rapidsmpf::MemoryType::DEVICE)}), + testing::Values(1, 2, 5, 10), // total_num_partitions + testing::Values(1, 9, 100, 100'000) // total_num_rows + ), + [](const testing::TestParamInfo& info) { + return std::to_string(info.index) + "__nparts_" + std::to_string(std::get<1>(info.param)) + + "__nrows_" + std::to_string(std::get<2>(info.param)); + }); + +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())); +} + +// Test that the same communicator can be used concurrently by multiple shufflers in +// separate threads +class ConcurrentShuffleTest : public cudf::test::BaseFixtureWithParam> { + public: + void SetUp() override + { + num_shufflers = std::get<0>(GetParam()); + total_num_partitions = static_cast(std::get<1>(GetParam())); + + // these resources will be used by multiple threads to instantiate shufflers + br = rapidsmpf::BufferResource::create(mr()); + stream = cudf::get_default_stream(); + } + + 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) + { + rapidsmpf::shuffler::Shuffler shuffler(GlobalEnvironment->comm_, + 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) + { + 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))); + })); + } + + for (auto& f : futures) { + ASSERT_NO_THROW(f.wait()); + } + } + + int num_shufflers; + rapidsmpf::shuffler::PartID total_num_partitions; + + rmm::cuda_stream_view stream; + std::shared_ptr br; +}; + +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(); })); +} + +// test different `num_shufflers` and `total_num_partitions`. +INSTANTIATE_TEST_SUITE_P(ConcurrentShuffle, + ConcurrentShuffleTest, + testing::Combine(testing::ValuesIn({1, 2, 4}), // num_shufflers + testing::ValuesIn({1, 10, 100}) // total_num_partitions + ), + [](const testing::TestParamInfo& info) { + return "num_shufflers_" + std::to_string(std::get<0>(info.param)) + + "__total_num_partitions_" + + std::to_string(std::get<1>(info.param)); + }); + +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(); + + // 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()}; + + // Control spilling by adjusting the DEVICE memory limit at runtime. + // `memory_available(DEVICE)` is computed as `limit - current_allocated()`, so a + // sufficiently large positive limit reliably keeps available memory > 0 (no spill), + // while a sufficiently large negative limit reliably keeps available memory < 0 + // (force spill), regardless of how many bytes are currently allocated from `mr`. + constexpr std::int64_t k_no_spill_limit = (1LL << 40); + constexpr std::int64_t k_force_spill_limit = -(1LL << 40); + auto br = rapidsmpf::BufferResource::create(mr, + rapidsmpf::PinnedMemoryResource::Disabled, + {{rapidsmpf::MemoryType::DEVICE, k_no_spill_limit}}, + std::nullopt // disable periodic spill check + ); + + // Create a communicator of size 1, such that each shuffler will run locally. + auto comm = GlobalEnvironment->split_comm(); + EXPECT_EQ(comm->nranks(), 1); + + // Create a shuffler and input chunks. + rapidsmpf::shuffler::Shuffler shuffler(comm, + 0, // op_id + 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 + + // Insert spills does nothing when device memory is available, we start + // with 2 device allocations. + EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); + shuffler.insert(std::move(input_chunks)); + // And we end with two 2 device allocations. + EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); + + // Let's force spilling. + br->set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_force_spill_limit); + + { + // Now extract triggers spilling of the partition not being extracted. + std::vector output_chunks = + cudf_streaming::integrations::unspill_partitions( + shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES); + EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1); + + // And insert also triggers spilling. We end up with zero device allocations. + std::unordered_map chunk; + chunk.emplace(0, std::move(output_chunks.at(0))); + shuffler.insert(std::move(chunk)); + EXPECT_EQ(mr.get_main_record().num_current_allocs(), 0); + } + + // Extract and unspill both partitions. + std::vector out0 = cudf_streaming::integrations::unspill_partitions( + shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES); + EXPECT_EQ(mr.get_main_record().num_current_allocs(), 1); + std::vector out1 = cudf_streaming::integrations::unspill_partitions( + shuffler.extract(1), br.get(), rapidsmpf::AllowOverbooking::YES); + EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); + + // Disable spilling and insert the first partition. + br->set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_no_spill_limit); + { + std::unordered_map chunk; + chunk.emplace(0, std::move(out0.at(0))); + shuffler.insert(std::move(chunk)); + } + EXPECT_EQ(mr.get_main_record().num_current_allocs(), 2); + + // Enable spilling and insert the second partition, which should trigger spilling + // of both the first partition already in the shuffler and the second partition + // that are being inserted. + br->set_memory_limit(rapidsmpf::MemoryType::DEVICE, k_force_spill_limit); + { + std::unordered_map chunk; + chunk.emplace(1, std::move(out1.at(0))); + shuffler.insert(std::move(chunk)); + } + EXPECT_EQ(mr.get_main_record().num_current_allocs(), 0); + shuffler.insert_finished(); +} + +TEST(FinishCounterTests, zero_local_partitions_immediately_finished) +{ + rapidsmpf::shuffler::detail::FinishCounter finish_counter( + /*nranks=*/2, /*n_local_partitions=*/0); + + EXPECT_TRUE(finish_counter.all_finished()); +} + +TEST(FinishCounterTests, nonzero_local_partitions_finishes_after_all_chunks) +{ + rapidsmpf::shuffler::detail::FinishCounter finish_counter( + /*nranks=*/1, /*n_local_partitions=*/2); + + EXPECT_FALSE(finish_counter.all_finished()); + + // One rank sends 3 chunks total. + finish_counter.move_goalpost(0, 3); + finish_counter.add_finished_chunk(); + finish_counter.add_finished_chunk(); + EXPECT_FALSE(finish_counter.all_finished()); + + finish_counter.add_finished_chunk(); + EXPECT_TRUE(finish_counter.all_finished()); +} + +TEST(FinishCounterTests, multi_rank_completion) +{ + auto comm = GlobalEnvironment->comm_; + + if (comm->rank() != 0) { GTEST_SKIP() << "Test only runs on rank 0"; } + + // Use nranks partitions so each rank owns exactly 1 partition (round robin). + auto out_nparts = rapidsmpf::safe_cast(comm->nranks()); + + auto local_partitions = rapidsmpf::shuffler::Shuffler::local_partitions( + comm, out_nparts, &rapidsmpf::shuffler::Shuffler::round_robin); + ASSERT_EQ(local_partitions.size(), 1); + + rapidsmpf::shuffler::detail::FinishCounter finish_counter(comm->nranks(), + local_partitions.size()); + + // Not finished yet. + EXPECT_FALSE(finish_counter.all_finished()); + + // For nranks ranks, each rank sends 1 data chunk + 1 control, so + // move_goalpost(rank, 2) per rank. + for (rapidsmpf::Rank r = 0; r < comm->nranks(); r++) { + finish_counter.move_goalpost(r, 2); + } + + // Add finished chunks: 1 data chunk per rank + 1 control per rank = 2 * nranks + for (rapidsmpf::Rank r = 0; r < comm->nranks(); r++) { + finish_counter.add_finished_chunk(); // data chunk + finish_counter.add_finished_chunk(); // control chunk + } + + EXPECT_TRUE(finish_counter.all_finished()); +} + +class FinishCounterMultithreadingTest + : public ::testing::TestWithParam> { + protected: + rapidsmpf::Rank const nranks{1}; // simulate a single rank + + std::unique_ptr finish_counter; + rapidsmpf::shuffler::PartID npartitions; + std::uint32_t nthreads; + + void SetUp() override + { + std::tie(npartitions, nthreads) = GetParam(); + + finish_counter = + std::make_unique(nranks, npartitions); + } + + void produce_data() + { + // Simulate nranks=1: one rank reports chunk count = npartitions + 1 + // (one data chunk per partition + 1 control message) + finish_counter->move_goalpost(rapidsmpf::Rank{0}, npartitions + 1); + for (rapidsmpf::shuffler::PartID i = 0; i <= npartitions; i++) { + finish_counter->add_finished_chunk(); + } + } +}; + +// Parametrize on number of partitions and number of consumer threads +INSTANTIATE_TEST_SUITE_P(FinishCounterMultithreadingTestP, + FinishCounterMultithreadingTest, + testing::Combine(testing::Values(1, 2, 100, 101), + testing::Values(1, 2, 3)), + [](const auto& info) { + return "npartitions_" + std::to_string(std::get<0>(info.param)) + + "__nthreads_" + std::to_string(std::get<1>(info.param)); + }); + +TEST_P(FinishCounterMultithreadingTest, concurrent_all_finished_check) +{ + produce_data(); + + std::atomic n_checks{0}; + std::vector> futures; + for (std::uint32_t tid = 0; tid < nthreads; tid++) { + futures.emplace_back(std::async(std::launch::async, [&, tid] { + for (std::uint32_t i = tid; i < npartitions; i += nthreads) { + EXPECT_TRUE(finish_counter->all_finished()); + n_checks.fetch_add(1, std::memory_order_relaxed); + } + })); + } + + EXPECT_NO_THROW(std::ranges::for_each(futures, [](auto& f) { f.get(); })); + + EXPECT_EQ(npartitions, n_checks); + EXPECT_TRUE(finish_counter->all_finished()); +} + +class ContiguousPartitionAssignmentTest + : public ::testing::TestWithParam { + protected: + void SetUp() override + { + comm = GlobalEnvironment->comm_; + nranks = comm->nranks(); + rank = comm->rank(); + total_num_partitions = GetParam(); + } + + std::shared_ptr comm; + rapidsmpf::Rank nranks; + rapidsmpf::Rank rank; + rapidsmpf::shuffler::PartID total_num_partitions; +}; + +INSTANTIATE_TEST_SUITE_P(PartitionAssignment, + ContiguousPartitionAssignmentTest, + testing::Values(1, 2, 3, 5, 7, 10, 16, 100), + [](const testing::TestParamInfo& info) { + return "nparts_" + std::to_string(info.param); + }); + +TEST_P(ContiguousPartitionAssignmentTest, contiguous) +{ + std::vector> rank_partitions(nranks); + for (rapidsmpf::shuffler::PartID pid = 0; pid < total_num_partitions; ++pid) { + auto owner = rapidsmpf::shuffler::Shuffler::contiguous(comm, pid, total_num_partitions); + EXPECT_GE(owner, 0); + EXPECT_LT(owner, nranks); + rank_partitions[owner].push_back(pid); + } + + // Each rank's partitions must be contiguous. + for (rapidsmpf::Rank r = 0; r < nranks; ++r) { + auto const& pids = rank_partitions[r]; + for (std::size_t i = 1; i < pids.size(); ++i) { + EXPECT_EQ(pids[i], pids[i - 1] + 1); + } + } + + // Concatenating all rank partitions should cover [0, total_num_partitions). + std::vector all_pids; + for (auto const& pids : rank_partitions) { + all_pids.insert(all_pids.end(), pids.begin(), pids.end()); + } + EXPECT_EQ(all_pids, iota_vector(total_num_partitions)); +} + +TEST(Shuffler, ShutdownWhilePaused) +{ + auto progress_thread = GlobalEnvironment->comm_->progress_thread(); + auto mr = cudf::get_current_device_resource_ref(); + + auto br = rapidsmpf::BufferResource::create(mr); + + auto shuffler = rapidsmpf::shuffler::Shuffler(GlobalEnvironment->comm_, 0, 1, br.get()); + + progress_thread->pause(); + EXPECT_FALSE(progress_thread->is_running()); + shuffler.insert_finished(); + // Progress thread must be running before shuffle shutdown, otherwise we have some + // orphan messages in the shuffle that are never sent/received. + progress_thread->resume(); + EXPECT_TRUE(progress_thread->is_running()); + 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 { + 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()); + + shuffler = std::make_unique( + GlobalEnvironment->comm_, 0, nparts, br.get()); + } + + void TearDown() override { shuffler.reset(); } + + void insert_chunks( + std::unordered_map&& chunks) + { + if (!chunks.empty()) { shuffler->insert(std::move(chunks)); } + shuffler->insert_finished(); + } + + void verify_extracted_chunks(auto expected_empty_fn) + { + EXPECT_NO_THROW(shuffler->wait(wait_timeout)); + for (auto pid : shuffler->local_partitions()) { + SCOPED_TRACE("pid: " + std::to_string(pid)); + std::vector chunks; + EXPECT_NO_THROW({ chunks = shuffler->extract(pid); }); + + if (expected_empty_fn(pid)) { + EXPECT_TRUE(chunks.empty()); + } else { + EXPECT_EQ(GlobalEnvironment->comm_->nranks(), chunks.size()); + } + } + } + + auto empty_packed_data() + { + return rapidsmpf::PackedData{std::make_unique>(), + br->move(std::make_unique(), stream)}; + } + + auto non_empty_packed_data() + { + return rapidsmpf::PackedData{ + std::make_unique>(10), + br->move(std::make_unique(10, stream), stream)}; + } + + rmm::cuda_stream_view stream; + std::shared_ptr br; + std::unique_ptr shuffler; +}; + +TEST_F(ExtractEmptyPartitionsTest, NoInsertions) +{ + insert_chunks({}); + EXPECT_NO_FATAL_FAILURE(verify_extracted_chunks([](auto) { return true; })); +} + +TEST_F(ExtractEmptyPartitionsTest, AllEmptyInsertions) +{ + std::unordered_map chunks; + for (rapidsmpf::shuffler::PartID pid = 0; pid < nparts; ++pid) { + chunks.emplace(pid, empty_packed_data()); + } + + insert_chunks(std::move(chunks)); + EXPECT_NO_FATAL_FAILURE(verify_extracted_chunks([](auto) { return true; })); +} + +TEST_F(ExtractEmptyPartitionsTest, SomeEmptyInsertions) +{ + std::unordered_map chunks; + for (rapidsmpf::shuffler::PartID pid = 0; pid < nparts; ++pid) { + if (pid % 3 == 0) { chunks.emplace(pid, empty_packed_data()); } + } + + insert_chunks(std::move(chunks)); + EXPECT_NO_FATAL_FAILURE(verify_extracted_chunks([](auto) { return true; })); +} + +TEST_F(ExtractEmptyPartitionsTest, SomeEmptyAndNonEmptyInsertions) +{ + std::unordered_map chunks; + for (rapidsmpf::shuffler::PartID pid = 0; pid < nparts; ++pid) { + if (pid % 3 == 0) { + chunks.emplace(pid, empty_packed_data()); + } else { + chunks.emplace(pid, non_empty_packed_data()); + } + } + + insert_chunks(std::move(chunks)); + EXPECT_NO_FATAL_FAILURE(verify_extracted_chunks([](auto pid) { return pid % 3 == 0; })); +} + +TEST(ShufflerTest, multiple_shutdowns) +{ + auto& comm = GlobalEnvironment->comm_; + auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + auto shuffler = + std::make_unique(comm, 0, comm->nranks(), br.get()); + + shuffler->insert_finished(); + EXPECT_NO_THROW(shuffler->wait(std::chrono::seconds(30))); + for (auto pid : shuffler->local_partitions()) { + std::ignore = shuffler->extract(pid); + } + + constexpr int n_threads = 10; + std::vector> futures; + for (int i = 0; i < n_threads; ++i) { + futures.emplace_back(std::async(std::launch::async, [&] { shuffler->shutdown(); })); + } + std::ranges::for_each(futures, [](auto& future) { future.get(); }); +} + +// 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()); + + // 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); + + { + 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; + } + std::ranges::for_each(insert_futures, [](auto& f) { f.get(); }); + shuffler.insert_finished(); + } + + auto local_pids = shuffler.local_partitions(); + std::vector> futures; + for (auto pid : local_pids) { + futures.push_back(std::async(std::launch::async, [&, pid] { + EXPECT_NO_THROW(shuffler.wait(wait_timeout)); + auto chunks = shuffler.extract(pid); + auto result = cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::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::ranges::for_each(futures, [](auto& f) { f.get(); }); +} + +// Test that reusing an OpID after a completed shuffle doesn't cause cross-matching of +// messages between the old and new shuffle. +// +// On rank 0 we inject a stream-ordered delay into device allocations so that received +// chunks stay "not ready" in the event loop. With small messages, other ranks can finish +// and move on to the next shuffle. Its messages will then be matched on rank 0 by the +// blocked previous shuffle, unless recv gating correctly prevents cross-talk. +TEST(Shuffler, opid_reuse) +{ + auto const& comm = GlobalEnvironment->comm_; + if (comm->nranks() == 1) { GTEST_SKIP() << "OpID reuse test requires multiple ranks"; } + + auto stream = cudf::get_default_stream(); + 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}; + + rmm::mr::cuda_memory_resource mr; + auto br = rapidsmpf::BufferResource::create(mr); + + // On rank 0, wrap the device MR with a delayed version for the shuffler. + std::unique_ptr delayed_mr; + std::shared_ptr delayed_br; + rapidsmpf::BufferResource* shuffler_br = br.get(); + if (comm->rank() == 0) { + delayed_mr = std::make_unique(mr, std::chrono::milliseconds(500)); + delayed_br = rapidsmpf::BufferResource::create(*delayed_mr); + 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; + } + }; + + 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( + cudf_streaming::integrations::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])); + } + }; + + rapidsmpf::shuffler::Shuffler shuffle1(comm, op_id, total_num_partitions, shuffler_br); + insert_data(shuffle1, 42); + shuffle1.insert_finished(); + EXPECT_NO_THROW(shuffle1.wait(wait_timeout)); + + rapidsmpf::shuffler::Shuffler shuffle2(comm, op_id, total_num_partitions, shuffler_br); + insert_data(shuffle2, 123); + shuffle2.insert_finished(); + EXPECT_NO_THROW(shuffle2.wait(wait_timeout)); + + validate_results(shuffle1, 42); + validate_results(shuffle2, 123); +} + +// Same as opid_reuse but with total_num_partitions=1, so only rank 0 owns a partition. +// All other ranks have empty local_partitions and empty recv loops. This exercises the +// edge case where non-partition-owning ranks must still correctly handle op_id reuse. +TEST(Shuffler, opid_reuse_with_empty_partitions) +{ + auto const& comm = GlobalEnvironment->comm_; + if (comm->nranks() == 1) { GTEST_SKIP() << "OpID reuse test requires multiple ranks"; } + + auto stream = cudf::get_default_stream(); + 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}; + + rmm::mr::cuda_memory_resource mr; + auto br = rapidsmpf::BufferResource::create(mr); + + // On rank 0, wrap the device MR with a delayed version for the shuffler. + std::unique_ptr delayed_mr; + std::shared_ptr delayed_br; + rapidsmpf::BufferResource* shuffler_br = br.get(); + if (comm->rank() == 0) { + delayed_mr = std::make_unique(mr, std::chrono::milliseconds(500)); + delayed_br = rapidsmpf::BufferResource::create(*delayed_mr); + 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)); + } + }; + + 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( + cudf_streaming::integrations::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])); + } + }; + + rapidsmpf::shuffler::Shuffler shuffle1(comm, op_id, total_num_partitions, shuffler_br); + insert_data(shuffle1, 42); + shuffle1.insert_finished(); + EXPECT_NO_THROW(shuffle1.wait(wait_timeout)); + + rapidsmpf::shuffler::Shuffler shuffle2(comm, op_id, total_num_partitions, shuffler_br); + insert_data(shuffle2, 123); + shuffle2.insert_finished(); + EXPECT_NO_THROW(shuffle2.wait(wait_timeout)); + + validate_results(shuffle1, 42); + validate_results(shuffle2, 123); +} diff --git a/cpp/libcudf_streaming/tests/test_shuffler_many_streams.cpp b/cpp/libcudf_streaming/tests/test_shuffler_many_streams.cpp new file mode 100644 index 000000000000..277826cd7c64 --- /dev/null +++ b/cpp/libcudf_streaming/tests/test_shuffler_many_streams.cpp @@ -0,0 +1,87 @@ +/** + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * reserved. SPDX-License-Identifier: Apache-2.0 + */ + +#include "environment.hpp" +#include "utils.hpp" + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +extern Environment* GlobalEnvironment; + +namespace { + +/** + * @brief Generate a random CUDA stream priority. + * + * @param random_generator A random number generator used to produce the priority. + * @return A valid CUDA stream priority in the device range. + */ +int gen_stream_priority(std::mt19937& random_generator) +{ + int least_priority = 0; // numerically larger (often 0) => lower priority + int greatest_priority = 0; // numerically smaller (often negative) => higher priority + RAPIDSMPF_CUDA_TRY(cudaDeviceGetStreamPriorityRange(&least_priority, &greatest_priority)); + int const num_priorities = least_priority - greatest_priority + 1; + std::uniform_int_distribution dist(0, num_priorities - 1); + return greatest_priority + dist(random_generator); +} + +// To expose unexpected deadlocks, we use a 30s timeout. In a normal run, the +// shuffle shouldn't get near 30s. +constexpr auto wait_timeout = std::chrono::seconds{30}; + +} // namespace + +TEST(ShufflerManyStreams, Test) +{ + std::mt19937 random_generator{42}; + constexpr std::size_t chunksize = 1 << 20; + constexpr int num_partitions = 100; + auto br = rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); + + // Create a CUDA stream for each partition. + // To stress-test stream handling, assign random priorities so streams are more + // likely to execute in mixed order. + std::array partition_streams{}; + for (rapidsmpf::shuffler::PartID pid = 0; pid < num_partitions; ++pid) { + RAPIDSMPF_CUDA_TRY(cudaStreamCreateWithPriority( + &partition_streams[pid], cudaStreamNonBlocking, gen_stream_priority(random_generator))); + } + + rapidsmpf::shuffler::Shuffler shuffler(GlobalEnvironment->comm_, 0, num_partitions, br.get()); + + std::unordered_map partitions; + for (rapidsmpf::shuffler::PartID pid = 0; pid < num_partitions; ++pid) { + partitions.insert({pid, generate_packed_data(chunksize, pid, partition_streams[pid], *br)}); + } + + shuffler.insert(std::move(partitions)); + shuffler.insert_finished(); + + EXPECT_NO_THROW(shuffler.wait(wait_timeout)); + for (auto const pid : shuffler.local_partitions()) { + auto partition_chunks = shuffler.extract(pid); + for (auto& chunk : partition_chunks) { + auto const stream = chunk.data->stream(); + EXPECT_NO_FATAL_FAILURE(validate_packed_data(std::move(chunk), chunksize, pid, stream, *br)); + } + } + + for (auto& stream : partition_streams) { + RAPIDSMPF_CUDA_TRY(cudaStreamDestroy(stream)); + } +} diff --git a/cpp/src/io/utilities/config_utils.cpp b/cpp/src/io/utilities/config_utils.cpp index f452babee592..c07ad634279a 100644 --- a/cpp/src/io/utilities/config_utils.cpp +++ b/cpp/src/io/utilities/config_utils.cpp @@ -23,7 +23,7 @@ void set_up_kvikio() // Workaround for https://github.com/rapidsai/cudf/issues/14140, where cuFileDriverOpen errors // out if no CUDA calls have been made before it. This is a no-op if the CUDA context is already // initialized. - cudaFree(nullptr); + CUDF_CUDA_TRY(cudaFree(nullptr)); auto const compat_mode = kvikio::getenv_or("KVIKIO_COMPAT_MODE", kvikio::CompatMode::ON); kvikio::defaults::set_compat_mode(compat_mode); diff --git a/dependencies.yaml b/dependencies.yaml index 0409d0c4cfe2..fc3d14073e75 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -24,11 +24,13 @@ files: - depends_on_librmm - depends_on_libnvcomp - depends_on_librapidsmpf + - depends_on_libucxx - depends_on_numba_cuda - depends_on_rapids_logger - depends_on_rapidsmpf - depends_on_ray - depends_on_rmm + - depends_on_ucxx - develop - docs - iwyu @@ -492,6 +494,7 @@ files: - depends_on_libcudf - depends_on_librmm - depends_on_librapidsmpf + - depends_on_libucxx py_run_libcudf_streaming: output: pyproject pyproject_dir: python/libcudf_streaming @@ -520,10 +523,12 @@ files: - build_python_common - depends_on_libcudf_streaming - depends_on_librapidsmpf + - depends_on_libucxx - depends_on_librmm - depends_on_pylibcudf - depends_on_rapidsmpf - depends_on_rmm + - depends_on_ucxx py_run_cudf_streaming: output: pyproject pyproject_dir: python/cudf_streaming @@ -757,6 +762,7 @@ dependencies: - output_types: [conda] packages: - cuda-cudart-dev + - cuda-cupti-dev - cuda-nvrtc-dev - cuda-nvtx-dev - cuda-sanitizer-api @@ -1338,6 +1344,56 @@ dependencies: - matrix: packages: - *rmm_unsuffixed + depends_on_libucxx: + common: + - output_types: conda + packages: + - &libucxx_unsuffixed libucxx==0.51.*,>=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: + - libucxx-cu12==0.51.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - libucxx-cu13==0.51.*,>=0.0.0a0 + - matrix: + packages: + - *libucxx_unsuffixed + depends_on_ucxx: + common: + - output_types: conda + packages: + - &ucxx_unsuffixed ucxx==0.51.*,>=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: + - ucxx-cu12==0.51.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - ucxx-cu13==0.51.*,>=0.0.0a0 + - matrix: + packages: + - *ucxx_unsuffixed depends_on_rapidsmpf: common: - output_types: conda diff --git a/python/cudf_streaming/cudf_streaming/examples/__init__.py b/python/cudf_streaming/cudf_streaming/examples/__init__.py new file mode 100644 index 000000000000..1c72bd2fabd3 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/examples/__init__.py @@ -0,0 +1,3 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +from __future__ import annotations diff --git a/python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py b/python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py new file mode 100644 index 000000000000..76faf1e0df53 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/examples/bulk_mpi_shuffle.py @@ -0,0 +1,542 @@ +# 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 mpi4py import MPI + +import rapidsmpf.bootstrap +import rapidsmpf.communicator.mpi +import rmm.mr +from cudf_streaming.integrations.partition import ( + partition_and_pack, + unpack_and_concat, + unspill_partitions, +) +from rapidsmpf.config import Options, get_environment_variables +from rapidsmpf.memory.buffer import MemoryType +from rapidsmpf.memory.buffer_resource import BufferResource +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 +from rmm.pylibrmm.stream import DEFAULT_STREAM + +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/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py b/python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py new file mode 100644 index 000000000000..508e6b184810 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/examples/bulk_ray_shuffle.py @@ -0,0 +1,459 @@ +# 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 + +import rmm.mr +from cudf_streaming.integrations.partition import ( + partition_and_pack, + unpack_and_concat, + unspill_partitions, +) +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.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") # noqa: T201 + 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/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.py b/python/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.py new file mode 100644 index 000000000000..06ebfd2bbd5f --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/examples/ray_shuffle_example.py @@ -0,0 +1,202 @@ +# 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 + +import rmm +from cudf_streaming.integrations.partition import ( + partition_and_pack, + unpack_and_concat, + unspill_partitions, +) +from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster +from rapidsmpf.memory.buffer_resource import BufferResource +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). + + rng = np.random.default_rng( + 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(rng.integers(0, 1000, self._num_rows)), + plc.Column.from_iterable_of_py( + (["cat", "dog"] * ((self._num_rows + 1) // 2))[ + : self._num_rows + ], + 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/cudf_streaming/cudf_streaming/examples/streaming_basic_example.py b/python/cudf_streaming/cudf_streaming/examples/streaming_basic_example.py new file mode 100644 index 000000000000..caebe577e015 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/examples/streaming_basic_example.py @@ -0,0 +1,154 @@ +# 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 + +import rmm.mr +from cudf_streaming.streaming.table_chunk import TableChunk +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 +from rmm.pylibrmm.stream import DEFAULT_STREAM + +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()}") # noqa: T201 diff --git a/python/cudf_streaming/cudf_streaming/tests/conftest.py b/python/cudf_streaming/cudf_streaming/tests/conftest.py index 6eb6cf7fc546..5a0f60fad3eb 100644 --- a/python/cudf_streaming/cudf_streaming/tests/conftest.py +++ b/python/cudf_streaming/cudf_streaming/tests/conftest.py @@ -71,6 +71,21 @@ def comm( MPI.COMM_WORLD.barrier() +@pytest.fixture +def device_mr() -> Generator[rmm.mr.CudaMemoryResource, None, None]: + """ + Fixture for creating a new cuda memory resource and making it the + current rmm resource temporarily. + """ + prior_mr = rmm.mr.get_current_device_resource() + try: + mr = rmm.mr.CudaMemoryResource() + rmm.mr.set_current_device_resource(mr) + yield mr + finally: + rmm.mr.set_current_device_resource(prior_mr) + + @pytest.fixture def stream() -> Stream: """CUDA stream for test operations.""" diff --git a/python/cudf_streaming/cudf_streaming/tests/test_allgather.py b/python/cudf_streaming/cudf_streaming/tests/test_allgather.py new file mode 100644 index 000000000000..010121e5d609 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_allgather.py @@ -0,0 +1,214 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for AllGather functionality.""" + +from __future__ import annotations + +from contextlib import nullcontext +from typing import TYPE_CHECKING + +import numpy as np +import pylibcudf as plc +import pytest +from pylibcudf.contiguous_split import pack + +from cudf_streaming.integrations.partition import ( + packed_data_from_cudf_packed_columns, + unpack_and_concat, +) +from rapidsmpf.coll import AllGather +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.testing import assert_eq + +if TYPE_CHECKING: + import rmm.mr + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.memory.packed_data import PackedData + from rmm.pylibrmm.stream import Stream + + +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 packed_data_from_cudf_packed_columns(packed_columns, stream, br) + + +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) + + +def gen_offset(i: int, r: int) -> int: + """Generate offset value like the C++ test: i * 10 + r.""" + return i * 10 + r + + +@pytest.mark.parametrize("n_elements", [0, 1, 10, 100]) +@pytest.mark.parametrize("n_inserts", [0, 1, 10]) +@pytest.mark.parametrize("ordered", [False, True]) +@pytest.mark.parametrize( + "use_context_manager", [True, False], ids=["context", "non-context"] +) +def test_basic_allgather( + comm: Communicator, + device_mr: rmm.mr.CudaMemoryResource, + stream: Stream, + n_elements: int, + n_inserts: int, + ordered: bool, + use_context_manager: bool, +) -> None: + """ + Test basic AllGather functionality. + + This test mirrors the C++ AllGatherTest::basic_allgather test. + Each rank inserts n_inserts pieces of data, then all ranks + should receive all data from all ranks. + """ + br = BufferResource(device_mr) + + # Create AllGather instance + allgather = AllGather( + comm=comm, + op_id=0, + br=br, + ) + + n_ranks = comm.nranks + this_rank = comm.rank + + cm = allgather if use_context_manager else nullcontext(allgather) + with cm as ag: + for i in range(n_inserts): + packed_data = generate_packed_data( + n_elements, gen_offset(i, this_rank), stream, br + ) + ag.insert(i, packed_data) + if not use_context_manager: + allgather.insert_finished() + + # Wait for completion and extract results + results = allgather.wait_and_extract(ordered=ordered) + + # Check results + if n_inserts > 0: + expected_total = n_inserts * n_ranks + assert len(results) == expected_total + + if ordered: + # Results should be ordered by rank and insertion order + # Expected pattern: + # rank0: offset(0,0), offset(1,0), ..., offset(n_inserts-1,0) + # rank1: offset(0,1), offset(1,1), ..., offset(n_inserts-1,1) + # ... + # rankN: offset(0,N), offset(1,N), ..., offset(n_inserts-1,N) + + for r in range(n_ranks): + for i in range(n_inserts): + result_idx = r * n_inserts + i + expected_offset = gen_offset(i, r) + validate_packed_data( + results[result_idx], + n_elements, + expected_offset, + stream, + br, + ) + 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 + + +def test_insert_finished_raises_in_context( + comm: Communicator, + device_mr: rmm.mr.CudaMemoryResource, +) -> None: + """Test that insert_finished raises when called inside a context manager.""" + br = BufferResource(device_mr) + ag = AllGather(comm=comm, op_id=0, br=br) + with ( + ag, + pytest.raises( + ValueError, match=r"Cannot call insert_finished.*within a context" + ), + ): + ag.insert_finished() + ag.wait_and_extract(ordered=True) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_integration_partition.py b/python/cudf_streaming/cudf_streaming/tests/test_integration_partition.py new file mode 100644 index 000000000000..788d4f2d7e91 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_integration_partition.py @@ -0,0 +1,140 @@ +# 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 + +from cudf_streaming.integrations.partition import ( + partition_and_pack, + spill_partitions, + split_and_pack, + unpack_and_concat, + unspill_partitions, +) +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.testing import assert_eq +from rmm.pylibrmm.stream import DEFAULT_STREAM + +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, + ) + + +@pytest.mark.parametrize("cols", [[[1, 2, 3], [2, 2, 1]], [[], []]]) +@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) +def test_spill_unspill_roundtrip( + 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, + ) + + # Spill roundtrip + spilled = spill_partitions(partitions.values(), 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) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_shuffler.py b/python/cudf_streaming/cudf_streaming/tests/test_shuffler.py new file mode 100644 index 000000000000..3bdd00c22841 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_shuffler.py @@ -0,0 +1,190 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# 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 + +from cudf_streaming.integrations.partition import ( + partition_and_pack, + unpack_and_concat, + unspill_partitions, +) +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.shuffler import ( + Shuffler, +) +from rapidsmpf.testing import assert_eq + +if TYPE_CHECKING: + import rmm.mr + from rapidsmpf.communicator.communicator import Communicator + from rmm.pylibrmm.stream import Stream + + +@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( + comm: Communicator, + device_mr: rmm.mr.CudaMemoryResource, + stream: Stream, + batch_size: int | None, + total_num_partitions: int, +) -> None: + 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 + rng = np.random.default_rng( + 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( + rng.integers(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, + total_num_partitions=total_num_partitions, + 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, + ) + shuffler.insert_chunks(packed_inputs) + + # Tell shuffler we are done adding data + 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, + ) + + shuffler.shutdown() + assert extracted_partitions == expected_partitions diff --git a/python/cudf_streaming/cudf_streaming/tests/test_sparse_alltoall.py b/python/cudf_streaming/cudf_streaming/tests/test_sparse_alltoall.py new file mode 100644 index 000000000000..7853a13b025f --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_sparse_alltoall.py @@ -0,0 +1,199 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for SparseAlltoall functionality.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pylibcudf as plc +import pytest + +from cudf_streaming.integrations.partition import ( + packed_data_from_cudf_packed_columns, + unpack_and_concat, +) +from rapidsmpf.coll.sparse_alltoall import SparseAlltoall +from rapidsmpf.memory.buffer_resource import BufferResource +from rapidsmpf.testing import assert_eq + +if TYPE_CHECKING: + import rmm.mr + from rapidsmpf.communicator.communicator import Communicator + from rapidsmpf.memory.packed_data import PackedData + from rmm.pylibrmm.stream import Stream + + +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 packed_data_from_cudf_packed_columns(packed_columns, stream, br) + + +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 = [] + if comm.rank > 0: + peers.append(comm.rank - 1) + if comm.rank + 1 < comm.nranks: + peers.append(comm.rank + 1) + return peers, peers + + +def make_offset(src: int, dst: int, ordinal: int) -> int: + """Encode the sender, receiver, and per-destination ordinal into payload data.""" + return src * 1000 + dst * 100 + ordinal * 10 + + +@pytest.mark.parametrize("n_inserts", [0, 1, 3]) +def test_basic( + comm: Communicator, + device_mr: rmm.mr.CudaMemoryResource, + stream: Stream, + n_inserts: int, +) -> None: + """ + Test a simple sparse exchange with immediate neighbors. + + Each rank sends `n_inserts` messages to each adjacent rank and then checks + that extraction by source returns the expected number of payloads. + """ + br = BufferResource(device_mr) + srcs, dsts = expected_peers(comm) + sparse_alltoall = SparseAlltoall( + comm=comm, + op_id=0, + br=br, + srcs=srcs, + dsts=dsts, + ) + + for dst in dsts: + for ordinal in range(n_inserts): + sparse_alltoall.insert( + dst, + generate_packed_data( + n_elements=4, + offset=make_offset(comm.rank, dst, ordinal), + stream=stream, + br=br, + ), + ) + + sparse_alltoall.insert_finished() + sparse_alltoall.wait() + + for src in srcs: + 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) + + +def test_non_participating_ranks( + comm: Communicator, + device_mr: rmm.mr.CudaMemoryResource, + stream: Stream, +) -> None: + if comm.nranks < 2: + pytest.skip("Need at least two ranks") + + br = BufferResource(device_mr) + + if comm.rank == 0: + srcs = [] + dsts = [1] + elif comm.rank == 1: + srcs = [0] + dsts = [] + else: + srcs = [] + dsts = [] + + sparse_alltoall = SparseAlltoall( + comm=comm, + op_id=1, + br=br, + srcs=srcs, + dsts=dsts, + ) + + if comm.rank == 0: + sparse_alltoall.insert(1, generate_packed_data(1, 11, stream, br)) + sparse_alltoall.insert(1, generate_packed_data(1, 29, stream, br)) + + sparse_alltoall.insert_finished() + sparse_alltoall.wait() + + 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 + ) + ] + ), + ) + + +def test_invalid_peers_raise( + comm: Communicator, + device_mr: rmm.mr.CudaMemoryResource, +) -> None: + rank = comm.rank + size = comm.nranks + br = BufferResource(device_mr) + for src, dst in [([], [rank]), ([rank], []), ([], [size]), ([size], [])]: + with pytest.raises( + IndexError, + match=r"SparseAlltoall invalid (source|destination) rank", + ): + SparseAlltoall(comm, 1, br=br, srcs=src, dsts=dst) + if size > 1: + peer = (rank + 1) % size + for src, dst in [([], [peer, peer]), ([peer, peer], [])]: + with pytest.raises( + ValueError, + match=r"SparseAlltoall (source|destination) rank list must be unique", + ): + SparseAlltoall(comm, 1, br=br, srcs=src, dsts=dst) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_streaming_allgather.py b/python/cudf_streaming/cudf_streaming/tests/test_streaming_allgather.py new file mode 100644 index 000000000000..5206bf0beb1c --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_streaming_allgather.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +from contextlib import nullcontext +from typing import TYPE_CHECKING + +import numpy as np +import pylibcudf as plc +import pytest + +from cudf_streaming.integrations.partition import ( + packed_data_from_cudf_packed_columns, + unpack_and_concat, +) +from cudf_streaming.streaming.table_chunk import TableChunk +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 + +if TYPE_CHECKING: + from collections.abc import Awaitable + + 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 + + +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( + packed_data_from_cudf_packed_columns( + plc.contiguous_split.pack(table, stream=stream), + stream, + context.br(), + ), + br=context.br(), + ) + for table in input_tables + ] + actors = [] + + ch1: Channel[PackedDataChunk] = context.create_channel() + actors.append( + push_to_channel( + context, ch1, [Message(i, chunk) for i, chunk in enumerate(inputs)] + ) + ) + + ch2: Channel[PackedDataChunk] = context.create_channel() + actors.append(allgather(context, comm, ch1, ch2, op_id, ordered=True)) + + actor, deferred = pull_from_channel(context, ch2) + 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) + + +@define_actor() +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( + packed_data_from_cudf_packed_columns( + plc.contiguous_split.pack(table, stream=stream), + stream, + context.br(), + ), + br=context.br(), + ), + ) + await ch.send(context, msg) + await ch.drain(context) + + +@define_actor() +async def allgather_and_concat( + context: Context, + comm: Communicator, + ch_in: Channel[PackedDataChunk], + ch_out: Channel[TableChunk], + op_id: int, + use_context_manager: bool, +) -> None: + gather = AllGather(context, comm, op_id) + cm = gather if use_context_manager else nullcontext(gather) + with cm as ag: + while (msg := await ch_in.recv(context)) is not None: + chunk = PackedDataChunk.from_message( + msg, br=context.br() + ).to_packed_data() + ag.insert(msg.sequence_number, chunk) + 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)) + await ch_out.drain(context) + + +@pytest.mark.parametrize( + "use_context_manager", [True, False], ids=["context", "non-context"] +) +def test_allgather_object_interface( + context: Context, + comm: Communicator, + use_context_manager: bool, +) -> None: + if comm.nranks != 1: + pytest.skip("Only support single-rank runs") + + ch_in: Channel[PackedDataChunk] = context.create_channel() + ch_out: Channel[TableChunk] = 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 + ) + ) + + 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) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_streaming_define_actor.py b/python/cudf_streaming/cudf_streaming/tests/test_streaming_define_actor.py new file mode 100644 index 000000000000..fca157f09c0a --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_streaming_define_actor.py @@ -0,0 +1,182 @@ +# 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 pylibcudf as plc +import pytest + +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 + + +@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) + ] + + +if TYPE_CHECKING: + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + from rmm.pylibrmm.stream import Stream + + +def test_send_table_chunks( + context: Context, stream: Stream, expects: list[plc.Table] +) -> None: + ch1: Channel[TableChunk] = context.create_channel() + + # The actor access `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) + + actor2, output = pull_from_channel(context, ch_in=ch1) + + run_actor_network( + context, + actors=[ + actor1(context, ch_out=ch1), + 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) + + +def test_shutdown(context: Context) -> None: + @define_actor() + async def actor1(ctx: Context, ch_out: Channel[TableChunk]) -> None: + await ch_out.shutdown(ctx) + # Calling shutdown multiple times is allowed. + await ch_out.shutdown(ctx) + + ch1: Channel[TableChunk] = context.create_channel() + actor2, output = pull_from_channel(context, ch_in=ch1) + + run_actor_network( + context, + actors=[ + actor1(context, ch_out=ch1), + actor2, + ], + ) + + assert output.release() == [] + + +def test_send_error(context: Context) -> None: + @define_actor() + async def actor1(ctx: Context, ch_out: Channel[TableChunk]) -> None: + raise RuntimeError("MyError") + + ch1: Channel[TableChunk] = context.create_channel() + actor2, output = pull_from_channel(context, ch_in=ch1) + + with pytest.RaisesGroup( + pytest.RaisesExc( + RuntimeError, + match="MyError", + ) + ): + run_actor_network( + context, + actors=[ + actor1(context, ch_out=ch1), + actor2, + ], + ) + + assert output.release() == [] + + +def test_recv_table_chunks( + context: Context, stream: Stream, expects: list[plc.Table] +) -> None: + table_chunks = [ + Message( + seq, + TableChunk.from_pylibcudf_table( + expect, stream, exclusive_view=False, br=context.br() + ), + ) + for seq, expect in enumerate(expects) + ] + + results: list[Message[TableChunk]] = [] + + @define_actor() + async def actor1(ctx: Context, ch_in: Channel[TableChunk]) -> None: + while True: + chunk = await ch_in.recv(context) + if chunk is None: + break + results.append(chunk) + + ch1: Channel[TableChunk] = context.create_channel() + + run_actor_network( + context, + actors=[ + push_to_channel(context, ch_out=ch1, messages=table_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) + + +@pytest.mark.filterwarnings("error") +def test_unawaited_actor_closed_coroutines_no_warning( + context: Context, +) -> None: + ch: Channel[ArbitraryChunk[int]] = context.create_channel() + + @define_actor() + async def my_actor( + ctx: Context, ch_out: Channel[ArbitraryChunk[int]] + ) -> None: + await ch_out.send(ctx, Message(0, ArbitraryChunk(42))) + await ch_out.drain(ctx) + + # Never awaited, just verifying no RuntimeWarning is emitted + actor = my_actor(context, ch_out=ch) + del actor diff --git a/python/cudf_streaming/cudf_streaming/tests/test_streaming_fanout.py b/python/cudf_streaming/cudf_streaming/tests/test_streaming_fanout.py new file mode 100644 index 000000000000..1ac22589ca1e --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_streaming_fanout.py @@ -0,0 +1,174 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""Tests for streaming fanout actor.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pylibcudf as plc +import pytest + +from cudf_streaming.streaming.table_chunk import TableChunk +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 + +_INT64 = plc.DataType(plc.TypeId.INT64) + + +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), + ), + ] + ) + + +if TYPE_CHECKING: + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + from rmm.pylibrmm.stream import Stream + + +@pytest.mark.parametrize( + "policy", [FanoutPolicy.BOUNDED, FanoutPolicy.UNBOUNDED] +) +def test_fanout_basic( + context: Context, stream: Stream, 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() + + # 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)) + + # 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) + + +@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 +) -> None: + """Test fanout with varying numbers of output channels.""" + # Create channels + ch_in: Channel[TableChunk] = context.create_channel() + chs_out: list[Channel[TableChunk]] = [ + context.create_channel() for _ in range(num_outputs) + ] + + if num_outputs == 1: + with pytest.raises(ValueError): + 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)) + + # Create actors + push_actor = push_to_channel(context, ch_in, messages) + fanout_actor = fanout(context, ch_in, chs_out, policy) + pull_actors = [] + outputs = [] + for ch_out in chs_out: + pull_actor, output = pull_from_channel(context, ch_out) + 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 + + +def test_fanout_empty_outputs(context: Context, stream: Stream) -> None: + """Test fanout with empty output list raises value error.""" + ch_in: Channel[TableChunk] = context.create_channel() + with pytest.raises(ValueError): + fanout(context, ch_in, [], FanoutPolicy.BOUNDED) + + +def test_fanout_policy_enum() -> None: + """Test that FanoutPolicy enum has correct values.""" + assert FanoutPolicy.BOUNDED == 0 + assert FanoutPolicy.UNBOUNDED == 1 + assert len(FanoutPolicy) == 2 diff --git a/python/cudf_streaming/cudf_streaming/tests/test_streaming_leaf_actor.py b/python/cudf_streaming/cudf_streaming/tests/test_streaming_leaf_actor.py new file mode 100644 index 000000000000..d60b252c4cf0 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_streaming_leaf_actor.py @@ -0,0 +1,54 @@ +# 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 pylibcudf as plc + +from cudf_streaming.streaming.table_chunk import TableChunk +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 + +if TYPE_CHECKING: + from rapidsmpf.streaming.core.channel import Channel + from rapidsmpf.streaming.core.context import Context + from rmm.pylibrmm.stream import Stream + + +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) + ] + ch1: Channel[TableChunk] = context.create_channel() + actor1 = push_to_channel(context, ch_out=ch1, messages=table_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) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_streaming_shuffler.py b/python/cudf_streaming/cudf_streaming/tests/test_streaming_shuffler.py new file mode 100644 index 000000000000..43fb6751f691 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_streaming_shuffler.py @@ -0,0 +1,303 @@ +# 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 cupy as cp +import numpy as np +import pylibcudf as plc +import pytest + +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 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 + +if TYPE_CHECKING: + from collections.abc import Awaitable + + 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 + from rmm.pylibrmm.stream import Stream + + +@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( + 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[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, + ) + ) + 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] + + nranks = comm.nranks + rank = comm.rank + expected_local = list( + range( + rank * num_partitions // nranks, + (rank + 1) * num_partitions // nranks, + ) + ) + assert set(received_pids) == set(expected_local) + assert len(received_pids) == len(expected_local) + + +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[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) diff --git a/python/cudf_streaming/cudf_streaming/tests/test_streaming_sparse_alltoall.py b/python/cudf_streaming/cudf_streaming/tests/test_streaming_sparse_alltoall.py new file mode 100644 index 000000000000..78df8b3adc78 --- /dev/null +++ b/python/cudf_streaming/cudf_streaming/tests/test_streaming_sparse_alltoall.py @@ -0,0 +1,98 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import asyncio +from typing import TYPE_CHECKING + +import numpy as np +import pylibcudf as plc +import pytest + +from cudf_streaming.integrations.partition import ( + packed_data_from_cudf_packed_columns, + unpack_and_concat, +) +from rapidsmpf.streaming.coll.sparse_alltoall import SparseAlltoall +from rapidsmpf.testing import assert_eq + +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 packed_data_from_cudf_packed_columns( + plc.contiguous_split.pack(table, stream=stream), + stream, + context.br(), + ) + + +def unpack_table(context: Context, packed_data: PackedData) -> plc.Table: + stream = context.get_stream_from_pool() + return unpack_and_concat([packed_data], stream, context.br()) + + +def test_sparse_alltoall_non_participating_ranks( + context: Context, + comm: Communicator, +) -> None: + if comm.nranks < 2: + pytest.skip("Need at least two ranks") + if comm.rank == 0: + srcs = [] + dsts = [1] + elif comm.rank == 1: + srcs = [0] + dsts = [] + else: + srcs = [] + dsts = [] + + exchange = SparseAlltoall( + context, + comm, + 0, + srcs, + dsts, + ) + + 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)) + ) + + 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 + ) + ] + ), + ) diff --git a/python/cudf_streaming/pyproject.toml b/python/cudf_streaming/pyproject.toml index fc31a16dae1c..f7817937ff0a 100644 --- a/python/cudf_streaming/pyproject.toml +++ b/python/cudf_streaming/pyproject.toml @@ -80,10 +80,12 @@ requires = [ "libcudf-streaming==26.8.*,>=0.0.0a0", "librapidsmpf==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", + "libucxx==0.51.*,>=0.0.0a0", "ninja", "pylibcudf==26.8.*,>=0.0.0a0", "rapidsmpf==26.8.*,>=0.0.0a0", "rmm==26.8.*,>=0.0.0a0", + "ucxx==0.51.*,>=0.0.0a0", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. [tool.scikit-build] diff --git a/python/libcudf_streaming/pyproject.toml b/python/libcudf_streaming/pyproject.toml index 21580a0de4c4..23549d7408c4 100644 --- a/python/libcudf_streaming/pyproject.toml +++ b/python/libcudf_streaming/pyproject.toml @@ -71,5 +71,6 @@ requires = [ "libcudf==26.8.*,>=0.0.0a0", "librapidsmpf==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", + "libucxx==0.51.*,>=0.0.0a0", "ninja", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.