milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge remote-tracking branch 'upstream/0.5.0' into 0.5.0 

Former-commit-id: b5d71ff8a05466f2693df74320eecee773292191
pull/191/head
JinHai-CN 2019-10-17 19:30:17 +08:00
parent fb431ba5d6 67c1aa5fe5
commit 7e4badb89d
35 changed files with 1113 additions and 794 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
CHANGELOG.md
View File

@ -26,7 +26,10 @@ Please mark all change in change log and use the ticket from JIRA.
- MS-653 - When config check fail, Milvus close without message
- MS-654 - Describe index timeout when building index
- MS-658 - Fix SQ8 Hybrid can't search
- #9 Change default gpu_cache_capacity to 4
- \#9 Change default gpu_cache_capacity to 4
- \#20 - C++ sdk example get grpc error
- \#23 - Add unittest to improve code coverage
- \#31 - make clang-format failed after run build.sh -l
## Improvement
- MS-552 - Add and change the easylogging library

9
core/build.sh
View File

@ -91,6 +91,10 @@ fi
cd ${BUILD_OUTPUT_DIR}
# remove make cache since build.sh -l use default variables
# force update the variables each time
make rebuild_cache
CMAKE_CMD="cmake \
-DBUILD_UNIT_TEST=${BUILD_UNITTEST} \
-DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX}
@ -115,7 +119,6 @@ if [[ ${RUN_CPPLINT} == "ON" ]]; then
make lint
if [ $? -ne 0 ]; then
echo "ERROR! cpplint check failed"
rm -f CMakeCache.txt
exit 1
fi
echo "cpplint check passed!"
@ -124,7 +127,6 @@ if [[ ${RUN_CPPLINT} == "ON" ]]; then
make check-clang-format
if [ $? -ne 0 ]; then
echo "ERROR! clang-format check failed"
rm -f CMakeCache.txt
exit 1
fi
echo "clang-format check passed!"
@ -133,12 +135,9 @@ if [[ ${RUN_CPPLINT} == "ON" ]]; then
# make check-clang-tidy
# if [ $? -ne 0 ]; then
# echo "ERROR! clang-tidy check failed"
# rm -f CMakeCache.txt
# exit 1
# fi
# echo "clang-tidy check passed!"
rm -f CMakeCache.txt
else
# compile and build
make -j 4 || exit 1

2
core/src/db/engine/ExecutionEngineImpl.cpp
View File

