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

This commit is contained in:
Philip Hyunsu Cho
2022-10-12 13:57:39 -07:00
committed by GitHub
parent 97a5b088a5
commit 2faa744aba
16 changed files with 190 additions and 117 deletions

View File

@@ -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();