Add ut when single thread

Signed-off-by: lixinguo <[email protected]>
milvus-io · Sep 27, 2022 · d2cc1de · d2cc1de
1 parent 75177e4
commit d2cc1de
Show file tree

Hide file tree

Showing 11 changed files with 171 additions and 3 deletions.
diff --git a/unittest/CMakeLists.txt b/unittest/CMakeLists.txt
@@ -39,6 +39,7 @@ endif ()
 
 set( UTIL_SRCS
         ${KNOWHERE_SOURCE_DIR}/unittest/utils.cpp
+        ${KNOWHERE_SOURCE_DIR}/unittest/ThreadChecker.cpp
         )
 
 set( ALL_INDEXING_SRCS

diff --git a/unittest/Helper.h b/unittest/Helper.h
@@ -35,8 +35,8 @@ constexpr int64_t K = 10;
 constexpr int64_t PINMEM = 1024 * 1024 * 200;
 constexpr int64_t TEMPMEM = 1024 * 1024 * 300;
 constexpr int64_t RESNUM = 2;
-constexpr int64_t BUILD_INDEX_OMP_NUM = 3;
-constexpr int64_t QUERY_OMP_NUM = 4;
+constexpr int64_t BUILD_INDEX_OMP_NUM = 17;
+constexpr int64_t QUERY_OMP_NUM = 13;
 
 class ParamGenerator {
  public:

diff --git a/unittest/ThreadChecker.cpp b/unittest/ThreadChecker.cpp
@@ -0,0 +1,50 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+#include "unittest/ThreadChecker.h"
+
+namespace knowhere::threadchecker {
+
+#define MAX_THREAD_NUM 100000
+#define MIN_THREAD_NUM 0
+#define THREAD_LENGTH 7  // the length of max_thread_num +1
+
+int32_t
+GetBuildOmpThread(const Config& conf) {
+    return CheckKeyInConfig(conf, meta::BUILD_INDEX_OMP_NUM) ? GetMetaBuildIndexOmpNum(conf) : omp_get_max_threads();
+}
+
+int32_t
+GetQueryOmpThread(const Config& conf) {
+    return CheckKeyInConfig(conf, meta::QUERY_OMP_NUM) ? GetMetaQueryOmpNum(conf) : omp_get_max_threads();
+}
+
+int32_t
+GetThreadNum(int pid) {
+    std::string cmd = "ps -p " + std::to_string(pid) + " -Tf | wc -l";
+    FILE* file = popen(cmd.c_str(), "r");
+    char fBuff[THREAD_LENGTH];
+    if (fgets(fBuff, sizeof(fBuff), file) == nullptr) {
+        KNOWHERE_THROW_MSG("could not open the file to get thread number");
+    }
+    pclose(file);
+    const std::size_t len = strlen(fBuff);
+    if (len > 0 && fBuff[len - 1] == '\n') {
+        fBuff[len - 1] = '\0';
+    }
+    int32_t ans = atoi(fBuff);
+    if (ans < MIN_THREAD_NUM || ans > MAX_THREAD_NUM) {
+        KNOWHERE_THROW_MSG("thread number is out of control");
+    }
+    return ans;
+}
+
+}  // namespace knowhere::threadchecker
diff --git a/unittest/ThreadChecker.h b/unittest/ThreadChecker.h
@@ -0,0 +1,32 @@
+// Copyright (C) 2019-2020 Zilliz. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under the License.
+
+#pragma once
+
+#include <omp.h>
+#include <stdio.h>
+#include <unistd.h>
+
+#include "common/Log.h"
+#include "knowhere/index/vector_index/helpers/IndexParameter.h"
+
+namespace knowhere::threadchecker {
+
+int32_t
+GetBuildOmpThread(const Config& conf);
+
+int32_t
+GetQueryOmpThread(const Config& conf);
+
+int32_t
+GetThreadNum(int pid);
+
+}  // namespace knowhere::threadchecker
diff --git a/unittest/test_annoy.cpp b/unittest/test_annoy.cpp
@@ -17,6 +17,7 @@
 #include "knowhere/index/vector_index/adapter/VectorAdapter.h"
 #include "knowhere/index/vector_index/helpers/IndexParameter.h"
 #include "unittest/Helper.h"
+#include "unittest/ThreadChecker.h"
 #include "unittest/utils.h"
 
 using ::testing::Combine;
@@ -146,3 +147,14 @@ TEST_P(AnnoyTest, annoy_slice) {
     auto result = index_->Query(query_dataset, conf_, nullptr);
     AssertAnns(result, nq, knowhere::GetMetaTopk(conf_));
 }
+
+TEST_P(AnnoyTest, annoy_omp) {
+    int pid = getpid();
+    int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
+    index_->BuildAll(base_dataset, conf_);
+    int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
+    auto result = index_->Query(query_dataset, conf_, nullptr);
+    int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
+}
diff --git a/unittest/test_binaryidmap.cpp b/unittest/test_binaryidmap.cpp
@@ -17,6 +17,7 @@
 #include "knowhere/index/vector_index/adapter/VectorAdapter.h"
 #include "unittest/Helper.h"
 #include "unittest/range_utils.h"
+#include "unittest/ThreadChecker.h"
 #include "unittest/utils.h"
 
 using ::testing::Combine;
@@ -268,3 +269,14 @@ TEST_P(BinaryIDMAPTest, binaryidmap_range_search_substructure) {
     auto qd = knowhere::GenDataset(nq, dim, xq_bin.data());
     ASSERT_ANY_THROW(index_->QueryByRange(qd, conf_, nullptr));
 }
+
+TEST_P(BinaryIDMAPTest, binaryidmap_omp) {
+    int pid = getpid();
+    int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
+    index_->BuildAll(base_dataset, conf_);
+    int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
+    auto result = index_->Query(query_dataset, conf_, nullptr);
+    int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
+}
diff --git a/unittest/test_binaryivf.cpp b/unittest/test_binaryivf.cpp
@@ -19,6 +19,7 @@
 #include "knowhere/index/vector_index/adapter/VectorAdapter.h"
 #include "unittest/Helper.h"
 #include "unittest/range_utils.h"
+#include "unittest/ThreadChecker.h"
 #include "unittest/utils.h"
 
 using ::testing::Combine;
@@ -228,3 +229,14 @@ TEST_P(BinaryIVFTest, binaryivf_range_search_substructure) {
     knowhere::SetMetaMetricType(conf_, knowhere::metric::SUBSTRUCTURE);
     ASSERT_ANY_THROW(index_->Train(base_dataset, conf_));
 }
+
+TEST_P(BinaryIVFTest, binaryivf_omp) {
+    int pid = getpid();
+    int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
+    index_->BuildAll(base_dataset, conf_);
+    int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
+    auto result = index_->Query(query_dataset, conf_, nullptr);
+    int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
+}
diff --git a/unittest/test_hnsw.cpp b/unittest/test_hnsw.cpp
@@ -20,6 +20,7 @@
 #include "knowhere/index/vector_index/helpers/IndexParameter.h"
 #include "unittest/Helper.h"
 #include "unittest/range_utils.h"
+#include "unittest/ThreadChecker.h"
 #include "unittest/utils.h"
 
 using ::testing::Combine;
@@ -315,3 +316,14 @@ TEST_P(HNSWTest, HNSW_range_trace_visit) {
 
     CheckFederResult(result, nb);
 }
+
+TEST_P(HNSWTest, HNSW_omp) {
+    int pid = getpid();
+    int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
+    index_->BuildAll(base_dataset, conf_);
+    int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
+    auto result = index_->Query(query_dataset, conf_, nullptr);
+    int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
+}
diff --git a/unittest/test_idmap.cpp b/unittest/test_idmap.cpp
@@ -25,6 +25,7 @@
 #include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
 #endif
 #include "unittest/range_utils.h"
+#include "unittest/ThreadChecker.h"
 #include "unittest/utils.h"
 #include "unittest/Helper.h"
 
@@ -258,6 +259,17 @@ TEST_P(IDMAPTest, idmap_range_search_ip) {
     knowhere::KnowhereConfig::SetBlasThreshold(old_blas_threshold);
 }
 
+TEST_P(IDMAPTest, idmap_omp) {
+    int pid = getpid();
+    int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
+    index_->BuildAll(base_dataset, conf_);
+    int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
+    auto result = index_->Query(query_dataset, conf_, nullptr);
+    int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
+}
+
 #ifdef KNOWHERE_GPU_VERSION
 TEST_P(IDMAPTest, idmap_copy) {
     ASSERT_TRUE(!xb.empty());

diff --git a/unittest/test_ivf.cpp b/unittest/test_ivf.cpp
@@ -34,6 +34,7 @@
 
 #include "unittest/Helper.h"
 #include "unittest/range_utils.h"
+#include "unittest/ThreadChecker.h"
 #include "unittest/utils.h"
 
 using ::testing::Combine;
@@ -216,6 +217,17 @@ TEST_P(IVFTest, ivf_range_search_ip) {
     }
 }
 
+TEST_P(IVFTest, ivf_omp) {
+    int pid = getpid();
+    int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
+    index_->BuildAll(base_dataset, conf_);
+    int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
+    auto result = index_->Query(query_dataset, conf_, nullptr);
+    int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
+}
+
 // TODO(linxj): deprecated
 #ifdef KNOWHERE_GPU_VERSION
 TEST_P(IVFTest, clone_test) {

diff --git a/unittest/test_ivf_nm.cpp b/unittest/test_ivf_nm.cpp
@@ -34,6 +34,7 @@
 #include "knowhere/utils/distances_simd.h"
 #include "unittest/Helper.h"
 #include "unittest/range_utils.h"
+#include "unittest/ThreadChecker.h"
 #include "unittest/utils.h"
 
 using ::testing::Combine;
@@ -258,4 +259,16 @@ TEST_P(IVFNMTest, ivfnm_get_meta) {
     knowhere::Config j2 = nlohmann::json::parse(json_id_set);
     ASSERT_NO_THROW(nlohmann::from_json(j2, id_set));
     std::cout << "id_set num = " << id_set.size() << std::endl;
-}
+}
+
+TEST_P(IVFNMTest, ivfnm_omp) {
+    int pid = getpid();
+    int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
+    index_->BuildAll(base_dataset, conf_);
+    int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
+    LoadRawData(index_, base_dataset, conf_);
+    auto result = index_->Query(query_dataset, conf_, nullptr);
+    int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
+    EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
+}