MS-626 Refactor DataObj to support cache any type data

Former-commit-id: 7795e47deb81657c484d4d90acb49adf9ed4501c

cpp/CHANGELOG.md

@@ -28,6 +28,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-609 - Update task construct function
 - MS-611 - Add resources validity check in ResourceMgr
 - MS-619 - Add optimizer class in scheduler
+- MS-626 - Refactor DataObj to support cache any type data
 
 ## New Feature
 

cpp/src/cache/Cache.inl

@@ -86,11 +86,11 @@ Cache<ItemObj>::insert(const std::string &key, const ItemObj &item) {
         //if key already exist, subtract old item size
         if (lru_.exists(key)) {
             const ItemObj &old_item = lru_.get(key);
-            usage_ -= old_item->size();
+            usage_ -= old_item->Size();
         }
 
         //plus new item size
-        usage_ += item->size();
+        usage_ += item->Size();
     }
 
     //if usage exceed capacity, free some items
@@ -106,7 +106,7 @@ Cache<ItemObj>::insert(const std::string &key, const ItemObj &item) {
         std::lock_guard<std::mutex> lock(mutex_);
 
         lru_.put(key, item);
-        SERVER_LOG_DEBUG << "Insert " << key << " size:" << item->size()
+        SERVER_LOG_DEBUG << "Insert " << key << " size:" << item->Size()
                          << " bytes into cache, usage: " << usage_ << " bytes";
     }
 }
@@ -120,9 +120,9 @@ Cache<ItemObj>::erase(const std::string &key) {
     }
 
     const ItemObj &old_item = lru_.get(key);
-    usage_ -= old_item->size();
+    usage_ -= old_item->Size();
 
-    SERVER_LOG_DEBUG << "Erase " << key << " size: " << old_item->size();
+    SERVER_LOG_DEBUG << "Erase " << key << " size: " << old_item->Size();
 
     lru_.erase(key);
 }
@@ -160,7 +160,7 @@ Cache<ItemObj>::free_memory() {
             auto &obj_ptr = it->second;
 
             key_array.emplace(key);
-            released_size += obj_ptr->size();
+            released_size += obj_ptr->Size();
             ++it;
         }
     }

cpp/src/cache/CpuCacheMgr.cpp

@@ -59,14 +59,10 @@ CpuCacheMgr::GetInstance() {
     return &s_mgr;
 }
 
-engine::VecIndexPtr
+DataObjPtr
 CpuCacheMgr::GetIndex(const std::string& key) {
     DataObjPtr obj = GetItem(key);
-    if (obj != nullptr) {
-        return obj->data();
-    }
-
-    return nullptr;
+    return obj;
 }
 
 }  // namespace cache

cpp/src/cache/CpuCacheMgr.h

@@ -35,7 +35,7 @@ class CpuCacheMgr : public CacheMgr<DataObjPtr> {
     static CpuCacheMgr*
     GetInstance();
 
-    engine::VecIndexPtr
+    DataObjPtr
     GetIndex(const std::string& key);
 };
 

cpp/src/cache/DataObj.h

@@ -17,7 +17,6 @@
 
 #pragma once
 
-#include "src/wrapper/VecIndex.h"
 
 #include <memory>
 
@@ -26,38 +25,9 @@ namespace cache {
 
 class DataObj {
  public:
-    explicit DataObj(const engine::VecIndexPtr& index) : index_(index) {
-    }
+    virtual int64_t
+    Size() = 0;
 
-    DataObj(const engine::VecIndexPtr& index, int64_t size) : index_(index), size_(size) {
-    }
-
-    engine::VecIndexPtr
-    data() {
-        return index_;
-    }
-
-    const engine::VecIndexPtr&
-    data() const {
-        return index_;
-    }
-
-    int64_t
-    size() const {
-        if (index_ == nullptr) {
-            return 0;
-        }
-
-        if (size_ > 0) {
-            return size_;
-        }
-
-        return index_->Count() * index_->Dimension() * sizeof(float);
-    }
-
- private:
-    engine::VecIndexPtr index_ = nullptr;
-    int64_t size_ = 0;
 };
 
 using DataObjPtr = std::shared_ptr<DataObj>;

cpp/src/cache/GpuCacheMgr.cpp

@@ -71,14 +71,10 @@ GpuCacheMgr::GetInstance(uint64_t gpu_id) {
     }
 }
 
-engine::VecIndexPtr
+DataObjPtr
 GpuCacheMgr::GetIndex(const std::string& key) {
     DataObjPtr obj = GetItem(key);
-    if (obj != nullptr) {
-        return obj->data();
-    }
-
-    return nullptr;
+    return obj;
 }
 
 }  // namespace cache

cpp/src/cache/GpuCacheMgr.h

