Support test variants (#742)

* TEST: Support test variants

* TEST: CR fix

* TEST: minor fix

* TEST/COMMON: CR fix

* NIXL/TEST: clang

* NIXl/TEST: clang fix

---------

Co-authored-by: Leonid Genkin <[email protected]>

Author: shasson5

src/plugins/ucx/ucx_backend.cpp

@@ -1102,8 +1102,12 @@ nixlUcxEngine::nixlUcxEngine(const nixlBackendInitParams &init_params)
         err_handling_mode = ucx_err_mode_from_string(err_handling_mode_it->second);
     }
 
+    const auto engine_config_it = custom_params->find("engine_config");
+    const auto engine_config =
+        (engine_config_it != custom_params->end()) ? engine_config_it->second : "";
+
     uc = std::make_unique<nixlUcxContext>(
-        devs, init_params.enableProgTh, num_workers, init_params.syncMode);
+        devs, init_params.enableProgTh, num_workers, init_params.syncMode, engine_config);
 
     for (size_t i = 0; i < num_workers; i++) {
         uws.emplace_back(std::make_unique<nixlUcxWorker>(*uc, err_handling_mode));

src/utils/ucx/ucx_utils.cpp

@@ -408,7 +408,8 @@ bool nixlUcxMtLevelIsSupported(const nixl_ucx_mt_t mt_type) noexcept
 nixlUcxContext::nixlUcxContext(std::vector<std::string> devs,
                                bool prog_thread,
                                unsigned long num_workers,
-                               nixl_thread_sync_t sync_mode) {
+                               nixl_thread_sync_t sync_mode,
+                               const std::string &engine_config) {
     ucp_params_t ucp_params;
     unsigned major_version, minor_version, release_number;
     ucp_get_version(&major_version, &minor_version, &release_number);
@@ -454,6 +455,18 @@ nixlUcxContext::nixlUcxContext(std::vector<std::string> devs,
         config.modify("MAX_COMPONENT_MDS", "32");
     }
 
+    std::string elem;
+    std::stringstream stream(engine_config);
+
+    while (std::getline(stream, elem, ',')) {
+        std::string_view elem_view = elem;
+        size_t pos = elem_view.find('=');
+
+        if (pos != std::string::npos) {
+            config.modify(elem_view.substr(0, pos), elem_view.substr(pos + 1));
+        }
+    }
+
     const auto status = ucp_init (&ucp_params, config.getUcpConfig(), &ctx);
     if (status != UCS_OK) {
         throw std::runtime_error ("Failed to create UCX context: " +

src/utils/ucx/ucx_utils.h

@@ -201,7 +201,8 @@ class nixlUcxContext {
     nixlUcxContext(std::vector<std::string> devices,
                    bool prog_thread,
                    unsigned long num_workers,
-                   nixl_thread_sync_t sync_mode);
+                   nixl_thread_sync_t sync_mode,
+                   const std::string &engine_conf = "");
     ~nixlUcxContext();
 
     /* Memory management */

test/gtest/common.h

@@ -25,6 +25,7 @@
 #include <stack>
 #include <optional>
 #include <mutex>
+#include "gtest/gtest.h"
 
 namespace gtest {
 constexpr const char *
@@ -102,6 +103,29 @@ class PortAllocator {
     uint16_t _max_port = MAX_PORT;
 };
 
+struct nixlTestParam {
+    std::string backendName;
+    bool progressThreadEnabled;
+    unsigned numWorkers;
+    unsigned numThreads;
+    std::string engineConfig;
+};
+
+using nixl_test_t = testing::TestWithParam<nixlTestParam>;
+
 } // namespace gtest
 
+#define NIXL_INSTANTIATE_TEST(_test_name,               \
+                              _test_case,               \
+                              _backend,                 \
+                              _progress_thread_enabled, \
+                              _num_workers,             \
+                              _num_threads,             \
+                              _engine_config)           \
+    INSTANTIATE_TEST_SUITE_P(                           \
+        _test_name,                                     \
+        _test_case,                                     \
+        testing::ValuesIn(std::vector<nixlTestParam>(   \
+            {{_backend, _progress_thread_enabled, _num_workers, _num_threads, _engine_config}})));
+
 #endif /* TEST_GTEST_COMMON_H */

test/gtest/test_transfer.cpp

@@ -96,9 +96,7 @@ class MemBuffer : std::shared_ptr<void> {
     const size_t size;
 };
 
-class TestTransfer :
-    // Tuple fields are: backend_name, enable_progress_thread, num_workers, num_threads
-    public testing::TestWithParam<std::tuple<std::string, bool, size_t, size_t>> {
+class TestTransfer : public nixl_test_t {
 protected:
     nixlAgentConfig
     getConfig(int listen_port, bool capture_telemetry) {
@@ -127,6 +125,7 @@ class TestTransfer :
             params["split_batch_size"] = "32";
         }
 
+        params["engine_config"] = GetParam().engineConfig;
         return params;
     }
 
@@ -165,22 +164,22 @@ class TestTransfer :
 
     std::string getBackendName() const
     {
-        return std::get<0>(GetParam());
+        return GetParam().backendName;
     }
 
     bool
     isProgressThreadEnabled() const {
-        return std::get<1>(GetParam());
+        return GetParam().progressThreadEnabled;
     }
 
     size_t
     getNumWorkers() const {
-        return std::get<2>(GetParam());
+        return GetParam().numWorkers;
     }
 
     size_t
     getNumThreads() const {
-        return std::get<3>(GetParam());
+        return GetParam().numThreads;
     }
 
     nixl_opt_args_t
@@ -718,14 +717,9 @@ TEST_P(TestTransfer, PrepGpuSignal) {
 #endif
 }
 
-INSTANTIATE_TEST_SUITE_P(ucx, TestTransfer, testing::Values(std::make_tuple("UCX", true, 2, 0)));
-INSTANTIATE_TEST_SUITE_P(ucx_no_pt,
-                         TestTransfer,
-                         testing::Values(std::make_tuple("UCX", false, 2, 0)));
-INSTANTIATE_TEST_SUITE_P(ucx_threadpool,
-                         TestTransfer,
-                         testing::Values(std::make_tuple("UCX", true, 6, 4)));
-INSTANTIATE_TEST_SUITE_P(ucx_threadpool_no_pt,
-                         TestTransfer,
-                         testing::Values(std::make_tuple("UCX", false, 6, 4)));
+NIXL_INSTANTIATE_TEST(ucx, TestTransfer, "UCX", true, 2, 0, "");
+NIXL_INSTANTIATE_TEST(ucx_no_pt, TestTransfer, "UCX", false, 2, 0, "");
+NIXL_INSTANTIATE_TEST(ucx_threadpool, TestTransfer, "UCX", true, 6, 4, "");
+NIXL_INSTANTIATE_TEST(ucx_threadpool_no_pt, TestTransfer, "UCX", false, 6, 4, "");
+
 } // namespace gtest