From 9926cc02b7ea8f70bc3bca9bed25762817a7eba5 Mon Sep 17 00:00:00 2001
From: Yu Kun <kun.yu@zilliz.com>
Date: Fri, 18 Oct 2019 17:45:29 +0800
Subject: [PATCH] Put union of index_build_device and search resources to
 gpu_pool

Former-commit-id: 92db6ea293737bc83c43ed565d7144d297196b99
---
 core/src/scheduler/SchedInst.cpp              | 25 +++++++++++++------
 .../scheduler/action/PushTaskToNeighbour.cpp  |  4 +--
 2 files changed, 19 insertions(+), 10 deletions(-)

diff --git a/core/src/scheduler/SchedInst.cpp b/core/src/scheduler/SchedInst.cpp
index 4d8fadd3..1bd1ca00 100644
--- a/core/src/scheduler/SchedInst.cpp
+++ b/core/src/scheduler/SchedInst.cpp
@@ -53,26 +53,35 @@ load_simple_config() {
     config.GetResourceConfigSearchResources(pool);
 
     // get resources
-    bool use_cpu_to_compute = false;
-    for (auto& resource : pool) {
-        if (resource == "cpu") {
-            use_cpu_to_compute = true;
-            break;
-        }
-    }
     auto gpu_ids = get_gpu_pool();
 
+    int32_t build_gpu_id;
+    config.GetResourceConfigIndexBuildDevice(build_gpu_id);
+
     // create and connect
     ResMgrInst::GetInstance()->Add(ResourceFactory::Create("disk", "DISK", 0, true, false));
 
     auto io = Connection("io", 500);
-    ResMgrInst::GetInstance()->Add(ResourceFactory::Create("cpu", "CPU", 0, true, use_cpu_to_compute));
+    ResMgrInst::GetInstance()->Add(ResourceFactory::Create("cpu", "CPU", 0, true, true));
     ResMgrInst::GetInstance()->Connect("disk", "cpu", io);
 
     auto pcie = Connection("pcie", 12000);
+    bool find_build_gpu_id = false;
     for (auto& gpu_id : gpu_ids) {
         ResMgrInst::GetInstance()->Add(ResourceFactory::Create(std::to_string(gpu_id), "GPU", gpu_id, true, true));
         ResMgrInst::GetInstance()->Connect("cpu", std::to_string(gpu_id), pcie);
+        if (build_gpu_id == gpu_id) {
+            find_build_gpu_id = true;
+        }
+    }
+
+    if (not find_build_gpu_id) {
+        ResMgrInst::GetInstance()->Add(ResourceFactory::Create(std::to_string(build_gpu_id),
+                                                               "GPU",
+                                                               build_gpu_id,
+                                                               true,
+                                                               true));
+        ResMgrInst::GetInstance()->Connect("cpu", std::to_string(build_gpu_id), pcie);
     }
 }
 
diff --git a/core/src/scheduler/action/PushTaskToNeighbour.cpp b/core/src/scheduler/action/PushTaskToNeighbour.cpp
index 4c815534..c64e81dc 100644
--- a/core/src/scheduler/action/PushTaskToNeighbour.cpp
+++ b/core/src/scheduler/action/PushTaskToNeighbour.cpp
@@ -187,8 +187,8 @@ Action::SpecifiedResourceLabelTaskScheduler(ResourceMgrWPtr res_mgr, ResourcePtr
             Status stat = config.GetResourceConfigIndexBuildDevice(build_index_gpu);
 
             bool find_gpu_res = false;
-            for (uint64_t i = 0; i < compute_resources.size(); ++i) {
-                if (res_mgr.lock()->GetResource(ResourceType::GPU, build_index_gpu) != nullptr) {
+            if (res_mgr.lock()->GetResource(ResourceType::GPU, build_index_gpu) != nullptr) {
+                for (uint64_t i = 0; i < compute_resources.size(); ++i) {
                     if (compute_resources[i]->name() ==
                         res_mgr.lock()->GetResource(ResourceType::GPU, build_index_gpu)->name()) {
                         find_gpu_res = true;
-- 
GitLab