Adding cudf streaming dist tests

cpp/libcudf_streaming/cmake/ConfigureOptionalCommunication.cmake

@@ -12,7 +12,7 @@ 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 +23,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

@@ -12,6 +12,14 @@ 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)
+endif()
+
 if(cudf_REQUIRES_CUDA)
   rapids_cuda_init_architectures(CUDF_STREAMING)
 

cpp/libcudf_streaming/tests/CMakeLists.txt

@@ -20,6 +20,33 @@ if(NOT EXISTS "${_cudf_streaming_gtests_link}")
   file(CREATE_LINK "${CUDF_STREAMING_BINARY_DIR}/gtests" "${_cudf_streaming_gtests_link}" SYMBOLIC)
 endif()
 
+# 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(cudf_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 "cudf_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(cudf_streaming_mpirun_test_add)
+
 add_library(cudf_streaming_test_sources OBJECT)
 set_target_properties(
   cudf_streaming_test_sources
@@ -38,6 +65,13 @@ target_link_libraries(
   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(
+  cudf_streaming_test_sources
+  PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:${CUDF_CUDA_FLAGS}>"
+)
 target_sources(
   cudf_streaming_test_sources
   PRIVATE streaming/test_table_chunk.cpp
@@ -81,6 +115,63 @@ rapids_test_add(
   INSTALL_COMPONENT_SET testing
 )
 
+# ##################################################################################################
+# MPI and UCXX based tests -------------------------------------------------------------------------
+# These executables reuse cudf_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(cudf_streaming_mpi_tests main/mpi.cpp)
+  set_target_properties(
+    cudf_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(
+    cudf_streaming_mpi_tests PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include"
+                                     "${CMAKE_CURRENT_SOURCE_DIR}"
+  )
+  target_link_libraries(
+    cudf_streaming_mpi_tests
+    PRIVATE cudf_streaming rapidsmpf::rapidsmpf cudf_streaming::optional_communication GTest::gmock
+            GTest::gtest $<TARGET_NAME_IF_EXISTS:conda_env> cudf_streaming_test_sources
+  )
+  cudf_streaming_mpirun_test_add(TEST_TARGET cudf_streaming_mpi_tests)
+
+  if(CUDF_STREAMING_HAVE_UCXX)
+    add_executable(cudf_streaming_ucxx_tests main/ucxx.cpp)
+    set_target_properties(
+      cudf_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(
+      cudf_streaming_ucxx_tests PRIVATE "${CUDF_STREAMING_SOURCE_DIR}/include"
+                                        "${CMAKE_CURRENT_SOURCE_DIR}"
+    )
+    target_link_libraries(
+      cudf_streaming_ucxx_tests
+      PRIVATE cudf_streaming rapidsmpf::rapidsmpf cudf_streaming::optional_communication
+              GTest::gmock GTest::gtest $<TARGET_NAME_IF_EXISTS:conda_env>
+              cudf_streaming_test_sources
+    )
+    cudf_streaming_mpirun_test_add(TEST_TARGET cudf_streaming_ucxx_tests)
+  else()
+    cudf_streaming_skip_optional_target(cudf_streaming_ucxx_tests "UCXX not found")
+  endif()
+else()
+  cudf_streaming_skip_optional_target(cudf_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

@@ -8,29 +8,42 @@
 
 #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 `main/*.cpp`, 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:
+ protected:
+  int argc_;
+  char** argv_;
   std::shared_ptr<rapidsmpf::Communicator> split_comm_{nullptr};
   rapidsmpf::config::Options options_;
 };

cpp/libcudf_streaming/tests/main/mpi.cpp

@@ -0,0 +1,82 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights
+ * reserved. SPDX-License-Identifier: Apache-2.0
+ */
+
+#include "../environment.hpp"
+
+#include <cudf_test/cudf_gtest.hpp>
+
+#include <mpi.h>
+#include <rapidsmpf/communicator/mpi.hpp>
+#include <rapidsmpf/config.hpp>
+#include <rapidsmpf/progress_thread.hpp>
+
+#include <memory>
+
+namespace {
+class MPIEnvironment : public Environment {
+ public:
+  using Environment::Environment;
+
+  [[nodiscard]] TestEnvironmentType type() const override { return TestEnvironmentType::MPI; }
+
+  void SetUp() override
+  {
+    rapidsmpf::mpi::init(&argc_, &argv_);
+
+    RAPIDSMPF_MPI(MPI_Comm_dup(MPI_COMM_WORLD, &mpi_comm_));
+
+    options_ = rapidsmpf::config::Options(rapidsmpf::config::get_environment_variables());
+
+    comm_ = std::make_shared<rapidsmpf::MPI>(
+      mpi_comm_, options_, std::make_shared<rapidsmpf::ProgressThread>());
+  }
+
+  void TearDown() override
+  {
+    split_comm_ = nullptr;  // Clean up the split communicator.
+    comm_       = nullptr;  // Clean up the communicator.
+
+    RAPIDSMPF_MPI(MPI_Comm_free(&mpi_comm_));
+    RAPIDSMPF_MPI(MPI_Finalize());
+  }
+
+  void barrier() override { RAPIDSMPF_MPI(MPI_Barrier(mpi_comm_)); }
+
+  std::shared_ptr<rapidsmpf::Communicator> split_comm() override
+  {
+    // Return cached split communicator if it exists
+    if (split_comm_ != nullptr) { return split_comm_; }
+
+    // Initialize configuration options from environment variables.
+    rapidsmpf::config::Options options{rapidsmpf::config::get_environment_variables()};
+
+    // Create the new split communicator
+    int rank;
+    RAPIDSMPF_MPI(MPI_Comm_rank(mpi_comm_, &rank));
+    MPI_Comm split_comm = MPI_COMM_NULL;
+    RAPIDSMPF_MPI(MPI_Comm_split(mpi_comm_, rank, 0, &split_comm));
+    return std::shared_ptr<rapidsmpf::MPI>(
+      new rapidsmpf::MPI(split_comm, options, comm_->progress_thread()),
+      // Don't leak the split handle.
+      [comm = split_comm](rapidsmpf::MPI* x) mutable {
+        delete x;
+        MPI_Comm_free(&comm);
+      });
+  }
+
+ private:
+  MPI_Comm mpi_comm_{MPI_COMM_NULL};
+};
+}  // namespace
+
+Environment* GlobalEnvironment = nullptr;
+
+int main(int argc, char** argv)
+{
+  ::testing::InitGoogleTest(&argc, argv);
+  GlobalEnvironment = new MPIEnvironment(argc, argv);
+  ::testing::AddGlobalTestEnvironment(GlobalEnvironment);
+  return RUN_ALL_TESTS();
+}

cpp/libcudf_streaming/tests/main/single.cpp

@@ -8,38 +8,44 @@
 #include <cudf_test/cudf_gtest.hpp>
 
 #include <rapidsmpf/communicator/single.hpp>
+#include <rapidsmpf/config.hpp>
 #include <rapidsmpf/progress_thread.hpp>
 
 #include <memory>
 
-Environment* GlobalEnvironment = nullptr;
+namespace {
+class SingleEnvironment : public Environment {
+ public:
+  using Environment::Environment;
 
-Environment::Environment(int, char**) {}
+  [[nodiscard]] TestEnvironmentType type() const override { return TestEnvironmentType::SINGLE; }
 
-TestEnvironmentType Environment::type() const { return TestEnvironmentType::SINGLE; }
+  void SetUp() override
+  {
+    options_ = rapidsmpf::config::Options(rapidsmpf::config::get_environment_variables());
+    comm_ =
+      std::make_shared<rapidsmpf::Single>(options_, std::make_shared<rapidsmpf::ProgressThread>());
+    split_comm_ = comm_;
+  }
 
-void Environment::SetUp()
-{
-  options_ = rapidsmpf::config::Options(rapidsmpf::config::get_environment_variables());
-  comm_ =
-    std::make_shared<rapidsmpf::Single>(options_, std::make_shared<rapidsmpf::ProgressThread>());
-  split_comm_ = comm_;
-}
+  void TearDown() override
+  {
+    split_comm_ = nullptr;
+    comm_       = nullptr;
+  }
 
-void Environment::TearDown()
-{
-  split_comm_ = nullptr;
-  comm_       = nullptr;
-}
+  void barrier() override {}
 
-void Environment::barrier() {}
+  std::shared_ptr<rapidsmpf::Communicator> split_comm() override { return split_comm_; }
+};
+}  // namespace
 
-std::shared_ptr<rapidsmpf::Communicator> Environment::split_comm() { return split_comm_; }
+Environment* GlobalEnvironment = nullptr;
 
 int main(int argc, char** argv)
 {
   ::testing::InitGoogleTest(&argc, argv);
-  GlobalEnvironment = new Environment(argc, argv);
+  GlobalEnvironment = new SingleEnvironment(argc, argv);
   ::testing::AddGlobalTestEnvironment(GlobalEnvironment);
   return RUN_ALL_TESTS();
 }