Commit 882c7bdf authored by jinhai's avatar jinhai
Browse files

Merge branch 'branch-0.5.0-yk' into 'branch-0.5.0' 

MS-603 Add BuildIndex to scheduler

See merge request megasearch/milvus!676

Former-commit-id: 5ce5e03956aa7279530e112f0595de4653a047a1
parents 076d2e9b 22371d4c

cpp/src/db/DBImpl.cpp

+78 −61
Original line number Diff line number Diff line
@@ -28,6 +28,7 @@ 
#include "scheduler/SchedInst.h"

#include "scheduler/job/DeleteJob.h"

#include "scheduler/job/SearchJob.h"

#include "scheduler/job/BuildIndexJob.h"

#include "utils/Log.h"

#include "utils/TimeRecorder.h"
@@ -898,18 +899,34 @@ DBImpl::BackgroundBuildIndex() { 
    meta::TableFilesSchema to_index_files;

    meta_ptr_->FilesToIndex(to_index_files);

    Status status;

    for (auto& file : to_index_files) {

        status = BuildIndex(file);

        if (!status.ok()) {

            ENGINE_LOG_ERROR << "Building index for " << file.id_ << " failed: " << status.ToString();

        }



        if (shutting_down_.load(std::memory_order_acquire)) {

            ENGINE_LOG_DEBUG << "Server will shutdown, skip build index action";

            break;

    scheduler::BuildIndexJobPtr

        job = std::make_shared<scheduler::BuildIndexJob>(0, meta_ptr_, options_);



    // step 2: put build index task to scheduler

    for (auto &file : to_index_files) {

        scheduler::TableFileSchemaPtr file_ptr = std::make_shared<meta::TableFileSchema>(file);

        job->AddToIndexFiles(file_ptr);

    }

    scheduler::JobMgrInst::GetInstance()->Put(job);

    job->WaitBuildIndexFinish();

    if (!job->GetStatus().ok()) {

        Status status = job->GetStatus();

        ENGINE_LOG_ERROR << "Building index failed: " << status.ToString();

    }



//    for (auto &file : to_index_files) {

//        status = BuildIndex(file);

//        if (!status.ok()) {

//            ENGINE_LOG_ERROR << "Building index for " << file.id_ << " failed: " << status.ToString();

//        }

//

//        if (shutting_down_.load(std::memory_order_acquire)) {

//            ENGINE_LOG_DEBUG << "Server will shutdown, skip build index action";

//            break;

//        }

//    }



    ENGINE_LOG_TRACE << "Background build index thread exit";

}

cpp/src/db/engine/ExecutionEngine.h

+3 −0
Original line number Diff line number Diff line
@@ -66,6 +66,9 @@ class ExecutionEngine { 
    virtual Status

    CopyToGpu(uint64_t device_id) = 0;



    virtual Status

    CopyToIndexFileToGpu(uint64_t device_id) = 0;



    virtual Status

    CopyToCpu() = 0;

cpp/src/db/engine/ExecutionEngineImpl.cpp

+11 −0
Original line number Diff line number Diff line
@@ -187,6 +187,17 @@ ExecutionEngineImpl::CopyToGpu(uint64_t device_id) { 
    return Status::OK();

}



Status

ExecutionEngineImpl::CopyToIndexFileToGpu(uint64_t device_id) {

    auto index = cache::GpuCacheMgr::GetInstance(device_id)->GetIndex(location_);

    bool already_in_cache = (index != nullptr);

    if (!already_in_cache) {

        cache::DataObjPtr obj = std::make_shared<cache::DataObj>(nullptr, PhysicalSize());

        milvus::cache::GpuCacheMgr::GetInstance(device_id)->InsertItem(location_, obj);

    }

    return Status::OK();

}



Status

ExecutionEngineImpl::CopyToCpu() {

    auto index = cache::CpuCacheMgr::GetInstance()->GetIndex(location_);

cpp/src/db/engine/ExecutionEngineImpl.h

+3 −0
Original line number Diff line number Diff line
@@ -58,6 +58,9 @@ class ExecutionEngineImpl : public ExecutionEngine { 
    Status

    CopyToGpu(uint64_t device_id) override;



    Status

    CopyToIndexFileToGpu(uint64_t device_id) override;



    Status

    CopyToCpu() override;

cpp/src/scheduler/TaskCreator.cpp

+24 −3
Original line number Diff line number Diff line
@@ -15,9 +15,12 @@ 
// specific language governing permissions and limitations

// under the License.



#include <src/scheduler/tasklabel/SpecResLabel.h>

#include "scheduler/TaskCreator.h"

#include "scheduler/tasklabel/BroadcastLabel.h"

#include "tasklabel/DefaultLabel.h"

#include "SchedInst.h"





namespace milvus {

namespace scheduler {
@@ -31,6 +34,9 @@ TaskCreator::Create(const JobPtr& job) { 
        case JobType::DELETE: {

            return Create(std::static_pointer_cast<DeleteJob>(job));

        }

        case JobType::BUILD: {

            return Create(std::static_pointer_cast<BuildIndexJob>(job));

        }

        default: {

            // TODO(wxyu): error

            return std::vector<TaskPtr>();
@@ -62,5 +68,20 @@ TaskCreator::Create(const DeleteJobPtr& job) { 
    return tasks;

}



std::vector<TaskPtr>

TaskCreator::Create(const BuildIndexJobPtr &job) {

    std::vector<TaskPtr> tasks;

    //TODO(yukun): remove "disk" hardcode here

    ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource("disk");



    for (auto &to_index_file : job->to_index_files()) {

        auto label = std::make_shared<SpecResLabel>(std::weak_ptr<Resource>(res_ptr));

        auto task = std::make_shared<XBuildIndexTask>(to_index_file.second, label);

        task->job_ = job;

        tasks.emplace_back(task);

    }

    return tasks;

}



}  // namespace scheduler

}  // namespace milvus