Skip to content

milvus
milvus

BaiXuePrincess / milvus
与 Fork 源项目一致

Fork自 milvus / milvus

通知 7
Star 4
Fork 0
milvus
milvus

体验新版 GitCode，发现更多精彩内容 >>

提交 f50dca75 编写于 作者: G groot
浏览文件
操作

format wrapper code 


Former-commit-id: 743ea20fd6a827c949296b16cf026e09e89f112a
上级 8e941b81
隐藏空白更改
内联 并排
Showing with 117 addition and 121 deletion
cpp/src/main.cpp
浏览文件 @ f50dca75
......@@ -27,15 +27,13 @@
#include "utils/CommonUtil.h"
#include "metrics/Metrics.h"
#include "server/Server.h"
#include "version.h"
#include "../version.h"
INITIALIZE_EASYLOGGINGPP
void print_help(const std::string &app_name);
using namespace zilliz::milvus;
int
main(int argc, char *argv[]) {
std::cout << std::endl << "Welcome to use Milvus by Zilliz!" << std::endl;
......@@ -103,6 +101,7 @@ main(int argc, char *argv[]) {
}
/* Handle Signal */
using ::zilliz::milvus;
signal(SIGHUP, server::SignalUtil::HandleSignal);
signal(SIGINT, server::SignalUtil::HandleSignal);
signal(SIGUSR1, server::SignalUtil::HandleSignal);
......
cpp/src/wrapper/DataTransfer.cpp
浏览文件 @ f50dca75
......@@ -16,45 +16,46 @@
// under the License.
#include "DataTransfer.h"
#include "wrapper/DataTransfer.h"
#include <vector>
#include <memory>
#include <utility>
namespace zilliz {
namespace milvus {
namespace engine {
using namespace zilliz::knowhere;
DatasetPtr
GenDatasetWithIds(const int64_t &nb, const int64_t &dim, const float *xb, const long *ids) {
knowhere::DatasetPtr
GenDatasetWithIds(const int64_t &nb, const int64_t &dim, const float *xb, const int64_t *ids) {
std::vector<int64_t> shape{nb, dim};
auto tensor = ConstructFloatTensor((uint8_t *) xb, nb * dim * sizeof(float), shape);
std::vector<TensorPtr> tensors{tensor};
std::vector<FieldPtr> tensor_fields{ConstructFloatField("data")};
auto tensor_schema = std::make_shared<Schema>(tensor_fields);
auto tensor = knowhere::ConstructFloatTensor((uint8_t *) xb, nb * dim * sizeof(float), shape);
std::vector<knowhere::TensorPtr> tensors{tensor};
std::vector<knowhere::FieldPtr> tensor_fields{knowhere::ConstructFloatField("data")};
auto tensor_schema = std::make_shared<knowhere::Schema>(tensor_fields);
auto id_array = ConstructInt64Array((uint8_t *) ids, nb * sizeof(int64_t));
std::vector<ArrayPtr> arrays{id_array};
std::vector<FieldPtr> array_fields{ConstructInt64Field("id")};
auto array_schema = std::make_shared<Schema>(tensor_fields);
auto id_array = knowhere::ConstructInt64Array((uint8_t *) ids, nb * sizeof(int64_t));
std::vector<knowhere::ArrayPtr> arrays{id_array};
std::vector<knowhere::FieldPtr> array_fields{knowhere::ConstructInt64Field("id")};
auto array_schema = std::make_shared<knowhere::Schema>(tensor_fields);
auto dataset = std::make_shared<Dataset>(std::move(arrays), array_schema,
auto dataset = std::make_shared<knowhere::Dataset>(std::move(arrays), array_schema,
std::move(tensors), tensor_schema);
return dataset;
}
DatasetPtr
knowhere::DatasetPtr
GenDataset(const int64_t &nb, const int64_t &dim, const float *xb) {
std::vector<int64_t> shape{nb, dim};
auto tensor = ConstructFloatTensor((uint8_t *) xb, nb * dim * sizeof(float), shape);
std::vector<TensorPtr> tensors{tensor};
std::vector<FieldPtr> tensor_fields{ConstructFloatField("data")};
auto tensor_schema = std::make_shared<Schema>(tensor_fields);
auto tensor = knowhere::ConstructFloatTensor((uint8_t *) xb, nb * dim * sizeof(float), shape);
std::vector<knowhere::TensorPtr> tensors{tensor};
std::vector<knowhere::FieldPtr> tensor_fields{knowhere::ConstructFloatField("data")};
auto tensor_schema = std::make_shared<knowhere::Schema>(tensor_fields);
auto dataset = std::make_shared<Dataset>(std::move(tensors), tensor_schema);
auto dataset = std::make_shared<knowhere::Dataset>(std::move(tensors), tensor_schema);
return dataset;
}
}
}
}
} // namespace engine
} // namespace milvus
} // namespace zilliz
cpp/src/wrapper/DataTransfer.h
浏览文件 @ f50dca75
......@@ -20,17 +20,16 @@
#include "knowhere/adapter/Structure.h"
namespace zilliz {
namespace milvus {
namespace engine {
extern zilliz::knowhere::DatasetPtr
GenDatasetWithIds(const int64_t &nb, const int64_t &dim, const float *xb, const long *ids);
GenDatasetWithIds(const int64_t &nb, const int64_t &dim, const float *xb, const int64_t *ids);
extern zilliz::knowhere::DatasetPtr
GenDataset(const int64_t &nb, const int64_t &dim, const float *xb);
}
}
}
} // namespace engine
} // namespace milvus
} // namespace zilliz
cpp/src/wrapper/KnowhereResource.cpp
浏览文件 @ f50dca75
......@@ -16,11 +16,15 @@
// under the License.
#include "KnowhereResource.h"
#include "wrapper/KnowhereResource.h"
#include "knowhere/index/vector_index/helpers/FaissGpuResourceMgr.h"
#include "server/Config.h"
#include <map>
#include <set>
#include <vector>
#include <string>
#include <utility>
namespace zilliz {
namespace milvus {
......@@ -31,20 +35,20 @@ constexpr int64_t M_BYTE = 1024 * 1024;
Status
KnowhereResource::Initialize() {
struct GpuResourceSetting {
int64_t pinned_memory = 300*M_BYTE;
int64_t temp_memory = 300*M_BYTE;
int64_t pinned_memory = 300 * M_BYTE;
int64_t temp_memory = 300 * M_BYTE;
int64_t resource_num = 2;
};
using GpuResourcesArray = std::map<int64_t , GpuResourceSetting>;
using GpuResourcesArray = std::map<int64_t, GpuResourceSetting>;
GpuResourcesArray gpu_resources;
Status s;
//get build index gpu resource
server::Config& config = server::Config::GetInstance();
server::Config &config = server::Config::GetInstance();
int32_t build_index_gpu;
s = config.GetDBConfigBuildIndexGPU(build_index_gpu);
if (!s.ok()) return s;
if (!s.ok()) return s;
gpu_resources.insert(std::make_pair(build_index_gpu, GpuResourceSetting()));
......@@ -64,7 +68,7 @@ KnowhereResource::Initialize() {
}
//init gpu resources
for(auto iter = gpu_resources.begin(); iter != gpu_resources.end(); ++iter) {
for (auto iter = gpu_resources.begin(); iter != gpu_resources.end(); ++iter) {
knowhere::FaissGpuResourceMgr::GetInstance().InitDevice(iter->first,
iter->second.pinned_memory,
iter->second.temp_memory,
......@@ -80,6 +84,6 @@ KnowhereResource::Finalize() {
return Status::OK();
}
}
}
}
\ No newline at end of file
} // namespace engine
} // namespace milvus
} // namespace zilliz
cpp/src/wrapper/KnowhereResource.h
浏览文件 @ f50dca75
......@@ -25,7 +25,7 @@ namespace milvus {
namespace engine {
class KnowhereResource {
public:
public:
static Status
Initialize();
......@@ -33,7 +33,6 @@ public:
Finalize();
};
}
}
}
} // namespace engine
} // namespace milvus
} // namespace zilliz
cpp/src/wrapper/VecImpl.cpp
浏览文件 @ f50dca75
......@@ -16,27 +16,24 @@
// under the License.
#include "wrapper/VecImpl.h"
#include "utils/Log.h"
#include "knowhere/index/vector_index/IndexIDMAP.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/common/Exception.h"
#include "knowhere/index/vector_index/helpers/Cloner.h"
#include "VecImpl.h"
#include "DataTransfer.h"
namespace zilliz {
namespace milvus {
namespace engine {
using namespace zilliz::knowhere;
Status
VecIndexImpl::BuildAll(const long &nb,
VecIndexImpl::BuildAll(const int64_t &nb,
const float *xb,
const long *ids,
const int64_t *ids,
const Config &cfg,
const long &nt,
const int64_t &nt,
const float *xt) {
try {
dim = cfg["dim"].as<int>();
......@@ -47,7 +44,7 @@ VecIndexImpl::BuildAll(const long &nb,
auto model = index_->Train(dataset, cfg);
index_->set_index_model(model);
index_->Add(dataset, cfg);
} catch (KnowhereException &e) {
} catch (knowhere::KnowhereException &e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (jsoncons::json_exception &e) {
......@@ -61,12 +58,12 @@ VecIndexImpl::BuildAll(const long &nb,
}
Status
VecIndexImpl::Add(const long &nb, const float *xb, const long *ids, const Config &cfg) {
VecIndexImpl::Add(const int64_t &nb, const float *xb, const int64_t *ids, const Config &cfg) {
try {
auto dataset = GenDatasetWithIds(nb, dim, xb, ids);
index_->Add(dataset, cfg);
} catch (KnowhereException &e) {
} catch (knowhere::KnowhereException &e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (jsoncons::json_exception &e) {
......@@ -80,7 +77,7 @@ VecIndexImpl::Add(const long &nb, const float *xb, const long *ids, const Config
}
Status
VecIndexImpl::Search(const long &nq, const float *xq, float *dist, long *ids, const Config &cfg) {
VecIndexImpl::Search(const int64_t &nq, const float *xq, float *dist, int64_t *ids, const Config &cfg) {
try {
auto k = cfg["k"].as<int>();
auto dataset = GenDataset(nq, dim, xq);
......@@ -116,8 +113,7 @@ VecIndexImpl::Search(const long &nq, const float *xq, float *dist, long *ids, co
// TODO(linxj): avoid copy here.
memcpy(ids, p_ids, sizeof(int64_t) * nq * k);
memcpy(dist, p_dist, sizeof(float) * nq * k);
} catch (KnowhereException &e) {
} catch (knowhere::KnowhereException &e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (jsoncons::json_exception &e) {
......@@ -186,7 +182,7 @@ VecIndexImpl::Clone() {
int64_t
VecIndexImpl::GetDeviceId() {
if (auto device_idx = std::dynamic_pointer_cast<GPUIndex>(index_)) {
if (auto device_idx = std::dynamic_pointer_cast<knowhere::GPUIndex>(index_)) {
return device_idx->GetGpuDevice();
}
// else
......@@ -195,22 +191,22 @@ VecIndexImpl::GetDeviceId() {
float *
BFIndex::GetRawVectors() {
auto raw_index = std::dynamic_pointer_cast<IDMAP>(index_);
auto raw_index = std::dynamic_pointer_cast<knowhere::IDMAP>(index_);
if (raw_index) { return raw_index->GetRawVectors(); }
return nullptr;
}
int64_t *
BFIndex::GetRawIds() {
return std::static_pointer_cast<IDMAP>(index_)->GetRawIds();
return std::static_pointer_cast<knowhere::IDMAP>(index_)->GetRawIds();
}
ErrorCode
BFIndex::Build(const Config &cfg) {
try {
dim = cfg["dim"].as<int>();
std::static_pointer_cast<IDMAP>(index_)->Train(cfg);
} catch (KnowhereException &e) {
std::static_pointer_cast<knowhere::IDMAP>(index_)->Train(cfg);
} catch (knowhere::KnowhereException &e) {
WRAPPER_LOG_ERROR << e.what();
return KNOWHERE_UNEXPECTED_ERROR;
} catch (jsoncons::json_exception &e) {
......@@ -224,19 +220,19 @@ BFIndex::Build(const Config &cfg) {
}
Status
BFIndex::BuildAll(const long &nb,
BFIndex::BuildAll(const int64_t &nb,
const float *xb,
const long *ids,
const int64_t *ids,
const Config &cfg,
const long &nt,
const int64_t &nt,
const float *xt) {
try {
dim = cfg["dim"].as<int>();
auto dataset = GenDatasetWithIds(nb, dim, xb, ids);
std::static_pointer_cast<IDMAP>(index_)->Train(cfg);
std::static_pointer_cast<knowhere::IDMAP>(index_)->Train(cfg);
index_->Add(dataset, cfg);
} catch (KnowhereException &e) {
} catch (knowhere::KnowhereException &e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (jsoncons::json_exception &e) {
......@@ -251,11 +247,11 @@ BFIndex::BuildAll(const long &nb,
// TODO(linxj): add lock here.
Status
IVFMixIndex::BuildAll(const long &nb,
IVFMixIndex::BuildAll(const int64_t &nb,
const float *xb,
const long *ids,
const int64_t *ids,
const Config &cfg,
const long &nt,
const int64_t &nt,
const float *xt) {
try {
dim = cfg["dim"].as<int>();
......@@ -267,7 +263,7 @@ IVFMixIndex::BuildAll(const long &nb,
index_->set_index_model(model);
index_->Add(dataset, cfg);
if (auto device_index = std::dynamic_pointer_cast<GPUIVF>(index_)) {
if (auto device_index = std::dynamic_pointer_cast<knowhere::GPUIVF>(index_)) {
auto host_index = device_index->CopyGpuToCpu(Config());
index_ = host_index;
type = ConvertToCpuIndexType(type);
......@@ -275,7 +271,7 @@ IVFMixIndex::BuildAll(const long &nb,
WRAPPER_LOG_ERROR << "Build IVFMIXIndex Failed";
return Status(KNOWHERE_ERROR, "Build IVFMIXIndex Failed");
}
} catch (KnowhereException &e) {
} catch (knowhere::KnowhereException &e) {
WRAPPER_LOG_ERROR << e.what();
return Status(KNOWHERE_UNEXPECTED_ERROR, e.what());
} catch (jsoncons::json_exception &e) {
......@@ -296,6 +292,6 @@ IVFMixIndex::Load(const zilliz::knowhere::BinarySet &index_binary) {
return Status::OK();
}
}
}
}
} // namespace engine
} // namespace milvus
} // namespace zilliz
cpp/src/wrapper/VecImpl.h
浏览文件 @ f50dca75
......@@ -21,6 +21,8 @@
#include "knowhere/index/vector_index/VectorIndex.h"
#include "VecIndex.h"
#include <utility>
#include <memory>
namespace zilliz {
namespace milvus {
......@@ -29,14 +31,15 @@ namespace engine {
class VecIndexImpl : public VecIndex {
public:
explicit VecIndexImpl(std::shared_ptr<zilliz::knowhere::VectorIndex> index, const IndexType &type)
: index_(std::move(index)), type(type) {};
: index_(std::move(index)), type(type) {
}
Status
BuildAll(const long &nb,
BuildAll(const int64_t &nb,
const float *xb,
const long *ids,
const int64_t *ids,
const Config &cfg,
const long &nt,
const int64_t &nt,
const float *xt) override;
VecIndexPtr
......@@ -55,7 +58,7 @@ class VecIndexImpl : public VecIndex {
Count() override;
Status
Add(const long &nb, const float *xb, const long *ids, const Config &cfg) override;
Add(const int64_t &nb, const float *xb, const int64_t *ids, const Config &cfg) override;
zilliz::knowhere::BinarySet
Serialize() override;
......@@ -70,7 +73,7 @@ class VecIndexImpl : public VecIndex {
GetDeviceId() override;
Status
Search(const long &nq, const float *xq, float *dist, long *ids, const Config &cfg) override;
Search(const int64_t &nq, const float *xq, float *dist, int64_t *ids, const Config &cfg) override;
protected:
int64_t dim = 0;
......@@ -83,14 +86,15 @@ class VecIndexImpl : public VecIndex {
class IVFMixIndex : public VecIndexImpl {
public:
explicit IVFMixIndex(std::shared_ptr<zilliz::knowhere::VectorIndex> index, const IndexType &type)
: VecIndexImpl(std::move(index), type) {};
: VecIndexImpl(std::move(index), type) {
}
Status
BuildAll(const long &nb,
BuildAll(const int64_t &nb,
const float *xb,
const long *ids,
const int64_t *ids,
const Config &cfg,
const long &nt,
const int64_t &nt,
const float *xt) override;
Status
......@@ -100,7 +104,8 @@ class IVFMixIndex : public VecIndexImpl {
class BFIndex : public VecIndexImpl {
public:
explicit BFIndex(std::shared_ptr<zilliz::knowhere::VectorIndex> index) : VecIndexImpl(std::move(index),
IndexType::FAISS_IDMAP) {};
IndexType::FAISS_IDMAP) {
}
ErrorCode
Build(const Config &cfg);
......@@ -109,17 +114,17 @@ class BFIndex : public VecIndexImpl {
GetRawVectors();
Status
BuildAll(const long &nb,
BuildAll(const int64_t &nb,
const float *xb,
const long *ids,
const int64_t *ids,
const Config &cfg,
const long &nt,
const int64_t &nt,
const float *xt) override;
int64_t *
GetRawIds();
};
}
}
}
} // namespace engine
} // namespace milvus
} // namespace zilliz
cpp/src/wrapper/VecIndex.cpp
浏览文件 @ f50dca75
......@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
#include "wrapper/VecIndex.h"
#include "knowhere/index/vector_index/IndexIVF.h"
#include "knowhere/index/vector_index/IndexGPUIVF.h"
#include "knowhere/index/vector_index/IndexIVFSQ.h"
......@@ -25,13 +26,11 @@
#include "knowhere/index/vector_index/IndexKDT.h"
#include "knowhere/index/vector_index/IndexNSG.h"
#include "knowhere/common/Exception.h"
#include "VecIndex.h"
#include "VecImpl.h"
#include "utils/Log.h"
#include <cuda.h>
namespace zilliz {
namespace milvus {
namespace engine {
......@@ -42,7 +41,7 @@ struct FileIOReader {
std::fstream fs;
std::string name;
FileIOReader(const std::string &fname);
explicit FileIOReader(const std::string &fname);
~FileIOReader();
......@@ -72,13 +71,11 @@ FileIOReader::operator()(void *ptr, size_t size, size_t pos) {
return 0;
}
struct FileIOWriter {
std::fstream fs;
std::string name;
FileIOWriter(const std::string &fname);
explicit FileIOWriter(const std::string &fname);
~FileIOWriter();
size_t operator()(void *ptr, size_t size);
};
......@@ -97,7 +94,6 @@ FileIOWriter::operator()(void *ptr, size_t size) {
fs.write(reinterpret_cast<char *>(ptr), size);
}
VecIndexPtr
GetVecIndexFactory(const IndexType &type, const Config &cfg) {
std::shared_ptr<zilliz::knowhere::VectorIndex> index;
......@@ -215,7 +211,7 @@ write_index(VecIndexPtr index, const std::string &location) {
FileIOWriter writer(location);
writer(&index_type, sizeof(IndexType));
for (auto &iter: binaryset.binary_map_) {
for (auto &iter : binaryset.binary_map_) {
auto meta = iter.first.c_str();
size_t meta_length = iter.first.length();
writer(&meta_length, sizeof(meta_length));
......@@ -242,15 +238,14 @@ write_index(VecIndexPtr index, const std::string &location) {
return Status::OK();
}
// TODO(linxj): redo here.
void
AutoGenParams(const IndexType &type, const long &size, zilliz::knowhere::Config &cfg) {
AutoGenParams(const IndexType &type, const int64_t &size, zilliz::knowhere::Config &cfg) {
auto nlist = cfg.get_with_default("nlist", 0);
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);
}
......@@ -341,7 +336,6 @@ ConvertToGpuIndexType(const IndexType &type) {
}
}
}
}
}
} // namespace engine
} // namespace milvus
} // namespace zilliz
cpp/src/wrapper/VecIndex.h
浏览文件 @ f50dca75
......@@ -25,7 +25,6 @@
#include "knowhere/common/Config.h"
#include "knowhere/common/BinarySet.h"
namespace zilliz {
namespace milvus {
namespace engine {
......@@ -55,24 +54,24 @@ using VecIndexPtr = std::shared_ptr<VecIndex>;
class VecIndex {
public:
virtual Status
BuildAll(const long &nb,
BuildAll(const int64_t &nb,
const float *xb,
const long *ids,
const int64_t *ids,
const Config &cfg,
const long &nt = 0,
const int64_t &nt = 0,
const float *xt = nullptr) = 0;
virtual Status
Add(const long &nb,
Add(const int64_t &nb,
const float *xb,
const long *ids,
const int64_t *ids,
const Config &cfg = Config()) = 0;
virtual Status
Search(const long &nq,
Search(const int64_t &nq,
const float *xq,
float *dist,
long *ids,
int64_t *ids,
const Config &cfg = Config()) = 0;
virtual VecIndexPtr
......@@ -117,7 +116,7 @@ extern VecIndexPtr
LoadVecIndex(const IndexType &index_type, const zilliz::knowhere::BinarySet &index_binary);
extern void
AutoGenParams(const IndexType &type, const long &size, Config &cfg);
AutoGenParams(const IndexType &type, const int64_t &size, Config &cfg);
extern void
ParameterValidation(const IndexType &type, Config &cfg);
......@@ -128,6 +127,6 @@ ConvertToCpuIndexType(const IndexType &type);
extern IndexType
ConvertToGpuIndexType(const IndexType &type);
}
}
}
} // namespace engine
} // namespace milvus
} // namespace zilliz
Markdown is supported
0% 或 .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑！
想要评论请 注册 或