Merge branch 'branch-25.10' into shuffler-callbacks

Author: nirandaperera

.github/workflows/pr.yaml

@@ -17,6 +17,7 @@ jobs:
       - conda-cpp-build
       - conda-cpp-linters
       - conda-cpp-tests
+      - conda-cpp-memcheck
       - conda-python-build
       - conda-python-tests
       - docs-build
@@ -121,6 +122,14 @@ jobs:
       build_type: pull-request
       container-options: "--cap-add CAP_SYS_PTRACE --shm-size=8g --ulimit=nofile=1000000:1000000"
       script: ci/test_cpp.sh
+  conda-cpp-memcheck:
+    secrets: inherit
+    needs: conda-cpp-build
+    uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10
+    with:
+      build_type: pull-request
+      script: "ci/test_cpp_memcheck.sh"
+      node_type: "gpu-l4-latest-1"
   conda-python-build:
     needs: conda-cpp-build
     secrets: inherit

.github/workflows/test.yaml

@@ -33,6 +33,14 @@ jobs:
       date: ${{ inputs.date }}
       script: ci/test_cpp.sh
       sha: ${{ inputs.sha }}
+  conda-cpp-memcheck:
+    secrets: inherit
+    needs: checks
+    uses: rapidsai/shared-workflows/.github/workflows/custom-job.yaml@branch-25.10
+    with:
+      build_type: pull-request
+      script: "ci/test_cpp_memcheck.sh"
+      node_type: "gpu-l4-latest-1"
   conda-python-tests:
     secrets: inherit
     uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-25.10

ci/test_cpp_memcheck.sh

@@ -0,0 +1,46 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025, NVIDIA CORPORATION & AFFILIATES.
+# SPDX-License-Identifier: Apache-2.0
+
+set -xeuo pipefail
+
+. /opt/conda/etc/profile.d/conda.sh
+
+CPP_CHANNEL=$(rapids-download-conda-from-github cpp)
+
+rapids-logger "Generate C++ testing dependencies"
+rapids-dependency-file-generator \
+  --output conda \
+  --file-key test_cpp \
+  --matrix "cuda=${RAPIDS_CUDA_VERSION%.*};arch=$(arch)" \
+  --prepend-channel "${CPP_CHANNEL}" \
+  | tee env.yaml
+
+rapids-mamba-retry env create --yes -f env.yaml -n test
+
+# Temporarily allow unbound variables for conda activation.
+set +u
+conda activate test
+set -u
+
+RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
+mkdir -p "${RAPIDS_TESTS_DIR}"
+
+rapids-print-env
+
+rapids-logger "Check GPU usage"
+nvidia-smi
+
+# Trap ERR so that `EXITCODE=1` is set when a command fails
+EXITCODE=0
+trap "EXITCODE=1" ERR
+set +e
+
+# Support customizing the ctests' install location
+cd "${INSTALL_PREFIX:-${CONDA_PREFIX:-/usr}}/bin/tests/librapidsmpf/"
+
+rapids-logger "Run librapidsmpf gtests with compute-sanitizer (Single Node)"
+compute-sanitizer --tool memcheck --track-stream-ordered-races=all gtests/single_tests --gtest_filter=-CuptiMonitorTest.*
+
+rapids-logger "Test script exiting with value: $EXITCODE"
+exit ${EXITCODE}

cpp/CMakeLists.txt

@@ -114,11 +114,6 @@ include(../cmake/thirdparty/get_nvtx.cmake)
 include(../cmake/thirdparty/get_rmm.cmake)
 include(../cmake/thirdparty/get_cudf.cmake)
 if(RAPIDSMPF_HAVE_UCXX)
