Add an in-memory collective communicator (#8494)

tests/cpp/collective/test_in_memory_communicator.cc

@@ -0,0 +1,112 @@
+/*!
+ * Copyright 2022 XGBoost contributors
+ */
+#include <dmlc/parameter.h>
+#include <gtest/gtest.h>
+
+#include <thread>
+
+#include "../../../src/collective/in_memory_communicator.h"
+
+namespace xgboost {
+namespace collective {
+
+class InMemoryCommunicatorTest : public ::testing::Test {
+ public:
+  static void VerifyAllreduce(int rank) {
+    InMemoryCommunicator comm{kWorldSize, rank};
+    int buffer[] = {1, 2, 3, 4, 5};
+    comm.AllReduce(buffer, sizeof(buffer) / sizeof(buffer[0]), DataType::kInt32, Operation::kSum);
+    int expected[] = {3, 6, 9, 12, 15};
+    for (auto i = 0; i < 5; i++) {
+      EXPECT_EQ(buffer[i], expected[i]);
+    }
+  }
+
+  static void VerifyBroadcast(int rank) {
+    InMemoryCommunicator comm{kWorldSize, rank};
+    if (rank == 0) {
+      std::string buffer{"hello"};
+      comm.Broadcast(&buffer[0], buffer.size(), 0);
+      EXPECT_EQ(buffer, "hello");
+    } else {
+      std::string buffer{"     "};
+      comm.Broadcast(&buffer[0], buffer.size(), 0);
+      EXPECT_EQ(buffer, "hello");
+    }
+  }
+
+ protected:
+  static int const kWorldSize{3};
+};
+
+TEST(InMemoryCommunicatorSimpleTest, ThrowOnWorldSizeTooSmall) {
+  auto construct = []() { InMemoryCommunicator comm{0, 0}; };
+  EXPECT_THROW(construct(), dmlc::Error);
+}
+
+TEST(InMemoryCommunicatorSimpleTest, ThrowOnRankTooSmall) {
+  auto construct = []() { InMemoryCommunicator comm{1, -1}; };
+  EXPECT_THROW(construct(), dmlc::Error);
+}
+
+TEST(InMemoryCommunicatorSimpleTest, ThrowOnRankTooBig) {
+  auto construct = []() { InMemoryCommunicator comm{1, 1}; };
+  EXPECT_THROW(construct(), dmlc::Error);
+}
+
+TEST(InMemoryCommunicatorSimpleTest, ThrowOnWorldSizeNotInteger) {
+  auto construct = []() {
+    Json config{JsonObject()};
+    config["in_memory_world_size"] = std::string("1");
+    config["in_memory_rank"] = Integer(0);
+    auto *comm = InMemoryCommunicator::Create(config);
+    delete comm;
+  };
+  EXPECT_THROW(construct(), dmlc::Error);
+}
+
+TEST(InMemoryCommunicatorSimpleTest, ThrowOnRankNotInteger) {
+  auto construct = []() {
+    Json config{JsonObject()};
+    config["in_memory_world_size"] = 1;
+    config["in_memory_rank"] = std::string("0");
+    auto *comm = InMemoryCommunicator::Create(config);
+    delete comm;
+  };
+  EXPECT_THROW(construct(), dmlc::Error);
+}
+
+TEST(InMemoryCommunicatorSimpleTest, GetWorldSizeAndRank) {
+  InMemoryCommunicator comm{1, 0};
+  EXPECT_EQ(comm.GetWorldSize(), 1);
+  EXPECT_EQ(comm.GetRank(), 0);
+}
+
+TEST(InMemoryCommunicatorSimpleTest, IsDistributed) {
+  InMemoryCommunicator comm{1, 0};
+  EXPECT_TRUE(comm.IsDistributed());
+}
+
+TEST_F(InMemoryCommunicatorTest, Allreduce) {
+  std::vector<std::thread> threads;
+  for (auto rank = 0; rank < kWorldSize; rank++) {
+    threads.emplace_back(std::thread(&InMemoryCommunicatorTest::VerifyAllreduce, rank));
+  }
+  for (auto &thread : threads) {
+    thread.join();
+  }
+}
+
+TEST_F(InMemoryCommunicatorTest, Broadcast) {
+  std::vector<std::thread> threads;
+  for (auto rank = 0; rank < kWorldSize; rank++) {
+    threads.emplace_back(std::thread(&InMemoryCommunicatorTest::VerifyBroadcast, rank));
+  }
+  for (auto &thread : threads) {
+    thread.join();
+  }
+}
+
+}  // namespace collective
+}  // namespace xgboost

tests/cpp/plugin/test_federated_server.cc

@@ -26,11 +26,6 @@ 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);
@@ -44,7 +39,6 @@ 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);
     }
@@ -68,11 +62,6 @@ class FederatedServerTest : public ::testing::Test {
     server_thread_->join();
   }
 
-  static void CheckAllgather(federated::FederatedClient& client, int rank) {
-    auto reply = client.Allgather("hello " + std::to_string(rank) + " ");
-    EXPECT_EQ(reply, "hello 0 hello 1 hello 2 ");
-  }
-
   static void CheckAllreduce(federated::FederatedClient& client) {
     int data[] = {1, 2, 3, 4, 5};
     std::string send_buffer(reinterpret_cast<char const*>(data), sizeof(data));
@@ -90,7 +79,7 @@ class FederatedServerTest : public ::testing::Test {
       send_buffer = "hello broadcast";
     }
     auto reply = client.Broadcast(send_buffer, 0);
-    EXPECT_EQ(reply, "hello broadcast");
+    EXPECT_EQ(reply, "hello broadcast") << "rank " << rank;
   }
 
   static int const kWorldSize{3};
@@ -99,16 +88,6 @@ 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++) {