Build libcudf_streaming comm benchmarks with partial MPI/UCXX availability

cpp/libcudf_streaming/CMakeLists.txt

@@ -61,29 +61,12 @@ message(VERBOSE "CUDF_STREAMING: Build examples: ${BUILD_EXAMPLES}")
 
 # add third party dependencies using CPM
 rapids_cpm_init()
+include(cmake/ConfigureOptionalCommunication.cmake)
 include(cmake/thirdparty/get_cudf.cmake)
 include(cmake/thirdparty/get_rapidsmpf.cmake)
+cudf_streaming_configure_optional_communication()
 include(../cmake/thirdparty/get_cucollections.cmake)
 
-if(BUILD_BENCHMARKS)
-  if(RAPIDSMPF_HAVE_UCXX)
-    rapids_find_package(ucxx REQUIRED)
-  else()
-    message(
-      STATUS
-        "CUDF_STREAMING: Skipping UCXX-dependent benchmarks (RAPIDSMPF_HAVE_UCXX=${RAPIDSMPF_HAVE_UCXX})"
-    )
-  endif()
-  if(RAPIDSMPF_HAVE_MPI)
-    rapids_find_package(MPI REQUIRED)
-  else()
-    message(
-      STATUS
-        "CUDF_STREAMING: Skipping MPI-dependent benchmarks (RAPIDSMPF_HAVE_MPI=${RAPIDSMPF_HAVE_MPI})"
-    )
-  endif()
-endif()
-
 # ##################################################################################################
 # * library target --------------------------------------------------------------------------------
 add_library(

cpp/libcudf_streaming/benchmarks/CMakeLists.txt

@@ -13,7 +13,7 @@ add_library(bench_utils INTERFACE)
 target_sources(bench_utils INTERFACE utils/random_data.cu)
 target_compile_options(bench_utils INTERFACE $<$<COMPILE_LANGUAGE:CUDA>:--expt-extended-lambda>)
 
-if(RAPIDSMPF_HAVE_MPI AND RAPIDSMPF_HAVE_UCXX)
+if(CUDF_STREAMING_HAVE_COMM)
   add_executable(bench_shuffle "bench_shuffle.cpp")
   set_target_properties(
     bench_shuffle
@@ -25,15 +25,18 @@ if(RAPIDSMPF_HAVE_MPI AND RAPIDSMPF_HAVE_UCXX)
                CUDA_STANDARD_REQUIRED ON
   )
   target_link_libraries(
-    bench_shuffle PRIVATE cudf_streaming rapidsmpf::rapidsmpf ucxx::ucxx MPI::MPI_CXX
-                          $<TARGET_NAME_IF_EXISTS:conda_env> bench_utils
+    bench_shuffle
+    PRIVATE cudf_streaming rapidsmpf::rapidsmpf cudf_streaming::optional_communication
+            $<TARGET_NAME_IF_EXISTS:conda_env> bench_utils
   )
   install(
     TARGETS bench_shuffle
     COMPONENT benchmarking
     DESTINATION bin/benchmarks/libcudf_streaming
     EXCLUDE_FROM_ALL
   )
+else()
+  cudf_streaming_skip_optional_target(bench_shuffle "requires MPI or UCXX")
 endif()
 
 add_executable(bench_partition "bench_partition.cpp")
@@ -79,6 +82,4 @@ install(
   EXCLUDE_FROM_ALL
 )
 
-if(RAPIDSMPF_HAVE_MPI AND RAPIDSMPF_HAVE_UCXX)
-  add_subdirectory(streaming)
-endif()
+add_subdirectory(streaming)

cpp/libcudf_streaming/benchmarks/bench_shuffle.cpp

@@ -4,14 +4,9 @@
  */
 
 #include <cudf_streaming/integrations/partition.hpp>
-#include <mpi.h>
 #include <rapidsmpf/bootstrap/bootstrap.hpp>
-#include <rapidsmpf/bootstrap/ucxx.hpp>
 #include <rapidsmpf/bootstrap/utils.hpp>
 #include <rapidsmpf/communicator/communicator.hpp>
-#include <rapidsmpf/communicator/mpi.hpp>
-#include <rapidsmpf/communicator/ucxx.hpp>
-#include <rapidsmpf/communicator/ucxx_utils.hpp>
 #include <rapidsmpf/error.hpp>
 #include <rapidsmpf/memory/spill.hpp>
 #include <rapidsmpf/nvtx.hpp>
@@ -30,6 +25,21 @@
 #include <rapidsmpf/cupti.hpp>
 #endif
 
+#ifdef CUDF_STREAMING_HAVE_MPI
+#include <mpi.h>
+#include <rapidsmpf/communicator/mpi.hpp>
+#endif
+
+#ifdef CUDF_STREAMING_HAVE_UCXX
+#include <rapidsmpf/bootstrap/ucxx.hpp>
+#include <rapidsmpf/communicator/ucxx.hpp>
+#endif
+
+#if defined(CUDF_STREAMING_HAVE_MPI) && defined(CUDF_STREAMING_HAVE_UCXX)
+#include <rapidsmpf/communicator/ucxx_utils.hpp>
+#endif
+
+#include "utils/comm.hpp"
 #include "utils/misc.hpp"
 #include "utils/random_data.hpp"
 #include "utils/rmm_utils.hpp"
@@ -38,17 +48,31 @@ class ArgumentParser {
  public:
   ArgumentParser(int argc, char* const* argv, bool use_mpi = true)
   {
-    int rank   = 0;
-    int nranks = 1;
+    int rank           = 0;
+    int nranks         = 1;
+    auto abort_or_exit = [&](int code) {
+#ifdef CUDF_STREAMING_HAVE_MPI
+      if (use_mpi) { RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, code)); }
+#endif
+      std::exit(code);
+    };
 
     if (use_mpi) {
+#ifdef CUDF_STREAMING_HAVE_MPI
       RAPIDSMPF_EXPECTS(rapidsmpf::mpi::is_initialized() == true, "MPI is not initialized");
 
       RAPIDSMPF_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &rank));
       RAPIDSMPF_MPI(MPI_Comm_size(MPI_COMM_WORLD, &nranks));