-  rapids_find_package(
-    UCX REQUIRED
-    BUILD_EXPORT_SET rapidsmpf-exports
-    INSTALL_EXPORT_SET rapidsmpf-exports
-  )
   rapids_find_package(
     ucxx REQUIRED
     BUILD_EXPORT_SET rapidsmpf-exports
@@ -292,17 +287,7 @@ endif()
 # ##################################################################################################
 # * add tests -------------------------------------------------------------------------------------
 if(RAPIDSMPF_BUILD_TESTS)
-  # include CTest module -- automatically calls enable_testing()
   include(CTest)
-
-  # ctest cuda memcheck
-  find_program(CUDA_SANITIZER compute-sanitizer)
-  set(MEMORYCHECK_COMMAND ${CUDA_SANITIZER})
-  set(MEMORYCHECK_TYPE CudaSanitizer)
-  set(CUDA_SANITIZER_COMMAND_OPTIONS "--tool memcheck")
-
-  # Always print verbose output when tests fail if run using `make test`.
-  list(APPEND CMAKE_CTEST_ARGUMENTS "--output-on-failure")
   add_subdirectory(tests)
 endif()
 

cpp/benchmarks/CMakeLists.txt

@@ -22,9 +22,8 @@ target_compile_options(
                         "$<$<COMPILE_LANGUAGE:CUDA>:${RAPIDSMPF_CUDA_FLAGS}>"
 )
 target_link_libraries(
-  bench_shuffle
-  PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx ucx::ucp $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>
-          $<TARGET_NAME_IF_EXISTS:conda_env> maybe_asan bench_utils
+  bench_shuffle PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>
+                        $<TARGET_NAME_IF_EXISTS:conda_env> maybe_asan bench_utils
 )
 install(
   TARGETS bench_shuffle
@@ -50,7 +49,7 @@ target_compile_options(
                      "$<$<COMPILE_LANGUAGE:CUDA>:${RAPIDSMPF_CUDA_FLAGS}>"
 )
 target_link_libraries(
-  bench_comm PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx ucx::ucp $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>
+  bench_comm PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>
                      $<TARGET_NAME_IF_EXISTS:conda_env> maybe_asan bench_utils
 )
 install(
@@ -84,7 +83,6 @@ target_link_libraries(
   bench_partition
   PRIVATE rapidsmpf::rapidsmpf
           ucxx::ucxx
-          ucx::ucp
           benchmark::benchmark
           benchmark::benchmark_main
           $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>

cpp/benchmarks/streaming/CMakeLists.txt

@@ -19,7 +19,7 @@ target_compile_options(
 )
 target_link_libraries(
   bench_streaming_shuffle
-  PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx ucx::ucp $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>
+  PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>
           $<TARGET_NAME_IF_EXISTS:conda_env> maybe_asan bench_utils
 )
 install(

cpp/examples/CMakeLists.txt

@@ -20,9 +20,8 @@ target_compile_options(
                           "$<$<COMPILE_LANGUAGE:CUDA>:${RAPIDSMPF_CUDA_FLAGS}>"
 )
 target_link_libraries(
-  example_shuffle
-  PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx ucx::ucp $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>
-          $<TARGET_NAME_IF_EXISTS:conda_env> maybe_asan
+  example_shuffle PRIVATE rapidsmpf::rapidsmpf ucxx::ucxx $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>
+                          $<TARGET_NAME_IF_EXISTS:conda_env> maybe_asan
 )
 install(
   TARGETS example_shuffle

cpp/tests/CMakeLists.txt

@@ -121,7 +121,6 @@ if(RAPIDSMPF_HAVE_MPI)
               GTest::gmock
               GTest::gtest
               ucxx::ucxx
-              ucx::ucp
               $<TARGET_NAME_IF_EXISTS:MPI::MPI_C>
               $<TARGET_NAME_IF_EXISTS:conda_env>
               maybe_asan

cpp/tests/test_rmm_resource_adaptor.cpp

@@ -37,8 +37,8 @@ struct throw_at_limit_resource final : public rmm::mr::device_memory_resource {
         return ptr;
     }
 
-    void do_deallocate(void* ptr, std::size_t, rmm::cuda_stream_view) override {
-        RAPIDSMPF_CUDA_TRY(cudaFree(ptr));
+    void do_deallocate(void* ptr, std::size_t, rmm::cuda_stream_view stream) override {
+        RAPIDSMPF_CUDA_TRY(cudaFreeAsync(ptr, stream.value()));
         allocs.erase(ptr);
     }
 

python/rapidsmpf/rapidsmpf/communicator/CMakeLists.txt

@@ -17,7 +17,7 @@ if(RAPIDSMPF_HAVE_UCXX)
   rapids_cython_create_modules(
     CXX
     SOURCE_FILES "${modules_need_ucxx}"
-    LINKED_LIBRARIES rapidsmpf::rapidsmpf PRIVATE ucxx::ucxx ucxx::python ucx::ucp maybe_asan
+    LINKED_LIBRARIES rapidsmpf::rapidsmpf PRIVATE ucxx::ucxx ucxx::python maybe_asan
   )
 endif()