From 1167e6c554f2fad47183f835bf0845e2cde5cd52 Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Sun, 24 Sep 2023 00:30:28 +0800 Subject: [PATCH] Limit the number of threads for external memory. (#9605) --- src/data/sparse_page_source.h | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/data/sparse_page_source.h b/src/data/sparse_page_source.h index b32c536af..9cb0e364f 100644 --- a/src/data/sparse_page_source.h +++ b/src/data/sparse_page_source.h @@ -177,15 +177,15 @@ class SparsePageSourceImpl : public BatchIteratorImpl { } // An heuristic for number of pre-fetched batches. We can make it part of BatchParam // to let user adjust number of pre-fetched batches when needed. - uint32_t constexpr kPreFetch = 3; - - size_t n_prefetch_batches = std::min(kPreFetch, n_batches_); + std::int32_t n_prefetches = std::max(nthreads_, 3); + std::int32_t n_prefetch_batches = + std::min(static_cast(n_prefetches), n_batches_); CHECK_GT(n_prefetch_batches, 0) << "total batches:" << n_batches_; std::size_t fetch_it = count_; exce_.Rethrow(); - for (std::size_t i = 0; i < n_prefetch_batches; ++i, ++fetch_it) { + for (std::int32_t i = 0; i < n_prefetch_batches; ++i, ++fetch_it) { fetch_it %= n_batches_; // ring if (ring_->at(fetch_it).valid()) { continue;