diff --git a/CHANGELOG.md b/CHANGELOG.md index 76c8a13c45a7387bda077dd5e76cfbcea066a58f..f5ed05522df01d4949c648338e29b7ad2e5466b1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,13 +8,16 @@ Please mark all change in change log and use the ticket from JIRA. - \#246 - Exclude src/external folder from code coverage for jenkin ci - \#248 - Reside src/external in thirdparty - \#316 - Some files not merged after vectors added +- \#327 - Search does not use GPU when index type is FLAT - \#340 - Test cases run failed on 0.6.0 - \#353 - Rename config.h.in to version.h.in +- \#374 - sdk_simple return empty result ## Feature - \#12 - Pure CPU version for Milvus - \#77 - Support table partition - \#226 - Experimental shards middleware for Milvus +- \#127 - Support new Index type IVFPQ ## Improvement - \#275 - Rename C++ SDK IndexType @@ -22,7 +25,7 @@ Please mark all change in change log and use the ticket from JIRA. - \#260 - C++ SDK README - \#314 - add Find FAISS in CMake - \#310 - Add Q&A for 'protocol https not supported or disable in libcurl' issue -- \#322 - Add option to enable / disable prometheus +- \#322 - Add option to enable / disable prometheus - \#358 - Add more information in build.sh and install.md ## Task diff --git a/core/src/db/engine/ExecutionEngine.h b/core/src/db/engine/ExecutionEngine.h index 51c77eb78ef13ba7624711adc528ffb1dc792e5c..86a014cf66ce2058b331372a0c41bd84acd2a766 100644 --- a/core/src/db/engine/ExecutionEngine.h +++ b/core/src/db/engine/ExecutionEngine.h @@ -26,6 +26,7 @@ namespace milvus { namespace engine { +// TODO(linxj): replace with VecIndex::IndexType enum class EngineType { INVALID = 0, FAISS_IDMAP = 1, @@ -33,7 +34,8 @@ enum class EngineType { FAISS_IVFSQ8, NSG_MIX, FAISS_IVFSQ8H, - MAX_VALUE = FAISS_IVFSQ8H, + FAISS_PQ, + MAX_VALUE = FAISS_PQ, }; enum class MetricType { diff --git a/core/src/db/engine/ExecutionEngineImpl.cpp b/core/src/db/engine/ExecutionEngineImpl.cpp index 19c699bda7ee9bab8f60c7eb9e7c3ba5cb8507b3..3a2b940cec01f20a1eb16ffc6b474b18da4c813d 100644 --- a/core/src/db/engine/ExecutionEngineImpl.cpp +++ b/core/src/db/engine/ExecutionEngineImpl.cpp @@ -116,6 +116,14 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) { index = GetVecIndexFactory(IndexType::FAISS_IVFSQ8_HYBRID); break; } + case EngineType::FAISS_PQ: { +#ifdef MILVUS_CPU_VERSION + index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_CPU); +#else + index = GetVecIndexFactory(IndexType::FAISS_IVFPQ_MIX); +#endif + break; + } default: { ENGINE_LOG_ERROR << "Unsupported index type"; return nullptr; diff --git a/core/src/grpc/gen-milvus/milvus.pb.h b/core/src/grpc/gen-milvus/milvus.pb.h index 53ed2db22ea41448717e9f660768a6dee33da74e..536365b23b88e954cfd7146ed88d55dab0ed3072 100644 --- a/core/src/grpc/gen-milvus/milvus.pb.h +++ b/core/src/grpc/gen-milvus/milvus.pb.h @@ -1321,7 +1321,8 @@ class RowRecord : void clear_vector_data(); float vector_data(int index) const; void set_vector_data(int index, float value); - void add_vector_data(float value); +// void add_vector_data(float value); + void add_vector_data(std::vector::const_iterator begin, std::vector::const_iterator end); const ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >& vector_data() const; ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >* @@ -1473,7 +1474,9 @@ class InsertParam : void clear_row_id_array(); ::PROTOBUF_NAMESPACE_ID::int64 row_id_array(int index) const; void set_row_id_array(int index, ::PROTOBUF_NAMESPACE_ID::int64 value); - void add_row_id_array(::PROTOBUF_NAMESPACE_ID::int64 value); +// void add_row_id_array(::PROTOBUF_NAMESPACE_ID::int64 value); + void add_row_id_array(std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator begin, + std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator end); const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >& row_id_array() const; ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >* @@ -2136,7 +2139,9 @@ class TopKQueryResult : void clear_ids(); ::PROTOBUF_NAMESPACE_ID::int64 ids(int index) const; void set_ids(int index, ::PROTOBUF_NAMESPACE_ID::int64 value); - void add_ids(::PROTOBUF_NAMESPACE_ID::int64 value); +// void add_ids(::PROTOBUF_NAMESPACE_ID::int64 value); + void add_ids(std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator begin, + std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator end); const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >& ids() const; ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >* @@ -2147,7 +2152,8 @@ class TopKQueryResult : void clear_distances(); float distances(int index) const; void set_distances(int index, float value); - void add_distances(float value); +// void add_distances(float value); + void add_distances(std::vector::const_iterator begin, std::vector::const_iterator end); const ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >& distances() const; ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >* @@ -3922,9 +3928,14 @@ inline void RowRecord::set_vector_data(int index, float value) { vector_data_.Set(index, value); // @@protoc_insertion_point(field_set:milvus.grpc.RowRecord.vector_data) } -inline void RowRecord::add_vector_data(float value) { - vector_data_.Add(value); - // @@protoc_insertion_point(field_add:milvus.grpc.RowRecord.vector_data) +//inline void RowRecord::add_vector_data(float value) { +// vector_data_.Add(value); +// // @@protoc_insertion_point(field_add:milvus.grpc.RowRecord.vector_data) +//} +inline void RowRecord::add_vector_data(std::vector::const_iterator begin, + std::vector::const_iterator end) { + vector_data_.Add(begin, end); + // @@protoc_insertion_point(field_add:milvus.grpc.RowRecord.vector_data) } inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >& RowRecord::vector_data() const { @@ -4037,9 +4048,14 @@ inline void InsertParam::set_row_id_array(int index, ::PROTOBUF_NAMESPACE_ID::in row_id_array_.Set(index, value); // @@protoc_insertion_point(field_set:milvus.grpc.InsertParam.row_id_array) } -inline void InsertParam::add_row_id_array(::PROTOBUF_NAMESPACE_ID::int64 value) { - row_id_array_.Add(value); - // @@protoc_insertion_point(field_add:milvus.grpc.InsertParam.row_id_array) +//inline void InsertParam::add_row_id_array(::PROTOBUF_NAMESPACE_ID::int64 value) { +// row_id_array_.Add(value); +// // @@protoc_insertion_point(field_add:milvus.grpc.InsertParam.row_id_array) +//} +inline void InsertParam::add_row_id_array(std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator begin, + std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator end) { + row_id_array_.Add(begin, end); + // @@protoc_insertion_point(field_add:milvus.grpc.InsertParam.row_id_array) } inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >& InsertParam::row_id_array() const { @@ -4588,9 +4604,14 @@ inline void TopKQueryResult::set_ids(int index, ::PROTOBUF_NAMESPACE_ID::int64 v ids_.Set(index, value); // @@protoc_insertion_point(field_set:milvus.grpc.TopKQueryResult.ids) } -inline void TopKQueryResult::add_ids(::PROTOBUF_NAMESPACE_ID::int64 value) { - ids_.Add(value); - // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.ids) +//inline void TopKQueryResult::add_ids(::PROTOBUF_NAMESPACE_ID::int64 value) { +// ids_.Add(value); +// // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.ids) +//} +inline void TopKQueryResult::add_ids(std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator begin, + std::vector<::PROTOBUF_NAMESPACE_ID::int64>::const_iterator end) { + ids_.Add(begin,end); + // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.ids) } inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< ::PROTOBUF_NAMESPACE_ID::int64 >& TopKQueryResult::ids() const { @@ -4618,9 +4639,13 @@ inline void TopKQueryResult::set_distances(int index, float value) { distances_.Set(index, value); // @@protoc_insertion_point(field_set:milvus.grpc.TopKQueryResult.distances) } -inline void TopKQueryResult::add_distances(float value) { - distances_.Add(value); - // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.distances) +//inline void TopKQueryResult::add_distances(float value) { +// distances_.Add(value); +// // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.distances) +//} +inline void TopKQueryResult::add_distances(std::vector::const_iterator begin, std::vector::const_iterator end) { + distances_.Add(begin, end); + // @@protoc_insertion_point(field_add:milvus.grpc.TopKQueryResult.distances) } inline const ::PROTOBUF_NAMESPACE_ID::RepeatedField< float >& TopKQueryResult::distances() const { diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp index b027539c37dd485a1af9f7841ba5092618b828ea..cbd4f4f09cf8c06544879f7a227e305cbf8ef963 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.cpp @@ -39,17 +39,19 @@ GPUIVFPQ::Train(const DatasetPtr& dataset, const Config& config) { GETTENSOR(dataset) - // TODO(linxj): set device here. - // TODO(linxj): set gpu resource here. - faiss::gpu::StandardGpuResources res; - faiss::gpu::GpuIndexIVFPQ device_index(&res, dim, build_cfg->nlist, build_cfg->m, build_cfg->nbits, - GetMetricType(build_cfg->metric_type)); // IP not support - device_index.train(rows, (float*)p_data); - - std::shared_ptr host_index = nullptr; - host_index.reset(faiss::gpu::index_gpu_to_cpu(&device_index)); - - return std::make_shared(host_index); + auto temp_resource = FaissGpuResourceMgr::GetInstance().GetRes(gpu_id_); + if (temp_resource != nullptr) { + ResScope rs(temp_resource, gpu_id_, true); + auto device_index = new faiss::gpu::GpuIndexIVFPQ(temp_resource->faiss_res.get(), dim, build_cfg->nlist, + build_cfg->m, build_cfg->nbits, + GetMetricType(build_cfg->metric_type)); // IP not support + device_index->train(rows, (float*)p_data); + std::shared_ptr host_index = nullptr; + host_index.reset(faiss::gpu::index_gpu_to_cpu(device_index)); + return std::make_shared(host_index); + } else { + KNOWHERE_THROW_MSG("Build IVFSQ can't get gpu resource"); + } } std::shared_ptr @@ -66,7 +68,14 @@ GPUIVFPQ::GenParams(const Config& config) { VectorIndexPtr GPUIVFPQ::CopyGpuToCpu(const Config& config) { - KNOWHERE_THROW_MSG("not support yet"); + std::lock_guard lk(mutex_); + + faiss::Index* device_index = index_.get(); + faiss::Index* host_index = faiss::gpu::index_gpu_to_cpu(device_index); + + std::shared_ptr new_index; + new_index.reset(host_index); + return std::make_shared(new_index); } } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h index 13ea1075ca5740bc6094d0ec2e59bcdd501f3d91..f699445d65acc9e7bee5e4e116cfbaa809851cfa 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexGPUIVFPQ.h @@ -18,6 +18,7 @@ #pragma once #include +#include #include "IndexGPUIVF.h" @@ -28,6 +29,10 @@ class GPUIVFPQ : public GPUIVF { explicit GPUIVFPQ(const int& device_id) : GPUIVF(device_id) { } + GPUIVFPQ(std::shared_ptr index, const int64_t& device_id, ResPtr& resource) + : GPUIVF(std::move(index), device_id, resource) { + } + IndexModelPtr Train(const DatasetPtr& dataset, const Config& config) override; diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp index 03acbf31d774730b9514bd6b46d93866f384b5da..841ef6324540ca9f592ad0b252d65f7b4fd9649d 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp @@ -17,11 +17,19 @@ #include #include +#ifdef MILVUS_GPU_VERSION +#include +#endif + #include #include #include "knowhere/adapter/VectorAdapter.h" #include "knowhere/common/Exception.h" +#ifdef MILVUS_GPU_VERSION +#include "knowhere/index/vector_index/IndexGPUIVF.h" +#include "knowhere/index/vector_index/IndexGPUIVFPQ.h" +#endif #include "knowhere/index/vector_index/IndexIVFPQ.h" namespace knowhere { @@ -60,4 +68,22 @@ IVFPQ::Clone_impl(const std::shared_ptr& index) { return std::make_shared(index); } +VectorIndexPtr +IVFPQ::CopyCpuToGpu(const int64_t& device_id, const Config& config) { +#ifdef MILVUS_GPU_VERSION + if (auto res = FaissGpuResourceMgr::GetInstance().GetRes(device_id)) { + ResScope rs(res, device_id, false); + auto gpu_index = faiss::gpu::index_cpu_to_gpu(res->faiss_res.get(), device_id, index_.get()); + + std::shared_ptr device_index; + device_index.reset(gpu_index); + return std::make_shared(device_index, device_id, res); + } else { + KNOWHERE_THROW_MSG("CopyCpuToGpu Error, can't get gpu_resource"); + } +#else + KNOWHERE_THROW_MSG("Calling IVFPQ::CopyCpuToGpu when we are using CPU version"); +#endif +} + } // namespace knowhere diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h index 69aaa5090bd46a5cc115f8f7bd177da4135bbcdd..fc50c68389000d7ff22d2d35ee4778fe089196e5 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h +++ b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.h @@ -34,6 +34,9 @@ class IVFPQ : public IVF { IndexModelPtr Train(const DatasetPtr& dataset, const Config& config) override; + VectorIndexPtr + CopyCpuToGpu(const int64_t& device_id, const Config& config) override; + protected: std::shared_ptr GenParams(const Config& config) override; diff --git a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp index 5ff2bfc2e323ebb669b47d57305342c519af5691..363e629eca5c655ac2f4d144713c8b4c504c831a 100644 --- a/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp +++ b/core/src/index/knowhere/knowhere/index/vector_index/helpers/Cloner.cpp @@ -49,7 +49,7 @@ CopyCpuToGpu(const VectorIndexPtr& index, const int64_t& device_id, const Config if (auto cpu_index = std::dynamic_pointer_cast(index)) { return cpu_index->CopyCpuToGpu(device_id, config); } else if (auto cpu_index = std::dynamic_pointer_cast(index)) { - KNOWHERE_THROW_MSG("IVFPQ not support transfer to gpu"); + return cpu_index->CopyCpuToGpu(device_id, config); } else if (auto cpu_index = std::dynamic_pointer_cast(index)) { return cpu_index->CopyCpuToGpu(device_id, config); } else if (auto cpu_index = std::dynamic_pointer_cast(index)) { diff --git a/core/src/index/unittest/test_ivf.cpp b/core/src/index/unittest/test_ivf.cpp index 5cb820df958fd826910b6c3dca40b5cd1ef43822..79ef1aa170f0a56e72e6f09bf9149e9e604a42a2 100644 --- a/core/src/index/unittest/test_ivf.cpp +++ b/core/src/index/unittest/test_ivf.cpp @@ -213,7 +213,7 @@ TEST_P(IVFTest, clone_test) { { // copy from gpu to cpu - std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ", "IVFSQHybrid"}; + std::vector support_idx_vec{"GPUIVF", "GPUIVFSQ", "GPUIVFPQ", "IVFSQHybrid"}; auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); if (finder != support_idx_vec.cend()) { EXPECT_NO_THROW({ @@ -238,7 +238,7 @@ TEST_P(IVFTest, clone_test) { { // copy to gpu - std::vector support_idx_vec{"IVF", "GPUIVF", "IVFSQ", "GPUIVFSQ"}; + std::vector support_idx_vec{"IVF", "GPUIVF", "IVFSQ", "GPUIVFSQ", "IVFPQ", "GPUIVFPQ"}; auto finder = std::find(support_idx_vec.cbegin(), support_idx_vec.cend(), index_type); if (finder != support_idx_vec.cend()) { EXPECT_NO_THROW({ diff --git a/core/src/scheduler/interface/interfaces.h b/core/src/scheduler/interface/interfaces.h index 9920e4f80af69133319a05e95c0aeb36f5b51d53..311560adc61e8ccc92506f09eb9b299223519628 100644 --- a/core/src/scheduler/interface/interfaces.h +++ b/core/src/scheduler/interface/interfaces.h @@ -33,8 +33,7 @@ namespace milvus { namespace interface { struct dumpable { - virtual ~dumpable() { - } + virtual ~dumpable() = default; virtual json Dump() const = 0; diff --git a/core/src/scheduler/optimizer/OnlyGPUPass.cpp b/core/src/scheduler/optimizer/OnlyGPUPass.cpp index 3fcda0e8a347bc6ca62890c78b394314f426c9e1..2a72f9757e9b10b4b08d09597d5454af2ca00cee 100644 --- a/core/src/scheduler/optimizer/OnlyGPUPass.cpp +++ b/core/src/scheduler/optimizer/OnlyGPUPass.cpp @@ -34,7 +34,8 @@ OnlyGPUPass::Run(const TaskPtr& task) { auto search_task = std::static_pointer_cast(task); if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFSQ8 && - search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFFLAT) { + search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFFLAT && + search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IDMAP) { return false; } diff --git a/core/src/sdk/grpc/ClientProxy.cpp b/core/src/sdk/grpc/ClientProxy.cpp index 4a84323d649cf9ed9d5abc24a690e7b5a5be3a8d..131d0bdc9943837814b3e6c7be8ee460004bb47f 100644 --- a/core/src/sdk/grpc/ClientProxy.cpp +++ b/core/src/sdk/grpc/ClientProxy.cpp @@ -32,13 +32,6 @@ UriCheck(const std::string& uri) { return (index != std::string::npos); } -void -CopyRowRecord(::milvus::grpc::RowRecord* target, const RowRecord& src) { - auto vector_data = target->mutable_vector_data(); - vector_data->Resize(static_cast(src.data.size()), 0.0); - memcpy(vector_data->mutable_data(), src.data.data(), src.data.size() * sizeof(float)); -} - Status ClientProxy::Connect(const ConnectParam& param) { std::string uri = param.ip_address + ":" + param.port; @@ -196,17 +189,14 @@ ClientProxy::Insert(const std::string& table_name, const std::string& partition_ for (auto& record : record_array) { ::milvus::grpc::RowRecord* grpc_record = insert_param.add_row_record_array(); - CopyRowRecord(grpc_record, record); + grpc_record->add_vector_data(record.data.begin(), record.data.end()); } // Single thread ::milvus::grpc::VectorIds vector_ids; if (!id_array.empty()) { /* set user's ids */ - auto row_ids = insert_param.mutable_row_id_array(); - row_ids->Reserve(static_cast(id_array.size())); - memcpy(row_ids->mutable_data(), id_array.data(), id_array.size() * sizeof(int64_t)); - + insert_param.add_row_id_array(id_array.begin(), id_array.end()); client_ptr_->Insert(vector_ids, insert_param, status); } else { client_ptr_->Insert(vector_ids, insert_param, status); @@ -236,7 +226,7 @@ ClientProxy::Search(const std::string& table_name, const std::vectoradd_vector_data(record.data.begin(), record.data.end()); } // step 2: convert range array diff --git a/core/src/server/grpc_impl/GrpcRequestTask.cpp b/core/src/server/grpc_impl/GrpcRequestTask.cpp index 0d685bed4bc3c6bfdd5f046c0f62f6d595064059..4c6178efeea9b899e3367c9b2b3b240deb07346c 100644 --- a/core/src/server/grpc_impl/GrpcRequestTask.cpp +++ b/core/src/server/grpc_impl/GrpcRequestTask.cpp @@ -682,11 +682,8 @@ SearchTask::OnExecute() { // step 7: construct result array topk_result_->set_row_num(record_count); - topk_result_->mutable_ids()->Resize(static_cast(result_ids.size()), 0); - memcpy(topk_result_->mutable_ids()->mutable_data(), result_ids.data(), result_ids.size() * sizeof(int64_t)); - topk_result_->mutable_distances()->Resize(static_cast(result_distances.size()), 0.0); - memcpy(topk_result_->mutable_distances()->mutable_data(), result_distances.data(), - result_distances.size() * sizeof(float)); + topk_result_->add_ids(result_ids.begin(), result_ids.end()); + topk_result_->add_distances(result_distances.begin(), result_distances.end()); // step 8: print time cost percent rc.RecordSection("construct result and send"); diff --git a/core/src/wrapper/ConfAdapter.cpp b/core/src/wrapper/ConfAdapter.cpp index 2dcf6bab7e1a0c9ccb354a5bcd46ebce9dffca27..461745f1fdbaeb5386eb5fa89b4b8c0e942d64bf 100644 --- a/core/src/wrapper/ConfAdapter.cpp +++ b/core/src/wrapper/ConfAdapter.cpp @@ -16,6 +16,7 @@ // under the License. #include "wrapper/ConfAdapter.h" +#include "WrapperException.h" #include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "utils/Log.h" @@ -76,7 +77,7 @@ IVFConfAdapter::MatchNlist(const int64_t& size, const int64_t& nlist) { if (size <= TYPICAL_COUNT / 16384 + 1) { // handle less row count, avoid nlist set to 0 return 1; - } else if (int(size / TYPICAL_COUNT) * nlist == 0) { + } else if (int(size / TYPICAL_COUNT) * nlist <= 0) { // calculate a proper nlist if nlist not specified or size less than TYPICAL_COUNT return int(size / TYPICAL_COUNT * 16384); } @@ -87,7 +88,11 @@ knowhere::Config IVFConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) { auto conf = std::make_shared(); conf->k = metaconf.k; - conf->nprobe = metaconf.nprobe; + + if (metaconf.nprobe <= 0) + conf->nprobe = 16; // hardcode here + else + conf->nprobe = metaconf.nprobe; switch (type) { case IndexType::FAISS_IVFFLAT_GPU: @@ -123,11 +128,47 @@ IVFPQConfAdapter::Match(const TempMetaConf& metaconf) { conf->metric_type = metaconf.metric_type; conf->gpu_id = conf->gpu_id; conf->nbits = 8; - conf->m = 8; + + if (!(conf->d % 4)) + conf->m = conf->d / 4; // compression radio = 16 + else if (!(conf->d % 2)) + conf->m = conf->d / 2; // compression radio = 8 + else if (!(conf->d % 3)) + conf->m = conf->d / 3; // compression radio = 12 + else + conf->m = conf->d; // same as SQ8, compression radio = 4 + MatchBase(conf); return conf; } +knowhere::Config +IVFPQConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) { + auto conf = std::make_shared(); + conf->k = metaconf.k; + + if (metaconf.nprobe <= 0) { + WRAPPER_LOG_ERROR << "The nprobe of PQ is wrong!"; + throw WrapperException("The nprobe of PQ is wrong!"); + } else { + conf->nprobe = metaconf.nprobe; + } + + return conf; +} + +int64_t +IVFPQConfAdapter::MatchNlist(const int64_t& size, const int64_t& nlist) { + if (size <= TYPICAL_COUNT / 16384 + 1) { + // handle less row count, avoid nlist set to 0 + return 1; + } else if (int(size / TYPICAL_COUNT) * nlist <= 0) { + // calculate a proper nlist if nlist not specified or size less than TYPICAL_COUNT + return int(size / TYPICAL_COUNT * 16384); + } + return nlist; +} + knowhere::Config NSGConfAdapter::Match(const TempMetaConf& metaconf) { auto conf = std::make_shared(); @@ -136,13 +177,14 @@ NSGConfAdapter::Match(const TempMetaConf& metaconf) { conf->metric_type = metaconf.metric_type; conf->gpu_id = conf->gpu_id; + double factor = metaconf.size / TYPICAL_COUNT; auto scale_factor = round(metaconf.dim / 128.0); scale_factor = scale_factor >= 4 ? 4 : scale_factor; - conf->nprobe = 6 + 10 * scale_factor; - conf->knng = 100 + 100 * scale_factor; - conf->search_length = 40 + 5 * scale_factor; - conf->out_degree = 50 + 5 * scale_factor; - conf->candidate_pool_size = 200 + 100 * scale_factor; + conf->nprobe = conf->nlist > 10000 ? conf->nlist * 0.02 : conf->nlist * 0.1; + conf->knng = (100 + 100 * scale_factor) * factor; + conf->search_length = (40 + 5 * scale_factor) * factor; + conf->out_degree = (50 + 5 * scale_factor) * factor; + conf->candidate_pool_size = (200 + 100 * scale_factor) * factor; MatchBase(conf); // WRAPPER_LOG_DEBUG << "nlist: " << conf->nlist @@ -156,6 +198,9 @@ NSGConfAdapter::MatchSearch(const TempMetaConf& metaconf, const IndexType& type) auto conf = std::make_shared(); conf->k = metaconf.k; conf->search_length = metaconf.search_length; + if (metaconf.search_length == TEMPMETA_DEFAULT_VALUE) { + conf->search_length = 30; // TODO(linxj): hardcode here. + } return conf; } diff --git a/core/src/wrapper/ConfAdapter.h b/core/src/wrapper/ConfAdapter.h index 4c8e528a2d63c8e9898ef0c1c906732e30e21e2f..5ec3d524860e95dfdc0770442acda1ac6bde860d 100644 --- a/core/src/wrapper/ConfAdapter.h +++ b/core/src/wrapper/ConfAdapter.h @@ -79,6 +79,13 @@ class IVFPQConfAdapter : public IVFConfAdapter { public: knowhere::Config Match(const TempMetaConf& metaconf) override; + + knowhere::Config + MatchSearch(const TempMetaConf& metaconf, const IndexType& type) override; + + protected: + static int64_t + MatchNlist(const int64_t& size, const int64_t& nlist); }; class NSGConfAdapter : public IVFConfAdapter { diff --git a/core/src/wrapper/ConfAdapterMgr.cpp b/core/src/wrapper/ConfAdapterMgr.cpp index b329588cf2f87aa9eb31f6e6a4bb5f2098ded8f7..d0eba0452934200c2e34298e326c6bb5849287e6 100644 --- a/core/src/wrapper/ConfAdapterMgr.cpp +++ b/core/src/wrapper/ConfAdapterMgr.cpp @@ -53,6 +53,7 @@ AdapterMgr::RegisterAdapter() { REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_CPU, ivfpq_cpu); REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_GPU, ivfpq_gpu); + REGISTER_CONF_ADAPTER(IVFPQConfAdapter, IndexType::FAISS_IVFPQ_MIX, ivfpq_mix); REGISTER_CONF_ADAPTER(NSGConfAdapter, IndexType::NSG_MIX, nsg_mix); } diff --git a/core/src/wrapper/VecIndex.cpp b/core/src/wrapper/VecIndex.cpp index 2323783c3d123adde72f427e06eee6f917e90062..81a13c60f90a4709d020f5c3d4bd664ffebb641e 100644 --- a/core/src/wrapper/VecIndex.cpp +++ b/core/src/wrapper/VecIndex.cpp @@ -145,6 +145,10 @@ GetVecIndexFactory(const IndexType& type, const Config& cfg) { index = std::make_shared(gpu_device); break; } + case IndexType::FAISS_IVFPQ_MIX: { + index = std::make_shared(gpu_device); + return std::make_shared(index, IndexType::FAISS_IVFPQ_MIX); + } case IndexType::FAISS_IVFSQ8_MIX: { index = std::make_shared(gpu_device); return std::make_shared(index, IndexType::FAISS_IVFSQ8_MIX); @@ -276,6 +280,10 @@ ConvertToCpuIndexType(const IndexType& type) { case IndexType::FAISS_IVFSQ8_MIX: { return IndexType::FAISS_IVFSQ8_CPU; } + case IndexType::FAISS_IVFPQ_GPU: + case IndexType::FAISS_IVFPQ_MIX: { + return IndexType::FAISS_IVFPQ_CPU; + } default: { return type; } } } @@ -291,9 +299,12 @@ ConvertToGpuIndexType(const IndexType& type) { case IndexType::FAISS_IVFSQ8_CPU: { return IndexType::FAISS_IVFSQ8_GPU; } + case IndexType::FAISS_IVFPQ_MIX: + case IndexType::FAISS_IVFPQ_CPU: { + return IndexType::FAISS_IVFPQ_GPU; + } default: { return type; } } } - } // namespace engine } // namespace milvus diff --git a/core/src/wrapper/VecIndex.h b/core/src/wrapper/VecIndex.h index cd9ba86951937e91dbfeaac4ac389c6ccb1b1e5e..efe01a25d7b40b5209649e18bffe5102e4cc917c 100644 --- a/core/src/wrapper/VecIndex.h +++ b/core/src/wrapper/VecIndex.h @@ -33,6 +33,7 @@ namespace engine { using Config = knowhere::Config; +// TODO(linxj): replace with string, Do refactor serialization enum class IndexType { INVALID = 0, FAISS_IDMAP = 1, @@ -47,6 +48,7 @@ enum class IndexType { FAISS_IVFSQ8_GPU, FAISS_IVFSQ8_HYBRID, // only support build on gpu. NSG_MIX, + FAISS_IVFPQ_MIX, }; class VecIndex; diff --git a/core/unittest/wrapper/test_wrapper.cpp b/core/unittest/wrapper/test_wrapper.cpp index ddfdfa2fad47bf80ba87e5aa8bd64b2aba0dc664..025601a1cd65d6d90cc8cb8aa92361d6e41bacef 100644 --- a/core/unittest/wrapper/test_wrapper.cpp +++ b/core/unittest/wrapper/test_wrapper.cpp @@ -29,33 +29,40 @@ INITIALIZE_EASYLOGGINGPP -using ::testing::Combine; using ::testing::TestWithParam; using ::testing::Values; +using ::testing::Combine; class KnowhereWrapperTest - : public DataGenBase, - public TestWithParam<::std::tuple> { + : public DataGenBase, + public TestWithParam<::std::tuple> { protected: - void - SetUp() override { + void SetUp() override { #ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(DEVICEID, PINMEM, TEMPMEM, RESNUM); #endif - std::string generator_type; std::tie(index_type, generator_type, dim, nb, nq, k) = GetParam(); GenData(dim, nb, nq, xb, xq, ids, k, gt_ids, gt_dis); + milvus::engine::TempMetaConf tempconf; + tempconf.metric_type = knowhere::METRICTYPE::L2; + tempconf.gpu_id = DEVICEID; + tempconf.size = nb; + tempconf.dim = dim; + tempconf.k = k; + tempconf.nprobe = 16; + index_ = GetVecIndexFactory(index_type); - conf = ParamGenerator::GetInstance().Gen(index_type); - conf->k = k; - conf->d = dim; - conf->gpu_id = DEVICEID; + conf = ParamGenerator::GetInstance().GenBuild(index_type, tempconf); + searchconf = ParamGenerator::GetInstance().GenSearchConf(index_type, tempconf); + +// conf->k = k; +// conf->d = dim; +// conf->gpu_id = DEVICEID; } - void - TearDown() override { + void TearDown() override { #ifdef MILVUS_GPU_VERSION knowhere::FaissGpuResourceMgr::GetInstance().Free(); #endif @@ -65,24 +72,27 @@ class KnowhereWrapperTest milvus::engine::IndexType index_type; milvus::engine::VecIndexPtr index_ = nullptr; knowhere::Config conf; + knowhere::Config searchconf; }; -INSTANTIATE_TEST_CASE_P( - WrapperParam, KnowhereWrapperTest, - Values( -//["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"] +INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest, + Values( + //["Index type", "Generator type", "dim", "nb", "nq", "k", "build config", "search config"] + #ifdef MILVUS_GPU_VERSION std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_GPU, "Default", DIM, NB, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 100000, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_MIX, "Default", 64, 1000, 10, 10), // std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, // 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_GPU, "Default", DIM, NB, 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_MIX, "Default", DIM, NB, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFPQ_MIX, "Default", 64, 1000, 10, 10), + // std::make_tuple(IndexType::NSG_MIX, "Default", 128, 250000, 10, 10), #endif // std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default", 128, 250000, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 100000, 10, 10), - std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 100000, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IDMAP, "Default", 64, 1000, 10, 10), + std::make_tuple(milvus::engine::IndexType::FAISS_IVFFLAT_CPU, "Default", 64, 1000, 10, 10), std::make_tuple(milvus::engine::IndexType::FAISS_IVFSQ8_CPU, "Default", DIM, NB, 10, 10))); TEST_P(KnowhereWrapperTest, BASE_TEST) { @@ -93,12 +103,11 @@ TEST_P(KnowhereWrapperTest, BASE_TEST) { std::vector res_dis(elems); index_->BuildAll(nb, xb.data(), ids.data(), conf); - index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); } #ifdef MILVUS_GPU_VERSION - TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { EXPECT_EQ(index_->GetType(), index_type); @@ -107,13 +116,13 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { std::vector res_dis(elems); index_->BuildAll(nb, xb.data(), ids.data(), conf); - index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); { 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); + dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); } AssertResult(res_ids, res_dis); } @@ -125,7 +134,7 @@ TEST_P(KnowhereWrapperTest, TO_GPU_TEST) { 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); + dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); } AssertResult(res_ids, res_dis); } @@ -139,7 +148,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { std::vector res_ids(elems); std::vector res_dis(elems); index_->BuildAll(nb, xb.data(), ids.data(), conf); - index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); { @@ -152,7 +161,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { std::vector res_ids(elems); std::vector res_dis(elems); - new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); } @@ -166,7 +175,7 @@ TEST_P(KnowhereWrapperTest, SERIALIZE_TEST) { std::vector res_ids(elems); std::vector res_dis(elems); - new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), conf); + new_index->Search(nq, xq.data(), res_dis.data(), res_ids.data(), searchconf); AssertResult(res_ids, res_dis); } } diff --git a/core/unittest/wrapper/utils.h b/core/unittest/wrapper/utils.h index c16b8fc930cc2439872a326b634d49001e8d68fa..1eaa7ae4ec7d70557acf8ceabbd3d85c8a1bb51a 100644 --- a/core/unittest/wrapper/utils.h +++ b/core/unittest/wrapper/utils.h @@ -15,18 +15,22 @@ // specific language governing permissions and limitations // under the License. + #pragma once #include -#include -#include -#include #include #include +#include +#include +#include + -#include "knowhere/index/vector_index/helpers/IndexParameter.h" #include "wrapper/VecIndex.h" #include "wrapper/utils.h" +#include "knowhere/index/vector_index/helpers/IndexParameter.h" +#include "wrapper/ConfAdapterMgr.h" +#include "wrapper/ConfAdapter.h" class DataGenBase; @@ -40,29 +44,31 @@ constexpr int64_t PINMEM = 1024 * 1024 * 200; constexpr int64_t TEMPMEM = 1024 * 1024 * 300; constexpr int64_t RESNUM = 2; -static const char* CONFIG_PATH = "/tmp/milvus_test"; -static const char* CONFIG_FILE = "/server_config.yaml"; +static const char *CONFIG_PATH = "/tmp/milvus_test"; +static const char *CONFIG_FILE = "/server_config.yaml"; class KnowhereTest : public ::testing::Test { protected: - void - SetUp() override; - void - TearDown() override; + void SetUp() override; + void TearDown() override; }; 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& xb, std::vector& xq, - std::vector& ids, const int& k, std::vector& gt_ids, std::vector& gt_dis); + virtual void GenData(const int& dim, + const int& nb, + const int& nq, + std::vector& xb, + std::vector& xq, + std::vector& ids, + const int& k, + std::vector& gt_ids, + std::vector& gt_dis); - void - AssertResult(const std::vector& ids, const std::vector& dis); + void AssertResult(const std::vector& ids, const std::vector& dis); int dim = DIM; int nb = NB; @@ -79,14 +85,22 @@ class DataGenBase { class ParamGenerator { public: - static ParamGenerator& - GetInstance() { + static ParamGenerator& GetInstance() { static ParamGenerator instance; return instance; } - knowhere::Config - Gen(const milvus::engine::IndexType& type) { + knowhere::Config GenSearchConf(const milvus::engine::IndexType& type, const milvus::engine::TempMetaConf& conf) { + auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type); + return adapter->MatchSearch(conf, type); + } + + knowhere::Config GenBuild(const milvus::engine::IndexType& type, const milvus::engine::TempMetaConf& conf) { + auto adapter = milvus::engine::AdapterMgr::GetInstance().GetAdapter(type); + return adapter->Match(conf); + } + + knowhere::Config Gen(const milvus::engine::IndexType& type) { switch (type) { case milvus::engine::IndexType::FAISS_IDMAP: { auto tempconf = std::make_shared(); @@ -113,34 +127,37 @@ class ParamGenerator { 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(); - // 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(); - // 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; - // } + case milvus::engine::IndexType::FAISS_IVFPQ_CPU: + case milvus::engine::IndexType::FAISS_IVFPQ_GPU: + case milvus::engine::IndexType::FAISS_IVFPQ_MIX: { + auto tempconf = std::make_shared(); + 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(); + 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 SanityCheck : public DataGenBase { + +//class SanityCheck : public DataGenBase { // public: // void GenData(const int &dim, const int &nb, const int &nq, float *xb, float *xq, long *ids, // const int &k, long *gt_ids, float *gt_dis) override; //}; + diff --git a/tests/milvus_python_test/conftest.py b/tests/milvus_python_test/conftest.py index 8bab824606685920d4f28758352d1eb5472f11b6..dc58063a451c382503e3a9ae78c955df3f5682cc 100644 --- a/tests/milvus_python_test/conftest.py +++ b/tests/milvus_python_test/conftest.py @@ -42,16 +42,22 @@ def connect(request): port = request.config.getoption("--port") milvus = Milvus() try: - milvus.connect(host=ip, port=port) - except: + status = milvus.connect(host=ip, port=port) + logging.getLogger().info(status) + if not status.OK(): + # try again + logging.getLogger().info("------------------------------------") + logging.getLogger().info("Try to connect again") + logging.getLogger().info("------------------------------------") + res = milvus.connect(host=ip, port=port) + except Exception as e: + logging.getLogger().error(str(e)) pytest.exit("Milvus server can not connected, exit pytest ...") - def fin(): try: milvus.disconnect() except: pass - request.addfinalizer(fin) return milvus @@ -129,4 +135,4 @@ def ip_table(request, connect): request.addfinalizer(teardown) - return table_name \ No newline at end of file + return table_name diff --git a/tests/milvus_python_test/test_connect.py b/tests/milvus_python_test/test_connect.py index 96ce1d3bdfc85b352b9162b1f7c4d65dbed79b39..dd7e80c1f9682c7dcd864a22a0ed61ab6c0532ef 100644 --- a/tests/milvus_python_test/test_connect.py +++ b/tests/milvus_python_test/test_connect.py @@ -249,8 +249,8 @@ class TestConnect: ''' milvus = Milvus() uri_value = "tcp://%s:19540" % args["ip"] - milvus.connect(host=args["ip"], port="", uri=uri_value) - assert milvus.connected() + with pytest.raises(Exception) as e: + milvus.connect(host=args["ip"], port="", uri=uri_value) def test_connect_param_priority_uri(self, args): '''