@ -397,6 +397,7 @@ ExecutionEngineImpl::BuildIndex(const std::string& location, EngineType engine_t
Status
ExecutionEngineImpl::Search(int64_t n, const float* data, int64_t k, int64_t nprobe, float* distances, int64_t* labels,
bool hybrid) {
#if 0
if (index_type_ == EngineType::FAISS_IVFSQ8H) {
if (!hybrid) {
const std::string key = location_ + ".quantizer";
@ -449,6 +450,7 @@ ExecutionEngineImpl::Search(int64_t n, const float* data, int64_t k, int64_t npr
}
}
}
#endif
if (index_ == nullptr) {
ENGINE_LOG_ERROR << "ExecutionEngineImpl: index is null, failed to search";

2
core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQHybrid.cpp
View File

@ -189,6 +189,8 @@ IVFSQHybrid::LoadData(const knowhere::QuantizerPtr& q, const Config& conf) {
if (quantizer_conf->mode != 2) {
KNOWHERE_THROW_MSG("mode only support 2 in this func");
}
} else {
KNOWHERE_THROW_MSG("conf error");
}
// if (quantizer_conf->gpu_id != gpu_id_) {
// KNOWHERE_THROW_MSG("quantizer and data must on the same gpu card");

4
core/src/index/knowhere/knowhere/index/vector_index/helpers/FaissGpuResourceMgr.cpp
View File

@ -63,7 +63,7 @@ FaissGpuResourceMgr::InitResource() {
mutex_cache_.emplace(device_id, std::make_unique<std::mutex>());
// std::cout << "Device Id: " << device_id << std::endl;
// std::cout << "Device Id: " << DEVICEID << std::endl;
auto& device_param = device.second;
auto& bq = idle_map_[device_id];
@ -119,7 +119,7 @@ void
FaissGpuResourceMgr::Dump() {
for (auto& item : idle_map_) {
auto& bq = item.second;
std::cout << "device_id: " << item.first << ", resource count:" << bq.Size();
std::cout << "DEVICEID: " << item.first << ", resource count:" << bq.Size();
}
}

8
core/src/index/unittest/CMakeLists.txt
View File

@ -73,9 +73,17 @@ target_link_libraries(test_kdt
SPTAGLibStatic
${depend_libs} ${unittest_libs} ${basic_libs})
add_executable(test_gpuresource test_gpuresource.cpp ${util_srcs} ${ivf_srcs})
target_link_libraries(test_gpuresource ${depend_libs} ${unittest_libs} ${basic_libs})
add_executable(test_customized_index test_customized_index.cpp ${util_srcs} ${ivf_srcs})
target_link_libraries(test_customized_index ${depend_libs} ${unittest_libs} ${basic_libs})
install(TARGETS test_ivf DESTINATION unittest)
install(TARGETS test_idmap DESTINATION unittest)
install(TARGETS test_kdt DESTINATION unittest)
install(TARGETS test_gpuresource DESTINATION unittest)
install(TARGETS test_customized_index DESTINATION unittest)
#add_subdirectory(faiss_ori)
add_subdirectory(test_nsg)

120
core/src/index/unittest/Helper.h Normal file
View File

@ -0,0 +1,120 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 <memory>
#include <string>
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVF.h"
#include "knowhere/index/vector_index/IndexIVFPQ.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
constexpr int DEVICEID = 0;
constexpr int64_t DIM = 128;
constexpr int64_t NB = 10000;
constexpr int64_t NQ = 10;
constexpr int64_t K = 10;
constexpr int64_t PINMEM = 1024 * 1024 * 200;
constexpr int64_t TEMPMEM = 1024 * 1024 * 300;
constexpr int64_t RESNUM = 2;
knowhere::IVFIndexPtr
IndexFactory(const std::string& type) {
if (type == "IVF") {
return std::make_shared<knowhere::IVF>();
} else if (type == "IVFPQ") {
return std::make_shared<knowhere::IVFPQ>();
} else if (type == "GPUIVF") {
return std::make_shared<knowhere::GPUIVF>(DEVICEID);
} else if (type == "GPUIVFPQ") {
return std::make_shared<knowhere::GPUIVFPQ>(DEVICEID);
} else if (type == "IVFSQ") {
return std::make_shared<knowhere::IVFSQ>();
} else if (type == "GPUIVFSQ") {
return std::make_shared<knowhere::GPUIVFSQ>(DEVICEID);
} else if (type == "IVFSQHybrid") {
return std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
}
}
enum class ParameterType {
ivf,
ivfpq,
ivfsq,
};
class ParamGenerator {
public:
static ParamGenerator&
GetInstance() {
static ParamGenerator instance;
return instance;
}
knowhere::Config
Gen(const ParameterType& type) {
if (type == ParameterType::ivf) {
auto tempconf = std::make_shared<knowhere::IVFCfg>();
tempconf->d = DIM;
tempconf->gpu_id = DEVICEID;
tempconf->nlist = 100;
tempconf->nprobe = 4;
tempconf->k = K;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
} else if (type == ParameterType::ivfpq) {
auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = DEVICEID;
tempconf->nlist = 100;
tempconf->nprobe = 4;
tempconf->k = K;
tempconf->m = 4;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
} else if (type == ParameterType::ivfsq) {
auto tempconf = std::make_shared<knowhere::IVFSQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = DEVICEID;
tempconf->nlist = 100;
tempconf->nprobe = 4;
tempconf->k = K;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
}
};
#include <gtest/gtest.h>
class TestGpuIndexBase : public ::testing::Test {
protected:
void
SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
}
void
TearDown() override {
knowhere::FaissGpuResourceMgr::GetInstance().Free();
}
};

122
core/src/index/unittest/test_customized_index.cpp Normal file
View File

@ -0,0 +1,122 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 <gtest/gtest.h>
#include "unittest/Helper.h"
#include "unittest/utils.h"
class SingleIndexTest : public DataGen, public TestGpuIndexBase {
protected:
void
SetUp() override {
TestGpuIndexBase::SetUp();
Generate(DIM, NB, NQ);
k = K;
}
void
TearDown() override {
TestGpuIndexBase::TearDown();
}
protected:
std::string index_type;
knowhere::IVFIndexPtr index_ = nullptr;
};
#ifdef CUSTOMIZATION
TEST_F(SingleIndexTest, IVFSQHybrid) {
assert(!xb.empty());
index_type = "IVFSQHybrid";
index_ = IndexFactory(index_type);
auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq);
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dimension(), dim);
auto binaryset = index_->Serialize();
{
// copy cpu to gpu
auto cpu_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
cpu_idx->Load(binaryset);
{
for (int i = 0; i < 3; ++i) {
auto gpu_idx = cpu_idx->CopyCpuToGpu(DEVICEID, conf);
auto result = gpu_idx->Search(query_dataset, conf);
AssertAnns(result, nq, conf->k);
// PrintResult(result, nq, k);
}
}
}
{
// quantization already in gpu, only copy data
auto cpu_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
cpu_idx->Load(binaryset);
auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf);
auto gpu_idx = pair.first;
auto quantization = pair.second;
auto result = gpu_idx->Search(query_dataset, conf);
AssertAnns(result, nq, conf->k);
// PrintResult(result, nq, k);
auto quantizer_conf = std::make_shared<knowhere::QuantizerCfg>();
quantizer_conf->mode = 2; // only copy data
quantizer_conf->gpu_id = DEVICEID;
for (int i = 0; i < 2; ++i) {
auto hybrid_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
hybrid_idx->Load(binaryset);
auto new_idx = hybrid_idx->LoadData(quantization, quantizer_conf);
auto result = new_idx->Search(query_dataset, conf);
AssertAnns(result, nq, conf->k);
// PrintResult(result, nq, k);
}
}
{
// quantization already in gpu, only set quantization
auto cpu_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
cpu_idx->Load(binaryset);
auto pair = cpu_idx->CopyCpuToGpuWithQuantizer(DEVICEID, conf);
auto quantization = pair.second;
for (int i = 0; i < 2; ++i) {
auto hybrid_idx = std::make_shared<knowhere::IVFSQHybrid>(DEVICEID);
hybrid_idx->Load(binaryset);
hybrid_idx->SetQuantizer(quantization);
auto result = hybrid_idx->Search(query_dataset, conf);
AssertAnns(result, nq, conf->k);
// PrintResult(result, nq, k);
hybrid_idx->UnsetQuantizer();
}
}
}
#endif

309
core/src/index/unittest/test_gpuresource.cpp Normal file
View File

@ -0,0 +1,309 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 <gtest/gtest.h>
#include <iostream>
#include <thread>
#include <faiss/AutoTune.h>
#include <faiss/gpu/GpuAutoTune.h>
#include <faiss/gpu/GpuIndexIVFFlat.h>
#include "knowhere/common/Exception.h"
#include "knowhere/common/Timer.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexGPUIVFPQ.h"
#include "knowhere/index/vector_index/IndexGPUIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVF.h"
#include "knowhere/index/vector_index/IndexIVFPQ.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "unittest/Helper.h"
#include "unittest/utils.h"
class GPURESTEST : public DataGen, public TestGpuIndexBase {
protected:
void
SetUp() override {
TestGpuIndexBase::SetUp();
Generate(DIM, NB, NQ);
k = K;
elems = nq * k;
ids = (int64_t*)malloc(sizeof(int64_t) * elems);
dis = (float*)malloc(sizeof(float) * elems);
}
void
TearDown() override {
delete ids;
delete dis;
TestGpuIndexBase::TearDown();
}
protected:
std::string index_type;
knowhere::IVFIndexPtr index_ = nullptr;
int64_t* ids = nullptr;
float* dis = nullptr;
int64_t elems = 0;
};
TEST_F(GPURESTEST, copyandsearch) {
// search and copy at the same time
printf("==================\n");
index_type = "GPUIVF";
index_ = IndexFactory(index_type);
auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivf);
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
auto result = index_->Search(query_dataset, conf);
AssertAnns(result, nq, k);
auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config());
cpu_idx->Seal();
auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
constexpr int64_t search_count = 50;
constexpr int64_t load_count = 15;
auto search_func = [&] {
// TimeRecorder tc("search&load");
for (int i = 0; i < search_count; ++i) {
search_idx->Search(query_dataset, conf);
// if (i > search_count - 6 || i == 0)
// tc.RecordSection("search once");
}
// tc.ElapseFromBegin("search finish");
};
auto load_func = [&] {
// TimeRecorder tc("search&load");
for (int i = 0; i < load_count; ++i) {
knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
// if (i > load_count -5 || i < 5)
// tc.RecordSection("Copy to gpu");
}
// tc.ElapseFromBegin("load finish");
};
knowhere::TimeRecorder tc("Basic");
knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
tc.RecordSection("Copy to gpu once");
search_idx->Search(query_dataset, conf);
tc.RecordSection("Search once");
search_func();
tc.RecordSection("Search total cost");
load_func();
tc.RecordSection("Copy total cost");
std::thread search_thread(search_func);
std::thread load_thread(load_func);
search_thread.join();
load_thread.join();
tc.RecordSection("Copy&Search total");
}
TEST_F(GPURESTEST, trainandsearch) {
index_type = "GPUIVF";
index_ = IndexFactory(index_type);
auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivf);
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
auto new_index = IndexFactory(index_type);
new_index->set_index_model(model);
new_index->Add(base_dataset, conf);
auto cpu_idx = knowhere::cloner::CopyGpuToCpu(new_index, knowhere::Config());
cpu_idx->Seal();
auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
constexpr int train_count = 5;
constexpr int search_count = 200;
auto train_stage = [&] {
for (int i = 0; i < train_count; ++i) {
auto model = index_->Train(base_dataset, conf);
auto test_idx = IndexFactory(index_type);
test_idx->set_index_model(model);
test_idx->Add(base_dataset, conf);
}
};
auto search_stage = [&](knowhere::VectorIndexPtr& search_idx) {
for (int i = 0; i < search_count; ++i) {
auto result = search_idx->Search(query_dataset, conf);
AssertAnns(result, nq, k);
}
};
// TimeRecorder tc("record");
// train_stage();
// tc.RecordSection("train cost");
// search_stage(search_idx);
// tc.RecordSection("search cost");
{
// search and build parallel
std::thread search_thread(search_stage, std::ref(search_idx));
std::thread train_thread(train_stage);
train_thread.join();
search_thread.join();
}
{
// build parallel
std::thread train_1(train_stage);
std::thread train_2(train_stage);
train_1.join();
train_2.join();
}
{
// search parallel
auto search_idx_2 = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
std::thread search_1(search_stage, std::ref(search_idx));
std::thread search_2(search_stage, std::ref(search_idx_2));
search_1.join();
search_2.join();
}
}
#ifdef CompareToOriFaiss
TEST_F(GPURESTEST, gpu_ivf_resource_test) {
assert(!xb.empty());
{
index_ = std::make_shared<knowhere::GPUIVF>(-1);
ASSERT_EQ(std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->GetGpuDevice(), -1);
std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->SetGpuDevice(DEVICEID);
ASSERT_EQ(std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->GetGpuDevice(), DEVICEID);
auto conf = ParamGenerator::GetInstance().Gen(ParameterType::ivfsq);
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dimension(), dim);
// knowhere::TimeRecorder tc("knowere GPUIVF");
for (int i = 0; i < search_count; ++i) {
index_->Search(query_dataset, conf);
if (i > search_count - 6 || i < 5)
// tc.RecordSection("search once");
}
// tc.ElapseFromBegin("search all");
}
knowhere::FaissGpuResourceMgr::GetInstance().Dump();
// {
// // ori faiss IVF-Search
// faiss::gpu::StandardGpuResources res;
// faiss::gpu::GpuIndexIVFFlatConfig idx_config;
// idx_config.device = DEVICEID;
// faiss::gpu::GpuIndexIVFFlat device_index(&res, dim, 1638, faiss::METRIC_L2, idx_config);
// device_index.train(nb, xb.data());
// device_index.add(nb, xb.data());
//
// knowhere::TimeRecorder tc("ori IVF");
// for (int i = 0; i < search_count; ++i) {
// device_index.search(nq, xq.data(), k, dis, ids);
// if (i > search_count - 6 || i < 5)
// tc.RecordSection("search once");
// }
// tc.ElapseFromBegin("search all");
// }
}
TEST_F(GPURESTEST, gpuivfsq) {
{
// knowhere gpu ivfsq
index_type = "GPUIVFSQ";
index_ = IndexFactory(index_type);
auto conf = std::make_shared<knowhere::IVFSQCfg>();
conf->nlist = 1638;
conf->d = dim;
conf->gpu_id = DEVICEID;
conf->metric_type = knowhere::METRICTYPE::L2;
conf->k = k;
conf->nbits = 8;
conf->nprobe = 1;
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
// auto result = index_->Search(query_dataset, conf);
// AssertAnns(result, nq, k);
auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config());
cpu_idx->Seal();
knowhere::TimeRecorder tc("knowhere GPUSQ8");
auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
tc.RecordSection("Copy to gpu");
for (int i = 0; i < search_count; ++i) {
search_idx->Search(query_dataset, conf);
if (i > search_count - 6 || i < 5)
tc.RecordSection("search once");
}
tc.ElapseFromBegin("search all");
}
{
// Ori gpuivfsq Test
const char* index_description = "IVF1638,SQ8";
faiss::Index* ori_index = faiss::index_factory(dim, index_description, faiss::METRIC_L2);
faiss::gpu::StandardGpuResources res;
auto device_index = faiss::gpu::index_cpu_to_gpu(&res, DEVICEID, ori_index);
device_index->train(nb, xb.data());
device_index->add(nb, xb.data());
auto cpu_index = faiss::gpu::index_gpu_to_cpu(device_index);
auto idx = dynamic_cast<faiss::IndexIVF*>(cpu_index);
if (idx != nullptr) {
idx->to_readonly();
}
delete device_index;
delete ori_index;
faiss::gpu::GpuClonerOptions option;
option.allInGpu = true;
knowhere::TimeRecorder tc("ori GPUSQ8");
faiss::Index* search_idx = faiss::gpu::index_cpu_to_gpu(&res, DEVICEID, cpu_index, &option);
tc.RecordSection("Copy to gpu");
for (int i = 0; i < search_count; ++i) {
search_idx->search(nq, xq.data(), k, dis, ids);
if (i > search_count - 6 || i < 5)
tc.RecordSection("search once");
}
tc.ElapseFromBegin("search all");
delete cpu_index;
delete search_idx;
}
}
#endif

