Commit 49105e44 authored by 余昆's avatar 余昆
Browse files

Add only GPU and only CPU version for IVF_SQ8 and IVF_FLAT 


Former-commit-id: 14f5b095df2988d5f65ecd1cb2628fd3a3791a14
parent 87b80491

CHANGELOG.md

+1 −0
Original line number Diff line number Diff line
@@ -11,6 +11,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

+4 −0
Original line number Diff line number Diff line
@@ -23,6 +23,8 @@ 
#include "Scheduler.h"

#include "optimizer/HybridPass.h"

#include "optimizer/LargeSQ8HPass.h"

#include "optimizer/OnlyCPUPass.h"

#include "optimizer/OnlyGPUPass.h"

#include "optimizer/Optimizer.h"



#include <memory>
@@ -96,6 +98,8 @@ class OptimizerInst { 
                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>());

                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/task/SearchTask.h"

#include "scheduler/tasklabel/SpecResLabel.h"

#include "scheduler/Utils.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 (gpu_id.empty()) {

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

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

        task->label() = label;

        return true;

    }

    return false;

}



}  // 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
+66 −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/task/SearchTask.h"

#include "scheduler/tasklabel/SpecResLabel.h"

#include "scheduler/Utils.h"

#include "server/Config.h"



namespace milvus {

namespace scheduler {



bool

OnlyGPUPass::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;

    }



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

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

    config.GetResourceConfigSearchResources(search_resources);

    for (auto &resource : search_resources) {

        if (resource == "cpu") {

            return false;

        }

    }



    auto gpu_id = get_gpu_pool();

    if (!gpu_id.empty()) {

        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;

    } else {

        return false;

    }



    if (specified_gpu_id_ < gpu_id.size() - 1) {

        ++specified_gpu_id_;

    } else {

        specified_gpu_id_ = 0;

    }

    return true;

}



}  // namespace scheduler

}  // namespace milvus