From 6fba3cda43adddb6370c11fd08aaaaef953582c1 Mon Sep 17 00:00:00 2001
From: "xj.lin" <xiaojun.lin@zilliz.com>
Date: Fri, 6 Sep 2019 15:55:23 +0800
Subject: [PATCH] MS-493: 1. FIX return cuda error when system exit 2. enable
 knowhere unittest

Former-commit-id: 6081d0da211e07fd74e5b1607518da879a09eb93
---
 .../include/knowhere/index/vector_index/gpu_ivf.h  |  7 ++++++-
 .../src/knowhere/index/vector_index/gpu_ivf.cpp    | 14 +++++++++++++-
 cpp/src/core/test/test_idmap.cpp                   |  5 +++++
 cpp/src/core/test/test_ivf.cpp                     |  6 ++++--
 cpp/src/server/Server.cpp                          |  2 ++
 cpp/unittest/CMakeLists.txt                        |  2 +-
 cpp/unittest/knowhere/knowhere_test.cpp            |  3 +++
 7 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h b/cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h
index 22c43292..3a62addf 100644
--- a/cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h
+++ b/cpp/src/core/include/knowhere/index/vector_index/gpu_ivf.h
@@ -9,7 +9,7 @@ namespace zilliz {
 namespace knowhere {
 
 struct Resource {
-    Resource(std::shared_ptr<faiss::gpu::StandardGpuResources> &r): faiss_res(r) {
+    explicit Resource(std::shared_ptr<faiss::gpu::StandardGpuResources> &r): faiss_res(r) {
         static int64_t global_id = 0;
         id = global_id++;
     }
@@ -32,6 +32,11 @@ class FaissGpuResourceMgr {
     static FaissGpuResourceMgr &
     GetInstance();
 
+    // Free gpu resource, avoid cudaGetDevice error when deallocate.
+    // this func should be invoke before main return
+    void
+    Free();
+
     void
     AllocateTempMem(ResPtr &resource, const int64_t& device_id, const int64_t& size);
 
diff --git a/cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp b/cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
index 26fcd6a7..0546674e 100644
--- a/cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
+++ b/cpp/src/core/src/knowhere/index/vector_index/gpu_ivf.cpp
@@ -282,7 +282,7 @@ void FaissGpuResourceMgr::InitResource() {
     for(auto& device : devices_params_) {
         auto& resource_vec = idle_[device.first];
 
-        for (int i = 0; i < device.second.resource_num; ++i) {
+        for (int64_t i = 0; i < device.second.resource_num; ++i) {
             auto res = std::make_shared<faiss::gpu::StandardGpuResources>();
 
             // TODO(linxj): enable set pinned memory
@@ -351,6 +351,18 @@ void FaissGpuResourceMgr::MoveToIdle(const int64_t &device_id, const ResPtr &res
     idle_[device_id].insert(it, res);
 }
 
+void FaissGpuResourceMgr::Free() {
+    for (auto &item : in_use_) {
+        auto& res_vec = item.second;
+        res_vec.clear();
+    }
+    for (auto &item : idle_) {
+        auto& res_vec = item.second;
+        res_vec.clear();
+    }
+    is_init = false;
+}
+
 void GPUIndex::SetGpuDevice(const int &gpu_id) {
     gpu_id_ = gpu_id;
 }
diff --git a/cpp/src/core/test/test_idmap.cpp b/cpp/src/core/test/test_idmap.cpp
index 25f8f762..90546429 100644
--- a/cpp/src/core/test/test_idmap.cpp
+++ b/cpp/src/core/test/test_idmap.cpp
@@ -26,6 +26,11 @@ class IDMAPTest : public DataGen, public ::testing::Test {
         Init_with_default();
         index_ = std::make_shared<IDMAP>();
     }
+
+    void TearDown() override {
+        FaissGpuResourceMgr::GetInstance().Free();
+    }
+
  protected:
     IDMAPPtr index_ = nullptr;
 };
diff --git a/cpp/src/core/test/test_ivf.cpp b/cpp/src/core/test/test_ivf.cpp
index 23847f7a..9ebcf767 100644
--- a/cpp/src/core/test/test_ivf.cpp
+++ b/cpp/src/core/test/test_ivf.cpp
@@ -7,13 +7,11 @@
 #include <gtest/gtest.h>
 
 #include <iostream>
-#include <sstream>
 #include <thread>
 
 #include <faiss/AutoTune.h>
 #include <faiss/gpu/GpuAutoTune.h>
 #include <faiss/gpu/GpuIndexIVFFlat.h>
-#include <faiss/gpu/GpuClonerOptions.h>
 
 #include "knowhere/index/vector_index/gpu_ivf.h"
 #include "knowhere/index/vector_index/ivf.h"
@@ -58,6 +56,9 @@ class IVFTest
         index_ = IndexFactory(index_type);
         FaissGpuResourceMgr::GetInstance().InitDevice(device_id, 1024*1024*200, 1024*1024*300, 2);
     }
+    void TearDown() override {
+        FaissGpuResourceMgr::GetInstance().Free();
+    }
 
  protected:
     std::string index_type;
@@ -369,6 +370,7 @@ class GPURESTEST
     void TearDown() override {
         delete ids;
         delete dis;
+        FaissGpuResourceMgr::GetInstance().Free();
     }
 
  protected:
diff --git a/cpp/src/server/Server.cpp b/cpp/src/server/Server.cpp
index 2df42791..955c55d9 100644
--- a/cpp/src/server/Server.cpp
+++ b/cpp/src/server/Server.cpp
@@ -19,6 +19,7 @@
 #include <unistd.h>
 #include <string.h>
 #include <src/scheduler/SchedInst.h>
+#include "knowhere/index/vector_index/gpu_ivf.h"
 
 #include "metrics/Metrics.h"
 #include "DBWrapper.h"
@@ -232,6 +233,7 @@ Server::StopService() {
     grpc::GrpcMilvusServer::StopService();
     DBWrapper::GetInstance().StopService();
     engine::StopSchedulerService();
+    knowhere::FaissGpuResourceMgr::GetInstance().Free(); // free gpu resource.
 }
 
 }
diff --git a/cpp/unittest/CMakeLists.txt b/cpp/unittest/CMakeLists.txt
index ff78178d..ac666c86 100644
--- a/cpp/unittest/CMakeLists.txt
+++ b/cpp/unittest/CMakeLists.txt
@@ -40,7 +40,7 @@ set(unittest_libs
 
 add_subdirectory(server)
 add_subdirectory(db)
-#add_subdirectory(knowhere)
+add_subdirectory(knowhere)
 add_subdirectory(metrics)
 #add_subdirectory(scheduler)
 #add_subdirectory(storage)
\ No newline at end of file
diff --git a/cpp/unittest/knowhere/knowhere_test.cpp b/cpp/unittest/knowhere/knowhere_test.cpp
index 8d252427..2bd4bddc 100644
--- a/cpp/unittest/knowhere/knowhere_test.cpp
+++ b/cpp/unittest/knowhere/knowhere_test.cpp
@@ -40,6 +40,9 @@ class KnowhereWrapperTest
 
         index_ = GetVecIndexFactory(index_type);
     }
+    void TearDown() override {
+        zilliz::knowhere::FaissGpuResourceMgr::GetInstance().Free();
+    }
 
     void AssertResult(const std::vector<long> &ids, const std::vector<float> &dis) {
         EXPECT_EQ(ids.size(), nq * k);
-- 
GitLab