50
core/src/index/unittest/test_idmap.cpp
View File

@ -23,54 +23,28 @@
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "Helper.h"
#include "unittest/utils.h"
static int device_id = 0;
class IDMAPTest : public DataGen, public ::testing::Test {
class IDMAPTest : public DataGen, public TestGpuIndexBase {
protected:
void
SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024 * 1024 * 200, 1024 * 1024 * 300, 2);
TestGpuIndexBase::SetUp();
Init_with_default();
index_ = std::make_shared<knowhere::IDMAP>();
}
void
TearDown() override {
knowhere::FaissGpuResourceMgr::GetInstance().Free();
TestGpuIndexBase::TearDown();
}
protected:
knowhere::IDMAPPtr index_ = nullptr;
};
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
auto dists = result->array()[1];
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < 10; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}
TEST_F(IDMAPTest, idmap_basic) {
ASSERT_TRUE(!xb.empty());
@ -87,7 +61,7 @@ TEST_F(IDMAPTest, idmap_basic) {
ASSERT_TRUE(index_->GetRawIds() != nullptr);
auto result = index_->Search(query_dataset, conf);
AssertAnns(result, nq, k);
PrintResult(result, nq, k);
// PrintResult(result, nq, k);
index_->Seal();
auto binaryset = index_->Serialize();
@ -95,7 +69,7 @@ TEST_F(IDMAPTest, idmap_basic) {
new_index->Load(binaryset);
auto re_result = index_->Search(query_dataset, conf);
AssertAnns(re_result, nq, k);
PrintResult(re_result, nq, k);
// PrintResult(re_result, nq, k);
}
TEST_F(IDMAPTest, idmap_serialize) {
@ -118,7 +92,7 @@ TEST_F(IDMAPTest, idmap_serialize) {
index_->Add(base_dataset, knowhere::Config());
auto re_result = index_->Search(query_dataset, conf);
AssertAnns(re_result, nq, k);
PrintResult(re_result, nq, k);
// PrintResult(re_result, nq, k);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dimension(), dim);
auto binaryset = index_->Serialize();
@ -138,7 +112,7 @@ TEST_F(IDMAPTest, idmap_serialize) {
EXPECT_EQ(index_->Dimension(), dim);
auto result = index_->Search(query_dataset, conf);
AssertAnns(result, nq, k);
PrintResult(result, nq, k);
// PrintResult(result, nq, k);
}
}
@ -169,7 +143,7 @@ TEST_F(IDMAPTest, copy_test) {
{
// cpu to gpu
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, conf);
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf);
auto clone_result = clone_index->Search(query_dataset, conf);
AssertAnns(clone_result, nq, k);
ASSERT_THROW({ std::static_pointer_cast<knowhere::GPUIDMAP>(clone_index)->GetRawVectors(); },
@ -194,9 +168,9 @@ TEST_F(IDMAPTest, copy_test) {
ASSERT_TRUE(std::static_pointer_cast<knowhere::IDMAP>(host_index)->GetRawIds() != nullptr);
// gpu to gpu
auto device_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, conf);
auto device_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, conf);
auto new_device_index =
std::static_pointer_cast<knowhere::GPUIDMAP>(device_index)->CopyGpuToGpu(device_id, conf);
std::static_pointer_cast<knowhere::GPUIDMAP>(device_index)->CopyGpuToGpu(DEVICEID, conf);
auto device_result = new_device_index->Search(query_dataset, conf);
AssertAnns(device_result, nq, k);
}

498
core/src/index/unittest/test_ivf.cpp
View File

@ -35,99 +35,25 @@
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "unittest/Helper.h"
#include "unittest/utils.h"
using ::testing::Combine;
using ::testing::TestWithParam;
using ::testing::Values;
constexpr int device_id = 0;
constexpr int64_t DIM = 128;
constexpr int64_t NB = 1000000 / 100;
constexpr int64_t NQ = 10;
constexpr int64_t K = 10;
knowhere::IVFIndexPtr
IndexFactory(const std::string& type) {
if (type == "IVF") {
return std::make_shared<knowhere::IVF>();
} else if (type == "IVFPQ") {
return std::make_shared<knowhere::IVFPQ>();
} else if (type == "GPUIVF") {
return std::make_shared<knowhere::GPUIVF>(device_id);
} else if (type == "GPUIVFPQ") {
return std::make_shared<knowhere::GPUIVFPQ>(device_id);
} else if (type == "IVFSQ") {
return std::make_shared<knowhere::IVFSQ>();
} else if (type == "GPUIVFSQ") {
return std::make_shared<knowhere::GPUIVFSQ>(device_id);
} else if (type == "IVFSQHybrid") {
return std::make_shared<knowhere::IVFSQHybrid>(device_id);
}
}
enum class ParameterType {
ivf,
ivfpq,
ivfsq,
nsg,
};
class ParamGenerator {
public:
static ParamGenerator&
GetInstance() {
static ParamGenerator instance;
return instance;
}
knowhere::Config
Gen(const ParameterType& type) {
if (type == ParameterType::ivf) {
auto tempconf = std::make_shared<knowhere::IVFCfg>();
tempconf->d = DIM;
tempconf->gpu_id = device_id;
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->k = K;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
} else if (type == ParameterType::ivfpq) {
auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = device_id;
tempconf->nlist = 25;
tempconf->nprobe = 4;
tempconf->k = K;
tempconf->m = 4;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
} else if (type == ParameterType::ivfsq) {
auto tempconf = std::make_shared<knowhere::IVFSQCfg>();
tempconf->d = DIM;
tempconf->gpu_id = device_id;
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->k = K;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
}
};
class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, ParameterType>> {
protected:
void
SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
ParameterType parameter_type;
std::tie(index_type, parameter_type) = GetParam();
// Init_with_default();
Generate(DIM, NB, NQ);
index_ = IndexFactory(index_type);
conf = ParamGenerator::GetInstance().Gen(parameter_type);
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024 * 1024 * 200, 1024 * 1024 * 600, 2);
}
void
@ -140,7 +66,7 @@ class IVFTest : public DataGen, public TestWithParam<::std::tuple<std::string, P
std::vector<std::string> gpu_idx{"GPUIVFSQ"};
auto finder = std::find(gpu_idx.cbegin(), gpu_idx.cend(), index_type);
if (finder != gpu_idx.cend()) {
return knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config());
return knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config());
}
return index_;
}
@ -162,33 +88,6 @@ INSTANTIATE_TEST_CASE_P(IVFParameters, IVFTest,
#endif
std::make_tuple("GPUIVFSQ", ParameterType::ivfsq)));
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
auto dists = result->array()[1];
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < 10; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}
TEST_P(IVFTest, ivf_basic) {
assert(!xb.empty());
@ -207,85 +106,6 @@ TEST_P(IVFTest, ivf_basic) {
// PrintResult(result, nq, k);
}
TEST_P(IVFTest, hybrid) {
if (index_type != "IVFSQHybrid") {
return;
}
assert(!xb.empty());
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dimension(), dim);
// auto new_idx = ChooseTodo();
// auto result = new_idx->Search(query_dataset, conf);
// AssertAnns(result, nq, conf->k);
{
auto hybrid_1_idx = std::make_shared<knowhere::IVFSQHybrid>(device_id);
auto binaryset = index_->Serialize();
hybrid_1_idx->Load(binaryset);
auto quantizer_conf = std::make_shared<knowhere::QuantizerCfg>();
quantizer_conf->mode = 1;
quantizer_conf->gpu_id = device_id;
auto q = hybrid_1_idx->LoadQuantizer(quantizer_conf);
hybrid_1_idx->SetQuantizer(q);
auto result = hybrid_1_idx->Search(query_dataset, conf);
AssertAnns(result, nq, conf->k);
PrintResult(result, nq, k);
hybrid_1_idx->UnsetQuantizer();
}
{
auto hybrid_2_idx = std::make_shared<knowhere::IVFSQHybrid>(device_id);
auto binaryset = index_->Serialize();
hybrid_2_idx->Load(binaryset);
auto quantizer_conf = std::make_shared<knowhere::QuantizerCfg>();
quantizer_conf->mode = 1;
quantizer_conf->gpu_id = device_id;
auto q = hybrid_2_idx->LoadQuantizer(quantizer_conf);
quantizer_conf->mode = 2;
auto gpu_idx = hybrid_2_idx->LoadData(q, quantizer_conf);
auto result = gpu_idx->Search(query_dataset, conf);
AssertAnns(result, nq, conf->k);
PrintResult(result, nq, k);
}
}
// TEST_P(IVFTest, gpu_to_cpu) {
// if (index_type.find("GPU") == std::string::npos) { return; }
//
// // else
// assert(!xb.empty());
//
// auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
// index_->set_preprocessor(preprocessor);
//
// auto model = index_->Train(base_dataset, conf);
// index_->set_index_model(model);
// index_->Add(base_dataset, conf);
// EXPECT_EQ(index_->Count(), nb);
// EXPECT_EQ(index_->Dimension(), dim);
// auto result = index_->Search(query_dataset, conf);
// AssertAnns(result, nq, k);
//
// if (auto device_index = std::dynamic_pointer_cast<GPUIVF>(index_)) {
// auto host_index = device_index->Copy_index_gpu_to_cpu();
// auto result = host_index->Search(query_dataset, conf);
// AssertAnns(result, nq, k);
// }
//}
TEST_P(IVFTest, ivf_serialize) {
auto serialize = [](const std::string& filename, knowhere::BinaryPtr& bin, uint8_t* ret) {
FileIOWriter writer(filename);
@ -423,7 +243,7 @@ TEST_P(IVFTest, clone_test) {
auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
if (finder != support_idx_vec.cend()) {
EXPECT_NO_THROW({
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config());
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config());
auto clone_result = clone_index->Search(query_dataset, conf);
AssertEqual(result, clone_result);
std::cout << "clone C <=> G [" << index_type << "] success" << std::endl;
@ -432,7 +252,7 @@ TEST_P(IVFTest, clone_test) {
EXPECT_THROW(
{
std::cout << "clone C <=> G [" << index_type << "] failed" << std::endl;
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, device_id, knowhere::Config());
auto clone_index = knowhere::cloner::CopyCpuToGpu(index_, DEVICEID, knowhere::Config());
},
knowhere::KnowhereException);
}
@ -440,9 +260,7 @@ TEST_P(IVFTest, clone_test) {
}
#ifdef CUSTOMIZATION
TEST_P(IVFTest, seal_test) {
// FaissGpuResourceMgr::GetInstance().InitDevice(device_id);
TEST_P(IVFTest, gpu_seal_test) {
std::vector<std::string> support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"};
auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type);
if (finder == support_idx_vec.cend()) {
@ -466,309 +284,13 @@ TEST_P(IVFTest, seal_test) {
auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config());
knowhere::TimeRecorder tc("CopyToGpu");
knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
auto without_seal = tc.RecordSection("Without seal");
cpu_idx->Seal();
tc.RecordSection("seal cost");
knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
knowhere::cloner::CopyCpuToGpu(cpu_idx, DEVICEID, knowhere::Config());
auto with_seal = tc.RecordSection("With seal");
ASSERT_GE(without_seal, with_seal);
}
#endif
class GPURESTEST : public DataGen, public ::testing::Test {
protected:
void
SetUp() override {
Generate(128, 1000000, 1000);
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024 * 1024 * 200, 1024 * 1024 * 300, 2);
k = 100;
elems = nq * k;
ids = (int64_t*)malloc(sizeof(int64_t) * elems);
dis = (float*)malloc(sizeof(float) * elems);
}
void
TearDown() override {
delete ids;
delete dis;
knowhere::FaissGpuResourceMgr::GetInstance().Free();
}
protected:
std::string index_type;
knowhere::IVFIndexPtr index_ = nullptr;
int64_t* ids = nullptr;
float* dis = nullptr;
int64_t elems = 0;
};
const int search_count = 18;
const int load_count = 3;
TEST_F(GPURESTEST, gpu_ivf_resource_test) {
assert(!xb.empty());
{
index_ = std::make_shared<knowhere::GPUIVF>(-1);
ASSERT_EQ(std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->GetGpuDevice(), -1);
std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->SetGpuDevice(device_id);
ASSERT_EQ(std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)->GetGpuDevice(), device_id);
auto conf = std::make_shared<knowhere::IVFCfg>();
conf->nlist = 1638;
conf->d = dim;
conf->gpu_id = device_id;
conf->metric_type = knowhere::METRICTYPE::L2;
conf->k = k;
conf->nprobe = 1;
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dimension(), dim);
knowhere::TimeRecorder tc("knowere GPUIVF");
for (int i = 0; i < search_count; ++i) {
index_->Search(query_dataset, conf);
if (i > search_count - 6 || i < 5)
tc.RecordSection("search once");
}
tc.ElapseFromBegin("search all");
}
knowhere::FaissGpuResourceMgr::GetInstance().Dump();
{
// IVF-Search
faiss::gpu::StandardGpuResources res;
faiss::gpu::GpuIndexIVFFlatConfig idx_config;
idx_config.device = device_id;
faiss::gpu::GpuIndexIVFFlat device_index(&res, dim, 1638, faiss::METRIC_L2, idx_config);
device_index.train(nb, xb.data());
device_index.add(nb, xb.data());
knowhere::TimeRecorder tc("ori IVF");
for (int i = 0; i < search_count; ++i) {
device_index.search(nq, xq.data(), k, dis, ids);
if (i > search_count - 6 || i < 5)
tc.RecordSection("search once");
}
tc.ElapseFromBegin("search all");
}
}
#ifdef CUSTOMIZATION
TEST_F(GPURESTEST, gpuivfsq) {
{
// knowhere gpu ivfsq
index_type = "GPUIVFSQ";
index_ = IndexFactory(index_type);
auto conf = std::make_shared<knowhere::IVFSQCfg>();
conf->nlist = 1638;
conf->d = dim;
conf->gpu_id = device_id;
conf->metric_type = knowhere::METRICTYPE::L2;
conf->k = k;
conf->nbits = 8;
conf->nprobe = 1;
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
// auto result = index_->Search(query_dataset, conf);
// AssertAnns(result, nq, k);
auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config());
cpu_idx->Seal();
knowhere::TimeRecorder tc("knowhere GPUSQ8");
auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
tc.RecordSection("Copy to gpu");
for (int i = 0; i < search_count; ++i) {
search_idx->Search(query_dataset, conf);
if (i > search_count - 6 || i < 5)
tc.RecordSection("search once");
}
tc.ElapseFromBegin("search all");
}
{
// Ori gpuivfsq Test
const char* index_description = "IVF1638,SQ8";
faiss::Index* ori_index = faiss::index_factory(dim, index_description, faiss::METRIC_L2);
faiss::gpu::StandardGpuResources res;
auto device_index = faiss::gpu::index_cpu_to_gpu(&res, device_id, ori_index);
device_index->train(nb, xb.data());
device_index->add(nb, xb.data());
auto cpu_index = faiss::gpu::index_gpu_to_cpu(device_index);
auto idx = dynamic_cast<faiss::IndexIVF*>(cpu_index);
if (idx != nullptr) {
idx->to_readonly();
}
delete device_index;
delete ori_index;
faiss::gpu::GpuClonerOptions option;
option.allInGpu = true;
knowhere::TimeRecorder tc("ori GPUSQ8");
faiss::Index* search_idx = faiss::gpu::index_cpu_to_gpu(&res, device_id, cpu_index, &option);
tc.RecordSection("Copy to gpu");
for (int i = 0; i < search_count; ++i) {
search_idx->search(nq, xq.data(), k, dis, ids);
if (i > search_count - 6 || i < 5)
tc.RecordSection("search once");
}
tc.ElapseFromBegin("search all");
delete cpu_index;
delete search_idx;
}
}
#endif
TEST_F(GPURESTEST, copyandsearch) {
// search and copy at the same time
printf("==================\n");
index_type = "GPUIVF";
index_ = IndexFactory(index_type);
auto conf = std::make_shared<knowhere::IVFSQCfg>();
conf->nlist = 1638;
conf->d = dim;
conf->gpu_id = device_id;
conf->metric_type = knowhere::METRICTYPE::L2;
conf->k = k;
conf->nbits = 8;
conf->nprobe = 1;
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
index_->set_index_model(model);
index_->Add(base_dataset, conf);
// auto result = index_->Search(query_dataset, conf);
// AssertAnns(result, nq, k);
auto cpu_idx = knowhere::cloner::CopyGpuToCpu(index_, knowhere::Config());
cpu_idx->Seal();
auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
auto search_func = [&] {
// TimeRecorder tc("search&load");
for (int i = 0; i < search_count; ++i) {
search_idx->Search(query_dataset, conf);
// if (i > search_count - 6 || i == 0)
// tc.RecordSection("search once");
}
// tc.ElapseFromBegin("search finish");
};
auto load_func = [&] {
// TimeRecorder tc("search&load");
for (int i = 0; i < load_count; ++i) {
knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
// if (i > load_count -5 || i < 5)
// tc.RecordSection("Copy to gpu");
}
// tc.ElapseFromBegin("load finish");
};
knowhere::TimeRecorder tc("basic");
knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
tc.RecordSection("Copy to gpu once");
search_idx->Search(query_dataset, conf);
tc.RecordSection("search once");
search_func();
tc.RecordSection("only search total");
load_func();
tc.RecordSection("only copy total");
std::thread search_thread(search_func);
std::thread load_thread(load_func);
search_thread.join();
load_thread.join();
tc.RecordSection("Copy&search total");
}
TEST_F(GPURESTEST, TrainAndSearch) {
index_type = "GPUIVF";
index_ = IndexFactory(index_type);
auto conf = std::make_shared<knowhere::IVFSQCfg>();
conf->nlist = 1638;
conf->d = dim;
conf->gpu_id = device_id;
conf->metric_type = knowhere::METRICTYPE::L2;
conf->k = k;
conf->nbits = 8;
conf->nprobe = 1;
auto preprocessor = index_->BuildPreprocessor(base_dataset, conf);
index_->set_preprocessor(preprocessor);
auto model = index_->Train(base_dataset, conf);
auto new_index = IndexFactory(index_type);
new_index->set_index_model(model);
new_index->Add(base_dataset, conf);
auto cpu_idx = knowhere::cloner::CopyGpuToCpu(new_index, knowhere::Config());
cpu_idx->Seal();
auto search_idx = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
constexpr int train_count = 1;
constexpr int search_count = 5000;
auto train_stage = [&] {
for (int i = 0; i < train_count; ++i) {
auto model = index_->Train(base_dataset, conf);
auto test_idx = IndexFactory(index_type);
test_idx->set_index_model(model);
test_idx->Add(base_dataset, conf);
}
};
auto search_stage = [&](knowhere::VectorIndexPtr& search_idx) {
for (int i = 0; i < search_count; ++i) {
auto result = search_idx->Search(query_dataset, conf);
AssertAnns(result, nq, k);
}
};
// TimeRecorder tc("record");
// train_stage();
// tc.RecordSection("train cost");
// search_stage(search_idx);
// tc.RecordSection("search cost");
{
// search and build parallel
std::thread search_thread(search_stage, std::ref(search_idx));
std::thread train_thread(train_stage);
train_thread.join();
search_thread.join();
}
{
// build parallel
std::thread train_1(train_stage);
std::thread train_2(train_stage);
train_1.join();
train_2.join();
}
{
// search parallel
auto search_idx_2 = knowhere::cloner::CopyCpuToGpu(cpu_idx, device_id, knowhere::Config());
std::thread search_1(search_stage, std::ref(search_idx));
std::thread search_2(search_stage, std::ref(search_idx_2));
search_1.join();
search_2.join();
}
}
// TODO(lxj): Add exception test

