Add Allgather to collective communicator (#8765)

* Add Allgather to collective communicator

tests/cpp/plugin/test_federated_communicator.cc

@@ -28,6 +28,11 @@ namespace collective {
 
 class FederatedCommunicatorTest : public ::testing::Test {
  public:
+  static void VerifyAllgather(int rank, const std::string& server_address) {
+    FederatedCommunicator comm{kWorldSize, rank, server_address};
+    CheckAllgather(comm, rank);
+  }
+
   static void VerifyAllreduce(int rank, const std::string& server_address) {
     FederatedCommunicator comm{kWorldSize, rank, server_address};
     CheckAllreduce(comm);
@@ -56,6 +61,15 @@ class FederatedCommunicatorTest : public ::testing::Test {
     server_thread_->join();
   }
 
+  static void CheckAllgather(FederatedCommunicator &comm, int rank) {
+    int buffer[kWorldSize] = {0, 0, 0};
+    buffer[rank] = rank;
+    comm.AllGather(buffer, sizeof(buffer));
+    for (auto i = 0; i < kWorldSize; i++) {
+      EXPECT_EQ(buffer[i], i);
+    }
+  }
+
   static void CheckAllreduce(FederatedCommunicator &comm) {
     int buffer[] = {1, 2, 3, 4, 5};
     comm.AllReduce(buffer, sizeof(buffer) / sizeof(buffer[0]), DataType::kInt32, Operation::kSum);
@@ -144,6 +158,17 @@ TEST(FederatedCommunicatorSimpleTest, IsDistributed) {
   EXPECT_TRUE(comm.IsDistributed());
 }
 
+TEST_F(FederatedCommunicatorTest, Allgather) {
+  std::vector<std::thread> threads;
+  for (auto rank = 0; rank < kWorldSize; rank++) {
+    threads.emplace_back(
+        std::thread(&FederatedCommunicatorTest::VerifyAllgather, rank, server_address_));
+  }
+  for (auto &thread : threads) {
+    thread.join();
+  }
+}
+
 TEST_F(FederatedCommunicatorTest, Allreduce) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {

tests/cpp/plugin/test_federated_server.cc

@@ -4,13 +4,13 @@
 #include <grpcpp/server_builder.h>
 #include <gtest/gtest.h>
 
+#include <ctime>
 #include <iostream>
 #include <thread>
-#include <ctime>
 
-#include "helpers.h"
 #include "federated_client.h"
 #include "federated_server.h"
+#include "helpers.h"
 
 namespace {
 
@@ -26,6 +26,11 @@ namespace xgboost {
 
 class FederatedServerTest : public ::testing::Test {
  public:
+  static void VerifyAllgather(int rank, const std::string& server_address) {
+    federated::FederatedClient client{server_address, rank};
+    CheckAllgather(client, rank);
+  }
+
   static void VerifyAllreduce(int rank, const std::string& server_address) {
     federated::FederatedClient client{server_address, rank};
     CheckAllreduce(client);
@@ -39,6 +44,7 @@ class FederatedServerTest : public ::testing::Test {
   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);
       CheckBroadcast(client, rank);
     }
@@ -62,6 +68,17 @@ class FederatedServerTest : public ::testing::Test {
     server_thread_->join();
   }
 
+  static void CheckAllgather(federated::FederatedClient& client, int rank) {
+    int data[kWorldSize] = {0, 0, 0};
+    data[rank] = rank;
+    std::string send_buffer(reinterpret_cast<char const*>(data), sizeof(data));
+    auto reply = client.Allgather(send_buffer);
+    auto const* result = reinterpret_cast<int const*>(reply.data());
+    for (auto i = 0; i < kWorldSize; i++) {
+      EXPECT_EQ(result[i], i);
+    }
+  }
+
   static void CheckAllreduce(federated::FederatedClient& client) {
     int data[] = {1, 2, 3, 4, 5};
     std::string send_buffer(reinterpret_cast<char const*>(data), sizeof(data));
@@ -88,6 +105,16 @@ class FederatedServerTest : public ::testing::Test {
   std::unique_ptr<grpc::Server> server_;
 };
 
+TEST_F(FederatedServerTest, Allgather) {
+  std::vector<std::thread> threads;
+  for (auto rank = 0; rank < kWorldSize; rank++) {
+    threads.emplace_back(std::thread(&FederatedServerTest::VerifyAllgather, rank, server_address_));
+  }
+  for (auto& thread : threads) {
+    thread.join();
+  }
+}
+
 TEST_F(FederatedServerTest, Allreduce) {
   std::vector<std::thread> threads;
   for (auto rank = 0; rank < kWorldSize; rank++) {