This is an automated email from the ASF dual-hosted git repository.
yiguolei pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/doris.git
The following commit(s) were added to refs/heads/master by this push:
new b325dcd122b [fix](be) Add fixed per-task ANN build thread budgets
(#70454)
b325dcd122b is described below
commit b325dcd122b5f50685881c1399f1eb904ea2efa3
Author: yiguolei <[email protected]>
AuthorDate: Fri Nov 13 10:51:43 2026 +0800
[fix](be) Add fixed per-task ANN build thread budgets (#70454)
### Background
Faiss ANN index training and vector insertion share a BE-wide OpenMP
thread budget controlled by `omp_threads_limit`. Each call currently
reserves half of the remaining budget, so concurrent calls can use
different OpenMP thread counts.
Maybe we meet OpenBLAS bug
https://github.com/OpenMathLib/OpenBLAS/issues/4418#issuecomment-1879923133
### What problem does this PR solve?
Issue Number: None
Related PR: None
PR Type: Enhancement
Problem Summary:
Make the OpenMP reservation for each ANN train/add call independent of
concurrent budget usage. Add `ann_index_build_threads_per_task`, cap its
effective value by the total budget, and wait until a full reservation
is available. This preserves aggregate admission control while making
per-call parallelism predictable. Most added lines are unit tests.
### Details
- **Before and after:** Previously, a total budget of 10 could produce
reservations of 5, 2, and 1 as calls arrived. With 4 threads per task,
two calls reserve 4 each and a third waits until 4 slots are free. Faiss
work executes outside the budget mutex; RAII releases reservations and
wakes waiters on completion or failure.
- **Configs and defaults:** Add non-dynamic BE configuration
`ann_index_build_threads_per_task`, default `-1`: use 1 thread below 32
available CPU cores and 4 otherwise. Positive explicit values are
accepted, zero and values below -1 are rejected. Cap the reservation by
`omp_threads_limit`. Keep the existing total-budget default of 80% of
available cores, at least 1. Both auto defaults use one CPU snapshot
captured on first use; a nonpositive core count logs a warning and falls
back to 16. Configuration changes require a BE restart.
- **Impact on existing users:** Default per-call parallelism changes,
which can change build latency and concurrency. The total configured ANN
budget is retained. A remainder smaller than one reservation stays
unused until another call finishes. This controls reserved OpenMP
parallelism, not a hard OS CPU quota.
- **Behavior change:** Yes; fixed full reservations replace
half-of-remaining reservations for both ANN training and vector
insertion. Other branches are unchanged.
### Testing
- Passed locally: `build-support/check-format.sh` with clang-format
16.0.5, `build-support/check-build-hygiene.sh`, and `git diff --check`.
- Added unit tests for configuration validation, explicit and automatic
reservations, capping to the total budget, stable CPU snapshots,
concurrent train/add admission, and reservation release after a training
exception.
- Attempted `./run-be-ut.sh --run
--filter='AnnOmpThreadBudgetTest.*:VectorSearchTest.OmpThreadBudgetNeverExceedsLimit'
-j 4` on macOS arm64. Environment initialization stopped at the Java
version check before compilation; the tests have not run. BE compilation
and clang-tidy have not run.
### Notes for Reviewers
This change was motivated by a persistent OpenBLAS GEMM hang with 12
logical jobs and six observed workers executing only even-numbered jobs.
Varying caller thread counts interacting with shared OpenBLAS
thread-count state is a source-supported hypothesis, not a reproduced
root cause. This PR removes varying ANN build reservations; it does not
claim to fix every possible OpenBLAS deadlock. The new tests exercise
admission and reservation behavior, not the native OpenBLAS race.
### Release note
Add the BE configuration `ann_index_build_threads_per_task` to control
threads reserved per ANN train/add call while retaining
`omp_threads_limit` as the total budget. The automatic per-call default
is 1 thread below 32 available CPU cores and 4 otherwise; calls wait
when a full reservation is unavailable.
### Check List (For Author)
- Behavior changed: Yes (see Details).
- Does this need documentation: Yes; document the new BE configuration
and its defaults in apache/doris-website as a follow-up.
- Pre-submission status: Partially verified: formatting and build
hygiene passed; compilation, unit tests, and clang-tidy remain
unverified.
### AI Disclosure
- Tool: Codex (GPT-6).
- Generated: Code, tests, and this description.
- Human review: The author reviewed and directed the budget design,
defaults, CPU snapshot behavior, and invalid-core fallback; full code
review is pending.
---
be/src/common/config.cpp | 9 +
be/src/common/config.h | 5 +
be/src/storage/index/ann/faiss_ann_index.cpp | 42 +++--
.../storage/index/ann/faiss_vector_index_test.cpp | 198 +++++++++++++++++++++
4 files changed, 244 insertions(+), 10 deletions(-)
diff --git a/be/src/common/config.cpp b/be/src/common/config.cpp
index 28066e9e4d8..a87e5bf2ae4 100644
--- a/be/src/common/config.cpp
+++ b/be/src/common/config.cpp
@@ -2131,8 +2131,17 @@ DEFINE_mInt64(load_reader_max_block_bytes, "67108864");
// Maximum number of OpenMP threads allowed for concurrent vector index builds.
// -1 means auto: use 80% of the available CPU cores.
+// Auto sizing shares the CPU count captured on first use with
ann_index_build_threads_per_task.
DEFINE_Int32(omp_threads_limit, "-1");
+// OpenMP threads reserved by each ANN index train/add call, capped by
omp_threads_limit.
+// -1 means auto: 1 thread below 32 available CPU cores, otherwise 4 threads.
+// The CPU count is captured once on first use and reused until BE restart.
+// A call waits until its entire reservation is available. Requires a BE
restart.
+DEFINE_Int32(ann_index_build_threads_per_task, "-1");
+DEFINE_Validator(ann_index_build_threads_per_task,
+ [](int32_t value) { return value == -1 || value > 0; });
+
// The capacity of segment partial column cache, used to cache column readers
for each segment.
DEFINE_mInt32(max_segment_partial_column_cache_size, "100");
diff --git a/be/src/common/config.h b/be/src/common/config.h
index 02ce11f50e5..47973f78421 100644
--- a/be/src/common/config.h
+++ b/be/src/common/config.h
@@ -2092,7 +2092,12 @@ DECLARE_mInt64(load_reader_max_block_bytes);
// Maximum number of OpenMP threads available for concurrent index builds.
// -1 means auto: use 80% of detected CPU cores.
+// Auto sizing shares the CPU count captured on first use with
ann_index_build_threads_per_task.
DECLARE_Int32(omp_threads_limit);
+// OpenMP threads per ANN train/add call, capped by the global budget above.
+// -1 means auto: 1 thread below 32 available CPU cores, otherwise 4 threads.
+// The CPU count is captured once on first use and reused until BE restart.
+DECLARE_Int32(ann_index_build_threads_per_task);
// The capacity of segment partial column cache, used to cache column readers
for each segment.
DECLARE_mInt32(max_segment_partial_column_cache_size);
// Cache for ANN index IVF on-disk list data.
diff --git a/be/src/storage/index/ann/faiss_ann_index.cpp
b/be/src/storage/index/ann/faiss_ann_index.cpp
index 821bba40397..3f7c15c7c5a 100644
--- a/be/src/storage/index/ann/faiss_ann_index.cpp
+++ b/be/src/storage/index/ann/faiss_ann_index.cpp
@@ -94,19 +94,24 @@ public:
ScopedIoCtxBinding& operator=(const ScopedIoCtxBinding&) = delete;
};
-// Guard that ensures the total OpenMP threads used by concurrent index builds
-// never exceed the configured omp_threads_limit.
+// Guard that keeps the total OpenMP thread reservations for concurrent index
builds
+// within omp_threads_limit. Each call reserves a fixed amount independent of
other builds.
class ScopedOmpThreadBudget {
public:
- // For each index build, reserve at most half of the remaining threads, at
least 1 thread.
ScopedOmpThreadBudget() {
std::unique_lock<std::mutex> lock(g_omp_thread_mutex);
- auto omp_threads_limit = get_omp_threads_limit();
- // Block until there is at least one OpenMP slot available under the
global cap.
- g_omp_thread_cv.wait(lock, [&] { return g_index_threads_in_use <
omp_threads_limit; });
- auto thread_cap = omp_threads_limit - g_index_threads_in_use;
- // Keep headroom for other concurrent index builds: take up to half of
remaining budget.
- _reserved_threads = std::max(1, thread_cap / 2);
+ const int omp_threads_limit = get_omp_threads_limit();
+ const int threads_per_task = config::ann_index_build_threads_per_task
== -1
+ ? (get_build_num_cores() < 32 ? 1
: 4)
+ :
config::ann_index_build_threads_per_task;
+ // Cap a reservation by the total budget so it can always fit once
other calls finish.
+ _reserved_threads = std::min(threads_per_task, omp_threads_limit);
+ // Wait for the whole reservation instead of shrinking it to the
remaining budget.
+ // For example, with a limit of 10 and 4 threads per call, a third
call waits even
+ // though 2 slots remain. Subtraction avoids overflow for large
configured limits.
+ g_omp_thread_cv.wait(lock, [&] {
+ return g_index_threads_in_use <= omp_threads_limit -
_reserved_threads;
+ });
g_index_threads_in_use += _reserved_threads;
DorisMetrics::instance()->ann_index_build_index_threads->increment(_reserved_threads);
omp_set_num_threads(_reserved_threads);
@@ -133,12 +138,29 @@ public:
if (config::omp_threads_limit > 0) {
return config::omp_threads_limit;
}
- int core_cap = std::max(1, CpuInfo::num_cores());
+ const int core_cap = get_build_num_cores();
// Use at most 80% of the available CPU cores.
return std::max(1, core_cap * 4 / 5);
}
private:
+ static int get_build_num_cores() {
+ // CpuInfo is initialized before any index build. Snapshot it on first
use so
+ // both auto defaults remain consistent even if CPU detection later
becomes dynamic.
+ // For example, a change from 32 cores to 1 must not reduce a later
call's reservation.
+ static const int num_cores = [] {
+ const int cores = CpuInfo::num_cores();
+ // An invalid CPU count must not abort an index build or produce a
zero budget.
+ if (cores <= 0) {
+ LOG(WARNING) << "Invalid CPU core count " << cores
+ << "; using 16 cores for ANN index build thread
budgets";
+ return 16;
+ }
+ return cores;
+ }();
+ return num_cores;
+ }
+
int _reserved_threads = 1;
};
diff --git a/be/test/storage/index/ann/faiss_vector_index_test.cpp
b/be/test/storage/index/ann/faiss_vector_index_test.cpp
index 370cd33e0c1..45c3dabb01f 100644
--- a/be/test/storage/index/ann/faiss_vector_index_test.cpp
+++ b/be/test/storage/index/ann/faiss_vector_index_test.cpp
@@ -15,10 +15,12 @@
// specific language governing permissions and limitations
// under the License.
+#include <faiss/IndexFlat.h>
#include <faiss/IndexHNSW.h>
#include <faiss/IndexIVFFlat.h>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
+#include <omp.h>
#include <algorithm>
#include <atomic>
@@ -35,12 +37,14 @@
#include "common/config.h"
#include "common/metrics/doris_metrics.h"
+#include "faiss/impl/FaissException.h"
#include "storage/index/ann/ann_index.h"
#include "storage/index/ann/ann_search_params.h"
#include "storage/index/ann/faiss_ann_index.h"
// metrics.h not used directly here
#include "storage/cache/ann_index_ivf_list_cache.h"
#include "storage/index/ann/vector_search_utils.h"
+#include "util/cpu_info.h"
#include "util/defer_op.h"
using namespace doris::segment_v2;
@@ -301,6 +305,200 @@ TEST_F(VectorSearchTest,
OmpThreadBudgetNeverExceedsLimit) {
EXPECT_EQ(budget_metric->value(), 0);
}
+namespace {
+
+class MockFaissBuildIndex : public faiss::IndexFlatL2 {
+public:
+ MockFaissBuildIndex() : faiss::IndexFlatL2(1) {}
+ MOCK_METHOD(void, train, (faiss::idx_t, const float*), (override));
+ MOCK_METHOD(void, add, (faiss::idx_t, const float*), (override));
+};
+
+class AnnOmpThreadBudgetTest : public VectorSearchTest {
+public:
+ void SetUp() override {
+ VectorSearchTest::SetUp();
+ _old_limit = config::omp_threads_limit;
+ _old_threads_per_task = config::ann_index_build_threads_per_task;
+ _old_num_cores = CpuInfo::num_cores_;
+ _old_omp_threads = omp_get_max_threads();
+ }
+
+ void TearDown() override {
+
EXPECT_EQ(DorisMetrics::instance()->ann_index_build_index_threads->value(), 0);
+ config::omp_threads_limit = _old_limit;
+ config::ann_index_build_threads_per_task = _old_threads_per_task;
+ CpuInfo::num_cores_ = _old_num_cores;
+ omp_set_num_threads(_old_omp_threads);
+ VectorSearchTest::TearDown();
+ }
+
+private:
+ int _old_limit = 0;
+ int _old_threads_per_task = 0;
+ int _old_num_cores = 0;
+ int _old_omp_threads = 0;
+};
+
+} // namespace
+
+TEST_F(AnnOmpThreadBudgetTest, ConfigurationRejectsInvalidReservations) {
+ const auto& field =
config::Register::_s_field_map->at("ann_index_build_threads_per_task");
+ EXPECT_FALSE(field.valmutable);
+ const auto& validate =
config::RegisterConfValidator::_s_field_validator->at(
+ "ann_index_build_threads_per_task");
+ for (int value : {-2, 0}) {
+ config::ann_index_build_threads_per_task = value;
+ EXPECT_FALSE(validate()) << value;
+ }
+ for (int value : {-1, 1, 4}) {
+ config::ann_index_build_threads_per_task = value;
+ EXPECT_TRUE(validate()) << value;
+ }
+}
+
+TEST_F(AnnOmpThreadBudgetTest, FixedAndAutomaticReservations) {
+ struct TestCase {
+ int total;
+ int per_task;
+ int expected;
+ };
+ const int auto_threads = CpuInfo::num_cores() < 32 ? 1 : 4;
+ const TestCase cases[] = {{10, 4, 4},
+ {2, 8, 2},
+ {100, -1, auto_threads},
+ {2, -1, std::min(2, auto_threads)},
+ {-1, -1, auto_threads}};
+ float value = 0;
+ for (const auto& test : cases) {
+ SCOPED_TRACE(testing::Message()
+ << "total=" << test.total << ", per_task=" <<
test.per_task);
+ config::omp_threads_limit = test.total;
+ config::ann_index_build_threads_per_task = test.per_task;
+ FaissVectorIndex index;
+ auto mock =
std::make_unique<testing::StrictMock<MockFaissBuildIndex>>();
+ auto check_budget = [&](faiss::idx_t, const float*) {
+ EXPECT_EQ(omp_get_max_threads(), test.expected);
+
EXPECT_EQ(DorisMetrics::instance()->ann_index_build_index_threads->value(),
+ test.expected);
+ };
+ EXPECT_CALL(*mock, train(0, &value)).WillOnce(check_budget);
+ EXPECT_CALL(*mock, add(0, &value)).WillOnce(check_budget);
+ index._index = std::move(mock);
+ EXPECT_TRUE(index.train(0, &value).ok());
+
EXPECT_EQ(DorisMetrics::instance()->ann_index_build_index_threads->value(), 0);
+ EXPECT_TRUE(index.add(0, &value).ok());
+
EXPECT_EQ(DorisMetrics::instance()->ann_index_build_index_threads->value(), 0);
+ }
+}
+
+TEST_F(AnnOmpThreadBudgetTest, AutomaticReservationsKeepFirstCpuSnapshot) {
+ config::omp_threads_limit = -1;
+ config::ann_index_build_threads_per_task = -1;
+ const int expected_threads = CpuInfo::num_cores() < 32 ? 1 : 4;
+ const int expected_limit = std::max(1, CpuInfo::num_cores() * 4 / 5);
+ float value = 0;
+ FaissVectorIndex index;
+ auto mock = std::make_unique<testing::StrictMock<MockFaissBuildIndex>>();
+ EXPECT_CALL(*mock, train(0,
&value)).Times(3).WillRepeatedly([&](faiss::idx_t, const float*) {
+ EXPECT_EQ(omp_get_max_threads(), expected_threads);
+
EXPECT_EQ(DorisMetrics::instance()->ann_index_build_index_threads->value(),
+ expected_threads);
+ });
+ EXPECT_CALL(*mock, add(0,
&value)).Times(2).WillRepeatedly([&](faiss::idx_t, const float*) {
+ EXPECT_EQ(omp_get_max_threads(), expected_limit);
+
EXPECT_EQ(DorisMetrics::instance()->ann_index_build_index_threads->value(),
expected_limit);
+ });
+ index._index = std::move(mock);
+ ASSERT_TRUE(index.train(0, &value).ok());
+ // Cross the 32-core default boundary in both directions after the first
reservation.
+ // Also request the whole budget to verify that auto total sizing uses the
same snapshot.
+ for (int cores : {1, 64}) {
+ CpuInfo::num_cores_ = cores;
+ config::ann_index_build_threads_per_task = -1;
+ EXPECT_TRUE(index.train(0, &value).ok());
+ config::ann_index_build_threads_per_task =
std::numeric_limits<int>::max();
+ EXPECT_TRUE(index.add(0, &value).ok());
+ }
+}
+
+TEST_F(AnnOmpThreadBudgetTest, ConcurrentCallsKeepFixedReservations) {
+ config::ann_index_build_threads_per_task = 4;
+ float value = 0;
+ auto* metric = DorisMetrics::instance()->ann_index_build_index_threads;
+ // Two calls fit in 10 slots. With 7 slots, the second must wait rather
than take 3.
+ for (int total : {10, 7}) {
+ SCOPED_TRACE(total);
+ config::omp_threads_limit = total;
+ std::promise<void> first_entered;
+ auto first_entered_future = first_entered.get_future();
+ std::promise<void> release_first;
+ auto release_first_future = release_first.get_future();
+ std::promise<void> second_started;
+ auto second_started_future = second_started.get_future();
+ std::promise<void> second_entered;
+ auto second_entered_future = second_entered.get_future();
+
+ FaissVectorIndex first;
+ auto first_mock =
std::make_unique<testing::StrictMock<MockFaissBuildIndex>>();
+ EXPECT_CALL(*first_mock, train(0, &value)).WillOnce([&](faiss::idx_t,
const float*) {
+ EXPECT_EQ(omp_get_max_threads(), 4);
+ EXPECT_EQ(metric->value(), 4);
+ first_entered.set_value();
+ release_first_future.wait();
+ });
+ first._index = std::move(first_mock);
+ auto first_result = std::async(std::launch::async, [&] { return
first.train(0, &value); });
+ EXPECT_EQ(first_entered_future.wait_for(std::chrono::seconds(5)),
+ std::future_status::ready);
+
+ FaissVectorIndex second;
+ auto second_mock =
std::make_unique<testing::StrictMock<MockFaissBuildIndex>>();
+ EXPECT_CALL(*second_mock, add(0, &value)).WillOnce([&](faiss::idx_t,
const float*) {
+ EXPECT_EQ(omp_get_max_threads(), 4);
+ EXPECT_EQ(metric->value(), total == 10 ? 8 : 4);
+ second_entered.set_value();
+ });
+ second._index = std::move(second_mock);
+ auto second_result = std::async(std::launch::async, [&] {
+ second_started.set_value();
+ return second.add(0, &value);
+ });
+ second_started_future.wait();
+ if (total == 10) {
+ EXPECT_EQ(second_entered_future.wait_for(std::chrono::seconds(5)),
+ std::future_status::ready);
+ } else {
+
EXPECT_EQ(second_entered_future.wait_for(std::chrono::milliseconds(100)),
+ std::future_status::timeout);
+ EXPECT_EQ(metric->value(), 4);
+ }
+ release_first.set_value();
+ EXPECT_TRUE(first_result.get().ok());
+ EXPECT_TRUE(second_result.get().ok());
+ EXPECT_EQ(metric->value(), 0);
+ }
+}
+
+TEST_F(AnnOmpThreadBudgetTest, FailedTrainingReleasesReservation) {
+ config::omp_threads_limit = 4;
+ config::ann_index_build_threads_per_task = 4;
+ float value = 0;
+ FaissVectorIndex index;
+ auto mock = std::make_unique<testing::StrictMock<MockFaissBuildIndex>>();
+ EXPECT_CALL(*mock, train(0, &value)).WillOnce([](faiss::idx_t, const
float*) {
+ throw faiss::FaissException("injected training failure");
+ });
+ EXPECT_CALL(*mock, add(0, &value)).WillOnce([](faiss::idx_t, const float*)
{
+ EXPECT_EQ(omp_get_max_threads(), 4);
+
EXPECT_EQ(DorisMetrics::instance()->ann_index_build_index_threads->value(), 4);
+ });
+ index._index = std::move(mock);
+ EXPECT_FALSE(index.train(0, &value).ok());
+
EXPECT_EQ(DorisMetrics::instance()->ann_index_build_index_threads->value(), 0);
+ EXPECT_TRUE(index.add(0, &value).ok());
+}
+
TEST_F(VectorSearchTest, CompareResultWithNativeFaiss1) {
const size_t iterations = 3;
// Create random number generator
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
