diff --git a/test/gtest/common.h b/test/gtest/common.h index a7c835f9fb..bfbc88ac3d 100644 --- a/test/gtest/common.h +++ b/test/gtest/common.h @@ -210,6 +210,7 @@ struct nixlTestParam { unsigned numWorkers; unsigned numThreads; std::string engineConfig; + bool sglEnabled = false; }; using nixl_test_t = testing::TestWithParam; @@ -237,4 +238,21 @@ checkPayload(const void *addr, size_t size, uint32_t index, uint32_t seed); testing::ValuesIn(std::vector( \ {{_backend, _progress_thread_enabled, _num_workers, _num_threads, _engine_config}}))); +#define NIXL_INSTANTIATE_TEST_SGL(_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({{_backend, \ + _progress_thread_enabled, \ + _num_workers, \ + _num_threads, \ + _engine_config, \ + true}}))); + #endif /* TEST_GTEST_COMMON_H */ diff --git a/test/gtest/test_transfer.cpp b/test/gtest/test_transfer.cpp index eb5732290d..0164d79cf2 100644 --- a/test/gtest/test_transfer.cpp +++ b/test/gtest/test_transfer.cpp @@ -156,6 +156,8 @@ class TestTransfer : public nixl_test_t { // Disabling Telemetry until the corresponding test env.addVar("NIXL_TELEMETRY_ENABLE", "n"); + env.addVar("NIXL_UCX_SGL_ENABLE", isSglEnabled() ? "y" : "n"); + if (getBackendName() == "MOONCAKE") { // TCP loopback on host memory: CI hosts may expose RDMA devices the // test cannot use, and CI builds Mooncake without GPU support. @@ -195,6 +197,11 @@ class TestTransfer : public nixl_test_t { return GetParam().numThreads; } + bool + isSglEnabled() const { + return GetParam().sglEnabled; + } + nixl_opt_args_t extra_params_ip(int remote) { nixl_opt_args_t extra_params; @@ -813,6 +820,14 @@ NIXL_INSTANTIATE_TEST(ucx_threadpool_no_pt, TestTransfer, "UCX", false, 6, 4, "" NIXL_INSTANTIATE_TEST(mooncake, TestTransfer, "MOONCAKE", false, 1, 0, ""); #endif +#ifdef HAVE_UCX_SGL_API + +NIXL_INSTANTIATE_TEST_SGL(ucx_sgl, TestTransfer, "UCX", true, 2, 0, ""); +NIXL_INSTANTIATE_TEST_SGL(ucx_sgl_no_pt, TestTransfer, "UCX", false, 2, 0, ""); +NIXL_INSTANTIATE_TEST_SGL(ucx_sgl_threadpool, TestTransfer, "UCX", true, 6, 4, ""); +NIXL_INSTANTIATE_TEST_SGL(ucx_sgl_threadpool_no_pt, TestTransfer, "UCX", false, 6, 4, ""); +#endif + NIXL_INSTANTIATE_TEST(ucx_telemetry, TestTransferTelemetry, "UCX", true, 2, 0, ""); NIXL_INSTANTIATE_TEST(ucx_telemetry_no_pt, TestTransferTelemetry, "UCX", false, 2, 0, ""); NIXL_INSTANTIATE_TEST(ucx_telemetry_threadpool, TestTransferTelemetry, "UCX", true, 6, 4, "");