Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add Allgather to collective communicator (#8765)

Browse Source 
* Add Allgather to collective communicator
This commit is contained in:
Rong Ou
2023-02-08 19:31:22 -08:00
committed by GitHub
parent 48cefa012e
commit cbf98cb9c6
14 changed files with 187 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
tests/cpp/plugin/test_federated_communicator.cc
View File

@@ -28,6 +28,11 @@ namespace collective {
class FederatedCommunicatorTest : public ::testing::Test {
public:
static void VerifyAllgather(int rank, const std::string& server_address) {
FederatedCommunicator comm{kWorldSize, rank, server_address};
CheckAllgather(comm, rank);
}
static void VerifyAllreduce(int rank, const std::string& server_address) {
FederatedCommunicator comm{kWorldSize, rank, server_address};
CheckAllreduce(comm);
@@ -56,6 +61,15 @@ class FederatedCommunicatorTest : public ::testing::Test {
server_thread_->join();
}
static void CheckAllgather(FederatedCommunicator &comm, int rank) {
int buffer[kWorldSize] = {0, 0, 0};
buffer[rank] = rank;
comm.AllGather(buffer, sizeof(buffer));
for (auto i = 0; i < kWorldSize; i++) {
EXPECT_EQ(buffer[i], i);
}
}
static void CheckAllreduce(FederatedCommunicator &comm) {
int buffer[] = {1, 2, 3, 4, 5};
comm.AllReduce(buffer, sizeof(buffer) / sizeof(buffer[0]), DataType::kInt32, Operation::kSum);
@@ -144,6 +158,17 @@ TEST(FederatedCommunicatorSimpleTest, IsDistributed) {
EXPECT_TRUE(comm.IsDistributed());
}
TEST_F(FederatedCommunicatorTest, Allgather) {
std::vector<std::thread> threads;
for (auto rank = 0; rank < kWorldSize; rank++) {
threads.emplace_back(
std::thread(&FederatedCommunicatorTest::VerifyAllgather, rank, server_address_));
}
for (auto &thread : threads) {
thread.join();
}
}
TEST_F(FederatedCommunicatorTest, Allreduce) {
std::vector<std::thread> threads;
for (auto rank = 0; rank < kWorldSize; rank++) {