Merge branch 'main' into enforce-abi3audit

Author: jameslamb

cpp/src/communicator/ucxx.cpp

@@ -18,10 +18,26 @@
 #include <rapidsmpf/error.hpp>
 #include <rapidsmpf/utils/misc.hpp>
 
+#include "ucxx_internal.hpp"
+
 namespace rapidsmpf {
 
 namespace ucxx {
 
+namespace detail {
+
+std::shared_ptr<::ucxx::Context> create_context(ProgressMode progress_mode) {
+    auto feature_flags = ::ucxx::Context::defaultFeatureFlags;
+    if (progress_mode == ProgressMode::Polling
+        || progress_mode == ProgressMode::ThreadPolling)
+    {
+        feature_flags &= ~static_cast<std::uint64_t>(UCP_FEATURE_WAKEUP);
+    }
+    return ::ucxx::contextBuilder(feature_flags).build();
+}
+
+}  // namespace detail
+
 namespace {
 
 /**
@@ -984,8 +1000,7 @@ std::unique_ptr<rapidsmpf::ucxx::InitializedRank> init(
         });
 
     auto create_worker = [progress_mode]() {
-        auto context =
-            ::ucxx::contextBuilder(::ucxx::Context::defaultFeatureFlags).build();
+        auto context = detail::create_context(progress_mode);
         auto worker = context->workerBuilder().build();
 
         RAPIDSMPF_EXPECTS(

cpp/src/communicator/ucxx_internal.hpp

@@ -0,0 +1,17 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <memory>
+
+#include <ucxx/context.h>
+
+#include <rapidsmpf/communicator/ucxx.hpp>
+
+namespace rapidsmpf::ucxx::detail {
+
+[[nodiscard]] std::shared_ptr<::ucxx::Context> create_context(ProgressMode progress_mode);
+
+}  // namespace rapidsmpf::ucxx::detail

cpp/tests/CMakeLists.txt

@@ -157,6 +157,8 @@ if(RAPIDSMPF_HAVE_MPI)
 
   if(RAPIDSMPF_HAVE_UCXX)
     add_executable(ucxx_tests main/ucxx.cpp)
+    target_sources(ucxx_tests PRIVATE test_ucxx.cpp)
+    target_include_directories(ucxx_tests PRIVATE "${PROJECT_SOURCE_DIR}/src")
     set_target_properties(
       ucxx_tests
       PROPERTIES RUNTIME_OUTPUT_DIRECTORY "${RAPIDSMPF_BINARY_DIR}/gtests"

cpp/tests/test_ucxx.cpp

@@ -0,0 +1,49 @@
+/**
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <cstdint>
+
+#include <gtest/gtest.h>
+#include <ucp/api/ucp.h>
+#include <ucxx/context.h>
+
+#include <rapidsmpf/communicator/ucxx.hpp>
+
+#include "communicator/ucxx_internal.hpp"
+
+namespace {
+
+struct ProgressModeFeatureFlagsParam {
+    rapidsmpf::ucxx::ProgressMode progress_mode;
+    bool wakeup_enabled;
+};
+
+class ProgressModeFeatureFlagsTest
+    : public ::testing::TestWithParam<ProgressModeFeatureFlagsParam> {};
+
+TEST_P(ProgressModeFeatureFlagsTest, CreatesContextWithExpectedFeatureFlags) {
+    auto const context =
+        rapidsmpf::ucxx::detail::create_context(GetParam().progress_mode);
+    auto const default_flags = ::ucxx::Context::defaultFeatureFlags;
+    auto const expected_flags =
+        GetParam().wakeup_enabled ? default_flags : default_flags & ~UCP_FEATURE_WAKEUP;
+
+    EXPECT_EQ(context->getFeatureFlags(), expected_flags);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    ProgressModes,
+    ProgressModeFeatureFlagsTest,
+    ::testing::Values(
+        ProgressModeFeatureFlagsParam{rapidsmpf::ucxx::ProgressMode::Blocking, true},
+        ProgressModeFeatureFlagsParam{rapidsmpf::ucxx::ProgressMode::Polling, false},
+        ProgressModeFeatureFlagsParam{
+            rapidsmpf::ucxx::ProgressMode::ThreadBlocking, true
+        },
+        ProgressModeFeatureFlagsParam{rapidsmpf::ucxx::ProgressMode::ThreadPolling, false}
+    )
+);
+
+}  // namespace