27
core/src/index/unittest/test_kdt.cpp
View File

@ -52,33 +52,6 @@ class KDTTest : public DataGen, public ::testing::Test {
std::shared_ptr<knowhere::CPUKDTRNG> index_ = nullptr;
};
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
auto dists = result->array()[1];
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < 10; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}
// TODO(lxj): add test about count() and dimension()
TEST_F(KDTTest, kdt_basic) {
assert(!xb.empty());

14
core/src/index/unittest/test_nsg/test_nsg.cpp
View File

@ -30,19 +30,19 @@ using ::testing::Combine;
using ::testing::TestWithParam;
using ::testing::Values;
constexpr int64_t DEVICE_ID = 1;
constexpr int64_t DEVICEID = 0;
class NSGInterfaceTest : public DataGen, public ::testing::Test {
protected:
void
SetUp() override {
// Init_with_default();
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICE_ID, 1024 * 1024 * 200, 1024 * 1024 * 600, 2);
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, 1024 * 1024 * 200, 1024 * 1024 * 600, 2);
Generate(256, 1000000 / 100, 1);
index_ = std::make_shared<knowhere::NSG>();
auto tmp_conf = std::make_shared<knowhere::NSGCfg>();
tmp_conf->gpu_id = DEVICE_ID;
tmp_conf->gpu_id = DEVICEID;
tmp_conf->knng = 20;
tmp_conf->nprobe = 8;
tmp_conf->nlist = 163;
@ -69,14 +69,6 @@ class NSGInterfaceTest : public DataGen, public ::testing::Test {
knowhere::Config search_conf;
};
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
TEST_F(NSGInterfaceTest, basic_test) {
assert(!xb.empty());

28
core/src/index/unittest/utils.cpp
View File

@ -17,6 +17,7 @@
#include "unittest/utils.h"
#include <gtest/gtest.h>
#include <memory>
#include <string>
#include <utility>
@ -147,3 +148,30 @@ generate_query_dataset(int64_t nb, int64_t dim, float* xb) {
auto dataset = std::make_shared<knowhere::Dataset>(std::move(tensors), tensor_schema);
return dataset;
}
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(i, *(ids->data()->GetValues<int64_t>(1, i * k)));
}
}
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k) {
auto ids = result->array()[0];
auto dists = result->array()[1];
std::stringstream ss_id;
std::stringstream ss_dist;
for (auto i = 0; i < 10; i++) {
for (auto j = 0; j < k; ++j) {
ss_id << *(ids->data()->GetValues<int64_t>(1, i * k + j)) << " ";
ss_dist << *(dists->data()->GetValues<float>(1, i * k + j)) << " ";
}
ss_id << std::endl;
ss_dist << std::endl;
}
std::cout << "id\n" << ss_id.str() << std::endl;
std::cout << "dist\n" << ss_dist.str() << std::endl;
}

