Fix get cpu num for segcore (#20120) (#20307)

Signed-off-by: zhagnlu <lu.zhang@zilliz.com>

Signed-off-by: zhagnlu <lu.zhang@zilliz.com>
Co-authored-by: zhagnlu <lu.zhang@zilliz.com>

internal/core/src/common/Common.cpp

@@ -21,6 +21,7 @@ namespace milvus {
 
 int64_t index_file_slice_size = DEFAULT_INDEX_FILE_SLICE_SIZE;
 int64_t thread_core_coefficient = DEFAULT_THREAD_CORE_COEFFICIENT;
+int cpu_num = DEFAULT_CPU_NUM;
 
 void
 SetIndexSliceSize(const int64_t size) {
@@ -34,4 +35,9 @@ SetThreadCoreCoefficient(const int64_t coefficient) {
     LOG_SEGCORE_DEBUG_ << "set thread pool core coefficient: " << thread_core_coefficient;
 }
 
+void
+SetCpuNum(const int num) {
+    cpu_num = num;
+}
+
 }  // namespace milvus

internal/core/src/common/Common.h

@@ -23,6 +23,7 @@ namespace milvus {
 
 extern int64_t index_file_slice_size;
 extern int64_t thread_core_coefficient;
+extern int cpu_num;
 
 void
 SetIndexSliceSize(const int64_t size);
@@ -30,4 +31,7 @@ SetIndexSliceSize(const int64_t size);
 void
 SetThreadCoreCoefficient(const int64_t coefficient);
 
+void
+SetCpuNum(const int core);
+
 }  // namespace milvus

internal/core/src/common/Consts.h

@@ -41,3 +41,5 @@ const int DEFAULT_DISK_INDEX_MAX_MEMORY_LIMIT = 2;  // gigabytes
 const int64_t DEFAULT_THREAD_CORE_COEFFICIENT = 50;
 
 const int64_t DEFAULT_INDEX_FILE_SLICE_SIZE = 4;  // megabytes
+
+const int DEFAULT_CPU_NUM = 1;

internal/core/src/common/init_c.cpp

@@ -23,7 +23,7 @@
 #include "common/Slice.h"
 #include "common/Common.h"
 
-std::once_flag flag1, flag2, flag3;
+std::once_flag flag1, flag2, flag3, flag4;
 
 void
 InitLocalRootPath(const char* root_path) {
@@ -43,3 +43,9 @@ InitThreadCoreCoefficient(const int64_t value) {
     std::call_once(
         flag3, [](int64_t value) { milvus::SetThreadCoreCoefficient(value); }, value);
 }
+
+void
+InitCpuNum(const int value) {
+    std::call_once(
+        flag4, [](int value) { milvus::SetCpuNum(value); }, value);
+}

internal/core/src/common/init_c.h

@@ -29,6 +29,9 @@ InitIndexSliceSize(const int64_t);
 void
 InitThreadCoreCoefficient(const int64_t);
 
+void
+InitCpuNum(const int);
+
 void
 InitLocalRootPath(const char*);
 

internal/core/src/storage/ThreadPool.h

@@ -34,7 +34,7 @@ namespace milvus {
 class ThreadPool {
  public:
     explicit ThreadPool(const int thread_core_coefficient) : shutdown_(false) {
-        auto thread_num = std::thread::hardware_concurrency() * thread_core_coefficient;
+        auto thread_num = cpu_num * thread_core_coefficient;
         LOG_SEGCORE_INFO_C << "Thread pool's worker num:" << thread_num;
         threads_ = std::vector<std::thread>(thread_num);
         Init();

internal/indexnode/indexnode.go

@@ -49,6 +49,7 @@ import (
 	"github.com/milvus-io/milvus/internal/proto/internalpb"
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/dependency"
+	"github.com/milvus-io/milvus/internal/util/hardware"
 	"github.com/milvus-io/milvus/internal/util/initcore"
 	"github.com/milvus-io/milvus/internal/util/paramtable"
 	"github.com/milvus-io/milvus/internal/util/sessionutil"
@@ -159,6 +160,9 @@ func (i *IndexNode) initKnowhere() {
 	cThreadCoreCoefficient := C.int64_t(Params.CommonCfg.ThreadCoreCoefficient)
 	C.InitThreadCoreCoefficient(cThreadCoreCoefficient)
 
+	cCpuNum := C.int(hardware.GetCPUNum())
+	C.InitCpuNum(cCpuNum)
+
 	initcore.InitLocalStorageConfig(&Params)
 }
 

internal/querynode/query_node.go

@@ -53,6 +53,7 @@ import (
 	"github.com/milvus-io/milvus/internal/types"
 	"github.com/milvus-io/milvus/internal/util/concurrency"
 	"github.com/milvus-io/milvus/internal/util/dependency"
+	"github.com/milvus-io/milvus/internal/util/hardware"
 	"github.com/milvus-io/milvus/internal/util/initcore"
 	"github.com/milvus-io/milvus/internal/util/lock"
 	"github.com/milvus-io/milvus/internal/util/metricsinfo"
@@ -221,6 +222,9 @@ func (node *QueryNode) InitSegcore() {
 	cThreadCoreCoefficient := C.int64_t(Params.CommonCfg.ThreadCoreCoefficient)
 	C.InitThreadCoreCoefficient(cThreadCoreCoefficient)
 
+	cCpuNum := C.int(hardware.GetCPUNum())
+	C.InitCpuNum(cCpuNum)
+
 	initcore.InitLocalStorageConfig(&Params)
 }