+#else
+      RAPIDSMPF_FAIL("MPI support is not available in this build", std::runtime_error);
+#endif
     } else {
       // When not using MPI, expect to be using bootstrap mode (rrun)
+#ifdef CUDF_STREAMING_HAVE_UCXX
       nranks = rapidsmpf::bootstrap::get_nranks();
+#else
+      RAPIDSMPF_FAIL("UCXX bootstrap support is not available in this build", std::runtime_error);
+#endif
     }
     try {
       int option;
@@ -58,7 +82,9 @@ class ArgumentParser {
             std::stringstream ss;
             ss << "Usage: " << argv[0] << " [options]\n"
                << "Options:\n"
-               << "  -C <comm>  Communicator {mpi, ucxx} (default: mpi)\n"
+               << "  -C <comm>  Communicator {"
+               << cudf_streaming::benchmarks::available_communicators()
+               << "} (default: " << comm_type << ")\n"
                << "  -r <num>   Number of runs (default: 1)\n"
                << "  -w <num>   Number of warmup runs (default: 0)\n"
                << "  -c <num>   Number of columns in the input tables "
@@ -88,23 +114,17 @@ class ArgumentParser {
 #endif
                << "  -h         Display this help message\n";
             if (rank == 0) { std::cerr << ss.str(); }
-            if (use_mpi) {
-              RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, 0));
-            } else {
-              std::exit(0);
-            }
+            abort_or_exit(0);
           } break;
           case 'C':
             comm_type = std::string{optarg};
-            if (!(comm_type == "mpi" || comm_type == "ucxx")) {
+            if (!cudf_streaming::benchmarks::is_communicator_available(comm_type)) {
               if (rank == 0) {
-                std::cerr << "-C (Communicator) must be one of {mpi, ucxx}" << std::endl;
-              }
-              if (use_mpi) {
-                RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1));
-              } else {
-                std::exit(-1);
+                std::cerr << "-C (Communicator) must be one of {"
+                          << cudf_streaming::benchmarks::available_communicators() << "}"
+                          << std::endl;
               }
