From 1dc7b321a45a178813350ea7c87d2d0ce57103c8 Mon Sep 17 00:00:00 2001
From: fishpenguin <kun.yu@zilliz.com>
Date: Thu, 31 Oct 2019 16:41:48 +0800
Subject: [PATCH] Read config once in optimizer

Former-commit-id: 8fee2a3cfa20cdc8b51e86867b7b3eb71dd2e759
---
 core/src/scheduler/SchedInst.h               | 14 ++++++++-
 core/src/scheduler/optimizer/OnlyCPUPass.cpp | 14 ++++-----
 core/src/scheduler/optimizer/OnlyGPUPass.cpp | 32 ++++++--------------
 core/src/scheduler/optimizer/OnlyGPUPass.h   |  3 +-
 4 files changed, 32 insertions(+), 31 deletions(-)

diff --git a/core/src/scheduler/SchedInst.h b/core/src/scheduler/SchedInst.h
index e758f378..a3048069 100644
--- a/core/src/scheduler/SchedInst.h
+++ b/core/src/scheduler/SchedInst.h
@@ -26,9 +26,11 @@
 #include "optimizer/OnlyCPUPass.h"
 #include "optimizer/OnlyGPUPass.h"
 #include "optimizer/Optimizer.h"
+#include "server/Config.h"
 
 #include <memory>
 #include <mutex>
+#include <string>
 #include <vector>
 
 namespace milvus {
@@ -95,11 +97,21 @@ class OptimizerInst {
         if (instance == nullptr) {
             std::lock_guard<std::mutex> lock(mutex_);
             if (instance == nullptr) {
+                server::Config& config = server::Config::GetInstance();
+                std::vector<std::string> search_resources;
+                bool has_cpu = false;
+                config.GetResourceConfigSearchResources(search_resources);
+                for (auto& resource : search_resources) {
+                    if (resource == "cpu") {
+                        has_cpu = true;
+                    }
+                }
+
                 std::vector<PassPtr> pass_list;
                 pass_list.push_back(std::make_shared<LargeSQ8HPass>());
                 pass_list.push_back(std::make_shared<HybridPass>());
                 pass_list.push_back(std::make_shared<OnlyCPUPass>());
-                pass_list.push_back(std::make_shared<OnlyGPUPass>());
+                pass_list.push_back(std::make_shared<OnlyGPUPass>(has_cpu));
                 instance = std::make_shared<Optimizer>(pass_list);
             }
         }
diff --git a/core/src/scheduler/optimizer/OnlyCPUPass.cpp b/core/src/scheduler/optimizer/OnlyCPUPass.cpp
index 2651a6e1..238a91a8 100644
--- a/core/src/scheduler/optimizer/OnlyCPUPass.cpp
+++ b/core/src/scheduler/optimizer/OnlyCPUPass.cpp
@@ -35,13 +35,13 @@ OnlyCPUPass::Run(const TaskPtr& task) {
     }
 
     auto gpu_id = get_gpu_pool();
-    if (gpu_id.empty()) {
-        ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
-        auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
-        task->label() = label;
-        return true;
-    }
-    return false;
+    if (not gpu_id.empty())
+        return false;
+
+    ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource("cpu");
+    auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
+    task->label() = label;
+    return true;
 }
 
 }  // namespace scheduler
diff --git a/core/src/scheduler/optimizer/OnlyGPUPass.cpp b/core/src/scheduler/optimizer/OnlyGPUPass.cpp
index f39ca1a0..3fcda0e8 100644
--- a/core/src/scheduler/optimizer/OnlyGPUPass.cpp
+++ b/core/src/scheduler/optimizer/OnlyGPUPass.cpp
@@ -20,14 +20,16 @@
 #include "scheduler/Utils.h"
 #include "scheduler/task/SearchTask.h"
 #include "scheduler/tasklabel/SpecResLabel.h"
-#include "server/Config.h"
 
 namespace milvus {
 namespace scheduler {
 
+OnlyGPUPass::OnlyGPUPass(bool has_cpu) : has_cpu_(has_cpu) {
+}
+
 bool
 OnlyGPUPass::Run(const TaskPtr& task) {
-    if (task->Type() != TaskType::SearchTask)
+    if (task->Type() != TaskType::SearchTask || has_cpu_)
         return false;
 
     auto search_task = std::static_pointer_cast<XSearchTask>(task);
@@ -36,29 +38,15 @@ OnlyGPUPass::Run(const TaskPtr& task) {
         return false;
     }
 
-    server::Config& config = server::Config::GetInstance();
-    std::vector<std::string> search_resources;
-    config.GetResourceConfigSearchResources(search_resources);
-    for (auto& resource : search_resources) {
-        if (resource == "cpu") {
-            return false;
-        }
-    }
-
     auto gpu_id = get_gpu_pool();
-    if (!gpu_id.empty()) {
-        ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, gpu_id[specified_gpu_id_]);
-        auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
-        task->label() = label;
-    } else {
+    if (gpu_id.empty())
         return false;
-    }
 
-    if (specified_gpu_id_ < gpu_id.size() - 1) {
-        ++specified_gpu_id_;
-    } else {
-        specified_gpu_id_ = 0;
-    }
+    ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, gpu_id[specified_gpu_id_]);
+    auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));
+    task->label() = label;
+
+    specified_gpu_id_ = specified_gpu_id_++ % gpu_id.size();
     return true;
 }
 
diff --git a/core/src/scheduler/optimizer/OnlyGPUPass.h b/core/src/scheduler/optimizer/OnlyGPUPass.h
index 75a5f9e4..10d909d3 100644
--- a/core/src/scheduler/optimizer/OnlyGPUPass.h
+++ b/core/src/scheduler/optimizer/OnlyGPUPass.h
@@ -34,7 +34,7 @@ namespace scheduler {
 
 class OnlyGPUPass : public Pass {
  public:
-    OnlyGPUPass() = default;
+    explicit OnlyGPUPass(bool has_cpu);
 
  public:
     bool
@@ -42,6 +42,7 @@ class OnlyGPUPass : public Pass {
 
  private:
     uint64_t specified_gpu_id_ = 0;
+    bool has_cpu_ = false;
 };
 
 using OnlyGPUPassPtr = std::shared_ptr<OnlyGPUPass>;
-- 
GitLab