提交 30b43758 编写于 作者: Y Yu Kun

add GpuCacheMgr and unittest


Former-commit-id: 57e2a59b5039fca56e10c56ee5474472d87faa9a
上级 a6ef8bec
...@@ -46,7 +46,8 @@ public: ...@@ -46,7 +46,8 @@ public:
double freemem_percent() const { return freemem_percent_; }; double freemem_percent() const { return freemem_percent_; };
void set_freemem_percent(double percent) { freemem_percent_ = percent; } void set_freemem_percent(double percent) { freemem_percent_ = percent; }
void set_gpu_ids(std::vector<uint64_t> gpu_ids) { gpu_ids_.assign(gpu_ids.begin(), gpu_ids.end()); } void set_gpu_ids(std::vector<uint64_t>& gpu_ids) { gpu_ids_ = gpu_ids; }
std::vector<uint64_t> gpu_ids() const { return gpu_ids_; } std::vector<uint64_t> gpu_ids() const { return gpu_ids_; }
size_t size() const; size_t size() const;
......
...@@ -56,6 +56,7 @@ engine::VecIndexPtr CacheMgr::GetIndex(const std::string& key) { ...@@ -56,6 +56,7 @@ engine::VecIndexPtr CacheMgr::GetIndex(const std::string& key) {
} }
void CacheMgr::InsertItem(const std::string& key, const DataObjPtr& data) { void CacheMgr::InsertItem(const std::string& key, const DataObjPtr& data) {
std::cout << "dashalk\n";
if(cache_ == nullptr) { if(cache_ == nullptr) {
SERVER_LOG_ERROR << "Cache doesn't exist"; SERVER_LOG_ERROR << "Cache doesn't exist";
return; return;
...@@ -130,6 +131,24 @@ void CacheMgr::SetCapacity(int64_t capacity) { ...@@ -130,6 +131,24 @@ void CacheMgr::SetCapacity(int64_t capacity) {
cache_->set_capacity(capacity); cache_->set_capacity(capacity);
} }
std::vector<uint64_t> CacheMgr::GpuIds() const {
if(cache_ == nullptr) {
SERVER_LOG_ERROR << "Cache doesn't exist";
std::vector<uint64_t> gpu_ids;
return gpu_ids;
}
return cache_->gpu_ids();
}
void CacheMgr::SetGpuIds(std::vector<uint64_t> gpu_ids){
if(cache_ == nullptr) {
SERVER_LOG_ERROR << "Cache doesn't exist";
return;
}
cache_->set_gpu_ids(gpu_ids);
}
} }
} }
} }
...@@ -33,6 +33,8 @@ public: ...@@ -33,6 +33,8 @@ public:
int64_t CacheUsage() const; int64_t CacheUsage() const;
int64_t CacheCapacity() const; int64_t CacheCapacity() const;
void SetCapacity(int64_t capacity); void SetCapacity(int64_t capacity);
std::vector<uint64_t > GpuIds() const;
void SetGpuIds(std::vector<uint64_t> gpu_ids);
protected: protected:
CacheMgr(); CacheMgr();
......
...@@ -16,6 +16,7 @@ private: ...@@ -16,6 +16,7 @@ private:
CpuCacheMgr(); CpuCacheMgr();
public: public:
//TODO: use smart pointer instead
static CacheMgr* GetInstance() { static CacheMgr* GetInstance() {
static CpuCacheMgr s_mgr; static CpuCacheMgr s_mgr;
return &s_mgr; return &s_mgr;
......
...@@ -13,35 +13,39 @@ namespace milvus { ...@@ -13,35 +13,39 @@ namespace milvus {
namespace cache { namespace cache {
std::mutex GpuCacheMgr::mutex_; std::mutex GpuCacheMgr::mutex_;
std::unordered_map<uint64_t, GpuCacheMgr*> GpuCacheMgr::instance_; std::unordered_map<uint64_t, GpuCacheMgrPtr> GpuCacheMgr::instance_;
namespace { namespace {
constexpr int64_t unit = 1024 * 1024 * 1024; constexpr int64_t unit = 1024 * 1024 * 1024;
}
GpuCacheMgr::GpuCacheMgr() { void parse_gpu_ids(std::string gpu_ids_str, std::vector<uint64_t>& gpu_ids) {
server::ConfigNode& config = server::ServerConfig::GetInstance().GetConfig(server::CONFIG_CACHE); for (auto i = 0; i < gpu_ids_str.length(); ) {
std::string gpu_ids_str = config.GetValue(server::CONFIG_GPU_IDS, "0,1"); if (gpu_ids_str[i] != ',') {
std::vector<uint64_t> gpu_ids; int id = 0;
for (auto i = 0; i < gpu_ids_str.length(); ) { while (gpu_ids_str[i] <= '9' && gpu_ids_str[i] >= '0') {
if (gpu_ids_str[i] != ',') { id = id * 10 + gpu_ids_str[i] - '0';
int id = 0; ++i;
while (gpu_ids_str[i] != ',') { }
id = id * 10 + gpu_ids_str[i] - '0'; gpu_ids.push_back(id);
} else {
++i; ++i;
} }
gpu_ids.push_back(id);
} else {
++i;
} }
} }
}
cache_->set_gpu_ids(gpu_ids); GpuCacheMgr::GpuCacheMgr() {
server::ConfigNode& config = server::ServerConfig::GetInstance().GetConfig(server::CONFIG_CACHE);
std::string gpu_ids_str = config.GetValue(server::CONFIG_GPU_IDS, "0,1");
int64_t cap = config.GetInt64Value(server::CONFIG_GPU_CACHE_CAPACITY, 1); int64_t cap = config.GetInt64Value(server::CONFIG_GPU_CACHE_CAPACITY, 2);
cap *= unit; cap *= unit;
cache_ = std::make_shared<Cache>(cap, 1UL<<32); cache_ = std::make_shared<Cache>(cap, 1UL<<32);
std::vector<uint64_t> gpu_ids;
parse_gpu_ids(gpu_ids_str, gpu_ids);
cache_->set_gpu_ids(gpu_ids);
double free_percent = config.GetDoubleValue(server::GPU_CACHE_FREE_PERCENT, 0.85); double free_percent = config.GetDoubleValue(server::GPU_CACHE_FREE_PERCENT, 0.85);
if (free_percent > 0.0 && free_percent <= 1.0) { if (free_percent > 0.0 && free_percent <= 1.0) {
cache_->set_freemem_percent(free_percent); cache_->set_freemem_percent(free_percent);
......
...@@ -6,33 +6,34 @@ ...@@ -6,33 +6,34 @@
#include "CacheMgr.h" #include "CacheMgr.h"
#include <unordered_map> #include <unordered_map>
#include <memory>
namespace zilliz { namespace zilliz {
namespace milvus { namespace milvus {
namespace cache { namespace cache {
class GpuCacheMgr;
using GpuCacheMgrPtr = std::shared_ptr<GpuCacheMgr>;
class GpuCacheMgr : public CacheMgr { class GpuCacheMgr : public CacheMgr {
private: public:
GpuCacheMgr(); GpuCacheMgr();
public: public:
static CacheMgr* GetInstance(uint64_t gpu_id) { static CacheMgr* GetInstance(uint64_t gpu_id) {
if (!instance_[gpu_id]) { if (instance_.find(gpu_id) == instance_.end()) {
std::lock_guard<std::mutex> lock(mutex_); std::lock_guard<std::mutex> lock(mutex_);
if(!instance_[gpu_id]) { instance_.insert(std::pair<uint64_t, GpuCacheMgrPtr>(gpu_id, std::make_shared<GpuCacheMgr>()));
instance_.insert(std::pair<uint64_t, GpuCacheMgr* >(gpu_id, new GpuCacheMgr())); // instance_[gpu_id] = std::make_shared<GpuCacheMgr>();
}
} }
return instance_.at(gpu_id); return instance_[gpu_id].get();
// static GpuCacheMgr s_mgr;
// return &s_mgr;
} }
void InsertItem(const std::string& key, const DataObjPtr& data) override; void InsertItem(const std::string& key, const DataObjPtr& data) override;
private: private:
static std::mutex mutex_; static std::mutex mutex_;
static std::unordered_map<uint64_t, GpuCacheMgr* > instance_; static std::unordered_map<uint64_t, GpuCacheMgrPtr> instance_;
}; };
} }
......
...@@ -9,7 +9,6 @@ ...@@ -9,7 +9,6 @@
#include "db/meta/MetaConsts.h" #include "db/meta/MetaConsts.h"
#include "db/Factories.h" #include "db/Factories.h"
#include "cache/CpuCacheMgr.h" #include "cache/CpuCacheMgr.h"
#include "cache/GpuCacheMgr.h
#include "utils/CommonUtil.h" #include "utils/CommonUtil.h"
#include <gtest/gtest.h> #include <gtest/gtest.h>
...@@ -438,9 +437,4 @@ TEST_F(DBTest2, DELETE_BY_RANGE_TEST) { ...@@ -438,9 +437,4 @@ TEST_F(DBTest2, DELETE_BY_RANGE_TEST) {
ConvertTimeRangeToDBDates(start_value, end_value, dates); ConvertTimeRangeToDBDates(start_value, end_value, dates);
db_->DeleteTable(TABLE_NAME, dates); db_->DeleteTable(TABLE_NAME, dates);
}
TEST_F(DBTest, GPU_CACHE_MGR_TEST) {
std::vector<uint64_t > gpu_ids = cache::
cache::CpuCacheMgr::GetInstance()->CacheUsage();
} }
\ No newline at end of file
...@@ -146,7 +146,7 @@ TEST(CacheTest, CPU_CACHE_TEST) { ...@@ -146,7 +146,7 @@ TEST(CacheTest, CPU_CACHE_TEST) {
} }
TEST(CacheTest, GPU_CACHE_TEST) { TEST(CacheTest, GPU_CACHE_TEST) {
cache::CacheMgr* gpu_mgr = cache::GpuCacheMgr::GetInstance(); cache::CacheMgr* gpu_mgr = cache::GpuCacheMgr::GetInstance(0);
const int dim = 256; const int dim = 256;
...@@ -164,6 +164,25 @@ TEST(CacheTest, GPU_CACHE_TEST) { ...@@ -164,6 +164,25 @@ TEST(CacheTest, GPU_CACHE_TEST) {
gpu_mgr->ClearCache(); gpu_mgr->ClearCache();
ASSERT_EQ(gpu_mgr->ItemCount(), 0); ASSERT_EQ(gpu_mgr->ItemCount(), 0);
gpu_mgr->SetCapacity(4096000000);
for (auto i = 0; i < 3; i++) {
MockVecIndex *mock_index = new MockVecIndex();
mock_index->ntotal_ = 1000000; //2G
engine::VecIndexPtr index(mock_index);
cache::DataObjPtr data_obj = std::make_shared<cache::DataObj>(index);
std::cout << data_obj->size() <<std::endl;
gpu_mgr->InsertItem("index_" + std::to_string(i), data_obj);
}
// ASSERT_EQ(gpu_mgr->ItemCount(), 2);
// auto obj0 = gpu_mgr->GetItem("index_0");
// ASSERT_EQ(obj0, nullptr);
// auto obj1 = gpu_mgr->GetItem("index_1");
// auto obj2 = gpu_mgr->GetItem("index_2");
gpu_mgr->ClearCache();
ASSERT_EQ(gpu_mgr->ItemCount(), 0);
} }
TEST(CacheTest, INVALID_TEST) { TEST(CacheTest, INVALID_TEST) {
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册