mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 05:18:52 +08:00
MS-442 make it work
Former-commit-id: 7525662e68797e3049b22b8c8c2549afcb20c4cc
This commit is contained in:
parent
400cfaee06
commit
cd8e16858f
@ -34,7 +34,7 @@ else()
|
||||
endif()
|
||||
message(STATUS "Build type = ${BUILD_TYPE}")
|
||||
|
||||
#add_definitions(-DNEW_SCHEDULER)
|
||||
add_definitions(-DNEW_SCHEDULER)
|
||||
|
||||
project(milvus VERSION "${MILVUS_VERSION}")
|
||||
project(milvus_engine LANGUAGES CUDA CXX)
|
||||
|
@ -66,24 +66,24 @@ resource_config:
|
||||
enable_loader: true
|
||||
enable_executor: true
|
||||
|
||||
gtx1060:
|
||||
type: GPU
|
||||
memory: 6
|
||||
device_id: 0
|
||||
enable_loader: true
|
||||
enable_executor: true
|
||||
# gtx1060:
|
||||
# type: GPU
|
||||
# memory: 6
|
||||
# device_id: 0
|
||||
# enable_loader: true
|
||||
# enable_executor: true
|
||||
|
||||
gtx1660:
|
||||
type: GPU
|
||||
memory: 6
|
||||
device_id: 1
|
||||
enable_loader: true
|
||||
enable_executor: true
|
||||
# gtx1660:
|
||||
# type: GPU
|
||||
# memory: 6
|
||||
# device_id: 1
|
||||
# enable_loader: false
|
||||
# enable_executor: false
|
||||
|
||||
# connection list, length: 0~N
|
||||
# format: -${resource_name}===${resource_name}
|
||||
connections:
|
||||
- ssda===cpu
|
||||
- cpu===gtx1060
|
||||
- cpu===gtx1660
|
||||
# - cpu===gtx1060
|
||||
# - cpu===gtx1660
|
||||
|
||||
|
@ -7,6 +7,7 @@
|
||||
#include "SchedInst.h"
|
||||
#include "server/ServerConfig.h"
|
||||
#include "ResourceFactory.h"
|
||||
#include "knowhere/index/vector_index/gpu_ivf.h"
|
||||
|
||||
namespace zilliz {
|
||||
namespace milvus {
|
||||
@ -36,8 +37,12 @@ SchedServInit() {
|
||||
device_id,
|
||||
enable_loader,
|
||||
enable_executor));
|
||||
|
||||
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(device_id);
|
||||
}
|
||||
|
||||
knowhere::FaissGpuResourceMgr::GetInstance().InitResource();
|
||||
|
||||
auto default_connection = Connection("default_connection", 500.0);
|
||||
auto connections = config.GetSequence(server::CONFIG_RESOURCE_CONNECTIONS);
|
||||
for (auto &conn : connections) {
|
||||
|
@ -241,8 +241,9 @@ server::KnowhereError IVFMixIndex::BuildAll(const long &nb,
|
||||
index_->Add(dataset, cfg);
|
||||
|
||||
if (auto device_index = std::dynamic_pointer_cast<GPUIVF>(index_)) {
|
||||
auto host_index = device_index->Copy_index_gpu_to_cpu();
|
||||
auto host_index = device_index->CopyGpuToCpu(Config());
|
||||
index_ = host_index;
|
||||
type = TransferToCpuIndexType(type);
|
||||
} else {
|
||||
WRAPPER_LOG_ERROR << "Build IVFMIXIndex Failed";
|
||||
}
|
||||
|
@ -106,6 +106,10 @@ VecIndexPtr GetVecIndexFactory(const IndexType &type) {
|
||||
index = std::make_shared<zilliz::knowhere::GPUIVFSQ>(0);
|
||||
return std::make_shared<IVFMixIndex>(index, IndexType::FAISS_IVFSQ8_MIX);
|
||||
}
|
||||
case IndexType::FAISS_IVFSQ8: {
|
||||
index = std::make_shared<zilliz::knowhere::IVFSQ>();
|
||||
break;
|
||||
}
|
||||
case IndexType::NSG_MIX: { // TODO(linxj): bug.
|
||||
index = std::make_shared<zilliz::knowhere::NSG>(0);
|
||||
break;
|
||||
@ -194,10 +198,10 @@ server::KnowhereError write_index(VecIndexPtr index, const std::string &location
|
||||
// TODO(linxj): redo here.
|
||||
void AutoGenParams(const IndexType &type, const long &size, zilliz::knowhere::Config &cfg) {
|
||||
auto nlist = cfg.get_with_default("nlist", 0);
|
||||
if (size <= TYPICAL_COUNT/16384 + 1) {
|
||||
if (size <= TYPICAL_COUNT / 16384 + 1) {
|
||||
//handle less row count, avoid nlist set to 0
|
||||
cfg["nlist"] = 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
|
||||
cfg["nlist"] = int(size / TYPICAL_COUNT * 16384);
|
||||
}
|
||||
@ -225,6 +229,20 @@ void AutoGenParams(const IndexType &type, const long &size, zilliz::knowhere::Co
|
||||
}
|
||||
}
|
||||
|
||||
IndexType TransferToCpuIndexType(const IndexType &type) {
|
||||
switch (type) {
|
||||
case IndexType::FAISS_IVFFLAT_MIX: {
|
||||
return IndexType::FAISS_IVFFLAT_CPU;
|
||||
}
|
||||
case IndexType::FAISS_IVFSQ8_MIX: {
|
||||
return IndexType::FAISS_IVFSQ8;
|
||||
}
|
||||
default: {
|
||||
return IndexType::INVALID;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
|
@ -32,6 +32,7 @@ enum class IndexType {
|
||||
FAISS_IVFPQ_GPU,
|
||||
SPTAG_KDT_RNT_CPU,
|
||||
FAISS_IVFSQ8_MIX,
|
||||
FAISS_IVFSQ8,
|
||||
NSG_MIX,
|
||||
};
|
||||
|
||||
@ -88,6 +89,8 @@ extern VecIndexPtr LoadVecIndex(const IndexType &index_type, const zilliz::knowh
|
||||
|
||||
extern void AutoGenParams(const IndexType& type, const long& size, Config& cfg);
|
||||
|
||||
extern IndexType TransferToCpuIndexType(const IndexType& type);
|
||||
|
||||
}
|
||||
}
|
||||
}
|
||||
|
@ -8,13 +8,14 @@
|
||||
#include <easylogging++.h>
|
||||
|
||||
#include <wrapper/knowhere/vec_index.h>
|
||||
#include "knowhere/index/vector_index/gpu_ivf.h"
|
||||
|
||||
#include "utils.h"
|
||||
|
||||
INITIALIZE_EASYLOGGINGPP
|
||||
|
||||
using namespace zilliz::milvus::engine;
|
||||
using namespace zilliz::knowhere;
|
||||
//using namespace zilliz::knowhere;
|
||||
|
||||
using ::testing::TestWithParam;
|
||||
using ::testing::Values;
|
||||
@ -66,8 +67,8 @@ class KnowhereWrapperTest
|
||||
Config train_cfg;
|
||||
Config search_cfg;
|
||||
|
||||
int dim = 64;
|
||||
int nb = 10000;
|
||||
int dim = 512;
|
||||
int nb = 1000000;
|
||||
int nq = 10;
|
||||
int k = 10;
|
||||
std::vector<float> xb;
|
||||
@ -94,27 +95,27 @@ INSTANTIATE_TEST_CASE_P(WrapperParam, KnowhereWrapperTest,
|
||||
// Config::object{{"nlist", 100}, {"dim", 64}},
|
||||
// Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 40}}
|
||||
//),
|
||||
std::make_tuple(IndexType::FAISS_IVFFLAT_MIX, "Default",
|
||||
64, 100000, 10, 10,
|
||||
Config::object{{"nlist", 1000}, {"dim", 64}, {"metric_type", "L2"}},
|
||||
Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 5}}
|
||||
),
|
||||
std::make_tuple(IndexType::FAISS_IDMAP, "Default",
|
||||
64, 100000, 10, 10,
|
||||
Config::object{{"dim", 64}, {"metric_type", "L2"}},
|
||||
Config::object{{"dim", 64}, {"k", 10}}
|
||||
),
|
||||
// std::make_tuple(IndexType::FAISS_IVFFLAT_MIX, "Default",
|
||||
// 64, 100000, 10, 10,
|
||||
// Config::object{{"nlist", 1000}, {"dim", 64}, {"metric_type", "L2"}},
|
||||
// Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 5}}
|
||||
// ),
|
||||
// std::make_tuple(IndexType::FAISS_IDMAP, "Default",
|
||||
// 64, 100000, 10, 10,
|
||||
// Config::object{{"dim", 64}, {"metric_type", "L2"}},
|
||||
// Config::object{{"dim", 64}, {"k", 10}}
|
||||
// ),
|
||||
std::make_tuple(IndexType::FAISS_IVFSQ8_MIX, "Default",
|
||||
64, 100000, 10, 10,
|
||||
Config::object{{"dim", 64}, {"nlist", 1000}, {"nbits", 8}, {"metric_type", "L2"}},
|
||||
Config::object{{"dim", 64}, {"k", 10}, {"nprobe", 5}}
|
||||
),
|
||||
std::make_tuple(IndexType::NSG_MIX, "Default",
|
||||
128, 250000, 10, 10,
|
||||
Config::object{{"dim", 128}, {"nlist", 8192}, {"nprobe", 16}, {"metric_type", "L2"},
|
||||
{"knng", 200}, {"search_length", 40}, {"out_degree", 60}, {"candidate_pool_size", 200}},
|
||||
Config::object{{"k", 10}, {"search_length", 20}}
|
||||
512, 1000000, 10, 10,
|
||||
Config::object{{"dim", 512}, {"nlist", 1000}, {"nbits", 8}, {"metric_type", "L2"}},
|
||||
Config::object{{"dim", 512}, {"k", 10}, {"nprobe", 5}}
|
||||
)
|
||||
// std::make_tuple(IndexType::NSG_MIX, "Default",
|
||||
// 128, 250000, 10, 10,
|
||||
// Config::object{{"dim", 128}, {"nlist", 8192}, {"nprobe", 16}, {"metric_type", "L2"},
|
||||
// {"knng", 200}, {"search_length", 40}, {"out_degree", 60}, {"candidate_pool_size", 200}},
|
||||
// Config::object{{"k", 10}, {"search_length", 20}}
|
||||
// )
|
||||
//std::make_tuple(IndexType::SPTAG_KDT_RNT_CPU, "Default",
|
||||
// 64, 10000, 10, 10,
|
||||
// Config::object{{"TPTNumber", 1}, {"dim", 64}},
|
||||
@ -135,6 +136,34 @@ TEST_P(KnowhereWrapperTest, base_test) {
|
||||
AssertResult(res_ids, res_dis);
|
||||
}
|
||||
|
||||
TEST_P(KnowhereWrapperTest, to_gpu_test) {
|
||||
EXPECT_EQ(index_->GetType(), index_type);
|
||||
|
||||
zilliz::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(0);
|
||||
zilliz::knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(1);
|
||||
|
||||
auto elems = nq * k;
|
||||
std::vector<int64_t> res_ids(elems);
|
||||
std::vector<float> res_dis(elems);
|
||||
|
||||
index_->BuildAll(nb, xb.data(), ids.data(), train_cfg);
|
||||
index_->Search(nq, xq.data(), res_dis.data(), res_ids.data(), search_cfg);
|
||||
AssertResult(res_ids, res_dis);
|
||||
{
|
||||
index_->CopyToGpu(1);
|
||||
}
|
||||
|
||||
std::string file_location = "/tmp/whatever";
|
||||
write_index(index_, file_location);
|
||||
auto new_index = read_index(file_location);
|
||||
|
||||
auto dev_idx = new_index->CopyToGpu(1);
|
||||
for (int i = 0; i < 10000; ++i) {
|
||||
dev_idx->Search(nq, xq.data(), res_dis.data(), res_ids.data(), search_cfg);
|
||||
}
|
||||
AssertResult(res_ids, res_dis);
|
||||
}
|
||||
|
||||
TEST_P(KnowhereWrapperTest, serialize) {
|
||||
EXPECT_EQ(index_->GetType(), index_type);
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user