From 7144b891da0820d137194a9d93df67b0627db6b2 Mon Sep 17 00:00:00 2001
From: Gao <chao.gao@zilliz.com>
Date: Tue, 10 Oct 2023 15:17:28 +0800
Subject: [PATCH] Limit faiss ivf index build thread num (#136)

Signed-off-by: chasingegg <chao.gao@zilliz.com>
---
 include/knowhere/comp/thread_pool.h | 10 ++++++++--
 src/index/ivf/ivf.cc                |  4 ++++
 2 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/include/knowhere/comp/thread_pool.h b/include/knowhere/comp/thread_pool.h
index 2d7417abe..a00920de6 100644
--- a/include/knowhere/comp/thread_pool.h
+++ b/include/knowhere/comp/thread_pool.h
@@ -161,8 +161,14 @@ class ThreadPool {
         int omp_before;
 
      public:
-        explicit ScopedOmpSetter(int num_threads = 1) : omp_before(omp_get_max_threads()) {
-            omp_set_num_threads(num_threads);
+        explicit ScopedOmpSetter(int num_threads = 0) {
+            if (global_build_thread_pool_size_ == 0) {  // this should not happen in prod
+                omp_before = omp_get_max_threads();
+            } else {
+                omp_before = global_build_thread_pool_size_;
+            }
+
+            omp_set_num_threads(num_threads <= 0 ? omp_before : num_threads);
         }
         ~ScopedOmpSetter() {
             omp_set_num_threads(omp_before);
diff --git a/src/index/ivf/ivf.cc b/src/index/ivf/ivf.cc
index 47be23d1e..3c23b0239 100644
--- a/src/index/ivf/ivf.cc
+++ b/src/index/ivf/ivf.cc
@@ -247,6 +247,8 @@ IvfIndexNode<T>::Train(const DataSet& dataset, const Config& cfg) {
     std::unique_ptr<ThreadPool::ScopedOmpSetter> setter;
     if (base_cfg.num_build_thread.has_value()) {
         setter = std::make_unique<ThreadPool::ScopedOmpSetter>(base_cfg.num_build_thread.value());
+    } else {
+        setter = std::make_unique<ThreadPool::ScopedOmpSetter>();
     }
 
     bool is_cosine = IsMetricType(base_cfg.metric_type.value(), knowhere::metric::COSINE);
@@ -356,6 +358,8 @@ IvfIndexNode<T>::Add(const DataSet& dataset, const Config& cfg) {
     std::unique_ptr<ThreadPool::ScopedOmpSetter> setter;
     if (base_cfg.num_build_thread.has_value()) {
         setter = std::make_unique<ThreadPool::ScopedOmpSetter>(base_cfg.num_build_thread.value());
+    } else {
+        setter = std::make_unique<ThreadPool::ScopedOmpSetter>();
     }
     try {
         if constexpr (std::is_same<faiss::IndexBinaryIVF, T>::value) {