MS-508 Update normal_test in scheduler

Former-commit-id: 343dae4a96cac46259094de9c1d025de3b3db53e

cpp/CHANGELOG.md

@@ -89,6 +89,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-488 - Improve code format in scheduler
 - MS-502 - Update tasktable_test in scheduler
 - MS-504 - Update node_test in scheduler
+- MS-508 - Update normal_test in scheduler
 
 ## New Feature
 - MS-343 - Implement ResourceMgr

cpp/src/scheduler/Scheduler.cpp

@@ -108,6 +108,7 @@ void
 Scheduler::OnFinishTask(const EventPtr &event) {
 }
 
+// TODO: refactor the function
 void
 Scheduler::OnLoadCompleted(const EventPtr &event) {
     auto load_completed_event = std::static_pointer_cast<LoadCompletedEvent>(event);
@@ -120,9 +121,12 @@ Scheduler::OnLoadCompleted(const EventPtr &event) {
                 if (not resource->HasExecutor() && load_completed_event->task_table_item_->Move()) {
                     auto task = load_completed_event->task_table_item_->task;
                     auto search_task = std::static_pointer_cast<XSearchTask>(task);
-                    auto location = search_task->index_engine_->GetLocation();
                     bool moved = false;
 
+                    // to support test task, REFACTOR
+                    if (auto index_engine = search_task->index_engine_) {
+                        auto location = index_engine->GetLocation();
+
                         for (auto i = 0; i < res_mgr_.lock()->GetNumGpuResource(); ++i) {
                             auto index = zilliz::milvus::cache::GpuCacheMgr::GetInstance(i)->GetIndex(location);
                             if (index != nullptr) {
@@ -132,6 +136,8 @@ Scheduler::OnLoadCompleted(const EventPtr &event) {
                                 break;
                             }
                         }
+                    }
+
                     if (not moved) {
                         Action::PushTaskToNeighbourRandomly(task, resource);
                     }

cpp/src/scheduler/Scheduler.h

@@ -21,6 +21,7 @@ namespace milvus {
 namespace engine {
 
 
+// TODO: refactor, not friendly to unittest, logical in framework code
 class Scheduler {
 public:
     explicit

cpp/src/scheduler/task/SearchTask.cpp

@@ -83,6 +83,7 @@ CollectFileMetrics(int file_type, size_t file_size) {
 
 XSearchTask::XSearchTask(TableFileSchemaPtr file)
     : Task(TaskType::SearchTask), file_(file) {
+    if (file_) {
         index_engine_ = EngineFactory::Build(file_->dimension_,
                                              file_->location_,
                                              (EngineType) file_->engine_type_,
@@ -90,6 +91,8 @@ XSearchTask::XSearchTask(TableFileSchemaPtr file)
                                              file_->nlist_);
     }
 
+}
+
 void
 XSearchTask::Load(LoadType type, uint8_t device_id) {
     server::TimeRecorder rc("");

cpp/src/scheduler/task/SearchTask.h

@@ -12,6 +12,7 @@ namespace zilliz {
 namespace milvus {
 namespace engine {
 
+// TODO: rewrite
 class XSearchTask : public Task {
 public:
     explicit

cpp/src/scheduler/task/Task.h

@@ -34,6 +34,7 @@ class Task;
 
 using TaskPtr = std::shared_ptr<Task>;
 
+// TODO: re-design
 class Task {
 public:
     explicit

cpp/src/scheduler/task/TestTask.cpp

@@ -22,10 +22,13 @@ TestTask::Load(LoadType type, uint8_t device_id) {
 
 void
 TestTask::Execute() {
+    {
         std::lock_guard<std::mutex> lock(mutex_);
         exec_count_++;
         done_ = true;
     }
+    cv_.notify_one();
+}
 
 void
 TestTask::Wait() {

cpp/unittest/scheduler/normal_test.cpp

@@ -2,6 +2,7 @@
 #include "scheduler/ResourceMgr.h"
 #include "scheduler/Scheduler.h"
 #include "scheduler/task/TestTask.h"
+#include "scheduler/tasklabel/DefaultLabel.h"
 #include "scheduler/SchedInst.h"
 #include "utils/Log.h"
 #include <gtest/gtest.h>
@@ -9,48 +10,44 @@
 
 using namespace zilliz::milvus::engine;
 
-TEST(normal_test, test1) {
+
+TEST(normal_test, inst_test) {
     // ResourceMgr only compose resources, provide unified event
-//    auto res_mgr = std::make_shared<ResourceMgr>();
     auto res_mgr = ResMgrInst::GetInstance();
-    auto disk = res_mgr->Add(ResourceFactory::Create("disk", "ssd", true, false));
-    auto cpu = res_mgr->Add(ResourceFactory::Create("cpu", "CPU", 0));
-    auto gpu1 = res_mgr->Add(ResourceFactory::Create("gpu", "gpu0", false, false));
-    auto gpu2 = res_mgr->Add(ResourceFactory::Create("gpu", "gpu2", false, false));
+
+    res_mgr->Add(ResourceFactory::Create("disk", "DISK", 0, true, false));
+    res_mgr->Add(ResourceFactory::Create("cpu", "CPU", 0, true, true));
 
     auto IO = Connection("IO", 500.0);
-    auto PCIE = Connection("IO", 11000.0);
-    res_mgr->Connect(disk, cpu, IO);
-    res_mgr->Connect(cpu, gpu1, PCIE);
-    res_mgr->Connect(cpu, gpu2, PCIE);
+    res_mgr->Connect("disk", "cpu", IO);
+
+    auto scheduler = SchedInst::GetInstance();
 
     res_mgr->Start();
-
-//    auto scheduler = new Scheduler(res_mgr);
-    auto scheduler = SchedInst::GetInstance();
     scheduler->Start();
 
     const uint64_t NUM_TASK = 1000;
     std::vector<std::shared_ptr<TestTask>> tasks;
     TableFileSchemaPtr dummy = nullptr;
 
+    auto disks = res_mgr->GetDiskResources();
+    ASSERT_FALSE(disks.empty());
+    if (auto observe = disks[0].lock()) {
         for (uint64_t i = 0; i < NUM_TASK; ++i) {
-        if (auto observe = disk.lock()) {
             auto task = std::make_shared<TestTask>(dummy);
+            task->label() = std::make_shared<DefaultLabel>();
             tasks.push_back(task);
             observe->task_table().Put(task);
         }
     }
 
-    sleep(1);
+    for (auto &task : tasks) {
+        task->Wait();
+        ASSERT_EQ(task->load_count_, 1);
+        ASSERT_EQ(task->exec_count_, 1);
+    }
 
     scheduler->Stop();
     res_mgr->Stop();
 
-    auto pcpu = cpu.lock();
-    for (uint64_t i = 0; i < NUM_TASK; ++i) {
-        auto task = std::static_pointer_cast<TestTask>(pcpu->task_table()[i]->task);
-        ASSERT_EQ(task->load_count_, 1);
-        ASSERT_EQ(task->exec_count_, 1);
-    }
 }