Fix flaky sparse page dmatrix test. (#6417)

This commit is contained in:
Jiaming Yuan 2020-11-20 19:15:45 +08:00 committed by GitHub
parent a7b42adb74
commit c120822a24
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -2,6 +2,8 @@
#include <dmlc/filesystem.h> #include <dmlc/filesystem.h>
#include <gtest/gtest.h> #include <gtest/gtest.h>
#include <xgboost/data.h> #include <xgboost/data.h>
#include <thread>
#include <future>
#include "../../../src/common/io.h" #include "../../../src/common/io.h"
#include "../../../src/data/adapter.h" #include "../../../src/data/adapter.h"
#include "../../../src/data/sparse_page_dmatrix.h" #include "../../../src/data/sparse_page_dmatrix.h"
@ -99,26 +101,32 @@ TEST(SparsePageDMatrix, ThreadSafetyException) {
std::unique_ptr<xgboost::DMatrix> dmat = std::unique_ptr<xgboost::DMatrix> dmat =
xgboost::CreateSparsePageDMatrix(kEntries, kPageSize, filename); xgboost::CreateSparsePageDMatrix(kEntries, kPageSize, filename);
std::atomic<bool> exception {false};
int threads = 1000; int threads = 1000;
std::vector<std::thread> waiting; std::vector<std::future<void>> waiting;
std::atomic<bool> exception {false};
for (int32_t i = 0; i < threads; ++i) { for (int32_t i = 0; i < threads; ++i) {
waiting.emplace_back([&]() { waiting.emplace_back(std::async(std::launch::async, [&]() {
try { try {
auto iter = dmat->GetBatches<SparsePage>().begin(); auto iter = dmat->GetBatches<SparsePage>().begin();
++iter; ++iter;
} catch (...) { } catch (...) {
exception = true; exception.store(true);
} }
}); }));
} }
for (auto& t : waiting) { using namespace std::chrono_literals;
t.join();
while (std::any_of(waiting.cbegin(), waiting.cend(), [](auto const &f) {
return f.wait_for(0ms) != std::future_status::ready;
})) {
std::this_thread::sleep_for(50ms);
} }
EXPECT_TRUE(exception);
CHECK(exception);
} }
// Multi-batches access // Multi-batches access