Commit 461e15ab authored by JinHai-CN's avatar JinHai-CN
Browse files

Merge remote-tracking branch 'main/0.5.1' into 0.5.1 


Former-commit-id: bd150ad4d45fda689015d7ce38221589d9a98bcf
parents 38189735 8bd0e32a

CHANGELOG.md

+1 −0
Original line number Diff line number Diff line
@@ -12,6 +12,7 @@ Please mark all change in change log and use the ticket from JIRA. 
- \#104 - test_scheduler core dump

- \#115 - Using new structure for tasktable

- \#139 - New config opion use_gpu_threshold

- \#146 - Add only GPU and only CPU version for IVF_SQ8 and IVF_FLAT



## Improvement

- \#64 - Improvement dump function in scheduler

core/src/scheduler/SchedInst.h

+16 −0
Original line number Diff line number Diff line
@@ -23,10 +23,14 @@ 
#include "Scheduler.h"

#include "optimizer/HybridPass.h"

#include "optimizer/LargeSQ8HPass.h"

#include "optimizer/OnlyCPUPass.h"

#include "optimizer/OnlyGPUPass.h"

#include "optimizer/Optimizer.h"

#include "server/Config.h"



#include <memory>

#include <mutex>

#include <string>

#include <vector>



namespace milvus {
@@ -93,9 +97,21 @@ class OptimizerInst { 
        if (instance == nullptr) {

            std::lock_guard<std::mutex> lock(mutex_);

            if (instance == nullptr) {

                server::Config& config = server::Config::GetInstance();

                std::vector<std::string> search_resources;

                bool has_cpu = false;

                config.GetResourceConfigSearchResources(search_resources);

                for (auto& resource : search_resources) {

                    if (resource == "cpu") {

                        has_cpu = true;

                    }

                }



                std::vector<PassPtr> pass_list;

                pass_list.push_back(std::make_shared<LargeSQ8HPass>());

                pass_list.push_back(std::make_shared<HybridPass>());

                pass_list.push_back(std::make_shared<OnlyCPUPass>());

                pass_list.push_back(std::make_shared<OnlyGPUPass>(has_cpu));

                instance = std::make_shared<Optimizer>(pass_list);

            }

        }

core/src/scheduler/optimizer/OnlyCPUPass.cpp

0 → 100644
+48 −0
Original line number Diff line number Diff line 
// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.



#include "scheduler/optimizer/OnlyCPUPass.h"

#include "scheduler/SchedInst.h"

#include "scheduler/Utils.h"

#include "scheduler/task/SearchTask.h"

#include "scheduler/tasklabel/SpecResLabel.h"



namespace milvus {

namespace scheduler {



bool

OnlyCPUPass::Run(const TaskPtr& task) {

    if (task->Type() != TaskType::SearchTask)

        return false;

    auto search_task = std::static_pointer_cast<XSearchTask>(task);

    if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFSQ8 &&

        search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFFLAT) {

        return false;

    }



    auto gpu_id = get_gpu_pool();

    if (not gpu_id.empty())

        return false;



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

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

    task->label() = label;

    return true;

}



}  // namespace scheduler

}  // namespace milvus

core/src/scheduler/optimizer/OnlyCPUPass.h

0 → 100644
+47 −0
Original line number Diff line number Diff line 
// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

#pragma once



#include <condition_variable>

#include <deque>

#include <list>

#include <memory>

#include <mutex>

#include <queue>

#include <string>

#include <thread>

#include <unordered_map>

#include <vector>



#include "Pass.h"



namespace milvus {

namespace scheduler {



class OnlyCPUPass : public Pass {

 public:

    OnlyCPUPass() = default;



 public:

    bool

    Run(const TaskPtr& task) override;

};



using OnlyCPUPassPtr = std::shared_ptr<OnlyCPUPass>;



}  // namespace scheduler

}  // namespace milvus

core/src/scheduler/optimizer/OnlyGPUPass.cpp

0 → 100644
+54 −0
Original line number Diff line number Diff line 
// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.



#include "scheduler/optimizer/OnlyGPUPass.h"

#include "scheduler/SchedInst.h"

#include "scheduler/Utils.h"

#include "scheduler/task/SearchTask.h"

#include "scheduler/tasklabel/SpecResLabel.h"



namespace milvus {

namespace scheduler {



OnlyGPUPass::OnlyGPUPass(bool has_cpu) : has_cpu_(has_cpu) {

}



bool

OnlyGPUPass::Run(const TaskPtr& task) {

    if (task->Type() != TaskType::SearchTask || has_cpu_)

        return false;



    auto search_task = std::static_pointer_cast<XSearchTask>(task);

    if (search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFSQ8 &&

        search_task->file_->engine_type_ != (int)engine::EngineType::FAISS_IVFFLAT) {

        return false;

    }



    auto gpu_id = get_gpu_pool();

    if (gpu_id.empty())

        return false;



    ResourcePtr res_ptr = ResMgrInst::GetInstance()->GetResource(ResourceType::GPU, gpu_id[specified_gpu_id_]);

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

    task->label() = label;



    specified_gpu_id_ = specified_gpu_id_++ % gpu_id.size();

    return true;

}



}  // namespace scheduler

}  // namespace milvus