Use in-memory communicator to test quantile (#8710)

This commit is contained in:
Rong Ou
2023-01-27 07:28:28 -08:00
committed by GitHub
parent 96e6b6beba
commit 8af98e30fc
8 changed files with 86 additions and 207 deletions

View File

@@ -338,12 +338,9 @@ TEST(GPUQuantile, MultiMerge) {
});
}
TEST(GPUQuantile, AllReduceBasic) {
// This test is supposed to run by a python test that setups the environment.
std::string msg {"Skipping AllReduce test"};
auto n_gpus = AllVisibleGPUs();
InitCommunicatorContext(msg, n_gpus);
auto world = collective::GetWorldSize();
namespace {
void TestAllReduceBasic(int32_t n_gpus) {
auto const world = collective::GetWorldSize();
if (world != 1) {
ASSERT_EQ(world, n_gpus);
} else {
@@ -420,13 +417,16 @@ TEST(GPUQuantile, AllReduceBasic) {
ASSERT_NEAR(single_node_data[i].wmin, distributed_data[i].wmin, Eps);
}
});
collective::Finalize();
}
} // anonymous namespace
TEST(GPUQuantile, AllReduceBasic) {
auto const n_gpus = AllVisibleGPUs();
RunWithInMemoryCommunicator(n_gpus, TestAllReduceBasic, n_gpus);
}
TEST(GPUQuantile, SameOnAllWorkers) {
std::string msg {"Skipping SameOnAllWorkers test"};
auto n_gpus = AllVisibleGPUs();
InitCommunicatorContext(msg, n_gpus);
namespace {
void TestSameOnAllWorkers(int32_t n_gpus) {
auto world = collective::GetWorldSize();
if (world != 1) {
ASSERT_EQ(world, n_gpus);
@@ -490,6 +490,12 @@ TEST(GPUQuantile, SameOnAllWorkers) {
}
});
}
} // anonymous namespace
TEST(GPUQuantile, SameOnAllWorkers) {
auto const n_gpus = AllVisibleGPUs();
RunWithInMemoryCommunicator(n_gpus, TestSameOnAllWorkers, n_gpus);
}
TEST(GPUQuantile, Push) {
size_t constexpr kRows = 100;