85 lines
2.4 KiB
C++
85 lines
2.4 KiB
C++
/*!
|
|
* Copyright 2022-2023 XGBoost contributors
|
|
*/
|
|
#pragma once
|
|
|
|
#include <grpcpp/server_builder.h>
|
|
#include <gtest/gtest.h>
|
|
#include <xgboost/json.h>
|
|
|
|
#include <random>
|
|
#include <thread> // for thread, sleep_for
|
|
|
|
#include "../../../plugin/federated/federated_server.h"
|
|
#include "../../../src/collective/communicator-inl.h"
|
|
|
|
namespace xgboost {
|
|
|
|
class ServerForTest {
|
|
std::string server_address_;
|
|
std::unique_ptr<std::thread> server_thread_;
|
|
std::unique_ptr<grpc::Server> server_;
|
|
|
|
public:
|
|
explicit ServerForTest(std::int32_t world_size) {
|
|
server_thread_.reset(new std::thread([this, world_size] {
|
|
grpc::ServerBuilder builder;
|
|
xgboost::federated::FederatedService service{world_size};
|
|
int selected_port;
|
|
builder.AddListeningPort("localhost:0", grpc::InsecureServerCredentials(), &selected_port);
|
|
builder.RegisterService(&service);
|
|
server_ = builder.BuildAndStart();
|
|
server_address_ = std::string("localhost:") + std::to_string(selected_port);
|
|
server_->Wait();
|
|
}));
|
|
}
|
|
|
|
~ServerForTest() {
|
|
server_->Shutdown();
|
|
server_thread_->join();
|
|
}
|
|
|
|
auto Address() const {
|
|
using namespace std::chrono_literals;
|
|
while (server_address_.empty()) {
|
|
std::this_thread::sleep_for(100ms);
|
|
}
|
|
return server_address_;
|
|
}
|
|
};
|
|
|
|
class BaseFederatedTest : public ::testing::Test {
|
|
protected:
|
|
void SetUp() override { server_ = std::make_unique<ServerForTest>(kWorldSize); }
|
|
|
|
void TearDown() override { server_.reset(nullptr); }
|
|
|
|
static int constexpr kWorldSize{3};
|
|
std::unique_ptr<ServerForTest> server_;
|
|
};
|
|
|
|
template <typename Function, typename... Args>
|
|
void RunWithFederatedCommunicator(int32_t world_size, std::string const& server_address,
|
|
Function&& function, Args&&... args) {
|
|
std::vector<std::thread> threads;
|
|
for (auto rank = 0; rank < world_size; rank++) {
|
|
threads.emplace_back([&, rank]() {
|
|
Json config{JsonObject()};
|
|
config["xgboost_communicator"] = String("federated");
|
|
config["federated_server_address"] = String(server_address);
|
|
config["federated_world_size"] = world_size;
|
|
config["federated_rank"] = rank;
|
|
xgboost::collective::Init(config);
|
|
|
|
std::forward<Function>(function)(std::forward<Args>(args)...);
|
|
|
|
xgboost::collective::Finalize();
|
|
});
|
|
}
|
|
for (auto& thread : threads) {
|
|
thread.join();
|
|
}
|
|
}
|
|
|
|
} // namespace xgboost
|