Skip to content
This repository was archived by the owner on Aug 16, 2023. It is now read-only.

Commit d2cc1de

Browse files
author
lixinguo
committed
Add ut when single thread
Signed-off-by: lixinguo <[email protected]>
1 parent 75177e4 commit d2cc1de

11 files changed

+171
-3
lines changed

‎unittest/CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ endif ()
3939

4040
set( UTIL_SRCS
4141
${KNOWHERE_SOURCE_DIR}/unittest/utils.cpp
42+
${KNOWHERE_SOURCE_DIR}/unittest/ThreadChecker.cpp
4243
)
4344

4445
set( ALL_INDEXING_SRCS

‎unittest/Helper.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -35,8 +35,8 @@ constexpr int64_t K = 10;
3535
constexpr int64_t PINMEM = 1024 * 1024 * 200;
3636
constexpr int64_t TEMPMEM = 1024 * 1024 * 300;
3737
constexpr int64_t RESNUM = 2;
38-
constexpr int64_t BUILD_INDEX_OMP_NUM = 3;
39-
constexpr int64_t QUERY_OMP_NUM = 4;
38+
constexpr int64_t BUILD_INDEX_OMP_NUM = 17;
39+
constexpr int64_t QUERY_OMP_NUM = 13;
4040

4141
class ParamGenerator {
4242
public:

‎unittest/ThreadChecker.cpp

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
4+
// with the License. You may obtain a copy of the License at
5+
//
6+
// http://www.apache.org/licenses/LICENSE-2.0
7+
//
8+
// Unless required by applicable law or agreed to in writing, software distributed under the License
9+
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
10+
// or implied. See the License for the specific language governing permissions and limitations under the License.
11+
12+
#include "unittest/ThreadChecker.h"
13+
14+
namespace knowhere::threadchecker {
15+
16+
#define MAX_THREAD_NUM 100000
17+
#define MIN_THREAD_NUM 0
18+
#define THREAD_LENGTH 7 // the length of max_thread_num +1
19+
20+
int32_t
21+
GetBuildOmpThread(const Config& conf) {
22+
return CheckKeyInConfig(conf, meta::BUILD_INDEX_OMP_NUM) ? GetMetaBuildIndexOmpNum(conf) : omp_get_max_threads();
23+
}
24+
25+
int32_t
26+
GetQueryOmpThread(const Config& conf) {
27+
return CheckKeyInConfig(conf, meta::QUERY_OMP_NUM) ? GetMetaQueryOmpNum(conf) : omp_get_max_threads();
28+
}
29+
30+
int32_t
31+
GetThreadNum(int pid) {
32+
std::string cmd = "ps -p " + std::to_string(pid) + " -Tf | wc -l";
33+
FILE* file = popen(cmd.c_str(), "r");
34+
char fBuff[THREAD_LENGTH];
35+
if (fgets(fBuff, sizeof(fBuff), file) == nullptr) {
36+
KNOWHERE_THROW_MSG("could not open the file to get thread number");
37+
}
38+
pclose(file);
39+
const std::size_t len = strlen(fBuff);
40+
if (len > 0 && fBuff[len - 1] == '\n') {
41+
fBuff[len - 1] = '\0';
42+
}
43+
int32_t ans = atoi(fBuff);
44+
if (ans < MIN_THREAD_NUM || ans > MAX_THREAD_NUM) {
45+
KNOWHERE_THROW_MSG("thread number is out of control");
46+
}
47+
return ans;
48+
}
49+
50+
} // namespace knowhere::threadchecker

‎unittest/ThreadChecker.h

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
// Copyright (C) 2019-2020 Zilliz. All rights reserved.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance
4+
// with the License. You may obtain a copy of the License at
5+
//
6+
// http://www.apache.org/licenses/LICENSE-2.0
7+
//
8+
// Unless required by applicable law or agreed to in writing, software distributed under the License
9+
// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
10+
// or implied. See the License for the specific language governing permissions and limitations under the License.
11+
12+
#pragma once
13+
14+
#include <omp.h>
15+
#include <stdio.h>
16+
#include <unistd.h>
17+
18+
#include "common/Log.h"
19+
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
20+
21+
namespace knowhere::threadchecker {
22+
23+
int32_t
24+
GetBuildOmpThread(const Config& conf);
25+
26+
int32_t
27+
GetQueryOmpThread(const Config& conf);
28+
29+
int32_t
30+
GetThreadNum(int pid);
31+
32+
} // namespace knowhere::threadchecker

‎unittest/test_annoy.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
1818
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
1919
#include "unittest/Helper.h"
20+
#include "unittest/ThreadChecker.h"
2021
#include "unittest/utils.h"
2122

2223
using ::testing::Combine;
@@ -146,3 +147,14 @@ TEST_P(AnnoyTest, annoy_slice) {
146147
auto result = index_->Query(query_dataset, conf_, nullptr);
147148
AssertAnns(result, nq, knowhere::GetMetaTopk(conf_));
148149
}
150+
151+
TEST_P(AnnoyTest, annoy_omp) {
152+
int pid = getpid();
153+
int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
154+
index_->BuildAll(base_dataset, conf_);
155+
int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
156+
EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
157+
auto result = index_->Query(query_dataset, conf_, nullptr);
158+
int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
159+
EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
160+
}

‎unittest/test_binaryidmap.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
1818
#include "unittest/Helper.h"
1919
#include "unittest/range_utils.h"
20+
#include "unittest/ThreadChecker.h"
2021
#include "unittest/utils.h"
2122

2223
using ::testing::Combine;
@@ -268,3 +269,14 @@ TEST_P(BinaryIDMAPTest, binaryidmap_range_search_substructure) {
268269
auto qd = knowhere::GenDataset(nq, dim, xq_bin.data());
269270
ASSERT_ANY_THROW(index_->QueryByRange(qd, conf_, nullptr));
270271
}
272+
273+
TEST_P(BinaryIDMAPTest, binaryidmap_omp) {
274+
int pid = getpid();
275+
int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
276+
index_->BuildAll(base_dataset, conf_);
277+
int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
278+
EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
279+
auto result = index_->Query(query_dataset, conf_, nullptr);
280+
int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
281+
EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
282+
}

‎unittest/test_binaryivf.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "knowhere/index/vector_index/adapter/VectorAdapter.h"
2020
#include "unittest/Helper.h"
2121
#include "unittest/range_utils.h"
22+
#include "unittest/ThreadChecker.h"
2223
#include "unittest/utils.h"
2324

2425
using ::testing::Combine;
@@ -228,3 +229,14 @@ TEST_P(BinaryIVFTest, binaryivf_range_search_substructure) {
228229
knowhere::SetMetaMetricType(conf_, knowhere::metric::SUBSTRUCTURE);
229230
ASSERT_ANY_THROW(index_->Train(base_dataset, conf_));
230231
}
232+
233+
TEST_P(BinaryIVFTest, binaryivf_omp) {
234+
int pid = getpid();
235+
int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
236+
index_->BuildAll(base_dataset, conf_);
237+
int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
238+
EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
239+
auto result = index_->Query(query_dataset, conf_, nullptr);
240+
int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
241+
EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
242+
}

‎unittest/test_hnsw.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
2121
#include "unittest/Helper.h"
2222
#include "unittest/range_utils.h"
23+
#include "unittest/ThreadChecker.h"
2324
#include "unittest/utils.h"
2425

2526
using ::testing::Combine;
@@ -315,3 +316,14 @@ TEST_P(HNSWTest, HNSW_range_trace_visit) {
315316

316317
CheckFederResult(result, nb);
317318
}
319+
320+
TEST_P(HNSWTest, HNSW_omp) {
321+
int pid = getpid();
322+
int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
323+
index_->BuildAll(base_dataset, conf_);
324+
int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
325+
EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
326+
auto result = index_->Query(query_dataset, conf_, nullptr);
327+
int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
328+
EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
329+
}

‎unittest/test_idmap.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
2626
#endif
2727
#include "unittest/range_utils.h"
28+
#include "unittest/ThreadChecker.h"
2829
#include "unittest/utils.h"
2930
#include "unittest/Helper.h"
3031

@@ -258,6 +259,17 @@ TEST_P(IDMAPTest, idmap_range_search_ip) {
258259
knowhere::KnowhereConfig::SetBlasThreshold(old_blas_threshold);
259260
}
260261

262+
TEST_P(IDMAPTest, idmap_omp) {
263+
int pid = getpid();
264+
int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
265+
index_->BuildAll(base_dataset, conf_);
266+
int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
267+
EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
268+
auto result = index_->Query(query_dataset, conf_, nullptr);
269+
int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
270+
EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
271+
}
272+
261273
#ifdef KNOWHERE_GPU_VERSION
262274
TEST_P(IDMAPTest, idmap_copy) {
263275
ASSERT_TRUE(!xb.empty());

‎unittest/test_ivf.cpp

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434

3535
#include "unittest/Helper.h"
3636
#include "unittest/range_utils.h"
37+
#include "unittest/ThreadChecker.h"
3738
#include "unittest/utils.h"
3839

3940
using ::testing::Combine;
@@ -216,6 +217,17 @@ TEST_P(IVFTest, ivf_range_search_ip) {
216217
}
217218
}
218219

220+
TEST_P(IVFTest, ivf_omp) {
221+
int pid = getpid();
222+
int32_t num_threads_before_build = knowhere::threadchecker::GetThreadNum(pid);
223+
index_->BuildAll(base_dataset, conf_);
224+
int32_t num_threads_after_build = knowhere::threadchecker::GetThreadNum(pid);
225+
EXPECT_GE(knowhere::threadchecker::GetBuildOmpThread(conf_), num_threads_after_build - num_threads_before_build + 1);
226+
auto result = index_->Query(query_dataset, conf_, nullptr);
227+
int32_t num_threads_after_query = knowhere::threadchecker::GetThreadNum(pid);
228+
EXPECT_GE(knowhere::threadchecker::GetQueryOmpThread(conf_), num_threads_after_query - num_threads_after_build + 1);
229+
}
230+
219231
// TODO(linxj): deprecated
220232
#ifdef KNOWHERE_GPU_VERSION
221233
TEST_P(IVFTest, clone_test) {

0 commit comments

Comments
 (0)