[CI] Test federated learning plugin in the CI (#8325)

tests/cpp/plugin/helpers.cc

@@ -0,0 +1,19 @@
+#include <chrono>
+#include <thread>
+#include <random>
+#include <cstdint>
+
+#include "helpers.h"
+
+using namespace std::chrono_literals;
+
+int GenerateRandomPort(int low, int high) {
+  // Ensure unique timestamp by introducing a small artificial delay
+  std::this_thread::sleep_for(100ms);
+  auto timestamp = static_cast<uint64_t>(std::chrono::duration_cast<std::chrono::milliseconds>(
+    std::chrono::system_clock::now().time_since_epoch()).count());
+  std::mt19937_64 rng(timestamp);
+  std::uniform_int_distribution<int> dist(low, high);
+  int port = dist(rng);
+  return port;
+}

tests/cpp/plugin/helpers.h

@@ -0,0 +1,10 @@
+/*!
+ * Copyright 2022 XGBoost contributors
+ */
+
+#ifndef XGBOOST_TESTS_CPP_PLUGIN_HELPERS_H_
+#define XGBOOST_TESTS_CPP_PLUGIN_HELPERS_H_
+
+int GenerateRandomPort(int low, int high);
+
+#endif  // XGBOOST_TESTS_CPP_PLUGIN_HELPERS_H_

tests/cpp/plugin/test_federated_adapter.cu

@@ -5,24 +5,36 @@
 #include <gtest/gtest.h>
 #include <thrust/host_vector.h>
 
+#include <iostream>
 #include <thread>
+#include <ctime>
 
+#include "./helpers.h"
 #include "../../../plugin/federated/federated_communicator.h"
 #include "../../../plugin/federated/federated_server.h"
 #include "../../../src/collective/device_communicator_adapter.cuh"
 
+namespace {
+
+std::string GetServerAddress() {
+  int port = GenerateRandomPort(50000, 60000);
+  std::string address = std::string("localhost:") + std::to_string(port);
+  return address;
+}
+
+}  // anonymous namespace
+
 namespace xgboost {
 namespace collective {
 
-std::string const kServerAddress{"localhost:56789"};  // NOLINT(cert-err58-cpp)
-
 class FederatedAdapterTest : public ::testing::Test {
  protected:
   void SetUp() override {
+    server_address_ = GetServerAddress();
     server_thread_.reset(new std::thread([this] {
       grpc::ServerBuilder builder;
       federated::FederatedService service{kWorldSize};
-      builder.AddListeningPort(kServerAddress, grpc::InsecureServerCredentials());
+      builder.AddListeningPort(server_address_, grpc::InsecureServerCredentials());
       builder.RegisterService(&service);
       server_ = builder.BuildAndStart();
       server_->Wait();
@@ -35,6 +47,7 @@ class FederatedAdapterTest : public ::testing::Test {
   }
 
   static int const kWorldSize{2};
+  std::string server_address_;
   std::unique_ptr<std::thread> server_thread_;
   std::unique_ptr<grpc::Server> server_;
 };
@@ -52,9 +65,10 @@ TEST(FederatedAdapterSimpleTest, ThrowOnInvalidCommunicator) {
 TEST_F(FederatedAdapterTest, DeviceAllReduceSum) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {
-    threads.emplace_back(std::thread([rank] {
-      FederatedCommunicator comm{kWorldSize, rank, kServerAddress};
-      DeviceCommunicatorAdapter adapter{rank, &comm};
+    threads.emplace_back(std::thread([rank, server_address=server_address_] {
+      FederatedCommunicator comm{kWorldSize, rank, server_address};
+      // Assign device 0 to all workers, since we run gtest in a single-GPU machine
+      DeviceCommunicatorAdapter adapter{0, &comm};
       int const count = 3;
       thrust::device_vector<double> buffer(count, 0);
       thrust::sequence(buffer.begin(), buffer.end());
@@ -74,9 +88,10 @@ TEST_F(FederatedAdapterTest, DeviceAllReduceSum) {
 TEST_F(FederatedAdapterTest, DeviceAllGatherV) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {
-    threads.emplace_back(std::thread([rank] {
-      FederatedCommunicator comm{kWorldSize, rank, kServerAddress};
-      DeviceCommunicatorAdapter adapter{rank, &comm};
+    threads.emplace_back(std::thread([rank, server_address=server_address_] {
+      FederatedCommunicator comm{kWorldSize, rank, server_address};
+      // Assign device 0 to all workers, since we run gtest in a single-GPU machine
+      DeviceCommunicatorAdapter adapter{0, &comm};
 
       int const count = rank + 2;
       thrust::device_vector<char> buffer(count, 0);

tests/cpp/plugin/test_federated_communicator.cc

@@ -5,34 +5,46 @@
 #include <grpcpp/server_builder.h>
 #include <gtest/gtest.h>
 
+#include <iostream>
 #include <thread>
+#include <ctime>
 
+#include "helpers.h"
 #include "../../../plugin/federated/federated_communicator.h"
 #include "../../../plugin/federated/federated_server.h"
 
+namespace {
+
+std::string GetServerAddress() {
+  int port = GenerateRandomPort(50000, 60000);
+  std::string address = std::string("localhost:") + std::to_string(port);
+  return address;
+}
+
+}  // anonymous namespace
+
 namespace xgboost {
 namespace collective {
 
-std::string const kServerAddress{"localhost:56789"};  // NOLINT(cert-err58-cpp)
-
 class FederatedCommunicatorTest : public ::testing::Test {
  public:
-  static void VerifyAllreduce(int rank) {
-    FederatedCommunicator comm{kWorldSize, rank, kServerAddress};
+  static void VerifyAllreduce(int rank, const std::string& server_address) {
+    FederatedCommunicator comm{kWorldSize, rank, server_address};
     CheckAllreduce(comm);
   }
 
-  static void VerifyBroadcast(int rank) {
-    FederatedCommunicator comm{kWorldSize, rank, kServerAddress};
+  static void VerifyBroadcast(int rank, const std::string& server_address) {
+    FederatedCommunicator comm{kWorldSize, rank, server_address};
     CheckBroadcast(comm, rank);
   }
 
  protected:
   void SetUp() override {
+    server_address_ = GetServerAddress();
     server_thread_.reset(new std::thread([this] {
       grpc::ServerBuilder builder;
       federated::FederatedService service{kWorldSize};
-      builder.AddListeningPort(kServerAddress, grpc::InsecureServerCredentials());
+      builder.AddListeningPort(server_address_, grpc::InsecureServerCredentials());
       builder.RegisterService(&service);
       server_ = builder.BuildAndStart();
       server_->Wait();
@@ -66,29 +78,40 @@ class FederatedCommunicatorTest : public ::testing::Test {
   }
 
   static int const kWorldSize{3};
+  std::string server_address_;
   std::unique_ptr<std::thread> server_thread_;
   std::unique_ptr<grpc::Server> server_;
 };
 
 TEST(FederatedCommunicatorSimpleTest, ThrowOnWorldSizeTooSmall) {
-  auto construct = []() { FederatedCommunicator comm{0, 0, kServerAddress, "", "", ""}; };
+  std::string server_address{GetServerAddress()};
+  auto construct = [server_address]() {
+    FederatedCommunicator comm{0, 0, server_address, "", "", ""};
+  };
   EXPECT_THROW(construct(), dmlc::Error);
 }
 
 TEST(FederatedCommunicatorSimpleTest, ThrowOnRankTooSmall) {
-  auto construct = []() { FederatedCommunicator comm{1, -1, kServerAddress, "", "", ""}; };
+  std::string server_address{GetServerAddress()};
+  auto construct = [server_address]() {
+    FederatedCommunicator comm{1, -1, server_address, "", "", ""};
+  };
   EXPECT_THROW(construct(), dmlc::Error);
 }
 
 TEST(FederatedCommunicatorSimpleTest, ThrowOnRankTooBig) {
-  auto construct = []() { FederatedCommunicator comm{1, 1, kServerAddress, "", "", ""}; };
+  std::string server_address{GetServerAddress()};
+  auto construct = [server_address]() {
+    FederatedCommunicator comm{1, 1, server_address, "", "", ""};
+  };
   EXPECT_THROW(construct(), dmlc::Error);
 }
 
 TEST(FederatedCommunicatorSimpleTest, ThrowOnWorldSizeNotInteger) {
-  auto construct = []() {
+  std::string server_address{GetServerAddress()};
+  auto construct = [server_address]() {
     Json config{JsonObject()};
-    config["federated_server_address"] = kServerAddress;
+    config["federated_server_address"] = server_address;
     config["federated_world_size"] = std::string("1");
     config["federated_rank"] = Integer(0);
     auto *comm = FederatedCommunicator::Create(config);
@@ -97,9 +120,10 @@ TEST(FederatedCommunicatorSimpleTest, ThrowOnWorldSizeNotInteger) {
 }
 
 TEST(FederatedCommunicatorSimpleTest, ThrowOnRankNotInteger) {
-  auto construct = []() {
+  std::string server_address{GetServerAddress()};
+  auto construct = [server_address]() {
     Json config{JsonObject()};
-    config["federated_server_address"] = kServerAddress;
+    config["federated_server_address"] = server_address;
     config["federated_world_size"] = 1;
     config["federated_rank"] = std::string("0");
     auto *comm = FederatedCommunicator::Create(config);
@@ -108,20 +132,23 @@ TEST(FederatedCommunicatorSimpleTest, ThrowOnRankNotInteger) {
 }
 
 TEST(FederatedCommunicatorSimpleTest, GetWorldSizeAndRank) {
-  FederatedCommunicator comm{6, 3, kServerAddress};
+  std::string server_address{GetServerAddress()};
+  FederatedCommunicator comm{6, 3, server_address};
   EXPECT_EQ(comm.GetWorldSize(), 6);
   EXPECT_EQ(comm.GetRank(), 3);
 }
 
 TEST(FederatedCommunicatorSimpleTest, IsDistributed) {
-  FederatedCommunicator comm{2, 1, kServerAddress};
+  std::string server_address{GetServerAddress()};
+  FederatedCommunicator comm{2, 1, server_address};
   EXPECT_TRUE(comm.IsDistributed());
 }
 
 TEST_F(FederatedCommunicatorTest, Allreduce) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {
-    threads.emplace_back(std::thread(&FederatedCommunicatorTest::VerifyAllreduce, rank));
+    threads.emplace_back(
+        std::thread(&FederatedCommunicatorTest::VerifyAllreduce, rank, server_address_));
   }
   for (auto &thread : threads) {
     thread.join();
@@ -131,7 +158,8 @@ TEST_F(FederatedCommunicatorTest, Allreduce) {
 TEST_F(FederatedCommunicatorTest, Broadcast) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {
-    threads.emplace_back(std::thread(&FederatedCommunicatorTest::VerifyBroadcast, rank));
+    threads.emplace_back(
+        std::thread(&FederatedCommunicatorTest::VerifyBroadcast, rank, server_address_));
   }
   for (auto &thread : threads) {
     thread.join();

tests/cpp/plugin/test_federated_server.cc

@@ -4,32 +4,45 @@
 #include <grpcpp/server_builder.h>
 #include <gtest/gtest.h>
 
+#include <iostream>
 #include <thread>
+#include <ctime>
 
+#include "helpers.h"
 #include "federated_client.h"
 #include "federated_server.h"
 
+namespace {
+
+std::string GetServerAddress() {
+  int port = GenerateRandomPort(50000, 60000);
+  std::string address = std::string("localhost:") + std::to_string(port);
+  return address;
+}
+
+}  // anonymous namespace
+
 namespace xgboost {
 
 class FederatedServerTest : public ::testing::Test {
  public:
-  static void VerifyAllgather(int rank) {
-    federated::FederatedClient client{kServerAddress, rank};
+  static void VerifyAllgather(int rank, const std::string& server_address) {
+    federated::FederatedClient client{server_address, rank};
     CheckAllgather(client, rank);
   }
 
-  static void VerifyAllreduce(int rank) {
-    federated::FederatedClient client{kServerAddress, rank};
+  static void VerifyAllreduce(int rank, const std::string& server_address) {
+    federated::FederatedClient client{server_address, rank};
     CheckAllreduce(client);
   }
 
-  static void VerifyBroadcast(int rank) {
-    federated::FederatedClient client{kServerAddress, rank};
+  static void VerifyBroadcast(int rank, const std::string& server_address) {
+    federated::FederatedClient client{server_address, rank};
     CheckBroadcast(client, rank);
   }
 
-  static void VerifyMixture(int rank) {
-    federated::FederatedClient client{kServerAddress, rank};
+  static void VerifyMixture(int rank, const std::string& server_address) {
+    federated::FederatedClient client{server_address, rank};
     for (auto i = 0; i < 10; i++) {
       CheckAllgather(client, rank);
       CheckAllreduce(client);
@@ -39,10 +52,11 @@ class FederatedServerTest : public ::testing::Test {
 
  protected:
   void SetUp() override {
+    server_address_ = GetServerAddress();
     server_thread_.reset(new std::thread([this] {
       grpc::ServerBuilder builder;
       federated::FederatedService service{kWorldSize};
-      builder.AddListeningPort(kServerAddress, grpc::InsecureServerCredentials());
+      builder.AddListeningPort(server_address_, grpc::InsecureServerCredentials());
       builder.RegisterService(&service);
       server_ = builder.BuildAndStart();
       server_->Wait();
@@ -80,17 +94,15 @@ class FederatedServerTest : public ::testing::Test {
   }
 
   static int const kWorldSize{3};
-  static std::string const kServerAddress;
+  std::string server_address_;
   std::unique_ptr<std::thread> server_thread_;
   std::unique_ptr<grpc::Server> server_;
 };
 
-std::string const FederatedServerTest::kServerAddress{"localhost:56789"};  // NOLINT(cert-err58-cpp)
-
 TEST_F(FederatedServerTest, Allgather) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {
-    threads.emplace_back(std::thread(&FederatedServerTest::VerifyAllgather, rank));
+    threads.emplace_back(std::thread(&FederatedServerTest::VerifyAllgather, rank, server_address_));
   }
   for (auto& thread : threads) {
     thread.join();
@@ -100,7 +112,7 @@ TEST_F(FederatedServerTest, Allgather) {
 TEST_F(FederatedServerTest, Allreduce) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {
-    threads.emplace_back(std::thread(&FederatedServerTest::VerifyAllreduce, rank));
+    threads.emplace_back(std::thread(&FederatedServerTest::VerifyAllreduce, rank, server_address_));
   }
   for (auto& thread : threads) {
     thread.join();
@@ -110,7 +122,7 @@ TEST_F(FederatedServerTest, Allreduce) {
 TEST_F(FederatedServerTest, Broadcast) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {
-    threads.emplace_back(std::thread(&FederatedServerTest::VerifyBroadcast, rank));
+    threads.emplace_back(std::thread(&FederatedServerTest::VerifyBroadcast, rank, server_address_));
   }
   for (auto& thread : threads) {
     thread.join();
@@ -120,7 +132,7 @@ TEST_F(FederatedServerTest, Broadcast) {
 TEST_F(FederatedServerTest, Mixture) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {
-    threads.emplace_back(std::thread(&FederatedServerTest::VerifyMixture, rank));
+    threads.emplace_back(std::thread(&FederatedServerTest::VerifyMixture, rank, server_address_));
   }
   for (auto& thread : threads) {
     thread.join();