[CI] Test federated learning plugin in the CI (#8325)
This commit is contained in:
committed by
GitHub
parent
97a5b088a5
commit
2faa744aba
@@ -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();
|
||||
|
||||
Reference in New Issue
Block a user