diff --git a/ci/build_docs.sh b/ci/build_docs.sh index de370189c..c54898b9b 100755 --- a/ci/build_docs.sh +++ b/ci/build_docs.sh @@ -50,7 +50,6 @@ set +e rapids-logger "Build CPP docs" pushd cpp/doxygen aws s3 cp s3://rapidsai-docs/librmm/html/"${RAPIDS_VERSION_MAJOR_MINOR}"/rmm.tag . || echo "Failed to download rmm Doxygen tag" -aws s3 cp s3://rapidsai-docs/libcudf/html/"${RAPIDS_VERSION_MAJOR_MINOR}"/libcudf.tag . || echo "Failed to download cudf Doxygen tag" doxygen Doxyfile diff --git a/ci/build_wheel_librapidsmpf.sh b/ci/build_wheel_librapidsmpf.sh index 42fb5f8cc..da29c1e1c 100755 --- a/ci/build_wheel_librapidsmpf.sh +++ b/ci/build_wheel_librapidsmpf.sh @@ -37,9 +37,6 @@ export SITE_PACKAGES ./ci/build_wheel.sh "${package_name}" "${package_dir}" python -m auditwheel repair \ - --exclude libcudf.so \ - --exclude libkvikio.so \ - --exclude libnvcomp.so.5 \ --exclude libnvidia-ml.so.1 \ --exclude librapids_logger.so \ --exclude librmm.so \ diff --git a/ci/build_wheel_rapidsmpf.sh b/ci/build_wheel_rapidsmpf.sh index d582e9f2d..11dceab71 100755 --- a/ci/build_wheel_rapidsmpf.sh +++ b/ci/build_wheel_rapidsmpf.sh @@ -46,7 +46,6 @@ export RAPIDS_PY_API ./ci/build_wheel.sh "${package_name}" "${package_dir}" --stable python -m auditwheel repair \ - --exclude libcudf.so \ --exclude libnvidia-ml.so.1 \ --exclude librapids_logger.so \ --exclude librmm.so \ diff --git a/ci/build_wheel_singlecomm.sh b/ci/build_wheel_singlecomm.sh index 703657f77..1873cb87c 100755 --- a/ci/build_wheel_singlecomm.sh +++ b/ci/build_wheel_singlecomm.sh @@ -44,9 +44,6 @@ export SITE_PACKAGES ./ci/build_wheel.sh "${package_name}" "${package_dir}" python -m auditwheel repair \ - --exclude libcudf.so \ - --exclude libkvikio.so \ - --exclude libnvcomp.so.5 \ --exclude libnvidia-ml.so.1 \ --exclude librapids_logger.so \ --exclude librmm.so \ @@ -86,9 +83,6 @@ export RAPIDS_PY_API ./ci/build_wheel.sh "${package_name_py}" "${package_dir_py}" --stable python -m auditwheel repair \ - --exclude libcudf.so \ - --exclude libkvikio.so \ - --exclude libnvcomp.so.5 \ --exclude libnvidia-ml.so.1 \ --exclude librapids_logger.so \ --exclude librmm.so \ diff --git a/ci/release/update-version.sh b/ci/release/update-version.sh index f12b25a8f..4d83a4d36 100755 --- a/ci/release/update-version.sh +++ b/ci/release/update-version.sh @@ -105,12 +105,9 @@ echo "${NEXT_UCXX_SHORT_TAG}.00" > UCXX_VERSION sed_runner "s/ucxx==.*/ucxx==${NEXT_UCXX_SHORT_TAG_PEP440}.*,>=0.0.0a0/g" dependencies.yaml DEPENDENCIES=( - cudf - libcudf librapidsmpf librapidsmpf-tests librmm - pylibcudf rapidsmpf rmm ) diff --git a/cmake/Modules/ConfigureCUDA.cmake b/cmake/Modules/ConfigureCUDA.cmake index a73acbb18..eb8b5b92f 100644 --- a/cmake/Modules/ConfigureCUDA.cmake +++ b/cmake/Modules/ConfigureCUDA.cmake @@ -1,6 +1,6 @@ # ================================================================================= # cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 # cmake-format: on # ================================================================================= @@ -31,7 +31,7 @@ list( RAPIDSMPF_CUDA_FLAGS -Xcompiler=-Wall,-Werror,-Wextra,-Wsign-conversion,-Wno-unknown-pragmas,-Wno-missing-field-initializers,-Wno-error=deprecated-declarations ) -# This warning needs to be suppressed because some parts of cudf instantiate templated CCCL +# This warning needs to be suppressed because some parts of rapidsmpf instantiate templated CCCL # functions in contexts where the resulting instantiations would have internal linkage (e.g. in # anonymous namespaces). In such contexts, the visibility attribute on the template is ignored, and # the compiler issues a warning. This is not a problem and will be fixed in future versions of CCCL. diff --git a/cmake/thirdparty/get_cucascade.cmake b/cmake/thirdparty/get_cucascade.cmake index 776ca51ea..da0fd6208 100644 --- a/cmake/thirdparty/get_cucascade.cmake +++ b/cmake/thirdparty/get_cucascade.cmake @@ -11,16 +11,16 @@ function(find_and_configure_cucascade) rapids_cpm_find( cuCascade 0.1.0 - GLOBAL_TARGETS cuCascade::cucascade cuCascade::cucascade_topology_discovery + GLOBAL_TARGETS cuCascade::cucascade_topology_discovery CPM_ARGS GIT_REPOSITORY https://github.com/NVIDIA/cuCascade.git GIT_TAG main - GIT_SHALLOW FALSE OPTIONS "CUCASCADE_BUILD_TESTS OFF" "CUCASCADE_BUILD_BENCHMARKS OFF" "CUCASCADE_BUILD_SHARED_LIBS OFF" "CUCASCADE_BUILD_STATIC_LIBS ON" "CUCASCADE_WARNINGS_AS_ERRORS OFF" + "CUCASCADE_TOPOLOGY_ONLY ON" EXCLUDE_FROM_ALL ) endfunction() diff --git a/cmake/thirdparty/get_cudf.cmake b/cmake/thirdparty/get_cudf.cmake deleted file mode 100644 index 2bf7164c7..000000000 --- a/cmake/thirdparty/get_cudf.cmake +++ /dev/null @@ -1,48 +0,0 @@ -# ================================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ================================================================================= - -# This function finds cudf and sets any additional necessary environment variables. -function(find_and_configure_cudf) - - if(TARGET cudf::cudf) - return() - endif() - - set(oneValueArgs VERSION GIT_REPO GIT_TAG USE_CUDF_STATIC EXCLUDE_FROM_ALL - PER_THREAD_DEFAULT_STREAM - ) - cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) - - set(cudf_global_targets cudf::cudf) - set(cudf_components "") - - if(BUILD_TESTS) - list(APPEND cudf_global_targets cudf::cudftestutil) - set(cudf_components COMPONENTS testing) - endif() - - rapids_cpm_find( - cudf ${PKG_VERSION} ${cudf_components} - GLOBAL_TARGETS ${cudf_global_targets} - BUILD_EXPORT_SET rapidsmpf-exports - INSTALL_EXPORT_SET rapidsmpf-exports - CPM_ARGS - GIT_REPOSITORY ${PKG_GIT_REPO} - GIT_TAG ${PKG_GIT_TAG} - GIT_SHALLOW TRUE SOURCE_SUBDIR cpp - OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_SHARED_LIBS ON" - "CUDF_BUILD_TESTUTIL ${BUILD_TESTS}" "CUDF_BUILD_STREAMS_TEST_UTIL OFF" - ) - - if(TARGET cudf) - set_property(TARGET cudf PROPERTY SYSTEM TRUE) - endif() -endfunction() -find_and_configure_cudf( - VERSION ${RAPIDS_VERSION} GIT_REPO https://github.com/rapidsai/cudf.git GIT_TAG - "${RAPIDS_BRANCH}" -) diff --git a/cmake/thirdparty/get_cudf_streaming.cmake b/cmake/thirdparty/get_cudf_streaming.cmake new file mode 100644 index 000000000..ad493298c --- /dev/null +++ b/cmake/thirdparty/get_cudf_streaming.cmake @@ -0,0 +1,33 @@ +# ================================================================================= +# cmake-format: off +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 +# cmake-format: on +# ================================================================================= + +# This function finds cudf_streaming for test/benchmark use only. It does NOT add cudf_streaming to +# any rapidsmpf export set. +function(find_and_configure_cudf_streaming) + + if(TARGET cudf_streaming::cudf_streaming) + return() + endif() + + set(oneValueArgs VERSION GIT_REPO GIT_TAG) + cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + + rapids_cpm_find( + cudf_streaming ${PKG_VERSION} + GLOBAL_TARGETS cudf_streaming::cudf_streaming + CPM_ARGS + GIT_REPOSITORY ${PKG_GIT_REPO} + GIT_TAG ${PKG_GIT_TAG} + GIT_SHALLOW TRUE SOURCE_SUBDIR cpp/libcudf_streaming + OPTIONS "BUILD_TESTS OFF" "BUILD_BENCHMARKS OFF" "BUILD_EXAMPLES OFF" + ) +endfunction() + +find_and_configure_cudf_streaming( + VERSION ${RAPIDS_VERSION} GIT_REPO https://github.com/rapidsai/cudf.git GIT_TAG + "${RAPIDS_BRANCH}" +) diff --git a/conda/environments/all_cuda-129_arch-aarch64.yaml b/conda/environments/all_cuda-129_arch-aarch64.yaml index 5833e2ad6..5184df872 100644 --- a/conda/environments/all_cuda-129_arch-aarch64.yaml +++ b/conda/environments/all_cuda-129_arch-aarch64.yaml @@ -28,7 +28,7 @@ dependencies: - gcc_linux-aarch64=14.* - gdb - ipython -- libcudf==26.8.*,>=0.0.0a0 +- libcudf-streaming==26.8.*,>=0.0.0a0 - libnuma - libpmix-devel >=5.0,<6.0 - librmm==26.8.*,>=0.0.0a0 @@ -47,7 +47,6 @@ dependencies: - psutil - pyarrow - pydata-sphinx-theme>=0.15.4 -- pylibcudf==26.8.*,>=0.0.0a0 - pytest - python>=3.11 - rapids-build-backend>=0.4.0,<0.5.0 diff --git a/conda/environments/all_cuda-129_arch-x86_64.yaml b/conda/environments/all_cuda-129_arch-x86_64.yaml index 4ea0682da..12271fda4 100644 --- a/conda/environments/all_cuda-129_arch-x86_64.yaml +++ b/conda/environments/all_cuda-129_arch-x86_64.yaml @@ -28,7 +28,7 @@ dependencies: - gcc_linux-64=14.* - gdb - ipython -- libcudf==26.8.*,>=0.0.0a0 +- libcudf-streaming==26.8.*,>=0.0.0a0 - libnuma - libpmix-devel >=5.0,<6.0 - librmm==26.8.*,>=0.0.0a0 @@ -47,7 +47,6 @@ dependencies: - psutil - pyarrow - pydata-sphinx-theme>=0.15.4 -- pylibcudf==26.8.*,>=0.0.0a0 - pytest - python>=3.11 - rapids-build-backend>=0.4.0,<0.5.0 diff --git a/conda/environments/all_cuda-132_arch-aarch64.yaml b/conda/environments/all_cuda-132_arch-aarch64.yaml index 9af102f64..ea4435c41 100644 --- a/conda/environments/all_cuda-132_arch-aarch64.yaml +++ b/conda/environments/all_cuda-132_arch-aarch64.yaml @@ -28,7 +28,7 @@ dependencies: - gcc_linux-aarch64=14.* - gdb - ipython -- libcudf==26.8.*,>=0.0.0a0 +- libcudf-streaming==26.8.*,>=0.0.0a0 - libnuma - libpmix-devel >=5.0,<6.0 - librmm==26.8.*,>=0.0.0a0 @@ -47,7 +47,6 @@ dependencies: - psutil - pyarrow - pydata-sphinx-theme>=0.15.4 -- pylibcudf==26.8.*,>=0.0.0a0 - pytest - python>=3.11 - rapids-build-backend>=0.4.0,<0.5.0 diff --git a/conda/environments/all_cuda-132_arch-x86_64.yaml b/conda/environments/all_cuda-132_arch-x86_64.yaml index eace9b6de..f5cde57c7 100644 --- a/conda/environments/all_cuda-132_arch-x86_64.yaml +++ b/conda/environments/all_cuda-132_arch-x86_64.yaml @@ -28,7 +28,7 @@ dependencies: - gcc_linux-64=14.* - gdb - ipython -- libcudf==26.8.*,>=0.0.0a0 +- libcudf-streaming==26.8.*,>=0.0.0a0 - libnuma - libpmix-devel >=5.0,<6.0 - librmm==26.8.*,>=0.0.0a0 @@ -47,7 +47,6 @@ dependencies: - psutil - pyarrow - pydata-sphinx-theme>=0.15.4 -- pylibcudf==26.8.*,>=0.0.0a0 - pytest - python>=3.11 - rapids-build-backend>=0.4.0,<0.5.0 diff --git a/conda/recipes/librapidsmpf/recipe.yaml b/conda/recipes/librapidsmpf/recipe.yaml index 3bd37d121..111af0c7e 100644 --- a/conda/recipes/librapidsmpf/recipe.yaml +++ b/conda/recipes/librapidsmpf/recipe.yaml @@ -29,7 +29,7 @@ cache: export CXXFLAGS=$(echo $CXXFLAGS | sed -E 's@\-fdebug\-prefix\-map[^ ]*@@g') set +x - ./build.sh -v -n --cmake-args=\"-DBUILD_CUPTI_SUPPORT=ON\" --no-clang-tidy librapidsmpf + ./build.sh -v -n --cmake-args="\"-DBUILD_CUPTI_SUPPORT=ON -DBUILD_CUDF_TESTS=ON\"" --no-clang-tidy librapidsmpf secrets: - AWS_ACCESS_KEY_ID - AWS_SECRET_ACCESS_KEY @@ -77,8 +77,9 @@ cache: - cuda-version =${{ cuda_version }} - cuda-cupti-dev - cuda-nvml-dev - - libcudf =${{ minor_version }} - libnuma + - libcudf =${{ minor_version }} + - libcudf-streaming =${{ minor_version }} - libpmix-devel >=5.0,<6.0 - librmm =${{ minor_version }} - libucxx ${{ ucxx_version }} @@ -92,7 +93,6 @@ outputs: script: content: | cmake --install cpp/build - cmake --install cpp/build --component=benchmarking cmake --install cpp/build --component=tools dynamic_linking: overlinking_behavior: "error" @@ -102,7 +102,6 @@ outputs: ignore: # See https://github.com/rapidsai/build-planning/issues/160 - lib/librapidsmpf.so - - lib/librapidsmpfndsh.so string: cuda${{ cuda_major }}_${{ date_string }}_${{ head_rev }} requirements: build: @@ -114,15 +113,14 @@ outputs: - cuda-version =${{ cuda_version }} - cuda-cupti-dev - cuda-nvml-dev - - libcudf =${{ minor_version }} - libpmix-devel >=5.0,<6.0 - libucxx ${{ ucxx_version }} - openmpi >=5.0 + - libnuma run: - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} - cuda-cupti - librmm =${{ minor_version }} - - libcudf =${{ minor_version }} - libpmix >=5.0,<6.0 - libucxx ${{ ucxx_version }} - openmpi >=5.0 # See @@ -135,8 +133,9 @@ outputs: by_name: - cuda-cupti - libcudf - - openmpi + - libcudf-streaming - librmm + - openmpi about: homepage: https://github.com/rapidsai/rapidsmpf license: Apache-2.0 @@ -148,7 +147,8 @@ outputs: build: script: content: | - cmake --install cpp/build --component testing + cmake --install cpp/build --component=testing + cmake --install cpp/build --component=benchmarking string: cuda${{ cuda_major }}_${{ date_string }}_${{ head_rev }} requirements: build: @@ -161,14 +161,17 @@ outputs: - cuda-nvml-dev - cuda-version =${{ cuda_version }} - libcudf =${{ minor_version }} + - libcudf-streaming =${{ minor_version }} - librmm =${{ minor_version }} - openmpi >=5.0 + - libnuma - ${{ pin_subpackage("librapidsmpf", exact=True) }} run: - ${{ pin_compatible("cuda-version", upper_bound="x", lower_bound="x") }} - cuda-cupti - - librmm =${{ minor_version }} - libcudf =${{ minor_version }} + - libcudf-streaming =${{ minor_version }} + - librmm =${{ minor_version }} - openmpi >=5.0 # See - libucxx ${{ ucxx_version }} ignore_run_exports: @@ -180,7 +183,6 @@ outputs: by_name: - cuda-cupti - librmm - - libcudf - openmpi about: homepage: https://github.com/rapidsai/rapidsmpf diff --git a/conda/recipes/rapidsmpf/recipe.yaml b/conda/recipes/rapidsmpf/recipe.yaml index 62e167a5b..cf408a96d 100644 --- a/conda/recipes/rapidsmpf/recipe.yaml +++ b/conda/recipes/rapidsmpf/recipe.yaml @@ -92,7 +92,6 @@ requirements: - mpi4py - openmpi >=5.0 - pip - - pylibcudf =${{ minor_version }} - python =${{ py_abi_min }} - python-abi3 ${{ py_abi_min }}.* - rapids-build-backend >=0.4.0,<0.5.0 @@ -108,7 +107,6 @@ requirements: - librapidsmpf =${{ version }} - mpi4py - openmpi >=5.0 - - pylibcudf =${{ minor_version }} - python - rmm =${{ minor_version }} - ucxx ${{ ucxx_version }} diff --git a/cpp/.clang-format b/cpp/.clang-format index 608db414d..f1e09bf26 100644 --- a/cpp/.clang-format +++ b/cpp/.clang-format @@ -58,7 +58,7 @@ IncludeCategories: Priority: 30 # CUDA includes - Regex: '^<(thrust|cub|cuda)/' Priority: 40 # CCCL includes - - Regex: '^<(cudf.*|rmm|cugraph|cuml|raft|kvikio|cucascade)' + - Regex: '^<(cudf|rmm|cugraph|cuml|raft|kvikio|cucascade)' Priority: 50 # RAPIDS includes - Regex: '^ - $ CUDA::cudart_static + PUBLIC rmm::rmm CCCL::CCCL $ $ + CUDA::cudart_static PRIVATE cuco::cuco cuCascade::cucascade_topology_discovery $<$:numa> @@ -342,7 +339,13 @@ target_compile_definitions( rapids_cuda_set_runtime(rapidsmpf USE_STATIC ON) -add_library(rapidsmpf::rapidsmpf ALIAS rapidsmpf) +# Guard: when BUILD_CUDF_TESTS is enabled, cudf_streaming is fetched via CPM and its own +# cmake/thirdparty/get_rapidsmpf.cmake calls rapids_cpm_find(rapidsmpf ...) which creates this alias +# target. Without the guard we get "cannot create ALIAS target because another target with the same +# name already exists" due to the circular dependency (rapidsmpf -> cudf_streaming -> rapidsmpf). +if(NOT TARGET rapidsmpf::rapidsmpf) + add_library(rapidsmpf::rapidsmpf ALIAS rapidsmpf) +endif() # ################################################################################################## # * linter configuration --------------------------------------------------------------------------- diff --git a/cpp/asan_suppressions.txt b/cpp/asan_suppressions.txt index b78fb1cb1..f77eaab5a 100644 --- a/cpp/asan_suppressions.txt +++ b/cpp/asan_suppressions.txt @@ -11,13 +11,6 @@ # Append `:print_suppressions=0` to LSAN_OPTIONS to silence the # "Suppressions used:" footer. -# cuDF process-lifetime singletons. Note: event_for_thread is defined in an -# anonymous namespace, which gets mangled into the symbolicated name between -# `cudf::detail::` and `event_for_thread`, so we match the leaf name only -# since LSan suppression matching is plain substring. -leak:event_for_thread -leak:cudf::detail::create_global_cuda_stream_pool - # RMM process-lifetime stream pool / streams owned by the global pool MR. leak:rmm::cuda_stream_pool::cuda_stream_pool leak:rmm::cuda_stream::cuda_stream diff --git a/cpp/benchmarks/CMakeLists.txt b/cpp/benchmarks/CMakeLists.txt index b94d5f1dd..5b4d715e8 100644 --- a/cpp/benchmarks/CMakeLists.txt +++ b/cpp/benchmarks/CMakeLists.txt @@ -5,34 +5,6 @@ # cmake-format: on # ================================================================================= -add_library(bench_utils INTERFACE) -target_sources(bench_utils INTERFACE utils/random_data.cu) -add_executable(bench_shuffle "bench_shuffle.cpp") -set_target_properties( - bench_shuffle - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - # For std:: support of __int128_t. Can be removed once using cuda::std - CXX_EXTENSIONS ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON -) -target_compile_options( - bench_shuffle PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" -) -target_link_libraries( - bench_shuffle PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx $ - $ maybe_asan bench_utils -) -install( - TARGETS bench_shuffle - COMPONENT benchmarking - DESTINATION bin/benchmarks/librapidsmpf - EXCLUDE_FROM_ALL -) - add_executable(bench_comm "bench_comm.cpp") set_target_properties( bench_comm @@ -50,7 +22,7 @@ target_compile_options( ) target_link_libraries( bench_comm PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx $ - $ maybe_asan bench_utils + $ maybe_asan ) install( TARGETS bench_comm @@ -63,39 +35,6 @@ install( include(${rapids-cmake-dir}/cpm/gbench.cmake) rapids_cpm_gbench(BUILD_STATIC) -add_executable(bench_partition "bench_partition.cpp") -set_target_properties( - bench_partition - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - # For std:: support of __int128_t. Can be removed once using cuda::std - CXX_EXTENSIONS ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON -) -target_compile_options( - bench_partition PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" -) -target_link_libraries( - bench_partition - PRIVATE rapidsmpf::rapidsmpf - ucxx::ucxx - benchmark::benchmark - benchmark::benchmark_main - $ - $ - maybe_asan - bench_utils -) -install( - TARGETS bench_partition - COMPONENT benchmarking - DESTINATION bin/benchmarks/librapidsmpf - EXCLUDE_FROM_ALL -) - add_executable(bench_memory_resources "bench_memory_resources.cpp") set_target_properties( bench_memory_resources @@ -113,7 +52,7 @@ target_compile_options( target_link_libraries( bench_memory_resources PRIVATE rapidsmpf::rapidsmpf rmm::rmm benchmark::benchmark benchmark::benchmark_main - $ maybe_asan bench_utils + $ maybe_asan ) install( TARGETS bench_memory_resources @@ -122,32 +61,102 @@ install( EXCLUDE_FROM_ALL ) -add_executable(bench_pack "bench_pack.cpp") -set_target_properties( - bench_pack - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - CXX_EXTENSIONS ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON - LINK_FLAGS "-Wl,--allow-shlib-undefined" -) -target_compile_options( - bench_pack PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" -) -target_link_libraries( - bench_pack PRIVATE rapidsmpf::rapidsmpf benchmark::benchmark benchmark::benchmark_main - $ maybe_asan bench_utils -) -install( - TARGETS bench_pack - COMPONENT benchmarking - DESTINATION bin/benchmarks/librapidsmpf - EXCLUDE_FROM_ALL -) +if(BUILD_CUDF_TESTS) + add_library(bench_utils INTERFACE) + target_sources(bench_utils INTERFACE utils/random_data.cu) + + add_executable(bench_shuffle "bench_shuffle.cpp") + set_target_properties( + bench_shuffle + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + ) + target_compile_options( + bench_shuffle PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" + "$<$:${RAPIDSMPF_CUDA_FLAGS}>" + ) + target_link_libraries( + bench_shuffle + PRIVATE rapidsmpf::rapidsmpf + ucxx::ucxx + cudf_streaming::cudf_streaming + $ + $ + maybe_asan + bench_utils + ) + install( + TARGETS bench_shuffle + COMPONENT benchmarking + DESTINATION bin/benchmarks/librapidsmpf + EXCLUDE_FROM_ALL + ) + + add_executable(bench_partition "bench_partition.cpp") + set_target_properties( + bench_partition + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + ) + target_compile_options( + bench_partition PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" + "$<$:${RAPIDSMPF_CUDA_FLAGS}>" + ) + target_link_libraries( + bench_partition + PRIVATE rapidsmpf::rapidsmpf + ucxx::ucxx + cudf_streaming::cudf_streaming + benchmark::benchmark + benchmark::benchmark_main + $ + $ + maybe_asan + bench_utils + ) + install( + TARGETS bench_partition + COMPONENT benchmarking + DESTINATION bin/benchmarks/librapidsmpf + EXCLUDE_FROM_ALL + ) + + add_executable(bench_pack "bench_pack.cpp") + set_target_properties( + bench_pack + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/benchmarks" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + LINK_FLAGS "-Wl,--allow-shlib-undefined" + ) + target_compile_options( + bench_pack PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" + "$<$:${RAPIDSMPF_CUDA_FLAGS}>" + ) + target_link_libraries( + bench_pack + PRIVATE rapidsmpf::rapidsmpf cudf_streaming::cudf_streaming benchmark::benchmark + benchmark::benchmark_main $ maybe_asan bench_utils + ) + install( + TARGETS bench_pack + COMPONENT benchmarking + DESTINATION bin/benchmarks/librapidsmpf + EXCLUDE_FROM_ALL + ) +endif() -if(RAPIDSMPF_HAVE_STREAMING) +if(BUILD_CUDF_TESTS AND RAPIDSMPF_HAVE_STREAMING) add_subdirectory(streaming) endif() diff --git a/cpp/benchmarks/bench_comm.cpp b/cpp/benchmarks/bench_comm.cpp index afc70faf7..7e1a4c60a 100644 --- a/cpp/benchmarks/bench_comm.cpp +++ b/cpp/benchmarks/bench_comm.cpp @@ -10,6 +10,8 @@ #include #include +#include +#include #include #include @@ -29,7 +31,6 @@ #endif #include "utils/misc.hpp" -#include "utils/random_data.hpp" #include "utils/rmm_utils.hpp" @@ -220,7 +221,11 @@ Duration run( auto [res, _] = br->reserve(MemoryType::DEVICE, args.msg_size * 2, AllowOverbooking::YES); auto buf = br->make_buffer(args.msg_size, stream, res); - random_fill(*buf, br->device_mr()); + buf->write_access( + [size = args.msg_size](std::byte* ptr, rmm::cuda_stream_view s) { + RAPIDSMPF_CUDA_TRY(cudaMemsetAsync(ptr, 0x42, size, s.value())); + } + ); send_bufs.push_back(std::move(buf)); recv_bufs.push_back(br->make_buffer(args.msg_size, stream, res)); } @@ -312,11 +317,11 @@ int main(int argc, char** argv) { } auto& log = comm->logger(); - rmm::cuda_stream_view stream = cudf::get_default_stream(); + rmm::cuda_stream_view stream = rmm::cuda_stream_default; args.pprint(*comm); set_current_rmm_resource(args.rmm_mr); - rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref(); + rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource_ref(); auto br = BufferResource::create( mr, PinnedMemoryResource::Disabled, diff --git a/cpp/benchmarks/bench_partition.cpp b/cpp/benchmarks/bench_partition.cpp index 909744f24..779a91274 100644 --- a/cpp/benchmarks/bench_partition.cpp +++ b/cpp/benchmarks/bench_partition.cpp @@ -12,12 +12,12 @@ #include #include #include +#include #include #include #include #include -#include #include // Helper function to create a table with a single int column @@ -69,7 +69,7 @@ static void BM_PartitionAndPack(benchmark::State& state) { std::vector columns_to_hash{0}; for (auto _ : state) { - auto pack_partitions = rapidsmpf::partition_and_pack( + auto pack_partitions = cudf_streaming::integrations::partition_and_pack( *table, columns_to_hash, num_partitions, @@ -121,7 +121,7 @@ static void BM_PartitionAndPackCurrentImpl(benchmark::State& state) { for (auto _ : state) { for (int i = 0; i < num_partitions; i++) { - auto pack_partitions = rapidsmpf::partition_and_pack( + auto pack_partitions = cudf_streaming::integrations::partition_and_pack( *table, columns_to_hash, total_npartitions, diff --git a/cpp/benchmarks/bench_shuffle.cpp b/cpp/benchmarks/bench_shuffle.cpp index 36710c4fb..c6c918209 100644 --- a/cpp/benchmarks/bench_shuffle.cpp +++ b/cpp/benchmarks/bench_shuffle.cpp @@ -10,6 +10,8 @@ #include #include +#include + #include #include #include @@ -18,7 +20,6 @@ #include #include #include -#include #include #include #include @@ -318,8 +319,8 @@ rapidsmpf::Duration do_run( shuffler.wait(); for (auto finished_partition : shuffler.local_partitions()) { auto packed_chunks = shuffler.extract(finished_partition); - auto output_partition = rapidsmpf::unpack_and_concat( - rapidsmpf::unspill_partitions( + auto output_partition = cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::unspill_partitions( std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES ), stream, @@ -338,7 +339,7 @@ rapidsmpf::Duration do_run( // thus we only check large shuffles). if (args.num_local_rows >= 1000000) { for (const auto& output_partition : output_partitions) { - auto [parts, owner] = rapidsmpf::partition_and_split( + auto [parts, owner] = cudf_streaming::integrations::partition_and_split( output_partition->view(), {0}, static_cast(total_num_partitions), @@ -452,7 +453,7 @@ rapidsmpf::Duration run_hash_partition_inline( generate_input_partitions(args, stream, br, std::identity{}); auto make_chunk_fn = [&](cudf::table const& partition) { - return rapidsmpf::partition_and_pack( + return cudf_streaming::integrations::partition_and_pack( partition, {0}, static_cast(total_num_partitions), @@ -497,7 +498,7 @@ rapidsmpf::Duration run_hash_partition_with_datagen( std::vector> input_partitions = generate_input_partitions(args, stream, br, [&](cudf::table&& table) { - return rapidsmpf::partition_and_pack( + return cudf_streaming::integrations::partition_and_pack( table, {0}, static_cast(total_num_partitions), diff --git a/cpp/benchmarks/streaming/CMakeLists.txt b/cpp/benchmarks/streaming/CMakeLists.txt index c33871926..e8c6542ae 100644 --- a/cpp/benchmarks/streaming/CMakeLists.txt +++ b/cpp/benchmarks/streaming/CMakeLists.txt @@ -20,8 +20,13 @@ target_compile_options( ) target_link_libraries( bench_streaming_shuffle - PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx $ - $ maybe_asan bench_utils + PRIVATE rapidsmpf::rapidsmpf + ucxx::ucxx + cudf_streaming::cudf_streaming + $ + $ + maybe_asan + bench_utils ) install( TARGETS bench_streaming_shuffle diff --git a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp index 8b3b20cd8..9806b49ec 100644 --- a/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp +++ b/cpp/benchmarks/streaming/bench_streaming_shuffle.cpp @@ -10,6 +10,10 @@ #include #include +#include +#include +#include + #include #include #include @@ -18,7 +22,6 @@ #include #include #include -#include #include #include #include @@ -26,8 +29,6 @@ #include #include #include -#include -#include #include #include @@ -281,7 +282,7 @@ rapidsmpf::Duration run( ); auto ch2 = ctx->create_channel(); actors.push_back( - rapidsmpf::streaming::actor::partition_and_pack( + cudf_streaming::streaming::actor::partition_and_pack( ctx, ch1, ch2, @@ -298,7 +299,9 @@ rapidsmpf::Duration run( ) ); auto ch4 = ctx->create_channel(); - actors.push_back(rapidsmpf::streaming::actor::unpack_and_concat(ctx, ch3, ch4)); + 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(); diff --git a/cpp/benchmarks/streaming/data_generator.hpp b/cpp/benchmarks/streaming/data_generator.hpp index 9bb53795c..d72e3ca0c 100644 --- a/cpp/benchmarks/streaming/data_generator.hpp +++ b/cpp/benchmarks/streaming/data_generator.hpp @@ -7,11 +7,11 @@ #include #include +#include #include #include #include -#include #include #include "../utils/random_data.hpp" @@ -58,15 +58,17 @@ inline Actor random_table_generator( 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->send( + cudf_streaming::streaming::to_message( + seq, + std::make_unique( + std::make_unique(random_table( + ncolumns, nrows, min_val, max_val, stream, ctx->br()->device_mr() + )), + stream + ) ) - )); + ); } co_await ch_out->drain(ctx->executor()); } diff --git a/cpp/benchmarks/streaming/ndsh/CMakeLists.txt b/cpp/benchmarks/streaming/ndsh/CMakeLists.txt index 047c78086..45e0c96a0 100644 --- a/cpp/benchmarks/streaming/ndsh/CMakeLists.txt +++ b/cpp/benchmarks/streaming/ndsh/CMakeLists.txt @@ -34,8 +34,9 @@ target_compile_options( "$<$:${RAPIDSMPF_CUDA_FLAGS}>" ) target_link_libraries( - rapidsmpfndsh PRIVATE rapidsmpf::rapidsmpf cuco::cuco $ - $ maybe_asan + rapidsmpfndsh + PRIVATE rapidsmpf::rapidsmpf cudf_streaming::cudf_streaming cuco::cuco + $ $ maybe_asan ) set(RAPIDSMPFNDSH_QUERIES q01 q03 q04 q09 q21 bench_read) @@ -55,8 +56,9 @@ foreach(query IN ITEMS ${RAPIDSMPFNDSH_QUERIES}) "$<$:${RAPIDSMPF_CUDA_FLAGS}>" ) target_link_libraries( - ${query} PRIVATE rapidsmpfndsh rapidsmpf::rapidsmpf $ - $ maybe_asan + ${query} + PRIVATE rapidsmpfndsh rapidsmpf::rapidsmpf cudf_streaming::cudf_streaming + $ $ maybe_asan ) endforeach() diff --git a/cpp/benchmarks/streaming/ndsh/bench_read.cpp b/cpp/benchmarks/streaming/ndsh/bench_read.cpp index 7d75f9687..5e1125f9b 100644 --- a/cpp/benchmarks/streaming/ndsh/bench_read.cpp +++ b/cpp/benchmarks/streaming/ndsh/bench_read.cpp @@ -20,6 +20,7 @@ #include #include #include +#include #include #include @@ -31,7 +32,6 @@ #include #include #include -#include #include #include "utils.hpp" @@ -56,7 +56,7 @@ rapidsmpf::streaming::Actor read_parquet( if (columns.has_value()) { options.set_column_names(*columns); } - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -75,8 +75,8 @@ rapidsmpf::streaming::Actor consume_channel_parallel( if (msg.empty()) { break; } - if (msg.holds()) { - auto chunk = co_await msg.release() + if (msg.holds()) { + auto chunk = co_await msg.release() .make_available(ctx); ctx->logger()->print( "Consumed chunk with ", diff --git a/cpp/benchmarks/streaming/ndsh/concatenate.cpp b/cpp/benchmarks/streaming/ndsh/concatenate.cpp index 39ce71720..e80799800 100644 --- a/cpp/benchmarks/streaming/ndsh/concatenate.cpp +++ b/cpp/benchmarks/streaming/ndsh/concatenate.cpp @@ -11,13 +11,13 @@ #include #include #include +#include #include #include #include #include #include -#include namespace rapidsmpf::ndsh { @@ -43,9 +43,9 @@ streaming::Actor concatenate( } if (messages.size() == 0) { co_await ch_out->send( - streaming::to_message( + cudf_streaming::streaming::to_message( 0, - std::make_unique( + std::make_unique( std::make_unique(), concat_stream ) ) @@ -53,7 +53,7 @@ streaming::Actor concatenate( } else if (messages.size() == 1) { co_await ch_out->send(std::move(messages[0])); } else { - std::vector chunks; + std::vector chunks; std::vector views; if (order == ConcatOrder::LINEARIZE) { std::ranges::sort(messages, std::less{}, [](auto&& msg) { @@ -63,13 +63,13 @@ streaming::Actor concatenate( chunks.reserve(messages.size()); views.reserve(messages.size()); for (auto&& msg : messages) { - auto chunk = - co_await msg.release().make_available(ctx); + 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( + auto result = std::make_unique( cudf::concatenate(views, concat_stream, ctx->br()->device_mr()), concat_stream ); cuda_stream_join( @@ -78,7 +78,9 @@ streaming::Actor concatenate( &event ); chunks.clear(); - co_await ch_out->send(streaming::to_message(0, std::move(result))); + co_await ch_out->send( + cudf_streaming::streaming::to_message(0, std::move(result)) + ); } co_await ch_out->drain(ctx->executor()); } diff --git a/cpp/benchmarks/streaming/ndsh/groupby.cpp b/cpp/benchmarks/streaming/ndsh/groupby.cpp index 5a13054c2..6c6337218 100644 --- a/cpp/benchmarks/streaming/ndsh/groupby.cpp +++ b/cpp/benchmarks/streaming/ndsh/groupby.cpp @@ -12,11 +12,11 @@ #include #include #include +#include #include #include #include -#include namespace rapidsmpf::ndsh { @@ -35,7 +35,10 @@ streaming::Actor chunkwise_group_by( if (msg.empty()) { break; } - auto chunk = co_await msg.release().make_available(ctx); + auto chunk = + co_await msg.release().make_available( + ctx + ); auto stream = chunk.stream(); auto table = chunk.table_view(); auto agg_requests = std::vector(); @@ -62,9 +65,9 @@ streaming::Actor chunkwise_group_by( std::ranges::move(a.results, std::back_inserter(result)); } co_await ch_out->send( - streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( std::make_unique(std::move(result)), stream ) ) diff --git a/cpp/benchmarks/streaming/ndsh/join.cpp b/cpp/benchmarks/streaming/ndsh/join.cpp index 65ff8be52..30470066e 100644 --- a/cpp/benchmarks/streaming/ndsh/join.cpp +++ b/cpp/benchmarks/streaming/ndsh/join.cpp @@ -19,13 +19,14 @@ #include #include #include +#include +#include #include #include #include #include #include -#include #include #include #include @@ -33,7 +34,6 @@ #include #include #include -#include namespace rapidsmpf::ndsh { @@ -49,7 +49,7 @@ coro::task broadcast( CudaEvent event; comm->logger()->print("Broadcast ", static_cast(tag)); if (comm->nranks() == 1) { - std::vector chunks; + std::vector chunks; std::vector views; auto gather_stream = ctx->br()->stream_pool().get_stream(); while (true) { @@ -57,15 +57,18 @@ coro::task broadcast( if (msg.empty()) { break; } - auto chunk = - co_await msg.release().make_available(ctx); + 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 streaming::to_message( - 0, std::make_unique(std::move(chunks[0])) + co_return cudf_streaming::streaming::to_message( + 0, + std::make_unique( + std::move(chunks[0]) + ) ); } else { RAPIDSMPF_EXPECTS(chunks.size() > 0, "No chunks in broadcast"); @@ -78,9 +81,11 @@ coro::task broadcast( std::ranges::single_view(gather_stream), &event ); - co_return streaming::to_message( + co_return cudf_streaming::streaming::to_message( 0, - std::make_unique(std::move(result), gather_stream) + std::make_unique( + std::move(result), gather_stream + ) ); } } else { @@ -91,8 +96,8 @@ coro::task broadcast( break; } // TODO: If this chunk is already in pack form, this is unnecessary. - auto chunk = - co_await msg.release().make_available(ctx); + 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( @@ -104,19 +109,19 @@ coro::task broadcast( gatherer.insert_finished(); auto result = co_await gatherer.extract_all(ordered); if (result.size() == 1) { - co_return streaming::to_message( + co_return cudf_streaming::streaming::to_message( 0, - std::make_unique( + std::make_unique( std::make_unique(std::move(result[0])) ) ); } else { auto stream = ctx->br()->stream_pool().get_stream(); - co_return streaming::to_message( + co_return cudf_streaming::streaming::to_message( 0, - std::make_unique( - unpack_and_concat( - unspill_partitions( + std::make_unique( + cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::unspill_partitions( std::move(result), ctx->br().get(), AllowOverbooking::YES ), stream, @@ -159,8 +164,8 @@ streaming::Actor broadcast( */ streaming::Message semi_join_chunk( std::shared_ptr ctx, - streaming::TableChunk const& left_chunk, - streaming::TableChunk&& right_chunk, + 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, @@ -205,9 +210,11 @@ streaming::Message semi_join_chunk( // Deallocation of the join indices will happen on chunk_stream, so add stream dep cuda_stream_join(left_chunk.stream(), chunk_stream); - return streaming::to_message( + return cudf_streaming::streaming::to_message( sequence, - std::make_unique(std::move(result_table), chunk_stream) + std::make_unique( + std::move(result_table), chunk_stream + ) ); } @@ -229,7 +236,7 @@ streaming::Message semi_join_chunk( */ streaming::Message inner_join_chunk( std::shared_ptr ctx, - streaming::TableChunk&& right_chunk, + cudf_streaming::streaming::TableChunk&& right_chunk, std::uint64_t sequence, cudf::hash_join& joiner, cudf::table_view build_carrier, @@ -282,9 +289,9 @@ streaming::Message inner_join_chunk( // 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 streaming::to_message( + return cudf_streaming::streaming::to_message( sequence, - std::make_unique( + std::make_unique( std::make_unique(std::move(result_columns)), chunk_stream ) ); @@ -307,7 +314,7 @@ streaming::Actor inner_join_broadcast( 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() + .release() .make_available(ctx) ); comm->logger()->print( @@ -341,7 +348,7 @@ streaming::Actor inner_join_broadcast( } co_await ch_out->send(inner_join_chunk( ctx, - right_msg.release(), + right_msg.release(), right_msg.sequence_number(), joiner, build_carrier, @@ -386,7 +393,8 @@ streaming::Actor inner_join_shuffle( ); // TODO: currently always using left as build table. auto build_chunk = - co_await left_msg.release().make_available(ctx); + 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), @@ -408,7 +416,7 @@ streaming::Actor inner_join_shuffle( } co_await ch_out->send(inner_join_chunk( ctx, - right_msg.release(), + right_msg.release(), left_msg.sequence_number(), joiner, build_carrier, @@ -436,7 +444,7 @@ streaming::Actor left_semi_join_broadcast_left( 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() + .release() .make_available(ctx); comm->logger()->print( "Left (probe) table has ", left_table.table_view().num_rows(), " rows" @@ -467,7 +475,8 @@ streaming::Actor left_semi_join_broadcast_left( // 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 right_msg.release() + .make_available(ctx); co_await ch_out->send(semi_join_chunk( ctx, left_table, @@ -516,9 +525,11 @@ streaming::Actor left_semi_join_shuffle( ); auto left_chunk = - co_await left_msg.release().make_available(ctx); + co_await left_msg.release() + .make_available(ctx); auto right_chunk = - co_await right_msg.release().make_available(ctx); + co_await right_msg.release() + .make_available(ctx); left_event.record(left_chunk.stream()); @@ -567,8 +578,11 @@ streaming::Actor shuffle( comm->logger()->debug("Shuffle: no more input"); break; } - auto chunk = co_await msg.release().make_available(ctx); - auto packed = partition_and_pack( + auto chunk = + co_await msg.release().make_available( + ctx + ); + auto packed = cudf_streaming::integrations::partition_and_pack( chunk.table_view(), keys, static_cast(num_partitions), @@ -584,11 +598,11 @@ streaming::Actor shuffle( auto packed_data = shuffler.extract(pid); auto stream = ctx->br()->stream_pool().get_stream(); co_await ch_out->send( - streaming::to_message( + cudf_streaming::streaming::to_message( pid, - std::make_unique( - unpack_and_concat( - unspill_partitions( + std::make_unique( + cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::unspill_partitions( std::move(packed_data), ctx->br().get(), AllowOverbooking::YES ), stream, diff --git a/cpp/benchmarks/streaming/ndsh/parquet_writer.cpp b/cpp/benchmarks/streaming/ndsh/parquet_writer.cpp index 7c566e998..2ffe49909 100644 --- a/cpp/benchmarks/streaming/ndsh/parquet_writer.cpp +++ b/cpp/benchmarks/streaming/ndsh/parquet_writer.cpp @@ -11,11 +11,11 @@ #include #include +#include #include #include #include -#include namespace rapidsmpf::ndsh { @@ -30,7 +30,8 @@ rapidsmpf::streaming::Actor write_parquet( 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 chunk = + co_await msg.release().make_available(ctx); auto table = chunk.table_view(); auto metadata = cudf::io::table_input_metadata(table); CudaEvent event; @@ -51,7 +52,10 @@ rapidsmpf::streaming::Actor write_parquet( if (msg.empty()) { break; } - chunk = co_await msg.release().make_available(ctx); + chunk = + co_await msg.release().make_available( + ctx + ); table = chunk.table_view(); RAPIDSMPF_EXPECTS( static_cast(table.num_columns()) == column_names.size(), diff --git a/cpp/benchmarks/streaming/ndsh/q01.cpp b/cpp/benchmarks/streaming/ndsh/q01.cpp index d4dd6975f..a84fdb5c9 100644 --- a/cpp/benchmarks/streaming/ndsh/q01.cpp +++ b/cpp/benchmarks/streaming/ndsh/q01.cpp @@ -21,6 +21,8 @@ #include #include #include +#include +#include #include #include @@ -30,8 +32,6 @@ #include #include #include -#include -#include #include #include "concatenate.hpp" @@ -79,7 +79,7 @@ rapidsmpf::streaming::Actor read_lineitem( : rapidsmpf::ndsh::make_date_filter( stream, date, "l_shipdate", cudf::ast::ast_operator::LESS_EQUAL ); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, @@ -133,7 +133,7 @@ rapidsmpf::streaming::Actor postprocess_group_by( (co_await ch_in->receive()).empty(), "Expecting concatenated input at this point" ); auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available(ctx); auto stream = chunk.stream(); auto columns = cudf::table{chunk.table_view(), stream, ctx->br()->device_mr()}.release(); @@ -166,9 +166,9 @@ rapidsmpf::streaming::Actor postprocess_group_by( ); columns.push_back(std::move(count)); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( std::make_unique(std::move(columns)), stream ) ) @@ -196,7 +196,9 @@ rapidsmpf::streaming::Actor select_columns_for_groupby( break; } auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); auto chunk_stream = chunk.stream(); auto sequence_number = msg.sequence_number(); auto table = chunk.table_view(); @@ -256,9 +258,9 @@ static __device__ void calculate_charge(double *charge, double discprice, double std::make_unique(discount, chunk_stream, ctx->br()->device_mr()) ); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( sequence_number, - std::make_unique( + std::make_unique( std::make_unique(std::move(result)), chunk_stream ) ) diff --git a/cpp/benchmarks/streaming/ndsh/q03.cpp b/cpp/benchmarks/streaming/ndsh/q03.cpp index 89a922a80..2a2e46f3f 100644 --- a/cpp/benchmarks/streaming/ndsh/q03.cpp +++ b/cpp/benchmarks/streaming/ndsh/q03.cpp @@ -28,23 +28,23 @@ #include #include #include +#include +#include +#include +#include +#include #include #include #include #include #include -#include -#include #include #include #include #include #include #include -#include -#include -#include #include #include "concatenate.hpp" @@ -69,7 +69,7 @@ rapidsmpf::streaming::Actor read_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 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)); @@ -90,7 +90,7 @@ rapidsmpf::streaming::Actor read_customer( *std::any_cast>(owner->at(1)) ) ); - return std::make_unique( + return std::make_unique( stream, *std::any_cast>(owner->back()), rapidsmpf::OwningWrapper(static_cast(owner), [](void* p) { @@ -98,7 +98,7 @@ rapidsmpf::streaming::Actor read_customer( }) ); }(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, @@ -142,7 +142,7 @@ rapidsmpf::streaming::Actor read_lineitem( : rapidsmpf::ndsh::make_date_filter( stream, date, "l_shipdate", cudf::ast::ast_operator::GREATER ); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, @@ -187,7 +187,7 @@ rapidsmpf::streaming::Actor read_orders( : rapidsmpf::ndsh::make_date_filter( stream, date, "o_orderdate", cudf::ast::ast_operator::LESS ); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, @@ -224,7 +224,9 @@ rapidsmpf::streaming::Actor select_columns_for_groupby( break; } auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); auto chunk_stream = chunk.stream(); auto sequence_number = msg.sequence_number(); auto table = chunk.table_view(); @@ -274,9 +276,9 @@ static __device__ void calculate_revenue(double *revenue, double extprice, doubl ) ); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( sequence_number, - std::make_unique( + std::make_unique( std::make_unique(std::move(result)), chunk_stream ) ) @@ -305,7 +307,9 @@ rapidsmpf::streaming::Actor top_k_by( break; } auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); auto const indices = cudf::sorted_order( chunk.table_view().select(keys), order, @@ -343,9 +347,9 @@ rapidsmpf::streaming::Actor top_k_by( ctx->br()->device_mr() ); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( 0, - std::make_unique( + std::make_unique( std::move(result), out_stream ) ) @@ -370,14 +374,16 @@ rapidsmpf::streaming::Actor fanout_bounded( break; } auto chunk = - co_await msg.release().make_available(ctx); + 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 = rapidsmpf::streaming::to_message( + auto msg1 = cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( std::make_unique( chunk.table_view().select(ch1_cols), chunk.stream(), @@ -391,9 +397,11 @@ rapidsmpf::streaming::Actor fanout_bounded( if (!ch2_out->is_shutdown()) { // TODO: We know here that ch2 wants the whole table. tasks.push_back(ch2_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique(std::move(chunk)) + std::make_unique( + std::move(chunk) + ) ) )); } @@ -473,7 +481,9 @@ int main(int argc, char** argv) { RAPIDSMPF_CUDA_TRY(cudaGetDevice(&device)); RAPIDSMPF_CUDA_TRY(cudaDeviceGetAttribute(&l2size, cudaDevAttrL2CacheSize, device)); auto const num_filter_blocks = - rapidsmpf::BloomFilter::fitting_num_blocks(static_cast(l2size)); + cudf_streaming::integrations::BloomFilter::fitting_num_blocks( + static_cast(l2size) + ); for (int i = 0; i < arguments.num_iterations; i++) { int op_id{0}; @@ -533,7 +543,7 @@ int main(int argc, char** argv) { {0}, customer_x_orders_input )); - auto bloom_filter = rapidsmpf::streaming::BloomFilter( + auto bloom_filter = cudf_streaming::streaming::BloomFilter( ctx, comm, cudf::DEFAULT_HASH_SEED, num_filter_blocks ); actors.push_back(bloom_filter.build( diff --git a/cpp/benchmarks/streaming/ndsh/q04.cpp b/cpp/benchmarks/streaming/ndsh/q04.cpp index a76c3fd80..2c9825ab4 100644 --- a/cpp/benchmarks/streaming/ndsh/q04.cpp +++ b/cpp/benchmarks/streaming/ndsh/q04.cpp @@ -28,19 +28,19 @@ #include #include #include +#include +#include +#include +#include #include #include #include -#include #include #include #include #include #include -#include -#include -#include #include #include "concatenate.hpp" @@ -95,7 +95,7 @@ rapidsmpf::streaming::Actor read_lineitem( }) .build(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -139,7 +139,7 @@ rapidsmpf::streaming::Actor read_orders( stream, start_date, end_date, "o_orderdate" ); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk, std::move(filter) ); } @@ -159,7 +159,9 @@ rapidsmpf::streaming::Actor filter_lineitem( break; } auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); auto chunk_stream = chunk.stream(); auto table = chunk.table_view(); @@ -176,9 +178,9 @@ rapidsmpf::streaming::Actor filter_lineitem( auto filtered_table = cudf::apply_boolean_mask(table.select({2}), mask->view(), chunk_stream, mr); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( std::move(filtered_table), chunk_stream ) ) @@ -205,16 +207,16 @@ rapidsmpf::streaming::Actor fanout_bounded( break; } auto chunk = - co_await msg.release().make_available( + 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 = rapidsmpf::streaming::to_message( + auto msg1 = cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( std::make_unique( chunk.table_view().select(ch1_cols), chunk.stream(), @@ -228,9 +230,11 @@ rapidsmpf::streaming::Actor fanout_bounded( if (!ch2_out->is_shutdown()) { // TODO: We know here that ch2 wants the whole table. tasks.push_back(ch2_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique(std::move(chunk)) + std::make_unique( + std::move(chunk) + ) ) )); } @@ -310,7 +314,9 @@ int main(int argc, char** argv) { RAPIDSMPF_CUDA_TRY(cudaGetDevice(&device)); RAPIDSMPF_CUDA_TRY(cudaDeviceGetAttribute(&l2size, cudaDevAttrL2CacheSize, device)); auto const num_filter_blocks = - rapidsmpf::BloomFilter::fitting_num_blocks(static_cast(l2size)); + cudf_streaming::integrations::BloomFilter::fitting_num_blocks( + static_cast(l2size) + ); for (int i = 0; i < arguments.num_iterations; i++) { rapidsmpf::OpID op_id{0}; @@ -366,7 +372,7 @@ int main(int argc, char** argv) { // Build bloom filter from filtered orders' o_orderkey auto bloom_filter_output = ctx->create_channel(); - auto bloom_filter = rapidsmpf::streaming::BloomFilter( + auto bloom_filter = cudf_streaming::streaming::BloomFilter( ctx, comm, cudf::DEFAULT_HASH_SEED, num_filter_blocks ); actors.push_back(bloom_filter.build( diff --git a/cpp/benchmarks/streaming/ndsh/q09.cpp b/cpp/benchmarks/streaming/ndsh/q09.cpp index a52ae40f2..dbe038209 100644 --- a/cpp/benchmarks/streaming/ndsh/q09.cpp +++ b/cpp/benchmarks/streaming/ndsh/q09.cpp @@ -25,6 +25,8 @@ #include #include #include +#include +#include #include #include @@ -34,8 +36,6 @@ #include #include #include -#include -#include #include "concatenate.hpp" #include "groupby.hpp" @@ -69,7 +69,7 @@ rapidsmpf::streaming::Actor read_lineitem( "l_suppkey"} ) .build(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -88,7 +88,7 @@ rapidsmpf::streaming::Actor read_nation( auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) .column_names({"n_name", "n_nationkey"}) .build(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -107,7 +107,7 @@ rapidsmpf::streaming::Actor read_orders( auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) .column_names({"o_orderdate", "o_orderkey"}) .build(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -126,7 +126,7 @@ rapidsmpf::streaming::Actor read_part( auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) .column_names({"p_partkey", "p_name"}) .build(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -145,7 +145,7 @@ rapidsmpf::streaming::Actor read_partsupp( auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) .column_names({"ps_partkey", "ps_suppkey", "ps_supplycost"}) .build(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -164,7 +164,7 @@ rapidsmpf::streaming::Actor read_supplier( auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) .column_names({"s_nationkey", "s_suppkey"}) .build(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -183,7 +183,9 @@ rapidsmpf::streaming::Actor filter_part( } co_await ctx->executor()->schedule(); auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); auto chunk_stream = chunk.stream(); auto table = chunk.table_view(); auto p_name = table.column(1); @@ -192,9 +194,9 @@ rapidsmpf::streaming::Actor filter_part( p_name, *static_cast(target.get()), chunk_stream, mr ); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( cudf::apply_boolean_mask( table.select({0}), mask->view(), chunk_stream, mr ), @@ -224,7 +226,9 @@ rapidsmpf::streaming::Actor select_columns( } co_await ctx->executor()->schedule(); auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); auto chunk_stream = chunk.stream(); auto sequence_number = msg.sequence_number(); auto table = chunk.table_view(); @@ -271,9 +275,9 @@ static __device__ void calculate_amount(double *amount, double discount, double ) ); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( sequence_number, - std::make_unique( + std::make_unique( std::make_unique(std::move(result)), chunk_stream ) ) @@ -302,7 +306,7 @@ rapidsmpf::streaming::Actor round_sum_profit( 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); + co_await msg.release().make_available(ctx); auto table = chunk.table_view(); #pragma GCC diagnostic push #pragma GCC diagnostic ignored "-Wdeprecated-declarations" @@ -315,9 +319,9 @@ rapidsmpf::streaming::Actor round_sum_profit( ctx->br()->device_mr() ); #pragma GCC diagnostic pop - auto result = rapidsmpf::streaming::to_message( + auto result = cudf_streaming::streaming::to_message( 0, - std::make_unique( + std::make_unique( std::make_unique( cudf::table_view({table.column(0), table.column(1), rounded->view()}), chunk.stream(), diff --git a/cpp/benchmarks/streaming/ndsh/q21.cpp b/cpp/benchmarks/streaming/ndsh/q21.cpp index 3b7e6a052..fb69d98ab 100644 --- a/cpp/benchmarks/streaming/ndsh/q21.cpp +++ b/cpp/benchmarks/streaming/ndsh/q21.cpp @@ -23,12 +23,15 @@ #include #include #include +#include +#include +#include +#include #include #include #include -#include #include #include #include @@ -38,9 +41,6 @@ #include #include #include -#include -#include -#include #include "concatenate.hpp" #include "groupby.hpp" @@ -72,7 +72,7 @@ rapidsmpf::streaming::Actor read_lineitem( if (latch != nullptr) { co_await *latch; } - co_return co_await rapidsmpf::streaming::actor::read_parquet( + co_return co_await cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -92,7 +92,7 @@ rapidsmpf::streaming::Actor read_nation( .column_names({"n_nationkey"}) .build(); // filter: "n_name" == "SAUDI ARABIA" - auto filter_expr = [&]() -> std::unique_ptr { + auto filter_expr = [&]() -> std::unique_ptr { auto stream = ctx->br()->stream_pool().get_stream(); auto owner = new std::vector; constexpr auto name = "SAUDI ARABIA"; @@ -116,7 +116,7 @@ rapidsmpf::streaming::Actor read_nation( *std::any_cast>(owner->at(1)) ) ); - return std::make_unique( + return std::make_unique( stream, *std::any_cast>(owner->back()), rapidsmpf::OwningWrapper(static_cast(owner), [](void* p) { @@ -124,7 +124,7 @@ rapidsmpf::streaming::Actor read_nation( }) ); }(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, @@ -150,7 +150,7 @@ rapidsmpf::streaming::Actor read_orders( .column_names({"o_orderkey"}) .build(); // filter: "o_orderstatus" == "F" - auto filter_expr = [&]() -> std::unique_ptr { + auto filter_expr = [&]() -> std::unique_ptr { auto stream = ctx->br()->stream_pool().get_stream(); auto owner = new std::vector; constexpr auto status = "F"; @@ -176,7 +176,7 @@ rapidsmpf::streaming::Actor read_orders( *std::any_cast>(owner->at(1)) ) ); - return std::make_unique( + return std::make_unique( stream, *std::any_cast>(owner->back()), rapidsmpf::OwningWrapper(static_cast(owner), [](void* p) { @@ -184,7 +184,7 @@ rapidsmpf::streaming::Actor read_orders( }) ); }(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, @@ -201,7 +201,7 @@ rapidsmpf::streaming::Actor read_orders_with_bloom_filter( std::shared_ptr bloom_filter_in, std::shared_ptr ch_out, std::vector filter_keys, - rapidsmpf::streaming::BloomFilter bloom_filter, + cudf_streaming::streaming::BloomFilter bloom_filter, std::size_t num_producers, cudf::size_type num_rows_per_chunk, std::string const input_directory @@ -253,7 +253,7 @@ rapidsmpf::streaming::Actor read_supplier( auto options = cudf::io::parquet_reader_options::builder(cudf::io::source_info(files)) .column_names({"s_suppkey", "s_nationkey", "s_name"}) .build(); - return rapidsmpf::streaming::actor::read_parquet( + return cudf_streaming::streaming::actor::read_parquet( ctx, comm, ch_out, num_producers, options, num_rows_per_chunk ); } @@ -271,7 +271,9 @@ rapidsmpf::streaming::Actor filter_lineitem( break; } auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); auto mask = cudf::binary_operation( chunk.table_view().column(2), @@ -282,9 +284,9 @@ rapidsmpf::streaming::Actor filter_lineitem( ctx->br()->device_mr() ); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( cudf::apply_boolean_mask( chunk.table_view().select({0, 1}), mask->view(), @@ -313,7 +315,9 @@ rapidsmpf::streaming::Actor filter_grouped_greater( break; } auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); auto mask = cudf::binary_operation( chunk.table_view().column(1), @@ -327,9 +331,9 @@ rapidsmpf::streaming::Actor filter_grouped_greater( ); latch->count_down(); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( cudf::apply_boolean_mask( chunk.table_view().select({0}), mask->view(), @@ -357,7 +361,9 @@ rapidsmpf::streaming::Actor filter_grouped_equal( break; } auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); auto mask = cudf::binary_operation( chunk.table_view().column(1), @@ -370,9 +376,9 @@ rapidsmpf::streaming::Actor filter_grouped_equal( ctx->br()->device_mr() ); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( cudf::apply_boolean_mask( chunk.table_view().select({0}), mask->view(), @@ -404,14 +410,16 @@ rapidsmpf::streaming::Actor fanout_bounded( break; } auto chunk = - co_await msg.release().make_available(ctx); + 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 = rapidsmpf::streaming::to_message( + auto msg1 = cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( std::make_unique( chunk.table_view().select(ch1_cols), chunk.stream(), @@ -425,9 +433,11 @@ rapidsmpf::streaming::Actor fanout_bounded( if (!ch2_out->is_shutdown()) { // TODO: We know here that ch2 wants the whole table. tasks.push_back(ch2_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique(std::move(chunk)) + std::make_unique( + std::move(chunk) + ) ) )); } @@ -466,13 +476,15 @@ rapidsmpf::streaming::Actor slice( break; } auto chunk = - co_await msg.release().make_available(ctx); + co_await msg.release().make_available( + ctx + ); if (global_start == global_end) { co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( cudf::empty_like(chunk.table_view()), chunk.stream() ) ) @@ -494,9 +506,9 @@ rapidsmpf::streaming::Actor slice( if (local_start == 0 && local_end == num_rows) { co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( std::move(chunk) ) ) @@ -508,9 +520,9 @@ rapidsmpf::streaming::Actor slice( ctx->br()->device_mr() ); co_await ch_out->send( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( std::move(sliced_table), chunk.stream() ) ) @@ -556,7 +568,7 @@ rapidsmpf::streaming::Actor populate_bloom_filter( std::shared_ptr ch_out, std::vector keys, rapidsmpf::OpID tag, - rapidsmpf::streaming::BloomFilter bloom_filter + cudf_streaming::streaming::BloomFilter bloom_filter ) { rapidsmpf::streaming::ShutdownAtExit c{ch_in, ch_out}; auto passthrough = ctx->create_channel(); @@ -568,19 +580,17 @@ rapidsmpf::streaming::Actor populate_bloom_filter( if (msg.empty()) { break; } - auto chunk = - co_await msg.release().make_available( - ctx - ); + 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( - rapidsmpf::streaming::to_message( + cudf_streaming::streaming::to_message( msg.sequence_number(), - std::make_unique( + std::make_unique( std::move(out), stream ) ) @@ -657,7 +667,9 @@ int main(int argc, char** argv) { RAPIDSMPF_CUDA_TRY(cudaGetDevice(&device)); RAPIDSMPF_CUDA_TRY(cudaDeviceGetAttribute(&l2size, cudaDevAttrL2CacheSize, device)); auto const num_filter_blocks = - rapidsmpf::BloomFilter::fitting_num_blocks(static_cast(l2size)); + cudf_streaming::integrations::BloomFilter::fitting_num_blocks( + static_cast(l2size) + ); for (int i = 0; i < arguments.num_iterations; i++) { int op_id{0}; std::vector actors; @@ -863,7 +875,7 @@ int main(int argc, char** argv) { ) ); auto bloom_output = ctx->create_channel(); - auto bloom_filter = rapidsmpf::streaming::BloomFilter( + 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. diff --git a/cpp/benchmarks/streaming/ndsh/sort.cpp b/cpp/benchmarks/streaming/ndsh/sort.cpp index 756d2bc65..78ebc8f2a 100644 --- a/cpp/benchmarks/streaming/ndsh/sort.cpp +++ b/cpp/benchmarks/streaming/ndsh/sort.cpp @@ -11,10 +11,10 @@ #include #include +#include #include #include -#include namespace rapidsmpf::ndsh { @@ -29,7 +29,7 @@ rapidsmpf::streaming::Actor chunkwise_sort_by( ) { streaming::ShutdownAtExit c{ch_in, ch_out}; co_await ctx->executor()->schedule(); - auto make_table = [&](streaming::TableChunk& chunk) { + auto make_table = [&](cudf_streaming::streaming::TableChunk& chunk) { if (std::ranges::equal(keys, values)) { return cudf::sort( chunk.table_view().select(keys), @@ -54,10 +54,15 @@ rapidsmpf::streaming::Actor chunkwise_sort_by( if (msg.empty()) { break; } - auto chunk = co_await msg.release().make_available(ctx); + 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()) + std::make_unique( + make_table(chunk), chunk.stream() + ) )); } co_await ch_out->drain(ctx->executor()); diff --git a/cpp/benchmarks/streaming/ndsh/utils.cpp b/cpp/benchmarks/streaming/ndsh/utils.cpp index 8716da3be..76237873a 100644 --- a/cpp/benchmarks/streaming/ndsh/utils.cpp +++ b/cpp/benchmarks/streaming/ndsh/utils.cpp @@ -18,6 +18,7 @@ #include #include +#include #include #include #include @@ -35,7 +36,6 @@ #include #include #include -#include namespace rapidsmpf::ndsh { namespace detail { @@ -114,9 +114,9 @@ streaming::Actor consume_channel( if (msg.empty()) { break; } - if (msg.holds()) { - auto chunk = - co_await msg.release().make_available(ctx); + if (msg.holds()) { + auto chunk = co_await msg.release() + .make_available(ctx); ctx->logger()->print( "Consumed chunk with ", chunk.table_view().num_rows(), diff --git a/cpp/benchmarks/streaming/ndsh/utils.hpp b/cpp/benchmarks/streaming/ndsh/utils.hpp index 737477d29..bd63c45b8 100644 --- a/cpp/benchmarks/streaming/ndsh/utils.hpp +++ b/cpp/benchmarks/streaming/ndsh/utils.hpp @@ -20,6 +20,8 @@ #include #include #include +#include +#include #include #include @@ -29,8 +31,6 @@ #include #include #include -#include -#include namespace rapidsmpf::ndsh { namespace detail { @@ -94,7 +94,7 @@ namespace detail { * @return Filter expression with proper lifetime management */ template -std::unique_ptr make_date_filter( +std::unique_ptr make_date_filter( rmm::cuda_stream_view stream, cuda::std::chrono::year_month_day date, std::string const& column_name, @@ -124,7 +124,7 @@ std::unique_ptr make_date_filter( *std::any_cast>(owner->at(1)) ) ); - return std::make_unique( + return std::make_unique( stream, *std::any_cast>(owner->back()), OwningWrapper(static_cast(owner), [](void* p) { @@ -149,7 +149,7 @@ std::unique_ptr make_date_filter( * @return Filter expression with proper lifetime management */ template -std::unique_ptr make_date_range_filter( +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, @@ -219,7 +219,7 @@ std::unique_ptr make_date_range_filter( ) ); - return std::make_unique( + return std::make_unique( stream, *std::any_cast>(owner->back()), OwningWrapper(static_cast(owner), [](void* p) { diff --git a/cpp/compute-sanitizer-suppressions.xml b/cpp/compute-sanitizer-suppressions.xml index 309d4d02f..5e5cd718f 100644 --- a/cpp/compute-sanitizer-suppressions.xml +++ b/cpp/compute-sanitizer-suppressions.xml @@ -1,73 +1,5 @@ - - InitcheckApiError - Error - - Host API uninitialized memory access - - - error - - .*/libcuda.so.* - - - .*/libcudart.so.* - - - .*/libcudart.so.* - - - .*/libcudart.so.* - - - cudf::detail::cuda_memcpy_async_impl - .*/libcudf.so - - - .*/libcudf.so - - - cudf::detail::contiguous_split_state::contiguous_split_state - .*/libcudf.so - - - - - InitcheckApiError - Error - - Host API uninitialized memory access - - - error - - .*/libcuda.so.* - - - .*/libcudart.so.* - - - .*/libcudart.so.* - - - cudaMemcpyAsync - .*/libcudart.so.* - - - cudf::detail::cuda_memcpy_async_impl - .*/libcudf.so - - - cudf::detail::contiguous_split - .*/libcudf.so - - - cudf::contiguous_split - .*/libcudf.so - - - InitcheckApiError @@ -119,45 +51,6 @@ - - InitcheckApiError - Error - - Host API uninitialized memory access - - - error - - .*/libcuda.so.* - - - .*/libcudart.so.* - - - .*/libcudart.so.* - - - .*/libcudart.so.* - - - cudf::detail::cuda_memcpy_async_impl - .*/libcudf.so - - - cudf::detail::contiguous_split - .*/libcudf.so - - - cudf::detail::pack - .*/libcudf.so - - - cudf::pack - .*/libcudf.so - - - - InitcheckApiError Error @@ -265,7 +158,7 @@ .*/librapidsmpf.so - rapidsmpf::streaming::TableChunk::spill_to_host + cudf_streaming::streaming::TableChunk::spill_to_host .*/streaming/cudf/table_chunk.cpp .*/librapidsmpf.so diff --git a/cpp/doxygen/Doxyfile b/cpp/doxygen/Doxyfile index 69dbf29cb..2852dd1ec 100644 --- a/cpp/doxygen/Doxyfile +++ b/cpp/doxygen/Doxyfile @@ -2219,8 +2219,7 @@ SKIP_FUNCTION_MACROS = YES # the path). If a tag file is not located in the directory in which doxygen is # run, you must also specify the path to the tagfile here. -TAGFILES = rmm.tag=https://docs.rapids.ai/api/librmm/$(RAPIDS_VERSION_MAJOR_MINOR) \ - cudf.tag=https://docs.rapids.ai/api/libcudf/$(RAPIDS_VERSION_MAJOR_MINOR) +TAGFILES = rmm.tag=https://docs.rapids.ai/api/librmm/$(RAPIDS_VERSION_MAJOR_MINOR) # When a file name is specified after GENERATE_TAGFILE, doxygen will create a # tag file that is based on the input files it reads. See section "Linking to diff --git a/cpp/examples/CMakeLists.txt b/cpp/examples/CMakeLists.txt index b73560fc9..8f107fe32 100644 --- a/cpp/examples/CMakeLists.txt +++ b/cpp/examples/CMakeLists.txt @@ -5,31 +5,33 @@ # cmake-format: on # ================================================================================= -add_executable(example_shuffle "example_shuffle.cpp" "../benchmarks/utils/random_data.cu") -set_target_properties( - example_shuffle - PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/examples" - CXX_STANDARD 20 - CXX_STANDARD_REQUIRED ON - # For std:: support of __int128_t. Can be removed once using cuda::std - CXX_EXTENSIONS ON - CUDA_STANDARD 20 - CUDA_STANDARD_REQUIRED ON -) -target_compile_options( - example_shuffle PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" - "$<$:${RAPIDSMPF_CUDA_FLAGS}>" -) -target_link_libraries( - example_shuffle PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx $ - $ maybe_asan -) -install( - TARGETS example_shuffle - COMPONENT testing - DESTINATION bin/examples/librapidsmpf - EXCLUDE_FROM_ALL -) +if(BUILD_CUDF_TESTS) + add_executable(example_shuffle "example_shuffle.cpp" "../benchmarks/utils/random_data.cu") + set_target_properties( + example_shuffle + PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/examples" + CXX_STANDARD 20 + CXX_STANDARD_REQUIRED ON + CXX_EXTENSIONS ON + CUDA_STANDARD 20 + CUDA_STANDARD_REQUIRED ON + ) + target_compile_options( + example_shuffle PRIVATE "$<$:${RAPIDSMPF_CXX_FLAGS}>" + "$<$:${RAPIDSMPF_CUDA_FLAGS}>" + ) + target_link_libraries( + example_shuffle + PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx cudf_streaming::cudf_streaming + $ $ maybe_asan + ) + install( + TARGETS example_shuffle + COMPONENT testing + DESTINATION bin/examples/librapidsmpf + EXCLUDE_FROM_ALL + ) +endif() if(RAPIDSMPF_HAVE_CUPTI) add_executable(example_cupti_monitor "example_cupti_monitor.cpp") diff --git a/cpp/examples/example_shuffle.cpp b/cpp/examples/example_shuffle.cpp index 144c3bd59..1e322b476 100644 --- a/cpp/examples/example_shuffle.cpp +++ b/cpp/examples/example_shuffle.cpp @@ -8,9 +8,10 @@ #include #include +#include + #include #include -#include #include #include #include @@ -76,7 +77,7 @@ int main(int argc, char** argv) { // each partition. The result is a mapping of `PartID`, globally unique partition // identifiers, to their packed partitions. std::unordered_map packed_inputs = - rapidsmpf::partition_and_pack( + cudf_streaming::integrations::partition_and_pack( local_input, {0}, // columns_to_hash static_cast(total_num_partitions), @@ -110,8 +111,8 @@ int main(int argc, char** argv) { // Unpack (deserialize) and concatenate the chunks into a single table using a // convenience function. local_outputs.push_back( - rapidsmpf::unpack_and_concat( - rapidsmpf::unspill_partitions( + cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::unspill_partitions( std::move(packed_chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, diff --git a/cpp/include/rapidsmpf/integrations/cudf/bloom_filter.hpp b/cpp/include/rapidsmpf/integrations/cudf/bloom_filter.hpp deleted file mode 100644 index fe267de00..000000000 --- a/cpp/include/rapidsmpf/integrations/cudf/bloom_filter.hpp +++ /dev/null @@ -1,142 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include - -#include -#include -#include -#include - -namespace rapidsmpf { - -/** - * @brief A bloom filter, used for approximate set membership queries. - */ -struct BloomFilter { - /** - * @brief Create a filter. - * - * @param num_blocks Number of blocks in the filter. - * @param seed Seed used for hashing each value. - * @param storage Storage to view as a bloom filter, must be appropriately - * initialized. - * @param stream CUDA stream for device operations. - */ - BloomFilter( - std::size_t num_blocks, - std::uint64_t seed, - void* storage, - rmm::cuda_stream_view stream - ); - - /** - * @brief Create a read-only filter. - * - * @param num_blocks Number of blocks in the filter. - * @param seed Seed used for hashing each value. - * @param storage View of storage, must be appropriately initialized. - * @param stream CUDA stream for device operations. - * - * @return A const-qualified bloom filter viewing the underlying storage. - */ - static const BloomFilter view( - std::size_t num_blocks, - std::uint64_t seed, - void const* storage, - rmm::cuda_stream_view stream - ); - - /** - * @brief Create uninitialized storage for a filter. - * - * @param num_blocks Number of blocks in the filter. - * @param stream CUDA stream for device operations. - * @param mr Memory resource for allocations. - * - * @return Unique pointer to a device buffer containing storage for the requested - * number of filter blocks. - */ - static std::unique_ptr storage( - std::size_t num_blocks, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr - ); - - /** - * @brief Add values to the filter. - * - * @param values_to_hash table of values to hash (with cudf::hashing::xxhash_64()) - * @param stream CUDA stream for allocations and device operations. - * @param mr Memory resource for allocations. - */ - void add( - cudf::table_view const& values_to_hash, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr - ); - - /** - * @brief Merge two filters, computing their union. - * - * @param other Other filter to merge into this one. - * @param stream CUDA stream for device operations. - * - * @throws std::logic_error If `other` is not compatible with this filter. - */ - void merge(BloomFilter const& other, rmm::cuda_stream_view stream); - - /** - * @brief Return a mask of which rows are contained in the filter. - * - * @param values Value to check for set membership - * @param stream CUDA stream for allocations and device operations. - * @param mr Memory resource for allocations. - * - * @return Mask vector to be used for filtering the table. - */ - [[nodiscard]] rmm::device_uvector contains( - cudf::table_view const& values, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr - ) const; - - /** - * @brief @return The stream the underlying storage is valid on. - */ - [[nodiscard]] rmm::cuda_stream_view stream() const noexcept; - - /** - * @brief @return Pointer to the underlying storage. - */ - [[nodiscard]] void* data() noexcept; - - /** - * @brief @return Cnst Pointer to the underlying storage. - */ - [[nodiscard]] void const* data() const noexcept; - - /** - * @brief @return Size in bytes of the underlying storage. - */ - [[nodiscard]] std::size_t size() const noexcept; - - /** - * @brief @return Number of blocks to use if the filter should fit in a given L2 cache - * size. - * - * @param l2size Size of the L2 cache in bytes. - */ - [[nodiscard]] static std::size_t fitting_num_blocks(std::size_t l2size) noexcept; - - private: - std::size_t num_blocks_; ///< Number of blocks used in the filter. - std::uint64_t seed_; ///< Seed used when hashing values. - void* storage_; ///< Backing storage. - rmm::cuda_stream_view stream_; ///< Stream storage is valid on. -}; - -} // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/integrations/cudf/partition.hpp b/cpp/include/rapidsmpf/integrations/cudf/partition.hpp deleted file mode 100644 index c55d58968..000000000 --- a/cpp/include/rapidsmpf/integrations/cudf/partition.hpp +++ /dev/null @@ -1,198 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include - -namespace rapidsmpf { - - -/** - * @brief Partitions rows from the input table into multiple output tables. - * - * @param table The table to partition. - * @param columns_to_hash Indices of input columns to hash. - * @param num_partitions The number of partitions. - * @param hash_function Hash function to use. - * @param seed Seed value to the hash function. - * @param stream CUDA stream used for device memory operations and kernel launches. - * @param br Buffer resource for memory allocations. - * @param allow_overbooking If true, allow overbooking (true by default) - * - * @return A vector of each partition and a table that owns the device memory. - * - * @throws std::out_of_range if index is `columns_to_hash` is invalid - * - * @see cudf::hash_partition - * @see cudf::split - */ -[[nodiscard]] std::pair, std::unique_ptr> -partition_and_split( - cudf::table_view const& table, - std::vector const& columns_to_hash, - int num_partitions, - cudf::hash_id hash_function, - std::uint32_t seed, - rmm::cuda_stream_view stream, - BufferResource* br, - AllowOverbooking allow_overbooking = AllowOverbooking::YES -); - - -/** - * @brief Partitions rows from the input table into multiple packed (serialized) tables. - * - * @param table The table to partition. - * @param columns_to_hash Indices of input columns to hash. - * @param num_partitions The number of partitions to use. - * @param hash_function Hash function to use. - * @param seed Seed value to the hash function. - * @param stream CUDA stream used for device memory operations and kernel launches. - * @param br Buffer resource for memory allocations. - * @param allow_overbooking If true, allow overbooking (true by default) - * // TODO: disable this by default https://github.com/rapidsmpf/rapidsmpf/issues/449 - * - * @return A map of partition IDs and their packed tables. - * - * @throws std::out_of_range if index is `columns_to_hash` is invalid - * - * @see unpack_and_concat - * @see cudf::hash_partition - * @see cudf::pack - */ -[[nodiscard]] std::unordered_map partition_and_pack( - cudf::table_view const& table, - std::vector const& columns_to_hash, - int num_partitions, - cudf::hash_id hash_function, - std::uint32_t seed, - rmm::cuda_stream_view stream, - BufferResource* br, - AllowOverbooking allow_overbooking = AllowOverbooking::YES -); - - -/** - * @brief Splits rows from the input table into multiple packed (serialized) tables. - * - * @param table The table to split and pack into partitions. - * @param splits The split points, equivalent to cudf::split(), i.e. one less than - * the number of result partitions. - * @param stream CUDA stream used for device memory operations and kernel launches. - * @param br Buffer resource for memory allocations. - * @param allow_overbooking If true, allow overbooking (true by default) - * // TODO: disable this by default https://github.com/rapidsmpf/rapidsmpf/issues/449 - * - * @return A map of partition IDs and their packed tables. - * - * @throws std::out_of_range if the splits are invalid. - * - * @see unpack_and_concat - * @see cudf::split - * @see partition_and_pack - */ -[[nodiscard]] std::unordered_map split_and_pack( - cudf::table_view const& table, - std::vector const& splits, - rmm::cuda_stream_view stream, - BufferResource* br, - AllowOverbooking allow_overbooking = AllowOverbooking::YES -); - - -/** - * @brief Unpack (deserialize) input partitions and concatenate them into a single table. - * - * Empty partitions are ignored. - * - * The unpacking of each partition is stream-ordered on that partition's own CUDA stream. - * The returned table is stream-ordered on the provided @p stream and synchronized with - * the unpacking. - * - * @param partitions Packed input tables (partitions). - * @param stream CUDA stream on which concatenation occurs and on which the resulting - * table is ordered. - * @param br Buffer resource used for memory allocations. - * @param allow_overbooking If true, allow overbooking (true by default). - * @return The concatenated table resulting from unpacking the input partitions. - * - * @throws rapidsmpf::reservation_error If the buffer resource cannot reserve enough - * memory to concatenate all partitions. - * @throws std::logic_error If the partitions are not in device memory. - * - * @see partition_and_pack - * @see cudf::unpack - * @see cudf::concatenate - */ -[[nodiscard]] std::unique_ptr unpack_and_concat( - std::vector&& partitions, - rmm::cuda_stream_view stream, - BufferResource* br, - AllowOverbooking allow_overbooking = AllowOverbooking::YES -); - -/** - * @brief Spill partitions from device memory to host memory. - * - * Moves the buffer of each `PackedData` from device memory to host memory using - * the provided buffer resource and the buffer's CUDA stream. Partitions that are - * already in host memory are passed through unchanged. - * - * For device-resident partitions, a host memory reservation is made before moving - * the buffer. If the reservation fails due to insufficient host memory, an exception - * is thrown. Overbooking is not allowed. - * - * @param partitions The partitions to spill. - * @param br Buffer resource used to reserve host memory and perform the move. - * - * @return A vector of `PackedData`, where each buffer resides in host memory. - * - * @throws rapidsmpf::reservation_error If host memory reservation fails. - */ -std::vector spill_partitions( - std::vector&& partitions, BufferResource* br -); - -/** - * @brief Move spilled partitions (i.e., packed tables in host memory) back to device - * memory. - * - * Each partition is inspected to determine whether its buffer resides in device memory. - * Buffers already in device memory are left untouched. Host-resident buffers are moved - * to device memory using the provided buffer resource and the buffer's CUDA stream. - * - * If insufficient device memory is available, the buffer resource's spill manager is - * invoked to free memory. If overbooking occurs and spilling fails to reclaim enough - * memory, behavior depends on the `allow_overbooking` flag. - * - * @param partitions The partitions to unspill, potentially containing host-resident data. - * @param br Buffer resource responsible for memory reservation and spills. - * @param allow_overbooking If false, ensures enough memory is freed to satisfy the - * reservation; otherwise, allows overbooking even if spilling was insufficient. - * - * @return A vector of `PackedData`, each with a buffer in device memory. - * - * @throws rapidsmpf::reservation_error If overbooking exceeds the amount spilled and - * `allow_overbooking` is false. - */ -std::vector unspill_partitions( - std::vector&& partitions, - BufferResource* br, - AllowOverbooking allow_overbooking -); - -} // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/integrations/cudf/utils.hpp b/cpp/include/rapidsmpf/integrations/cudf/utils.hpp deleted file mode 100644 index c6401fd9c..000000000 --- a/cpp/include/rapidsmpf/integrations/cudf/utils.hpp +++ /dev/null @@ -1,83 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#pragma once - -#include -#include - -#include -#include -#include -#include - -namespace rapidsmpf { - -/** - * @brief Converts the element at a specific index in a `cudf::column_view` to a string. - * - * @param col The column view containing the data. - * @param index The index of the element to convert. - * @param stream CUDA stream used for device memory operations and kernel launches. - * @param mr Memory resource for device memory allocation. - * @return A string representation of the element at the specified index. - */ -std::string str( - cudf::column_view col, - cudf::size_type index, - rmm::cuda_stream_view stream = cudf::get_default_stream(), - rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref() -); - -/** - * @brief Converts all elements in a `cudf::column_view` to a string. - * - * @param col The column view containing the data. - * @param stream CUDA stream used for device memory operations and kernel launches. - * @param mr Memory resource for device memory allocation. - * @return A string representation of all elements in the column. - */ -std::string str( - cudf::column_view col, - rmm::cuda_stream_view stream = cudf::get_default_stream(), - rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref() -); - -/** - * @brief Converts all rows in a `cudf::table_view` to a string. - * - * @param tbl The table view containing the data. - * @param stream CUDA stream used for device memory operations and kernel launches. - * @param mr Memory resource for device memory allocation. - * @return A string representation of all rows in the table. - */ -std::string str( - cudf::table_view tbl, - rmm::cuda_stream_view stream = cudf::get_default_stream(), - rmm::device_async_resource_ref mr = cudf::get_current_device_resource_ref() -); - -/** - * @brief Estimate the memory usage of a column. - * - * @param col The column to estimate the memory usage of. - * @param stream CUDA stream used for device memory operations and kernel launches. - * @return The estimated memory usage of the column. - */ -std::size_t estimated_memory_usage( - cudf::column_view const& col, rmm::cuda_stream_view stream -); - -/** - * @brief Estimate the memory usage of a table. - * - * @param tbl The table to estimate the memory usage of. - * @param stream CUDA stream used for device memory operations and kernel launches. - * @return The estimated memory usage of the table. - */ -std::size_t estimated_memory_usage( - cudf::table_view const& tbl, rmm::cuda_stream_view stream -); - -} // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/memory/cuda_memcpy_async.hpp b/cpp/include/rapidsmpf/memory/cuda_memcpy_async.hpp index 5bd4b5d3d..a9e97538a 100644 --- a/cpp/include/rapidsmpf/memory/cuda_memcpy_async.hpp +++ b/cpp/include/rapidsmpf/memory/cuda_memcpy_async.hpp @@ -5,12 +5,103 @@ #pragma once #include +#include + +#include -#include #include namespace rapidsmpf { +/** + * @brief Asynchronously copies a batch of buffers using the most efficient available API. + * + * On CUDA 13.0+ with a non-default stream, uses `cudaMemcpyBatchAsync` with + * `cudaMemcpySrcAccessOrderStream`, which defers reading the source buffers until + * the stream reaches each copy. This enables true asynchronous copies from pageable + * host memory on modern systems with HMM/ATS support. + * + * Falls back to per-copy `cudaMemcpyAsync` on older CUDA versions or when the default + * stream is used. + * + * @param dsts Host pointer to a list of destination pointers. + * @param srcs Host pointer to a list of source pointers. + * @param sizes Host pointer to a list of sizes (bytes). + * @param count Number of entries in dsts, srcs, sizes. + * @param stream CUDA stream on which copies are enqueued. + * @return cudaError_t CUDA error code. + */ +[[nodiscard]] inline cudaError_t cuda_memcpy_batch_async( + void* const* dsts, + void const* const* srcs, + std::size_t const* sizes, + std::size_t count, + rmm::cuda_stream_view stream +) { +#if CUDART_VERSION >= 13000 + if (!stream.is_default()) { + // Filter out invalid copies; cudaMemcpyBatchAsync does not support + // nullptr dst/src or size==0. + auto is_invalid = [&](std::size_t i) { + return dsts[i] == nullptr || srcs[i] == nullptr || sizes[i] == 0; + }; + + std::vector valid_dsts; + std::vector valid_srcs; + std::vector valid_sizes; + + bool has_invalid = false; + for (std::size_t i = 0; i < count; ++i) { + if (is_invalid(i)) { + has_invalid = true; + break; + } + } + + if (has_invalid) { + valid_dsts.reserve(count); + valid_srcs.reserve(count); + valid_sizes.reserve(count); + for (std::size_t i = 0; i < count; ++i) { + if (dsts[i] != nullptr && srcs[i] != nullptr && sizes[i] != 0) { + valid_dsts.push_back(dsts[i]); + valid_srcs.push_back(srcs[i]); + valid_sizes.push_back(sizes[i]); + } + } + if (valid_dsts.empty()) { + return cudaSuccess; + } + dsts = valid_dsts.data(); + srcs = valid_srcs.data(); + sizes = valid_sizes.data(); + count = valid_dsts.size(); + } + + cudaMemcpyAttributes attrs = { + .srcAccessOrder = cudaMemcpySrcAccessOrderStream, + .flags = cudaMemcpyFlagPreferOverlapWithCompute + }; + std::size_t attrs_idxs = 0; + return cudaMemcpyBatchAsync( + dsts, srcs, sizes, count, &attrs, &attrs_idxs, 1, stream.value() + ); + } +#endif // CUDART_VERSION >= 13000 + for (std::size_t i = 0; i < count; ++i) { + if (dsts[i] == nullptr || srcs[i] == nullptr || sizes[i] == 0) { + continue; + } + cudaError_t status = cudaMemcpyAsync( + dsts[i], srcs[i], sizes[i], cudaMemcpyDefault, stream.value() + ); + if (status != cudaSuccess) { + return status; + } + } + return cudaSuccess; +} + /** * @brief Asynchronously copies memory between host and/or device buffers. * @@ -44,7 +135,11 @@ namespace rapidsmpf { [[nodiscard]] inline cudaError_t cuda_memcpy_async( void* dst, void const* src, std::size_t count, rmm::cuda_stream_view stream ) { - return cudf::detail::memcpy_async(dst, src, count, stream); + if (count == 0) { + return cudaSuccess; + } + void const* src_ptr = src; + return cuda_memcpy_batch_async(&dst, &src_ptr, &count, 1, stream); } } // namespace rapidsmpf diff --git a/cpp/include/rapidsmpf/streaming/coll/shuffler.hpp b/cpp/include/rapidsmpf/streaming/coll/shuffler.hpp index b8ef5d0f2..f962a8a71 100644 --- a/cpp/include/rapidsmpf/streaming/coll/shuffler.hpp +++ b/cpp/include/rapidsmpf/streaming/coll/shuffler.hpp @@ -8,10 +8,10 @@ #include #include +#include #include #include #include -#include namespace rapidsmpf::streaming { diff --git a/cpp/include/rapidsmpf/streaming/cudf/bloom_filter.hpp b/cpp/include/rapidsmpf/streaming/cudf/bloom_filter.hpp deleted file mode 100644 index 97c4258ad..000000000 --- a/cpp/include/rapidsmpf/streaming/cudf/bloom_filter.hpp +++ /dev/null @@ -1,107 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include -#include - -#include - -#include -#include -#include -#include - -namespace rapidsmpf::streaming { - -/** - * @brief Utility managing construction and use of a bloom filter. - * - * This class provides methods to build a bloom filter from a stream of `TableChunk`s and - * then apply that filter to a different stream. - * - * A bloom filter is a fixed size probabilistic data structure that provides approximate - * set membership queries with no false negatives. That is, let `A` be some set and `f(A)` - * be the bloom filter representation of that set. Then, for all `a ∈ A` it holds that `a - * ∈ f(A)`. Conversely, there is a false positive rate that increases with the number of - * distinct values inserted into the bloom filter, and decreases with the number of filter - * blocks. That is, for any given bloom filter, there exists `a ∉ A` such that `a ∈ f(A)`. - * - * See https://arxiv.org/pdf/2512.15595 for details on the GPU implementation used. - * - * We use bloom filters to provide runtime pre-filtering of tables during shuffle-based - * joins. We gather the keys that will match from the build side and use those to - * pre-filter the probe side before shuffling. - */ -struct BloomFilter { - /** - * @brief Construct storage for a bloom filter. - * - * @param ctx Streaming context. - * @param comm Communicator for the collective operation. - * @param seed Hash seed used when hashing values into the filter. - * @param num_filter_blocks Number of blocks in the filter. - */ - explicit BloomFilter( - std::shared_ptr ctx, - std::shared_ptr comm, - std::uint64_t seed, - std::size_t num_filter_blocks - ) noexcept - : ctx_{std::move(ctx)}, - comm_{std::move(comm)}, - seed_{seed}, - num_filter_blocks_{num_filter_blocks} {} - - /** - * @brief Gets the communicator associated with this BloomFilter. - * - * @return Shared pointer to communicator. - */ - [[nodiscard]] std::shared_ptr const& comm() const noexcept { - return comm_; - } - - /** - * @brief Build a bloom filter from the input channel. - * - * @param ch_in Input channel of `TableChunk`s to build bloom filter for. - * @param ch_out Output channel receiving a single message containing the bloom - * filter. - * @param tag Disambiguating tag to combine filters across ranks. - * @return Coroutine representing the construction of the bloom filter. - */ - [[nodiscard]] Actor build( - std::shared_ptr ch_in, std::shared_ptr ch_out, OpID tag - ); - - /** - * @brief Apply a bloom filter to an input channel. - * - * @param bloom_filter Channel containing the bloom filter (a single message). - * @param ch_in Input channel of `TableChunk`s to apply bloom filter to. - * @param ch_out Output channel receiving filtered `TableChunk`s. - * @param keys Indices selecting the key columns for the hash fingerprint - * - * @note The application of the bloom filter expects _exactly one_ message to come - * through the `bloom_filter` channel, which must be drained after that message is - * sent. - * - * @return Coroutine representing the application of the bloom filter. - */ - [[nodiscard]] Actor apply( - std::shared_ptr bloom_filter, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys - ); - - private: - std::shared_ptr ctx_{}; - std::shared_ptr comm_{}; - std::uint64_t seed_{}; - std::size_t num_filter_blocks_{}; -}; -} // namespace rapidsmpf::streaming diff --git a/cpp/include/rapidsmpf/streaming/cudf/channel_metadata.hpp b/cpp/include/rapidsmpf/streaming/cudf/channel_metadata.hpp deleted file mode 100644 index 0e27c3f63..000000000 --- a/cpp/include/rapidsmpf/streaming/cudf/channel_metadata.hpp +++ /dev/null @@ -1,251 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights - * reserved. SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include -#include -#include -#include - -#include - -#include - -#include -#include -#include -#include - -namespace rapidsmpf::streaming { - -/** - * @brief Hash partitioning scheme. - * - * Rows are distributed by `hash(columns[column_indices]) % modulus`. - */ -struct HashScheme { - std::vector column_indices; ///< Column indices to hash on. - int modulus; ///< Hash modulus (number of partitions). - - /** - * @brief Equality comparison. - * @return True if both schemes are equal. - */ - bool operator==(HashScheme const&) const = default; -}; - -/** - * @brief A single sort key: column index, sort direction, and null placement. - */ -struct OrderKey { - cudf::size_type column_index; ///< Column to sort on. - cudf::order order; ///< ASCENDING or DESCENDING. - cudf::null_order null_order; ///< BEFORE or AFTER. - - /** - * @brief Equality comparison. - * @return True if all fields are equal. - */ - bool operator==(OrderKey const&) const = default; - - /** - * @brief Inequality comparison. - * @return True if any field is different. - */ - bool operator!=(OrderKey const&) const = default; -}; - -/** - * @brief Order-based partitioning scheme for sorted/range-partitioned data. - * - * Data is partitioned by value ranges based on predetermined boundaries. - * For N partitions, there are N-1 boundary rows: - * - Partition 0: values < boundaries[0] - * - Partition i (0 < i < N-1): boundaries[i-1] <= values < boundaries[i] - * - Partition N-1: values >= boundaries[N-2] - * - * `keys[i]` is the i-th sort column; ordering is lexicographic by `keys[0]`, - * then `keys[1]`, and so on. - * - * When `boundaries` is set, its columns must align with `keys` - * (same count and compatible dtypes). Mismatched dtypes are a usage error. - * - * `strict_boundaries`: when true, every row in a chunk belongs to a single partition's - * half-open key range (partition keys do not straddle chunk interiors). When false, - * a chunk may contain keys spanning multiple partitions. - */ -struct OrderScheme { - std::vector keys; ///< Sort keys (column, order, null_order per entry). - std::shared_ptr boundaries; ///< N-1 boundary rows for N partitions. - /// See struct-level note on `strict_boundaries` semantics. - bool strict_boundaries{false}; - - /// @brief Default constructor. Produces an invalid (empty) scheme. - OrderScheme() = default; - - /** - * @brief Construct a validated OrderScheme. - * - * @param keys Non-empty sort keys; size must equal `boundaries->shape().second`. - * @param boundaries Non-null, device-resident boundary table (N-1 rows for N - * partitions). Accepts a `unique_ptr` via implicit conversion. - * @param strict_boundaries See struct-level doc. Defaults to false. - * @throws std::invalid_argument if `keys` is empty, `boundaries` is null or not - * device-resident, or `keys.size() != boundaries->shape().second`. - */ - OrderScheme( - std::vector keys, - std::shared_ptr boundaries, - bool strict_boundaries = false - ); - - /** - * @brief Return a new OrderScheme with updated key column indices, sharing - * boundaries. - * - * The new key count must match the existing boundary column count. - * - * @param new_keys Replacement sort keys; size must equal - * `boundaries->shape().second`. - * @return A new OrderScheme with `new_keys` and the same `boundaries` and - * `strict_boundaries`. - * @throws std::invalid_argument if `new_keys` is empty or size mismatches boundaries. - */ - [[nodiscard]] OrderScheme with_keys(std::vector new_keys) const; - - /** - * @brief Check whether boundary values are aligned with another scheme. - * - * @param other The OrderScheme to compare against. - * @param br Buffer resource used for temporary allocations during comparison. - * @return True when both schemes have matching boundary values and strict_boundaries - * attributes, and the schemes are otherwise compatible (same order and null_order). - */ - [[nodiscard]] bool boundaries_aligned_with( - OrderScheme const& other, rapidsmpf::BufferResource& br - ) const; -}; - -/** - * @brief Partitioning specification for a single hierarchical level. - * - * Represents how data is partitioned at one level of the hierarchy - * (e.g., inter-rank or local). Use the static factory methods to construct. - * - * - `none()`: No partitioning information at this level. - * - `inherit()`: Partitioning is inherited from the parent level unchanged. - * - `from_hash(h)`: Explicit hash partitioning with the given scheme. - * - `from_order(o)`: Explicit order/range partitioning with the given scheme. - */ -struct PartitioningSpec { - /** - * @brief Type tag for PartitioningSpec. - */ - enum class Type : std::uint8_t { - NONE, ///< No partitioning information at this level. - INHERIT, ///< Partitioning is inherited from parent level unchanged. - HASH, ///< Hash partitioning. - ORDER, ///< Order/range partitioning. - }; - - Type type = Type::NONE; ///< The type of partitioning. - std::optional hash; ///< Valid only when type == HASH. - std::optional order; ///< Valid only when type == ORDER. - - /** - * @brief Create a spec indicating no partitioning information. - * @return A PartitioningSpec with type NONE. - */ - static PartitioningSpec none() { - return {}; - } - - /** - * @brief Create a spec indicating partitioning passes through from parent. - * @return A PartitioningSpec with type INHERIT. - */ - static PartitioningSpec inherit() { - return {.type = Type::INHERIT, .hash = std::nullopt, .order = std::nullopt}; - } - - /** - * @brief Create a spec for hash partitioning. - * @param h The hash scheme to use. - * @return A PartitioningSpec with type HASH. - */ - static PartitioningSpec from_hash(HashScheme h) { - return {.type = Type::HASH, .hash = std::move(h), .order = std::nullopt}; - } - - /** - * @brief Create a spec for order/range partitioning. - * @param o The order scheme to use. `o.keys` must be non-empty; otherwise - * throws `std::invalid_argument`. - * @return A PartitioningSpec with type ORDER. - */ - static PartitioningSpec from_order(OrderScheme o); -}; - -/** - * @brief Hierarchical partitioning metadata for a data stream. - * - * Describes how data flowing through a channel is partitioned at multiple - * levels of the system hierarchy. Each level corresponds to a communicator - * used to shuffle data at that level: - * - * - `inter_rank`: Distribution across ranks, corresponding to the primary - * communicator (e.g., `Context::comm()`). Shuffle operations at this level - * move data between ranks. - * - `local`: Distribution within a rank, corresponding to a single-rank - * communicator. Operations at this level repartition data locally without - * network communication. - */ -struct Partitioning { - /// Distribution across ranks (corresponds to primary communicator). - PartitioningSpec inter_rank; - /// Distribution within a rank (corresponds to local/single communicator). - PartitioningSpec local; -}; - -/** - * @brief Channel-level metadata describing the data stream. - * - * Contains information about chunk counts, partitioning, and duplication - * status for the data flowing through a channel. - */ -struct ChannelMetadata { - std::uint64_t local_count{}; ///< Local chunk-count estimate for this rank. - Partitioning partitioning; ///< How the data is partitioned. - bool duplicated{}; ///< Whether data is duplicated on all workers. - - /// @brief Default constructor. - ChannelMetadata() = default; - - /** - * @brief Construct metadata with specified values. - * - * @param local_count Local chunk count. - * @param partitioning Partitioning metadata (default: no partitioning). - * @param duplicated Whether data is duplicated (default: false). - */ - ChannelMetadata( - std::uint64_t local_count, Partitioning partitioning = {}, bool duplicated = false - ) - : local_count(local_count), - partitioning(std::move(partitioning)), - duplicated(duplicated) {} -}; - -/** - * @brief Wrap a `ChannelMetadata` into a `Message`. - * - * @param sequence_number Ordering identifier for the message. - * @param m The metadata to wrap. - * @return A `Message` encapsulating the metadata as its payload. - */ -Message to_message(std::uint64_t sequence_number, std::unique_ptr m); - -} // namespace rapidsmpf::streaming diff --git a/cpp/include/rapidsmpf/streaming/cudf/parquet.hpp b/cpp/include/rapidsmpf/streaming/cudf/parquet.hpp deleted file mode 100644 index 4d7d1707c..000000000 --- a/cpp/include/rapidsmpf/streaming/cudf/parquet.hpp +++ /dev/null @@ -1,64 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include - -namespace rapidsmpf::streaming { - -/** - * @brief Filter ast expression with lifetime/stream management. - */ -struct Filter { - rmm::cuda_stream_view stream; ///< Stream the filter's scalars are valid on. - cudf::ast::expression& filter; ///< Filter expression. - OwningWrapper owner{}; ///< Owner of all objects in the filter. -}; - -namespace actor { -/** - * @brief Asynchronously read parquet files into an output channel. - * - * @note This is a collective operation, all ranks named by the execution context's - * communicator will participate. All ranks must specify the same set of options. - * Behaviour is undefined if a `read_parquet` actor appears only on a subset of the ranks - * named by the communicator, or the options differ between ranks. - * - * @param ctx The execution context to use. - * @param comm Communicator for distributing files across ranks. - * @param ch_out Channel to which `TableChunk`s are sent. - * @param num_producers Number of concurrent producer tasks. - * @param options Template reader options. The files within will be picked apart and used - * to reconstruct new options for each read chunk. The options should therefore specify - * the read options "as-if" one were reading the whole input in one go. - * @param num_rows_per_chunk Target (maximum) number of rows any sent `TableChunk` should - * have. - * @param filter Optional filter expression to apply to the read. - * - * @return Streaming actor representing the asynchronous read. - */ -Actor read_parquet( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::io::parquet_reader_options options, - // TODO: use byte count, not row count? - cudf::size_type num_rows_per_chunk, - std::unique_ptr filter = nullptr -); -} // namespace actor -} // namespace rapidsmpf::streaming diff --git a/cpp/include/rapidsmpf/streaming/cudf/partition.hpp b/cpp/include/rapidsmpf/streaming/cudf/partition.hpp deleted file mode 100644 index 068c0eb1d..000000000 --- a/cpp/include/rapidsmpf/streaming/cudf/partition.hpp +++ /dev/null @@ -1,84 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include -#include - -#include - -#include -#include -#include - -namespace rapidsmpf::streaming { - -namespace actor { - -/** - * @brief Asynchronously partitions input tables into multiple packed (serialized) tables. - * - * This is a streaming version of `rapidsmpf::partition_and_split` that operates on table - * chunks using channels. - * - * It receives tables from an input channel, partitions each row into one of - * `num_partitions` based on a hash of the selected columns, packs the resulting - * partitions, and sends them to an output channel. - * - * @param ctx The actor context to use. - * @param ch_in Input channel providing `TableChunk`s to partition. - * @param ch_out Output channel to which `PartitionMapChunk`s are sent. - * @param columns_to_hash Indices of input columns to hash. - * @param num_partitions The number of partitions to use. - * @param hash_function Hash function to use for partitioning. - * @param seed Seed value for the hash function. - * - * @return Streaming actor representing the asynchronous partitioning and packing - * operation. - * - * @throws std::out_of_range if any index in `columns_to_hash` is invalid. - * - * @see rapidsmpf::partition_and_split - */ -Actor partition_and_pack( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector columns_to_hash, - int num_partitions, - cudf::hash_id hash_function, - std::uint32_t seed -); - - -/** - * @brief Asynchronously unpacks and concatenates packed partitions. - * - * This is a streaming version of `rapidsmpf::unpack_and_concat` that operates on - * packed partition chunks using channels. - * - * It receives packed partitions from the input channel, deserializes and concatenates - * them, and sends the resulting tables to the output channel. Empty partitions are - * ignored. - * - * @param ctx The actor context to use. - * @param ch_in Input channel providing packed partitions as PartitionMapChunk or - * PartitionVectorChunk. - * @param ch_out Output channel to which unpacked and concatenated tables table are sent. - * - * @return Streaming actor representing the asynchronous unpacking and concatenation - * operation. - * - * @see rapidsmpf::unpack_and_concat - */ -Actor unpack_and_concat( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -); - -} // namespace actor -} // namespace rapidsmpf::streaming diff --git a/cpp/include/rapidsmpf/streaming/cudf/table_chunk.hpp b/cpp/include/rapidsmpf/streaming/cudf/table_chunk.hpp deleted file mode 100644 index 0050a5ecd..000000000 --- a/cpp/include/rapidsmpf/streaming/cudf/table_chunk.hpp +++ /dev/null @@ -1,337 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include -#include -#include -#include -#include - -#include -#include -#include -#include - -#include - -#include -#include -#include -#include -#include -#include - -namespace rapidsmpf::streaming { - -/** - * @brief A unit of table data in a streaming pipeline. - * - * Represents either an unpacked `cudf::table` or a `PackedData`. - * - * TableChunks may be initially unavailable (e.g., if the data is packed or spilled), - * and can be made available (i.e., materialized to device memory) on demand. - */ -class TableChunk { - public: - /** - * @brief Indicates whether the TableChunk holds an exclusive or shared view - * of the underlying table data. - * - * This boolean enum is used to explicitly express ownership semantics - * when constructing a TableChunk from a `cudf::table_view`. - * - * - `ExclusiveView::YES`: The TableChunk has exclusive ownership of - * the table's device memory and are considered spillable. - * - * - `ExclusiveView::NO`: The TableChunk is a non-owning view of data - * managed elsewhere. The memory may be shared or externally owned, - * and the chunk is therefore not spillable. - */ - enum class ExclusiveView : bool { - NO, - YES, - }; - - /** - * @brief Construct a TableChunk from a device table. - * - * @param table Device-resident table. - * @param stream The CUDA stream on which the table was created. - */ - TableChunk(std::unique_ptr table, rmm::cuda_stream_view stream); - - /** - * @brief Construct a TableChunk from a device table view. - * - * The TableChunk does not take ownership of the underlying data; instead, the - * provided @p owner object is kept alive for the lifetime of the TableChunk. - * The caller is responsible for ensuring that the underlying device memory - * referenced by @p table_view remains valid during this period. - * - * This constructor is typically used when creating a TableChunk from Python, - * where @p owner is used to keep the corresponding Python object alive until - * the TableChunk is destroyed. - * - * @param table_view Device-resident table view. - * @param stream CUDA stream on which the table was created. - * @param owner Object owning the memory backing @p table_view. This object will be - * destroyed last when the TableChunk is destroyed or spilled. - * @param exclusive_view Specifies whether this TableChunk has exclusive ownership - * semantics over the underlying table data: - * - When `ExclusiveView::YES`, the following guarantees must hold: - * - The @p table_view is the sole representation of the table. - * - The @p owner exclusively owns the table memory. - * These guarantees allow the TableChunk to be spillable and ensure that - * destroying @p owner will correctly free the associated device memory. - * - When `ExclusiveView::NO`, the chunk is considered a non-owning view and - * is therefore not spillable. - */ - TableChunk( - cudf::table_view table_view, - rmm::cuda_stream_view stream, - OwningWrapper&& owner, - ExclusiveView exclusive_view - ); - - /** - * @brief Construct a TableChunk from a packed data blob. - * - * The packed data's CUDA stream will be associated the new table chunk. - * - * @param packed_data Serialized host/device data with metadata. - */ - TableChunk(std::unique_ptr packed_data); - - ~TableChunk() = default; - - /** - * @brief Move constructor - * - * @note After this call `other.is_available() == false`. - * @param other The TableChunk to move from. - */ - TableChunk(TableChunk&& other) noexcept; - - /** - * @brief Move assignment - * - * @note After this call `other.is_available() == false`. - * @param other The TableChunk to move from. - * @return Reference to this. - */ - TableChunk& operator=(TableChunk&& other) noexcept; - TableChunk(TableChunk const&) = delete; - TableChunk& operator=(TableChunk const&) = delete; - - /** - * @brief Returns the CUDA stream on which this table chunk was created. - * - * @return The CUDA stream view. - */ - [[nodiscard]] rmm::cuda_stream_view stream() const noexcept; - - /** - * @brief Number of bytes allocated for the data in the specified memory type. - * - * @param mem_type The memory type to query. - * @return Number of bytes allocated. - */ - [[nodiscard]] std::size_t data_alloc_size(MemoryType mem_type) const; - - /** - * @brief Indicates whether the underlying cudf table data is fully available in - * device memory. - * - * @return `true` if the table is already available; otherwise, `false`. - */ - [[nodiscard]] bool is_available() const noexcept; - - /** - * @brief Returns the estimated cost (in bytes) of making the table available. - * - * Currently, only device memory cost is tracked. - * - * @return The cost in bytes. - */ - [[nodiscard]] std::size_t make_available_cost() const noexcept; - - /** - * @brief Moves this table chunk into a new one with its cudf table made available. - * - * As part of the move, a copy or unpack may be performed, the associated CUDA - * stream is used. - * - * @param reservation Memory reservation for allocations if needed. - * @return A new TableChunk with data available on device. - * - * @note After this call, the current object is in a moved-from state; - * only reassignment, movement, or destruction are valid. - */ - [[nodiscard]] TableChunk make_available(MemoryReservation& reservation); - - /** - * @brief Moves this table chunk into a new one with its cudf table made available. - * - * Takes ownership of the memory reservation and consumes it entirely as part - * of making the data available on device. The full reservation is considered - * used, even if the actual allocation requires fewer bytes. - * - * @param reservation Memory reservation to be consumed for allocations. - * @return A new TableChunk with data available on device. - * - * @note After this call, the current object is in a moved-from state; only - * reassignment, movement, or destruction are valid. - */ - [[nodiscard]] TableChunk make_available(MemoryReservation&& reservation); - - /** - * @brief Move this table chunk into a new one with its cudf table made available. - * - * This variant of make_available() is a coroutine that may suspend if device - * memory is not immediately available. - * - * @note After this call, the current object is in a moved-from state; only - * reassignment, movement, or destruction are valid. - * - * @param ctx Streaming context used to access the memory reservation mechanism. - * @param net_memory_delta Estimated change in memory usage after the reservation - * is granted and all work using the returned `TableChunk` has completed. See - * `MemoryReserveOrWait::reserve_or_wait` for details. - * @return A new `TableChunk` that is available on device. - * - * @throws std::runtime_error If shutdown occurs before the reservation can be - * processed. - * @throws std::overflow_error If no progress is possible within the timeout and - * overbooking is disabled. - */ - [[nodiscard]] coro::task make_available( - std::shared_ptr ctx, - std::int64_t net_memory_delta = MemoryReserveOrWait::missing_net_memory_delta - ); - - /** - * @brief Returns a view of the underlying table. - * - * The table must be available in device memory. - * - * @return cudf::table_view representing the table. - * - * @throws std::invalid_argument if `is_available() == false`. - */ - [[nodiscard]] cudf::table_view table_view() const; - - /** - * @brief Indicates whether this table chunk can be spilled to device memory. - * - * A table chunk is considered spillable if it owns its underlying memory. This is - * true when it was created from one of the following: - * - A device-owning source such as a `cudf::table`, `cudf::packed_columns`, or - * `PackedData`. - * - A `cudf::table_view` constructed with `is_exclusive_view == true`, indicating - * that the view is the sole representation of the underlying data and that its - * owner exclusively manages the table's memory. - * - * In contrast, chunks constructed from non-exclusive `cudf::table_view` instances are - * non-owning views of externally managed memory and therefore not spillable. - * - * To spill a table chunk from device to host memory, first call `copy()` to create a - * host-side copy, then delete or overwrite the original device chunk. If - * `is_spillable() == true`, destroying the original device chunk will release the - * associated device memory. - * - * @return `true` if the table chunk owns its memory and can be spilled; otherwise - * `false`. - */ - [[nodiscard]] bool is_spillable() const; - - /** - * @brief Create a deep copy of the table chunk. - * - * Allocates new memory for all buffers in the table using the specified - * `reservation`, which determines the target memory type (e.g., host or device). - * As a consequence, the `is_available()` status may differ in the new copy. For - * example, copying an available table chunk from device to host memory will result - * in an unavailable copy. - * - * @param reservation Memory reservation used to track and limit allocations. - * @return A new `TableChunk` instance containing copies of all buffers and metadata. - * - * @throws rapidsmpf::reservation_error If the total allocation size exceeds the - * available reservation. - */ - [[nodiscard]] TableChunk copy(MemoryReservation& reservation) const; - - /** - * @brief Convert this table chunk to a `PackedData`, avoiding unnecessary copies. - * - * If the chunk's data is already in packed form (e.g., it arrived over the network - * or was constructed from a `PackedData`), the packed data is moved out directly - * with no copy. Otherwise the table is serialized via `cudf::pack()`. - * - * @param br Buffer resource used for the device memory resource when packing - * is required. - * @return A unique pointer to the resulting `PackedData`. - * - * @throws std::invalid_argument If the data is not already packed and - * `is_available() == false`. - * - * @note After this call, this object is in a moved-from state; only reassignment, - * movement, or destruction are valid. - * - * @note No memory reservation is required. If the data is already in packed form, - * no allocation occurs. If packing is required, `cudf::pack()` allocates device - * memory that is not tracked via a reservation. - */ - [[nodiscard]] std::unique_ptr into_packed_data(BufferResource* br) &&; - - /** - * @brief Return the shape of the table stored by the table chunk. - * - * @return Pair of number of rows and number of columns. - */ - [[nodiscard]] std::pair shape() const noexcept; - - private: - ///< @brief Optional owning object if the TableChunk was constructed from a - ///< table_view. - OwningWrapper owner_{}; - - // At most, one of the following unique pointers is non-null. If all of them are null, - // the TableChunk is a non-owning view. - // TODO: use a variant and drop the unique pointers? - std::unique_ptr table_; - std::unique_ptr packed_data_; - - // Has value iff this TableChunk is available. - std::optional table_view_; - - // Zero initialized data allocation size (one for each memory type). - std::array data_alloc_size_ = {}; - std::size_t make_available_cost_; // For now, only device memory cost is tracked. - - rmm::cuda_stream_view stream_; - bool is_spillable_; -}; - -/** - * @brief Generate a content description for a `TableChunk`. - * - * @param obj The object's content to describe. - * @return A new content description. - */ -ContentDescription get_content_description(TableChunk const& obj); - -/** - * @brief Wrap a `TableChunk` into a `Message`. - * - * @param sequence_number Ordering identifier for the message. - * @param chunk The chunk to wrap into a message. - * @return A `Message` encapsulating the provided chunk as its payload. - */ -Message to_message(std::uint64_t sequence_number, std::unique_ptr chunk); - -} // namespace rapidsmpf::streaming diff --git a/cpp/src/integrations/cudf/bloom_filter.cu b/cpp/src/integrations/cudf/bloom_filter.cu deleted file mode 100644 index ad169657d..000000000 --- a/cpp/src/integrations/cudf/bloom_filter.cu +++ /dev/null @@ -1,168 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include - -#include - -// cuco headers have sign-conversion issues; suppress for the host compiler -#ifdef __clang__ -#pragma clang diagnostic push -#pragma clang diagnostic ignored "-Wsign-conversion" -#elif defined(__GNUC__) -#pragma GCC diagnostic push -#pragma GCC diagnostic ignored "-Wsign-conversion" -#endif -#include -#include -#include -#include -#ifdef __clang__ -#pragma clang diagnostic pop -#elif defined(__GNUC__) -#pragma GCC diagnostic pop -#endif - -#include -#include - -#include -#include -#include -#include - -#include -#include -#include - -namespace rapidsmpf { - -namespace { -using KeyType = std::uint64_t; - -using BloomFilterRefType = cuco::bloom_filter_ref< - KeyType, - cuco::extent, - cuco::thread_scope_device, - cuco::arrow_filter_policy>; -using StorageType = BloomFilterRefType::filter_block_type; - -} // namespace - -BloomFilter::BloomFilter( - std::size_t num_blocks, - std::uint64_t seed, - void* storage, - rmm::cuda_stream_view stream -) - : num_blocks_{num_blocks}, seed_{seed}, storage_{storage}, stream_{stream} { - // TODO: use an aligned allocator adaptor to ensure this holds. - // Today all RMM device allocators guarantee at least 256 byte alignment, but that is - // an implementation detail. - RAPIDSMPF_EXPECTS( - reinterpret_cast(storage_) % std::alignment_of_v - == 0, - "Allocation for bloom filter is not aligned." - ); -} - -BloomFilter const BloomFilter::view( - std::size_t num_blocks, - std::uint64_t seed, - void const* storage, - rmm::cuda_stream_view stream -) { - // const-cast is safe because the returned object is also const and therefore can't - // call methods that throw away constness. - return BloomFilter(num_blocks, seed, const_cast(storage), stream); -} - -std::unique_ptr BloomFilter::storage( - std::size_t num_blocks, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -) { - return std::make_unique( - num_blocks * sizeof(StorageType), stream, mr - ); -} - -void BloomFilter::add( - cudf::table_view const& values_to_hash, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -) { - RAPIDSMPF_NVTX_FUNC_RANGE(); - auto filter_ref = BloomFilterRefType{ - static_cast(storage_), num_blocks_, cuco::thread_scope_device, {} - }; - auto hashes = cudf::hashing::xxhash_64(values_to_hash, seed_, stream, mr); - auto hash_view = hashes->view(); - RAPIDSMPF_EXPECTS( - hash_view.type().id() == cudf::type_to_id(), - "Hash values do not have correct type" - ); - filter_ref.add_async(hash_view.begin(), hash_view.end(), stream); -} - -void BloomFilter::merge(BloomFilter const& other, rmm::cuda_stream_view stream) { - RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_EXPECTS( - num_blocks_ == other.num_blocks_, "Mismatching number of blocks in filters" - ); - auto ref_this = BloomFilterRefType{ - static_cast(storage_), num_blocks_, cuco::thread_scope_device, {} - }; - auto ref_other = BloomFilterRefType{ - static_cast(other.storage_), - num_blocks_, - cuco::thread_scope_device, - {} - }; - ref_this.merge_async(ref_other, stream); -} - -rmm::device_uvector BloomFilter::contains( - cudf::table_view const& values, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -) const { - RAPIDSMPF_NVTX_FUNC_RANGE(); - auto filter_ref = BloomFilterRefType{ - static_cast(storage_), num_blocks_, cuco::thread_scope_device, {} - }; - auto hashes = cudf::hashing::xxhash_64(values, seed_, stream, mr); - auto view = hashes->view(); - rmm::device_uvector result{static_cast(view.size()), stream, mr}; - filter_ref.contains_async( - view.begin(), view.end(), result.begin(), stream - ); - return result; -} - -std::size_t BloomFilter::fitting_num_blocks(std::size_t l2size) noexcept { - return (l2size * 2) / (3 * sizeof(StorageType)); -} - -rmm::cuda_stream_view BloomFilter::stream() const noexcept { - return stream_; -} - -void* BloomFilter::data() noexcept { - return storage_; -} - -void const* BloomFilter::data() const noexcept { - return storage_; -} - -std::size_t BloomFilter::size() const noexcept { - return num_blocks_ * sizeof(StorageType); -} - -} // namespace rapidsmpf diff --git a/cpp/src/integrations/cudf/partition.cpp b/cpp/src/integrations/cudf/partition.cpp deleted file mode 100644 index b423c0803..000000000 --- a/cpp/src/integrations/cudf/partition.cpp +++ /dev/null @@ -1,256 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace rapidsmpf { - -std::pair, std::unique_ptr> -partition_and_split( - cudf::table_view const& table, - std::vector const& columns_to_hash, - int num_partitions, - cudf::hash_id hash_function, - std::uint32_t seed, - rmm::cuda_stream_view stream, - BufferResource* br, - AllowOverbooking allow_overbooking -) { - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); - if (table.num_rows() == 0) { - // Return views of a copy of the empty `table`. - auto owner = std::make_unique(table, stream, br->device_mr()); - return { - std::vector( - safe_cast(num_partitions), owner->view() - ), - std::move(owner) - }; - } - - // hash_partition does a deep-copy. Therefore, we need to reserve memory for - // at least the size of the table. - auto reservation = br->reserve_device_memory_and_spill( - estimated_memory_usage(table, stream), allow_overbooking - ); - auto [partition_table, offsets] = cudf::hash_partition( - table, - columns_to_hash, - num_partitions, - hash_function, - seed, - stream, - br->device_mr() - ); - reservation.clear(); - - // Notice, the offset argument for split() and hash_partition() doesn't align. - // hash_partition() returns the start offset of each partition thus we have to - // skip the first offset. See: . - auto partition_offsets = - cudf::host_span(offsets.data() + 1, offsets.size() - 2); - - // split does not make any copies. - auto tbl_partitioned = - cudf::split(partition_table->view(), partition_offsets, stream); - - return {std::move(tbl_partitioned), std::move(partition_table)}; -} - -std::unordered_map partition_and_pack( - cudf::table_view const& table, - std::vector const& columns_to_hash, - int num_partitions, - cudf::hash_id hash_function, - std::uint32_t seed, - rmm::cuda_stream_view stream, - BufferResource* br, - AllowOverbooking allow_overbooking -) { - RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); - RAPIDSMPF_EXPECTS(num_partitions > 0, "Need to split to at least one partition"); - if (table.num_rows() == 0) { - auto splits = - std::vector(safe_cast(num_partitions - 1), 0); - return split_and_pack(table, splits, stream, br, allow_overbooking); - } - - // hash_partition does a deep-copy. Therefore, we need to reserve memory for - // at least the size of the table. - auto reservation = br->reserve_device_memory_and_spill( - estimated_memory_usage(table, stream), allow_overbooking - ); - auto [reordered, split_points] = cudf::hash_partition( - table, - columns_to_hash, - num_partitions, - hash_function, - seed, - stream, - br->device_mr() - ); - reservation.clear(); - std::vector splits(split_points.begin() + 1, split_points.end() - 1); - return split_and_pack(reordered->view(), splits, stream, br, allow_overbooking); -} - -std::unordered_map split_and_pack( - cudf::table_view const& table, - std::vector const& splits, - rmm::cuda_stream_view stream, - BufferResource* br, - AllowOverbooking allow_overbooking -) { - RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); - std::unordered_map ret; - - // contiguous split does a deep-copy. Therefore, we need to reserve memory for - // at least the size of the table. - auto reservation = br->reserve_device_memory_and_spill( - cudf::packed_size(table, stream, br->device_mr()), allow_overbooking - ); - auto packed = cudf::contiguous_split(table, splits, stream, br->device_mr()); - reservation.clear(); - ret.reserve(packed.size()); - for (shuffler::PartID i = 0; safe_cast(i) < packed.size(); i++) { - auto pack = std::move(packed[i].data); - ret.emplace( - i, - PackedData( - std::move(pack.metadata), br->move(std::move(pack.gpu_data), stream) - ) - ); - } - return ret; -} - -std::unique_ptr unpack_and_concat( - std::vector&& partitions, - rmm::cuda_stream_view stream, - BufferResource* br, - AllowOverbooking allow_overbooking -) { - RAPIDSMPF_NVTX_FUNC_RANGE(); - RAPIDSMPF_MEMORY_PROFILE(br->statistics(), br->device_mr()); - - // Let's find the total size of the partitions and how much of the packed data we - // need to move to device memory (unspill). - std::size_t total_size = 0; - std::size_t non_device_size = 0; - for (auto& packed_data : partitions) { - if (!packed_data.empty()) { - std::size_t size = packed_data.data->size; - total_size += size; - if (packed_data.data->mem_type() != MemoryType::DEVICE) { - non_device_size += size; - } - } - } - - std::vector unpacked; - std::vector references; - std::vector packed_data_streams; - unpacked.reserve(partitions.size()); - references.reserve(partitions.size()); - packed_data_streams.reserve(partitions.size()); - - // Reserve device memory for the unspill AND the cudf::unpack() calls. - auto reservation = br->reserve_device_memory_and_spill( - total_size + non_device_size, allow_overbooking - ); - for (auto& packed_data : partitions) { - if (!packed_data.empty()) { - if (packed_data.data->size > 0) { // No need to sync empty buffers. - packed_data_streams.push_back(packed_data.data->stream()); - } - unpacked.push_back( - cudf::unpack(references.emplace_back( - std::move(packed_data.metadata), - br->move_to_device_buffer(std::move(packed_data.data), reservation) - )) - ); - } - } - reservation.clear(); - - // We need to synchronize `stream` with the packed_data and update their - // underlying device buffers to use `stream` going forward. This ensures - // the packed data are not deallocated before we have a chance to - // concatenate them on `stream`. - cuda_stream_join(std::views::single(stream), packed_data_streams); - for (cudf::packed_columns& packed_columns : references) { - packed_columns.gpu_data->set_stream(stream); - } - - reservation = br->reserve_device_memory_and_spill(total_size, allow_overbooking); - return cudf::concatenate(unpacked, stream, br->device_mr()); -} - -std::vector spill_partitions( - std::vector&& partitions, BufferResource* br -) { - // Sum the total size of all packed data in device memory. - std::size_t device_size{0}; - for (auto& [_, data] : partitions) { - if (data->mem_type() == MemoryType::DEVICE) { - device_size += data->size; - } - } - // Spill each partition to host memory. - auto reservation = br->reserve_or_fail(device_size, SPILL_TARGET_MEMORY_TYPES); - std::vector ret; - ret.reserve(partitions.size()); - for (auto& [metadata, data] : partitions) { - ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation)); - } - return ret; -} - -std::vector unspill_partitions( - std::vector&& partitions, - BufferResource* br, - AllowOverbooking allow_overbooking -) { - auto statistics = br->statistics(); - // Sum the total size of all packed data not in device memory already. - std::size_t non_device_size{0}; - for (auto& [_, data] : partitions) { - if (data->mem_type() != MemoryType::DEVICE) { - non_device_size += data->size; - } - } - - // Unspill each partition. - auto reservation = - br->reserve_device_memory_and_spill(non_device_size, allow_overbooking); - std::vector ret; - ret.reserve(partitions.size()); - for (auto& [metadata, data] : partitions) { - ret.emplace_back(std::move(metadata), br->move(std::move(data), reservation)); - } - - return ret; -} -} // namespace rapidsmpf diff --git a/cpp/src/integrations/cudf/utils.cpp b/cpp/src/integrations/cudf/utils.cpp deleted file mode 100644 index 1d441edd5..000000000 --- a/cpp/src/integrations/cudf/utils.cpp +++ /dev/null @@ -1,151 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -#include -#include - -namespace rapidsmpf { - -namespace { -struct str_cudf_column_scalar_fn { - template - requires(cudf::is_numeric()) - std::string operator()( - cudf::column_view col, - cudf::size_type index, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr - ) { - std::unique_ptr scalar = cudf::get_element(col, index, stream, mr); - auto typed_scalar = static_cast const*>(scalar.get()); - T val = typed_scalar->value(stream); - return std::to_string(val); - } - - template - requires(!cudf::is_numeric()) - std::string operator()( - cudf::column_view /* col */, - cudf::size_type /* index */, - rmm::cuda_stream_view /* stream */, - rmm::device_async_resource_ref /* mr */ - ) { - RAPIDSMPF_FAIL("not implemented"); - } -}; - -struct cudf_column_data_size_fn { - template - requires(cudf::is_fixed_width()) - std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view) { - return safe_cast(col.size()) * cudf::size_of(col.type()) - + bitmask_size(col); - } - - // string type specialization - template - requires(std::is_same_v) - std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view stream) { - cudf::strings_column_view sv(col); - return safe_cast(sv.chars_size(stream)) + bitmask_size(col); - } - - // compound type specialization except string - template - requires(!std::is_same_v && cudf::is_compound()) - std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view) { - // compound types (except string) ie. list, dict, structs dont have a - // content::data buffer. Data is stored in children columns. So, just return the - // bitmask size. - return bitmask_size(col); - } - - template - std::size_t operator()(cudf::column_view const& col, rmm::cuda_stream_view) { - RAPIDSMPF_FAIL("not implemented for type: " + cudf::type_to_name(col.type())); - } - - static std::size_t bitmask_size(cudf::column_view const& col) { - return col.nullable() ? cudf::bitmask_allocation_size_bytes(col.size()) : 0; - } -}; - -} // namespace - -std::string str( - cudf::column_view col, - cudf::size_type index, - rmm::cuda_stream_view stream, - rmm::device_async_resource_ref mr -) { - return cudf::type_dispatcher( - col.type(), str_cudf_column_scalar_fn{}, col, index, stream, mr - ); -} - -std::string str( - cudf::column_view col, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr -) { - std::stringstream ss; - ss << "Column(["; - for (cudf::size_type i = 0; i < col.size(); ++i) { - ss << str(col, i, stream, mr) << ", "; - } - ss << (col.size() == 0 ? "])" : "\b\b])"); - return ss.str(); -} - -std::string str( - cudf::table_view tbl, rmm::cuda_stream_view stream, rmm::device_async_resource_ref mr -) { - std::stringstream ss; - ss << "Table(["; - for (auto col : tbl) { - ss << str(col, stream, mr) << ", "; - } - ss << (tbl.num_columns() == 0 ? "])" : "\b\b])"); - return ss.str(); -} - -std::size_t estimated_memory_usage( - cudf::column_view const& col, rmm::cuda_stream_view stream -) { - return std::transform_reduce( - col.child_begin(), - col.child_end(), - cudf::type_dispatcher(col.type(), cudf_column_data_size_fn{}, col, stream), - std::plus{}, - [&stream](cudf::column_view const& child) { - return estimated_memory_usage(child, stream); - } - ); -} - -std::size_t estimated_memory_usage( - cudf::table_view const& tbl, rmm::cuda_stream_view stream -) { - return std::transform_reduce( - tbl.begin(), - tbl.end(), - std::size_t{0}, - std::plus{}, - [&stream](cudf::column_view const& col) { - return estimated_memory_usage(col, stream); - } - ); -} - -} // namespace rapidsmpf diff --git a/cpp/src/streaming/cudf/bloom_filter.cpp b/cpp/src/streaming/cudf/bloom_filter.cpp deleted file mode 100644 index ac41e10ee..000000000 --- a/cpp/src/streaming/cudf/bloom_filter.cpp +++ /dev/null @@ -1,156 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace rapidsmpf::streaming { -Actor BloomFilter::build( - std::shared_ptr ch_in, std::shared_ptr ch_out, OpID tag -) { - ShutdownAtExit c{ch_in, ch_out}; - co_await ctx_->executor()->schedule(); - co_await ch_in->shutdown_metadata(); - co_await ch_out->shutdown_metadata(); - auto const& br = ctx_->br(); - auto mr = br->device_mr(); - auto filter_stream = br->stream_pool().get_stream(); - CudaEvent event; - auto storage = rapidsmpf::BloomFilter::storage(num_filter_blocks_, filter_stream, mr); - RAPIDSMPF_CUDA_TRY( - cudaMemsetAsync(storage->data(), 0, storage->size(), filter_stream) - ); - auto filter = - rapidsmpf::BloomFilter(num_filter_blocks_, seed_, storage->data(), filter_stream); - CudaEvent build_event; - build_event.record(filter_stream); - while (!ch_out->is_shutdown()) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = msg.release(); - chunk = co_await chunk.make_available( - ctx_, -safe_cast(chunk.data_alloc_size(MemoryType::DEVICE)) - ); - // Filter is allocated on `filter_stream`, but we run the additions on the chunk's - // stream. The addition modifies global memory but we can safely launch two - // kernels doing that concurrently because the updates are atomic. - build_event.stream_wait(chunk.stream()); - filter.add(chunk.table_view(), chunk.stream(), mr); - cuda_stream_join(filter_stream, chunk.stream(), &event); - } - if (comm_->nranks() > 1) { - auto reducer = streaming::AllReduce( - ctx_, - comm_, - br->move(std::move(storage), filter_stream), - br->move( - rapidsmpf::BloomFilter::storage(num_filter_blocks_, filter_stream, mr), - filter_stream - ), - tag, - [num_blocks = num_filter_blocks_, - seed = seed_](Buffer const* left, Buffer* right) { - right->write_access([&](std::byte* out_bytes, - rmm::cuda_stream_view stream) { - auto const in = rapidsmpf::BloomFilter::view( - num_blocks, seed, left->data(), stream - ); - rapidsmpf::BloomFilter(num_blocks, seed, out_bytes, stream) - .merge(in, stream); - }); - } - ); - auto result = co_await reducer.extract(); - auto [res, _] = br->reserve(MemoryType::DEVICE, 0, AllowOverbooking::YES); - storage = br->move_to_device_buffer(std::move(result.second), res); - } - co_await ch_out->send(Message{0, std::move(storage), {}, {}}); - co_await ch_out->drain(ctx_->executor()); -} - -Actor BloomFilter::apply( - std::shared_ptr bloom_filter, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys -) { - streaming::ShutdownAtExit c{bloom_filter, ch_in, ch_out}; - co_await ctx_->executor()->schedule(); - auto storage = (co_await bloom_filter->receive()).release(); - RAPIDSMPF_EXPECTS( - (co_await bloom_filter->receive()).empty(), - "Bloom filter channel contained more than one message" - ); - auto stream = storage.stream(); - CudaEvent event; - auto filter = - rapidsmpf::BloomFilter(num_filter_blocks_, seed_, storage.data(), stream); - auto meta = co_await ch_in->receive_metadata(); - if (!meta.empty()) { - co_await ch_out->send_metadata(std::move(meta)); - } - while (!ch_out->is_shutdown()) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - auto chunk = msg.release(); - chunk = co_await chunk.make_available( - ctx_, -safe_cast(chunk.data_alloc_size(MemoryType::DEVICE)) - ); - auto chunk_stream = chunk.stream(); - cuda_stream_join(chunk_stream, stream, &event); - // Reservation for the mask construction and guess at output size. - auto res = co_await ctx_->memory(MemoryType::DEVICE) - ->reserve_or_wait( - safe_cast(chunk.table_view().num_rows()) - // TODO: no magic numbers: the hashing algorithm in - // `contains` below returns an int64 column. - * (1 + sizeof(std::int64_t)) - // Guess at how selective the filter is. - + chunk.data_alloc_size(MemoryType::DEVICE) / 2, - 0 - ); - auto mask = filter.contains( - chunk.table_view().select(keys), chunk_stream, ctx_->br()->device_mr() - ); - cuda_stream_join(stream, chunk_stream, &event); - RAPIDSMPF_EXPECTS( - mask.size() == static_cast(chunk.table_view().num_rows()), - "Invalid mask size" - ); - auto mask_view = cudf::column_view{ - cudf::data_type{cudf::type_id::BOOL8}, - safe_cast(mask.size()), - mask.data(), - {}, - 0 - }; - auto result = cudf::apply_boolean_mask( - chunk.table_view(), mask_view, chunk_stream, ctx_->br()->device_mr() - ); - std::ignore = std::move(chunk); - std::ignore = std::move(res); - co_await ch_out->send(to_message( - msg.sequence_number(), - std::make_unique(std::move(result), chunk_stream) - )); - } - co_await ch_out->drain(ctx_->executor()); -} -} // namespace rapidsmpf::streaming diff --git a/cpp/src/streaming/cudf/channel_metadata.cpp b/cpp/src/streaming/cudf/channel_metadata.cpp deleted file mode 100644 index fc6bec396..000000000 --- a/cpp/src/streaming/cudf/channel_metadata.cpp +++ /dev/null @@ -1,121 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include - -#include -#include -#include -#include - -#include -#include -#include -#include - -namespace rapidsmpf::streaming { - -OrderScheme::OrderScheme( - std::vector keys, - std::shared_ptr boundaries, - bool strict_boundaries -) - : keys{std::move(keys)}, - boundaries{std::move(boundaries)}, - strict_boundaries{strict_boundaries} { - RAPIDSMPF_EXPECTS( - !this->keys.empty(), "OrderScheme: keys must not be empty", std::invalid_argument - ); - RAPIDSMPF_EXPECTS( - this->boundaries != nullptr, - "OrderScheme: boundaries must not be null", - std::invalid_argument - ); - RAPIDSMPF_EXPECTS( - this->boundaries->is_available(), - "OrderScheme: boundaries must be device-resident", - std::invalid_argument - ); - RAPIDSMPF_EXPECTS( - this->keys.size() == static_cast(this->boundaries->shape().second), - "OrderScheme: number of keys must match number of boundary columns", - std::invalid_argument - ); -} - -PartitioningSpec PartitioningSpec::from_order(OrderScheme o) { - return {.type = Type::ORDER, .hash = std::nullopt, .order = std::move(o)}; -} - -OrderScheme OrderScheme::with_keys(std::vector new_keys) const { - return OrderScheme(std::move(new_keys), boundaries, strict_boundaries); -} - -bool OrderScheme::boundaries_aligned_with( - OrderScheme const& other, rapidsmpf::BufferResource& br -) const { - if (strict_boundaries != other.strict_boundaries - || boundaries->shape() != other.boundaries->shape()) - { - return false; - } - if (!std::equal( - keys.begin(), - keys.end(), - other.keys.begin(), - [](OrderKey const& a, OrderKey const& b) { - return a.order == b.order && a.null_order == b.null_order; - } - )) - { - return false; - } - if (boundaries->shape().first == 0) { - return true; - } - auto const lhs = boundaries->table_view(); - auto const rhs = other.boundaries->table_view(); - auto const stream = boundaries->stream(); - cuda_stream_join(stream, other.boundaries->stream()); - for (cudf::size_type i = 0; i < lhs.num_columns(); ++i) { - auto eq = cudf::binary_operation( - lhs.column(i), - rhs.column(i), - cudf::binary_operator::NULL_EQUALS, - cudf::data_type{cudf::type_id::BOOL8}, - stream, - br.device_mr() - ); - auto result = cudf::reduce( - eq->view(), - *cudf::make_all_aggregation(), - cudf::data_type{cudf::type_id::BOOL8}, - stream, - br.device_mr() - ); - auto& scalar = static_cast&>(*result); - if (!scalar.value(stream)) { - return false; - } - } - return true; -} - -Message to_message(std::uint64_t sequence_number, std::unique_ptr m) { - return Message{ - sequence_number, - std::move(m), - {}, - [](Message const& msg, MemoryReservation& /* reservation */) -> Message { - auto copy = std::make_unique(msg.get()); - return Message{msg.sequence_number(), std::move(copy), {}, msg.copy_cb()}; - } - }; -} - -} // namespace rapidsmpf::streaming diff --git a/cpp/src/streaming/cudf/parquet.cpp b/cpp/src/streaming/cudf/parquet.cpp deleted file mode 100644 index 2edd16625..000000000 --- a/cpp/src/streaming/cudf/parquet.cpp +++ /dev/null @@ -1,445 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace rapidsmpf::streaming::actor { -namespace { - -/** - * @brief Per-context cache for file-backed messages. - * - * FileCache caches file read results by storing message copies in the associated - * Context's SpillableMessages instance. By tying cached data to the Context, the - * lifetime of cached entries matches the lifetime of the Context itself. - * - * Each cache instance is scoped to a single Context and is shared across callers - * using that Context. - * - * The cache is thread-safe. - */ -class FileCache { - public: - struct Key { - std::vector filepaths; - std::int64_t skip_rows; - std::size_t skip_bytes; - std::optional num_rows; - std::optional num_bytes; - std::optional> column_names; - std::optional> column_indices; - std::vector> row_groups; - - // Lexicographical comparison of all data members. - auto operator<=>(Key const&) const = default; - }; - - /** - * @brief Construct a FileCache. - * - * @param mem_type Memory type used for cache storage. - */ - FileCache(MemoryType mem_type = MemoryType::HOST) : mem_type_{mem_type} {} - - /** - * @brief Insert a message into the cache. - * - * The message is copied into the memory type configured for this cache - * and stored in the associated Context's SpillableMessages instance. - * - * @param ctx Streaming context. - * @param key Cache key identifying the message. - * @param msg Message to cache. - * @return True if the message was inserted, false if the key already existed. - */ - bool insert(std::shared_ptr ctx, Key key, Message const& msg) { - auto reservation = ctx->br()->reserve_or_fail(msg.copy_cost(), mem_type_); - auto msg_copy = msg.copy(reservation); - - std::lock_guard lock(mutex_); - if (cache_.contains(key)) { - return false; - } - cache_.emplace( - std::move(key), ctx->spillable_messages()->insert(std::move(msg_copy)) - ); - return true; - } - - /** - * @brief Retrieve a cached message. - * - * If the key exists, the cached message is copied out of spillable storage - * using newly reserved memory, prioritizing memory types in `MEMORY_TYPES` - * order. - * - * @param ctx Streaming context. - * @param key Cache key to look up. - * @return The cached message, or std::nullopt if the key is not present. - */ - std::optional get(std::shared_ptr ctx, Key const& key) const { - auto& stats = *ctx->statistics(); - - stats.add_report_entry( - "unbounded_file_read_cache hits", - {"unbounded_file_read_cache hits"}, - Statistics::Formatter::HitRate - ); - - SpillableMessages::MessageId mid; - { - std::lock_guard lock(mutex_); - auto it = cache_.find(key); - if (it == cache_.end()) { - stats.add_stat("unbounded_file_read_cache hits", 0); - return std::nullopt; - } - mid = it->second; - } - auto const size = - ctx->spillable_messages()->get_content_description(mid).content_size(); - - stats.add_stat("unbounded_file_read_cache hits", 1); - stats.add_bytes_stat("unbounded_file_read_cache saved", size); - auto reservation = ctx->br()->reserve_or_fail(size, MEMORY_TYPES); - return ctx->spillable_messages()->copy(mid, reservation); - } - - /** - * @brief Get the FileCache instance for a Context. - * - * Each Context has exactly one FileCache instance for the lifetime of the - * process. If the `unbounded_file_read_cache` option is disabled, this - * function returns nullptr. - * - * @param ctx Context used to identify the cache instance. The same Context must - * be used for all subsequent insert and get operations. - * @return Shared pointer to the per-context FileCache, or nullptr if the cache - * is disabled. - */ - static std::shared_ptr instance(std::shared_ptr ctx) { - static std::mutex mutex; - static std::unordered_map> instances; - - std::lock_guard lock(mutex); - auto const id = ctx->uid(); - auto it = instances.find(id); - if (it != instances.end()) { - return it->second; - } - - // Get the memory type of the file cache, if enabled. - auto const mem_type = ctx->options().get>( - "unbounded_file_read_cache", [](auto const& s) -> std::optional { - auto val = parse_optional(s); - if (!val.has_value() || val->empty()) { - return std::nullopt; - } - return parse_string(s); - } - ); - - if (mem_type.has_value()) { - auto ret = std::make_shared(*mem_type); - instances.emplace(id, ret); - return ret; - } - return nullptr; - } - - private: - mutable std::mutex mutex_; - std::map cache_; - MemoryType mem_type_; -}; - -/** - * @brief Read a single chunk from a parquet source. - * - * @param ctx The execution context to use. - * @param stream The stream on which to read the chunk. - * @param options The parquet reader options describing the data to read. - * @param sequence_number The ordered chunk id to reconstruct original ordering of the - * data. - * @return Message representing the read chunk. - */ -Message read_parquet_chunk( - std::shared_ptr ctx, - rmm::cuda_stream_view stream, - cudf::io::parquet_reader_options options, - std::uint64_t sequence_number -) { - auto do_read_parquet = [&]() -> Message { - return to_message( - sequence_number, - std::make_unique( - cudf::io::read_parquet(options, stream, ctx->br()->device_mr()).tbl, - stream - ) - ); - }; - - auto file_cache = FileCache::instance(ctx); - if (file_cache == nullptr) { - return do_read_parquet(); - } - - FileCache::Key key{ - .filepaths = options.get_source().filepaths(), - .skip_rows = options.get_skip_rows(), - .skip_bytes = options.get_skip_bytes(), - .num_rows = options.get_num_rows(), - .num_bytes = options.get_num_bytes(), - .column_names = options.get_column_names(), - .column_indices = options.get_column_indices(), - .row_groups = options.get_row_groups() - }; - - auto msg = file_cache->get(ctx, key); - if (msg.has_value()) { - return std::move(*msg); - } - - auto ret = do_read_parquet(); - file_cache->insert(ctx, key, ret); - return ret; -} - -struct ChunkDesc { - std::uint64_t sequence_number; - std::int64_t skip_rows; - std::int64_t num_rows; - cudf::io::source_info source; -}; - -/** - * @brief Read chunks and send them to an output channel. - * - * @param ctx Execution context to use. - * @param ch_out Channel to send output to. - * @param options Template reader options. - * @param chunks List of chunks from the input files to read. Processed in order. - * @param idx Index of the next chunk to process. - * - * @return Coroutine representing the processing of all chunks. - */ -Actor produce_chunks( - std::shared_ptr ctx, - std::shared_ptr ch_out, - std::vector& chunks, - cudf::io::parquet_reader_options options -) { - // ShutdownAtExit c{ch_out}; - co_await ctx->executor()->schedule(); - for (auto& chunk : chunks) { - cudf::io::parquet_reader_options chunk_options{options}; - chunk_options.set_skip_rows(chunk.skip_rows); - chunk_options.set_num_rows(chunk.num_rows); - chunk_options.set_source(chunk.source); - auto stream = ctx->br()->stream_pool().get_stream(); - auto ticket = co_await ch_out->acquire(); - if (!ticket.has_value()) { - // Semaphore (and hence output channel) shutdown - break; - } - // Having acquire a ticket, let's move to a new thread. - co_await ctx->executor()->schedule(); - // TODO: This reads the metadata ntasks times. - // See https://github.com/rapidsai/cudf/issues/20311 - auto [msg, exception] = [&]() -> std::pair { - try { - return { - read_parquet_chunk(ctx, stream, chunk_options, chunk.sequence_number), - nullptr - }; - } catch (...) { - return {Message{}, std::current_exception()}; - } - }(); - if (exception != nullptr) { - co_await ch_out->shutdown(); - std::rethrow_exception(exception); - } - auto sent = co_await ticket->send(std::move(msg)); - if (!sent) { - // Output channel is shutdown, no need for more reads. - break; - } - } - co_await ch_out->drain(ctx->executor()); -} -} // namespace - -Actor read_parquet( - std::shared_ptr ctx, - std::shared_ptr comm, - std::shared_ptr ch_out, - std::size_t num_producers, - cudf::io::parquet_reader_options options, - cudf::size_type num_rows_per_chunk, - std::unique_ptr filter -) { - ShutdownAtExit c{ch_out}; - co_await ctx->executor()->schedule(); - auto const size = safe_cast(comm->nranks()); - auto const rank = safe_cast(comm->rank()); - auto source = options.get_source(); - RAPIDSMPF_EXPECTS( - source.type() == cudf::io::io_type::FILEPATH, "Only implemented for file sources" - ); - // TODO: To handle this we need a prefix scan across all the ranks of the total - // number of rows that would be read by previous ranks. - RAPIDSMPF_EXPECTS( - size == 1 || !options.get_num_rows().has_value(), - "Reading subset of rows not yet supported in multi-rank execution" - ); - // TODO: To handle this we need a prefix scan across all the ranks of the total - // number of rows that would be read by previous ranks. - RAPIDSMPF_EXPECTS( - size == 1 || options.get_skip_rows() == 0, - "Skipping rows not yet supported in multi-rank execution" - ); - auto files = source.filepaths(); - RAPIDSMPF_EXPECTS(files.size() > 0, "Must have at least one file to read"); - RAPIDSMPF_EXPECTS( - !options.get_filter().has_value(), - "Do not set filter on options, use the filter argument" - ); - if (filter != nullptr) { - options.set_filter(filter->filter); - // Let's just join all the possible streams here rather than inducing cross-stream - // deps in the tasks - cuda_stream_join( - std::ranges::transform_view( - std::ranges::iota_view( - std::size_t{0}, ctx->br()->stream_pool().get_pool_size() - ), - [&](auto i) { return ctx->br()->stream_pool().get_stream(i); } - ), - std::ranges::single_view(filter->stream) - ); - } - // TODO: Handle case where multiple ranks are reading from a single file. - auto const files_per_rank = - safe_cast(files.size() / size + (rank < (files.size() % size))); - auto const file_offset = safe_cast( - rank * (files.size() / size) + std::min(rank, files.size() % size) - ); - auto local_files = std::vector( - files.begin() + file_offset, files.begin() + file_offset + files_per_rank - ); - std::uint64_t sequence_number = 0; - std::vector> chunks_per_producer(num_producers); - auto const num_files = local_files.size(); - // Estimate number of rows per file - std::size_t files_per_chunk = 1; - if (num_files > 1) { - auto nrows = - cudf::io::read_parquet_metadata(cudf::io::source_info(local_files[0])) - .num_rows(); - files_per_chunk = - safe_cast(std::max(num_rows_per_chunk / nrows, 1l)); - } - auto to_skip = options.get_skip_rows(); - auto to_read = - options.get_num_rows().value_or(std::numeric_limits::max()); - for (std::size_t file_offset = 0; file_offset < num_files; - file_offset += files_per_chunk) - { - std::vector chunk_files; - auto const nchunk_files = std::min(num_files - file_offset, files_per_chunk); - std::ranges::copy_n( - local_files.begin() + safe_cast(file_offset), - safe_cast(nchunk_files), - std::back_inserter(chunk_files) - ); - auto source = cudf::io::source_info(chunk_files); - // Must read [skip_rows, skip_rows + num_rows) from full fileset - auto chunk_rows = cudf::io::read_parquet_metadata(source).num_rows() - to_skip; - auto chunk_skip_rows = to_skip; - // If the chunk is larger than the number rows we need to skip, on the next - // iteration we don't need to skip any more rows, otherwise we must skip the - // remainder. - to_skip = std::max(0l, -chunk_rows); - while (chunk_rows > 0 && to_read > 0) { - auto rows_read = std::min( - {safe_cast(num_rows_per_chunk), chunk_rows, to_read} - ); - chunks_per_producer[sequence_number % num_producers].emplace_back( - sequence_number, chunk_skip_rows, rows_read, source - ); - sequence_number++; - to_read = std::max(0l, to_read - rows_read); - chunk_skip_rows += rows_read; - chunk_rows -= rows_read; - } - } - if (std::ranges::all_of(chunks_per_producer, [](auto&& v) { return v.empty(); })) { - if (local_files.size() > 0) { - // If we're on the hook to read some files, but the skip_rows/num_rows setup - // meant our slice was empty, send an empty table of correct shape. - // Anyone with no files will just immediately close their output channel. - auto empty_opts = options; - empty_opts.set_source(cudf::io::source_info(local_files[0])); - empty_opts.set_skip_rows(0); - empty_opts.set_num_rows(0); - co_await ctx->executor()->schedule(ch_out->send(read_parquet_chunk( - ctx, ctx->br()->stream_pool().get_stream(), std::move(empty_opts), 0 - ))); - } - } else { - std::vector read_tasks; - read_tasks.reserve(1 + num_producers); - auto lineariser = Lineariser(ctx, ch_out, num_producers); - auto queues = lineariser.get_queues(); - for (std::size_t i = 0; i < num_producers; i++) { - read_tasks.push_back( - produce_chunks(ctx, queues[i], chunks_per_producer[i], options) - ); - } - read_tasks.push_back(lineariser.drain()); - coro_results(co_await coro::when_all(std::move(read_tasks))); - } - co_await ch_out->drain(ctx->executor()); - if (filter != nullptr) { - // Let's just join all the possible streams here rather than inducing cross-stream - // deps in the tasks - cuda_stream_join( - std::ranges::single_view(filter->stream), - std::ranges::transform_view( - std::ranges::iota_view( - std::size_t{0}, ctx->br()->stream_pool().get_pool_size() - ), - [&](auto i) { return ctx->br()->stream_pool().get_stream(i); } - ) - ); - } -} -} // namespace rapidsmpf::streaming::actor diff --git a/cpp/src/streaming/cudf/partition.cpp b/cpp/src/streaming/cudf/partition.cpp deleted file mode 100644 index 33268739a..000000000 --- a/cpp/src/streaming/cudf/partition.cpp +++ /dev/null @@ -1,101 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ -#include - -#include - -#include -#include -#include -#include - -namespace rapidsmpf::streaming::actor { - - -Actor partition_and_pack( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector columns_to_hash, - int num_partitions, - cudf::hash_id hash_function, - std::uint32_t seed -) { - 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 table = msg.release(); - auto reservation = ctx->br()->reserve_device_memory_and_spill( - table.make_available_cost(), AllowOverbooking::NO - ); - auto tbl = table.make_available(reservation); - - PartitionMapChunk partition_map{ - .data = rapidsmpf::partition_and_pack( - tbl.table_view(), - columns_to_hash, - num_partitions, - hash_function, - seed, - tbl.stream(), - ctx->br().get() - ) - }; - - co_await ch_out->send(to_message( - msg.sequence_number(), - std::make_unique(std::move(partition_map)) - )); - } - co_await ch_out->drain(ctx->executor()); -} - -Actor unpack_and_concat( - std::shared_ptr ctx, - std::shared_ptr ch_in, - std::shared_ptr ch_out -) { - ShutdownAtExit c{ch_in, ch_out}; - co_await ctx->executor()->schedule(); - while (true) { - auto msg = co_await ch_in->receive(); - if (msg.empty()) { - break; - } - - // If receiving a partition map, we convert it to a vector and discard - // partition IDs. - std::uint64_t seq = msg.sequence_number(); - std::vector data; - if (msg.holds()) { - auto partition_map = msg.release(); - data = to_vector(std::move(partition_map.data)); - } else { - auto partition_vec = msg.release(); - data = std::move(partition_vec.data); - } - // Get a stream for the concatenated table chunk. - auto stream = ctx->br()->stream_pool().get_stream(); - - std::unique_ptr ret = rapidsmpf::unpack_and_concat( - rapidsmpf::unspill_partitions( - std::move(data), ctx->br().get(), AllowOverbooking::NO - ), - stream, - ctx->br().get() - ); - co_await ch_out->send( - to_message(seq, std::make_unique(std::move(ret), stream)) - ); - } - co_await ch_out->drain(ctx->executor()); -} - -} // namespace rapidsmpf::streaming::actor diff --git a/cpp/src/streaming/cudf/table_chunk.cpp b/cpp/src/streaming/cudf/table_chunk.cpp deleted file mode 100644 index afef6385f..000000000 --- a/cpp/src/streaming/cudf/table_chunk.cpp +++ /dev/null @@ -1,314 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include -#include - -#include -#include -#include -#include - -namespace rapidsmpf::streaming { - -TableChunk::TableChunk(std::unique_ptr table, rmm::cuda_stream_view stream) - : table_{std::move(table)}, stream_{stream}, is_spillable_{true} { - RAPIDSMPF_EXPECTS( - table_ != nullptr, "table pointer cannot be null", std::invalid_argument - ); - table_view_ = table_->view(); - data_alloc_size_[static_cast(MemoryType::DEVICE)] = table_->alloc_size(); - make_available_cost_ = 0; -} - -TableChunk::TableChunk( - cudf::table_view table_view, - rmm::cuda_stream_view stream, - OwningWrapper&& owner, - ExclusiveView exclusive_view -) - : owner_{std::move(owner)}, - table_view_{table_view}, - stream_{stream}, - is_spillable_{static_cast(exclusive_view)} { - data_alloc_size_[static_cast(MemoryType::DEVICE)] = cudf::packed_size( - table_view, stream_, rmm::mr::get_current_device_resource_ref() - ); - make_available_cost_ = 0; -} - -TableChunk::TableChunk(std::unique_ptr packed_data) - : packed_data_{std::move(packed_data)}, - stream_{packed_data_->data->stream()}, - is_spillable_{true} { - RAPIDSMPF_EXPECTS( - packed_data_ != nullptr, - "packed data pointer cannot be null", - std::invalid_argument - ); - RAPIDSMPF_EXPECTS( - !packed_data_->empty(), "packed data cannot be empty", std::invalid_argument - ); - data_alloc_size_[static_cast(packed_data_->data->mem_type())] = - packed_data_->data->size; - if (packed_data_->data->mem_type() != MemoryType::DEVICE) { - make_available_cost_ = packed_data_->data->size; - } else { - // table data is in device memory. We can trivially unpack it and make it - // available. - table_view_ = cudf::unpack( - packed_data_->metadata->data(), - reinterpret_cast(packed_data_->data->data()) - ); - make_available_cost_ = 0; - } -} - -TableChunk::TableChunk(TableChunk&& other) noexcept - : owner_(std::move(other.owner_)), - table_(std::move(other.table_)), - packed_data_(std::move(other.packed_data_)), - table_view_(std::exchange(other.table_view_, std::nullopt)), - data_alloc_size_(other.data_alloc_size_), - make_available_cost_(other.make_available_cost_), - stream_(other.stream_), - is_spillable_(other.is_spillable_) {} - -TableChunk& TableChunk::operator=(TableChunk&& other) noexcept { - if (this != &other) { - owner_ = std::move(other.owner_); - table_ = std::move(other.table_); - packed_data_ = std::move(other.packed_data_); - table_view_ = std::exchange(other.table_view_, std::nullopt); - data_alloc_size_ = other.data_alloc_size_; - make_available_cost_ = other.make_available_cost_; - stream_ = other.stream_; - is_spillable_ = other.is_spillable_; - } - return *this; -} - -rmm::cuda_stream_view TableChunk::stream() const noexcept { - return stream_; -} - -std::size_t TableChunk::data_alloc_size(MemoryType mem_type) const { - return data_alloc_size_.at(static_cast(mem_type)); -} - -bool TableChunk::is_available() const noexcept { - return table_view_.has_value(); -} - -std::size_t TableChunk::make_available_cost() const noexcept { - return make_available_cost_; -} - -TableChunk TableChunk::make_available(MemoryReservation& reservation) { - if (is_available()) { - return std::move(*this); - } - // Table chunk is not available. This means that the table data is not in device - // memory. We need to move the table data to device memory using a device reservation. - RAPIDSMPF_EXPECTS( - reservation.mem_type() == MemoryType::DEVICE, - "device memory reservation is required" - ); - RAPIDSMPF_EXPECTS(packed_data_ != nullptr, "packed data pointer cannot be null"); - auto packed_data = std::move(packed_data_); - packed_data->data = reservation.br()->move(std::move(packed_data->data), reservation); - return TableChunk{std::move(packed_data)}; -} - -TableChunk TableChunk::make_available(MemoryReservation&& reservation) { - MemoryReservation& res = reservation; - return make_available(res); -} - -coro::task TableChunk::make_available( - std::shared_ptr ctx, std::int64_t net_memory_delta -) { - co_return make_available( - co_await reserve_memory(ctx, make_available_cost(), net_memory_delta) - ); -} - -cudf::table_view TableChunk::table_view() const { - RAPIDSMPF_EXPECTS( - is_available(), - "the table view is unavailable, please make sure it is " - "unspilled and unpacked (see `make_available`).", - std::invalid_argument - ); - return table_view_.value(); -} - -bool TableChunk::is_spillable() const { - return is_spillable_; -} - -TableChunk TableChunk::copy(MemoryReservation& reservation) const { - // This method handles the two possible cases. Note that - // `!is_available() && packed_data_ == nullptr` is an invalid state, so the - // remaining valid combinations collapse into: - // - // 1. The chunk is available and not yet packed. The table is copied/packed - // into the reservation-specified memory type using libcudf: - // a. DEVICE - cudf-copy table_view() into device memory. - // b. PINNED_HOST - cudf::pack table_view() directly into pinned memory. - // c. HOST - cudf::pack table_view() into intermediate device - // memory and then copy to host memory. - // - // 2. The chunk data is already packed (packed_data_ != nullptr). - // Use buffer_copy() to copy the packed data into the reservation- - // specified memory type. The original memory type of the chunk does - // not matter. - BufferResource* br = reservation.br(); - - // If the table view is available and the table is not packed, we can use libcudf to - // copy the table in device memory, or pack it to pinned/ host memory. Else, fall - // through to case 2 (ie. use buffer_copy). - if (is_available() && packed_data_ == nullptr) { - switch (reservation.mem_type()) { - case MemoryType::DEVICE: // Case 1a. - { - // Use libcudf to copy the table_view(). - auto const nbytes = data_alloc_size(MemoryType::DEVICE); - auto statistics = br->statistics(); - StreamOrderedTiming timing{stream(), statistics}; - auto table = std::make_unique( - table_view(), stream(), br->device_mr() - ); - statistics->record_copy( - MemoryType::DEVICE, MemoryType::DEVICE, nbytes, std::move(timing) - ); - // And update the provided `reservation`. - br->release(reservation, nbytes); - return TableChunk(std::move(table), stream()); - } - case MemoryType::PINNED_HOST: // Case 1b. - { - StreamOrderedTiming timing{stream(), br->statistics()}; - - // use cudf pack with pinned mr - auto packed_pinned = cudf::pack(table_view(), stream(), br->pinned_mr()); - auto nbytes = packed_pinned.gpu_data->size(); - - br->statistics()->record_copy( - MemoryType::DEVICE, MemoryType::PINNED_HOST, nbytes, std::move(timing) - ); - // update the provided `reservation` - br->release(reservation, nbytes); - auto host_buffer = br->move(std::move(packed_pinned.gpu_data), stream()); - return TableChunk( - std::make_unique( - std::move(packed_pinned.metadata), std::move(host_buffer) - ) - ); - } - case MemoryType::HOST: // Case 1c. - { - // We use libcudf's pack() to serialize `table_view()` into a - // packed_columns and then we move the packed_columns' gpu_data to a - // new host buffer. - // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently, - // `cudf::pack()` allocates device memory we haven't reserved. - auto packed_columns = cudf::pack(table_view(), stream(), br->device_mr()); - auto packed_data = std::make_unique( - std::move(packed_columns.metadata), - br->move(std::move(packed_columns.gpu_data), stream()) - ); - - // Handle the case where `cudf::pack` allocates slightly more than the - // input size. This can occur because cudf uses aligned allocations, - // which may exceed the requested size. To accommodate this, we - // allow some wiggle room. - if (packed_data->data->size > reservation.size()) { - auto const wiggle_room = - 1024 * static_cast(table_view().num_columns()); - if (packed_data->data->size <= reservation.size() + wiggle_room) { - reservation = br->reserve( - reservation.mem_type(), - packed_data->data->size, - AllowOverbooking::YES - ) - .first; - } - } - packed_data->data = br->move(std::move(packed_data->data), reservation); - return TableChunk(std::move(packed_data)); - } - default: - RAPIDSMPF_FAIL("MemoryType: unknown"); - } - } - // `!is_available() && packed_data_ == nullptr` is an invalid state, so - // reaching this point implies `packed_data_ != nullptr`. - RAPIDSMPF_EXPECTS(packed_data_ != nullptr, "something went wrong"); - - // Case 2. The chunk data is already packed (packed_data_ != nullptr). We need - // to copy the packed data into the reservation-specified memory type. - auto const nbytes = packed_data_->data->size; - auto metadata = std::make_unique>(*packed_data_->metadata); - auto data = br->make_buffer(nbytes, packed_data_->stream(), reservation); - buffer_copy(br->statistics(), *data, *packed_data_->data, nbytes); - return TableChunk(std::make_unique(std::move(metadata), std::move(data))); -} - -std::unique_ptr TableChunk::into_packed_data(BufferResource* br) && { - if (packed_data_) { - table_view_ = std::nullopt; - return std::move(packed_data_); - } - RAPIDSMPF_EXPECTS( - is_available(), "TableChunk must be available; call make_available() first" - ); - // TODO: use `cudf::chunked_pack()` with a bounce buffer. Currently, - // `cudf::pack()` allocates device memory we haven't reserved. - auto packed_columns = cudf::pack(table_view_.value(), stream_, br->device_mr()); - table_view_ = std::nullopt; - return std::make_unique( - std::move(packed_columns.metadata), - br->move(std::move(packed_columns.gpu_data), stream_) - ); -} - -std::pair TableChunk::shape() const noexcept { - if (packed_data_ != nullptr) { - auto view = cudf::packed_metadata_view(*packed_data_->metadata); - return {view.num_rows(), view.num_columns()}; - } - return {table_view_->num_rows(), table_view_->num_columns()}; -} - -ContentDescription get_content_description(TableChunk const& obj) { - ContentDescription ret{ - obj.is_spillable() ? ContentDescription::Spillable::YES - : ContentDescription::Spillable::NO - }; - for (auto mem_type : MEMORY_TYPES) { - ret.content_size(mem_type) = obj.data_alloc_size(mem_type); - } - return ret; -} - -Message to_message(std::uint64_t sequence_number, std::unique_ptr chunk) { - auto cd = get_content_description(*chunk); - return Message{ - sequence_number, - std::move(chunk), - cd, - [](Message const& msg, MemoryReservation& reservation) -> Message { - auto const& self = msg.get(); - auto chunk = std::make_unique(self.copy(reservation)); - auto cd = get_content_description(*chunk); - return Message{msg.sequence_number(), std::move(chunk), cd, msg.copy_cb()}; - } - }; -} - -} // namespace rapidsmpf::streaming diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index c439df259..dd3f4a8de 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -42,9 +42,9 @@ endfunction(rapidsmpf_mpirun_test_add) # ${RAPIDSMPF_BINARY_DIR}/gtests) is created to seamlessly run ctest from ${RAPIDSMPF_BINARY_DIR}. file(CREATE_LINK "${RAPIDSMPF_BINARY_DIR}/gtests" "${CMAKE_CURRENT_BINARY_DIR}/gtests" SYMBOLIC) -# Use an object library for test sources to avoid recompiling (especially cudf_tests) them for each -# executable. Unlike STATIC libraries, OBJECT libraries always include all their object files when -# linked, which is required for GTest tests that register via global constructors. +# Use an object library for test sources to avoid recompiling them for each executable. Unlike +# STATIC libraries, OBJECT libraries always include all their object files when linked, which is +# required for GTest tests that register via global constructors. add_library(test_sources OBJECT) set_target_properties( test_sources @@ -62,8 +62,7 @@ target_compile_options( ) target_link_libraries( test_sources - PRIVATE rapidsmpf::rapidsmpf cudf::cudftestutil cudf::cudftestutil_impl - $<$:numa> maybe_asan + PRIVATE rapidsmpf::rapidsmpf $<$:numa> maybe_asan PUBLIC GTest::gmock GTest::gtest ) disable_sign_conversion_warning(test_sources) @@ -78,20 +77,16 @@ target_sources( test_communicator.cpp test_config.cpp test_cuda_stream.cpp - test_cudf_utils.cpp test_cupti_monitor.cpp test_error_macros.cpp test_host_buffer.cpp test_memory_resources.cpp test_metadata_payload_exchange.cpp test_misc.cpp - test_partition.cpp test_pausable_thread_loop.cpp test_progress_thread.cpp test_rmm_resource_adaptor.cpp test_sparse_alltoall.cpp - test_shuffler_many_streams.cpp - test_shuffler.cpp test_spill_manager.cpp test_statistics.cpp test_stream_ordered_timing.cpp @@ -103,25 +98,36 @@ target_sources( if(RAPIDSMPF_HAVE_STREAMING) target_sources( test_sources - PRIVATE streaming/test_allgather.cpp - streaming/test_allreduce.cpp + PRIVATE streaming/test_allreduce.cpp streaming/test_channel.cpp streaming/test_error_handling.cpp streaming/test_fanout.cpp - streaming/test_leaf_actor.cpp streaming/test_lineariser.cpp streaming/test_memory_reserve_or_wait.cpp streaming/test_message.cpp - streaming/test_partition.cpp - streaming/test_channel_metadata.cpp - streaming/test_read_parquet.cpp - streaming/test_shuffler.cpp streaming/test_sparse_alltoall.cpp streaming/test_spillable_messages.cpp - streaming/test_table_chunk.cpp ) endif() +# cudf-dependent test sources, gated behind BUILD_CUDF_TESTS +if(BUILD_CUDF_TESTS) + target_sources( + test_sources PRIVATE test_partition.cpp test_shuffler.cpp test_shuffler_many_streams.cpp + ) + target_link_libraries( + test_sources PRIVATE cudf_streaming::cudf_streaming cudf::cudftestutil cudf::cudftestutil_impl + ) + target_compile_definitions(test_sources PRIVATE RAPIDSMPF_HAVE_CUDF) + + if(RAPIDSMPF_HAVE_STREAMING) + target_sources( + test_sources PRIVATE streaming/test_allgather.cpp streaming/test_leaf_actor.cpp + streaming/test_shuffler.cpp + ) + endif() +endif() + if(RAPIDSMPF_HAVE_MPI) add_executable(mpi_tests main/mpi.cpp) set_target_properties( diff --git a/cpp/tests/streaming/base_streaming_fixture.hpp b/cpp/tests/streaming/base_streaming_fixture.hpp index 23cbf2b6e..311423aa6 100644 --- a/cpp/tests/streaming/base_streaming_fixture.hpp +++ b/cpp/tests/streaming/base_streaming_fixture.hpp @@ -7,7 +7,7 @@ #include #include -#include +#include #include #include @@ -42,7 +42,7 @@ class BaseStreamingFixture : public ::testing::Test { env_vars["num_streaming_threads"] = std::to_string(num_streaming_threads); rapidsmpf::config::Options options(std::move(env_vars)); - stream = cudf::get_default_stream(); + stream = rmm::cuda_stream_view{}; br = rapidsmpf::BufferResource::create( mr_cuda, rapidsmpf::PinnedMemoryResource::Disabled, std::move(memory_limits) ); diff --git a/cpp/tests/streaming/test_allgather.cpp b/cpp/tests/streaming/test_allgather.cpp index 8bf97f02f..6805f3100 100644 --- a/cpp/tests/streaming/test_allgather.cpp +++ b/cpp/tests/streaming/test_allgather.cpp @@ -10,11 +10,12 @@ #include #include +#include + #include #include #include -#include #include #include #include diff --git a/cpp/tests/streaming/test_channel_metadata.cpp b/cpp/tests/streaming/test_channel_metadata.cpp deleted file mode 100644 index 6c97ce3c5..000000000 --- a/cpp/tests/streaming/test_channel_metadata.cpp +++ /dev/null @@ -1,239 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights - * reserved. SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include - -#include - -#include -#include -#include -#include -#include - -#include -#include -#include - -using namespace rapidsmpf::streaming; - -class StreamingChannelMetadata : public ::testing::Test {}; - -TEST_F(StreamingChannelMetadata, HashScheme) { - HashScheme h{{0, 1}, 16}; - EXPECT_EQ(h.column_indices.size(), 2); - EXPECT_EQ(h.column_indices[0], 0); - EXPECT_EQ(h.column_indices[1], 1); - EXPECT_EQ(h.modulus, 16); - - // Equality - EXPECT_EQ(h, (HashScheme{{0, 1}, 16})); - EXPECT_NE(h, (HashScheme{{0, 1}, 32})); - EXPECT_NE(h, (HashScheme{{2}, 16})); -} - -TEST_F(StreamingChannelMetadata, OrderSchemeCtorRejectsEmptyKeys) { - EXPECT_THROW(static_cast(OrderScheme({}, nullptr)), std::invalid_argument); -} - -TEST_F(StreamingChannelMetadata, OrderSchemeCtorRejectsNullBoundaries) { - EXPECT_THROW( - static_cast( - OrderScheme({{0, cudf::order::ASCENDING, cudf::null_order::BEFORE}}, nullptr) - ), - std::invalid_argument - ); -} - -TEST_F(StreamingChannelMetadata, PartitioningSpec) { - // None - auto spec_none = PartitioningSpec::none(); - EXPECT_EQ(spec_none.type, PartitioningSpec::Type::NONE); - - // Inherit - auto spec_inherit = PartitioningSpec::inherit(); - EXPECT_EQ(spec_inherit.type, PartitioningSpec::Type::INHERIT); - - // Hash - auto spec_hash = PartitioningSpec::from_hash(HashScheme{{0}, 16}); - EXPECT_EQ(spec_hash.type, PartitioningSpec::Type::HASH); - EXPECT_EQ(spec_hash.hash->column_indices[0], 0); - EXPECT_EQ(spec_hash.hash->modulus, 16); - - // Type checks (operator== removed; use field comparisons) - EXPECT_EQ(spec_none.type, PartitioningSpec::Type::NONE); - EXPECT_EQ(spec_inherit.type, PartitioningSpec::Type::INHERIT); - EXPECT_EQ(spec_hash.type, PartitioningSpec::Type::HASH); - EXPECT_NE(spec_none.type, spec_inherit.type); - EXPECT_EQ(spec_hash.hash->modulus, 16); - EXPECT_NE((PartitioningSpec::from_hash(HashScheme{{0}, 32}).hash->modulus), 16); -} - -TEST_F(StreamingChannelMetadata, PartitioningScenarios) { - // Default construction - Partitioning p_default{}; - EXPECT_EQ(p_default.inter_rank.type, PartitioningSpec::Type::NONE); - EXPECT_EQ(p_default.local.type, PartitioningSpec::Type::NONE); - - // Direct global shuffle: inter_rank=Hash, local=Inherit - Partitioning p_global{ - PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit() - }; - EXPECT_EQ(p_global.inter_rank.type, PartitioningSpec::Type::HASH); - EXPECT_EQ(p_global.local.type, PartitioningSpec::Type::INHERIT); - EXPECT_EQ(p_global.inter_rank.hash->modulus, 16); - - // Two-stage shuffle: inter_rank=Hash(nranks), local=Hash(N_l) - Partitioning p_twostage{ - PartitioningSpec::from_hash(HashScheme{{0}, 4}), - PartitioningSpec::from_hash(HashScheme{{0}, 8}) - }; - EXPECT_EQ(p_twostage.inter_rank.hash->modulus, 4); - EXPECT_EQ(p_twostage.local.hash->modulus, 8); - - // Field comparisons (Partitioning::operator== removed) - { - Partitioning p_same{ - PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit() - }; - EXPECT_EQ(p_global.inter_rank.type, p_same.inter_rank.type); - EXPECT_EQ(p_global.inter_rank.hash->modulus, p_same.inter_rank.hash->modulus); - EXPECT_EQ(p_global.local.type, p_same.local.type); - } - EXPECT_NE(p_global.inter_rank.hash->modulus, p_twostage.inter_rank.hash->modulus); -} - -TEST_F(StreamingChannelMetadata, ChannelMetadata) { - // Full construction - use std::move to avoid GCC false positive on vector copy - Partitioning p{ - PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit() - }; - ChannelMetadata m{4, std::move(p), true}; - EXPECT_EQ(m.local_count, 4); - EXPECT_EQ(m.partitioning.inter_rank.type, PartitioningSpec::Type::HASH); - EXPECT_EQ(m.partitioning.local.type, PartitioningSpec::Type::INHERIT); - EXPECT_TRUE(m.duplicated); - - // Minimal construction - ChannelMetadata m_minimal{4}; - EXPECT_EQ(m_minimal.local_count, 4); - EXPECT_FALSE(m_minimal.duplicated); - - // Equality - create fresh partitionings and move them - ChannelMetadata m_same{ - 4, - Partitioning{ - PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit() - }, - true - }; - ChannelMetadata m_diff{ - 8, - Partitioning{ - PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit() - }, - true - }; - // Field comparisons (ChannelMetadata::operator== removed) - EXPECT_EQ(m.local_count, m_same.local_count); - EXPECT_EQ(m.duplicated, m_same.duplicated); - EXPECT_EQ( - m.partitioning.inter_rank.hash->modulus, - m_same.partitioning.inter_rank.hash->modulus - ); - EXPECT_NE(m.local_count, m_diff.local_count); -} - -TEST_F(StreamingChannelMetadata, MessageRoundTrip) { - // ChannelMetadata round-trip - Partitioning part{ - PartitioningSpec::from_hash(HashScheme{{0}, 16}), PartitioningSpec::inherit() - }; - auto m = std::make_unique(4, std::move(part), false); - auto msg_m = to_message(99, std::move(m)); - EXPECT_EQ(msg_m.sequence_number(), 99); - EXPECT_TRUE(msg_m.holds()); - auto released = msg_m.release(); - EXPECT_EQ(released.local_count, 4); - EXPECT_FALSE(released.duplicated); - EXPECT_EQ(released.partitioning.inter_rank.hash->modulus, 16); - EXPECT_TRUE(msg_m.empty()); -} - -class StreamingChannelMetadataGPU : public ::testing::Test { - protected: - rmm::cuda_stream_view stream{cudf::get_default_stream()}; - std::shared_ptr br = - rapidsmpf::BufferResource::create(cudf::get_current_device_resource_ref()); - - std::shared_ptr make_chunk(std::vector vals) { - rmm::device_buffer buf(vals.data(), vals.size() * sizeof(int32_t), stream); - auto col = std::make_unique( - cudf::data_type{cudf::type_id::INT32}, - static_cast(vals.size()), - std::move(buf), - rmm::device_buffer{}, - 0 - ); - std::vector> cols; - cols.push_back(std::move(col)); - return std::make_shared( - std::make_unique(std::move(cols)), stream - ); - } -}; - -TEST_F(StreamingChannelMetadataGPU, OrderSchemeReplaceKeys) { - OrderKey k0{0, cudf::order::ASCENDING, cudf::null_order::BEFORE}; - OrderKey k5{5, cudf::order::DESCENDING, cudf::null_order::AFTER}; - - auto b = make_chunk({100, 200}); - OrderScheme o1({k0}, b); - auto o2 = o1.with_keys({k5}); - - EXPECT_EQ(o2.keys[0].column_index, 5); - EXPECT_EQ(o2.keys[0].order, cudf::order::DESCENDING); - EXPECT_EQ(o2.strict_boundaries, o1.strict_boundaries); - EXPECT_EQ(o2.boundaries->shape(), o1.boundaries->shape()); - EXPECT_EQ(o2.boundaries.get(), b.get()); - EXPECT_NE(o1.keys[0].column_index, o2.keys[0].column_index); - - EXPECT_THROW(static_cast(o1.with_keys({k0, k5})), std::invalid_argument); -} - -TEST_F(StreamingChannelMetadataGPU, OrderSchemeBoundariesAlignedWith) { - OrderKey k0{0, cudf::order::ASCENDING, cudf::null_order::BEFORE}; - OrderKey k3{3, cudf::order::ASCENDING, cudf::null_order::BEFORE}; - - OrderScheme o1({k0}, make_chunk({100, 200})); - OrderScheme o2({k0}, make_chunk({100, 200})); - EXPECT_TRUE(o1.boundaries_aligned_with(o2, *br)); - - OrderScheme o_shifted({k3}, make_chunk({100, 200})); - EXPECT_TRUE(o1.boundaries_aligned_with(o_shifted, *br)); - - OrderScheme o_strict({k0}, make_chunk({100, 200}), /*strict=*/true); - EXPECT_FALSE(o1.boundaries_aligned_with(o_strict, *br)); - - OrderScheme o_diff({k0}, make_chunk({100, 300})); - EXPECT_FALSE(o1.boundaries_aligned_with(o_diff, *br)); -} - -TEST_F(StreamingChannelMetadataGPU, PartitioningSpecOrder) { - OrderKey k0{0, cudf::order::ASCENDING, cudf::null_order::BEFORE}; - OrderScheme o({k0}, make_chunk({100, 200})); - - auto spec = PartitioningSpec::from_order(o); - EXPECT_EQ(spec.type, PartitioningSpec::Type::ORDER); - EXPECT_TRUE(spec.order.has_value()); - EXPECT_EQ(spec.order->keys[0].column_index, 0); - - // Type checks only (PartitioningSpec::operator== removed; ORDER value comparison - // requires boundaries_aligned_with on the OrderScheme directly) - EXPECT_EQ(spec.type, PartitioningSpec::Type::ORDER); - EXPECT_NE(spec.type, PartitioningSpec::from_hash(HashScheme{{0}, 16}).type); - EXPECT_NE(spec.type, PartitioningSpec::none().type); -} diff --git a/cpp/tests/streaming/test_fanout.cpp b/cpp/tests/streaming/test_fanout.cpp index 3653f643e..d0394a9d7 100644 --- a/cpp/tests/streaming/test_fanout.cpp +++ b/cpp/tests/streaming/test_fanout.cpp @@ -9,8 +9,6 @@ #include #include -#include - #include #include diff --git a/cpp/tests/streaming/test_leaf_actor.cpp b/cpp/tests/streaming/test_leaf_actor.cpp index 37ec6b088..cd5e90290 100644 --- a/cpp/tests/streaming/test_leaf_actor.cpp +++ b/cpp/tests/streaming/test_leaf_actor.cpp @@ -10,6 +10,7 @@ #include #include +#include #include #include @@ -21,7 +22,6 @@ #include #include #include -#include #include "../utils.hpp" #include "base_streaming_fixture.hpp" @@ -48,15 +48,17 @@ TEST_F(StreamingLeafTasks, PushAndPullChunks) { { 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 + inputs.emplace_back( + cudf_streaming::streaming::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))); @@ -71,7 +73,8 @@ TEST_F(StreamingLeafTasks, PushAndPullChunks) { 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() + outputs[i].get().table_view(), + expects[i].view() ); } } diff --git a/cpp/tests/streaming/test_message.cpp b/cpp/tests/streaming/test_message.cpp index ba2a9454f..73a223f77 100644 --- a/cpp/tests/streaming/test_message.cpp +++ b/cpp/tests/streaming/test_message.cpp @@ -6,8 +6,8 @@ #include #include -#include -#include +#include +#include #include #include @@ -18,8 +18,8 @@ using namespace rapidsmpf::streaming; class StreamingMessage : public ::testing::Test { protected: void SetUp() override { - br = BufferResource::create(cudf::get_current_device_resource_ref()); - stream = cudf::get_default_stream(); + br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); + stream = rmm::cuda_stream_view{}; } std::shared_ptr br; diff --git a/cpp/tests/streaming/test_partition.cpp b/cpp/tests/streaming/test_partition.cpp deleted file mode 100644 index b43dc62d2..000000000 --- a/cpp/tests/streaming/test_partition.cpp +++ /dev/null @@ -1,147 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "../utils.hpp" -#include "base_streaming_fixture.hpp" - -using namespace rapidsmpf; -using namespace rapidsmpf::streaming; -namespace actor = rapidsmpf::streaming::actor; - -using StreamingPartition = BaseStreamingFixture; - -TEST_F(StreamingPartition, PackUnpackRoundTrip) { - int const num_partitions = 5; - int const num_rows = 100; - int const num_chunks = 10; - std::int64_t const seed = 42; - constexpr cudf::hash_id hash_function = cudf::hash_id::HASH_MURMUR3; - - std::vector expects; - for (int i = 0; i < num_chunks; ++i) { - expects.push_back(random_table_with_index(seed + i, num_rows, 0, 10)); - } - - 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 - ) - )); - } - - // Create and run the streaming pipeline. - std::vector outputs; - { - std::vector actors; - auto ch1 = ctx->create_channel(); - actors.push_back(actor::push_to_channel(ctx, ch1, std::move(inputs))); - - auto ch2 = ctx->create_channel(); - actors.push_back( - actor::partition_and_pack( - ctx, ch1, ch2, {1}, num_partitions, hash_function, seed - ) - ); - - auto ch3 = ctx->create_channel(); - actors.push_back(actor::unpack_and_concat(ctx, ch2, ch3)); - - actors.push_back(actor::pull_from_channel(ctx, ch3, 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); - auto output = outputs[i].release(); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT( - sort_table(output.table_view()), sort_table(expects[i].view()) - ); - } -} - -TEST_F(StreamingPartition, PartitionMapChunkToMessage) { - constexpr std::uint64_t seq = 42; - std::unordered_map data; - data.emplace(0, generate_packed_data(10, 0, stream, *br)); - data.emplace(1, generate_packed_data(10, 10, stream, *br)); - auto chunk = std::make_unique(std::move(data)); - - Message m = to_message(seq, std::move(chunk)); - EXPECT_FALSE(m.empty()); - EXPECT_TRUE(m.holds()); - EXPECT_TRUE(m.content_description().spillable()); - EXPECT_EQ(m.content_description().content_size(MemoryType::HOST), 0); - EXPECT_EQ(m.content_description().content_size(MemoryType::DEVICE), 80); - EXPECT_EQ(m.sequence_number(), seq); - - auto res = br->reserve_or_fail(m.copy_cost(), MemoryType::DEVICE); - Message m2 = m.copy(res); - EXPECT_EQ(res.size(), 0); - EXPECT_FALSE(m2.empty()); - EXPECT_TRUE(m2.holds()); - EXPECT_TRUE(m2.content_description().spillable()); - EXPECT_EQ(m2.content_description().content_size(MemoryType::HOST), 0); - EXPECT_EQ(m2.content_description().content_size(MemoryType::DEVICE), 80); - - auto chunk2 = m2.release(); - validate_packed_data(std::move(chunk2.data.at(0)), 10, 0, stream, *br); - validate_packed_data(std::move(chunk2.data.at(1)), 10, 10, stream, *br); -} - -TEST_F(StreamingPartition, PartitionMapChunkContentDescription) { - // Create a packed data, one in device and one in host memory. - std::unordered_map data; - auto pack1 = generate_packed_data(5, 0, stream, *br); - auto pack1_size = pack1.data->size; - auto pack2_size = pack1.data->size * 2; - auto res = br->reserve_or_fail(pack2_size, MemoryType::HOST); - auto pack2 = generate_packed_data(10, 0, stream, *br).copy(res); - data.emplace(0, std::move(pack1)); - data.emplace(1, std::move(pack2)); - - auto chunk = std::make_unique(std::move(data)); - auto cd = get_content_description(*chunk); - EXPECT_TRUE(cd.spillable()); - EXPECT_EQ(cd.content_size(MemoryType::DEVICE), pack1_size); - EXPECT_EQ(cd.content_size(MemoryType::HOST), pack2_size); -} - -TEST_F(StreamingPartition, PartitionVectorChunkContentDescription) { - // Create a packed data, one in device and one in host memory. - std::vector data; - auto pack1 = generate_packed_data(5, 0, stream, *br); - auto pack1_size = pack1.data->size; - auto pack2_size = pack1.data->size * 2; - auto res = br->reserve_or_fail(pack2_size, MemoryType::HOST); - auto pack2 = generate_packed_data(10, 0, stream, *br).copy(res); - data.push_back(std::move(pack1)); - data.push_back(std::move(pack2)); - - auto chunk = std::make_unique(std::move(data)); - auto cd = get_content_description(*chunk); - EXPECT_TRUE(cd.spillable()); - EXPECT_EQ(cd.content_size(MemoryType::DEVICE), pack1_size); - EXPECT_EQ(cd.content_size(MemoryType::HOST), pack2_size); -} diff --git a/cpp/tests/streaming/test_read_parquet.cpp b/cpp/tests/streaming/test_read_parquet.cpp deleted file mode 100644 index 1586bb7d1..000000000 --- a/cpp/tests/streaming/test_read_parquet.cpp +++ /dev/null @@ -1,284 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -#include "base_streaming_fixture.hpp" - - -using namespace rapidsmpf; -using namespace rapidsmpf::streaming; - -class StreamingReadParquet : public BaseStreamingFixture { - protected: - void SetUp() override { - BaseStreamingFixture::SetUp(); - constexpr int nfiles = 10; - constexpr int nrows = 10; - - temp_dir = std::filesystem::temp_directory_path() / "rapidsmpf_read_parquet_test"; - - for (int i = 0; i < nfiles; ++i) { - std::ostringstream filename_stream; - filename_stream << std::setw(3) << std::setfill('0') << i << ".pq"; - std::filesystem::path filepath = temp_dir / filename_stream.str(); - source_files.push_back(filepath.string()); - } - - if (GlobalEnvironment->comm_->rank() == 0) { - std::filesystem::create_directories(temp_dir); - - int start = 0; - for (auto& file : source_files) { - auto values = std::ranges::iota_view(start, start + nrows); - cudf::test::fixed_width_column_wrapper col( - values.begin(), values.end() - ); - - std::vector> columns; - columns.push_back(col.release()); - auto table = std::make_unique(std::move(columns)); - - cudf::io::sink_info sink{file}; - auto options = - cudf::io::parquet_writer_options::builder(sink, table->view()) - .build(); - cudf::io::write_parquet(options); - start += nrows + nrows / 2; - } - } - - GlobalEnvironment->barrier(); - } - - void TearDown() override { - GlobalEnvironment->barrier(); - - if (GlobalEnvironment->comm_->rank() == 0 && std::filesystem::exists(temp_dir)) { - std::filesystem::remove_all(temp_dir); - } - - BaseStreamingFixture::TearDown(); - } - - [[nodiscard]] cudf::io::source_info get_source_info(bool truncate_file_list) const { - if (truncate_file_list) { - std::vector files( - source_files.begin(), source_files.begin() + 2 - ); - return cudf::io::source_info(source_files); - } else { - return cudf::io::source_info(source_files); - } - } - - std::filesystem::path temp_dir; - std::vector source_files; -}; - -using ReadParquetParams = - std::tuple, std::optional, bool, bool>; - -class StreamingReadParquetParams - : public StreamingReadParquet, - public ::testing::WithParamInterface {}; - -INSTANTIATE_TEST_SUITE_P( - ReadParquetCombinations, - StreamingReadParquetParams, - ::testing::Combine( - // skip_rows - ::testing::Values( - std::nullopt, - std::optional{7}, - std::optional{19}, - std::optional{113} - ), - // num_rows - ::testing::Values( - std::nullopt, - std::optional{0}, - std::optional{3}, - std::optional{31}, - std::optional{83} - ), - // use_filter - ::testing::Values(false, true), - // truncate file list - ::testing::Values(false, true) - ), - [](const ::testing::TestParamInfo& info) { - auto const& skip_rows = std::get<0>(info.param); - auto const& num_rows = std::get<1>(info.param); - auto const& use_filter = std::get<2>(info.param); - auto const& truncate_file_list = std::get<3>(info.param); - std::string result = "skip_rows_"; - result += skip_rows.has_value() ? std::to_string(skip_rows.value()) : "none"; - result += "_num_rows_"; - result += num_rows.has_value() ? std::to_string(num_rows.value()) : "all"; - if (use_filter) { - result += "_with_filter"; - } else { - result += "_no_filter"; - } - if (truncate_file_list) { - result += "_one_file"; - } else { - result += "_all_files"; - } - return result; - } -); - -TEST_P(StreamingReadParquetParams, ReadParquet) { - auto [skip_rows, num_rows, use_filter, truncate_file_list] = GetParam(); - auto source = get_source_info(truncate_file_list); - - auto options = cudf::io::parquet_reader_options::builder(source).build(); - if (skip_rows.has_value()) { - options.set_skip_rows(skip_rows.value()); - } - if (num_rows.has_value()) { - options.set_num_rows(num_rows.value()); - } - auto filter_expr = [&]() -> std::unique_ptr { - if (!use_filter) { - return nullptr; - } - auto stream = ctx->br()->stream_pool().get_stream(); - auto owner = new std::vector; - owner->push_back( - std::make_shared>(15, true, stream) - ); - owner->push_back( - std::make_shared( - *std::any_cast>>( - owner->at(0) - ) - ) - ); - owner->push_back(std::make_shared(0)); - owner->push_back( - std::make_shared( - cudf::ast::ast_operator::LESS, - *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); - }) - ); - }(); - auto expected = [&]() { - if (filter_expr != nullptr) { - auto expected_options = options; - expected_options.set_filter(filter_expr->filter); - filter_expr->stream.synchronize(); - auto expected = cudf::io::read_parquet(expected_options).tbl; - filter_expr->stream.synchronize(); - return expected; - } else { - return cudf::io::read_parquet(options).tbl; - } - }(); - auto ch = ctx->create_channel(); - std::vector actors; - - actors.push_back( - actor::read_parquet( - ctx, GlobalEnvironment->comm_, ch, 4, options, 3, std::move(filter_expr) - ) - ); - - std::vector messages; - actors.push_back(actor::pull_from_channel(ctx, ch, messages)); - - if (GlobalEnvironment->comm_->nranks() > 1 - && (skip_rows.value_or(0) > 0 || num_rows.has_value())) - { - // We don't yet implement skip_rows/num_rows in multi-rank mode - EXPECT_THROW(run_actor_network(std::move(actors)), std::logic_error); - return; - } - run_actor_network(std::move(actors)); - - coll::AllGather allgather( - GlobalEnvironment->comm_, - /* op_id = */ 0, - br.get() - ); - - for (auto& msg : messages) { - auto chunk = msg.release(); - auto seq = msg.sequence_number(); - auto [reservation, _] = br->reserve( - MemoryType::DEVICE, chunk.make_available_cost(), AllowOverbooking::YES - ); - chunk = chunk.make_available(reservation); - auto packed_columns = - cudf::pack(chunk.table_view(), chunk.stream(), br->device_mr()); - auto packed_data = PackedData{ - std::move(packed_columns.metadata), - br->move(std::move(packed_columns.gpu_data), chunk.stream()) - }; - - allgather.insert(seq, std::move(packed_data)); - } - - allgather.insert_finished(); - - // May as well check on all ranks, so we also mildly exercise the allgather. - auto gathered_packed_data = allgather.wait_and_extract(coll::AllGather::Ordered::YES); - auto result = unpack_and_concat( - std::move(gathered_packed_data), rmm::cuda_stream_default, br.get() - ); - EXPECT_EQ(result->num_rows(), expected->num_rows()); - EXPECT_EQ(result->num_columns(), expected->num_columns()); - EXPECT_EQ(result->num_columns(), 1); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(result->view(), expected->view()); -} diff --git a/cpp/tests/streaming/test_shuffler.cpp b/cpp/tests/streaming/test_shuffler.cpp index 66d899eb9..7bfa893a5 100644 --- a/cpp/tests/streaming/test_shuffler.cpp +++ b/cpp/tests/streaming/test_shuffler.cpp @@ -7,18 +7,18 @@ #include #include +#include +#include +#include #include #include #include -#include #include #include #include #include #include -#include -#include #include "../utils.hpp" #include "base_streaming_fixture.hpp" @@ -76,23 +76,25 @@ class StreamingShuffler : public BaseStreamingShuffle, 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 + input_chunks.emplace_back( + cudf_streaming::streaming::to_message( + i, + std::make_unique( + std::make_unique( + cudf::slice( + full_input_table, + {static_cast(i * chunk_size), + static_cast((i + 1) * chunk_size)}, + stream + ) + .at(0), + stream, + ctx->br()->device_mr() + ), + stream + ) ) - )); + ); } // Create and run the streaming pipeline. @@ -104,7 +106,7 @@ class StreamingShuffler : public BaseStreamingShuffle, auto ch2 = ctx->create_channel(); actors.push_back( - actor::partition_and_pack( + cudf_streaming::streaming::actor::partition_and_pack( ctx, ch1, ch2, {1}, num_partitions, hash_function, seed ) ); @@ -113,7 +115,9 @@ class StreamingShuffler : public BaseStreamingShuffle, actors.emplace_back(make_shuffler_actor_fn(ch2, ch3)); auto ch4 = ctx->create_channel(); - actors.push_back(actor::unpack_and_concat(ctx, ch3, ch4)); + actors.push_back( + cudf_streaming::streaming::actor::unpack_and_concat(ctx, ch3, ch4) + ); actors.push_back(actor::pull_from_channel(ctx, ch4, output_chunks)); @@ -150,7 +154,9 @@ class StreamingShuffler : public BaseStreamingShuffle, // 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()); + output_chunks_as_views.push_back( + chunk.get().table_view() + ); } auto result_table = cudf::concatenate(output_chunks_as_views); diff --git a/cpp/tests/streaming/test_table_chunk.cpp b/cpp/tests/streaming/test_table_chunk.cpp deleted file mode 100644 index 5d42285bd..000000000 --- a/cpp/tests/streaming/test_table_chunk.cpp +++ /dev/null @@ -1,538 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - - -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include "../utils.hpp" -#include "base_streaming_fixture.hpp" - - -using namespace rapidsmpf; -using namespace rapidsmpf::streaming; - -class StreamingTableChunk : public BaseStreamingFixture, - public ::testing::WithParamInterface { - protected: - void SetUp() override { - rapidsmpf::config::Options options( - rapidsmpf::config::get_environment_variables() - ); - - std::unordered_map memory_limits{}; - auto stream_pool = std::make_shared( - 16, rmm::cuda_stream::flags::non_blocking - ); - stream = cudf::get_default_stream(); - br = rapidsmpf::BufferResource::create( - mr_cuda, // device_mr - rapidsmpf::PinnedMemoryResource::make_if_available(), // pinned_mr - memory_limits, // memory_limits - std::chrono::milliseconds{1}, // periodic_spill_check - stream_pool, // stream_pool - Statistics::disabled() // statistics - ); - ctx = std::make_shared( - options, GlobalEnvironment->comm_->logger(), br - ); - } - - rmm::cuda_stream_view stream; - rmm::mr::cuda_memory_resource mr_cuda; - std::shared_ptr br; - std::shared_ptr ctx; -}; - -TEST_F(StreamingTableChunk, FromTable) { - constexpr unsigned int num_rows = 100; - constexpr std::int64_t seed = 1337; - - cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); - - TableChunk chunk{std::make_unique(expect), stream}; - EXPECT_EQ(chunk.stream().value(), stream.value()); - EXPECT_TRUE(chunk.is_available()); - EXPECT_TRUE(chunk.is_spillable()); - EXPECT_EQ(chunk.make_available_cost(), 0); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk.table_view(), expect); - - auto chunk2 = chunk.make_available( - br->reserve_or_fail(chunk.make_available_cost(), MemoryType::DEVICE) - ); - EXPECT_FALSE(chunk.is_available()); - EXPECT_TRUE(chunk2.is_available()); - EXPECT_TRUE(chunk2.is_spillable()); - EXPECT_EQ(chunk2.make_available_cost(), 0); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk2.table_view(), expect); -} - -TEST_F(StreamingTableChunk, TableChunkOwner) { - constexpr unsigned int num_rows = 100; - constexpr std::int64_t seed = 1337; - constexpr std::uint64_t seq = 42; - - cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); - // Static because the deleter function is a void(*)(void*) which precludes the use of - // a lambda with captures. - static std::size_t num_deletions{0}; - auto deleter = [](void* p) { - num_deletions++; - delete static_cast(p); - }; - auto make_chunk = [&](TableChunk::ExclusiveView exclusive_view) { - return TableChunk{ - expect, stream, OwningWrapper(new int, deleter), exclusive_view - }; - }; - auto check_chunk = [&](TableChunk const& chunk, bool is_spillable) { - EXPECT_EQ(chunk.stream().value(), stream.value()); - EXPECT_TRUE(chunk.is_available()); - EXPECT_EQ(chunk.is_spillable(), is_spillable); - EXPECT_EQ(chunk.make_available_cost(), 0); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk.table_view(), expect); - }; - { - auto chunk = make_chunk(TableChunk::ExclusiveView::NO); - check_chunk(chunk, false); - EXPECT_EQ(num_deletions, 0); - } - EXPECT_EQ(num_deletions, 1); - { - auto msg = to_message( - seq, std::make_unique(make_chunk(TableChunk::ExclusiveView::NO)) - ); - EXPECT_EQ(num_deletions, 1); - } - EXPECT_EQ(num_deletions, 2); - { - auto msg = to_message( - seq, std::make_unique(make_chunk(TableChunk::ExclusiveView::YES)) - ); - auto chunk = msg.release(); - check_chunk(chunk, true); - EXPECT_EQ(num_deletions, 2); - } - EXPECT_EQ(num_deletions, 3); - { - auto chunk = make_chunk(TableChunk::ExclusiveView::YES); - check_chunk(chunk, true); - auto res = br->reserve_or_fail( - chunk.data_alloc_size(MemoryType::DEVICE), MemoryType::DEVICE - ); - // This is like spilling since the original `chunk` is ExclusiveView::YES and - // overwritten. - chunk = chunk.copy(res); - EXPECT_EQ(num_deletions, 4); - } -} - -TEST_F(StreamingTableChunk, FromPackedDataOnDevice) { - constexpr unsigned int num_rows = 100; - constexpr std::int64_t seed = 1337; - - cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); - auto packed_columns = cudf::pack(expect, stream); - - auto packed_data = std::make_unique( - std::move(packed_columns.metadata), - br->move(std::move(packed_columns.gpu_data), stream) - ); - TableChunk chunk{std::move(packed_data)}; - - EXPECT_EQ(chunk.stream().value(), stream.value()); - // chunk was created from packed data on device, so it is available and make available - // cost is 0. - EXPECT_TRUE(chunk.is_available()); - EXPECT_TRUE(chunk.is_spillable()); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, chunk.table_view()); - EXPECT_EQ(chunk.make_available_cost(), 0); - - auto chunk2 = chunk.make_available( - br->reserve_or_fail(chunk.make_available_cost(), MemoryType::DEVICE) - ); - EXPECT_FALSE(chunk.is_available()); - EXPECT_TRUE(chunk2.is_available()); - EXPECT_TRUE(chunk2.is_spillable()); - EXPECT_EQ(chunk2.make_available_cost(), 0); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk2.table_view(), expect); -} - -INSTANTIATE_TEST_SUITE_P( - StreamingTableChunkWithSpillTargets, - StreamingTableChunk, - ::testing::ValuesIn(rapidsmpf::SPILL_TARGET_MEMORY_TYPES), - [](testing::TestParamInfo const& info) { - return std::string{rapidsmpf::to_string(info.param)}; - } -); - -TEST_P(StreamingTableChunk, FromPackedDataOn) { - auto const spill_mem_type = GetParam(); - if (spill_mem_type == MemoryType::PINNED_HOST - && !is_pinned_memory_resources_supported()) - { - GTEST_SKIP() << "MemoryType::PINNED_HOST isn't supported on the system."; - } - - constexpr unsigned int num_rows = 100; - constexpr std::int64_t seed = 1337; - - cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); - auto packed_columns = cudf::pack(expect, stream); - std::size_t const size = packed_columns.gpu_data->size(); - - // Move the gpu_data to a Buffer (still device memory). - auto gpu_data_on_device = br->move(std::move(packed_columns.gpu_data), stream); - - // Copy the GPU data to the current spill target memory type. - auto [res, _] = br->reserve(spill_mem_type, size, AllowOverbooking::YES); - auto gpu_data_in_spill_memory = br->move(std::move(gpu_data_on_device), res); - - auto packed_data = std::make_unique( - std::move(packed_columns.metadata), std::move(gpu_data_in_spill_memory) - ); - TableChunk chunk{std::move(packed_data)}; - - EXPECT_EQ(chunk.stream().value(), stream.value()); - EXPECT_FALSE(chunk.is_available()); - EXPECT_TRUE(chunk.is_spillable()); - EXPECT_THROW(std::ignore = chunk.table_view(), std::invalid_argument); - EXPECT_EQ(chunk.make_available_cost(), size); - - auto chunk2 = chunk.make_available( - br->reserve_or_fail(chunk.make_available_cost(), MemoryType::DEVICE) - ); - EXPECT_FALSE(chunk.is_available()); - EXPECT_TRUE(chunk2.is_available()); - EXPECT_TRUE(chunk2.is_spillable()); - EXPECT_EQ(chunk2.make_available_cost(), 0); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk2.table_view(), expect); -} - -TEST_F(StreamingTableChunk, DeviceToDeviceCopy) { - constexpr unsigned int num_rows = 100; - constexpr std::int64_t seed = 1337; - - auto expect = random_table_with_index(seed, num_rows, 0, 10); - - rapidsmpf::streaming::TableChunk chunk{std::make_unique(expect), stream}; - EXPECT_TRUE(chunk.is_available()); - - auto res = br->reserve_or_fail( - chunk.data_alloc_size(MemoryType::DEVICE), MemoryType::DEVICE - ); - auto chunk2 = chunk.copy(res); - - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk2.table_view(), expect); -} - -TEST_F(StreamingTableChunk, ShapeOnAvailableAndSpilledChunk) { - constexpr unsigned int num_rows = 64; - constexpr std::int64_t seed = 2025; - - cudf::table expect = random_table_with_index(seed, num_rows, 0, 5); - auto const expected_shape = std::pair{ - expect.num_rows(), expect.num_columns() - }; - - TableChunk device_chunk{std::make_unique(expect), stream}; - EXPECT_TRUE(device_chunk.is_available()); - EXPECT_EQ(device_chunk.shape(), expected_shape); - - auto [res, _] = br->reserve( - MemoryType::HOST, - device_chunk.data_alloc_size(MemoryType::DEVICE), - AllowOverbooking::YES - ); - auto host_chunk = device_chunk.copy(res); - - EXPECT_FALSE(host_chunk.is_available()); - EXPECT_EQ(host_chunk.shape(), expected_shape); - - device_chunk = host_chunk.make_available( - br->reserve_or_fail(host_chunk.make_available_cost(), MemoryType::DEVICE) - ); - EXPECT_TRUE(device_chunk.is_available()); - EXPECT_EQ(device_chunk.shape(), expected_shape); -} - -TEST_P(StreamingTableChunk, DeviceToHostRoundTripCopy) { - auto const spill_mem_type = GetParam(); - if (spill_mem_type == MemoryType::PINNED_HOST - && !is_pinned_memory_resources_supported()) - { - GTEST_SKIP() << "MemoryType::PINNED_HOST isn't supported on the system."; - } - - constexpr unsigned int num_rows = 64; - constexpr std::int64_t seed = 2025; - - auto expect = random_table_with_index(seed, num_rows, 0, 5); - - TableChunk dev_chunk{std::make_unique(expect), stream}; - EXPECT_TRUE(dev_chunk.is_available()); - EXPECT_TRUE(dev_chunk.is_spillable()); - EXPECT_EQ(dev_chunk.stream().value(), stream.value()); - EXPECT_EQ(dev_chunk.make_available_cost(), 0); - { - auto cd = get_content_description(dev_chunk); - EXPECT_EQ(cd.spillable(), dev_chunk.is_spillable()); - for (auto mem_type : MEMORY_TYPES) { - EXPECT_EQ(cd.content_size(mem_type), dev_chunk.data_alloc_size(mem_type)); - } - } - - // Copy to host memory -> new chunk should be unavailable. - auto host_res = br->reserve_or_fail( - dev_chunk.data_alloc_size(MemoryType::DEVICE), spill_mem_type - ); - auto host_copy = dev_chunk.copy(host_res); - EXPECT_FALSE(host_copy.is_available()); - EXPECT_TRUE(host_copy.is_spillable()); - EXPECT_EQ(host_copy.stream().value(), stream.value()); - EXPECT_GT(host_copy.make_available_cost(), 0); - { - auto cd = get_content_description(host_copy); - EXPECT_EQ(cd.spillable(), host_copy.is_spillable()); - for (auto mem_type : MEMORY_TYPES) { - EXPECT_EQ(cd.content_size(mem_type), host_copy.data_alloc_size(mem_type)); - } - } - - // Host to host copy. - auto host_res2 = - br->reserve_or_fail(host_copy.data_alloc_size(spill_mem_type), spill_mem_type); - auto host_copy2 = host_copy.copy(host_res2); - EXPECT_FALSE(host_copy2.is_available()); - EXPECT_TRUE(host_copy2.is_spillable()); - EXPECT_EQ(host_copy2.stream().value(), stream.value()); - EXPECT_EQ(host_copy2.make_available_cost(), host_copy.make_available_cost()); - { - auto cd = get_content_description(host_copy2); - EXPECT_EQ(cd.spillable(), host_copy2.is_spillable()); - for (auto mem_type : MEMORY_TYPES) { - EXPECT_EQ(cd.content_size(mem_type), host_copy2.data_alloc_size(mem_type)); - } - } - - // Bring the new host copy back to device and verify equality. - auto dev_res = br->reserve_or_fail( - host_copy2.data_alloc_size(spill_mem_type), MemoryType::DEVICE - ); - auto dev_back = host_copy2.make_available(dev_res); - EXPECT_TRUE(dev_back.is_available()); - EXPECT_TRUE(dev_back.is_spillable()); - EXPECT_EQ(dev_back.stream().value(), stream.value()); - EXPECT_EQ(dev_back.make_available_cost(), 0); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(dev_back.table_view(), expect); - { - auto cd = get_content_description(dev_back); - EXPECT_EQ(cd.spillable(), dev_back.is_spillable()); - for (auto mem_type : MEMORY_TYPES) { - EXPECT_EQ(cd.content_size(mem_type), dev_back.data_alloc_size(mem_type)); - } - } - - // Sanity check: a second device copy should also remain equivalent. - auto dev_res2 = br->reserve_or_fail( - dev_back.data_alloc_size(MemoryType::DEVICE), MemoryType::DEVICE - ); - auto dev_copy2 = dev_back.copy(dev_res2); - EXPECT_TRUE(dev_copy2.is_available()); - EXPECT_EQ(dev_copy2.make_available_cost(), 0); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(dev_copy2.table_view(), expect); - { - auto cd = get_content_description(dev_copy2); - EXPECT_EQ(cd.spillable(), dev_copy2.is_spillable()); - for (auto mem_type : MEMORY_TYPES) { - EXPECT_EQ(cd.content_size(mem_type), dev_copy2.data_alloc_size(mem_type)); - } - } -} - -TEST_F(StreamingTableChunk, ToMessageRoundTrip) { - constexpr unsigned int num_rows = 64; - constexpr std::int64_t seed = 2025; - constexpr std::uint64_t seq = 7; - - auto expect = random_table_with_index(seed, num_rows, 0, 5); - auto chunk = - std::make_unique(std::make_unique(expect), stream); - - Message m = to_message(seq, std::move(chunk)); - EXPECT_FALSE(m.empty()); - EXPECT_TRUE(m.holds()); - EXPECT_TRUE(m.content_description().spillable()); - EXPECT_EQ(m.content_description().content_size(MemoryType::HOST), 0); - EXPECT_EQ(m.content_description().content_size(MemoryType::DEVICE), 1024); - EXPECT_EQ(m.sequence_number(), seq); - - // Deep-copy: device to host. - auto reservation = br->reserve_or_fail(m.copy_cost(), MemoryType::HOST); - Message m2 = m.copy(reservation); - EXPECT_EQ(reservation.size(), 0); - EXPECT_FALSE(m2.empty()); - EXPECT_TRUE(m2.holds()); - EXPECT_TRUE(m2.content_description().spillable()); - EXPECT_EQ(m2.content_description().content_size(MemoryType::HOST), 1024); - EXPECT_EQ(m2.content_description().content_size(MemoryType::DEVICE), 0); - EXPECT_EQ(m2.sequence_number(), seq); - - // Deep-copy: host to host. - reservation = br->reserve_or_fail(m2.copy_cost(), MemoryType::HOST); - Message m3 = m.copy(reservation); - EXPECT_EQ(reservation.size(), 0); - EXPECT_FALSE(m3.empty()); - EXPECT_TRUE(m3.holds()); - EXPECT_TRUE(m3.content_description().spillable()); - EXPECT_EQ(m3.content_description().content_size(MemoryType::HOST), 1024); - EXPECT_EQ(m3.content_description().content_size(MemoryType::DEVICE), 0); - EXPECT_EQ(m3.sequence_number(), seq); - - // Copy the chunk back to device and verify. - { - auto chunk = m3.release(); - auto res = br->reserve_or_fail(chunk.make_available_cost(), MemoryType::DEVICE); - chunk = chunk.make_available(res); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(chunk.table_view(), expect); - } - - // Deep-copy: host to device. - reservation = br->reserve_or_fail(m2.copy_cost(), MemoryType::DEVICE); - Message m4 = m.copy(reservation); - EXPECT_EQ(reservation.size(), 0); - EXPECT_FALSE(m4.empty()); - EXPECT_TRUE(m4.holds()); - EXPECT_TRUE(m4.content_description().spillable()); - EXPECT_EQ(m4.content_description().content_size(MemoryType::HOST), 0); - EXPECT_EQ(m4.content_description().content_size(MemoryType::DEVICE), 1024); - EXPECT_EQ(m4.sequence_number(), seq); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(m4.get().table_view(), expect); - - // Deep-copy: device to device. - reservation = br->reserve_or_fail(m4.copy_cost(), MemoryType::DEVICE); - Message m5 = m.copy(reservation); - EXPECT_EQ(reservation.size(), 0); - EXPECT_FALSE(m5.empty()); - EXPECT_TRUE(m5.holds()); - EXPECT_TRUE(m5.content_description().spillable()); - EXPECT_EQ(m5.content_description().content_size(MemoryType::HOST), 0); - EXPECT_EQ(m5.content_description().content_size(MemoryType::DEVICE), 1024); - EXPECT_EQ(m5.sequence_number(), seq); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(m5.get().table_view(), expect); -} - -TEST_F(StreamingTableChunk, ToMessageNotSpillable) { - constexpr unsigned int num_rows = 100; - constexpr std::int64_t seed = 1337; - constexpr std::uint64_t seq = 42; - - cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); - - auto deleter = [](void* p) { delete static_cast(p); }; - auto chunk = std::make_unique( - expect, stream, OwningWrapper(new int, deleter), TableChunk::ExclusiveView::NO - ); - - Message m = to_message(seq, std::move(chunk)); - EXPECT_FALSE(m.empty()); - EXPECT_TRUE(m.holds()); - EXPECT_FALSE(m.content_description().spillable()); - EXPECT_EQ(m.content_description().content_size(MemoryType::HOST), 0); - EXPECT_EQ( - m.content_description().content_size(MemoryType::DEVICE), - cudf::packed_size( - expect.view(), stream, rmm::mr::get_current_device_resource_ref() - ) - ); - // packed size is greater than or equal to the alloc size due to buffer alignments. - EXPECT_GE( - m.content_description().content_size(MemoryType::DEVICE), expect.alloc_size() - ); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT(m.get().table_view(), expect); -} - -TEST_F(StreamingTableChunk, ToPackedDataFromPackedChunk) { - constexpr unsigned int num_rows = 100; - constexpr std::int64_t seed = 1337; - - cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); - auto packed_columns = cudf::pack(expect, stream); - TableChunk chunk{std::make_unique( - std::move(packed_columns.metadata), - br->move(std::move(packed_columns.gpu_data), stream) - )}; - EXPECT_TRUE(chunk.is_available()); - - auto packed = std::move(chunk).into_packed_data(br.get()); - EXPECT_FALSE(chunk.is_available()); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT( - expect, TableChunk{std::move(packed)}.table_view() - ); -} - -TEST_F(StreamingTableChunk, ToPackedDataFromTable) { - constexpr unsigned int num_rows = 100; - constexpr std::int64_t seed = 1337; - - cudf::table expect = random_table_with_index(seed, num_rows, 0, 10); - TableChunk chunk{std::make_unique(expect), stream}; - EXPECT_TRUE(chunk.is_available()); - - auto packed = std::move(chunk).into_packed_data(br.get()); - EXPECT_FALSE(chunk.is_available()); - CUDF_TEST_EXPECT_TABLES_EQUIVALENT( - expect, TableChunk{std::move(packed)}.table_view() - ); -} - -TEST_F(StreamingTableChunk, ToMessageUnalignedSize) { - constexpr unsigned int num_rows = 5; - constexpr std::int64_t seed = 2025; - constexpr std::uint64_t seq = 7; - - auto expect = random_table_with_index(seed, num_rows, 0, 5); - auto chunk = - std::make_unique(std::make_unique(expect), stream); - - Message m = to_message(seq, std::move(chunk)); - EXPECT_EQ(m.sequence_number(), seq); - EXPECT_FALSE(m.empty()); - EXPECT_TRUE(m.holds()); - EXPECT_TRUE(m.content_description().spillable()); - EXPECT_EQ(m.content_description().content_size(MemoryType::HOST), 0); - EXPECT_EQ(m.content_description().content_size(MemoryType::DEVICE), 80); - EXPECT_EQ(m.copy_cost(), 80); - - // Deep copy: device → host. - // Note: `m.copy_cost() == 80`, but cudf performs 128-byte aligned allocations. - // This means `m.copy_cost()` is not always sufficient; however, TableChunk.copy() - // accounts for this alignment internally. - auto reservation = br->reserve_or_fail(m.copy_cost(), MemoryType::HOST); - Message m2 = m.copy(reservation); - EXPECT_EQ(reservation.size(), 0); - EXPECT_FALSE(m2.empty()); - EXPECT_TRUE(m2.holds()); - EXPECT_TRUE(m2.content_description().spillable()); - EXPECT_EQ(m2.copy_cost(), 128); - EXPECT_EQ(m2.content_description().content_size(MemoryType::HOST), 128); - EXPECT_EQ(m2.content_description().content_size(MemoryType::DEVICE), 0); - EXPECT_EQ(m2.sequence_number(), seq); -} diff --git a/cpp/tests/test_allgather.cpp b/cpp/tests/test_allgather.cpp index 99472b499..4cb3447b8 100644 --- a/cpp/tests/test_allgather.cpp +++ b/cpp/tests/test_allgather.cpp @@ -10,8 +10,7 @@ #include -#include -#include +#include #include #include @@ -33,7 +32,7 @@ extern Environment* GlobalEnvironment; class BaseAllGatherTest : public ::testing::Test { protected: void SetUp() override { - stream = cudf::get_default_stream(); + stream = rmm::cuda_stream_view{}; br = rapidsmpf::BufferResource::create(rmm::mr::cuda_memory_resource{}); } @@ -361,7 +360,7 @@ TEST_F(BaseAllGatherTest, opid_reuse) { // (the largest chunk, since none covers 100 alone). The second must search for a chunk // >= 10 and pick the 20-byte chunk, totalling 110. TEST(PostBox, spill_uses_remaining_amount) { - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_view{}; auto mr = std::make_unique(); auto br = rapidsmpf::BufferResource::create(*mr); diff --git a/cpp/tests/test_buffer.cpp b/cpp/tests/test_buffer.cpp index 742339c6f..a0814c497 100644 --- a/cpp/tests/test_buffer.cpp +++ b/cpp/tests/test_buffer.cpp @@ -12,10 +12,10 @@ #include -#include #include #include #include +#include #include #include @@ -54,7 +54,7 @@ class BufferRebindStreamTest : public ::testing::TestWithParam { } br = BufferResource::create( - cudf::get_current_device_resource_ref(), + rmm::mr::get_current_device_resource_ref(), PinnedMemoryResource::make_if_available(), std::unordered_map{}, std::nullopt, diff --git a/cpp/tests/test_buffer_resource.cpp b/cpp/tests/test_buffer_resource.cpp index 95cb74edc..a2e94445b 100644 --- a/cpp/tests/test_buffer_resource.cpp +++ b/cpp/tests/test_buffer_resource.cpp @@ -9,11 +9,10 @@ #include -#include -#include -#include -#include +#include #include +#include +#include #include #include @@ -55,7 +54,7 @@ std::unique_ptr zeros( TEST(BufferResource, ReservationOverbooking) { // Create a buffer resource that always reports 10 KiB of available device memory. auto br = BufferResource::create( - cudf::get_current_device_resource_ref(), + rmm::mr::get_current_device_resource_ref(), PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}} ); @@ -122,7 +121,7 @@ TEST(BufferResource, ReservationReleasing) { // Create a buffer resource that always reports 10 KiB of available host and device // memory. auto br = BufferResource::create( - cudf::get_current_device_resource_ref(), + rmm::mr::get_current_device_resource_ref(), PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, 10_KiB}, {MemoryType::HOST, 10_KiB}} ); @@ -171,7 +170,7 @@ TEST(BufferResource, ReservationReleasing) { TEST(BufferResource, MemoryLimit) { rmm::mr::cuda_memory_resource mr_cuda; - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_view{}; // Create a buffer resource that limits available device memory to 10 KiB. auto br = BufferResource::create( @@ -299,7 +298,7 @@ TEST(BufferResource, AllocStatistics) { std::make_shared(1, rmm::cuda_stream::flags::non_blocking), stats ); - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_view{}; constexpr std::size_t device_size = 4_KiB; constexpr std::size_t pinned_size = 8_KiB; @@ -419,8 +418,8 @@ TEST_F(BufferResourceReserveOrFailTest, MultipleTypes) { class BaseBufferResourceCopyTest : public ::testing::Test { protected: void SetUp() override { - br = BufferResource::create(cudf::get_current_device_resource_ref()); - stream = cudf::get_default_stream(); + br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); + stream = rmm::cuda_stream_view{}; // initialize the host pattern host_pattern.resize(buffer_size); @@ -631,7 +630,7 @@ class BufferResourceDifferentResourcesTest : public ::testing::Test { protected: void SetUp() override { buffer_size = 1_KiB; - stream = cudf::get_default_stream(); + stream = rmm::cuda_stream_view{}; // Host pattern for initialization and verification host_pattern.resize(buffer_size); @@ -792,9 +791,9 @@ TEST_F(BufferCopyEdgeCases, SameBufferIsDisallowed) { TEST(BufferResource, DeviceMrKeepsBufferResourceAlive) { constexpr std::size_t N = 1024; - auto br = BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); std::weak_ptr weak_br = br; - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_view{}; // Construct a device_buffer using the BR memory resource. Internally, // `rmm::device_buffer` stores the resource as an owning `cuda::mr::any_resource`, @@ -823,7 +822,7 @@ TEST(OwningResourceAdaptor, CopyThrowsWhenBackRefExpired) { std::weak_ptr weak_br; { - auto br = BufferResource::create(cudf::get_current_device_resource_ref()); + auto br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); weak_br = br; } ASSERT_TRUE(weak_br.expired()); diff --git a/cpp/tests/test_chunk.cpp b/cpp/tests/test_chunk.cpp index 79ea8ab9d..8aa014836 100644 --- a/cpp/tests/test_chunk.cpp +++ b/cpp/tests/test_chunk.cpp @@ -9,8 +9,9 @@ #include #include -#include -#include +#include +#include +#include #include #include @@ -24,8 +25,8 @@ using namespace rapidsmpf::shuffler::detail; class ChunkTest : public ::testing::Test { protected: void SetUp() override { - br = BufferResource::create(cudf::get_current_device_resource_ref()); - stream = cudf::get_default_stream(); + br = BufferResource::create(rmm::mr::get_current_device_resource_ref()); + stream = rmm::cuda_stream_view{}; } std::shared_ptr br; @@ -66,8 +67,7 @@ TEST_P(ChunkFromPackedDataTest, RoundTrip) { std::vector{1, 2, 3, 4} ); - auto data = - std::make_unique(data_size, cudf::get_default_stream()); + auto data = std::make_unique(data_size, rmm::cuda_stream_view{}); if (data_size > 0) { std::vector host_data(data_size); std::iota(host_data.begin(), host_data.end(), std::uint8_t{5}); diff --git a/cpp/tests/test_config.cpp b/cpp/tests/test_config.cpp index c78d840cf..6cf9ada50 100644 --- a/cpp/tests/test_config.cpp +++ b/cpp/tests/test_config.cpp @@ -8,7 +8,6 @@ #include #include -#include #include #include diff --git a/cpp/tests/test_cudf_utils.cpp b/cpp/tests/test_cudf_utils.cpp deleted file mode 100644 index d5cf30555..000000000 --- a/cpp/tests/test_cudf_utils.cpp +++ /dev/null @@ -1,185 +0,0 @@ -/** - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include - -#include - -#include -#include - -#include - -using namespace rapidsmpf; - -class BaseEstimatedMemoryUsageTest : public ::testing::Test { - protected: - void SetUp() override { - stream = cudf::get_default_stream(); - } - - rmm::cuda_stream_view stream; -}; - -/** - * @brief Templated test suite for testing estimated_memory_usage with different column - * types - */ -template -class EstimatedMemoryUsageTest : public BaseEstimatedMemoryUsageTest {}; - -// Define the types to test -using ColumnTypes = ::testing:: - Types; - -TYPED_TEST_SUITE(EstimatedMemoryUsageTest, ColumnTypes); - -TYPED_TEST(EstimatedMemoryUsageTest, FixedWidthColumnMemoryUsage) { - using T = TypeParam; - // Test with different sizes - std::vector test_sizes = {0, 1, 10, 100, 1000, 1000000}; - - for (auto size : test_sizes) { - SCOPED_TRACE("test size: " + std::to_string(size)); - std::vector data(size); - - cudf::test::fixed_width_column_wrapper wrapper(data.begin(), data.end()); - auto column = wrapper.release(); - - std::size_t exp = column->alloc_size(); - std::size_t est = rapidsmpf::estimated_memory_usage(column->view(), this->stream); - - EXPECT_EQ(exp, est); - } -} - -/** - * @brief Test suite for string column memory usage estimation - */ -TEST_F(BaseEstimatedMemoryUsageTest, StringType) { - // Test with different string data - std::vector> test_cases = { - {}, // Empty column - {"hello"}, // Single string - {"hello", "world", "test"}, // Multiple strings - { - "", "a", "very long string that should take more memory", "short" - }, // Mixed lengths - std::vector(100, "repeated string") // Many repeated strings - }; - - for (const auto& data : test_cases) { - // Create a string column - cudf::test::strings_column_wrapper wrapper(data.begin(), data.end()); - auto column = wrapper.release(); - - std::size_t exp = column->alloc_size(); - std::size_t est = rapidsmpf::estimated_memory_usage(column->view(), stream); - - EXPECT_EQ(exp, est); - } -} - -/** - * @brief Test suite for list column memory usage estimation - */ -TEST_F(BaseEstimatedMemoryUsageTest, ListType) { - // Test with different list data - std::vector> test_cases = { - {}, // Empty column - {1, 2, 3}, // Single list - {1, 2, 3, 4, 5, 6}, // Multiple values - {0, 1, 2, 3, 4, 5}, // Mixed values - std::vector(100, 42) // Many repeated values - }; - - for (const auto& data : test_cases) { - // Create a list column - cudf::test::lists_column_wrapper wrapper(data.begin(), data.end()); - auto column = wrapper.release(); - - std::size_t exp = column->alloc_size(); - std::size_t est = rapidsmpf::estimated_memory_usage(column->view(), stream); - - EXPECT_EQ(exp, est); - } -} - -/** - * @brief Test suite for struct column memory usage estimation - */ -TEST_F(BaseEstimatedMemoryUsageTest, StructType) { - // Test with different struct data configurations - std::vector>> test_cases = { - {}, // Empty struct column - {{std::make_pair(1, "hello")}}, // Single struct - {{std::make_pair(1, "hello"), std::make_pair(2, "world")}}, // Two structs - {{std::make_pair(0, ""), - std::make_pair(100, "very long string"), - std::make_pair(42, "short")}}, // Mixed data - std::vector>( - 50, std::make_pair(42, "repeated") - ) // Many repeated structs - }; - - for (const auto& data : test_cases) { - // Create struct columns for each field - std::vector int_data; - std::vector string_data; - - for (const auto& item : data) { - int_data.push_back(item.first); - string_data.push_back(item.second); - } - - cudf::test::fixed_width_column_wrapper int_wrapper( - int_data.begin(), int_data.end() - ); - cudf::test::strings_column_wrapper string_wrapper( - string_data.begin(), string_data.end() - ); - - std::vector> children; - children.push_back(int_wrapper.release()); - children.push_back(string_wrapper.release()); - - cudf::test::structs_column_wrapper wrapper(std::move(children)); - auto column = wrapper.release(); - - std::size_t exp = column->alloc_size(); - std::size_t est = rapidsmpf::estimated_memory_usage(column->view(), stream); - - EXPECT_EQ(exp, est); - } -} - -/** - * @brief Test suite for dictionary column memory usage estimation - */ -TEST_F(BaseEstimatedMemoryUsageTest, DictionaryType) { - // Test with different dictionary data - std::vector> test_cases = { - {}, // Empty column - {"hello"}, // Single value - {"hello", "world", "test", "hello", "world"}, // Repeated values - { - "", "a", "very long string", "short", "a", "very long string" - }, // Mixed with repetition - std::vector(100, "repeated") // Many repeated values - }; - - for (const auto& data : test_cases) { - // Create a dictionary column - cudf::test::dictionary_column_wrapper wrapper( - data.begin(), data.end() - ); - auto column = wrapper.release(); - - std::size_t exp = column->alloc_size(); - std::size_t est = rapidsmpf::estimated_memory_usage(column->view(), stream); - - EXPECT_EQ(exp, est); - } -} diff --git a/cpp/tests/test_host_buffer.cpp b/cpp/tests/test_host_buffer.cpp index cb2adcd0b..ccd9e7eae 100644 --- a/cpp/tests/test_host_buffer.cpp +++ b/cpp/tests/test_host_buffer.cpp @@ -296,7 +296,7 @@ std::size_t discover_pinned_pool_actual_size( TEST(PinnedResource, max_pool_size_limit) { // Ensure CUDA device context is initialized (required for pinned memory pools). RAPIDSMPF_CUDA_TRY(cudaFree(nullptr)); - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_view{}; // Create a PinnedMemoryResource with max pool size of 1 MiB; driver may round up. auto pinned_mr = rapidsmpf::PinnedMemoryResource::make_if_available( diff --git a/cpp/tests/test_partition.cpp b/cpp/tests/test_partition.cpp index b4e68a247..2100890f7 100644 --- a/cpp/tests/test_partition.cpp +++ b/cpp/tests/test_partition.cpp @@ -9,13 +9,13 @@ #include #include +#include +#include #include #include #include #include -#include -#include #include #include #include @@ -25,6 +25,7 @@ #include "utils.hpp" using namespace rapidsmpf; +using namespace cudf_streaming::integrations; class NumOfPartitions : public cudf::test::BaseFixtureWithParam> {}; @@ -49,7 +50,7 @@ TEST_P(NumOfPartitions, partition_and_pack) { cudf::table expect = random_table_with_index(seed, static_cast(num_rows), 0, 10); - auto chunks = rapidsmpf::partition_and_pack( + auto chunks = cudf_streaming::integrations::partition_and_pack( expect, {1}, num_partitions, hash_fn, seed, stream, br.get() ); @@ -60,8 +61,9 @@ TEST_P(NumOfPartitions, partition_and_pack) { } EXPECT_EQ(chunks_vector.size(), num_partitions); - auto result = - rapidsmpf::unpack_and_concat(std::move(chunks_vector), stream, br.get()); + auto result = cudf_streaming::integrations::unpack_and_concat( + std::move(chunks_vector), stream, br.get() + ); // Compare the input table with the result. We ignore the row order by // sorting by their index (first column). @@ -82,7 +84,8 @@ TEST_P(NumOfPartitions, split_and_pack) { splits.emplace_back(i * num_rows / num_partitions); } - auto chunks = rapidsmpf::split_and_pack(expect, splits, stream, br.get()); + auto chunks = + cudf_streaming::integrations::split_and_pack(expect, splits, stream, br.get()); // Convert to a vector (restoring the original order). std::vector chunks_vector; @@ -91,8 +94,9 @@ TEST_P(NumOfPartitions, split_and_pack) { } EXPECT_EQ(chunks_vector.size(), num_partitions); - auto result = - rapidsmpf::unpack_and_concat(std::move(chunks_vector), stream, br.get()); + auto result = cudf_streaming::integrations::unpack_and_concat( + std::move(chunks_vector), stream, br.get() + ); // Compare the input table with the result. CUDF_TEST_EXPECT_TABLES_EQUIVALENT(expect, *result); @@ -118,13 +122,16 @@ TEST_F(SpillingTest, SpillUnspillRoundtripPreservesDataAndMetadata) { input.push_back(create_packed_data(metadata, payload, stream, br.get())); // Device -> Device (moves data) - auto on_gpu = unspill_partitions(std::move(input), br.get(), AllowOverbooking::YES); + auto on_gpu = cudf_streaming::integrations::unspill_partitions( + std::move(input), br.get(), AllowOverbooking::YES + ); ASSERT_EQ(on_gpu.size(), 1); EXPECT_EQ(on_gpu[0].data->mem_type(), rapidsmpf::MemoryType::DEVICE); EXPECT_EQ(*on_gpu[0].metadata, metadata); // Device -> Host - auto back_on_host = spill_partitions(std::move(on_gpu), br.get()); + auto back_on_host = + cudf_streaming::integrations::spill_partitions(std::move(on_gpu), br.get()); ASSERT_EQ(back_on_host.size(), 1); EXPECT_EQ(back_on_host[0].data->mem_type(), rapidsmpf::MemoryType::HOST); EXPECT_EQ(*back_on_host[0].metadata, metadata); diff --git a/cpp/tests/test_progress_thread.cpp b/cpp/tests/test_progress_thread.cpp index 1435fe2ec..68ecceea3 100644 --- a/cpp/tests/test_progress_thread.cpp +++ b/cpp/tests/test_progress_thread.cpp @@ -8,8 +8,6 @@ #include #include -#include - #include #include @@ -17,8 +15,8 @@ using rapidsmpf::ProgressThread; -class ProgressThreadEvents - : public cudf::test::BaseFixtureWithParam> {}; +class ProgressThreadEvents : public ::testing::TestWithParam> { +}; // test different `num_threads` and `num_functions`. INSTANTIATE_TEST_SUITE_P( diff --git a/cpp/tests/test_rmm_resource_adaptor.cpp b/cpp/tests/test_rmm_resource_adaptor.cpp index 56bccef4b..54508cce0 100644 --- a/cpp/tests/test_rmm_resource_adaptor.cpp +++ b/cpp/tests/test_rmm_resource_adaptor.cpp @@ -17,6 +17,7 @@ #include #include #include +#include #include #include @@ -220,7 +221,7 @@ TEST(ScopedMemoryRecord, AddScopeMergesSiblingScopesCorrectly) { } TEST(RmmResourceAdaptor, EmptyScopedMemoryRecord) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; mr.begin_scoped_memory_record(); auto scope = mr.end_scoped_memory_record(); @@ -231,7 +232,7 @@ TEST(RmmResourceAdaptor, EmptyScopedMemoryRecord) { } TEST(RmmResourceAdaptorScopedMemory, SingleScopedAllocationTracksCorrectly) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; mr.begin_scoped_memory_record(); void* p = mr.allocate_sync(1_MiB); @@ -246,7 +247,7 @@ TEST(RmmResourceAdaptorScopedMemory, SingleScopedAllocationTracksCorrectly) { } TEST(RmmResourceAdaptorScopedMemory, NestedScopedAllocationsMerged) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; mr.begin_scoped_memory_record(); // Outer @@ -272,7 +273,7 @@ TEST(RmmResourceAdaptorScopedMemory, NestedScopedAllocationsMerged) { } TEST(RmmResourceAdaptorScopedMemory, NestedScopedTracksAllocsAndDeallocs) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; mr.begin_scoped_memory_record(); // Outer @@ -299,7 +300,7 @@ TEST(RmmResourceAdaptorScopedMemory, NestedScopedTracksAllocsAndDeallocs) { } TEST(RmmResourceAdaptorScopedMemory, NestedDeallocationYieldsNegativeStats) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; // Allocate in outer scope mr.begin_scoped_memory_record(); // Outer @@ -327,7 +328,7 @@ TEST(RmmResourceAdaptorScopedMemory, MultiThreadedScopedAllocations) { constexpr int num_allocs_per_thread = 8; constexpr std::size_t alloc_size = 1_MiB; - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; std::vector threads; std::vector> allocations(num_threads); std::vector records(num_threads); @@ -396,7 +397,7 @@ TEST(RmmResourceAdaptorScopedMemory, CrossThreadNestedScopesNotMerged) { constexpr std::size_t outer_alloc_size = 1_MiB; constexpr std::size_t inner_alloc_size = 2_MiB; - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; void* outer_alloc = nullptr; void* inner_alloc = nullptr; rapidsmpf::ScopedMemoryRecord inner_record; diff --git a/cpp/tests/test_shuffler.cpp b/cpp/tests/test_shuffler.cpp index 024820a90..cd3dc9d21 100644 --- a/cpp/tests/test_shuffler.cpp +++ b/cpp/tests/test_shuffler.cpp @@ -12,12 +12,12 @@ #include +#include #include #include #include #include -#include #include #include #include @@ -153,7 +153,7 @@ void test_shuffler( // 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] = rapidsmpf::partition_and_split( + auto [expect_partitions, owner] = cudf_streaming::integrations::partition_and_split( full_input_table, {1}, static_cast(total_num_partitions), @@ -185,7 +185,7 @@ void test_shuffler( // 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 = rapidsmpf::partition_and_pack( + auto packed_chunks = cudf_streaming::integrations::partition_and_pack( slice, {1}, static_cast(total_num_partitions), @@ -206,8 +206,8 @@ void test_shuffler( EXPECT_NO_THROW(shuffler.wait(wait_timeout)); for (auto finished_partition : shuffler.local_partitions()) { auto packed_chunks = shuffler.extract(finished_partition); - auto result = rapidsmpf::unpack_and_concat( - rapidsmpf::unspill_partitions( + auto result = cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::unspill_partitions( std::move(packed_chunks), br, rapidsmpf::AllowOverbooking::YES ), stream, @@ -433,7 +433,7 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { br.get() ); cudf::table input_table = random_table_with_index(seed, 1000, 0, 10); - auto input_chunks = rapidsmpf::partition_and_pack( + auto input_chunks = cudf_streaming::integrations::partition_and_pack( input_table, {1}, total_num_partitions, @@ -456,9 +456,10 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { { // Now extract triggers spilling of the partition not being extracted. - std::vector output_chunks = rapidsmpf::unspill_partitions( - shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES - ); + 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. @@ -469,13 +470,15 @@ TEST(Shuffler, SpillOnInsertAndExtraction) { } // Extract and unspill both partitions. - std::vector out0 = rapidsmpf::unspill_partitions( - shuffler.extract(0), br.get(), rapidsmpf::AllowOverbooking::YES - ); + 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 = rapidsmpf::unspill_partitions( - shuffler.extract(1), br.get(), rapidsmpf::AllowOverbooking::YES - ); + 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. @@ -849,7 +852,7 @@ TEST(Shuffler, concurrent_wait) { 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] = rapidsmpf::partition_and_split( + auto [expected, owner] = cudf_streaming::integrations::partition_and_split( full_input, {1}, static_cast(total_num_partitions), @@ -876,7 +879,7 @@ TEST(Shuffler, concurrent_wait) { auto slice = cudf::slice(full_input, {row_offset, row_end}).at(0); insert_futures.push_back(std::async(std::launch::async, [&, slice] { shuffler.insert( - rapidsmpf::partition_and_pack( + cudf_streaming::integrations::partition_and_pack( slice, {1}, static_cast(total_num_partitions), @@ -901,8 +904,8 @@ TEST(Shuffler, concurrent_wait) { futures.push_back(std::async(std::launch::async, [&, pid] { EXPECT_NO_THROW(shuffler.wait(wait_timeout)); auto chunks = shuffler.extract(pid); - auto result = rapidsmpf::unpack_and_concat( - rapidsmpf::unspill_partitions( + auto result = cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::unspill_partitions( std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, @@ -966,7 +969,7 @@ TEST(Shuffler, opid_reuse) { row_end = full_input.num_rows(); } auto slice = cudf::slice(full_input, {row_offset, row_end}).at(0); - auto packed = rapidsmpf::partition_and_pack( + auto packed = cudf_streaming::integrations::partition_and_pack( slice, {1}, static_cast(total_num_partitions), @@ -985,7 +988,7 @@ TEST(Shuffler, opid_reuse) { 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] = rapidsmpf::partition_and_split( + auto [expected, owner] = cudf_streaming::integrations::partition_and_split( full_input, {1}, static_cast(total_num_partitions), @@ -997,8 +1000,8 @@ TEST(Shuffler, opid_reuse) { ); for (auto pid : shuffler.local_partitions()) { auto chunks = shuffler.extract(pid); - auto result = rapidsmpf::unpack_and_concat( - rapidsmpf::unspill_partitions( + auto result = cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::unspill_partitions( std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, @@ -1065,7 +1068,7 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { if (rapidsmpf::shuffler::Shuffler::round_robin(comm, 0, total_num_partitions) == comm->rank()) { - auto packed = rapidsmpf::partition_and_pack( + auto packed = cudf_streaming::integrations::partition_and_pack( full_input, {1}, static_cast(total_num_partitions), @@ -1082,7 +1085,7 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { 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] = rapidsmpf::partition_and_split( + auto [expected, owner] = cudf_streaming::integrations::partition_and_split( full_input, {1}, static_cast(total_num_partitions), @@ -1094,8 +1097,8 @@ TEST(Shuffler, opid_reuse_with_empty_partitions) { ); for (auto pid : shuffler.local_partitions()) { auto chunks = shuffler.extract(pid); - auto result = rapidsmpf::unpack_and_concat( - rapidsmpf::unspill_partitions( + auto result = cudf_streaming::integrations::unpack_and_concat( + cudf_streaming::integrations::unspill_partitions( std::move(chunks), br.get(), rapidsmpf::AllowOverbooking::YES ), stream, diff --git a/cpp/tests/test_shuffler_many_streams.cpp b/cpp/tests/test_shuffler_many_streams.cpp index 43c3f6402..1043611c1 100644 --- a/cpp/tests/test_shuffler_many_streams.cpp +++ b/cpp/tests/test_shuffler_many_streams.cpp @@ -9,13 +9,13 @@ #include #include +#include #include #include #include #include #include -#include #include #include "environment.hpp" diff --git a/cpp/tests/test_spill_manager.cpp b/cpp/tests/test_spill_manager.cpp index 401b26f2f..b3abdab17 100644 --- a/cpp/tests/test_spill_manager.cpp +++ b/cpp/tests/test_spill_manager.cpp @@ -6,11 +6,8 @@ #include -#include -#include -#include -#include #include +#include #include #include @@ -29,7 +26,7 @@ TEST(SpillManager, SpillFunction) { // currently configured limit. std::int64_t mem_available = 10_KiB; auto br = BufferResource::create( - cudf::get_current_device_resource_ref(), + rmm::mr::get_current_device_resource_ref(), rapidsmpf::PinnedMemoryResource::Disabled, {{MemoryType::DEVICE, mem_available}} ); diff --git a/cpp/tests/test_statistics.cpp b/cpp/tests/test_statistics.cpp index 5b2ecc2f1..1fa8bb208 100644 --- a/cpp/tests/test_statistics.cpp +++ b/cpp/tests/test_statistics.cpp @@ -12,7 +12,9 @@ #include #include -#include +#include +#include +#include #include #include @@ -178,10 +180,10 @@ TEST_F(StatisticsTest, ReportSorting) { } TEST_F(StatisticsTest, MemoryProfiler) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; auto pinned_mr = rapidsmpf::PinnedMemoryResource::make_if_available(); auto stats = rapidsmpf::Statistics::create(); - auto stream = cudf::get_default_stream(); + auto stream = rmm::cuda_stream_view{}; // Outer scope { @@ -273,7 +275,7 @@ TEST_F(StatisticsTest, MemoryProfiler) { } TEST_F(StatisticsTest, MemoryProfilerDisabled) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; auto stats = rapidsmpf::Statistics::disabled(); { auto const& records = stats->get_memory_records(); @@ -309,7 +311,7 @@ TEST_F(StatisticsTest, MemoryProfilerDisabled) { // the dtor early-returned and the frame stayed on the stack. // 3. A follow-up recorder works correctly against the balanced stack. TEST_F(StatisticsTest, MemoryProfilerToggledMidScope) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; auto stats = rapidsmpf::Statistics::create(); { @@ -335,7 +337,7 @@ TEST_F(StatisticsTest, MemoryProfilerToggledMidScope) { } TEST_F(StatisticsTest, MemoryProfilerMacro) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; auto stats = rapidsmpf::Statistics::create(); { RAPIDSMPF_MEMORY_PROFILE(stats, mr); @@ -350,7 +352,7 @@ TEST_F(StatisticsTest, MemoryProfilerMacro) { } TEST_F(StatisticsTest, MemoryProfilerMacroDisabled) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; auto stats = rapidsmpf::Statistics::disabled(); { RAPIDSMPF_MEMORY_PROFILE(stats, mr); @@ -390,7 +392,7 @@ TEST_F(StatisticsTest, InvalidStatNames) { } TEST_F(StatisticsTest, InvalidMemoryRecordNames) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; auto stats = rapidsmpf::Statistics::create(); std::ignore = stats->create_memory_recorder(mr, "bad\"name"); std::ostringstream ss; @@ -398,7 +400,7 @@ TEST_F(StatisticsTest, InvalidMemoryRecordNames) { } TEST_F(StatisticsTest, JsonMemoryRecords) { - rapidsmpf::RmmResourceAdaptor mr{cudf::get_current_device_resource_ref()}; + rapidsmpf::RmmResourceAdaptor mr{rmm::mr::get_current_device_resource_ref()}; auto stats = rapidsmpf::Statistics::create(); { auto rec = stats->create_memory_recorder(mr, "alloc"); diff --git a/cpp/tests/utils.hpp b/cpp/tests/utils.hpp index 0aa23ff85..bc51e460b 100644 --- a/cpp/tests/utils.hpp +++ b/cpp/tests/utils.hpp @@ -4,10 +4,12 @@ */ #pragma once +#include #include #include #include #include +#include #include #include #include @@ -21,10 +23,14 @@ #include +#ifdef RAPIDSMPF_HAVE_CUDF +#include #include #include #include #include +#endif // RAPIDSMPF_HAVE_CUDF + #include #include @@ -99,15 +105,6 @@ template } template -[[nodiscard]] inline std::unique_ptr iota_column( - std::size_t nrows, T start = 0 -) { - std::vector vec = iota_vector(nrows, start); - cudf::test::fixed_width_column_wrapper ret(vec.begin(), vec.end()); - return ret.release(); -} - -template [[nodiscard]] inline std::vector random_vector( std::int64_t seed, std::size_t nelem, @@ -121,6 +118,16 @@ template return ret; } +#ifdef RAPIDSMPF_HAVE_CUDF +template +[[nodiscard]] inline std::unique_ptr iota_column( + std::size_t nrows, T start = 0 +) { + std::vector vec = iota_vector(nrows, start); + cudf::test::fixed_width_column_wrapper ret(vec.begin(), vec.end()); + return ret.release(); +} + [[nodiscard]] inline std::unique_ptr random_column( std::int64_t seed, std::size_t nrows, @@ -160,6 +167,7 @@ template ) { return sort_table(table->view(), column_indices); } +#endif // RAPIDSMPF_HAVE_CUDF /// @brief Create a PackedData object from a host buffer [[nodiscard]] inline rapidsmpf::PackedData create_packed_data( diff --git a/dependencies.yaml b/dependencies.yaml index f59e9db86..b3a35bbd3 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -19,10 +19,8 @@ files: - clang_tidy - cuda - cuda_version - - depends_on_libcudf - depends_on_librmm - depends_on_libucxx - - depends_on_pylibcudf - depends_on_rmm - depends_on_ucxx - depends_on_ray @@ -46,10 +44,8 @@ files: - clang_tidy - cuda - cuda_version - - depends_on_libcudf - depends_on_librmm - depends_on_libucxx - - depends_on_pylibcudf - depends_on_rmm - depends_on_ucxx - py_version @@ -71,7 +67,6 @@ files: output: none includes: - cuda_version - - depends_on_libcudf - depends_on_librapidsmpf - depends_on_librmm - depends_on_rapidsmpf @@ -80,6 +75,8 @@ files: - py_version - run_rapidsmpf - test_python + - depends_on_pylibcudf + - depends_on_cudf_streaming checks: output: none includes: @@ -90,7 +87,6 @@ files: includes: - build-cpp - build-universal - - depends_on_libcudf - depends_on_librmm - depends_on_libucxx - clang_tidy @@ -125,7 +121,6 @@ files: includes: - build-universal - build-nvml - - depends_on_libcudf - depends_on_librmm - depends_on_libucxx py_run_librapidsmpf: @@ -134,7 +129,6 @@ files: extras: table: project includes: - - depends_on_libcudf - depends_on_librmm - depends_on_libucxx py_build_rapidsmpf: @@ -153,10 +147,8 @@ files: includes: - build-universal - build-python - - depends_on_libcudf - depends_on_librapidsmpf - depends_on_librmm - - depends_on_pylibcudf - depends_on_rmm - depends_on_ucxx py_run_rapidsmpf: @@ -167,7 +159,6 @@ files: includes: - run_rapidsmpf - depends_on_rmm - - depends_on_pylibcudf - depends_on_librapidsmpf - depends_on_ucxx py_test_rapidsmpf: @@ -179,6 +170,8 @@ files: includes: - depends_on_ray - test_python + - depends_on_pylibcudf + - depends_on_cudf_streaming channels: - rapidsai-nightly - rapidsai @@ -397,6 +390,7 @@ dependencies: - libnuma - openmpi >=5.0 # See - valgrind + - libcudf-streaming==26.8.*,>=0.0.0a0 test_cpp_ndsh: common: - output_types: conda @@ -518,56 +512,6 @@ dependencies: - sphinx>=8.1.0 - sphinx-autobuild - sphinx-copybutton - depends_on_libcudf: - common: - - output_types: conda - packages: - - &libcudf_unsuffixed libcudf==26.8.*,>=0.0.0a0 - - output_types: requirements - packages: - # pip recognizes the index as a global option for the requirements.txt file - - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple - specific: - - output_types: [requirements, pyproject] - matrices: - - matrix: - cuda: "12.*" - cuda_suffixed: "true" - packages: - - libcudf-cu12==26.8.*,>=0.0.0a0 - - matrix: - cuda: "13.*" - cuda_suffixed: "true" - packages: - - libcudf-cu13==26.8.*,>=0.0.0a0 - - matrix: - packages: - - *libcudf_unsuffixed - depends_on_pylibcudf: - common: - - output_types: conda - packages: - - &pylibcudf_unsuffixed pylibcudf==26.8.*,>=0.0.0a0 - - output_types: requirements - packages: - # pip recognizes the index as a global option for the requirements.txt file - - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple - specific: - - output_types: [requirements, pyproject] - matrices: - - matrix: - cuda: "12.*" - cuda_suffixed: "true" - packages: - - pylibcudf-cu12==26.8.*,>=0.0.0a0 - - matrix: - cuda: "13.*" - cuda_suffixed: "true" - packages: - - pylibcudf-cu13==26.8.*,>=0.0.0a0 - - matrix: - packages: - - *pylibcudf_unsuffixed depends_on_librmm: common: - output_types: conda @@ -676,3 +620,33 @@ dependencies: - output_types: [requirements, pyproject] packages: - ray>=2.55.1 + depends_on_pylibcudf: + common: + - output_types: conda + packages: + - &pylibcudf_unsuffixed pylibcudf==26.8.*,>=0.0.0a0 + - output_types: requirements + packages: + # pip recognizes the index as a global option for the requirements.txt file + - --extra-index-url=https://pypi.anaconda.org/rapidsai-wheels-nightly/simple + specific: + - output_types: [requirements, pyproject] + matrices: + - matrix: + cuda: "12.*" + cuda_suffixed: "true" + packages: + - pylibcudf-cu12==26.8.*,>=0.0.0a0 + - matrix: + cuda: "13.*" + cuda_suffixed: "true" + packages: + - pylibcudf-cu13==26.8.*,>=0.0.0a0 + - matrix: + packages: + - *pylibcudf_unsuffixed + depends_on_cudf_streaming: + common: + - output_types: conda + packages: + - cudf-streaming==26.8.*,>=0.0.0a0 diff --git a/docs/source/background/streaming-engine.md b/docs/source/background/streaming-engine.md index bc4a22600..500b6bea7 100644 --- a/docs/source/background/streaming-engine.md +++ b/docs/source/background/streaming-engine.md @@ -56,9 +56,8 @@ operation in the query plan. Parallelism is obtained by using a multi-threaded executor to handle the concurrent actors that thus result. For use with data processing engines, we provide a number of utility actors -that layer a streaming (out of core) execution model over the -GPU-accelerated [libcudf](https://docs.rapids.ai/api/libcudf/stable/) -library. +that layer a streaming (out of core) execution model over +GPU-accelerated libraries. ``` +------+ +--------+ +--------+ +------+ diff --git a/docs/source/conf.py b/docs/source/conf.py index 9b0cb9345..839a3a2fb 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -79,7 +79,6 @@ intersphinx_mapping = { "python": ("https://docs.python.org/3", None), - "cudf": ("https://docs.rapids.ai/api/cudf/stable/", None), "ray": ("https://docs.ray.io/en/latest/", None), } diff --git a/docs/source/getting-started.md b/docs/source/getting-started.md index 35058af55..7e363eccb 100644 --- a/docs/source/getting-started.md +++ b/docs/source/getting-started.md @@ -1,7 +1,7 @@ # Getting Started Building rapidsmpf from source is recommended when running nightly/upstream versions, -since dependencies on non-ABI-stable libraries (e.g., pylibcudf) could cause temporary +since dependencies on non-ABI-stable libraries could cause temporary breakage leading to issues such as segmentation faults. Stable versions can be installed from conda or pip packages. diff --git a/docs/source/index.md b/docs/source/index.md index cac46d6a9..fba58f6d3 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -1,7 +1,7 @@ # RapidsMPF documentation Building high-performance GPU pipelines is hard. Each stage must move data efficiently between GPUs and processes, synchronize work, and manage limited device memory. -RapidsMPF, including its nascent Streaming Engine and Out-of-Core (OOC) shuffle, provides a unified framework for asynchronous, multi-GPU pipelines using simple streaming primitives — {term}`Channel`s, {term}`Actor`s, and {term}`Message`s — built on RAPIDS components: [rmm](https://docs.rapids.ai/api/rmm/nightly), [libcudf](https://docs.rapids.ai/api/libcudf/nightly/), and [ucxx](https://docs.rapids.ai/api/ucxx/nightly/). +RapidsMPF, including its nascent Streaming Engine and Out-of-Core (OOC) shuffle, provides a unified framework for asynchronous, multi-GPU pipelines using simple streaming primitives — {term}`Channel`s, {term}`Actor`s, and {term}`Message`s — built on RAPIDS components: [rmm](https://docs.rapids.ai/api/rmm/nightly) and [ucxx](https://docs.rapids.ai/api/ucxx/nightly/). RapidsMPF's design leverages Explicit Data Parallelism (SPMD-style coordination) combined with a local CSP-style streaming model, enabling the engine to overlap I/O, computation, and communication. This makes it possible to handle out-of-core processing efficiently (via {term}`Spilling`) and integrate seamlessly with frontend query engines such as Polars. The result is clean, composable, and scalable GPU streaming — from single-node prototypes to large-scale, multi-GPU deployments. See the {doc}`glossary` for definitions of key concepts. diff --git a/docs/source/python/api.md b/docs/source/python/api.md index ff9334dfc..488a1c7e8 100644 --- a/docs/source/python/api.md +++ b/docs/source/python/api.md @@ -21,20 +21,6 @@ libraries. :members: ``` -### cuDF - -```{eval-rst} -.. automodule:: rapidsmpf.integrations.cudf - :members: -``` - -#### Partition - -```{eval-rst} -.. automodule:: rapidsmpf.integrations.cudf.partition - :members: -``` - ## Shuffler ```{eval-rst} diff --git a/docs/source/python/index.md b/docs/source/python/index.md index ff109cfea..7a2fef201 100644 --- a/docs/source/python/index.md +++ b/docs/source/python/index.md @@ -17,8 +17,6 @@ frameworks. The Python API includes ready-to-use integrations with: - **Ray** (`rapidsmpf.integrations.ray`) — use RapidsMPF within Ray tasks and actors. -- **cuDF** (`rapidsmpf.integrations.cudf`) — partition and pack/unpack cuDF tables for - use with the Shuffler. ```{toctree} --- diff --git a/pyproject.toml b/pyproject.toml index 53e5ac52f..8d22591b5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -153,7 +153,7 @@ required-imports = ["from __future__ import annotations"] [tool.ruff.lint.isort.sections] rapidsmpf = ["rapidsmpf"] -rapids = ["rmm", "pylibcudf", "cudf"] +rapids = ["rmm"] [tool.ruff.format] docstring-code-format = true diff --git a/python/librapidsmpf/librapidsmpf/load.py b/python/librapidsmpf/librapidsmpf/load.py index 2888b311f..36c56448f 100644 --- a/python/librapidsmpf/librapidsmpf/load.py +++ b/python/librapidsmpf/librapidsmpf/load.py @@ -1,5 +1,5 @@ # noqa: D100 -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. # SPDX-License-Identifier: Apache-2.0 # from __future__ import annotations @@ -43,14 +43,12 @@ def load_library() -> list: try: # these libraries must be loaded before librapidsmpf because # librapidsmpf references their symbols - import libcudf import librmm import libucxx import rapids_logger librmm.load_library() rapids_logger.load_library() - libcudf.load_library() libucxx.load_library() except ModuleNotFoundError: # 'librapidsmpf' has a runtime dependency on 'librmm' et al.. However, diff --git a/python/librapidsmpf/pyproject.toml b/python/librapidsmpf/pyproject.toml index 5fbf80edc..af15c62ed 100644 --- a/python/librapidsmpf/pyproject.toml +++ b/python/librapidsmpf/pyproject.toml @@ -19,7 +19,6 @@ authors = [ license = "Apache-2.0" requires-python = ">=3.11" dependencies = [ - "libcudf==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", "libucxx==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`. @@ -47,7 +46,6 @@ matrix-entry = "cuda_suffixed=true" requires = [ "cmake>=4.0", "cuda-toolkit[nvml]==13.*", - "libcudf==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", "libucxx==0.51.*,>=0.0.0a0", "ninja", diff --git a/python/rapidsmpf/pyproject.toml b/python/rapidsmpf/pyproject.toml index 2d38b632a..b9e8a6d59 100644 --- a/python/rapidsmpf/pyproject.toml +++ b/python/rapidsmpf/pyproject.toml @@ -21,7 +21,6 @@ requires-python = ">=3.11" dependencies = [ "cuda-python>=13.0.1,<14.0", "librapidsmpf==26.8.*,>=0.0.0a0", - "pylibcudf==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`. @@ -42,6 +41,7 @@ test = [ "cupy-cuda13x>=13.6.0,!=14.0.0,!=14.1.0", "numpy >=1.23,<3.0", "psutil", + "pylibcudf==26.8.*,>=0.0.0a0", "pytest", "ray>=2.55.1", ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`. @@ -56,11 +56,9 @@ matrix-entry = "cuda_suffixed=true" requires = [ "cmake>=4.0", "cython>=3.2.2", - "libcudf==26.8.*,>=0.0.0a0", "librapidsmpf==26.8.*,>=0.0.0a0", "librmm==26.8.*,>=0.0.0a0", "ninja", - "pylibcudf==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`. diff --git a/python/rapidsmpf/rapidsmpf/CMakeLists.txt b/python/rapidsmpf/rapidsmpf/CMakeLists.txt index ba3d7cd27..99ae4e67a 100644 --- a/python/rapidsmpf/rapidsmpf/CMakeLists.txt +++ b/python/rapidsmpf/rapidsmpf/CMakeLists.txt @@ -25,7 +25,6 @@ add_subdirectory(bootstrap) add_subdirectory(coll) add_subdirectory(memory) add_subdirectory(communicator) -add_subdirectory(integrations/cudf) add_subdirectory(rrun) add_subdirectory(utils) diff --git a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py b/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py index c57af8cce..7e2dc9201 100644 --- a/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py +++ b/python/rapidsmpf/rapidsmpf/benchmarks/streaming_benchmark.py @@ -10,11 +10,11 @@ from typing import TYPE_CHECKING import cupy as cp +import pylibcudf as plc from mpi4py import MPI +from pylibcudf.contiguous_split import pack -import pylibcudf as plc import rmm.mr -from pylibcudf.contiguous_split import pack from rmm.pylibrmm.stream import DEFAULT_STREAM import rapidsmpf.bootstrap @@ -159,7 +159,7 @@ def streaming_shuffle( # NOTE: This will require part_size amount of GPU memory. chunks: dict[int, PackedData] = {} for i in range(output_nparts): - chunks[i] = PackedData.from_cudf_packed_columns( + chunks[i] = PackedData.from_cudf_packed_columns( # type: ignore[attr-defined] pack(dummy_table), DEFAULT_STREAM, br ) diff --git a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py index 149217c05..d907b1023 100644 --- a/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/bulk_mpi_shuffle.py @@ -10,20 +10,20 @@ from pathlib import Path from typing import TYPE_CHECKING +import pylibcudf as plc +from cudf_streaming.integrations.partition import ( + partition_and_pack, + unpack_and_concat, + unspill_partitions, +) from mpi4py import MPI -import pylibcudf as plc import rmm.mr from rmm.pylibrmm.stream import DEFAULT_STREAM import rapidsmpf.bootstrap import rapidsmpf.communicator.mpi from rapidsmpf.config import Options, get_environment_variables -from rapidsmpf.integrations.cudf.partition import ( - partition_and_pack, - unpack_and_concat, - unspill_partitions, -) from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.progress_thread import ProgressThread diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py index dcdf5f3fd..4b63df070 100644 --- a/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py +++ b/python/rapidsmpf/rapidsmpf/examples/ray/bulk_ray_shuffle.py @@ -11,16 +11,16 @@ from pathlib import Path from typing import TYPE_CHECKING -import ray - import pylibcudf as plc -import rmm.mr - -from rapidsmpf.integrations.cudf.partition import ( +import ray +from cudf_streaming.integrations.partition import ( partition_and_pack, unpack_and_concat, unspill_partitions, ) + +import rmm.mr + from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster from rapidsmpf.memory.buffer import MemoryType from rapidsmpf.memory.buffer_resource import BufferResource diff --git a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py b/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py index 526b1c2f6..34c3dc0fc 100644 --- a/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py +++ b/python/rapidsmpf/rapidsmpf/examples/ray/ray_shuffle_example.py @@ -8,16 +8,16 @@ import math import numpy as np -import ray - import pylibcudf as plc -import rmm - -from rapidsmpf.integrations.cudf.partition import ( +import ray +from cudf_streaming.integrations.partition import ( partition_and_pack, unpack_and_concat, unspill_partitions, ) + +import rmm + from rapidsmpf.integrations.ray import RapidsMPFActor, setup_ray_ucxx_cluster from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.shuffler import Shuffler diff --git a/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py b/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py index 0f10b9735..483a43c73 100644 --- a/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py +++ b/python/rapidsmpf/rapidsmpf/examples/streaming/basic_example.py @@ -7,6 +7,8 @@ from typing import TYPE_CHECKING import pylibcudf +from cudf_streaming.streaming.table_chunk import TableChunk + import rmm.mr from rmm.pylibrmm.stream import DEFAULT_STREAM @@ -23,7 +25,6 @@ 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 rapidsmpf.streaming.cudf.table_chunk import TableChunk if TYPE_CHECKING: from collections.abc import Awaitable diff --git a/python/rapidsmpf/rapidsmpf/integrations/cudf/CMakeLists.txt b/python/rapidsmpf/rapidsmpf/integrations/cudf/CMakeLists.txt deleted file mode 100644 index 709bcf881..000000000 --- a/python/rapidsmpf/rapidsmpf/integrations/cudf/CMakeLists.txt +++ /dev/null @@ -1,14 +0,0 @@ -# ================================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ================================================================================= - -set(cython_modules partition.pyx) - -rapids_cython_create_modules( - CXX - SOURCE_FILES "${cython_modules}" - LINKED_LIBRARIES rapidsmpf::rapidsmpf maybe_asan -) diff --git a/python/rapidsmpf/rapidsmpf/integrations/cudf/__init__.pxd b/python/rapidsmpf/rapidsmpf/integrations/cudf/__init__.pxd deleted file mode 100644 index a9c2f923c..000000000 --- a/python/rapidsmpf/rapidsmpf/integrations/cudf/__init__.pxd +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 diff --git a/python/rapidsmpf/rapidsmpf/integrations/cudf/__init__.py b/python/rapidsmpf/rapidsmpf/integrations/cudf/__init__.py deleted file mode 100644 index 8fb40eedf..000000000 --- a/python/rapidsmpf/rapidsmpf/integrations/cudf/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Collection of cuDF specific functions.""" diff --git a/python/rapidsmpf/rapidsmpf/integrations/cudf/partition.pxd b/python/rapidsmpf/rapidsmpf/integrations/cudf/partition.pxd deleted file mode 100644 index a9c2f923c..000000000 --- a/python/rapidsmpf/rapidsmpf/integrations/cudf/partition.pxd +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 diff --git a/python/rapidsmpf/rapidsmpf/integrations/cudf/partition.pyi b/python/rapidsmpf/rapidsmpf/integrations/cudf/partition.pyi deleted file mode 100644 index 55c0d315f..000000000 --- a/python/rapidsmpf/rapidsmpf/integrations/cudf/partition.pyi +++ /dev/null @@ -1,41 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from collections.abc import Iterable - -from pylibcudf.table import Table -from rmm.pylibrmm.stream import Stream - -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.memory.packed_data import PackedData - -def partition_and_pack( - table: Table, - columns_to_hash: Iterable[int], - num_partitions: int, - stream: Stream, - br: BufferResource, -) -> dict[int, PackedData]: ... -def split_and_pack( - table: Table, - splits: Iterable[int], - stream: Stream, - br: BufferResource, -) -> dict[int, PackedData]: ... -def unpack_and_concat( - partitions: Iterable[PackedData], - stream: Stream, - br: BufferResource, -) -> Table: ... -def spill_partitions( - partitions: Iterable[PackedData], - *, - br: BufferResource, -) -> list[PackedData]: ... -def unspill_partitions( - partitions: Iterable[PackedData], - *, - br: BufferResource, - allow_overbooking: bool, -) -> list[PackedData]: ... diff --git a/python/rapidsmpf/rapidsmpf/integrations/cudf/partition.pyx b/python/rapidsmpf/rapidsmpf/integrations/cudf/partition.pyx deleted file mode 100644 index c74b23158..000000000 --- a/python/rapidsmpf/rapidsmpf/integrations/cudf/partition.pyx +++ /dev/null @@ -1,362 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Partitioning of cuDF tables.""" - -from cython.operator cimport dereference as deref -from cython.operator cimport postincrement -from libc.stdint cimport uint32_t -from libcpp cimport bool as bool_t -from libcpp.memory cimport make_unique, unique_ptr -from libcpp.unordered_map cimport unordered_map -from libcpp.utility cimport move -from libcpp.vector cimport vector -from pylibcudf.libcudf.table.table cimport table as cpp_table -from pylibcudf.libcudf.table.table_view cimport table_view -from pylibcudf.libcudf.types cimport size_type -from pylibcudf.table cimport Table -from rmm.librmm.cuda_stream_view cimport cuda_stream_view -from rmm.pylibrmm.stream cimport Stream - -from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.memory.buffer_resource cimport (AllowOverbooking, - BufferResource, - cpp_BufferResource) -from rapidsmpf.memory.packed_data cimport (PackedData, cpp_PackedData, - packed_data_vector_to_list) - - -cdef extern from "" nogil: - int cpp_HASH_MURMUR3"cudf::hash_id::HASH_MURMUR3" - uint32_t cpp_DEFAULT_HASH_SEED"cudf::DEFAULT_HASH_SEED", - - cdef unordered_map[uint32_t, cpp_PackedData] cpp_partition_and_pack \ - "rapidsmpf::partition_and_pack"( - const table_view& table, - const vector[size_type] &columns_to_hash, - int num_partitions, - int hash_function, - uint32_t seed, - cuda_stream_view stream, - cpp_BufferResource* br, - ) except +ex_handler - - cdef unordered_map[uint32_t, cpp_PackedData] cpp_split_and_pack \ - "rapidsmpf::split_and_pack"( - const table_view& table, - const vector[size_type] &splits, - cuda_stream_view stream, - cpp_BufferResource* br, - ) except +ex_handler - - -def partition_and_pack( - Table table not None, - columns_to_hash, - int num_partitions, - Stream stream not None, - BufferResource br not None, -): - """ - Partition rows from the input table into multiple packed (serialized) tables. - - Parameters - ---------- - table - The input table to partition. - columns_to_hash - Indices of the input columns to use for hashing. - num_partitions - The number of partitions to create. - stream - The CUDA stream used for memory operations. - br - Buffer resource for memory allocations. - - Returns - ------- - A dictionary where the keys are partition IDs and the values are packed tables. - - Raises - ------ - IndexError - If any index in ``columns_to_hash`` is invalid. - - See Also - -------- - rapidsmpf.integrations.cudf.partition.unpack_and_concat - pylibcudf.partitioning.hash_partition - pylibcudf.contiguous_split.pack - rapidsmpf.integrations.cudf.partition.split_and_pack - """ - cdef cuda_stream_view _stream = stream.view() - cdef cpp_BufferResource* _br = br.ptr() - cdef vector[size_type] _columns_to_hash = tuple(columns_to_hash) - cdef unordered_map[uint32_t, cpp_PackedData] _ret - cdef table_view tbl = table.view() - with nogil: - _ret = cpp_partition_and_pack( - tbl, - _columns_to_hash, - num_partitions, - cpp_HASH_MURMUR3, - cpp_DEFAULT_HASH_SEED, - _stream, - _br, - ) - ret = {} - cdef unordered_map[uint32_t, cpp_PackedData].iterator it = _ret.begin() - while(it != _ret.end()): - ret[deref(it).first] = PackedData.from_librapidsmpf( - make_unique[cpp_PackedData](move(deref(it).second)), - br, - ) - postincrement(it) - return ret - - -def split_and_pack( - Table table not None, - splits, - Stream stream not None, - BufferResource br not None, -): - """ - Split rows from the input table into multiple packed (serialized) tables. - - Parameters - ---------- - table - The input table to split and pack. The table cannot be empty (the - split points would not be valid). - splits - The split points, one less than the number of result partitions. - stream - The CUDA stream used for memory operations. - br - Buffer resource for memory allocations. - - Returns - ------- - A map of partition IDs and their packed tables. - - Raises - ------ - IndexError - If the splits are out of range for ``[0, len(table)]``. - - See Also - -------- - rapidsmpf.integrations.cudf.partition.unpack_and_concat - pylibcudf.copying.split - rapidsmpf.integrations.cudf.partition.partition_and_pack - """ - cdef cuda_stream_view _stream = stream.view() - cdef cpp_BufferResource* _br = br.ptr() - cdef vector[size_type] _splits = tuple(splits) - cdef unordered_map[uint32_t, cpp_PackedData] _ret - cdef table_view tbl = table.view() - with nogil: - _ret = cpp_split_and_pack( - tbl, - _splits, - _stream, - _br, - ) - ret = {} - cdef unordered_map[uint32_t, cpp_PackedData].iterator it = _ret.begin() - while(it != _ret.end()): - ret[deref(it).first] = PackedData.from_librapidsmpf( - make_unique[cpp_PackedData](move(deref(it).second)), - br, - ) - postincrement(it) - return ret - - -cdef extern from "" nogil: - cdef unique_ptr[cpp_table] cpp_unpack_and_concat \ - "rapidsmpf::unpack_and_concat"( - vector[cpp_PackedData] partition, - cuda_stream_view stream, - cpp_BufferResource* br, - ) except +ex_handler - - -# Help function to convert an iterable of `PackedData` to a vector of -# `cpp_PackedData`. -cdef vector[cpp_PackedData] _partitions_py_to_cpp(partitions): - cdef vector[cpp_PackedData] ret - for part in partitions: - if not (part).c_obj: - raise ValueError("PackedData was empty") - ret.push_back(move(deref((part).c_obj))) - return move(ret) - - -def unpack_and_concat( - partitions, - Stream stream not None, - BufferResource br not None, -): - """ - Unpack (deserialize) input partitions and concatenate them into a single table. - - Empty partitions are ignored. - - The unpacking of each partition is stream-ordered on that partition's own CUDA - stream. The returned table is stream-ordered on the provided ``stream`` and - synchronized with the unpacking. - - Notes - ----- - The input partitions are released and left empty on return. - - Parameters - ---------- - partitions - Packed input tables (partitions). - stream - CUDA stream on which concatenation occurs and on which the resulting - table is ordered. - br - Buffer resource used for memory allocations. - - Returns - ------- - The concatenated table resulting from unpacking the input partitions. - - Raises - ------ - ReservationError - If the buffer resource cannot reserve enough memory to concatenate all - partitions. - - See Also - -------- - rapidsmpf.integrations.cudf.partition.partition_and_pack - """ - cdef cuda_stream_view _stream = stream.view() - cdef cpp_BufferResource* _br = br.ptr() - cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) - cdef unique_ptr[cpp_table] _ret - with nogil: - _ret = cpp_unpack_and_concat( - move(_partitions), - _stream, - _br, - ) - return Table.from_libcudf(move(_ret), stream, br._device_mr) - - -cdef extern from "" nogil: - cdef vector[cpp_PackedData] cpp_spill_partitions \ - "rapidsmpf::spill_partitions"( - vector[cpp_PackedData] partitions, - cpp_BufferResource* br, - ) except +ex_handler - - -def spill_partitions( - partitions, - BufferResource br not None, -): - """ - Spill partitions from device memory to host memory. - - Moves the buffer of each ``PackedData`` from device memory to host memory using - the provided buffer resource and the buffer's CUDA stream. Partitions already - in host memory are returned unchanged. - - For device-resident partitions, a host memory reservation is made before moving - the buffer. If the reservation fails due to insufficient host memory, an - exception is raised. Overbooking is not allowed. - - The input partitions are released and are left empty on return. - - Parameters - ---------- - partitions - The partitions to spill. - br - Buffer resource used to reserve host memory and perform the move. - - Returns - ------- - A list of partitions whose buffers reside in host memory. - - Raises - ------ - ReservationError - If host memory reservation fails. - """ - cdef cpp_BufferResource* _br = br.ptr() - cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) - cdef vector[cpp_PackedData] _ret - with nogil: - _ret = cpp_spill_partitions( - move(_partitions), - _br, - ) - return packed_data_vector_to_list(move(_ret), br) - - -cdef extern from "" nogil: - cdef vector[cpp_PackedData] cpp_unspill_partitions \ - "rapidsmpf::unspill_partitions"( - vector[cpp_PackedData] partitions, - cpp_BufferResource* br, - AllowOverbooking allow_overbooking, - ) except +ex_handler - - -def unspill_partitions( - partitions, - BufferResource br not None, - bool_t allow_overbooking, -): - """ - Move spilled partitions (i.e., packed tables in host memory) back to device memory. - - Each partition is inspected to determine whether its buffer resides in device - memory. Buffers already in device memory are left untouched. Host-resident buffers - are moved to device memory using the provided buffer resource and the buffer's CUDA - stream. - - If insufficient device memory is available, the buffer resource's spill manager is - invoked to free memory. If overbooking occurs and spilling fails to reclaim enough - memory, behavior depends on ``allow_overbooking``. - - The input partitions are released and are left empty on return. - - Parameters - ---------- - partitions - The partitions to unspill, potentially containing host-resident data. - br - Buffer resource responsible for memory reservation and spills. - allow_overbooking - If False, ensures enough memory is freed to satisfy the reservation; - otherwise, allows overbooking even if spilling was insufficient. - - Returns - ------- - A list of partitions whose buffers reside in device memory. - - Raises - ------ - ReservationError - If overbooking exceeds the amount spilled and ``allow_overbooking is False``. - """ - cdef cpp_BufferResource* _br = br.ptr() - cdef vector[cpp_PackedData] _partitions = _partitions_py_to_cpp(partitions) - cdef vector[cpp_PackedData] _ret - cdef AllowOverbooking ab = ( - AllowOverbooking.YES if allow_overbooking else AllowOverbooking.NO - ) - with nogil: - _ret = cpp_unspill_partitions( - move(_partitions), - _br, - ab, - ) - return packed_data_vector_to_list(move(_ret), br) diff --git a/python/rapidsmpf/rapidsmpf/memory/packed_data.pyi b/python/rapidsmpf/rapidsmpf/memory/packed_data.pyi index 1e7450f58..f044f4d48 100644 --- a/python/rapidsmpf/rapidsmpf/memory/packed_data.pyi +++ b/python/rapidsmpf/rapidsmpf/memory/packed_data.pyi @@ -1,22 +1,10 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -from typing import Self - -from pylibcudf.contiguous_split import PackedColumns as CudfPackedColumns -from rmm.pylibrmm.stream import Stream - from rapidsmpf.memory.buffer_resource import BufferResource class PackedData: def __init__(self) -> None: ... @classmethod - def from_cudf_packed_columns( - cls: type[Self], - packed_columns: CudfPackedColumns, - stream: Stream, - br: BufferResource, - ) -> Self: ... - @classmethod def from_host_bytes( cls, data: bytes | bytearray, br: BufferResource ) -> PackedData: ... diff --git a/python/rapidsmpf/rapidsmpf/memory/packed_data.pyx b/python/rapidsmpf/rapidsmpf/memory/packed_data.pyx index a778b2e6f..124faa089 100644 --- a/python/rapidsmpf/rapidsmpf/memory/packed_data.pyx +++ b/python/rapidsmpf/rapidsmpf/memory/packed_data.pyx @@ -1,15 +1,12 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 -from cython.operator cimport dereference as deref from libc.stdint cimport uint8_t from libcpp.memory cimport make_unique, unique_ptr from libcpp.utility cimport move from libcpp.vector cimport vector -from pylibcudf.contiguous_split cimport PackedColumns from rmm.librmm.cuda_stream_view cimport cuda_stream_view from rmm.librmm.device_buffer cimport device_buffer -from rmm.pylibrmm.stream cimport Stream from rapidsmpf._detail.exception_handling cimport ex_handler from rapidsmpf.memory.buffer_resource cimport (BufferResource, @@ -102,52 +99,6 @@ cdef class PackedData: self._br = br return self - @classmethod - def from_cudf_packed_columns( - cls, - PackedColumns packed_columns not None, - Stream stream not None, - BufferResource br not None, - ): - """ - Constructs a PackedData from CudfPackedColumns by taking the ownership of the - data and releasing ``packed_columns``. - - Parameters - ---------- - packed_columns - Packed data containing metadata and GPU data buffers - - Returns - ------- - A new PackedData instance containing the packed columns data - - Raises - ------ - ValueError - If the PackedColumns object is empty (has been released already). - """ - cdef cuda_stream_view _stream = stream.view() - cdef cpp_BufferResource* _br = br.ptr() - cdef PackedData ret = cls.__new__(cls) - with nogil: - if not (packed_columns.c_obj != NULL and - deref(packed_columns.c_obj).metadata and - deref(packed_columns.c_obj).gpu_data): - raise ValueError("Cannot release empty PackedColumns") - - # we cannot use packed_columns.release() because it returns a tuple of - # memoryview and gpumemoryview, and we need to take ownership of the - # underlying buffers - ret.c_obj = cpp_packed_data_from_buffers( - move(deref(packed_columns.c_obj).metadata), - move(deref(packed_columns.c_obj).gpu_data), - _stream, - _br, - ) - ret._br = br - return ret - def __init__(self): """Initialize an empty PackedData instance.""" pass diff --git a/python/rapidsmpf/rapidsmpf/streaming/CMakeLists.txt b/python/rapidsmpf/rapidsmpf/streaming/CMakeLists.txt index 929eaec92..0ba75ffe0 100644 --- a/python/rapidsmpf/rapidsmpf/streaming/CMakeLists.txt +++ b/python/rapidsmpf/rapidsmpf/streaming/CMakeLists.txt @@ -9,4 +9,3 @@ add_subdirectory(_detail) add_subdirectory(chunks) add_subdirectory(coll) add_subdirectory(core) -add_subdirectory(cudf) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/CMakeLists.txt b/python/rapidsmpf/rapidsmpf/streaming/cudf/CMakeLists.txt deleted file mode 100644 index c20c4b0c5..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/CMakeLists.txt +++ /dev/null @@ -1,14 +0,0 @@ -# ================================================================================= -# cmake-format: off -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -# cmake-format: on -# ================================================================================= - -set(cython_modules bloom_filter.pyx channel_metadata.pyx parquet.pyx partition.pyx table_chunk.pyx) - -rapids_cython_create_modules( - CXX - SOURCE_FILES "${cython_modules}" - LINKED_LIBRARIES rapidsmpf::rapidsmpf maybe_asan MODULE_PREFIX streaming_cudf_ -) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/__init__.pxd b/python/rapidsmpf/rapidsmpf/streaming/cudf/__init__.pxd deleted file mode 100644 index a9c2f923c..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/__init__.pxd +++ /dev/null @@ -1,2 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/__init__.py b/python/rapidsmpf/rapidsmpf/streaming/cudf/__init__.py deleted file mode 100644 index c3c7fac52..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/__init__.py +++ /dev/null @@ -1,23 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Submodule for streaming cudf operations.""" - -from __future__ import annotations - -from rapidsmpf.streaming.cudf.channel_metadata import ( - ChannelMetadata, - HashScheme, - OrderKey, - OrderScheme, - Partitioning, -) -from rapidsmpf.streaming.cudf.table_chunk import TableChunk - -__all__ = [ - "ChannelMetadata", - "HashScheme", - "OrderKey", - "OrderScheme", - "Partitioning", - "TableChunk", -] diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/bloom_filter.pxd b/python/rapidsmpf/rapidsmpf/streaming/cudf/bloom_filter.pxd deleted file mode 100644 index 49d7b3c09..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/bloom_filter.pxd +++ /dev/null @@ -1,44 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from libc.stddef cimport size_t -from libc.stdint cimport uint64_t -from libcpp.memory cimport shared_ptr, unique_ptr - -from rapidsmpf.communicator.communicator cimport Communicator, cpp_Communicator -from rapidsmpf.streaming.core.context cimport cpp_Context - - -cdef extern from "" nogil: - cdef cppclass cpp_BloomFilter "rapidsmpf::streaming::BloomFilter": - cpp_BloomFilter( - shared_ptr[cpp_Context] ctx, - shared_ptr[cpp_Communicator] comm, - uint64_t seed, - size_t num_filter_blocks, - ) noexcept - const shared_ptr[cpp_Communicator]& comm() noexcept - - -cdef extern from "" nogil: - size_t cpp_fitting_num_blocks \ - "rapidsmpf::BloomFilter::fitting_num_blocks"(size_t l2size) noexcept - - -cdef class BloomFilter: - """ - Streaming bloom filter construction and application. - - Parameters - ---------- - ctx - Streaming context. - comm - The communicator the bloom filter construction is collective over. - seed - Seed used for hashing values into the bloom filter. - num_filter_blocks - Number of blocks used to size the filter. - """ - cdef unique_ptr[cpp_BloomFilter] _handle - cdef Communicator _comm diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/bloom_filter.pyi b/python/rapidsmpf/rapidsmpf/streaming/cudf/bloom_filter.pyi deleted file mode 100644 index c7f074230..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/bloom_filter.pyi +++ /dev/null @@ -1,43 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from collections.abc import Iterable -from typing import Self - -from rapidsmpf.communicator.communicator import Communicator -from rapidsmpf.streaming.core.channel import Channel -from rapidsmpf.streaming.core.context import Context -from rapidsmpf.streaming.core.message import Message -from rapidsmpf.streaming.cudf.table_chunk import TableChunk - -class BloomFilterChunk: - # Note: if you go looking for this type in the cython bindings, you - # won't find it. This is purely to provide for better type-checking of - # the generic Channel argument to BloomFilter.build/apply below. - @classmethod - def from_message(cls: type[Self], message: Message[Self]) -> Self: ... - def into_message(self, sequence_number: int, message: Message[Self]) -> None: ... - -class BloomFilter: - def __init__( - self, ctx: Context, comm: Communicator, seed: int, num_filter_blocks: int - ) -> None: ... - @property - def comm(self) -> Communicator: ... - @staticmethod - def fitting_num_blocks(l2size: int) -> int: ... - async def build( - self, - ctx: Context, - ch_in: Channel[TableChunk], - ch_out: Channel[BloomFilterChunk], - tag: int, - ) -> None: ... - async def apply( - self, - ctx: Context, - bloom_filter: Channel[BloomFilterChunk], - ch_in: Channel[TableChunk], - ch_out: Channel[TableChunk], - keys: Iterable[int], - ) -> None: ... diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/bloom_filter.pyx b/python/rapidsmpf/rapidsmpf/streaming/cudf/bloom_filter.pyx deleted file mode 100644 index d5cd4be39..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/bloom_filter.pyx +++ /dev/null @@ -1,247 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from cpython.object cimport PyObject -from cpython.ref cimport Py_INCREF -from cython.operator cimport dereference as deref -from libc.stddef cimport size_t -from libc.stdint cimport int32_t, uint64_t -from libcpp.memory cimport make_unique, shared_ptr -from libcpp.utility cimport move -from libcpp.vector cimport vector -from pylibcudf.libcudf.types cimport size_type - -from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.communicator.communicator cimport Communicator -from rapidsmpf.owning_wrapper cimport cpp_OwningWrapper -from rapidsmpf.streaming._detail.libcoro_spawn_task cimport cpp_set_py_future -from rapidsmpf.streaming.chunks.utils cimport py_deleter -from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel -from rapidsmpf.streaming.core.context cimport Context, cpp_Context - -import asyncio - - -cdef extern from * nogil: - """ - namespace { - void cpp_bloom_filter_build( - std::shared_ptr ctx, - rapidsmpf::streaming::BloomFilter& bloom_filter, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - int32_t tag, - void (*cpp_set_py_future)(void*, const char *), - rapidsmpf::OwningWrapper py_future - ) { - RAPIDSMPF_EXPECTS( - ctx->executor()->spawn_detached( - cython_libcoro_task_wrapper( - cpp_set_py_future, - std::move(py_future), - bloom_filter.build( - std::move(ch_in), - std::move(ch_out), - tag - ) - ) - ), - "libcoro's spawn_detached() failed to spawn task" - ); - } - } // namespace - """ - void cpp_bloom_filter_build( - shared_ptr[cpp_Context] ctx, - cpp_BloomFilter& bloom_filter, - shared_ptr[cpp_Channel] ch_in, - shared_ptr[cpp_Channel] ch_out, - int32_t tag, - void (*cpp_set_py_future)(void*, const char *), - cpp_OwningWrapper py_future - ) except +ex_handler - - -cdef extern from * nogil: - """ - namespace { - void cpp_bloom_filter_apply( - std::shared_ptr ctx, - rapidsmpf::streaming::BloomFilter& bloom_filter, - std::shared_ptr bloom_filter_ch, - std::shared_ptr ch_in, - std::shared_ptr ch_out, - std::vector keys, - void (*cpp_set_py_future)(void*, const char *), - rapidsmpf::OwningWrapper py_future - ) { - RAPIDSMPF_EXPECTS( - ctx->executor()->spawn_detached( - cython_libcoro_task_wrapper( - cpp_set_py_future, - std::move(py_future), - bloom_filter.apply( - std::move(bloom_filter_ch), - std::move(ch_in), - std::move(ch_out), - std::move(keys) - ) - ) - ), - "libcoro's spawn_detached() failed to spawn task" - ); - } - } // namespace - """ - void cpp_bloom_filter_apply( - shared_ptr[cpp_Context] ctx, - cpp_BloomFilter& bloom_filter, - shared_ptr[cpp_Channel] bloom_filter_ch, - shared_ptr[cpp_Channel] ch_in, - shared_ptr[cpp_Channel] ch_out, - vector[size_type] keys, - void (*cpp_set_py_future)(void*, const char *), - cpp_OwningWrapper py_future - ) except +ex_handler - - -cdef class BloomFilter: - """ - Streaming bloom filter construction and application. - - Parameters - ---------- - ctx - Streaming context. - comm - The communicator the bloom filter construction is collective over. - seed - Seed used for hashing values into the bloom filter. - num_filter_blocks - Number of blocks used to size the filter. - """ - - def __init__( - self, - Context ctx not None, - Communicator comm not None, - uint64_t seed, - size_t num_filter_blocks, - ): - self._comm = comm - with nogil: - self._handle = make_unique[cpp_BloomFilter]( - ctx._handle, - comm._handle, - seed, - num_filter_blocks, - ) - - def __dealloc__(self): - with nogil: - self._handle.reset() - - @property - def comm(self): - """ - Get the communicator used by the bloom filter. - - Returns - ------- - The communicator. - """ - return self._comm - - @staticmethod - def fitting_num_blocks(size_t l2size): - """ - Return the number of blocks needed to fit within an L2 cache size. - - Parameters - ---------- - l2size - Size of the L2 cache in bytes. - - Returns - ------- - Number of blocks to use in the filter. - """ - cdef size_t ret - with nogil: - ret = cpp_fitting_num_blocks(l2size) - return ret - - async def build( - self, - Context ctx not None, - Channel ch_in not None, - Channel ch_out not None, - int32_t tag, - ): - """ - Build a bloom filter from input table chunks. - - Parameters - ---------- - ctx - The current streaming context. - ch_in - Input channel of ``TableChunk`` objects. - ch_out - Output channel receiving a single bloom filter message. - tag - Disambiguating tag to combine filters across ranks. - """ - ret = asyncio.get_running_loop().create_future() - Py_INCREF(ret) - with nogil: - cpp_bloom_filter_build( - ctx._handle, - deref(self._handle), - ch_in._handle, - ch_out._handle, - tag, - cpp_set_py_future, - move(cpp_OwningWrapper(ret, py_deleter)), - ) - await ret - - async def apply( - self, - Context ctx not None, - Channel bloom_filter not None, - Channel ch_in not None, - Channel ch_out not None, - keys, - ): - """ - Apply a bloom filter to incoming table chunks. - - Parameters - ---------- - ctx - The current streaming context. - bloom_filter - Channel containing the bloom filter (a single message). - ch_in - Input channel of ``TableChunk`` objects to filter. - ch_out - Output channel receiving filtered ``TableChunk`` objects. - keys - Indices selecting the key columns for hash fingerprints. - """ - cdef vector[size_type] c_keys = tuple(keys) - ret = asyncio.get_running_loop().create_future() - Py_INCREF(ret) - with nogil: - cpp_bloom_filter_apply( - ctx._handle, - deref(self._handle), - bloom_filter._handle, - ch_in._handle, - ch_out._handle, - move(c_keys), - cpp_set_py_future, - move(cpp_OwningWrapper(ret, py_deleter)), - ) - await ret diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pxd b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pxd deleted file mode 100644 index 98c778a79..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pxd +++ /dev/null @@ -1,130 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from libc.stdint cimport int32_t, uint64_t -from libcpp cimport bool as bool_t -from libcpp.memory cimport shared_ptr, unique_ptr -from libcpp.optional cimport optional -from libcpp.utility cimport pair -from libcpp.vector cimport vector -from pylibcudf.libcudf.types cimport null_order as cpp_null_order -from pylibcudf.libcudf.types cimport order as cpp_order -from rmm.librmm.cuda_stream_view cimport cuda_stream_view - -from rapidsmpf.memory.buffer_resource cimport cpp_BufferResource -from rapidsmpf.streaming.core.message cimport cpp_Message -from rapidsmpf.streaming.cudf.table_chunk cimport TableChunk, cpp_TableChunk - - -cdef extern from "" \ - namespace "rapidsmpf::streaming" nogil: - - cdef cppclass cpp_HashScheme "rapidsmpf::streaming::HashScheme": - vector[int32_t] column_indices - int modulus - cpp_HashScheme() except + - cpp_HashScheme(vector[int32_t], int) except + - bool_t operator==(const cpp_HashScheme&) - - cdef cppclass cpp_OrderKey "rapidsmpf::streaming::OrderKey": - cpp_OrderKey() noexcept - cpp_OrderKey(int32_t, cpp_order, cpp_null_order) noexcept - int32_t column_index - cpp_order order - cpp_null_order null_order - bool_t operator==(const cpp_OrderKey&) noexcept - - cdef cppclass cpp_OrderScheme "rapidsmpf::streaming::OrderScheme": - cpp_OrderScheme() noexcept - cpp_OrderScheme( - vector[cpp_OrderKey], unique_ptr[cpp_TableChunk], bool_t - ) except + - vector[cpp_OrderKey] keys - shared_ptr[cpp_TableChunk] boundaries - bool_t strict_boundaries - cpp_OrderScheme with_keys(vector[cpp_OrderKey]) except + - bool_t boundaries_aligned_with( - const cpp_OrderScheme&, const cpp_BufferResource& - ) except + - - cdef cppclass cpp_PartitioningSpec "rapidsmpf::streaming::PartitioningSpec": - enum cpp_Type "rapidsmpf::streaming::PartitioningSpec::Type": - NONE "rapidsmpf::streaming::PartitioningSpec::Type::NONE" - INHERIT "rapidsmpf::streaming::PartitioningSpec::Type::INHERIT" - HASH "rapidsmpf::streaming::PartitioningSpec::Type::HASH" - ORDER "rapidsmpf::streaming::PartitioningSpec::Type::ORDER" - - cpp_Type type - optional[cpp_HashScheme] hash - optional[cpp_OrderScheme] order - - @staticmethod - cpp_PartitioningSpec none() - - @staticmethod - cpp_PartitioningSpec inherit() - - @staticmethod - cpp_PartitioningSpec from_hash(cpp_HashScheme) - - @staticmethod - cpp_PartitioningSpec from_order(cpp_OrderScheme) - - cdef cppclass cpp_Partitioning "rapidsmpf::streaming::Partitioning": - cpp_PartitioningSpec inter_rank - cpp_PartitioningSpec local - cpp_Partitioning() except + - cpp_Partitioning(const cpp_Partitioning&) except + - - cdef cppclass cpp_ChannelMetadata "rapidsmpf::streaming::ChannelMetadata": - uint64_t local_count - cpp_Partitioning partitioning - bool_t duplicated - cpp_ChannelMetadata( - uint64_t, - cpp_Partitioning, - bool_t - ) except + - - cpp_Message cpp_to_message_channel_metadata "rapidsmpf::streaming::to_message" ( - uint64_t, unique_ptr[cpp_ChannelMetadata] - ) except + - - -cdef class HashScheme: - cdef cpp_HashScheme _handle - - @staticmethod - cdef HashScheme from_cpp(cpp_HashScheme scheme) - - -cdef class OrderKey: - cdef cpp_OrderKey _handle - - @staticmethod - cdef OrderKey from_cpp(cpp_OrderKey key) - - -cdef class OrderScheme: - cdef cpp_OrderScheme _handle - - @staticmethod - cdef OrderScheme from_cpp(cpp_OrderScheme scheme) - - -cdef class Partitioning: - cdef cpp_Partitioning _handle - - @staticmethod - cdef Partitioning from_cpp(cpp_Partitioning data) - - -cdef class ChannelMetadata: - cdef unique_ptr[cpp_ChannelMetadata] _handle - - @staticmethod - cdef ChannelMetadata from_handle(unique_ptr[cpp_ChannelMetadata] handle) - - cdef const cpp_ChannelMetadata* handle_ptr(self) except NULL - - cdef unique_ptr[cpp_ChannelMetadata] release_handle(self) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi deleted file mode 100644 index 09b075ad0..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyi +++ /dev/null @@ -1,88 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Type stubs for channel_metadata module.""" - -from __future__ import annotations - -from collections.abc import Sequence -from dataclasses import dataclass -from typing import Literal, Self - -import pylibcudf as plc - -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.streaming.core.message import Message -from rapidsmpf.streaming.cudf.table_chunk import TableChunk - -class HashScheme: - def __init__(self, column_indices: Sequence[int], modulus: int) -> None: ... - @property - def column_indices(self) -> tuple[int, ...]: ... - @property - def modulus(self) -> int: ... - def __eq__(self, other: object) -> bool: ... - def __repr__(self) -> str: ... - -@dataclass(frozen=True, slots=True) -class OrderKey: - """Sort key: column index, direction, and null ordering.""" - - column_index: int - order: plc.types.Order - null_order: plc.types.NullOrder - -class OrderScheme: - def __init__( - self, - keys: Sequence[OrderKey], - boundaries: TableChunk, - *, - strict_boundaries: bool = False, - ) -> None: ... - @property - def keys(self) -> tuple[OrderKey, ...]: ... - @property - def strict_boundaries(self) -> bool: ... - @property - def num_boundaries(self) -> int: ... - def get_boundaries(self, br: BufferResource) -> TableChunk: ... - def with_keys(self, new_keys: Sequence[OrderKey]) -> OrderScheme: ... - def boundaries_aligned_with( - self, other: OrderScheme, br: BufferResource - ) -> bool: ... - def __repr__(self) -> str: ... - -PartitioningSpecValue = HashScheme | OrderScheme | None | Literal["inherit"] - -class Partitioning: - def __init__( - self, - inter_rank: PartitioningSpecValue = None, - local: PartitioningSpecValue = None, - ) -> None: ... - @property - def inter_rank(self) -> PartitioningSpecValue: ... - @property - def local(self) -> PartitioningSpecValue: ... - def __eq__(self, other: object) -> bool: ... - def __repr__(self) -> str: ... - -class ChannelMetadata: - def __init__( - self, - local_count: int, - *, - partitioning: Partitioning | None = None, - duplicated: bool = False, - ) -> None: ... - @classmethod - def from_message(cls: type[Self], message: Message[Self]) -> ChannelMetadata: ... - def into_message(self, sequence_number: int, message: Message[Self]) -> None: ... - @property - def local_count(self) -> int: ... - @property - def partitioning(self) -> Partitioning: ... - @property - def duplicated(self) -> bool: ... - def __eq__(self, other: object) -> bool: ... - def __repr__(self) -> str: ... diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx b/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx deleted file mode 100644 index 28753f1f6..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/channel_metadata.pyx +++ /dev/null @@ -1,405 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Channel metadata types for streaming pipelines.""" - -from cython.operator cimport dereference as deref -from libc.stdint cimport int32_t, uint64_t -from libcpp.memory cimport make_unique, unique_ptr -from libcpp.utility cimport move -from libcpp.vector cimport vector -from pylibcudf.libcudf.types cimport null_order as cpp_null_order -from pylibcudf.libcudf.types cimport order as cpp_order -from pylibcudf.table cimport Table -from rmm.pylibrmm.stream cimport Stream - -from rapidsmpf.memory.buffer_resource cimport BufferResource -from rapidsmpf.streaming.core.message cimport Message -from rapidsmpf.streaming.cudf.table_chunk cimport TableChunk, cpp_TableChunk - - -cdef extern from * nogil: - """ - #include - #include - - static std::unique_ptr - cpp_channel_metadata_from_message(rapidsmpf::streaming::Message msg) { - return std::make_unique( - msg.release() - ); - } - """ - unique_ptr[cpp_ChannelMetadata] cpp_channel_metadata_from_message( - cpp_Message - ) except + - - -cdef class HashScheme: - """Hash partitioning scheme: rows distributed by hash(column_indices) % modulus.""" - - def __init__(self, object column_indices, int modulus): - """ - Parameters - ---------- - column_indices - Column indices to hash on. - modulus - Number of hash partitions. - """ - cdef vector[int32_t] cols - for c in column_indices: - cols.push_back(c) - self._handle = cpp_HashScheme(cols, modulus) - - @staticmethod - cdef HashScheme from_cpp(cpp_HashScheme scheme): - cdef HashScheme ret = HashScheme.__new__(HashScheme) - ret._handle = move(scheme) - return ret - - @property - def column_indices(self) -> tuple: - """Column indices used for hashing.""" - return tuple(self._handle.column_indices) - - @property - def modulus(self) -> int: - """Number of hash partitions.""" - return self._handle.modulus - - def __eq__(self, other): - if not isinstance(other, HashScheme): - return NotImplemented - return self._handle == (other)._handle - - def __repr__(self): - return f"HashScheme({self.column_indices!r}, {self.modulus})" - - -cdef class OrderKey: - """A single sort key: column index, sort direction, and null placement.""" - - def __init__( - self, - int column_index, - cpp_order order, - cpp_null_order null_order, - ): - """ - Parameters - ---------- - column_index - Zero-based index of the sort column. - order - Sort direction (ascending or descending). - null_order - Null placement (before or after non-null values). - """ - self._handle = cpp_OrderKey(column_index, order, null_order) - - @staticmethod - cdef OrderKey from_cpp(cpp_OrderKey key): - cdef OrderKey ret = OrderKey.__new__(OrderKey) - ret._handle = key - return ret - - @property - def column_index(self) -> int: - """Zero-based index of the sort column.""" - return self._handle.column_index - - @property - def order(self): - """Sort direction (ascending or descending).""" - return self._handle.order - - @property - def null_order(self): - """Null placement (before or after non-null values).""" - return self._handle.null_order - - def __eq__(self, other): - if not isinstance(other, OrderKey): - return NotImplemented - return self._handle == (other)._handle - - def __repr__(self): - return f"OrderKey({self.column_index}, {self.order!r}, {self.null_order!r})" - - -cdef class OrderScheme: - """Order-based partitioning scheme for sorted/range-partitioned data. - - Data is partitioned by value ranges based on predetermined boundaries. - For N partitions, there are N-1 boundary rows. - - Parameters - ---------- - keys - Sequence of ``OrderKey`` objects (one per sort column). - boundaries - Optional ``TableChunk`` of N-1 boundary rows for N partitions. - strict_boundaries - When true, every row in a chunk falls in a single partition's half-open key - range (keys do not straddle chunk interiors). See the C++ ``OrderScheme`` docs. - Default false. - """ - - def __init__( - self, - object keys, - TableChunk boundaries not None, - *, - bint strict_boundaries = False, - ): - cdef vector[cpp_OrderKey] cpp_keys - for key in keys: - cpp_keys.push_back((key)._handle) - if cpp_keys.empty(): - raise ValueError("OrderScheme: keys must not be empty") - self._handle = cpp_OrderScheme( - move(cpp_keys), move(boundaries.release_handle()), strict_boundaries - ) - - @staticmethod - cdef OrderScheme from_cpp(cpp_OrderScheme scheme): - cdef OrderScheme ret = OrderScheme.__new__(OrderScheme) - ret._handle = move(scheme) - return ret - - @property - def keys(self) -> tuple: - """Sort keys, one per sort column.""" - cdef int i - cdef int n = self._handle.keys.size() - return tuple(OrderKey.from_cpp(self._handle.keys[i]) for i in range(n)) - - @property - def strict_boundaries(self) -> bool: - """Same semantics as the C++ ``OrderScheme::strict_boundaries`` field.""" - return self._handle.strict_boundaries - - @property - def num_boundaries(self) -> int: - """Number of boundary rows (N-1 for N partitions).""" - return self._handle.boundaries.get().shape().first - - def get_boundaries(self, BufferResource br not None) -> TableChunk: - """ - Return the boundary rows. - - Parameters - ---------- - br - Buffer resource to associate with the returned table chunk. - - Returns - ------- - TableChunk - A non-exclusive view of the boundary rows owned by this scheme. - """ - cdef const cpp_TableChunk* chunk = self._handle.boundaries.get() - cdef Stream stream = Stream._from_cudaStream_t(chunk.stream().value()) - tbl = Table.from_table_view_of_arbitrary( - chunk.table_view(), owner=self, stream=stream - ) - return TableChunk.from_pylibcudf_table( - tbl, stream, exclusive_view=False, br=br - ) - - def with_keys(self, object new_keys) -> OrderScheme: - """Return a new ``OrderScheme`` with updated key column indices.""" - cdef vector[cpp_OrderKey] cpp_keys - for key in new_keys: - cpp_keys.push_back((key)._handle) - return OrderScheme.from_cpp(self._handle.with_keys(move(cpp_keys))) - - def boundaries_aligned_with( - self, OrderScheme other not None, BufferResource br not None - ) -> bool: - """ - Check whether boundary values are aligned with another scheme. - - Parameters - ---------- - other - The scheme to compare against. - br - Buffer resource for temporary allocations during comparison. - """ - return self._handle.boundaries_aligned_with(other._handle, deref(br.ptr())) - - def __repr__(self): - return ( - f"OrderScheme({self.keys!r}, " - f"strict_boundaries={self.strict_boundaries})" - ) - - -cdef void _apply_spec(cpp_PartitioningSpec& spec, obj) except *: - """Set *spec* in-place from a Python value.""" - if obj is None: - spec = cpp_PartitioningSpec.none() - elif obj == "inherit": - spec = cpp_PartitioningSpec.inherit() - elif isinstance(obj, HashScheme): - spec = cpp_PartitioningSpec.from_hash((obj)._handle) - elif isinstance(obj, OrderScheme): - spec = cpp_PartitioningSpec.from_order((obj)._handle) - else: - raise TypeError( - f"Expected HashScheme, OrderScheme, None, or 'inherit', " - f"got {type(obj).__name__}" - ) - - -cdef object _from_spec(const cpp_PartitioningSpec& spec): - """Convert PartitioningSpec (by reference) to a Python object.""" - if spec.type == cpp_PartitioningSpec.cpp_Type.NONE: - return None - elif spec.type == cpp_PartitioningSpec.cpp_Type.INHERIT: - return "inherit" - elif spec.type == cpp_PartitioningSpec.cpp_Type.HASH: - return HashScheme.from_cpp(deref(spec.hash)) - elif spec.type == cpp_PartitioningSpec.cpp_Type.ORDER: - return OrderScheme.from_cpp(deref(spec.order)) # copies out of optional - else: - raise ValueError("Unknown PartitioningSpec.Type") - - -cdef class Partitioning: - """ - Hierarchical partitioning metadata for a data stream. - - Parameters - ---------- - inter_rank - Distribution across ranks. Can be a HashScheme, OrderScheme, None, - or 'inherit'. - local - Distribution within a rank. Can be a HashScheme, OrderScheme, None, - or 'inherit'. - """ - - def __init__(self, inter_rank=None, local=None): - _apply_spec(self._handle.inter_rank, inter_rank) - _apply_spec(self._handle.local, local) - - @staticmethod - cdef Partitioning from_cpp(cpp_Partitioning data): - cdef Partitioning ret = Partitioning.__new__(Partitioning) - ret._handle = move(data) - return ret - - @property - def inter_rank(self): - """Inter-rank partitioning spec.""" - return _from_spec(self._handle.inter_rank) - - @property - def local(self): - """Intra-rank (local) partitioning spec.""" - return _from_spec(self._handle.local) - - def __repr__(self): - return f"Partitioning(inter_rank={self.inter_rank!r}, local={self.local!r})" - - -cdef class ChannelMetadata: - """ - Channel-level metadata describing a data stream. - - Parameters - ---------- - local_count - Estimated number of chunks for this rank. - partitioning - How the data is partitioned (default: no partitioning). - duplicated - Whether data is duplicated on all workers (default: False). - """ - - def __init__( - self, - int local_count, - *, - partitioning: Partitioning | None = None, - bint duplicated = False, - ): - if local_count < 0: - raise ValueError(f"local_count must be non-negative, got {local_count}") - - cdef cpp_Partitioning part - if partitioning is not None: - part = (partitioning)._handle - - self._handle = make_unique[cpp_ChannelMetadata]( - local_count, part, duplicated - ) - - def __dealloc__(self): - with nogil: - self._handle.reset() - - @staticmethod - cdef ChannelMetadata from_handle(unique_ptr[cpp_ChannelMetadata] handle): - cdef ChannelMetadata ret = ChannelMetadata.__new__(ChannelMetadata) - ret._handle = move(handle) - return ret - - @staticmethod - def from_message(Message message not None): - """Construct by consuming a Message (message becomes empty).""" - return ChannelMetadata.from_handle( - cpp_channel_metadata_from_message(move(message._handle)) - ) - - def into_message(self, uint64_t sequence_number, Message message not None): - """ - Move this ChannelMetadata into a Message. - - Parameters - ---------- - sequence_number - Ordering identifier for the message. - message - Empty message that will take ownership of this metadata. - """ - if not message.empty(): - raise ValueError("cannot move into a non-empty message") - message._handle = cpp_to_message_channel_metadata( - sequence_number, move(self.release_handle()) - ) - - cdef const cpp_ChannelMetadata* handle_ptr(self) except NULL: - """Return pointer to underlying handle, raising if released.""" - if not self._handle: - raise ValueError("ChannelMetadata is uninitialized, has it been released?") - return self._handle.get() - - @property - def local_count(self) -> int: - """Estimated number of chunks for this rank.""" - return self.handle_ptr().local_count - - @property - def partitioning(self) -> Partitioning: - """How the data is partitioned.""" - return Partitioning.from_cpp(self.handle_ptr().partitioning) - - @property - def duplicated(self) -> bool: - """Whether data is duplicated on all workers.""" - return self.handle_ptr().duplicated - - def __repr__(self): - return ( - f"ChannelMetadata(local_count={self.local_count}, " - f"partitioning={self.partitioning!r}, " - f"duplicated={self.duplicated})" - ) - - cdef unique_ptr[cpp_ChannelMetadata] release_handle(self): - if not self._handle: - raise ValueError("is uninitialized, has it been released?") - return move(self._handle) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/parquet.pyi b/python/rapidsmpf/rapidsmpf/streaming/cudf/parquet.pyi deleted file mode 100644 index 142992696..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/parquet.pyi +++ /dev/null @@ -1,25 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from pylibcudf.expressions import Expression -from pylibcudf.io.parquet import ParquetReaderOptions -from rmm.pylibrmm.stream import Stream - -from rapidsmpf.communicator.communicator import Communicator -from rapidsmpf.streaming.core.actor import CppActor -from rapidsmpf.streaming.core.channel import Channel -from rapidsmpf.streaming.core.context import Context -from rapidsmpf.streaming.cudf.table_chunk import TableChunk - -class Filter: - def __init__(self, stream: Stream, expression: Expression) -> None: ... - -def read_parquet( - ctx: Context, - comm: Communicator, - ch_out: Channel[TableChunk], - num_producers: int, - options: ParquetReaderOptions, - num_rows_per_chunk: int, - filter: Filter | None = None, -) -> CppActor: ... diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/parquet.pyx b/python/rapidsmpf/rapidsmpf/streaming/cudf/parquet.pyx deleted file mode 100644 index 37fb948ff..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/parquet.pyx +++ /dev/null @@ -1,145 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from cpython.object cimport PyObject -from cpython.ref cimport Py_INCREF -from cython.operator cimport dereference as deref -from libc.stddef cimport size_t -from libcpp.memory cimport make_unique, shared_ptr, unique_ptr -from libcpp.utility cimport move -from pylibcudf.expressions cimport Expression -from pylibcudf.io.parquet cimport ParquetReaderOptions -from pylibcudf.libcudf.expressions cimport expression -from pylibcudf.libcudf.io.parquet cimport parquet_reader_options -from pylibcudf.libcudf.types cimport size_type -from rmm.librmm.cuda_stream_view cimport cuda_stream_view -from rmm.pylibrmm.stream cimport Stream - -from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.communicator.communicator cimport Communicator, cpp_Communicator -from rapidsmpf.streaming.chunks.arbitrary cimport cpp_OwningWrapper -from rapidsmpf.streaming.chunks.utils cimport py_deleter -from rapidsmpf.streaming.core.actor cimport CppActor, cpp_Actor -from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel -from rapidsmpf.streaming.core.context cimport Context, cpp_Context - - -cdef extern from "" nogil: - cdef cppclass cpp_Filter "rapidsmpf::streaming::Filter": - cpp_Filter(cuda_stream_view, expression, cpp_OwningWrapper) - - cdef cpp_Actor cpp_read_parquet \ - "rapidsmpf::streaming::actor::read_parquet"( - shared_ptr[cpp_Context] ctx, - shared_ptr[cpp_Communicator] comm, - shared_ptr[cpp_Channel] ch_out, - size_t num_producers, - parquet_reader_options options, - size_type num_rows_per_chunk, - unique_ptr[cpp_Filter], - ) except +ex_handler - - -cdef class Filter: - """ - A filter expression for parquet reads. - - Parameters - ---------- - stream - The stream any scalars in the expression are valid on. - expression - The filter expression - - Notes - ----- - The object safely manages the lifetime of the expressions when called - from C++ coroutines, so it is safe to drop the expression passed in on - the python side. - """ - cdef unique_ptr[cpp_Filter] _handle - - def __init__(self, Stream stream not None, Expression filter not None): - Py_INCREF(filter) - self._handle = make_unique[cpp_Filter]( - stream.view(), - deref(filter.c_obj), - cpp_OwningWrapper( - filter, py_deleter - ) - ) - - cdef unique_ptr[cpp_Filter] release_handle(self): - """ - Move the owning C++ handle out of the object. - - Returns - ------- - unique_ptr to the C++ Filter object. - - Raises - ------ - ValueError - If this Filter has already been used and the handle is already released. - """ - if not self._handle: - raise ValueError("Filter is uninitialized, has it been released?") - return move(self._handle) - - def __dealloc__(self): - with nogil: - self._handle.reset() - - -def read_parquet( - Context ctx not None, - Communicator comm not None, - Channel ch_out not None, - size_t num_producers, - ParquetReaderOptions options not None, - size_type num_rows_per_chunk, - Filter filter = None, -): - """ - Create a streaming actor to read from parquet. - - Parameters - ---------- - ctx - Streaming execution context. - comm - The communicator. - ch_out - Output channel to receive the TableChunks. - num_producers - Number of concurrent producers of output chunks. - options - Reader options. - num_rows_per_chunk - Target (maximum) number of rows per output chunk. - filter - Optional filter object. If provided, is consumed by this function - and not subsequently usable. - - Notes - ----- - This is a collective operation, all ranks participating via the - communicator must call it with the same options. - """ - cdef cpp_Actor _ret - cdef unique_ptr[cpp_Filter] c_filter - if filter is not None: - c_filter = move(filter.release_handle()) - with nogil: - _ret = cpp_read_parquet( - ctx._handle, - comm._handle, - ch_out._handle, - num_producers, - options.c_obj, - num_rows_per_chunk, - move(c_filter) - ) - return CppActor.from_handle( - make_unique[cpp_Actor](move(_ret)), owner=None - ) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/partition.pyi b/python/rapidsmpf/rapidsmpf/streaming/cudf/partition.pyi deleted file mode 100644 index 7801dffb6..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/partition.pyi +++ /dev/null @@ -1,25 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from collections.abc import Iterable - -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 rapidsmpf.streaming.cudf.table_chunk import TableChunk - -def partition_and_pack( - ctx: Context, - ch_in: Channel[TableChunk], - ch_out: Channel[PartitionMapChunk], - columns_to_hash: Iterable[int], - num_partitions: int, -) -> CppActor: ... -def unpack_and_concat( - ctx: Context, - ch_in: Channel[PartitionMapChunk] - | Channel[PartitionVectorChunk] - | Channel[PartitionMapChunk | PartitionVectorChunk], - ch_out: Channel[TableChunk], -) -> CppActor: ... diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/partition.pyx b/python/rapidsmpf/rapidsmpf/streaming/cudf/partition.pyx deleted file mode 100644 index 3fd199bf7..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/partition.pyx +++ /dev/null @@ -1,150 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - - -from libc.stdint cimport uint32_t -from libcpp.memory cimport make_unique, shared_ptr -from libcpp.utility cimport move -from libcpp.vector cimport vector -from pylibcudf.libcudf.types cimport size_type - -from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.streaming.core.actor cimport CppActor, cpp_Actor -from rapidsmpf.streaming.core.channel cimport Channel, cpp_Channel -from rapidsmpf.streaming.core.context cimport Context, cpp_Context - - -cdef extern from "" nogil: - int cpp_HASH_MURMUR3"cudf::hash_id::HASH_MURMUR3" - uint32_t cpp_DEFAULT_HASH_SEED"cudf::DEFAULT_HASH_SEED", - cdef cpp_Actor cpp_partition_and_pack \ - "rapidsmpf::streaming::actor::partition_and_pack"( - shared_ptr[cpp_Context] ctx, - shared_ptr[cpp_Channel] ch_in, - shared_ptr[cpp_Channel] ch_out, - vector[size_type] columns_to_hash, - int num_partitions, - int hash_function, - uint32_t seed, - ) except +ex_handler - cdef cpp_Actor cpp_unpack_and_concat \ - "rapidsmpf::streaming::actor::unpack_and_concat"( - shared_ptr[cpp_Context] ctx, - shared_ptr[cpp_Channel] ch_in, - shared_ptr[cpp_Channel] ch_out, - ) except +ex_handler - - -def partition_and_pack( - Context ctx not None, - Channel ch_in not None, - Channel ch_out not None, - object columns_to_hash not None, - int num_partitions, -): - """ - Asynchronously partition and pack table chunks. - - This is the streaming equivalent of - :func:`rapidsmpf.integrations.cudf.partition.partition_and_split()`, - operating on incoming table chunks via channels. - - Each incoming table from `ch_in` is partitioned into `num_partitions` outputs - based on a hash of the specified columns. Each partition is then serialized - (packed) and sent to the output channel `ch_out`. - - Parameters - ---------- - ctx - The streaming actor context used to create and manage the asynchronous task. - ch_in - Input channel that provides ``TableChunk`` objects to partition. - ch_out - Output channel to which packed partitions (``PartitionMapChunk`` objects) - are sent. - columns_to_hash - Indices of input columns to hash when computing partition assignments. - num_partitions - Number of output partitions to create. - - Returns - ------- - A streaming actor representing the asynchronous partitioning and packing operation. - - Raises - ------ - ValueError - If any index in ``columns_to_hash`` is invalid. - - See Also - -------- - rapidsmpf.integrations.cudf.partition.partition_and_pack - Non-streaming variant operating on static tables. - rapidsmpf.streaming.cudf.partition.unpack_and_concat - The inverse operation that unpacks and concatenates packed partitions. - """ - cdef vector[size_type] _columns_to_hash = tuple(columns_to_hash) - cdef cpp_Actor _ret - with nogil: - _ret = cpp_partition_and_pack( - ctx._handle, - ch_in._handle, - ch_out._handle, - _columns_to_hash, - num_partitions, - cpp_HASH_MURMUR3, - cpp_DEFAULT_HASH_SEED, - ) - return CppActor.from_handle( - make_unique[cpp_Actor](move(_ret)), owner = None - ) - - -def unpack_and_concat( - Context ctx not None, - Channel ch_in not None, - Channel ch_out not None, -): - """ - Asynchronously unpack and concatenate packed partitions. - - This is the streaming equivalent of - :func:`rapidsmpf.integrations.cudf.partition.unpack_and_concat()`, - operating on packed partition chunks via channels. - - The function receives packed partitions from `ch_in`, deserializes them, - concatenates the partitions belonging to the same logical table, and sends the - resulting tables to `ch_out`. Empty partitions are automatically ignored. - - Parameters - ---------- - ctx - The streaming actor context used to manage asynchronous execution. - ch_in - Input channel providing packed partitions (``PartitionMapChunk`` or - ``PartitionVectorChunk``). - ch_out - Output channel receiving the unpacked and concatenated ``TableChunk`` objects. - - Returns - ------- - A streaming actor representing the asynchronous unpacking and concatenation - operation. - - See Also - -------- - rapidsmpf.integrations.cudf.partition.unpack_and_concat - Non-streaming version. - rapidsmpf.streaming.cudf.partition.partition_and_pack - The inverse operation that partitions and packs tables into partitions. - """ - cdef cpp_Actor _ret - with nogil: - _ret = cpp_unpack_and_concat( - ctx._handle, - ch_in._handle, - ch_out._handle, - ) - return CppActor.from_handle( - make_unique[cpp_Actor](move(_ret)), owner = None - ) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pxd b/python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pxd deleted file mode 100644 index 5b3bc9cd2..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pxd +++ /dev/null @@ -1,46 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from libc.stddef cimport size_t -from libc.stdint cimport uint64_t -from libcpp cimport bool as bool_t -from libcpp.memory cimport unique_ptr -from libcpp.utility cimport pair -from pylibcudf.libcudf.table.table_view cimport table_view as cpp_table_view -from pylibcudf.libcudf.types cimport size_type -from rmm.librmm.cuda_stream_view cimport cuda_stream_view -from rmm.pylibrmm.stream cimport Stream - -from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.memory.buffer cimport MemoryType -from rapidsmpf.memory.buffer_resource cimport (BufferResource, - cpp_BufferResource) -from rapidsmpf.memory.memory_reservation cimport cpp_MemoryReservation -from rapidsmpf.memory.packed_data cimport cpp_PackedData - - -cdef extern from "" nogil: - cdef cppclass cpp_TableChunk "rapidsmpf::streaming::TableChunk": - cpp_TableChunk(unique_ptr[cpp_PackedData]) except +ex_handler - cuda_stream_view stream() noexcept - size_t data_alloc_size(MemoryType mem_type) except +ex_handler - bool_t is_available() noexcept - size_t make_available_cost() noexcept - cpp_table_view table_view() except +ex_handler - bool_t is_spillable() noexcept - cpp_TableChunk copy(cpp_MemoryReservation& reservation) except +ex_handler - pair[size_type, size_type] shape() noexcept - unique_ptr[cpp_PackedData] into_packed_data( - cpp_BufferResource* br - ) except +ex_handler - -cdef class TableChunk: - cdef unique_ptr[cpp_TableChunk] _handle - # Keep the BufferResource alive as long as this object is so that when this - # object is deallocated the associated stream and memory resource are still alive. - cdef BufferResource _br - - @staticmethod - cdef TableChunk from_handle(unique_ptr[cpp_TableChunk] handle, BufferResource br) - cdef const cpp_TableChunk* handle_ptr(self) - cdef unique_ptr[cpp_TableChunk] release_handle(self) diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pyi b/python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pyi deleted file mode 100644 index 07098d899..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pyi +++ /dev/null @@ -1,70 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -from __future__ import annotations - -from collections.abc import Iterable -from typing import Self, overload - -from pylibcudf.table import Table -from rmm.pylibrmm.stream import Stream - -from rapidsmpf.memory.buffer import MemoryType -from rapidsmpf.memory.buffer_resource import BufferResource -from rapidsmpf.memory.memory_reservation import MemoryReservation -from rapidsmpf.memory.packed_data import PackedData -from rapidsmpf.streaming.core.context import Context -from rapidsmpf.streaming.core.message import Message - -class TableChunk: - @staticmethod - def from_pylibcudf_table( - table: Table, - stream: Stream, - *, - exclusive_view: bool, - br: BufferResource, - ) -> TableChunk: ... - @staticmethod - def from_packed_data(pd: PackedData, br: BufferResource) -> TableChunk: ... - @classmethod - def from_message( - cls: type[Self], message: Message[Self], br: BufferResource - ) -> Self: ... - def into_message(self, sequence_number: int, message: Message[Self]) -> None: ... - @property - def stream(self) -> Stream: ... - def data_alloc_size(self, mem_type: MemoryType | None = None) -> int: ... - def is_available(self) -> bool: ... - def make_available_cost(self) -> int: ... - def make_available(self, reservation: MemoryReservation) -> TableChunk: ... - async def make_available_or_wait( - self, ctx: Context, *, net_memory_delta: int - ) -> TableChunk: ... - def make_available_and_spill( - self, br: BufferResource, *, allow_overbooking: bool - ) -> TableChunk: ... - def table_view(self) -> Table: ... - def is_spillable(self) -> bool: ... - def copy(self, reservation: MemoryReservation) -> TableChunk: ... - def into_packed_data(self, br: BufferResource) -> PackedData: ... - @property - def shape(self) -> tuple[int, int]: ... - -@overload -async def make_table_chunks_available_or_wait( - context: Context, - chunks: TableChunk, - *, - reserve_extra: int, - net_memory_delta: int, - allow_overbooking: bool | None = None, -) -> tuple[TableChunk, MemoryReservation]: ... -@overload -async def make_table_chunks_available_or_wait( - context: Context, - chunks: Iterable[TableChunk], - *, - reserve_extra: int, - net_memory_delta: int, - allow_overbooking: bool | None = None, -) -> tuple[list[TableChunk], MemoryReservation]: ... diff --git a/python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pyx b/python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pyx deleted file mode 100644 index 8f7bb97d0..000000000 --- a/python/rapidsmpf/rapidsmpf/streaming/cudf/table_chunk.pyx +++ /dev/null @@ -1,663 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from cpython.object cimport PyObject -from cython cimport no_gc_clear -from cython.operator cimport dereference as deref -from libc.stdint cimport int64_t, uint64_t -from libcpp.memory cimport make_unique, unique_ptr -from libcpp.utility cimport move -from pylibcudf.libcudf.table.table_view cimport table_view as cpp_table_view -from pylibcudf.table cimport Table - -from rapidsmpf._detail.exception_handling cimport ex_handler -from rapidsmpf.memory.buffer_resource cimport (BufferResource, - cpp_BufferResource) -from rapidsmpf.memory.memory_reservation cimport (MemoryReservation, - cpp_MemoryReservation) -from rapidsmpf.memory.packed_data cimport PackedData -# Need the header include for inline C++ code -from rapidsmpf.owning_wrapper cimport cpp_OwningWrapper # no-cython-lint -from rapidsmpf.streaming.chunks.utils cimport py_deleter -from rapidsmpf.streaming.core.context cimport Context -from rapidsmpf.streaming.core.message cimport Message, cpp_Message - -from rapidsmpf.memory.buffer import MemoryType as py_MemoryType -from rapidsmpf.streaming.core.memory_reserve_or_wait import reserve_memory - - -cdef extern from "" nogil: - cpp_Message cpp_to_message"rapidsmpf::streaming::to_message"\ - (uint64_t sequence_number, unique_ptr[cpp_TableChunk]) except +ex_handler - - -cdef extern from * nogil: - """ - namespace { - std::unique_ptr - cpp_release_table_chunk_from_message( - rapidsmpf::streaming::Message &&msg - ) { - return std::make_unique( - msg.release() - ); - } - - std::unique_ptr cpp_from_table_view_with_owner( - cudf::table_view view, - rmm::cuda_stream_view stream, - PyObject *owner, - void(*py_deleter)(void *), - bool exclusive_view - ) { - // Called holding the gil. - // Decref is done by the deleter. - Py_XINCREF(owner); - return std::make_unique( - view, - stream, - rapidsmpf::OwningWrapper(owner, py_deleter), - exclusive_view ? - rapidsmpf::streaming::TableChunk::ExclusiveView::YES - : rapidsmpf::streaming::TableChunk::ExclusiveView::NO - ); - } - - std::unique_ptr cpp_table_make_available( - std::unique_ptr &&table, - rapidsmpf::MemoryReservation* reservation - ) { - return std::make_unique( - table->make_available(*reservation) - ); - } - - std::unique_ptr cpp_table_copy( - std::unique_ptr const& table, - rapidsmpf::MemoryReservation* reservation - ) { - return std::make_unique( - table->copy(*reservation) - ); - } - } // namespace - """ - unique_ptr[cpp_TableChunk] cpp_release_table_chunk_from_message( - cpp_Message - ) except +ex_handler - unique_ptr[cpp_TableChunk] cpp_from_table_view_with_owner(...) except +ex_handler - unique_ptr[cpp_TableChunk] cpp_table_make_available( - unique_ptr[cpp_TableChunk], cpp_MemoryReservation* - ) except +ex_handler - unique_ptr[cpp_TableChunk] cpp_table_copy( - unique_ptr[cpp_TableChunk], cpp_MemoryReservation* - ) except +ex_handler - - -@no_gc_clear -cdef class TableChunk: - """ - A unit of table data in a streaming pipeline. - - Represents either an unpacked pylibcudf table, a packed (serialized) table, - or `rapidsmpf.memory.packed_data.PackedData`. - - A TableChunk may be initially unavailable (e.g., if the data is packed or - spilled), and can be made available (i.e., materialized to device memory) - on demand. - - Use the factory functions `from_pylibcudf_table` and `from_message` to - create a new table chunk. - """ - def __init__(self): - raise ValueError("use the `from_*` factory functions") - - def __dealloc__(self): - with nogil: - self._handle.reset() - - @staticmethod - cdef TableChunk from_handle( - unique_ptr[cpp_TableChunk] handle, BufferResource br, - ): - """ - Construct a TableChunk from an existing C++ handle. - - Parameters - ---------- - handle - A unique pointer to a C++ TableChunk. - br - A BufferResource to keep alive. - - Returns - ------- - A new TableChunk wrapping the given handle. - """ - cdef TableChunk ret = TableChunk.__new__(TableChunk) - ret._handle = move(handle) - ret._br = br - return ret - - @staticmethod - def from_pylibcudf_table( - Table table not None, - Stream stream not None, - *, - bool_t exclusive_view, - BufferResource br not None, - ): - """ - Construct a TableChunk from a pylibcudf Table. - - Parameters - ---------- - table - A pylibcudf Table to wrap as a TableChunk. - stream - The CUDA stream on which this chunk was created. - exclusive_view - Indicates that this TableChunk has exclusive ownership semantics for the - underlying table view. - - When ``True``, the following guarantees must hold: - - The pylibcudf Table is the sole representation of the table data, - i.e. no views exist. - - The Table object exclusively owns the table's device memory. - - These guarantees allow the TableChunk to be spillable and ensure that - when the owner is destroyed, the underlying device memory is correctly - freed. - - Returns - ------- - A new TableChunk wrapping the given pylibcudf Table. - - Notes - ----- - The returned TableChunk maintains a reference to ``table`` to ensure - its underlying buffers remain valid for the lifetime of the chunk. - This reference is managed by the underlying C++ object, so it - persists even when the chunk is transferred through Channels. - - """ - cdef cuda_stream_view _stream = stream.view() - cdef cpp_table_view view = table.view() - return TableChunk.from_handle( - cpp_from_table_view_with_owner( - view, - _stream, - table, - py_deleter, - exclusive_view, - ), - br, - ) - - @staticmethod - def from_packed_data(PackedData pd not None, BufferResource br not None): - """ - Construct a TableChunk from packed data. - - Parameters - ---------- - pd - The PackedData object - - Returns - ------- - A new TableChunk owning the packed data. - - Notes - ----- - This takes ownership of the data in the PackedData object, which is left empty. - """ - return TableChunk.from_handle(make_unique[cpp_TableChunk](move(pd.c_obj)), br) - - @staticmethod - def from_message(Message message not None, BufferResource br not None): - """ - Construct a TableChunk by consuming a Message. - - Parameters - ---------- - message - Message containing a TableChunk. The message is released and is empty - after this call. - - Returns - ------- - A new TableChunk extracted from the given message. - """ - return TableChunk.from_handle( - cpp_release_table_chunk_from_message(move(message._handle)), - br, - ) - - def into_message(self, uint64_t sequence_number, Message message not None): - """ - Move this TableChunk into an empty Message. - - This method is not typically called directly. Instead, it is invoked by - `Message.__init__()` when creating a new Message with this TableChunk as - its payload. - - Parameters - ---------- - sequence_number - Ordering identifier for the message. - message - Message object that will take ownership of this TableChunk. - - Raises - ------ - ValueError - If the provided message is not empty. - - Warnings - -------- - The original table chunk is released and must not be used after this call. - """ - if not message.empty(): - raise ValueError("cannot move into a non-empty message") - message._handle = cpp_to_message( - sequence_number, move(self.release_handle()) - ) - - cdef const cpp_TableChunk* handle_ptr(self): - """ - Return a pointer to the underlying C++ TableChunk. - - Returns - ------- - Raw pointer to the underlying C++ object. - - Raises - ------ - ValueError - If the TableChunk is uninitialized. - """ - if not self._handle: - raise ValueError("TableChunk is uninitialized, has it been released?") - return self._handle.get() - - cdef unique_ptr[cpp_TableChunk] release_handle(self): - """ - Release ownership of the underlying C++ TableChunk. - - After this call, the current object is in a moved-from state and - must not be accessed. - - Returns - ------- - Unique pointer to the underlying C++ object. - - Raises - ------ - ValueError - If the TableChunk is uninitialized. - """ - if not self._handle: - raise ValueError("TableChunk is uninitialized, has it been released?") - return move(self._handle) - - @property - def stream(self): - """ - Return the CUDA stream on which this chunk was created. - - Returns - ------- - Stream - The CUDA stream. - """ - return Stream._from_cudaStream_t( - deref(self.handle_ptr()).stream().value() - ) - - def data_alloc_size(self, mem_type=None): - """ - Number of bytes allocated for the data in the specified memory type. - - Parameters - ---------- - mem_type - The memory type to query. If None, returns the total size across - all memory types. - - Returns - ------- - Number of bytes allocated. - """ - if mem_type is None: - return sum(self.data_alloc_size(m) for m in py_MemoryType) - return deref(self.handle_ptr()).data_alloc_size(mem_type) - - def is_available(self): - """ - Indicates whether the underlying table data is fully available in - device memory. - - Returns - ------- - True if the table is already available; otherwise, False. - """ - return deref(self.handle_ptr()).is_available() - - def make_available_cost(self): - """ - Return the estimated cost (in bytes) of making the table available. - - Currently, only device memory usage is accounted for in this estimate. - - Returns - ------- - The estimated cost in bytes. - """ - return deref(self.handle_ptr()).make_available_cost() - - def make_available(self, MemoryReservation reservation not None): - """ - Move this table chunk into a new one with its data made available. - - As part of the move, a copy or unpack operation may be performed, - using the associated CUDA stream for execution. - - Parameters - ---------- - reservation - Memory reservation used for allocations, if making data available - is needed. - - Returns - ------- - A new table chunk with its data available on device. - - Warnings - -------- - The original table chunk is released and must not be used after this call. - """ - cdef cpp_MemoryReservation* res = reservation._handle.get() - cdef unique_ptr[cpp_TableChunk] handle = self.release_handle() - cdef unique_ptr[cpp_TableChunk] ret - with nogil: - ret = cpp_table_make_available(move(handle), res) - return TableChunk.from_handle(move(ret), self._br) - - async def make_available_or_wait( - self, Context ctx not None, *, int64_t net_memory_delta - ): - """ - Move this table chunk into a new one with its data made available. - - This is an asynchronous variant of :meth:`make_available`. The coroutine may - suspend if the required device memory is not immediately available and - resumes once a memory reservation has been granted or an error condition is - reached. - - Parameters - ---------- - ctx - Streaming context used to access the memory reservation mechanism. - net_memory_delta - Estimated change in memory usage after the reservation is granted and - all work using the returned table chunk has completed. - - Returns - ------- - A new table chunk with its data available on device. - - Raises - ------ - RuntimeError - If shutdown occurs before the reservation can be processed. - OverflowError - If no progress is possible within the timeout and overbooking is - disabled. - - Warnings - -------- - The original table chunk is released and must not be used after this call. - """ - return self.make_available( - await reserve_memory( - ctx, self.make_available_cost(), net_memory_delta=net_memory_delta - ) - ) - - def make_available_and_spill( - self, BufferResource br not None, *, allow_overbooking - ): - """ - Make this table chunk available on device, spilling other data if necessary. - - Ensures that the data backing this table chunk is made available in device - memory. If there is insufficient free memory to complete the operation, the - buffer resource may spill other data until enough space has been freed. - - Parameters - ---------- - br - Buffer resource used for allocations and spill management. - allow_overbooking - Whether the memory reservation may temporarily exceed the current - allocation limit. - - Returns - ------- - A new table chunk with its data made available on device. - - Raises - ------ - ReservationError - If the allocation or spilling process fails to free enough memory. - - Warnings - -------- - The original table chunk is released and must not be used after this call. - - Examples - -------- - >>> # Make the data of an existing chunk available on device - >>> chunk = chunk.make_available_and_spill(br, allow_overbooking=False) - >>> chunk.table_view() - """ - - cdef MemoryReservation res = br.reserve_device_memory_and_spill( - self.make_available_cost(), - allow_overbooking=allow_overbooking - ) - return self.make_available(res) - - def table_view(self): - """ - Returns a view of the underlying pylibcudf table. - - The table must be available in device memory. - - Returns - ------- - A view of the underlying table. The view holds a reference to this - `TableChunk` to ensure it remains alive. - - Raises - ------ - ValueError - If ``self.is_available() is False``. - """ - cdef const cpp_TableChunk* handle = self.handle_ptr() - cdef cpp_table_view ret - with nogil: - ret = deref(handle).table_view() - return Table.from_table_view_of_arbitrary(ret, owner=self, stream=self.stream) - - def is_spillable(self): - """ - Indicates whether this chunk can be spilled. - - A chunk is considered spillable if it was created from one of the following: - - A message (via ``.from_message()``). - - An exclusive pylibcudf table (via - ``.from_pylibcudf_table(..., exclusive_view=True)``). - - Both of these creation paths imply device-owning semantics, meaning the - TableChunk owns its underlying memory and can safely be spilled to host memory. - - Returns - ------- - True if the table chunk can be spilled, otherwise, False. - """ - return deref(self.handle_ptr()).is_spillable() - - def copy(self, MemoryReservation reservation not None): - """ - Create a deep copy of this table chunk. - - All buffers are allocated for the new table chunk using the provided - memory reservation, which also determines the target memory type of - the copy. - - Parameters - ---------- - reservation - Memory reservation to consume for allocating the buffers of the - new table chunk. - - Returns - ------- - TableChunk - A new table chunk containing a deep copy of this chunk's data and - metadata. - """ - cdef unique_ptr[cpp_TableChunk] ret - cdef cpp_MemoryReservation* res = reservation._handle.get() - with nogil: - ret = cpp_table_copy(self._handle, res) - return TableChunk.from_handle(move(ret), self._br) - - def into_packed_data(self, BufferResource br not None): - """ - Convert this table chunk to a PackedData, avoiding unnecessary copies. - - If the chunk's data is already in packed form (e.g., it arrived over the - network or was constructed from a :class:`PackedData`), the packed data is - moved out directly with no copy. Otherwise the table is serialized via - ``cudf.pack()``. - - Parameters - ---------- - br - Buffer resource used when packing is required. - - Returns - ------- - PackedData - The resulting packed data. - - Raises - ------ - ValueError - If the data is not already packed and the table is not available - (i.e., ``is_available() == False``). - - Warnings - -------- - The original table chunk is released and must not be used after this call. - """ - cdef unique_ptr[cpp_PackedData] result - cdef cpp_BufferResource* _br = br.ptr() - cdef unique_ptr[cpp_TableChunk] handle = self.release_handle() - with nogil: - result = move(deref(handle)).into_packed_data(_br) - return PackedData.from_librapidsmpf(move(result), br) - - @property - def shape(self): - """Return the shape of the table in this TableChunk. - - Returns - ------- - Tuple of shape ``(num_rows, num_columns)```. - """ - return deref(self._handle).shape() - - -async def make_table_chunks_available_or_wait( - Context ctx not None, - chunks, - *, - size_t reserve_extra, - int64_t net_memory_delta, - allow_overbooking=None, -): - """ - Make one or more table chunks available, waiting on a memory reservation if needed. - - This helper combines :meth:`TableChunk.make_available` with :func:`reserve_memory`. - It computes the device-memory cost of making the provided table chunks available, - reserves that amount (plus ``reserve_extra``), and then returns new table chunks - whose data are available on device. - - The coroutine may suspend if the required device memory is not immediately - available and resumes once a memory reservation has been granted or an error - condition is reached. The behavior when the progress timeout expires depends - on whether overbooking is allowed. - - Parameters - ---------- - ctx - Streaming context used to access the memory reservation mechanism. - chunks - A table chunk or an iterable of table chunks to make available on device. - reserve_extra - Additional bytes to include in the reservation beyond the aggregated - availability cost of ``chunks``. - net_memory_delta - Estimated change in memory usage after the reservation is granted and all - work using the returned table chunks has completed. This value is used as - a heuristic to prioritize eligible requests. - allow_overbooking - Whether to allow overbooking if no progress is possible. - - If ``True``, the reservation may overbook memory when no further - progress can be made. If ``False``, the call fails when no progress - is possible. - - If ``None``, the behavior is determined by the configuration option - ``"allow_overbooking_by_default"``, which is read via ``ctx.options()``. - - Returns - ------- - A tuple containing: - - If a chunk is provided: the new table chunk with its data available on device. - - If multiple chunks are provided: a list of new table chunks with their data - available on device. - - The memory reservation used with a remaining size of @p reserve_extra. - - Raises - ------ - RuntimeError - If shutdown occurs before the reservation can be processed. - OverflowError - If no progress is possible within the timeout and overbooking is disabled. - - Warnings - -------- - The original table chunks are released and must not be used after this call. - """ - # Handle both single chunk and iterable of chunks - input_chunks = chunks - if isinstance(input_chunks, TableChunk): - chunks = (input_chunks,) - - size = sum(chunk.make_available_cost() for chunk in chunks) - res = await reserve_memory( - ctx, - size + reserve_extra, - net_memory_delta=net_memory_delta, - mem_type=MemoryType.DEVICE, - allow_overbooking=allow_overbooking, - ) - available_chunks = [chunk.make_available(res) for chunk in chunks] - if isinstance(input_chunks, TableChunk): - return available_chunks[0], res - else: - return available_chunks, res diff --git a/python/rapidsmpf/rapidsmpf/testing.py b/python/rapidsmpf/rapidsmpf/testing.py index e18746675..60207f286 100644 --- a/python/rapidsmpf/rapidsmpf/testing.py +++ b/python/rapidsmpf/rapidsmpf/testing.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING import pylibcudf + from rmm.pylibrmm.stream import DEFAULT_STREAM if TYPE_CHECKING: diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_allgather.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_allgather.py index 884e21b41..d8d6ed019 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_allgather.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_allgather.py @@ -7,20 +7,23 @@ from typing import TYPE_CHECKING import numpy as np +import pylibcudf as plc import pytest -import pylibcudf as plc +pytest.importorskip("cudf_streaming") +from cudf_streaming.integrations.partition import unpack_and_concat +from cudf_streaming.streaming.table_chunk import TableChunk -from rapidsmpf.integrations.cudf.partition import unpack_and_concat from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.chunks.packed_data import PackedDataChunk from rapidsmpf.streaming.coll.allgather import AllGather, allgather from rapidsmpf.streaming.core.actor import define_actor, run_actor_network from rapidsmpf.streaming.core.leaf_actor import pull_from_channel, push_to_channel from rapidsmpf.streaming.core.message import Message -from rapidsmpf.streaming.cudf.table_chunk import TableChunk from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + if TYPE_CHECKING: from collections.abc import Awaitable @@ -49,7 +52,7 @@ def test_allgather_actor(context: Context, comm: Communicator) -> None: ] inputs = [ PackedDataChunk.from_packed_data( - PackedData.from_cudf_packed_columns( + PackedData.from_cudf_packed_columns( # type: ignore[attr-defined] plc.contiguous_split.pack(table, stream=stream), stream, context.br(), @@ -104,7 +107,7 @@ async def generate_inputs( msg = Message( i, PackedDataChunk.from_packed_data( - PackedData.from_cudf_packed_columns( + PackedData.from_cudf_packed_columns( # type: ignore[attr-defined] plc.contiguous_split.pack(table, stream=stream), stream, context.br(), diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_bloom_filter.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_bloom_filter.py deleted file mode 100644 index da72f8d23..000000000 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_bloom_filter.py +++ /dev/null @@ -1,161 +0,0 @@ -# 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 pytest - -import pylibcudf as plc - -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.streaming.cudf import ChannelMetadata -from rapidsmpf.streaming.cudf.bloom_filter import BloomFilter -from rapidsmpf.streaming.cudf.table_chunk import TableChunk -from rapidsmpf.testing import assert_eq - -if TYPE_CHECKING: - from collections.abc import Awaitable - - from rmm.pylibrmm.stream import Stream - - from rapidsmpf.communicator.communicator import Communicator - from rapidsmpf.memory.buffer_resource import BufferResource - from rapidsmpf.streaming.core.actor import CppActor - from rapidsmpf.streaming.core.channel import Channel - from rapidsmpf.streaming.core.context import Context - from rapidsmpf.streaming.cudf.bloom_filter import BloomFilterChunk - - -def make_table(values: np.ndarray, stream: Stream, br: BufferResource) -> TableChunk: - table = plc.Table([plc.Column.from_array(values, stream=stream)]) - return TableChunk.from_pylibcudf_table(table, stream, exclusive_view=True, br=br) - - -@define_actor() -async def add_metadata( - ctx: Context, ch_in: Channel[TableChunk], ch_out: Channel[TableChunk] -) -> None: - await ch_out.send_metadata(ctx, Message(0, ChannelMetadata(1))) - await ch_out.drain_metadata(ctx) - while (msg := await ch_in.recv(ctx)) is not None: - await ch_out.send(ctx, msg) - await ch_out.drain(ctx) - - -@define_actor() -async def receive_metadata( - ctx: Context, ch_in: Channel[TableChunk], ch_out: Channel[TableChunk] -) -> None: - m = await ch_in.recv_metadata(ctx) - assert m is not None - meta = ChannelMetadata.from_message(m) - assert meta.local_count == 1 - while (msg := await ch_in.recv(ctx)) is not None: - await ch_out.send(ctx, msg) - await ch_out.drain(ctx) - - -@define_actor() -async def bloom_pipeline( - ctx: Context, - bloom: BloomFilter, - ch_build: Channel[TableChunk], - ch_probe: Channel[TableChunk], - ch_out: Channel[TableChunk], -) -> None: - ch_filter: Channel[BloomFilterChunk] = ctx.create_channel() - await asyncio.gather( - bloom.build(ctx, ch_in=ch_build, ch_out=ch_filter, tag=0), - bloom.apply( - ctx, - bloom_filter=ch_filter, - ch_in=ch_probe, - ch_out=ch_out, - keys=(0,), - ), - ) - - -def run_bloom_filter_pipeline( - context: Context, - comm: Communicator, - build_table: TableChunk, - probe_table: TableChunk, - *, - seed: int = 42, - l2size: int = 1 << 20, -) -> list[Message]: - bloom = BloomFilter( - context, - comm, - seed=seed, - num_filter_blocks=BloomFilter.fitting_num_blocks(l2size), - ) - - build_msg = Message(0, build_table) - probe_msg = Message(0, probe_table) - - ch_build: Channel[TableChunk] = context.create_channel() - ch_probe: Channel[TableChunk] = context.create_channel() - ch_probe_meta: Channel[TableChunk] = context.create_channel() - ch_out_meta: Channel[TableChunk] = context.create_channel() - ch_out: Channel[TableChunk] = context.create_channel() - - actors: list[CppActor | Awaitable[None]] = [ - push_to_channel(context, ch_build, [build_msg]), - push_to_channel(context, ch_probe, [probe_msg]), - add_metadata(context, ch_probe, ch_probe_meta), - bloom_pipeline(context, bloom, ch_build, ch_probe_meta, ch_out_meta), - receive_metadata(context, ch_out_meta, ch_out), - ] - pull_actor, deferred = pull_from_channel(context, ch_out) - actors.append(pull_actor) - run_actor_network(context, actors=actors) - return deferred.release() - - -def test_bloom_filter_roundtrip(context: Context, comm: Communicator) -> None: - if comm.nranks != 1: - pytest.skip("Only support single-rank runs") - - stream = context.get_stream_from_pool() - values = np.arange(10, dtype=np.int32) - build_table = make_table(values, stream=stream, br=context.br()) - probe_table = make_table(values, stream=stream, br=context.br()) - messages = run_bloom_filter_pipeline(context, comm, build_table, probe_table) - assert len(messages) == 1 - - result = TableChunk.from_message(messages[0], br=context.br()) - expected = plc.Table([plc.Column.from_array(values, stream=result.stream)]) - result.stream.synchronize() - assert_eq(result.table_view(), expected) - - -def test_bloom_filter_empty_build_filters_all( - context: Context, comm: Communicator -) -> None: - if comm.nranks != 1: - pytest.skip("Only support single-rank runs") - - stream = context.get_stream_from_pool() - build_table = make_table( - np.array([], dtype=np.int32), stream=stream, br=context.br() - ) - probe_table = make_table( - np.arange(5, dtype=np.int32), stream=stream, br=context.br() - ) - messages = run_bloom_filter_pipeline(context, comm, build_table, probe_table) - assert len(messages) == 1 - - result = TableChunk.from_message(messages[0], br=context.br()) - expected = plc.Table( - [plc.Column.from_array(np.array([], dtype=np.int32), stream=result.stream)] - ) - result.stream.synchronize() - assert_eq(result.table_view(), expected) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py deleted file mode 100644 index 77bdcae78..000000000 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_channel_metadata.py +++ /dev/null @@ -1,390 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 -"""Tests for streaming metadata types (Partitioning and ChannelMetadata).""" - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import pytest - -import pylibcudf as plc - -from rapidsmpf.streaming.core.message import Message -from rapidsmpf.streaming.cudf import ( - ChannelMetadata, - HashScheme, - OrderKey, - OrderScheme, - Partitioning, - TableChunk, -) - -if TYPE_CHECKING: - from rapidsmpf.streaming.core.context import Context - - -def _make_boundaries(context: Context, table: plc.Table) -> TableChunk: - stream = context.get_stream_from_pool() - return TableChunk.from_pylibcudf_table( - table, - stream, - exclusive_view=False, - br=context.br(), - ) - - -def _two_key_order_scheme( - context: Context, *, strict_boundaries: bool = False -) -> OrderScheme: - """Two-key OrderScheme with a 1-row boundary table (2 partitions).""" - boundaries = _make_boundaries( - context, - plc.Table( - [ - plc.Column.from_iterable_of_py([100], plc.DataType(plc.TypeId.INT64)), - plc.Column.from_iterable_of_py( - ["abc"], plc.DataType(plc.TypeId.STRING) - ), - ] - ), - ) - return OrderScheme( - [ - OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), - OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), - ], - boundaries, - strict_boundaries=strict_boundaries, - ) - - -def test_hash_scheme() -> None: - """Test HashScheme construction, properties, equality, and repr.""" - h1 = HashScheme((0, 1), 16) - assert h1.column_indices == (0, 1) - assert h1.modulus == 16 - assert repr(h1) == "HashScheme((0, 1), 16)" - - # Equality - assert h1 == HashScheme((0, 1), 16) - assert h1 != HashScheme((0, 1), 32) - assert h1 != HashScheme((2,), 16) - - -def test_order_key() -> None: - k = OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE) - assert k.column_index == 0 - assert k.order == plc.types.Order.ASCENDING - assert k.null_order == plc.types.NullOrder.BEFORE - assert k == OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE) - assert k != OrderKey(1, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE) - assert "OrderKey" in repr(k) - - -def test_order_scheme(context: Context) -> None: - """Test OrderScheme construction, properties, equality, and repr.""" - o1 = _two_key_order_scheme(context) - assert o1.keys == ( - OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), - OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), - ) - assert not o1.strict_boundaries - assert o1.num_boundaries == 1 - assert "OrderScheme" in repr(o1) - - assert o1.boundaries_aligned_with(_two_key_order_scheme(context), context.br()) - - o_strict = _two_key_order_scheme(context, strict_boundaries=True) - assert o_strict.strict_boundaries - assert not o1.boundaries_aligned_with(o_strict, context.br()) - assert o_strict.boundaries_aligned_with( - _two_key_order_scheme(context, strict_boundaries=True), context.br() - ) - - with pytest.raises(TypeError, match="OrderKey"): - OrderScheme( - [(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE)], # type: ignore[arg-type, list-item] - _make_boundaries( - context, - plc.Table( - [ - plc.Column.from_iterable_of_py( - [0], plc.DataType(plc.TypeId.INT64) - ) - ] - ), - ), - ) - - with pytest.raises(ValueError, match="empty"): - OrderScheme( - [], - _make_boundaries( - context, - plc.Table( - [ - plc.Column.from_iterable_of_py( - [0], plc.DataType(plc.TypeId.INT64) - ) - ] - ), - ), - ) - - -def test_order_scheme_get_boundaries(context: Context) -> None: - scheme = _two_key_order_scheme(context) - chunk = scheme.get_boundaries(context.br()) - assert chunk.table_view().num_columns() == 2 - assert chunk.table_view().num_rows() == 1 - scheme2 = OrderScheme( - scheme.keys, - chunk, - strict_boundaries=scheme.strict_boundaries, - ) - assert scheme2.boundaries_aligned_with(scheme, context.br()) - - -def test_order_scheme_with_keys(context: Context) -> None: - """with_keys shares boundaries and updates column indices.""" - o1 = _two_key_order_scheme(context) - new_keys = [ - OrderKey(5, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), - OrderKey(3, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), - ] - o2 = o1.with_keys(new_keys) - assert o2.keys[0].column_index == 5 - assert o2.keys[1].column_index == 3 - assert o2.num_boundaries == o1.num_boundaries - assert o2.strict_boundaries == o1.strict_boundaries - # Schemes with different key indices but shared boundaries are boundary-aligned - assert o1.boundaries_aligned_with(o2, context.br()) - - -def test_order_scheme_boundaries_aligned_with(context: Context) -> None: - """boundaries_aligned_with performs value-level boundary comparison, ignoring key indices.""" - df = plc.Table( - [ - plc.Column.from_iterable_of_py([100, 200], plc.DataType(plc.TypeId.INT64)), - plc.Column.from_iterable_of_py( - ["abc", "xyz"], plc.DataType(plc.TypeId.STRING) - ), - ] - ) - keys = [ - OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), - OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), - ] - o1 = OrderScheme(keys, _make_boundaries(context, df)) - o2 = OrderScheme(keys, _make_boundaries(context, df)) - assert o1.boundaries_aligned_with(o2, context.br()) - - # Different key column indices but same boundary values → still aligned - shifted_keys = [ - OrderKey(2, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), - OrderKey(3, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), - ] - o_shifted = OrderScheme(shifted_keys, _make_boundaries(context, df)) - assert o1.boundaries_aligned_with(o_shifted, context.br()) - - # Different boundary values → not aligned (shape matches, values differ) - df_diff = plc.Table( - [ - plc.Column.from_iterable_of_py([100, 300], plc.DataType(plc.TypeId.INT64)), - plc.Column.from_iterable_of_py( - ["abc", "xyz"], plc.DataType(plc.TypeId.STRING) - ), - ] - ) - o3 = OrderScheme(keys, _make_boundaries(context, df_diff)) - assert not o1.boundaries_aligned_with(o3, context.br()) - - # Different strict_boundaries → not aligned - o_strict = OrderScheme(keys, _make_boundaries(context, df), strict_boundaries=True) - assert not o1.boundaries_aligned_with(o_strict, context.br()) - - -def test_order_scheme_key_column_mismatch(context: Context) -> None: - """OrderScheme rejects key/column count mismatch.""" - boundaries = _make_boundaries( - context, - plc.Table( - [plc.Column.from_iterable_of_py([0], plc.DataType(plc.TypeId.INT64))] - ), - ) - with pytest.raises(ValueError, match="keys must match"): - OrderScheme( - [ - OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), - OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), - ], - boundaries, # 1 column, but 2 keys - ) - - -def test_partitioning_scenarios(context: Context) -> None: - """Test various partitioning configurations.""" - # Default / None - p_default = Partitioning() - assert p_default.inter_rank is None - assert p_default.local is None - assert Partitioning(None, None).inter_rank is None - assert Partitioning(None, None).local is None - - # Direct global shuffle: inter_rank=Hash, local=Aligned - p_global = Partitioning(HashScheme((0,), 16), "inherit") - assert p_global.inter_rank == HashScheme((0,), 16) - assert p_global.local == "inherit" - - # Two-stage shuffle: inter_rank=Hash(nranks), local=Hash(N_l) - p_twostage = Partitioning(HashScheme((0,), 4), HashScheme((0,), 8)) - assert p_twostage.inter_rank == HashScheme((0,), 4) - assert p_twostage.local == HashScheme((0,), 8) - - # Order-based partitioning (range partitioned / sorted) - order_scheme = _two_key_order_scheme(context) - p_ordered = Partitioning(order_scheme, "inherit") - assert isinstance(p_ordered.inter_rank, OrderScheme) - assert p_ordered.inter_rank.boundaries_aligned_with(order_scheme, context.br()) - assert p_ordered.local == "inherit" - - # Mixed: inter_rank=Order, local=Hash - p_mixed = Partitioning( - _two_key_order_scheme(context), - HashScheme((1,), 8), - ) - assert isinstance(p_mixed.inter_rank, OrderScheme) - assert isinstance(p_mixed.local, HashScheme) - - # Repr - assert "Partitioning" in repr(p_global) - assert "inter_rank" in repr(p_global) - - # Invalid type - with pytest.raises(TypeError): - Partitioning("invalid", None) # type: ignore[arg-type] - - -def test_channel_metadata() -> None: - """Test ChannelMetadata construction and properties.""" - # Basic construction - m = ChannelMetadata(local_count=4) - assert m.local_count == 4 - assert not m.duplicated - - # With partitioning and duplicated - p = Partitioning(HashScheme((0,), 16), "inherit") - m_full = ChannelMetadata(local_count=4, partitioning=p, duplicated=True) - assert m_full.partitioning.inter_rank == HashScheme((0,), 16) - assert m_full.partitioning.local == "inherit" - assert m_full.duplicated - - # Field comparisons (ChannelMetadata.__eq__ removed) - m2 = ChannelMetadata(local_count=4) - assert m.local_count == m2.local_count - assert m.duplicated == m2.duplicated - assert ChannelMetadata(local_count=8).local_count != m.local_count - assert "local_count=4" in repr(m) - - # Validation - with pytest.raises(ValueError, match="local_count must be non-negative"): - ChannelMetadata(local_count=-1) - - -def test_message_roundtrip() -> None: - """Test ChannelMetadata can round-trip through Message.""" - m = ChannelMetadata( - local_count=4, - partitioning=Partitioning(HashScheme((0,), 16), "inherit"), - duplicated=True, - ) - msg_m = Message(99, m) - assert msg_m.sequence_number == 99 - got_m = ChannelMetadata.from_message(msg_m) - assert got_m.local_count == 4 - assert got_m.duplicated - assert got_m.partitioning.inter_rank == HashScheme((0,), 16) - assert msg_m.empty() - - -def test_message_roundtrip_with_order_scheme(context: Context) -> None: - """Test ChannelMetadata with OrderScheme can round-trip through Message.""" - table = plc.Table( - [ - plc.Column.from_iterable_of_py([100, 200], plc.DataType(plc.TypeId.INT64)), - plc.Column.from_iterable_of_py( - ["abc", "xyz"], plc.DataType(plc.TypeId.STRING) - ), - ] - ) - boundaries = _make_boundaries(context, table) - order_scheme = OrderScheme( - [ - OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), - OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), - ], - boundaries, - strict_boundaries=True, - ) - m = ChannelMetadata( - local_count=8, - partitioning=Partitioning(order_scheme, "inherit"), - duplicated=True, - ) - msg_m = Message(42, m) - assert msg_m.sequence_number == 42 - got_m = ChannelMetadata.from_message(msg_m) - assert got_m.local_count == 8 - assert got_m.duplicated - assert isinstance(got_m.partitioning.inter_rank, OrderScheme) - assert got_m.partitioning.inter_rank.keys == ( - OrderKey(0, plc.types.Order.ASCENDING, plc.types.NullOrder.BEFORE), - OrderKey(1, plc.types.Order.DESCENDING, plc.types.NullOrder.AFTER), - ) - assert got_m.partitioning.local == "inherit" - assert got_m.partitioning.inter_rank.strict_boundaries - assert got_m.partitioning.inter_rank.num_boundaries == 2 - assert got_m.partitioning.inter_rank.boundaries_aligned_with( - order_scheme, context.br() - ) - assert msg_m.empty() - - -def test_order_scheme_roundtrip_from_metadata(context: Context) -> None: - """An OrderScheme read back from ChannelMetadata can be re-used in a new Partitioning.""" - src = ChannelMetadata( - local_count=1, - partitioning=Partitioning(_two_key_order_scheme(context), "inherit"), - ) - scheme = src.partitioning.inter_rank - assert isinstance(scheme, OrderScheme) - - p2 = Partitioning(scheme, None) - assert isinstance(p2.inter_rank, OrderScheme) - assert p2.inter_rank.boundaries_aligned_with( - _two_key_order_scheme(context), context.br() - ) - - -def test_access_after_move_raises() -> None: - """Test that accessing a released ChannelMetadata raises ValueError.""" - m = ChannelMetadata( - local_count=4, - partitioning=Partitioning(HashScheme((0,), 16), "inherit"), - ) - # Move into a message (releases the handle) - _ = Message(0, m) - - # Accessing any property should raise ValueError - with pytest.raises(ValueError, match="uninitialized"): - _ = m.local_count - - with pytest.raises(ValueError, match="uninitialized"): - _ = m.partitioning - - with pytest.raises(ValueError, match="uninitialized"): - _ = m.duplicated - - with pytest.raises(ValueError, match="uninitialized"): - repr(m) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_define_actor.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_define_actor.py index b780bb028..955fc0fb8 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_define_actor.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_define_actor.py @@ -5,17 +5,20 @@ from typing import TYPE_CHECKING +import pylibcudf as plc import pytest -import pylibcudf as plc +pytest.importorskip("cudf_streaming") +from cudf_streaming.streaming.table_chunk import TableChunk from rapidsmpf.streaming.chunks.arbitrary import ArbitraryChunk from rapidsmpf.streaming.core.actor import define_actor, run_actor_network from rapidsmpf.streaming.core.leaf_actor import pull_from_channel, push_to_channel from rapidsmpf.streaming.core.message import Message -from rapidsmpf.streaming.cudf.table_chunk import TableChunk from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + @pytest.fixture def expects() -> list[plc.Table]: diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_examples.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_examples.py index 3ca9611de..d67a73bce 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_examples.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_examples.py @@ -1,8 +1,12 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. # SPDX-License-Identifier: Apache-2.0 from __future__ import annotations -from rapidsmpf.examples.streaming import basic_example +import pytest + +cudf = pytest.importorskip("cudf") + +from rapidsmpf.examples.streaming import basic_example # noqa: E402 def test_basic_streaming_example() -> None: diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_fanout.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_fanout.py index 996103860..c41b641f8 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_fanout.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_fanout.py @@ -7,17 +7,20 @@ from typing import TYPE_CHECKING +import pylibcudf as plc import pytest -import pylibcudf as plc +pytest.importorskip("cudf_streaming") +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.streaming.cudf.table_chunk import TableChunk from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + _INT64 = plc.DataType(plc.TypeId.INT64) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_leaf_actor.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_leaf_actor.py index 35587379a..af5d36ee7 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_leaf_actor.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_leaf_actor.py @@ -6,13 +6,18 @@ from typing import TYPE_CHECKING import pylibcudf as plc +import pytest + +pytest.importorskip("cudf_streaming") +from cudf_streaming.streaming.table_chunk import TableChunk from rapidsmpf.streaming.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.streaming.cudf.table_chunk import TableChunk from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + if TYPE_CHECKING: from rmm.pylibrmm.stream import Stream diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_partition.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_partition.py deleted file mode 100644 index 1d95b56e9..000000000 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_partition.py +++ /dev/null @@ -1,84 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import pytest - -import pylibcudf as plc - -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.streaming.cudf.partition import partition_and_pack, unpack_and_concat -from rapidsmpf.streaming.cudf.table_chunk import TableChunk -from rapidsmpf.testing import assert_eq - -if TYPE_CHECKING: - from rmm.pylibrmm.stream import Stream - - from rapidsmpf.streaming.chunks.partition import PartitionMapChunk - from rapidsmpf.streaming.core.channel import Channel - from rapidsmpf.streaming.core.context import Context - - -@pytest.mark.parametrize("num_partitions", [1, 2, 3, 10]) -def test_partition_and_pack_unpack( - context: Context, stream: Stream, num_partitions: int -) -> None: - expects = [ - plc.Table( - [ - plc.Column.from_iterable_of_py( - [1, 2, 3], plc.DataType(plc.TypeId.INT64) - ), - plc.Column.from_iterable_of_py( - [2, 1, 1], plc.DataType(plc.TypeId.INT64) - ), - ] - ), - plc.Table( - [ - plc.Column.from_iterable_of_py([], plc.DataType(plc.TypeId.INT64)), - plc.Column.from_iterable_of_py([], plc.DataType(plc.TypeId.INT64)), - ] - ), - ] - 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) - - ch2: Channel[PartitionMapChunk] = context.create_channel() - actor2 = partition_and_pack( - context, - ch_in=ch1, - ch_out=ch2, - columns_to_hash=(1,), - num_partitions=num_partitions, - ) - - ch3: Channel[TableChunk] = context.create_channel() - actor3 = unpack_and_concat( - context, - ch_in=ch2, - ch_out=ch3, - ) - - actor4, output = pull_from_channel(context, ch_in=ch3) - run_actor_network(context, actors=(actor1, actor2, actor3, actor4)) - - 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, sort_rows=0) diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_read_parquet.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_read_parquet.py deleted file mode 100644 index 087443997..000000000 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_read_parquet.py +++ /dev/null @@ -1,167 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -import itertools -from typing import TYPE_CHECKING - -import numpy as np -import pytest - -import pylibcudf as plc - -from rapidsmpf.streaming.core.actor import run_actor_network -from rapidsmpf.streaming.core.leaf_actor import pull_from_channel -from rapidsmpf.streaming.cudf.parquet import Filter, read_parquet -from rapidsmpf.streaming.cudf.table_chunk import TableChunk - -if TYPE_CHECKING: - from typing import Literal - - from rmm.pylibrmm.stream import Stream - - from rapidsmpf.communicator.communicator import Communicator - from rapidsmpf.streaming.core.actor import CppActor - from rapidsmpf.streaming.core.channel import Channel - from rapidsmpf.streaming.core.context import Context - - -@pytest.fixture(scope="module") -def source( - tmp_path_factory: pytest.TempPathFactory, -) -> plc.io.SourceInfo: - path = tmp_path_factory.mktemp("read_parquet") - - nrows = 10 - start = 0 - sources = [] - for i in range(10): - table = plc.Table( - [plc.Column.from_array(np.arange(start, start + nrows, dtype="int32"))] - ) - # gaps in the column numbering we produce - start += nrows + nrows // 2 - filename = path / f"{i:3d}.pq" - sink = plc.io.SinkInfo([filename]) - options = plc.io.parquet.ParquetWriterOptions.builder(sink, table).build() - plc.io.parquet.write_parquet(options) - sources.append(filename) - return plc.io.SourceInfo(sources) - - -def make_filter(stream: Stream) -> plc.expressions.Expression: - return plc.expressions.Operation( - plc.expressions.ASTOperator.LESS, - plc.expressions.ColumnReference(0), - plc.expressions.Literal( - plc.Scalar.from_py(15, dtype=plc.DataType(plc.TypeId.INT32), stream=stream) - ), - ) - - -def make_producer( - context: Context, - comm: Communicator, - ch: Channel[TableChunk], - options: plc.io.parquet.ParquetReaderOptions, - *, - use_filter: bool, -) -> CppActor: - if use_filter: - fstream = context.get_stream_from_pool() - return read_parquet( - context, comm, ch, 4, options, 3, Filter(fstream, make_filter(fstream)) - ) - else: - return read_parquet(context, comm, ch, 4, options, 3) - - -def get_expected( - ctx: Context, - source: plc.io.SourceInfo, - skip_rows: int | Literal["none"], - num_rows: int | Literal["all"], - *, - use_filter: bool, -) -> plc.Table: - options = plc.io.parquet.ParquetReaderOptions.builder(source).build() - - if skip_rows != "none": - options.set_skip_rows(skip_rows) - if num_rows != "all": - options.set_num_rows(num_rows) - if use_filter: - fstream = ctx.get_stream_from_pool() - filter = make_filter(fstream) - fstream.synchronize() - options.set_filter(filter) - - expected = plc.io.parquet.read_parquet(options).tbl - - if use_filter: - fstream.synchronize() - return expected - - -@pytest.mark.parametrize( - "skip_rows", ["none", 7, 19, 113], ids=lambda s: f"skip_rows_{s}" -) -@pytest.mark.parametrize("num_rows", ["all", 0, 3, 31, 83], ids=lambda s: f"nrows_{s}") -@pytest.mark.parametrize("use_filter", [False, True]) -def test_read_parquet( - context: Context, - comm: Communicator, - source: plc.io.SourceInfo, - skip_rows: int | Literal["none"], - num_rows: int | Literal["all"], - use_filter: bool, # noqa: FBT001 -) -> None: - if comm.nranks != 1: - pytest.skip("Only support single-rank runs") - - ch: Channel[TableChunk] = context.create_channel() - - options = plc.io.parquet.ParquetReaderOptions.builder(source).build() - - if skip_rows != "none": - options.set_skip_rows(skip_rows) - if num_rows != "all": - options.set_num_rows(num_rows) - - producer = make_producer(context, comm, ch, options, use_filter=use_filter) - - consumer, deferred_messages = pull_from_channel(context, ch) - - run_actor_network(context, actors=[producer, consumer]) - - messages = deferred_messages.release() - assert all( - m1.sequence_number < m2.sequence_number - for m1, m2 in itertools.pairwise(messages) - ) - chunks = [TableChunk.from_message(m, br=context.br()) for m in messages] - for chunk in chunks: - chunk.stream.synchronize() - - got = plc.concatenate.concatenate([chunk.table_view() for chunk in chunks]) - for chunk in chunks: - chunk.stream.synchronize() - - expected = get_expected(context, source, skip_rows, num_rows, use_filter=use_filter) - - assert got.num_rows() == expected.num_rows() - assert got.num_columns() == expected.num_columns() - assert got.num_columns() == 1 - - all_equal = plc.reduce.reduce( - plc.binaryop.binary_operation( - got.columns()[0], - expected.columns()[0], - plc.binaryop.BinaryOperator.EQUAL, - plc.DataType(plc.TypeId.BOOL8), - ), - plc.aggregation.all(), - plc.DataType(plc.TypeId.BOOL8), - ) - assert all_equal.to_py() diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py index ee9257293..6f6fd6a9b 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_shuffler.py @@ -7,11 +7,17 @@ import cupy as cp import numpy as np +import pylibcudf as plc import pytest -import pylibcudf as plc +pytest.importorskip("cudf_streaming") +from cudf_streaming.integrations.partition import split_and_pack, unpack_and_concat +from cudf_streaming.streaming.partition import ( + partition_and_pack, + unpack_and_concat as streaming_unpack_and_concat, +) +from cudf_streaming.streaming.table_chunk import TableChunk -from rapidsmpf.integrations.cudf.partition import split_and_pack, unpack_and_concat from rapidsmpf.shuffler import PartitionAssignment from rapidsmpf.streaming.coll.shuffler import ( ShufflerAsync, @@ -20,13 +26,10 @@ 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.streaming.cudf.partition import ( - partition_and_pack, - unpack_and_concat as streaming_unpack_and_concat, -) -from rapidsmpf.streaming.cudf.table_chunk import TableChunk from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + if TYPE_CHECKING: from collections.abc import Awaitable diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_sparse_alltoall.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_sparse_alltoall.py index 2efd16d71..db2bbe566 100644 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_sparse_alltoall.py +++ b/python/rapidsmpf/rapidsmpf/tests/streaming/test_sparse_alltoall.py @@ -7,15 +7,18 @@ from typing import TYPE_CHECKING import numpy as np +import pylibcudf as plc import pytest -import pylibcudf as plc +pytest.importorskip("cudf_streaming") +from cudf_streaming.integrations.partition import unpack_and_concat -from rapidsmpf.integrations.cudf.partition import unpack_and_concat from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.streaming.coll.sparse_alltoall import SparseAlltoall from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + if TYPE_CHECKING: from rapidsmpf.communicator.communicator import Communicator from rapidsmpf.streaming.core.context import Context @@ -24,7 +27,7 @@ def make_packed_data(context: Context, values: np.ndarray) -> PackedData: stream = context.get_stream_from_pool() table = plc.Table([plc.Column.from_array(values, stream=stream)]) - return PackedData.from_cudf_packed_columns( + return PackedData.from_cudf_packed_columns( # type: ignore[attr-defined, no-any-return] plc.contiguous_split.pack(table, stream=stream), stream, context.br(), diff --git a/python/rapidsmpf/rapidsmpf/tests/streaming/test_table_chunk.py b/python/rapidsmpf/rapidsmpf/tests/streaming/test_table_chunk.py deleted file mode 100644 index 12343de34..000000000 --- a/python/rapidsmpf/rapidsmpf/tests/streaming/test_table_chunk.py +++ /dev/null @@ -1,620 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. -# SPDX-License-Identifier: Apache-2.0 - -from __future__ import annotations - -from typing import TYPE_CHECKING - -import cupy -import pytest - -import pylibcudf as plc - -from rapidsmpf.cuda_stream import is_equal_streams -from rapidsmpf.memory.buffer import MemoryType -from rapidsmpf.memory.content_description import ContentDescription -from rapidsmpf.memory.packed_data import PackedData -from rapidsmpf.streaming.core.actor import define_actor, run_actor_network -from rapidsmpf.streaming.core.message import Message -from rapidsmpf.streaming.core.spillable_messages import SpillableMessages -from rapidsmpf.streaming.cudf.table_chunk import ( - TableChunk, - make_table_chunks_available_or_wait, -) -from rapidsmpf.testing import assert_eq - -if TYPE_CHECKING: - from rmm.pylibrmm.stream import Stream - - from rapidsmpf.streaming.core.context import Context - - -def random_table(nbytes: int) -> plc.Table: - assert nbytes % 4 == 0 - return plc.Table( - [plc.Column.from_array(cupy.random.random(nbytes // 4, dtype=cupy.float32))] - ) - - -@pytest.mark.parametrize( - "exclusive_view", - [True, False], -) -def test_roundtrip(context: Context, stream: Stream, *, exclusive_view: bool) -> None: - seq = 42 - expect = random_table(1024) - table_chunk = TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=exclusive_view, br=context.br() - ) - assert is_equal_streams(table_chunk.stream, stream) - assert table_chunk.is_available() - assert table_chunk.make_available_cost() == 0 - assert table_chunk.is_spillable() == exclusive_view - assert_eq(expect, table_chunk.table_view()) - - # Message roundtrip check. - msg1 = Message(seq, table_chunk) - assert msg1.sequence_number == seq - assert msg1.get_content_description() == ContentDescription( - content_sizes={ - MemoryType.DEVICE: 1024, - MemoryType.PINNED_HOST: 0, - MemoryType.HOST: 0, - }, - spillable=exclusive_view, - ) - - # Make a copy of msg1 in host memory. - assert msg1.copy_cost() == 1024 - res, _ = context.br().reserve(MemoryType.HOST, 1024, allow_overbooking=True) - msg2 = msg1.copy(res) - assert res.size == 0 - - # msg1 is availabe - table_chunk2 = TableChunk.from_message(msg1, br=context.br()) - assert is_equal_streams(table_chunk2.stream, stream) - assert table_chunk2.is_available() - assert table_chunk2.make_available_cost() == 0 - assert_eq(expect, table_chunk2.table_view()) - - # Make a copy of msg2 back to device memory. - assert msg2.copy_cost() == 1024 - res, _ = context.br().reserve(MemoryType.DEVICE, 1024, allow_overbooking=True) - msg3 = msg2.copy(res) - assert res.size == 0 - - # msg2 is on host and is not availabe - table_chunk3 = TableChunk.from_message(msg2, br=context.br()) - assert is_equal_streams(table_chunk3.stream, stream) - assert not table_chunk3.is_available() - assert table_chunk3.make_available_cost() == 1024 - # but we can make its table available using `make_available()`. - res, _ = context.br().reserve(MemoryType.DEVICE, 1024, allow_overbooking=True) - table_chunk4 = table_chunk3.make_available(res) - assert is_equal_streams(table_chunk4.stream, stream) - assert table_chunk4.is_available() - assert table_chunk4.make_available_cost() == 0 - assert_eq(expect, table_chunk4.table_view()) - - # msg3 is on device (was created by copying the host msg2). During the copy this - # is made available trivially. - table_chunk5 = TableChunk.from_message(msg3, br=context.br()) - assert is_equal_streams(table_chunk5.stream, stream) - assert table_chunk5.is_available() - # and it cost no device memory to make available. - assert table_chunk5.make_available_cost() == 0 - res, _ = context.br().reserve(MemoryType.DEVICE, 0, allow_overbooking=True) - table_chunk6 = table_chunk5.make_available(res) - assert table_chunk6.is_available() - assert table_chunk6.make_available_cost() == 0 - assert_eq(expect, table_chunk6.table_view()) - - -def test_copy_roundtrip(context: Context, stream: Stream) -> None: - for nrows, ncols in [(1, 1), (1000, 100), (1, 1000)]: - expect = plc.Table( - [ - plc.Column.from_array(cupy.random.random(nrows, dtype=cupy.float32)) - for _ in range(ncols) - ] - ) - - tbl1 = TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ) - res, _ = context.br().reserve( - MemoryType.HOST, - tbl1.data_alloc_size(MemoryType.DEVICE), - allow_overbooking=True, - ) - tbl2 = tbl1.copy(res) - res, _ = context.br().reserve( - MemoryType.DEVICE, tbl2.make_available_cost(), allow_overbooking=True - ) - tbl3 = tbl2.make_available(res) - assert_eq(expect, tbl3.table_view()) - - -def test_spillable_messages(context: Context, stream: Stream) -> None: - seq = 42 - df1 = random_table(1024) - df2 = random_table(2048) - - sm = SpillableMessages(context.br()) - sm.insert( - Message( - seq, - TableChunk.from_pylibcudf_table( - df1, stream, exclusive_view=True, br=context.br() - ), - ) - ) - assert sm.get_content_descriptions() == { - 0: ContentDescription( - content_sizes={ - MemoryType.DEVICE: 1024, - MemoryType.PINNED_HOST: 0, - MemoryType.HOST: 0, - }, - spillable=True, - ) - } - sm.insert( - Message( - seq, - TableChunk.from_pylibcudf_table( - df2, stream, exclusive_view=False, br=context.br() - ), - ) - ) - assert sm.get_content_descriptions() == { - 0: ContentDescription( - content_sizes={ - MemoryType.DEVICE: 1024, - MemoryType.PINNED_HOST: 0, - MemoryType.HOST: 0, - }, - spillable=True, - ), - 1: ContentDescription( - content_sizes={ - MemoryType.DEVICE: 2048, - MemoryType.PINNED_HOST: 0, - MemoryType.HOST: 0, - }, - spillable=False, - ), - } - assert sm.spill(mid=0, br=context.br()) == 1024 - assert sm.get_content_descriptions() == { - 0: ContentDescription( - content_sizes={ - MemoryType.DEVICE: 0, - MemoryType.PINNED_HOST: 0, - MemoryType.HOST: 1024, - }, - spillable=True, - ), - 1: ContentDescription( - content_sizes={ - MemoryType.DEVICE: 2048, - MemoryType.PINNED_HOST: 0, - MemoryType.HOST: 0, - }, - spillable=False, - ), - } - assert sm.spill(mid=1, br=context.br()) == 0 - assert sm.get_content_descriptions() == { - 0: ContentDescription( - content_sizes={ - MemoryType.DEVICE: 0, - MemoryType.PINNED_HOST: 0, - MemoryType.HOST: 1024, - }, - spillable=True, - ), - 1: ContentDescription( - content_sizes={ - MemoryType.DEVICE: 2048, - MemoryType.PINNED_HOST: 0, - MemoryType.HOST: 0, - }, - spillable=False, - ), - } - - # Extract, make available, and check table chunk 1. - df1_got = TableChunk.from_message(sm.extract(mid=0), br=context.br()) - res, _ = context.br().reserve( - MemoryType.DEVICE, df1_got.make_available_cost(), allow_overbooking=True - ) - df1_got = df1_got.make_available(res) - assert_eq(df1, df1_got.table_view()) - - with pytest.raises(IndexError, match="Invalid key"): - sm.extract(mid=0) - - df2_got = TableChunk.from_message(sm.extract(mid=1), br=context.br()) - df2_got = df2_got.make_available_and_spill(context.br(), allow_overbooking=True) - assert_eq(df2, df2_got.table_view()) - assert sm.get_content_descriptions() == {} - - -def test_spillable_messages_by_context(context: Context, stream: Stream) -> None: - seq = 42 - expect = random_table(1024) - - mid = context.spillable_messages().insert( - Message( - seq, - TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ), - ) - ) - assert context.spillable_messages().get_content_descriptions() == { - 0: ContentDescription( - content_sizes={ - MemoryType.DEVICE: 1024, - MemoryType.PINNED_HOST: 0, - MemoryType.HOST: 0, - }, - spillable=True, - ) - } - got = TableChunk.from_message( - context.spillable_messages().extract(mid=mid), br=context.br() - ) - assert_eq(expect, got.table_view()) - - -def test_make_available_or_wait_already_available( - context: Context, stream: Stream -) -> None: - expect = random_table(1024) - chunk = TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ) - result_holder: list[TableChunk] = [] - - @define_actor() - async def test_actor(ctx: Context) -> None: - result = await chunk.make_available_or_wait(ctx, net_memory_delta=0) - result_holder.append(result) - - run_actor_network(context, actors=[test_actor(context)]) - assert_eq(expect, result_holder[0].table_view()) - - -@pytest.mark.parametrize("net_memory_delta", [0, 512]) -def test_make_available_or_wait_from_host( - context: Context, - stream: Stream, - *, - net_memory_delta: int, -) -> None: - expect = random_table(1024) - device_chunk = TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ) - res, _ = context.br().reserve( - MemoryType.HOST, - device_chunk.data_alloc_size(MemoryType.DEVICE), - allow_overbooking=True, - ) - host_chunk = device_chunk.copy(res) - result_holder: list[TableChunk] = [] - - @define_actor() - async def test_actor(ctx: Context) -> None: - result = await host_chunk.make_available_or_wait( - ctx, net_memory_delta=net_memory_delta - ) - result_holder.append(result) - - run_actor_network(context, actors=[test_actor(context)]) - assert_eq(expect, result_holder[0].table_view()) - - -def test_data_alloc_size(context: Context, stream: Stream) -> None: - # Create a table chunk on device memory. - expect = random_table(1024) - device_chunk = TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ) - - # Check device memory size. - assert device_chunk.data_alloc_size(MemoryType.DEVICE) == 1024 - assert device_chunk.data_alloc_size(MemoryType.HOST) == 0 - assert device_chunk.data_alloc_size(MemoryType.PINNED_HOST) == 0 - - # Check that None returns the total across all memory types. - total_size = device_chunk.data_alloc_size(None) - assert total_size == 1024 - assert total_size == ( - device_chunk.data_alloc_size(MemoryType.DEVICE) - + device_chunk.data_alloc_size(MemoryType.HOST) - + device_chunk.data_alloc_size(MemoryType.PINNED_HOST) - ) - - # Check that calling without arguments (default None) works the same. - assert device_chunk.data_alloc_size() == 1024 - assert device_chunk.data_alloc_size() == device_chunk.data_alloc_size(None) - - # Copy to host memory and verify memory distribution. - res, _ = context.br().reserve( - MemoryType.HOST, - device_chunk.data_alloc_size(MemoryType.DEVICE), - allow_overbooking=True, - ) - host_chunk = device_chunk.copy(res) - - assert host_chunk.data_alloc_size(MemoryType.DEVICE) == 0 - assert host_chunk.data_alloc_size(MemoryType.HOST) == 1024 - assert host_chunk.data_alloc_size(MemoryType.PINNED_HOST) == 0 - - # Check that None still returns the correct total. - total_size = host_chunk.data_alloc_size(None) - assert total_size == 1024 - assert total_size == ( - host_chunk.data_alloc_size(MemoryType.DEVICE) - + host_chunk.data_alloc_size(MemoryType.HOST) - + host_chunk.data_alloc_size(MemoryType.PINNED_HOST) - ) - - # Verify default parameter works after copy too. - assert host_chunk.data_alloc_size() == 1024 - assert host_chunk.data_alloc_size() == host_chunk.data_alloc_size(None) - - -@pytest.mark.parametrize("from_pack", [False, True], ids=["from_table", "from_pack"]) -def test_shape_accessor(context: Context, stream: Stream, from_pack: bool) -> None: # noqa: FBT001 - nrows = 64 - expect = plc.Table( - [ - plc.Column.from_iterable_of_py( - ("abc" for _ in range(nrows)), stream=stream - ), - plc.Column.from_iterable_of_py(range(nrows), stream=stream), - ] - ) - expected_shape = (expect.num_rows(), expect.num_columns()) - - if from_pack: - pd = PackedData.from_cudf_packed_columns( - plc.contiguous_split.pack(expect, stream), stream, context.br() - ) - device_chunk = TableChunk.from_packed_data(pd, br=context.br()) - else: - device_chunk = TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ) - assert device_chunk.is_available() - assert device_chunk.shape == expected_shape - - res, _ = context.br().reserve( - MemoryType.HOST, - device_chunk.data_alloc_size(MemoryType.DEVICE), - allow_overbooking=True, - ) - host_chunk = device_chunk.copy(res) - assert not host_chunk.is_available() - assert host_chunk.shape == expected_shape - - res, _ = context.br().reserve( - MemoryType.DEVICE, host_chunk.make_available_cost(), allow_overbooking=True - ) - device_chunk = host_chunk.make_available(res) - assert device_chunk.is_available() - assert device_chunk.shape == expected_shape - - -@pytest.mark.parametrize("from_pack", [False, True], ids=["from_table", "from_pack"]) -def test_into_packed_data(context: Context, stream: Stream, from_pack: bool) -> None: # noqa: FBT001 - expect = random_table(1024) - if from_pack: - pd = PackedData.from_cudf_packed_columns( - plc.contiguous_split.pack(expect, stream), stream, context.br() - ) - chunk = TableChunk.from_packed_data(pd, br=context.br()) - else: - chunk = TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ) - assert chunk.is_available() - - result = chunk.into_packed_data(context.br()) - assert isinstance(result, PackedData) - - # Wrap the PackedData back into a TableChunk and verify contents. - result_chunk = TableChunk.from_packed_data(result, br=context.br()) - assert result_chunk.is_available() - assert_eq(expect, result_chunk.table_view()) - - -@pytest.mark.parametrize("chunk_location", ["device", "host"]) -def test_make_table_chunks_available_or_wait_single_chunk( - context: Context, - stream: Stream, - *, - chunk_location: str, -) -> None: - expect = random_table(1024) - device_chunk = TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ) - - if chunk_location == "host": - res_holder, _ = context.br().reserve( - MemoryType.HOST, - device_chunk.data_alloc_size(MemoryType.DEVICE), - allow_overbooking=True, - ) - chunk = device_chunk.copy(res_holder) - else: - chunk = device_chunk - - result_holder: list[tuple] = [] - - @define_actor() - async def test_actor(ctx: Context) -> None: - result_chunk, res = await make_table_chunks_available_or_wait( - ctx, chunk, reserve_extra=0, net_memory_delta=0 - ) - result_holder.append((result_chunk, res)) - - run_actor_network(context, actors=[test_actor(context)]) - chunk, res = result_holder[0] - assert chunk.is_available() - assert_eq(expect, chunk.table_view()) - # Reservation should be consumed by making the chunk available. - assert res.size == 0 - - -@pytest.mark.parametrize("num_chunks", [1, 2, 3, 5]) -def test_make_table_chunks_available_or_wait_multiple_chunks( - context: Context, - stream: Stream, - *, - num_chunks: int, -) -> None: - # Create multiple chunks with different sizes. - sizes = [1024, 2048, 512, 768, 1536][:num_chunks] - expects = [random_table(size) for size in sizes] - - # Create host chunks. - device_chunks = [ - TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ) - for expect in expects - ] - - host_chunks = [] - for device_chunk in device_chunks: - res, _ = context.br().reserve( - MemoryType.HOST, - device_chunk.data_alloc_size(MemoryType.DEVICE), - allow_overbooking=True, - ) - host_chunks.append(device_chunk.copy(res)) - - result_holder: list[tuple] = [] - - @define_actor() - async def test_actor(ctx: Context) -> None: - chunks, res = await make_table_chunks_available_or_wait( - ctx, - host_chunks, - reserve_extra=0, - net_memory_delta=0, - ) - result_holder.append((chunks, res)) - - run_actor_network(context, actors=[test_actor(context)]) - chunks, res = result_holder[0] - assert len(chunks) == num_chunks - assert all(chunk.is_available() for chunk in chunks) - for i, expect in enumerate(expects): - assert_eq(expect, chunks[i].table_view()) - # Reservation should be consumed. - assert res.size == 0 - - -@pytest.mark.parametrize( - "reserve_extra,net_memory_delta,allow_overbooking", - [ - # Test reserve_extra variations. - (0, 0, None), - (512, 0, None), - (1024, 0, None), - # Test net_memory_delta variations. - (0, -1024, None), - (0, 512, None), - (0, 2048, None), - # Test allow_overbooking variations. - (0, 0, True), - (0, 0, False), - ], -) -def test_make_table_chunks_available_or_wait( - context: Context, - stream: Stream, - *, - reserve_extra: int, - net_memory_delta: int, - allow_overbooking: bool | None, -) -> None: - expect = random_table(1024) - device_chunk = TableChunk.from_pylibcudf_table( - expect, stream, exclusive_view=True, br=context.br() - ) - res_holder, _ = context.br().reserve( - MemoryType.HOST, - device_chunk.data_alloc_size(MemoryType.DEVICE), - allow_overbooking=True, - ) - host_chunk = device_chunk.copy(res_holder) - result_holder: list[tuple] = [] - - @define_actor() - async def test_actor(ctx: Context) -> None: - chunk, res = await make_table_chunks_available_or_wait( - ctx, - host_chunk, - reserve_extra=reserve_extra, - net_memory_delta=net_memory_delta, - allow_overbooking=allow_overbooking, - ) - result_holder.append((chunk, res)) - - run_actor_network(context, actors=[test_actor(context)]) - chunk, res = result_holder[0] - assert chunk.is_available() - assert_eq(expect, chunk.table_view()) - # Reservation should have reserve_extra bytes remaining. - assert res.size == reserve_extra - - -def test_make_table_chunks_available_or_wait_mixed_availability( - context: Context, stream: Stream -) -> None: - expect1 = random_table(1024) - expect2 = random_table(2048) - - # First chunk is already available on device. - available_chunk = TableChunk.from_pylibcudf_table( - expect1, stream, exclusive_view=True, br=context.br() - ) - - # Second chunk is on host memory. - device_chunk2 = TableChunk.from_pylibcudf_table( - expect2, stream, exclusive_view=True, br=context.br() - ) - res2, _ = context.br().reserve( - MemoryType.HOST, - device_chunk2.data_alloc_size(MemoryType.DEVICE), - allow_overbooking=True, - ) - host_chunk = device_chunk2.copy(res2) - result_holder: list[tuple] = [] - - @define_actor() - async def test_actor(ctx: Context) -> None: - chunks, res = await make_table_chunks_available_or_wait( - ctx, - [available_chunk, host_chunk], - reserve_extra=0, - net_memory_delta=0, - ) - result_holder.append((chunks, res)) - - run_actor_network(context, actors=[test_actor(context)]) - chunks, res = result_holder[0] - assert len(chunks) == 2 - assert all(chunk.is_available() for chunk in chunks) - assert_eq(expect1, chunks[0].table_view()) - assert_eq(expect2, chunks[1].table_view()) - # Only the host chunk required device memory. - assert res.size == 0 diff --git a/python/rapidsmpf/rapidsmpf/tests/test_allgather.py b/python/rapidsmpf/rapidsmpf/tests/test_allgather.py index 93ea25f34..244bb6c1b 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_allgather.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_allgather.py @@ -8,17 +8,20 @@ from typing import TYPE_CHECKING import numpy as np +import pylibcudf as plc import pytest -import pylibcudf as plc +pytest.importorskip("cudf_streaming") +from cudf_streaming.integrations.partition import unpack_and_concat from pylibcudf.contiguous_split import pack from rapidsmpf.coll import AllGather -from rapidsmpf.integrations.cudf.partition import unpack_and_concat from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + if TYPE_CHECKING: import rmm.mr from rmm.pylibrmm.stream import Stream @@ -54,7 +57,7 @@ def generate_packed_data( values = np.arange(offset, offset + n_elements, dtype=np.int32) table = plc.Table([plc.Column.from_array(values, stream=stream)]) packed_columns = pack(table, stream=stream) - return PackedData.from_cudf_packed_columns(packed_columns, stream, br) + return PackedData.from_cudf_packed_columns(packed_columns, stream, br) # type: ignore[attr-defined, no-any-return] def validate_packed_data( diff --git a/python/rapidsmpf/rapidsmpf/tests/test_examples.py b/python/rapidsmpf/rapidsmpf/tests/test_examples.py index 128bb66e3..f46be682e 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_examples.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_examples.py @@ -6,13 +6,16 @@ from typing import TYPE_CHECKING import numpy as np +import pylibcudf as plc import pytest -import pylibcudf as plc +pytest.importorskip("cudf_streaming") from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + MPI = pytest.importorskip("mpi4py.MPI") from rapidsmpf.examples.bulk_mpi_shuffle import bulk_mpi_shuffle # noqa: E402 diff --git a/python/rapidsmpf/rapidsmpf/tests/test_partition.py b/python/rapidsmpf/rapidsmpf/tests/test_partition.py index 9caa7f46c..8c37b5bed 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_partition.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_partition.py @@ -5,21 +5,25 @@ from typing import TYPE_CHECKING import numpy as np -import pytest - import pylibcudf as plc -from rmm.pylibrmm.stream import DEFAULT_STREAM +import pytest -from rapidsmpf.integrations.cudf.partition import ( +pytest.importorskip("cudf_streaming") +from cudf_streaming.integrations.partition import ( partition_and_pack, spill_partitions, split_and_pack, unpack_and_concat, unspill_partitions, ) + +from rmm.pylibrmm.stream import DEFAULT_STREAM + from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + if TYPE_CHECKING: import rmm.mr diff --git a/python/rapidsmpf/rapidsmpf/tests/test_ray.py b/python/rapidsmpf/rapidsmpf/tests/test_ray.py index 728f85212..cbdcc4cee 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_ray.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_ray.py @@ -14,6 +14,7 @@ import pytest ray = pytest.importorskip("ray") +cudf = pytest.importorskip("cudf") if TYPE_CHECKING: from collections.abc import Generator diff --git a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py index 6fcbd89c1..c6c8d08f7 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_shuffler.py @@ -6,21 +6,24 @@ from typing import TYPE_CHECKING import numpy as np -import pytest - import pylibcudf as plc +import pytest -from rapidsmpf.integrations.cudf.partition import ( +pytest.importorskip("cudf_streaming") +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 +cudf = pytest.importorskip("cudf") + if TYPE_CHECKING: import rmm.mr from rmm.pylibrmm.stream import Stream diff --git a/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py b/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py index 3990751df..1cd77c47b 100644 --- a/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py +++ b/python/rapidsmpf/rapidsmpf/tests/test_sparse_alltoall.py @@ -7,16 +7,19 @@ from typing import TYPE_CHECKING import numpy as np +import pylibcudf as plc import pytest -import pylibcudf as plc +pytest.importorskip("cudf_streaming") +from cudf_streaming.integrations.partition import unpack_and_concat from rapidsmpf.coll.sparse_alltoall import SparseAlltoall -from rapidsmpf.integrations.cudf.partition import unpack_and_concat from rapidsmpf.memory.buffer_resource import BufferResource from rapidsmpf.memory.packed_data import PackedData from rapidsmpf.testing import assert_eq +cudf = pytest.importorskip("cudf") + if TYPE_CHECKING: import rmm.mr from rmm.pylibrmm.stream import Stream @@ -31,7 +34,7 @@ def generate_packed_data( values = np.arange(offset, offset + n_elements, dtype=np.int32) table = plc.Table([plc.Column.from_array(values, stream=stream)]) packed_columns = plc.contiguous_split.pack(table, stream=stream) - return PackedData.from_cudf_packed_columns(packed_columns, stream, br) + return PackedData.from_cudf_packed_columns(packed_columns, stream, br) # type: ignore[attr-defined, no-any-return] def unpack_table(