Adding cudf streaming dist tests

ci/check_style.sh

@@ -1,5 +1,5 @@
 #!/bin/bash
-# SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
 set -euo pipefail
@@ -28,4 +28,5 @@ mkdir -p "$(dirname "${RAPIDS_CMAKE_FORMAT_FILE}")"
 wget -O ${RAPIDS_CMAKE_FORMAT_FILE} "${FORMAT_FILE_URL}"
 
 # Run pre-commit checks
+pre-commit clean
 pre-commit run --all-files --show-diff-on-failure

conda/recipes/libcudf/recipe.yaml

@@ -1,4 +1,4 @@
-# SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 schema_version: 1
 
@@ -437,7 +437,7 @@ outputs:
           - libnvcomp
     tests:
       - script:
-        - test -f $PREFIX/bin/gtests/libcudf_streaming/cudf_streaming_single_tests
+        - test -f $PREFIX/bin/gtests/libcudf_streaming/libcudf_streaming_single_tests
         - test -f $PREFIX/bin/libcudf_streaming_bench_pack
         - test -f $PREFIX/bin/libcudf_streaming_bench_partition
         - test -f $PREFIX/bin/libcudf_streaming_bench_ndsh_bench_read

cpp/libcudf_streaming/cmake/ConfigureOptionalCommunication.cmake

@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
@@ -12,7 +12,10 @@ set(CUDF_STREAMING_FOUND_UCXX OFF)
 set(CUDF_STREAMING_HAVE_MPI OFF)
 set(CUDF_STREAMING_HAVE_UCXX OFF)
 
