Fix memory leak in IVF indexes (#4325)

* fix memory lead in IVF indexes Signed-off-by: shengjun.li <[email protected]> * [skip ci] modify change log Signed-off-by: shengjun.li <[email protected]>
milvus-io · Dec 1, 2020 · bf5fdc3 · bf5fdc3
1 parent 99245c0
commit bf5fdc3
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 24 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,18 +4,19 @@ Please mark all change in change log and use the issue from GitHub
 
 # Milvus 0.10.4 (TBD)
 ## Bug
--   \#3626 Fix server crash when searching with IVF_PQ on GPU.
--   \#3903 The performance of IVF_SQ8H in 0.10.3 is degraded.
--   \#3906 Change DeleteTask state when it is loaded to avoid server crash.
--   \#4012 Milvus hangs when continually creating and dropping partitions.
--   \#4075 Improve performance for create large amount of partitions
--   \#4174 Search out of memory: CPU2GPU1 with index flat
+-   \#3626 Fix server crash when searching with IVF_PQ on GPU
+-   \#3903 The performance of IVF_SQ8H in 0.10.3 is degraded
+-   \#3906 Fix the delete task state to avoid server crash
+-   \#4012 Milvus hangs when continually creating and dropping partitions
+-   \#4174 Fix out of memory caused by too many data loaded to GPU
+-   \#4318 Fix memory leak in IVF indexes
 
 ## Feature
--   \#3773 Support IVF_PQ to run on FPGA.
+-   \#3773 Support IVF_PQ to run on FPGA
 
 ## Improvement
--   \#3775 Improve search performance in the case that no item deleted.
+-   \#3775 Improve search performance in the case that no item deleted
+-   \#4075 Improve performance for create large amount of partitions
 
 ## Task
 

diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexBinaryIVF.cpp
@@ -157,8 +157,9 @@ BinaryIVF::Train(const DatasetPtr& dataset_ptr, const Config& config) {
     faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
     faiss::IndexBinary* coarse_quantizer = new faiss::IndexBinaryFlat(dim, metric_type);
     auto index = std::make_shared<faiss::IndexBinaryIVF>(coarse_quantizer, dim, nlist, metric_type);
-    index->train(rows, (uint8_t*)p_data);
-    index->add_with_ids(rows, (uint8_t*)p_data, p_ids);
+    index->own_fields = true;
+    index->train(rows, static_cast<const uint8_t*>(p_data));
+    index->add_with_ids(rows, static_cast<const uint8_t*>(p_data), p_ids);
     index_ = index;
 }
 

diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVF.cpp
@@ -68,11 +68,13 @@ void
 IVF::Train(const DatasetPtr& dataset_ptr, const Config& config) {
     GETTENSOR(dataset_ptr)
 
+    int64_t nlist = config[IndexParams::nlist].get<int64_t>();
     faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
     faiss::Index* coarse_quantizer = new faiss::IndexFlat(dim, metric_type);
-    int64_t nlist = config[IndexParams::nlist].get<int64_t>();
-    index_ = std::shared_ptr<faiss::Index>(new faiss::IndexIVFFlat(coarse_quantizer, dim, nlist, metric_type));
-    index_->train(rows, (float*)p_data);
+    auto index = std::make_shared<faiss::IndexIVFFlat>(coarse_quantizer, dim, nlist, metric_type);
+    index->own_fields = true;
+    index->train(rows, reinterpret_cast<const float*>(p_data));
+    index_ = index;
 }
 
 void

diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFPQ.cpp
@@ -38,11 +38,12 @@ IVFPQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
 
     faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
     faiss::Index* coarse_quantizer = new faiss::IndexFlat(dim, metric_type);
-    index_ = std::shared_ptr<faiss::Index>(new faiss::IndexIVFPQ(
-        coarse_quantizer, dim, config[IndexParams::nlist].get<int64_t>(), config[IndexParams::m].get<int64_t>(),
-        config[IndexParams::nbits].get<int64_t>(), metric_type));
-
-    index_->train(rows, (float*)p_data);
+    auto index = std::make_shared<faiss::IndexIVFPQ>(coarse_quantizer, dim, config[IndexParams::nlist].get<int64_t>(),
+                                                     config[IndexParams::m].get<int64_t>(),
+                                                     config[IndexParams::nbits].get<int64_t>(), metric_type);
+    index->own_fields = true;
+    index->train(rows, reinterpret_cast<const float*>(p_data));
+    index_ = index;
 }
 
 VecIndexPtr

diff --git a/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp b/core/src/index/knowhere/knowhere/index/vector_index/IndexIVFSQ.cpp
@@ -16,6 +16,7 @@
 #include <faiss/gpu/GpuAutoTune.h>
 #include <faiss/gpu/GpuCloner.h>
 #endif
+#include <faiss/IndexFlat.h>
 #include <faiss/IndexScalarQuantizer.h>
 #include <faiss/clone_index.h>
 #include <faiss/index_factory.h>
@@ -36,12 +37,13 @@ void
 IVFSQ::Train(const DatasetPtr& dataset_ptr, const Config& config) {
     GETTENSOR(dataset_ptr)
 
-    std::stringstream index_type;
-    index_type << "IVF" << config[IndexParams::nlist] << ","
-               << "SQ" << config[IndexParams::nbits];
-    index_ = std::shared_ptr<faiss::Index>(
-        faiss::index_factory(dim, index_type.str().c_str(), GetMetricType(config[Metric::TYPE].get<std::string>())));
-    index_->train(rows, (float*)p_data);
+    faiss::MetricType metric_type = GetMetricType(config[Metric::TYPE].get<std::string>());
+    faiss::Index* coarse_quantizer = new faiss::IndexFlat(dim, metric_type);
+    auto index = std::make_shared<faiss::IndexIVFScalarQuantizer>(
+        coarse_quantizer, dim, config[IndexParams::nlist].get<int64_t>(), faiss::QuantizerType::QT_8bit, metric_type);
+    index->own_fields = true;
+    index->train(rows, reinterpret_cast<const float*>(p_data));
+    index_ = index;
 }
 
 VecIndexPtr