Fix multi-threaded gtests (#9148)

This commit is contained in:
Rong Ou 2023-05-10 04:15:32 -07:00 committed by GitHub
parent e4129ed6ee
commit 52311dcec9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 120 additions and 50 deletions

View File

@ -222,15 +222,15 @@ void InMemoryHandler::Handle(char const* input, std::size_t bytes, std::string*
std::unique_lock<std::mutex> lock(mutex_);
LOG(INFO) << functor.name << " rank " << rank << ": waiting for current sequence number";
LOG(DEBUG) << functor.name << " rank " << rank << ": waiting for current sequence number";
cv_.wait(lock, [this, sequence_number] { return sequence_number_ == sequence_number; });
LOG(INFO) << functor.name << " rank " << rank << ": handling request";
LOG(DEBUG) << functor.name << " rank " << rank << ": handling request";
functor(input, bytes, &buffer_);
received_++;
if (received_ == world_size_) {
LOG(INFO) << functor.name << " rank " << rank << ": all requests received";
LOG(DEBUG) << functor.name << " rank " << rank << ": all requests received";
output->assign(buffer_);
sent_++;
lock.unlock();
@ -238,15 +238,15 @@ void InMemoryHandler::Handle(char const* input, std::size_t bytes, std::string*
return;
}
LOG(INFO) << functor.name << " rank " << rank << ": waiting for all clients";
LOG(DEBUG) << functor.name << " rank " << rank << ": waiting for all clients";
cv_.wait(lock, [this] { return received_ == world_size_; });
LOG(INFO) << functor.name << " rank " << rank << ": sending reply";
LOG(DEBUG) << functor.name << " rank " << rank << ": sending reply";
output->assign(buffer_);
sent_++;
if (sent_ == world_size_) {
LOG(INFO) << functor.name << " rank " << rank << ": all replies sent";
LOG(DEBUG) << functor.name << " rank " << rank << ": all replies sent";
sent_ = 0;
received_ = 0;
buffer_.clear();

View File

@ -166,7 +166,7 @@ BatchSet<GHistIndexMatrix> SimpleDMatrix::GetGradientIndex(Context const* ctx,
}
if (!gradient_index_ || detail::RegenGHist(batch_param_, param)) {
// GIDX page doesn't exist, generate it
LOG(INFO) << "Generating new Gradient Index.";
LOG(DEBUG) << "Generating new Gradient Index.";
// These places can ask for a CSR gidx:
// - CPU Hist: the ctx must be on CPU.
// - IterativeDMatrix::InitFromCPU: The ctx must be on CPU.

View File

@ -26,6 +26,60 @@ class InMemoryCommunicatorTest : public ::testing::Test {
static void Allgather(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
VerifyAllgather(comm, rank);
}
static void AllreduceMax(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
VerifyAllreduceMax(comm, rank);
}
static void AllreduceMin(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
VerifyAllreduceMin(comm, rank);
}
static void AllreduceSum(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
VerifyAllreduceSum(comm);
}
static void AllreduceBitwiseAND(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
VerifyAllreduceBitwiseAND(comm, rank);
}
static void AllreduceBitwiseOR(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
VerifyAllreduceBitwiseOR(comm, rank);
}
static void AllreduceBitwiseXOR(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
VerifyAllreduceBitwiseXOR(comm, rank);
}
static void Broadcast(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
VerifyBroadcast(comm, rank);
}
static void Mixture(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
for (auto i = 0; i < 5; i++) {
VerifyAllgather(comm, rank);
VerifyAllreduceMax(comm, rank);
VerifyAllreduceMin(comm, rank);
VerifyAllreduceSum(comm);
VerifyAllreduceBitwiseAND(comm, rank);
VerifyAllreduceBitwiseOR(comm, rank);
VerifyAllreduceBitwiseXOR(comm, rank);
VerifyBroadcast(comm, rank);
}
}
protected:
static void VerifyAllgather(InMemoryCommunicator &comm, int rank) {
char buffer[kWorldSize] = {'a', 'b', 'c'};
buffer[rank] = '0' + rank;
comm.AllGather(buffer, kWorldSize);
@ -34,8 +88,7 @@ class InMemoryCommunicatorTest : public ::testing::Test {
}
}
static void AllreduceMax(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
static void VerifyAllreduceMax(InMemoryCommunicator &comm, int rank) {
int buffer[] = {1 + rank, 2 + rank, 3 + rank, 4 + rank, 5 + rank};
comm.AllReduce(buffer, sizeof(buffer) / sizeof(buffer[0]), DataType::kInt32, Operation::kMax);
int expected[] = {3, 4, 5, 6, 7};
@ -44,8 +97,7 @@ class InMemoryCommunicatorTest : public ::testing::Test {
}
}
static void AllreduceMin(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
static void VerifyAllreduceMin(InMemoryCommunicator &comm, int rank) {
int buffer[] = {1 + rank, 2 + rank, 3 + rank, 4 + rank, 5 + rank};
comm.AllReduce(buffer, sizeof(buffer) / sizeof(buffer[0]), DataType::kInt32, Operation::kMin);
int expected[] = {1, 2, 3, 4, 5};
@ -54,8 +106,7 @@ class InMemoryCommunicatorTest : public ::testing::Test {
}
}
static void AllreduceSum(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
static void VerifyAllreduceSum(InMemoryCommunicator &comm) {
int buffer[] = {1, 2, 3, 4, 5};
comm.AllReduce(buffer, sizeof(buffer) / sizeof(buffer[0]), DataType::kInt32, Operation::kSum);
int expected[] = {3, 6, 9, 12, 15};
@ -64,16 +115,14 @@ class InMemoryCommunicatorTest : public ::testing::Test {
}
}
static void AllreduceBitwiseAND(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
static void VerifyAllreduceBitwiseAND(InMemoryCommunicator &comm, int rank) {
std::bitset<2> original(rank);
auto buffer = original.to_ulong();
comm.AllReduce(&buffer, 1, DataType::kUInt32, Operation::kBitwiseAND);
EXPECT_EQ(buffer, 0UL);
}
static void AllreduceBitwiseOR(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
static void VerifyAllreduceBitwiseOR(InMemoryCommunicator &comm, int rank) {
std::bitset<2> original(rank);
auto buffer = original.to_ulong();
comm.AllReduce(&buffer, 1, DataType::kUInt32, Operation::kBitwiseOR);
@ -82,8 +131,7 @@ class InMemoryCommunicatorTest : public ::testing::Test {
EXPECT_EQ(actual, expected);
}
static void AllreduceBitwiseXOR(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
static void VerifyAllreduceBitwiseXOR(InMemoryCommunicator &comm, int rank) {
std::bitset<3> original(rank * 2);
auto buffer = original.to_ulong();
comm.AllReduce(&buffer, 1, DataType::kUInt32, Operation::kBitwiseXOR);
@ -92,8 +140,7 @@ class InMemoryCommunicatorTest : public ::testing::Test {
EXPECT_EQ(actual, expected);
}
static void Broadcast(int rank) {
InMemoryCommunicator comm{kWorldSize, rank};
static void VerifyBroadcast(InMemoryCommunicator &comm, int rank) {
if (rank == 0) {
std::string buffer{"hello"};
comm.Broadcast(&buffer[0], buffer.size(), 0);
@ -105,7 +152,6 @@ class InMemoryCommunicatorTest : public ::testing::Test {
}
}
protected:
static int const kWorldSize{3};
};
@ -173,5 +219,7 @@ TEST_F(InMemoryCommunicatorTest, AllreduceBitwiseXOR) { Verify(&AllreduceBitwise
TEST_F(InMemoryCommunicatorTest, Broadcast) { Verify(&Broadcast); }
TEST_F(InMemoryCommunicatorTest, Mixture) { Verify(&Mixture); }
} // namespace collective
} // namespace xgboost

View File

@ -497,9 +497,7 @@ inline std::int32_t AllThreadsForTest() { return Context{}.Threads(); }
template <typename Function, typename... Args>
void RunWithInMemoryCommunicator(int32_t world_size, Function&& function, Args&&... args) {
std::vector<std::thread> threads;
for (auto rank = 0; rank < world_size; rank++) {
threads.emplace_back([&, rank]() {
auto run = [&](auto rank) {
Json config{JsonObject()};
config["xgboost_communicator"] = String("in-memory");
config["in_memory_world_size"] = world_size;
@ -509,11 +507,22 @@ void RunWithInMemoryCommunicator(int32_t world_size, Function&& function, Args&&
std::forward<Function>(function)(std::forward<Args>(args)...);
xgboost::collective::Finalize();
});
};
#if defined(_OPENMP)
#pragma omp parallel num_threads(world_size)
{
auto rank = omp_get_thread_num();
run(rank);
}
#else
std::vector<std::thread> threads;
for (auto rank = 0; rank < world_size; rank++) {
threads.emplace_back(run, rank);
}
for (auto& thread : threads) {
thread.join();
}
#endif
}
class DeclareUnifiedDistributedTest(MetricTest) : public ::testing::Test {

View File

@ -3,6 +3,7 @@
*/
#pragma once
#include <dmlc/omp.h>
#include <grpcpp/server_builder.h>
#include <gtest/gtest.h>
#include <xgboost/json.h>
@ -61,9 +62,7 @@ class BaseFederatedTest : public ::testing::Test {
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]() {
auto run = [&](auto rank) {
Json config{JsonObject()};
config["xgboost_communicator"] = String("federated");
config["federated_server_address"] = String(server_address);
@ -74,11 +73,22 @@ void RunWithFederatedCommunicator(int32_t world_size, std::string const& server_
std::forward<Function>(function)(std::forward<Args>(args)...);
xgboost::collective::Finalize();
});
};
#if defined(_OPENMP)
#pragma omp parallel num_threads(world_size)
{
auto rank = omp_get_thread_num();
run(rank);
}
#else
std::vector<std::thread> threads;
for (auto rank = 0; rank < world_size; rank++) {
threads.emplace_back(run, rank);
}
for (auto& thread : threads) {
thread.join();
}
#endif
}
} // namespace xgboost

View File

@ -90,13 +90,16 @@ void TestColumnSplit(int32_t rows, bst_feature_t cols, RegTree const& expected_t
param.Init(Args{});
updater->Update(&param, p_gradients.get(), sliced.get(), position, {&tree});
EXPECT_EQ(tree.NumExtraNodes(), 10);
EXPECT_EQ(tree[0].SplitIndex(), 1);
ASSERT_EQ(tree.NumExtraNodes(), 10);
ASSERT_EQ(tree[0].SplitIndex(), 1);
EXPECT_NE(tree[tree[0].LeftChild()].SplitIndex(), 0);
EXPECT_NE(tree[tree[0].RightChild()].SplitIndex(), 0);
ASSERT_NE(tree[tree[0].LeftChild()].SplitIndex(), 0);
ASSERT_NE(tree[tree[0].RightChild()].SplitIndex(), 0);
EXPECT_EQ(tree, expected_tree);
FeatureMap fmap;
auto json = tree.DumpModel(fmap, false, "json");
auto expected_json = expected_tree.DumpModel(fmap, false, "json");
ASSERT_EQ(json, expected_json);
}
} // anonymous namespace