Add an in-memory collective communicator (#8494)

This commit is contained in:
Rong Ou
2022-11-30 08:24:12 -08:00
committed by GitHub
parent 157e98edf7
commit a8255ea678
15 changed files with 577 additions and 277 deletions

View File

@@ -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++) {