-if(BUILD_BENCHMARKS OR BUILD_EXAMPLES)
+if(BUILD_BENCHMARKS
+   OR BUILD_EXAMPLES
+   OR BUILD_TESTS
+)
   rapids_find_package(
     MPI QUIET
     BUILD_EXPORT_SET cudf_streaming-exports
@@ -23,7 +26,7 @@ if(BUILD_BENCHMARKS OR BUILD_EXAMPLES)
   endif()
 endif()
 
-if(BUILD_BENCHMARKS)
+if(BUILD_BENCHMARKS OR BUILD_TESTS)
   rapids_find_package(
     ucxx QUIET
     BUILD_EXPORT_SET cudf_streaming-exports

cpp/libcudf_streaming/cmake/thirdparty/get_cudf.cmake

@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
@@ -12,6 +12,20 @@ set(CUDF_STREAMING_MIN_VERSION
 )
 find_and_configure_cudf(${CUDF_STREAMING_MIN_VERSION} cudf_streaming-exports)
 
+# The cudf-dependent tests link against cudf::cudftestutil{,_impl}, which live in cudf's optional
+# `testing` component. The transitive find_dependency(cudf) triggered by find_and_configure_cudf
+# does not request it, so request it explicitly here when building tests. This avoids having to
+# build cudf from source: the component is provided by the installed (e.g. conda) cudf package.
+if(BUILD_TESTS AND NOT TARGET cudf::cudftestutil)
+  find_package(cudf ${CUDF_STREAMING_MIN_VERSION} REQUIRED COMPONENTS testing)
+
+  foreach(_target cudf::cudftestutil cudf::cudftestutil_impl)
+    if(NOT TARGET ${_target})
+      message(WARNING "cudf testing target NOT available: ${_target}")
+    endif()
+  endforeach()
+endif()
+
 if(cudf_REQUIRES_CUDA)
   rapids_cuda_init_architectures(CUDF_STREAMING)
 

cpp/libcudf_streaming/tests/CMakeLists.txt

@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
@@ -20,30 +20,63 @@ if(NOT EXISTS "${_cudf_streaming_gtests_link}")
   file(CREATE_LINK "${CUDF_STREAMING_BINARY_DIR}/gtests" "${_cudf_streaming_gtests_link}" SYMBOLIC)
 endif()
 
-add_library(cudf_streaming_test_sources OBJECT)
+# Create a separate test case for each of these n_ranks.
+set(nranks_to_run 1 2 5)
+
+# A helper function to create a test case for each parallelism (defined by nranks_to_run).
+# TEST_TARGET is the name of the test target (required). This adds len(nranks_to_run) test cases
+# identified by <TEST_TARGET>_<n_ranks>. NOTE: When running ctest, the TEST_TARGET executable should
+# be located in the ./gtests/ directory relative to the PWD.
+function(libcudf_streaming_mpirun_test_add)
+  set(options) # no options
+  set(one_value TEST_TARGET)
+  set(multi_value) # no multi_value args
+  cmake_parse_arguments(_MPIRUN_TEST "${options}" "${one_value}" "${multi_value}" ${ARGN})
+
+  if(NOT DEFINED _MPIRUN_TEST_TEST_TARGET)
+    message(FATAL_ERROR "libcudf_streaming_mpirun_test_add called without a test prefix")
+  endif()
+
+  message(STATUS "Adding mpirun test: ${_MPIRUN_TEST_TEST_TARGET} nranks: ${nranks_to_run}")
+  foreach(np IN ITEMS ${nranks_to_run})
+    # add the test using the relative path of the target
+    add_test(
+      NAME "${_MPIRUN_TEST_TEST_TARGET}_${np}" COMMAND mpirun --map-by node --bind-to none -np
+                                                       ${np} "gtests/${_MPIRUN_TEST_TEST_TARGET}"
+    )
+  endforeach(np)
+endfunction(libcudf_streaming_mpirun_test_add)
+
+add_library(libcudf_streaming_test_sources OBJECT)
 set_target_properties(
-  cudf_streaming_test_sources
+  libcudf_streaming_test_sources
   PROPERTIES CXX_STANDARD 20
              CXX_STANDARD_REQUIRED ON
              CXX_EXTENSIONS ON
              CUDA_STANDARD 20
              CUDA_STANDARD_REQUIRED ON
 )
 target_include_directories(
-  cudf_streaming_test_sources PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include"
-                                      "${CMAKE_CURRENT_SOURCE_DIR}"
+  libcudf_streaming_test_sources PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include"
+                                         "${CMAKE_CURRENT_SOURCE_DIR}"
 )
 target_compile_options(
-  cudf_streaming_test_sources
+  libcudf_streaming_test_sources
   PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:${CUDF_CUDA_FLAGS};--expt-extended-lambda>
 )
 target_link_libraries(
-  cudf_streaming_test_sources
+  libcudf_streaming_test_sources
   PRIVATE cudf_streaming rapidsmpf::rapidsmpf cudf::cudftestutil cudf::cudftestutil_impl
   PUBLIC GTest::gmock GTest::gtest
 )
+# cudf::cudftestutil_impl injects cudf test-utility .cu sources (via INTERFACE_SOURCES) that are
+# compiled into this object library. Those sources use device lambdas, so they require the same CUDA
+# flags cudf builds them with (notably --expt-extended-lambda from CUDF_CUDA_FLAGS).
+target_compile_options(
+  libcudf_streaming_test_sources PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:${CUDF_CUDA_FLAGS}>"
+)
 target_sources(
-  cudf_streaming_test_sources
+  libcudf_streaming_test_sources
   PRIVATE streaming/test_table_chunk.cpp
           streaming/test_read_parquet.cpp
           streaming/test_channel_metadata.cpp
@@ -55,9 +88,9 @@ target_sources(
           test_shuffler.cpp
 )
 
-add_executable(cudf_streaming_single_tests main/single.cpp)
+add_executable(libcudf_streaming_single_tests main/single.cpp)
 set_target_properties(
-  cudf_streaming_single_tests
+  libcudf_streaming_single_tests
   PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CUDF_STREAMING_BINARY_DIR}/gtests"
              INSTALL_RPATH "\$ORIGIN/../../../lib"
              CXX_STANDARD 20
@@ -67,22 +100,84 @@ set_target_properties(
              CUDA_STANDARD_REQUIRED ON
 )
 target_include_directories(
-  cudf_streaming_single_tests PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include"
-                                      "${CMAKE_CURRENT_SOURCE_DIR}"
+  libcudf_streaming_single_tests PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include"
+                                         "${CMAKE_CURRENT_SOURCE_DIR}"
 )
 target_link_libraries(
-  cudf_streaming_single_tests
+  libcudf_streaming_single_tests
   PRIVATE cudf_streaming rapidsmpf::rapidsmpf GTest::gmock GTest::gtest
-          $<TARGET_NAME_IF_EXISTS:conda_env> cudf_streaming_test_sources
+          $<TARGET_NAME_IF_EXISTS:conda_env> libcudf_streaming_test_sources
 )
 rapids_test_add(
-  NAME cudf_streaming_single_tests
-  COMMAND cudf_streaming_single_tests
+  NAME libcudf_streaming_single_tests
+  COMMAND libcudf_streaming_single_tests
   GPUS 1
   PERCENT 100
   INSTALL_COMPONENT_SET testing
 )
 
+# ##################################################################################################
+# MPI and UCXX based tests -------------------------------------------------------------------------
+# These executables reuse libcudf_streaming_test_sources but provide a different test environment
+# main (main/mpi.cpp / main/ucxx.cpp), and are launched via mpirun for multiple ranks.
+if(CUDF_STREAMING_HAVE_MPI)
+  add_executable(libcudf_streaming_mpi_tests main/mpi.cpp)
+  set_target_properties(
+    libcudf_streaming_mpi_tests
+    PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CUDF_STREAMING_BINARY_DIR}/gtests"
+               INSTALL_RPATH "\$ORIGIN/../../../lib"
+               CXX_STANDARD 20
+               CXX_STANDARD_REQUIRED ON
+               CXX_EXTENSIONS ON
+               CUDA_STANDARD 20
+               CUDA_STANDARD_REQUIRED ON
+  )
+  target_include_directories(
+    libcudf_streaming_mpi_tests PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include"
+                                        "${CMAKE_CURRENT_SOURCE_DIR}"
+  )
+  target_link_libraries(
+    libcudf_streaming_mpi_tests
+    PRIVATE cudf_streaming rapidsmpf::rapidsmpf cudf_streaming::optional_communication GTest::gmock
+            GTest::gtest $<TARGET_NAME_IF_EXISTS:conda_env> libcudf_streaming_test_sources
+  )
+  libcudf_streaming_mpirun_test_add(TEST_TARGET libcudf_streaming_mpi_tests)
+
+  if(CUDF_STREAMING_HAVE_UCXX)
+    add_executable(libcudf_streaming_ucxx_tests main/ucxx.cpp)
+    set_target_properties(
+      libcudf_streaming_ucxx_tests
+      PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${CUDF_STREAMING_BINARY_DIR}/gtests"
+                 INSTALL_RPATH "\$ORIGIN/../../../lib"
+                 CXX_STANDARD 20
+                 CXX_STANDARD_REQUIRED ON
+                 CXX_EXTENSIONS ON
+                 CUDA_STANDARD 20
+                 CUDA_STANDARD_REQUIRED ON
+    )
+    target_include_directories(
+      libcudf_streaming_ucxx_tests PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include"
+                                           "${CMAKE_CURRENT_SOURCE_DIR}"
+    )
+    target_link_libraries(
+      libcudf_streaming_ucxx_tests
+      PRIVATE cudf_streaming
+              rapidsmpf::rapidsmpf
+              cudf_streaming::optional_communication
+              ucxx::ucxx
+              GTest::gmock
+              GTest::gtest
+              $<TARGET_NAME_IF_EXISTS:conda_env>
+              libcudf_streaming_test_sources
+    )
+    libcudf_streaming_mpirun_test_add(TEST_TARGET libcudf_streaming_ucxx_tests)
+  else()
+    cudf_streaming_skip_optional_target(libcudf_streaming_ucxx_tests "UCXX not found")
+  endif()
+else()
+  cudf_streaming_skip_optional_target(libcudf_streaming_mpi_tests "MPI not found")
+endif()
+
 rapids_test_install_relocatable(
   INSTALL_COMPONENT_SET testing DESTINATION bin/gtests/libcudf_streaming
 )

cpp/libcudf_streaming/tests/environment.hpp

@@ -1,37 +1,50 @@
-/**
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights
- * reserved. SPDX-License-Identifier: Apache-2.0
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
  */
+
 #pragma once
 
 #include <cudf_test/cudf_gtest.hpp>
 
 #include <rapidsmpf/communicator/communicator.hpp>
 
+#include <memory>
+
 enum class TestEnvironmentType : int {
+  MPI,
+  UCXX,
   SINGLE,
 };
 
+/**
+ * Base test environment that exposes a communicator to the test suite. Concrete,
+ * communicator-specific environments (single/MPI/UCXX) derive from this class and are defined in
+ * the corresponding file in main dir, where the communicator-specific state (e.g. the MPI
+ * communicator) and headers live. Exactly one of them is linked into each test executable.
+ */
 class Environment : public ::testing::Environment {
  public:
-  Environment(int argc, char** argv);
+  Environment(int argc, char** argv) : argc_(argc), argv_(argv) {}
+  ~Environment() override = default;
 
-  void SetUp() override;
+  void SetUp() override = 0;
 
-  void TearDown() override;
+  void TearDown() override = 0;
 
-  void barrier();
+  virtual void barrier() = 0;
 
-  [[nodiscard]] TestEnvironmentType type() const;
+  [[nodiscard]] virtual TestEnvironmentType type() const = 0;
 
-  constexpr rapidsmpf::config::Options& options() { return options_; }
+  rapidsmpf::config::Options& options() { return options_; }
 
-  std::shared_ptr<rapidsmpf::Communicator> split_comm();
+  virtual std::shared_ptr<rapidsmpf::Communicator> split_comm() = 0;
 
   std::shared_ptr<rapidsmpf::Communicator> comm_;
 
- private:
-  std::shared_ptr<rapidsmpf::Communicator> split_comm_{nullptr};
+ protected:
+  int argc_;
+  char** argv_;
   rapidsmpf::config::Options options_;
 };