Limit faiss ivf index build thread num (#136)

Signed-off-by: chasingegg <[email protected]>
zilliztech · Oct 10, 2023 · 7144b89 · 7144b89
1 parent 6d743e2
commit 7144b89
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git 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
@@ -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) {