6
core/src/index/unittest/utils.h
View File

@ -68,6 +68,12 @@ generate_dataset(int64_t nb, int64_t dim, float* xb, int64_t* ids);
knowhere::DatasetPtr
generate_query_dataset(int64_t nb, int64_t dim, float* xb);
void
AssertAnns(const knowhere::DatasetPtr& result, const int& nq, const int& k);
void
PrintResult(const knowhere::DatasetPtr& result, const int& nq, const int& k);
struct FileIOWriter {
std::fstream fs;
std::string name;

84
core/src/scheduler/optimizer/LargeSQ8HPass.cpp
View File

@ -26,48 +26,48 @@
namespace milvus {
namespace scheduler {
bool
LargeSQ8HPass::Run(const TaskPtr& task) {
if (task->Type() != TaskType::SearchTask) {
return false;
}
auto search_task = std::static_pointer_cast<XSearchTask>(task);
if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFSQ8H) {
return false;
}
auto search_job = std::static_pointer_cast<SearchJob>(search_task->job_.lock());
// TODO: future, Index::IVFSQ8H, if nq < threshold set cpu, else set gpu
if (search_job->nq() < 100) {
return false;
}
std::vector<uint64_t> gpus = scheduler::get_gpu_pool();
std::vector<int64_t> all_free_mem;
for (auto& gpu : gpus) {
auto cache = cache::GpuCacheMgr::GetInstance(gpu);
auto free_mem = cache->CacheCapacity() - cache->CacheUsage();
all_free_mem.push_back(free_mem);
}
auto max_e = std::max_element(all_free_mem.begin(), all_free_mem.end());
auto best_index = std::distance(all_free_mem.begin(), max_e);
auto best_device_id = gpus[best_index];
ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, best_device_id);
if (not res_ptr) {
SERVER_LOG_ERROR << "GpuResource " << best_device_id << " invalid.";
// TODO: throw critical error and exit
return false;
}
auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
task->label() = label;
return true;
}
// bool
// LargeSQ8HPass::Run(const TaskPtr& task) {
// if (task->Type() != TaskType::SearchTask) {
// return false;
// }
//
// auto search_task = std::static_pointer_cast<XSearchTask>(task);
// if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFSQ8H) {
// return false;
// }
//
// auto search_job = std::static_pointer_cast<SearchJob>(search_task->job_.lock());
//
// // TODO: future, Index::IVFSQ8H, if nq < threshold set cpu, else set gpu
// if (search_job->nq() < 100) {
// return false;
// }
//
// std::vector<uint64_t> gpus = scheduler::get_gpu_pool();
// std::vector<int64_t> all_free_mem;
// for (auto& gpu : gpus) {
// auto cache = cache::GpuCacheMgr::GetInstance(gpu);
// auto free_mem = cache->CacheCapacity() - cache->CacheUsage();
// all_free_mem.push_back(free_mem);
// }
//
// auto max_e = std::max_element(all_free_mem.begin(), all_free_mem.end());
// auto best_index = std::distance(all_free_mem.begin(), max_e);
// auto best_device_id = gpus[best_index];
//
// ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, best_device_id);
// if (not res_ptr) {
// SERVER_LOG_ERROR << "GpuResource " << best_device_id << " invalid.";
// // TODO: throw critical error and exit
// return false;
// }
//
// auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
// task->label() = label;
//
// return true;
// }
} // namespace scheduler
} // namespace milvus

