MS-394 Update scheduler unittest

Former-commit-id: b4f0fe18d00b17c47a39e2a06cdd1119ae00d7fc

cpp/CHANGELOG.md

@@ -37,6 +37,7 @@ Please mark all change in change log and use the ticket from JIRA.
 - MS-389 - Add clone interface in Task
 - MS-390 - Update resource construct function
 - MS-391 - Add PushTaskToNeighbourHasExecutor action
+- MS-394 - Update scheduler unittest
 
 ## New Feature
 - MS-343 - Implement ResourceMgr

cpp/src/scheduler/action/PushTaskToNeighbour.cpp

@@ -13,7 +13,7 @@ namespace milvus {
 namespace engine {
 
 void
-next(std::list<ResourcePtr> neighbours, std::list<ResourcePtr>::iterator &it) {
+next(std::list<ResourcePtr> &neighbours, std::list<ResourcePtr>::iterator &it) {
     it++;
     if (neighbours.end() == it) {
         it = neighbours.begin();

cpp/src/scheduler/task/TestTask.cpp

@@ -6,6 +6,7 @@
 
 #include "TestTask.h"
 
+
 namespace zilliz {
 namespace milvus {
 namespace engine {
@@ -17,7 +18,23 @@ TestTask::Load(LoadType type, uint8_t device_id) {
 
 void
 TestTask::Execute() {
+    std::lock_guard<std::mutex> lock(mutex_);
     exec_count_++;
+    done_ = true;
+}
+
+TaskPtr
+TestTask::Clone() {
+    auto ret = std::make_shared<TestTask>();
+    ret->load_count_ = load_count_;
+    ret->exec_count_ = exec_count_;
+    return ret;
+}
+
+void
+TestTask::Wait() {
+    std::unique_lock<std::mutex> lock(mutex_);
+    cv_.wait(lock, [&] { return done_; });
 }
 
 }

cpp/src/scheduler/task/TestTask.h

@@ -23,9 +23,19 @@ public:
     void
     Execute() override;
 
+    TaskPtr
+    Clone() override;
+
+    void
+    Wait();
+
 public:
-    uint64_t load_count_;
-    uint64_t exec_count_;
+    uint64_t load_count_ = 0;
+    uint64_t exec_count_ = 0;
+
+    bool done_ = false;
+    std::mutex mutex_;
+    std::condition_variable cv_;
 };
 
 

cpp/src/wrapper/knowhere/vec_impl.cpp

@@ -144,7 +144,9 @@ VecIndexPtr VecIndexImpl::CopyToGpu(const int64_t &device_id, const Config &cfg)
 
     // TODO(linxj): update type
     auto gpu_index = zilliz::knowhere::CopyCpuToGpu(index_, device_id, cfg);
-    return std::make_shared<VecIndexImpl>(gpu_index, type);
+    auto new_index = std::make_shared<VecIndexImpl>(gpu_index, type);
+    new_index->dim = dim;
+    return new_index;
 }
 
 // TODO(linxj): rename copytocpu => copygputocpu

cpp/unittest/scheduler/normal_test.cpp

@@ -13,10 +13,10 @@ TEST(normal_test, test1) {
     // 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"));
+    auto disk = res_mgr->Add(ResourceFactory::Create("disk", "ssd", true, false));
     auto cpu = res_mgr->Add(ResourceFactory::Create("cpu"));
-    auto gpu1 = res_mgr->Add(ResourceFactory::Create("gpu"));
-    auto gpu2 = res_mgr->Add(ResourceFactory::Create("gpu"));
+    auto gpu1 = res_mgr->Add(ResourceFactory::Create("gpu", "gpu0", false, false));
+    auto gpu2 = res_mgr->Add(ResourceFactory::Create("gpu", "gpu2", false, false));
 
     auto IO = Connection("IO", 500.0);
     auto PCIE = Connection("IO", 11000.0);
@@ -30,7 +30,7 @@ TEST(normal_test, test1) {
     auto scheduler = SchedInst::GetInstance();
     scheduler->Start();
 
-    const uint64_t NUM_TASK = 100;
+    const uint64_t NUM_TASK = 1000;
     std::vector<std::shared_ptr<TestTask>> tasks;
     for (uint64_t i = 0; i < NUM_TASK; ++i) {
         if (auto observe = disk.lock()) {
@@ -45,8 +45,10 @@ TEST(normal_test, test1) {
     scheduler->Stop();
     res_mgr->Stop();
 
-    for (uint64_t i = 0 ; i < NUM_TASK; ++i) {
-        ASSERT_EQ(tasks[i]->load_count_, 1);
-        ASSERT_EQ(tasks[i]->exec_count_, 1);
+    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);
     }
 }