+              abort_or_exit(-1);
             }
             break;
           case 'r': parse_integer(num_runs, optarg); break;
@@ -122,11 +142,7 @@ class ArgumentParser {
                              "{cuda, pool, async, managed}"
                           << std::endl;
               }
-              if (use_mpi) {
-                RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1));
-              } else {
-                std::exit(-1);
-              }
+              abort_or_exit(-1);
             }
             break;
           case 'l': parse_integer(device_mem_limit_mb, optarg); break;
@@ -141,24 +157,14 @@ class ArgumentParser {
             enable_cupti_monitoring = true;
             break;
 #endif
-          case '?':
-            if (use_mpi) {
-              RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1));
-            } else {
-              std::exit(-1);
-            }
-            break;
+          case '?': abort_or_exit(-1); break;
           default: RAPIDSMPF_FAIL("unknown option", std::invalid_argument);
         }
       }
       if (optind < argc) { RAPIDSMPF_FAIL("unknown option", std::invalid_argument); }
     } catch (std::exception const& e) {
       if (rank == 0) { std::cerr << "Error parsing arguments: " << e.what() << std::endl; }
-      if (use_mpi) {
-        RAPIDSMPF_MPI(MPI_Abort(MPI_COMM_WORLD, -1));
-      } else {
-        std::exit(-1);
-      }
+      abort_or_exit(-1);
     }
 
     local_nbytes = num_columns * num_local_rows * num_local_partitions * sizeof(std::int32_t);
@@ -212,7 +218,7 @@ class ArgumentParser {
   rapidsmpf::shuffler::PartID num_local_partitions{1};
   rapidsmpf::shuffler::PartID num_output_partitions{1};
   std::string rmm_mr{"pool"};
-  std::string comm_type{"mpi"};
+  std::string comm_type{cudf_streaming::benchmarks::default_communicator()};
   std::uint64_t local_nbytes;
   std::uint64_t total_nbytes;
   bool enable_output_discard{false};
@@ -229,9 +235,22 @@ void barrier(std::shared_ptr<rapidsmpf::Communicator>& comm)
 {
   bool use_bootstrap = rapidsmpf::bootstrap::is_running_with_rrun();
   if (!use_bootstrap) {
+#ifdef CUDF_STREAMING_HAVE_MPI
     RAPIDSMPF_MPI(MPI_Barrier(MPI_COMM_WORLD));
+#else
+    RAPIDSMPF_FAIL("MPI barrier requested, but MPI support is not available in this build",
+                   std::runtime_error);
+#endif
   } else {
-    std::dynamic_pointer_cast<rapidsmpf::ucxx::UCXX>(comm)->barrier();
+#ifdef CUDF_STREAMING_HAVE_UCXX
+    auto ucxx = std::dynamic_pointer_cast<rapidsmpf::ucxx::UCXX>(comm);
+    RAPIDSMPF_EXPECTS(
+      ucxx != nullptr, "Expected UCXX communicator when using bootstrap mode", std::runtime_error);
+    ucxx->barrier();
+#else
+    RAPIDSMPF_FAIL("UCXX bootstrap barrier requested, but UCXX support is not available",
+                   std::runtime_error);
+#endif
   }
 }
 
@@ -449,10 +468,22 @@ int main(int argc, char** argv)
   // and ucxx communicators when not using bootstrap mode.
   int provided = 0;
   if (!use_bootstrap) {
+#ifdef CUDF_STREAMING_HAVE_MPI
     RAPIDSMPF_MPI(MPI_Init_thread(&argc, &argv, MPI_THREAD_MULTIPLE, &provided));
 
     RAPIDSMPF_EXPECTS(provided == MPI_THREAD_MULTIPLE,
                       "didn't get the requested thread level support: MPI_THREAD_MULTIPLE");
+#else
+    std::cerr << "Error: this build has no MPI support. Use UCXX bootstrap mode or build with MPI."
+              << std::endl;
+    return 1;
+#endif
+  } else {
+#ifndef CUDF_STREAMING_HAVE_UCXX
+    std::cerr << "Error: this build has no UCXX support. Bootstrap mode is unavailable."
+              << std::endl;
+    return 1;
+#endif
   }
 
   ArgumentParser args{argc, argv, !use_bootstrap};
@@ -484,6 +515,7 @@ int main(int argc, char** argv)
   std::shared_ptr<rapidsmpf::Communicator> comm;
   auto progress_thread = std::make_shared<rapidsmpf::ProgressThread>(stats);
   if (args.comm_type == "mpi") {
+#ifdef CUDF_STREAMING_HAVE_MPI
     if (use_bootstrap) {
       std::cerr << "Error: MPI communicator requires MPI initialization. Don't use with "
                    "rrun or unset RRUN_RANK."
@@ -492,15 +524,29 @@ int main(int argc, char** argv)
     }
     rapidsmpf::mpi::init(&argc, &argv);
     comm = std::make_shared<rapidsmpf::MPI>(MPI_COMM_WORLD, options, progress_thread);
+#else
+    std::cerr << "Error: MPI communicator is not available in this build." << std::endl;
+    return 1;
+#endif
   } else if (args.comm_type == "ucxx") {
+#ifdef CUDF_STREAMING_HAVE_UCXX
     if (use_bootstrap) {
       // Launched with rrun - use bootstrap backend
       comm = rapidsmpf::bootstrap::create_ucxx_comm(
         progress_thread, rapidsmpf::bootstrap::BackendType::AUTO, options);
     } else {
+#ifdef CUDF_STREAMING_HAVE_MPI
       // Launched with mpirun - use MPI bootstrap
       comm = rapidsmpf::ucxx::init_using_mpi(MPI_COMM_WORLD, options, progress_thread);
+#else
+      std::cerr << "Error: UCXX without MPI support requires bootstrap mode." << std::endl;
+      return 1;
+#endif
     }
+#else
+    std::cerr << "Error: UCXX communicator is not available in this build." << std::endl;
+    return 1;
+#endif
   } else {
     std::cerr << "Error: Unknown communicator type: " << args.comm_type << std::endl;
     return 1;
@@ -606,6 +652,8 @@ int main(int argc, char** argv)
   }
 #endif
 
+#ifdef CUDF_STREAMING_HAVE_MPI
   if (!use_bootstrap) { RAPIDSMPF_MPI(MPI_Finalize()); }
+#endif
   return 0;
 }

