Merge branch 'branch-0.4.0' into 'branch-0.4.0'

- MS-403 - Add GpuCacheMgr

See merge request megasearch/milvus!409

Former-commit-id: 41664e8a524d3f156c5bb977cc6225d9f3624c9e

cpp/CHANGELOG.md

@@ -40,6 +40,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-394 - Update scheduler unittest
 - MS-400 - Add timestamp record in task state change function
 - MS-402 - Add dump implementation for TaskTableItem
+- MS-403 - Add GpuCacheMgr
 
 ## New Feature
 - MS-343 - Implement ResourceMgr

cpp/src/cache/Cache.h

@@ -46,9 +46,6 @@ public:
 
     double freemem_percent() const { return freemem_percent_; };
     void set_freemem_percent(double percent) { freemem_percent_ = percent; }
-    void set_gpu_ids(std::vector<uint64_t>& gpu_ids) { gpu_ids_ = gpu_ids; }
-
-    std::vector<uint64_t> gpu_ids() const { return gpu_ids_; }
 
     size_t size() const;
     bool exists(const std::string& key);
@@ -63,7 +60,6 @@ private:
     int64_t usage_;
     int64_t capacity_;
     double freemem_percent_;
-    std::vector<uint64_t> gpu_ids_;
 
     LRU<std::string, CacheObjPtr> lru_;
     mutable std::mutex mutex_;

cpp/src/cache/CacheMgr.cpp

@@ -56,7 +56,6 @@ engine::VecIndexPtr CacheMgr::GetIndex(const std::string& key) {
 }
 
 void CacheMgr::InsertItem(const std::string& key, const DataObjPtr& data) {
-    std::cout << "dashalk\n";
     if(cache_ == nullptr) {
         SERVER_LOG_ERROR << "Cache doesn't exist";
         return;
@@ -131,24 +130,6 @@ void CacheMgr::SetCapacity(int64_t 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);
-}
-
 }
 }
 }

cpp/src/cache/CacheMgr.h

@@ -33,8 +33,6 @@ public:
     int64_t CacheUsage() const;
     int64_t CacheCapacity() const;
     void SetCapacity(int64_t capacity);
-    std::vector<uint64_t > GpuIds() const;
-    void SetGpuIds(std::vector<uint64_t> gpu_ids);
 
 protected:
     CacheMgr();

cpp/src/cache/GpuCacheMgr.cpp

@@ -4,6 +4,7 @@
 // Proprietary and confidential.
 ////////////////////////////////////////////////////////////////////////////////
 
+#include <sstream>
 #include "utils/Log.h"
 #include "GpuCacheMgr.h"
 #include "server/ServerConfig.h"
@@ -18,34 +19,39 @@ std::unordered_map<uint64_t, GpuCacheMgrPtr> GpuCacheMgr::instance_;
 namespace {
     constexpr int64_t unit = 1024 * 1024 * 1024;
 
-    void parse_gpu_ids(std::string gpu_ids_str, std::vector<uint64_t>& gpu_ids) {
-        for (auto i = 0; i < gpu_ids_str.length(); ) {
-            if (gpu_ids_str[i] != ',') {
-                int id = 0;
-                while (gpu_ids_str[i] <= '9' && gpu_ids_str[i] >= '0') {
-                    id = id * 10 + gpu_ids_str[i] - '0';
-                    ++i;
-                }
-                gpu_ids.push_back(id);
-            } else {
-                ++i;
+    std::vector<uint64_t> load() {
+        server::ConfigNode& config = server::ServerConfig::GetInstance().GetConfig(server::CONFIG_CACHE);
+        std::string gpu_ids_str = config.GetValue(server::CONFIG_GPU_IDS, "0,1");
+
+        std::vector<uint64_t > gpu_ids;
+
+        std::stringstream ss(gpu_ids_str);
+        for (int i; ss >> i;) {
+            gpu_ids.push_back(i);
+            if (ss.peek() == ',') {
+                ss.ignore();
             }
         }
+        return gpu_ids;
     }
 }
 
+
+bool GpuCacheMgr::GpuIdInConfig(uint64_t gpu_id) {
+    static std::vector<uint64_t > ids = load();
+    for (auto id : ids) {
+        if (gpu_id == id) return true;
+    }
+    return false;
+}
+
 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, 2);
     cap *= unit;
     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);
     if (free_percent > 0.0 && free_percent <= 1.0) {
         cache_->set_freemem_percent(free_percent);

cpp/src/cache/GpuCacheMgr.h

@@ -19,12 +19,18 @@ class GpuCacheMgr : public CacheMgr {
 public:
     GpuCacheMgr();
 
-public:
+    static bool GpuIdInConfig(uint64_t gpu_id);
+
     static CacheMgr* GetInstance(uint64_t gpu_id) {
         if (instance_.find(gpu_id) == instance_.end()) {
             std::lock_guard<std::mutex> lock(mutex_);
-            instance_.insert(std::pair<uint64_t, GpuCacheMgrPtr>(gpu_id, std::make_shared<GpuCacheMgr>()));
-//            instance_[gpu_id] = std::make_shared<GpuCacheMgr>();
+            if (instance_.find(gpu_id) == instance_.end()) {
+                if (GpuIdInConfig(gpu_id)) {
+                    instance_.insert(std::pair<uint64_t, GpuCacheMgrPtr>(gpu_id, std::make_shared<GpuCacheMgr>()));
+                } else {
+                    return nullptr;
+                }
+            }
         }
         return instance_[gpu_id].get();
     }

cpp/src/sdk/examples/grpcsimple/src/ClientTest.cpp

@@ -263,9 +263,6 @@ ClientTest::Test(const std::string& address, const std::string& port) {
                 search_record_array.push_back(
                         std::make_pair(record_ids[SEARCH_TARGET], record_array[SEARCH_TARGET]));
             }
-            int64_t row_count;
-            conn->CountTable(TABLE_NAME, row_count);
-            std::cout << "\t" << TABLE_NAME << "(" << row_count << " rows)" << std::endl;
         }
     }
 

cpp/src/server/grpc_impl/GrpcRequestTask.cpp

@@ -459,7 +459,6 @@ InsertTask::OnExecute() {
                                         std::to_string(table_info.dimension_);
                 return SetError(error_code, error_msg);
             }
-
             memcpy(&vec_f[i * table_info.dimension_],
                    insert_param_.row_record_array(i).vector_data().data(),
                    table_info.dimension_ * sizeof(float));

cpp/unittest/server/cache_test.cpp

@@ -165,8 +165,8 @@ TEST(CacheTest, GPU_CACHE_TEST) {
     gpu_mgr->ClearCache();
     ASSERT_EQ(gpu_mgr->ItemCount(), 0);
 
-    gpu_mgr->SetCapacity(4096000000);
     for (auto i = 0; i < 3; i++) {
+        // TODO: use gpu index to mock
         MockVecIndex *mock_index = new MockVecIndex();
         mock_index->ntotal_ = 1000000;  //2G
         engine::VecIndexPtr index(mock_index);