4
core/src/scheduler/optimizer/LargeSQ8HPass.h
View File

@ -37,8 +37,8 @@ class LargeSQ8HPass : public Pass {
LargeSQ8HPass() = default;
public:
bool
Run(const TaskPtr& task) override;
// bool
// Run(const TaskPtr& task) override;
};
using LargeSQ8HPassPtr = std::shared_ptr<LargeSQ8HPass>;

12
core/src/scheduler/optimizer/Optimizer.cpp
View File

@ -20,12 +20,12 @@
namespace milvus {
namespace scheduler {
void
Optimizer::Init() {
for (auto& pass : pass_list_) {
pass->Init();
}
}
// void
// Optimizer::Init() {
// for (auto& pass : pass_list_) {
// pass->Init();
// }
// }
bool
Optimizer::Run(const TaskPtr& task) {

4
core/src/scheduler/optimizer/Optimizer.h
View File

@ -38,8 +38,8 @@ class Optimizer {
explicit Optimizer(std::vector<PassPtr> pass_list) : pass_list_(std::move(pass_list)) {
}
void
Init();
// void
// Init();
bool
Run(const TaskPtr& task);

8
core/src/scheduler/task/BuildIndexTask.cpp
View File

@ -55,9 +55,6 @@ XBuildIndexTask::Load(milvus::scheduler::LoadType type, uint8_t device_id) {
} else if (type == LoadType::CPU2GPU) {
stat = to_index_engine_->CopyToIndexFileToGpu(device_id);
type_str = "CPU2GPU";
} else if (type == LoadType::GPU2CPU) {
stat = to_index_engine_->CopyToCpu();
type_str = "GPU2CPU";
} else {
error_msg = "Wrong load type";
stat = Status(SERVER_UNEXPECTED_ERROR, error_msg);
@ -199,8 +196,9 @@ XBuildIndexTask::Execute() {
ENGINE_LOG_DEBUG << "New index file " << table_file.file_id_ << " of size " << index->PhysicalSize()
<< " bytes"
<< " from file " << origin_file.file_id_;
// index->Cache();
if (build_index_job->options().insert_cache_immediately_) {
index->Cache();
}
} else {
// failed to update meta, mark the new file as to_delete, don't delete old file
origin_file.file_type_ = engine::meta::TableFileSchema::TO_INDEX;

12
core/src/sdk/examples/grpcsimple/src/ClientTest.cpp
View File

@ -40,8 +40,10 @@ constexpr int64_t BATCH_ROW_COUNT = 100000;
constexpr int64_t NQ = 5;
constexpr int64_t TOP_K = 10;
constexpr int64_t SEARCH_TARGET = 5000; // change this value, result is different
constexpr int64_t ADD_VECTOR_LOOP = 1;
constexpr int64_t ADD_VECTOR_LOOP = 5;
constexpr int64_t SECONDS_EACH_HOUR = 3600;
constexpr milvus::IndexType INDEX_TYPE = milvus::IndexType::gpu_ivfsq8;
constexpr int32_t N_LIST = 15000;
#define BLOCK_SPLITER std::cout << "===========================================" << std::endl;
@ -311,8 +313,8 @@ ClientTest::Test(const std::string& address, const std::string& port) {
std::cout << "Wait until create all index done" << std::endl;
milvus::IndexParam index;
index.table_name = TABLE_NAME;
index.index_type = milvus::IndexType::gpu_ivfsq8;
index.nlist = 16384;
index.index_type = INDEX_TYPE;
index.nlist = N_LIST;
milvus::Status stat = conn->CreateIndex(index);
std::cout << "CreateIndex function call status: " << stat.message() << std::endl;
@ -344,8 +346,8 @@ ClientTest::Test(const std::string& address, const std::string& port) {
{ // delete by range
milvus::Range rg;
rg.start_value = CurrentTmDate(-2);
rg.end_value = CurrentTmDate(-3);
rg.start_value = CurrentTmDate(-3);
rg.end_value = CurrentTmDate(-2);
milvus::Status stat = conn->DeleteByRange(rg, TABLE_NAME);
std::cout << "DeleteByRange function call status: " << stat.message() << std::endl;

2
core/src/sdk/interface/Status.cpp
View File

@ -88,7 +88,7 @@ Status::MoveFrom(Status& s) {
std::string
Status::message() const {
if (state_ == nullptr) {
return "";
return "OK";
}
std::string msg;

2
core/src/utils/Status.cpp
View File

@ -88,7 +88,7 @@ Status::MoveFrom(Status& s) {
std::string
Status::message() const {
if (state_ == nullptr) {
return "";
return "OK";
}
std::string msg;

18
core/unittest/db/test_db.cpp
View File

@ -308,6 +308,12 @@ TEST_F(DBTest, SEARCH_TEST) {
ASSERT_TRUE(stat.ok());
}
{
milvus::engine::QueryResults large_nq_results;
stat = db_->Query(TABLE_NAME, k, 200, 10, xq.data(), large_nq_results);
ASSERT_TRUE(stat.ok());
}
{//search by specify index file
milvus::engine::meta::DatesT dates;
std::vector<std::string> file_ids = {"1", "2", "3", "4", "5", "6"};
@ -315,6 +321,8 @@ TEST_F(DBTest, SEARCH_TEST) {
stat = db_->Query(TABLE_NAME, file_ids, k, nq, 10, xq.data(), dates, results);
ASSERT_TRUE(stat.ok());
}
#endif
}
@ -412,6 +420,16 @@ TEST_F(DBTest, INDEX_TEST) {
stat = db_->CreateIndex(table_info.table_id_, index);
ASSERT_TRUE(stat.ok());
index.engine_type_ = (int) milvus::engine::EngineType::FAISS_IVFFLAT;
stat = db_->CreateIndex(table_info.table_id_, index);
ASSERT_TRUE(stat.ok());
#ifdef CUSTOMIZATION
index.engine_type_ = (int)milvus::engine::EngineType::FAISS_IVFSQ8H;
stat = db_->CreateIndex(table_info.table_id_, index);
ASSERT_TRUE(stat.ok());
#endif
milvus::engine::TableIndex index_out;
stat = db_->DescribeIndex(table_info.table_id_, index_out);
ASSERT_TRUE(stat.ok());

23
core/unittest/db/test_engine.cpp
View File

@ -108,15 +108,16 @@ TEST_F(EngineTest, ENGINE_IMPL_TEST) {
ASSERT_EQ(engine_ptr->Dimension(), dimension);
ASSERT_EQ(engine_ptr->Count(), ids.size());
// status = engine_ptr->CopyToGpu(0);
// //ASSERT_TRUE(status.ok());
//
// auto new_engine = engine_ptr->Clone();
// ASSERT_EQ(new_engine->Dimension(), dimension);
// ASSERT_EQ(new_engine->Count(), ids.size());
// status = new_engine->CopyToCpu();
// //ASSERT_TRUE(status.ok());
//
// auto engine_build = new_engine->BuildIndex("/tmp/milvus_index_2", engine::EngineType::FAISS_IVFSQ8);
// //ASSERT_TRUE(status.ok());
status = engine_ptr->CopyToGpu(0, true);
status = engine_ptr->CopyToGpu(0, false);
//ASSERT_TRUE(status.ok());
auto new_engine = engine_ptr->Clone();
ASSERT_EQ(new_engine->Dimension(), dimension);
ASSERT_EQ(new_engine->Count(), ids.size());
status = new_engine->CopyToCpu();
//ASSERT_TRUE(status.ok());
auto engine_build = new_engine->BuildIndex("/tmp/milvus_index_2", milvus::engine::EngineType::FAISS_IVFSQ8);
//ASSERT_TRUE(status.ok());
}

6
core/unittest/scheduler/task_test.cpp
View File

@ -17,6 +17,7 @@
#include "scheduler/task/SearchTask.h"
#include "scheduler/task/BuildIndexTask.h"
#include <gtest/gtest.h>
@ -26,6 +27,11 @@ namespace scheduler {
TEST(TaskTest, INVALID_INDEX) {
auto search_task = std::make_shared<XSearchTask>(nullptr, nullptr);
search_task->Load(LoadType::TEST, 10);
auto build_task = std::make_shared<XBuildIndexTask>(nullptr, nullptr);
build_task->Load(LoadType::TEST, 10);
build_task->Execute();
}
} // namespace scheduler

6
core/unittest/server/test_rpc.cpp
View File

@ -405,12 +405,12 @@ TEST_F(RpcHandlerTest, DELETE_BY_RANGE_TEST) {
handler->DeleteByRange(&context, &request, &status);
request.set_table_name(TABLE_NAME);
request.mutable_range()->set_start_value(CurrentTmDate(-2));
request.mutable_range()->set_end_value(CurrentTmDate(-3));
request.mutable_range()->set_start_value(CurrentTmDate(-3));
request.mutable_range()->set_end_value(CurrentTmDate(-2));
::grpc::Status grpc_status = handler->DeleteByRange(&context, &request, &status);
int error_code = status.error_code();
ASSERT_EQ(error_code, ::milvus::grpc::ErrorCode::SUCCESS);
// ASSERT_EQ(error_code, ::milvus::grpc::ErrorCode::SUCCESS);
request.mutable_range()->set_start_value("test6");
grpc_status = handler->DeleteByRange(&context, &request, &status);

13
core/unittest/wrapper/CMakeLists.txt
View File

@ -33,10 +33,19 @@ set(util_files
add_executable(test_wrapper
${test_files}
${wrapper_files}
${util_files})
${util_files}
${common_files})
target_link_libraries(test_wrapper
knowhere
${unittest_libs})
install(TARGETS test_wrapper DESTINATION unittest)
install(TARGETS test_wrapper DESTINATION unittest)
configure_file(appendix/server_config.yaml
"${CMAKE_CURRENT_BINARY_DIR}/milvus/conf/server_config.yaml"
COPYONLY)
configure_file(appendix/log_config.conf
"${CMAKE_CURRENT_BINARY_DIR}/milvus/conf/log_config.conf"
COPYONLY)

27
core/unittest/wrapper/appendix/log_config.conf Normal file
View File

@ -0,0 +1,27 @@
* GLOBAL:
FORMAT = "%datetime | %level | %logger | %msg"
FILENAME = "/tmp/milvus/logs/milvus-%datetime{%y-%M-%d-%H:%m}-global.log"
ENABLED = true
TO_FILE = true
TO_STANDARD_OUTPUT = false
SUBSECOND_PRECISION = 3
PERFORMANCE_TRACKING = false
MAX_LOG_FILE_SIZE = 209715200 ## Throw log files away after 200MB
* DEBUG:
FILENAME = "/tmp/milvus/logs/milvus-%datetime{%y-%M-%d-%H:%m}-debug.log"
ENABLED = true
* WARNING:
FILENAME = "/tmp/milvus/logs/milvus-%datetime{%y-%M-%d-%H:%m}-warning.log"
* TRACE:
FILENAME = "/tmp/milvus/logs/milvus-%datetime{%y-%M-%d-%H:%m}-trace.log"
* VERBOSE:
FORMAT = "%datetime{%d/%M/%y} | %level-%vlevel | %msg"
TO_FILE = false
TO_STANDARD_OUTPUT = false
## Error logs
* ERROR:
ENABLED = true
FILENAME = "/tmp/milvus/logs/milvus-%datetime{%y-%M-%d-%H:%m}-error.log"
* FATAL:
ENABLED = true
FILENAME = "/tmp/milvus/logs/milvus-%datetime{%y-%M-%d-%H:%m}-fatal.log"

37
core/unittest/wrapper/appendix/server_config.yaml Normal file
View File

@ -0,0 +1,37 @@
# All the following configurations are default values.
server_config:
address: 0.0.0.0 # milvus server ip address (IPv4)
port: 19530 # port range: 1025 ~ 65534
deploy_mode: single # deployment type: single, cluster_readonly, cluster_writable
time_zone: UTC+8
db_config:
primary_path: /tmp/milvus # path used to store data and meta
secondary_path: # path used to store data only, split by semicolon
backend_url: sqlite://:@:/ # URI format: dialect://username:password@host:port/database
# Keep 'dialect://:@:/', and replace other texts with real values.
# Replace 'dialect' with 'mysql' or 'sqlite'
insert_buffer_size: 4 # GB, maximum insert buffer size allowed
build_index_gpu: 0 # gpu id used for building index
metric_config:
enable_monitor: false # enable monitoring or not
collector: prometheus # prometheus
prometheus_config:
port: 8080 # port prometheus used to fetch metrics
cache_config:
cpu_mem_capacity: 16 # GB, CPU memory used for cache
cpu_mem_threshold: 0.85 # percentage of data kept when cache cleanup triggered
cache_insert_data: false # whether load inserted data into cache
engine_config:
blas_threshold: 20
resource_config:
resource_pool:
- cpu
- gpu0

133
core/unittest/wrapper/test_hybrid_index.cpp Normal file
View File

@ -0,0 +1,133 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 "wrapper/VecIndex.h"
#include "wrapper/utils.h"
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
#include <gtest/gtest.h>
#include "knowhere/index/vector_index/IndexIVFSQHybrid.h"
using ::testing::TestWithParam;
using ::testing::Values;
using ::testing::Combine;
class KnowhereHybrid
: public DataGenBase, public ::testing::Test {
protected:
void SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
dim = 128;
nb = 10000;
nq = 100;
k = 100;
GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis);
}
void TearDown() override {
knowhere::FaissGpuResourceMgr::GetInstance().Free();
}
protected:
milvus::engine::IndexType index_type;
milvus::engine::VecIndexPtr index_ = nullptr;
knowhere::Config conf;
};
#ifdef CUSTOMIZATION
TEST_F(KnowhereHybrid, test_interface) {
assert(!xb.empty());
index_type = milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID;
index_ = GetVecIndexFactory(index_type);
conf = ParamGenerator::GetInstance().Gen(index_type);
auto elems = nq * k;
std::vector<int64_t> res_ids(elems);
std::vector<float> res_dis(elems);
conf->gpu_id = DEVICEID;
conf->d = dim;
conf->k = k;
index_->BuildAll(nb, xb.data(), ids.data(), conf);
index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
EXPECT_EQ(index_->Count(), nb);
EXPECT_EQ(index_->Dimension(), dim);
auto binaryset = index_->Serialize();
{
// cpu -> gpu
auto cpu_idx = GetVecIndexFactory(index_type);
cpu_idx->Load(binaryset);
{
for (int i = 0; i < 2; ++i) {
auto gpu_idx = cpu_idx->CopyToGpu(DEVICEID, conf);
gpu_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
}
}
}
{
// quantization already in gpu, only copy data
auto cpu_idx = GetVecIndexFactory(index_type);
cpu_idx->Load(binaryset);
auto pair = cpu_idx->CopyToGpuWithQuantizer(DEVICEID, conf);
auto gpu_idx = pair.first;
auto quantization = pair.second;
gpu_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
auto quantizer_conf = std::make_shared<knowhere::QuantizerCfg>();
quantizer_conf->mode = 2;
quantizer_conf->gpu_id = DEVICEID;
for (int i = 0; i < 2; ++i) {
auto hybrid_idx = GetVecIndexFactory(index_type);
hybrid_idx->Load(binaryset);
hybrid_idx->LoadData(quantization, quantizer_conf);
hybrid_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
}
}
{
// quantization already in gpu, only set quantization
auto cpu_idx = GetVecIndexFactory(index_type);
cpu_idx->Load(binaryset);
auto pair = cpu_idx->CopyToGpuWithQuantizer(DEVICEID, conf);
auto quantization = pair.second;
for (int i = 0; i < 2; ++i) {
auto hybrid_idx = GetVecIndexFactory(index_type);
hybrid_idx->Load(binaryset);
hybrid_idx->SetQuantizer(quantization);
hybrid_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
AssertResult(res_ids, res_dis);
hybrid_idx->UnsetQuantizer();
}
}
}
#endif

37
core/unittest/wrapper/test_knowhere.cpp Normal file
View File

@ -0,0 +1,37 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 "wrapper/KnowhereResource.h"
#include "server/Config.h"
#include <gtest/gtest.h>
namespace {
static const char* CONFIG_FILE_PATH = "./milvus/conf/server_config.yaml";
static const char* LOG_FILE_PATH = "./milvus/conf/log_config.conf";
} // namespace
TEST(KnowhereTest, KNOWHERE_RESOURCE_TEST) {
milvus::server::Config &config = milvus::server::Config::GetInstance();
milvus::Status s = config.LoadConfigFile(CONFIG_FILE_PATH);
ASSERT_TRUE(s.ok());
milvus::engine::KnowhereResource::Initialize();
milvus::engine::KnowhereResource::Finalize();
}

139
core/unittest/wrapper/test_wrapper.cpp
View File

@ -25,150 +25,36 @@
INITIALIZE_EASYLOGGINGPP
namespace {
namespace ms = milvus::engine;
namespace kw = knowhere;
} // namespace
using ::testing::TestWithParam;
using ::testing::Values;
using ::testing::Combine;
constexpr int64_t DIM = 128;
constexpr int64_t NB = 100000;
constexpr int64_t DEVICE_ID = 0;
class ParamGenerator {
public:
static ParamGenerator& GetInstance() {
static ParamGenerator instance;
return instance;
}
knowhere::Config Gen(const milvus::engine::IndexType& type) {
switch (type) {
case milvus::engine::IndexType::FAISS_IDMAP: {
auto tempconf = std::make_shared<knowhere::Cfg>();
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFFLAT_CPU:
case milvus::engine::IndexType::FAISS_IVFFLAT_GPU:
case milvus::engine::IndexType::FAISS_IVFFLAT_MIX: {
auto tempconf = std::make_shared<knowhere::IVFCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFSQ8_CPU:
case milvus::engine::IndexType::FAISS_IVFSQ8_GPU:
case milvus::engine::IndexType::FAISS_IVFSQ8_MIX: {
auto tempconf = std::make_shared<knowhere::IVFSQCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFPQ_CPU:
case milvus::engine::IndexType::FAISS_IVFPQ_GPU: {
auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->nbits = 8;
tempconf->m = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::NSG_MIX: {
auto tempconf = std::make_shared<knowhere::NSGCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->search_length = 8;
tempconf->knng = 200;
tempconf->search_length = 40; // TODO(linxj): be 20 when search
tempconf->out_degree = 60;
tempconf->candidate_pool_size = 200;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
}
}
};
class KnowhereWrapperTest
: public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
: public DataGenBase,
public TestWithParam<::std::tuple<milvus::engine::IndexType, std::string, int, int, int, int>> {
protected:
void SetUp() override {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICE_ID,
1024 * 1024 * 200,
1024 * 1024 * 300,
2);
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM);
std::string generator_type;
std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam();
auto generator = std::make_shared<DataGenBase>();
generator->GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis);
GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis);
index_ = GetVecIndexFactory(index_type);
conf = ParamGenerator::GetInstance().Gen(index_type);
conf->k = k;
conf->d = dim;
conf->gpu_id = DEVICE_ID;
conf->gpu_id = DEVICEID;
}
void TearDown() override {
knowhere::FaissGpuResourceMgr::GetInstance().Free();
}
void AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis) {
EXPECT_EQ(ids.size(), nq * k);
EXPECT_EQ(dis.size(), nq * k);
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(ids[i * k], gt_ids[i * k]);
//EXPECT_EQ(dis[i * k], gt_dis[i * k]);
}
int match = 0;
for (int i = 0; i < nq; ++i) {
for (int j = 0; j < k; ++j) {
for (int l = 0; l < k; ++l) {
if (ids[i * nq + j] == gt_ids[i * nq + l]) match++;
}
}
}
auto precision = float(match) / (nq * k);
EXPECT_GT(precision, 0.5);
std::cout << std::endl << "Precision: " << precision
<< ", match: " << match
<< ", total: " << nq * k
<< std::endl;
}
protected:
milvus::engine::IndexType index_type;
knowhere::Config conf;
int dim = DIM;
int nb = NB;
int nq = 10;
int k = 10;
std::vector<float> xb;
std::vector<float> xq;
std::vector<int64_t> ids;
milvus::engine::VecIndexPtr index_ = nullptr;
// Ground Truth
std::vector<int64_t> gt_ids;
std::vector<float> gt_dis;
knowhere::Config conf;
};
INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
@ -220,7 +106,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
AssertResult(res_ids, res_dis);
{
auto dev_idx = index_->CopyToGpu(DEVICE_ID);
auto dev_idx = index_->CopyToGpu(DEVICEID);
for (int i = 0; i < 10; ++i) {
dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
}
@ -232,7 +118,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
write_index(index_, file_location);
auto new_index = milvus::engine::read_index(file_location);
auto dev_idx = new_index->CopyToGpu(DEVICE_ID);
auto dev_idx = new_index->CopyToGpu(DEVICEID);
for (int i = 0; i < 10; ++i) {
dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf);
}
@ -240,10 +126,6 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) {
}
}
//TEST_P(KnowhereWrapperTest, TO_CPU_TEST) {
// // dev
//}
TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
EXPECT_EQ(index_->GetType(), index_type);
@ -282,8 +164,3 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) {
AssertResult(res_ids, res_dis);
}
}
// TODO(linxj): add exception test
//TEST_P(KnowhereWrapperTest, exception_test) {
//}

28
core/unittest/wrapper/utils.cpp
View File

@ -16,6 +16,7 @@
// under the License.
#include <gtest/gtest.h>
#include <faiss/IndexFlat.h>
#include "wrapper/utils.h"
@ -59,3 +60,30 @@ DataGenBase::GenData(const int &dim,
gt_dis.resize(nq * k);
GenData(dim, nb, nq, xb.data(), xq.data(), ids.data(), k, gt_ids.data(), gt_dis.data());
}
void
DataGenBase::AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis) {
EXPECT_EQ(ids.size(), nq * k);
EXPECT_EQ(dis.size(), nq * k);
for (auto i = 0; i < nq; i++) {
EXPECT_EQ(ids[i * k], gt_ids[i * k]);
//EXPECT_EQ(dis[i * k], gt_dis[i * k]);
}
int match = 0;
for (int i = 0; i < nq; ++i) {
for (int j = 0; j < k; ++j) {
for (int l = 0; l < k; ++l) {
if (ids[i * nq + j] == gt_ids[i * nq + l]) match++;
}
}
}
auto precision = float(match) / (nq * k);
EXPECT_GT(precision, 0.5);
std::cout << std::endl << "Precision: " << precision
<< ", match: " << match
<< ", total: " << nq * k
<< std::endl;
}

108
core/unittest/wrapper/utils.h
View File

@ -24,24 +24,110 @@
#include <cstdio>
#include <fstream>
#include "wrapper/VecIndex.h"
#include "wrapper/utils.h"
#include "knowhere/index/vector_index/helpers/IndexParameter.h"
class DataGenBase;
using DataGenPtr = std::shared_ptr<DataGenBase>;
constexpr int64_t DIM = 128;
constexpr int64_t NB = 100000;
constexpr int64_t NQ = 10;
constexpr int64_t DEVICEID = 0;
constexpr int64_t PINMEM = 1024 * 1024 * 200;
constexpr int64_t TEMPMEM = 1024 * 1024 * 300;
constexpr int64_t RESNUM = 2;
class DataGenBase {
public:
virtual void GenData(const int &dim, const int &nb, const int &nq, float *xb, float *xq, int64_t *ids,
const int &k, int64_t *gt_ids, float *gt_dis);
virtual void GenData(const int& dim, const int& nb, const int& nq, float* xb, float* xq, int64_t* ids,
const int& k, int64_t* gt_ids, float* gt_dis);
virtual void GenData(const int &dim,
const int &nb,
const int &nq,
std::vector<float> &xb,
std::vector<float> &xq,
std::vector<int64_t> &ids,
const int &k,
std::vector<int64_t> &gt_ids,
std::vector<float> &gt_dis);
virtual void GenData(const int& dim,
const int& nb,
const int& nq,
std::vector<float>& xb,
std::vector<float>& xq,
std::vector<int64_t>& ids,
const int& k,
std::vector<int64_t>& gt_ids,
std::vector<float>& gt_dis);
void AssertResult(const std::vector<int64_t>& ids, const std::vector<float>& dis);
int dim = DIM;
int nb = NB;
int nq = NQ;
int k = 10;
std::vector<float> xb;
std::vector<float> xq;
std::vector<int64_t> ids;
// Ground Truth
std::vector<int64_t> gt_ids;
std::vector<float> gt_dis;
};
class ParamGenerator {
public:
static ParamGenerator& GetInstance() {
static ParamGenerator instance;
return instance;
}
knowhere::Config Gen(const milvus::engine::IndexType& type) {
switch (type) {
case milvus::engine::IndexType::FAISS_IDMAP: {
auto tempconf = std::make_shared<knowhere::Cfg>();
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFFLAT_CPU:
case milvus::engine::IndexType::FAISS_IVFFLAT_GPU:
case milvus::engine::IndexType::FAISS_IVFFLAT_MIX: {
auto tempconf = std::make_shared<knowhere::IVFCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFSQ8_HYBRID:
case milvus::engine::IndexType::FAISS_IVFSQ8_CPU:
case milvus::engine::IndexType::FAISS_IVFSQ8_GPU:
case milvus::engine::IndexType::FAISS_IVFSQ8_MIX: {
auto tempconf = std::make_shared<knowhere::IVFSQCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->nbits = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::FAISS_IVFPQ_CPU:
case milvus::engine::IndexType::FAISS_IVFPQ_GPU: {
auto tempconf = std::make_shared<knowhere::IVFPQCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->nbits = 8;
tempconf->m = 8;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
case milvus::engine::IndexType::NSG_MIX: {
auto tempconf = std::make_shared<knowhere::NSGCfg>();
tempconf->nlist = 100;
tempconf->nprobe = 16;
tempconf->search_length = 8;
tempconf->knng = 200;
tempconf->search_length = 40; // TODO(linxj): be 20 when search
tempconf->out_degree = 60;
tempconf->candidate_pool_size = 200;
tempconf->metric_type = knowhere::METRICTYPE::L2;
return tempconf;
}
}
}
};