cpp/libcudf_streaming/benchmarks/streaming/CMakeLists.txt

@@ -5,7 +5,7 @@
 # cmake-format: on
 # =============================================================================
 
-if(RAPIDSMPF_HAVE_MPI AND RAPIDSMPF_HAVE_UCXX)
+if(CUDF_STREAMING_HAVE_COMM)
   add_executable(bench_streaming_shuffle "bench_streaming_shuffle.cpp")
   set_target_properties(
     bench_streaming_shuffle
@@ -16,15 +16,18 @@ if(RAPIDSMPF_HAVE_MPI AND RAPIDSMPF_HAVE_UCXX)
                CUDA_STANDARD_REQUIRED ON
   )
   target_link_libraries(
-    bench_streaming_shuffle PRIVATE cudf_streaming rapidsmpf::rapidsmpf ucxx::ucxx MPI::MPI_CXX
-                                    $<TARGET_NAME_IF_EXISTS:conda_env> bench_utils
+    bench_streaming_shuffle
+    PRIVATE cudf_streaming rapidsmpf::rapidsmpf cudf_streaming::optional_communication
+            $<TARGET_NAME_IF_EXISTS:conda_env> bench_utils
   )
   install(
     TARGETS bench_streaming_shuffle
     COMPONENT benchmarking
     DESTINATION bin/benchmarks/libcudf_streaming
     EXCLUDE_FROM_ALL
   )
-
-  add_subdirectory(ndsh)
+else()
+  cudf_streaming_skip_optional_target(bench_streaming_shuffle "requires MPI or UCXX")
 endif()
+
+add_subdirectory(ndsh)