@@ -35,7 +35,7 @@ class GpuCacheMgr : public CacheMgr<DataObjPtr> {
     static GpuCacheMgr*
     GetInstance(uint64_t gpu_id);
 
-    engine::VecIndexPtr
+    DataObjPtr
     GetIndex(const std::string& key);
 
  private:

cpp/src/db/engine/ExecutionEngineImpl.cpp

@@ -91,6 +91,60 @@ ExecutionEngineImpl::CreatetVecIndex(EngineType type) {
     return index;
 }
 
+void
+ExecutionEngineImpl::HybridLoad() {
+//    if (index_type_ != EngineType::FAISS_IVFSQ8Hybrid) {
+//        return;
+//    }
+//
+//    const std::string key = location_ + ".quantizer";
+//    std::vector<uint64_t> gpus;
+//
+//    // cache hit
+//    {
+//        int64_t selected = -1;
+//        void* quantizer = nullptr;
+//        for (auto& gpu : gpus) {
+//            auto cache = cache::GpuCacheMgr::GetInstance(gpu);
+//            if (auto quan = cache->GetIndex(key)) {
+//                selected = gpu;
+//                quantizer = quan;
+//            }
+//        }
+//
+//        if (selected != -1) {
+//            // set quantizer into index;
+//            return;
+//        }
+//    }
+//
+//    // cache miss
+//    {
+//        std::vector<int64_t> all_free_mem;
+//        for (auto& gpu : gpus) {
+//            auto cache = cache::GpuCacheMgr::GetInstance(gpu);
+//            auto free_mem = cache->CacheCapacity() - cache->CacheUsage();
+//            all_free_mem.push_back(free_mem);
+//        }
+//
+//        auto max_e = std::max_element(all_free_mem.begin(), all_free_mem.end());
+//        auto best = std::distance(all_free_mem.begin(), max_e);
+//
+//        // load to best device;
+//        // cache quantizer
+//    }
+//
+//    // if index_type == Hybrid
+//
+//    // 1. quantizer in which gpu
+//
+//    // 2.1 which gpu cache best
+//
+//    // 2.2 load to that gpu cache
+//
+//    // set quantizer into index
+}
+
 Status
 ExecutionEngineImpl::AddWithIds(int64_t n, const float* xdata, const int64_t* xids) {
     auto status = index_->Add(n, xdata, xids);
@@ -133,7 +187,7 @@ ExecutionEngineImpl::Serialize() {
 
 Status
 ExecutionEngineImpl::Load(bool to_cache) {
-    index_ = cache::CpuCacheMgr::GetInstance()->GetIndex(location_);
+    index_ = std::static_pointer_cast<VecIndex>(cache::CpuCacheMgr::GetInstance()->GetIndex(location_));
     bool already_in_cache = (index_ != nullptr);
     if (!already_in_cache) {
         try {
@@ -161,7 +215,7 @@ ExecutionEngineImpl::Load(bool to_cache) {
 
 Status
 ExecutionEngineImpl::CopyToGpu(uint64_t device_id) {
-    auto index = cache::GpuCacheMgr::GetInstance(device_id)->GetIndex(location_);
+    auto index = std::static_pointer_cast<VecIndex>(cache::GpuCacheMgr::GetInstance(device_id)->GetIndex(location_));
     bool already_in_cache = (index != nullptr);
     if (already_in_cache) {
         index_ = index;
@@ -189,7 +243,7 @@ ExecutionEngineImpl::CopyToGpu(uint64_t device_id) {
 
 Status
 ExecutionEngineImpl::CopyToCpu() {
-    auto index = cache::CpuCacheMgr::GetInstance()->GetIndex(location_);
+    auto index = std::static_pointer_cast<VecIndex>(cache::CpuCacheMgr::GetInstance()->GetIndex(location_));
     bool already_in_cache = (index != nullptr);
     if (already_in_cache) {
         index_ = index;
@@ -322,7 +376,7 @@ ExecutionEngineImpl::Search(int64_t n, const float* data, int64_t k, int64_t npr
 
 Status
 ExecutionEngineImpl::Cache() {
-    cache::DataObjPtr obj = std::make_shared<cache::DataObj>(index_, PhysicalSize());
+    cache::DataObjPtr obj = std::static_pointer_cast<cache::DataObj>(index_);
     milvus::cache::CpuCacheMgr::GetInstance()->InsertItem(location_, obj);
 
     return Status::OK();
@@ -330,7 +384,7 @@ ExecutionEngineImpl::Cache() {
 
 Status
 ExecutionEngineImpl::GpuCache(uint64_t gpu_id) {
-    cache::DataObjPtr obj = std::make_shared<cache::DataObj>(index_, PhysicalSize());
+    cache::DataObjPtr obj = std::static_pointer_cast<cache::DataObj>(index_);
     milvus::cache::GpuCacheMgr::GetInstance(gpu_id)->InsertItem(location_, obj);
 
     return Status::OK();

cpp/src/db/engine/ExecutionEngineImpl.h

@@ -104,6 +104,9 @@ class ExecutionEngineImpl : public ExecutionEngine {
     VecIndexPtr
     Load(const std::string& location);
 
+    void
+    HybridLoad();
+
  protected:
     VecIndexPtr index_ = nullptr;
     EngineType index_type_;

cpp/src/wrapper/VecIndex.cpp

@@ -34,6 +34,11 @@
 namespace milvus {
 namespace engine {
 
+int64_t
+VecIndex::Size() {
+    return Count() * Dimension() * sizeof(float);
+}
+
 struct FileIOReader {
     std::fstream fs;
     std::string name;

cpp/src/wrapper/VecIndex.h

@@ -20,6 +20,7 @@
 #include <memory>
 #include <string>
 
+#include "cache/DataObj.h"
 #include "knowhere/common/BinarySet.h"
 #include "knowhere/common/Config.h"
 #include "utils/Status.h"
@@ -48,7 +49,7 @@ class VecIndex;
 
 using VecIndexPtr = std::shared_ptr<VecIndex>;
 
-class VecIndex {
+class VecIndex : public cache::DataObj {
  public:
     virtual Status
     BuildAll(const int64_t& nb, const float* xb, const int64_t* ids, const Config& cfg, const int64_t& nt = 0,
@@ -81,6 +82,9 @@ class VecIndex {
     virtual int64_t
     Count() = 0;
 
+    int64_t
+    Size() override;
+
     virtual knowhere::BinarySet
     Serialize() = 0;
 

cpp/unittest/scheduler/test_scheduler.cpp

@@ -141,7 +141,7 @@ insert_dummy_index_into_gpu_cache(uint64_t device_id) {
     mock_index->ntotal_ = 1000;
     engine::VecIndexPtr index(mock_index);
 
-    cache::DataObjPtr obj = std::make_shared<cache::DataObj>(index);
+    cache::DataObjPtr obj = std::static_pointer_cast<cache::DataObj>(index);
 
     cache::GpuCacheMgr::GetInstance(device_id)->InsertItem("location", obj);
 }

cpp/unittest/server/test_cache.cpp

@@ -145,7 +145,7 @@ TEST(CacheTest, CPU_CACHE_TEST) {
     for (uint64_t i = 0; i < item_count; i++) {
         //each vector is 1k byte, total size less than 1G
         ms::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 1000000);
-        ms::cache::DataObjPtr data_obj = std::make_shared<ms::cache::DataObj>(mock_index);
+        ms::cache::DataObjPtr data_obj = std::static_pointer_cast<ms::cache::DataObj>(mock_index);
         cpu_mgr->InsertItem("index_" + std::to_string(i), data_obj);
     }
     ASSERT_LT(cpu_mgr->ItemCount(), g_num);
@@ -169,7 +169,7 @@ TEST(CacheTest, CPU_CACHE_TEST) {
 
         //each vector is 1k byte, total size less than 6G
         ms::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 6000000);
-        ms::cache::DataObjPtr data_obj = std::make_shared<ms::cache::DataObj>(mock_index);
+        ms::cache::DataObjPtr data_obj = std::static_pointer_cast<ms::cache::DataObj>(mock_index);
         cpu_mgr->InsertItem("index_6g", data_obj);
         ASSERT_TRUE(cpu_mgr->ItemExists("index_6g"));
     }
@@ -183,7 +183,7 @@ TEST(CacheTest, GPU_CACHE_TEST) {
     for (int i = 0; i < 20; i++) {
         //each vector is 1k byte
         ms::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 1000);
-        ms::cache::DataObjPtr data_obj = std::make_shared<ms::cache::DataObj>(mock_index);
+        ms::cache::DataObjPtr data_obj = std::static_pointer_cast<ms::cache::DataObj>(mock_index);
         gpu_mgr->InsertItem("index_" + std::to_string(i), data_obj);
     }
 
@@ -196,8 +196,8 @@ TEST(CacheTest, GPU_CACHE_TEST) {
         // TODO(myh): use gpu index to mock
         //each vector is 1k byte, total size less than 2G
         ms::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 2000000);
-        ms::cache::DataObjPtr data_obj = std::make_shared<ms::cache::DataObj>(mock_index);
-        std::cout << data_obj->size() << std::endl;
+        ms::cache::DataObjPtr data_obj = std::static_pointer_cast<ms::cache::DataObj>(mock_index);
+        std::cout << data_obj->Size() << std::endl;
         gpu_mgr->InsertItem("index_" + std::to_string(i), data_obj);
     }
 
@@ -227,7 +227,7 @@ TEST(CacheTest, INVALID_TEST) {
         for (int i = 0; i < 20; i++) {
             //each vector is 1k byte
             ms::engine::VecIndexPtr mock_index = std::make_shared<MockVecIndex>(256, 2);
-            ms::cache::DataObjPtr data_obj = std::make_shared<ms::cache::DataObj>(mock_index);
+            ms::cache::DataObjPtr data_obj = std::static_pointer_cast<ms::cache::DataObj>(mock_index);
             mgr.InsertItem("index_" + std::to_string(i), data_obj);
         }
         ASSERT_EQ(mgr.GetItem("index_0"), nullptr);