share the pool externally to ensure that other tasks can be scheduled (#1192)

Signed-off-by: xianliang.li <xianliang.li@zilliz.com>

Author: foxspy

cmake/libs/libcardinal.cmake

@@ -1,5 +1,5 @@
 # Use short SHA1 as version
-set(CARDINAL_VERSION v2.5.7 )
+set(CARDINAL_VERSION v2.5.8)
 set(CARDINAL_REPO_URL "https://github.com/zilliztech/cardinal.git")
 
 set(CARDINAL_REPO_DIR "${CMAKE_CURRENT_SOURCE_DIR}/thirdparty/cardinal")

include/knowhere/comp/task.h

@@ -12,6 +12,9 @@
 #include <functional>
 #include <memory>
 #include <vector>
+
+#include "folly/executors/CPUThreadPoolExecutor.h"
+
 namespace knowhere {
 
 void
@@ -26,6 +29,13 @@ size_t
 GetSearchThreadPoolSize();
 size_t
 GetBuildThreadPoolSize();
+
+folly::CPUThreadPoolExecutor&
+GetSearchThreadPool();
+
+folly::CPUThreadPoolExecutor&
+GetBuildThreadPool();
+
 class ThreadPool {
  public:
     class ScopedOmpSetter {

include/knowhere/comp/thread_pool.h

@@ -141,6 +141,11 @@ class ThreadPool {
         return pool_.getPendingTaskCount();
     }
 
+    folly::CPUThreadPoolExecutor&
+    GetPool() {
+        return pool_;
+    }
+
     void
     SetNumThreads(uint32_t num_threads) {
         if (num_threads == 0) {

python/knowhere/knowhere.i

@@ -37,7 +37,6 @@ typedef uint64_t size_t;
 #include <knowhere/comp/brute_force.h>
 #include <knowhere/comp/knowhere_config.h>
 #include <knowhere/comp/local_file_manager.h>
-#include <knowhere/comp/task.h>
 #include <knowhere/comp/index_param.h>
 #include <fstream>
 #include <string>
@@ -691,12 +690,12 @@ SetSimdType(const std::string type) {
 
 void
 SetBuildThreadPool(uint32_t num_threads) {
-    knowhere::InitBuildThreadPool(num_threads);
+    knowhere::KnowhereConfig::SetBuildThreadPoolSize(num_threads);
 }
 
 void
 SetSearchThreadPool(uint32_t num_threads) {
-    knowhere::InitSearchThreadPool(num_threads);
+    knowhere::KnowhereConfig::SetSearchThreadPoolSize(num_threads);
 }
 
 %}

src/common/thread/thread.cc

@@ -18,6 +18,7 @@
 #include <thread>
 #include <utility>
 
+#include "folly/executors/CPUThreadPoolExecutor.h"
 #include "knowhere/comp/thread_pool.h"
 
 namespace knowhere {
@@ -57,6 +58,16 @@ InitBuildThreadPool(uint32_t num_threads) {
     ThreadPool::InitGlobalBuildThreadPool(num_threads);
 }
 
+folly::CPUThreadPoolExecutor&
+GetBuildThreadPool() {
+    return ThreadPool::GetGlobalBuildThreadPool()->GetPool();
+}
+
+folly::CPUThreadPoolExecutor&
+GetSearchThreadPool() {
+    return ThreadPool::GetGlobalSearchThreadPool()->GetPool();
+}
+
 void
 InitSearchThreadPool(uint32_t num_threads) {
     ThreadPool::InitGlobalSearchThreadPool(num_threads);