8bad677c2f
* Update collective implementation. - Cleanup resource during `Finalize` to avoid handling threads in destructor. - Calculate the size for allgather automatically. - Use simple allgather for small (smaller than the number of worker) allreduce.
94 lines
3.4 KiB
C++
94 lines
3.4 KiB
C++
/**
|
|
* Copyright 2022-2023, XGBoost contributors
|
|
*/
|
|
#include <gtest/gtest.h>
|
|
#include <xgboost/span.h> // for Span
|
|
|
|
#include <array> // for array
|
|
|
|
#include "../../../../src/common/type.h" // for EraseType
|
|
#include "../../collective/test_worker.h" // for SocketTest
|
|
#include "federated_coll.h" // for FederatedColl
|
|
#include "federated_comm.h" // for FederatedComm
|
|
#include "test_worker.h" // for TestFederated
|
|
|
|
namespace xgboost::collective {
|
|
namespace {
|
|
class FederatedCollTest : public SocketTest {};
|
|
} // namespace
|
|
|
|
TEST_F(FederatedCollTest, Allreduce) {
|
|
std::int32_t n_workers = std::min(std::thread::hardware_concurrency(), 3u);
|
|
TestFederated(n_workers, [=](std::shared_ptr<FederatedComm> comm, std::int32_t) {
|
|
std::array<std::int32_t, 5> buffer = {1, 2, 3, 4, 5};
|
|
std::array<std::int32_t, 5> expected;
|
|
std::transform(buffer.cbegin(), buffer.cend(), expected.begin(),
|
|
[=](auto i) { return i * n_workers; });
|
|
|
|
auto coll = std::make_shared<FederatedColl>();
|
|
auto rc = coll->Allreduce(*comm, common::EraseType(common::Span{buffer.data(), buffer.size()}),
|
|
ArrayInterfaceHandler::kI4, Op::kSum);
|
|
ASSERT_TRUE(rc.OK());
|
|
for (auto i = 0; i < 5; i++) {
|
|
ASSERT_EQ(buffer[i], expected[i]);
|
|
}
|
|
});
|
|
}
|
|
|
|
TEST_F(FederatedCollTest, Broadcast) {
|
|
std::int32_t n_workers = std::min(std::thread::hardware_concurrency(), 3u);
|
|
TestFederated(n_workers, [=](std::shared_ptr<FederatedComm> comm, std::int32_t) {
|
|
FederatedColl coll{};
|
|
auto rc = Success();
|
|
if (comm->Rank() == 0) {
|
|
std::string buffer{"hello"};
|
|
rc = coll.Broadcast(*comm, common::EraseType(common::Span{buffer.data(), buffer.size()}), 0);
|
|
ASSERT_EQ(buffer, "hello");
|
|
} else {
|
|
std::string buffer{" "};
|
|
rc = coll.Broadcast(*comm, common::EraseType(common::Span{buffer.data(), buffer.size()}), 0);
|
|
ASSERT_EQ(buffer, "hello");
|
|
}
|
|
ASSERT_TRUE(rc.OK());
|
|
});
|
|
}
|
|
|
|
TEST_F(FederatedCollTest, Allgather) {
|
|
std::int32_t n_workers = std::min(std::thread::hardware_concurrency(), 3u);
|
|
TestFederated(n_workers, [=](std::shared_ptr<FederatedComm> comm, std::int32_t) {
|
|
FederatedColl coll{};
|
|
|
|
std::vector<std::int32_t> buffer(n_workers, 0);
|
|
buffer[comm->Rank()] = comm->Rank();
|
|
auto rc = coll.Allgather(*comm, common::EraseType(common::Span{buffer.data(), buffer.size()}));
|
|
ASSERT_TRUE(rc.OK());
|
|
for (auto i = 0; i < n_workers; i++) {
|
|
ASSERT_EQ(buffer[i], i);
|
|
}
|
|
});
|
|
}
|
|
|
|
TEST_F(FederatedCollTest, AllgatherV) {
|
|
std::int32_t n_workers = 2;
|
|
TestFederated(n_workers, [=](std::shared_ptr<FederatedComm> comm, std::int32_t) {
|
|
FederatedColl coll{};
|
|
|
|
std::vector<std::string_view> inputs{"Federated", " Learning!!!"};
|
|
std::vector<std::int64_t> recv_segments(inputs.size() + 1, 0);
|
|
std::string r;
|
|
std::vector<std::int64_t> sizes{static_cast<std::int64_t>(inputs[0].size()),
|
|
static_cast<std::int64_t>(inputs[1].size())};
|
|
r.resize(sizes[0] + sizes[1]);
|
|
|
|
auto rc = coll.AllgatherV(
|
|
*comm,
|
|
common::EraseType(common::Span{inputs[comm->Rank()].data(), inputs[comm->Rank()].size()}),
|
|
common::Span{sizes.data(), sizes.size()}, recv_segments,
|
|
common::EraseType(common::Span{r.data(), r.size()}), AllgatherVAlgo::kRing);
|
|
|
|
EXPECT_EQ(r, "Federated Learning!!!");
|
|
ASSERT_TRUE(rc.OK());
|
|
});
|
|
}
|
|
} // namespace xgboost::collective
|