diff --git a/cpp/CHANGELOG.md b/cpp/CHANGELOG.md index 232ea03cfb..4e671387ee 100644 --- a/cpp/CHANGELOG.md +++ b/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 diff --git a/cpp/src/scheduler/action/PushTaskToNeighbour.cpp b/cpp/src/scheduler/action/PushTaskToNeighbour.cpp index 5b9cb0f6ba..f253e4c612 100644 --- a/cpp/src/scheduler/action/PushTaskToNeighbour.cpp +++ b/cpp/src/scheduler/action/PushTaskToNeighbour.cpp @@ -13,7 +13,7 @@ namespace milvus { namespace engine { void -next(std::list neighbours, std::list::iterator &it) { +next(std::list &neighbours, std::list::iterator &it) { it++; if (neighbours.end() == it) { it = neighbours.begin(); diff --git a/cpp/src/scheduler/task/TestTask.cpp b/cpp/src/scheduler/task/TestTask.cpp index a974482e52..527e1b0fb1 100644 --- a/cpp/src/scheduler/task/TestTask.cpp +++ b/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 lock(mutex_); exec_count_++; + done_ = true; +} + +TaskPtr +TestTask::Clone() { + auto ret = std::make_shared(); + ret->load_count_ = load_count_; + ret->exec_count_ = exec_count_; + return ret; +} + +void +TestTask::Wait() { + std::unique_lock lock(mutex_); + cv_.wait(lock, [&] { return done_; }); } } diff --git a/cpp/src/scheduler/task/TestTask.h b/cpp/src/scheduler/task/TestTask.h index 5f49d2e31e..0eb77504c8 100644 --- a/cpp/src/scheduler/task/TestTask.h +++ b/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_; }; diff --git a/cpp/src/wrapper/knowhere/vec_impl.cpp b/cpp/src/wrapper/knowhere/vec_impl.cpp index f6bdd82618..b0fb4c0799 100644 --- a/cpp/src/wrapper/knowhere/vec_impl.cpp +++ b/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(gpu_index, type); + auto new_index = std::make_shared(gpu_index, type); + new_index->dim = dim; + return new_index; } // TODO(linxj): rename copytocpu => copygputocpu diff --git a/cpp/unittest/scheduler/normal_test.cpp b/cpp/unittest/scheduler/normal_test.cpp index 27d16a1b6b..e3e791c02c 100644 --- a/cpp/unittest/scheduler/normal_test.cpp +++ b/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(); 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> 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(pcpu->task_table()[i]->task); + ASSERT_EQ(task->load_count_, 1); + ASSERT_EQ(task->exec_count_, 1); } }