Resource.cpp 6.7 KB
Newer Older
J
jinhai 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17
// 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.

S
starlord 已提交
18
#include "scheduler/resource/Resource.h"
19
#include "scheduler/SchedInst.h"
S
starlord 已提交
20
#include "scheduler/Utils.h"
X
xj.lin 已提交
21

S
starlord 已提交
22
#include <iostream>
W
wxyu 已提交
23
#include <limits>
S
starlord 已提交
24
#include <utility>
X
xj.lin 已提交
25 26

namespace milvus {
W
wxyu 已提交
27
namespace scheduler {
X
xj.lin 已提交
28

S
starlord 已提交
29 30
std::ostream&
operator<<(std::ostream& out, const Resource& resource) {
31 32 33 34
    out << resource.Dump();
    return out;
}

W
wxyu 已提交
35 36 37 38 39 40 41 42 43 44 45 46
std::string
ToString(ResourceType type) {
    switch (type) {
        case ResourceType::DISK: {
            return "DISK";
        }
        case ResourceType::CPU: {
            return "CPU";
        }
        case ResourceType::GPU: {
            return "GPU";
        }
W
wxyu 已提交
47
        default: { return "UNKNOWN"; }
W
wxyu 已提交
48 49 50
    }
}

S
starlord 已提交
51
Resource::Resource(std::string name, ResourceType type, uint64_t device_id, bool enable_loader, bool enable_executor)
X
xj.lin 已提交
52 53
    : name_(std::move(name)),
      type_(type),
54
      device_id_(device_id),
55
      enable_loader_(enable_loader),
W
wxyu 已提交
56 57
      enable_executor_(enable_executor) {
    // register subscriber in tasktable
W
wxyu 已提交
58 59 60 61 62 63
    task_table_.RegisterSubscriber([&] {
        if (subscriber_) {
            auto event = std::make_shared<TaskTableUpdatedEvent>(shared_from_this());
            subscriber_(std::static_pointer_cast<Event>(event));
        }
    });
X
xj.lin 已提交
64 65
}

W
wxyu 已提交
66 67
void
Resource::Start() {
W
wxyu 已提交
68
    running_ = true;
69 70 71 72 73 74
    if (enable_loader_) {
        loader_thread_ = std::thread(&Resource::loader_function, this);
    }
    if (enable_executor_) {
        executor_thread_ = std::thread(&Resource::executor_function, this);
    }
X
xj.lin 已提交
75 76
}

W
wxyu 已提交
77 78
void
Resource::Stop() {
X
xj.lin 已提交
79
    running_ = false;
80 81 82 83 84 85 86 87
    if (enable_loader_) {
        WakeupLoader();
        loader_thread_.join();
    }
    if (enable_executor_) {
        WakeupExecutor();
        executor_thread_.join();
    }
X
xj.lin 已提交
88 89
}

W
wxyu 已提交
90 91
void
Resource::WakeupLoader() {
92 93 94 95
    {
        std::lock_guard<std::mutex> lock(load_mutex_);
        load_flag_ = true;
    }
X
xj.lin 已提交
96 97 98
    load_cv_.notify_one();
}

W
wxyu 已提交
99 100
void
Resource::WakeupExecutor() {
101 102 103 104
    {
        std::lock_guard<std::mutex> lock(exec_mutex_);
        exec_flag_ = true;
    }
W
wxyu 已提交
105 106 107
    exec_cv_.notify_one();
}

W
wxyu 已提交
108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123
json
Resource::Dump() const {
    json ret{
        {"device_id", device_id_},
        {"name", name_},
        {"type", ToString(type_)},
        {"task_average_cost", TaskAvgCost()},
        {"task_total_cost", total_cost_},
        {"total_tasks", total_task_},
        {"running", running_},
        {"enable_loader", enable_loader_},
        {"enable_executor", enable_executor_},
    };
    return ret;
}

W
wxyu 已提交
124 125
uint64_t
Resource::NumOfTaskToExec() {
W
wxyu 已提交
126
    return task_table_.TaskToExecute();
W
wxyu 已提交
127 128
}

S
starlord 已提交
129 130
TaskTableItemPtr
Resource::pick_task_load() {
131
    auto indexes = task_table_.PickToLoad(10);
X
xj.lin 已提交
132 133
    for (auto index : indexes) {
        // try to set one task loading, then return
S
starlord 已提交
134
        if (task_table_.Load(index))
135
            return task_table_.at(index);
X
xj.lin 已提交
136 137 138 139 140
        // else try next
    }
    return nullptr;
}

S
starlord 已提交
141 142
TaskTableItemPtr
Resource::pick_task_execute() {
143
    auto indexes = task_table_.PickToExecute(std::numeric_limits<uint64_t>::max());
X
xj.lin 已提交
144 145
    for (auto index : indexes) {
        // try to set one task executing, then return
W
wxyu 已提交
146
        if (task_table_[index]->task->label()->Type() == TaskLabelType::SPECIFIED_RESOURCE) {
Y
Yu Kun 已提交
147 148
            if (task_table_[index]->task->path().Last() != name()) {
                continue;
149 150
            }
        }
Y
Yu Kun 已提交
151 152

        if (task_table_.Execute(index)) {
153
            return task_table_.at(index);
Y
Yu Kun 已提交
154
        }
Y
Yu Kun 已提交
155 156 157 158 159 160 161 162 163
        //        if (task_table_[index]->task->label()->Type() == TaskLabelType::SPECIFIED_RESOURCE) {
        //            if (task_table_.Get(index)->task->path().Current() == task_table_.Get(index)->task->path().Last()
        //            &&
        //                task_table_.Get(index)->task->path().Last() == name()) {
        //                if (task_table_.Execute(index)) {
        //                    return task_table_.Get(index);
        //                }
        //            }
        //        }
X
xj.lin 已提交
164 165 166 167 168
        // else try next
    }
    return nullptr;
}

S
starlord 已提交
169 170
void
Resource::loader_function() {
X
xj.lin 已提交
171 172
    while (running_) {
        std::unique_lock<std::mutex> lock(load_mutex_);
S
starlord 已提交
173
        load_cv_.wait(lock, [&] { return load_flag_; });
W
wxyu 已提交
174
        load_flag_ = false;
175
        lock.unlock();
176 177 178 179 180
        while (true) {
            auto task_item = pick_task_load();
            if (task_item == nullptr) {
                break;
            }
181
            LoadFile(task_item->task);
182
            task_item->Loaded();
183 184 185 186
            if (task_item->from) {
                task_item->from->Moved();
                task_item->from = nullptr;
            }
W
wxyu 已提交
187
            if (subscriber_) {
188
                auto event = std::make_shared<LoadCompletedEvent>(shared_from_this(), task_item);
189
                subscriber_(std::static_pointer_cast<Event>(event));
W
wxyu 已提交
190
            }
X
xj.lin 已提交
191 192 193 194
        }
    }
}

S
starlord 已提交
195 196
void
Resource::executor_function() {
W
wxyu 已提交
197
    if (subscriber_) {
198 199
        auto event = std::make_shared<StartUpEvent>(shared_from_this());
        subscriber_(std::static_pointer_cast<Event>(event));
W
wxyu 已提交
200
    }
X
xj.lin 已提交
201 202
    while (running_) {
        std::unique_lock<std::mutex> lock(exec_mutex_);
S
starlord 已提交
203
        exec_cv_.wait(lock, [&] { return exec_flag_; });
W
wxyu 已提交
204
        exec_flag_ = false;
205
        lock.unlock();
206 207 208 209 210
        while (true) {
            auto task_item = pick_task_execute();
            if (task_item == nullptr) {
                break;
            }
Y
Yu Kun 已提交
211

212
            auto start = get_current_timestamp();
213
            Process(task_item->task);
214
            auto finish = get_current_timestamp();
Y
Yu Kun 已提交
215 216 217
            ++total_task_;
            total_cost_ += finish - start;

218
            task_item->Executed();
219 220 221 222

            if (task_item->task->Type() == TaskType::BuildIndexTask) {
                BuildMgrInst::GetInstance()->Put();
                ResMgrInst::GetInstance()->GetResource("cpu")->WakeupLoader();
Y
Yu Kun 已提交
223
                ResMgrInst::GetInstance()->GetResource("disk")->WakeupLoader();
224 225
            }

W
wxyu 已提交
226
            if (subscriber_) {
227 228
                auto event = std::make_shared<FinishTaskEvent>(shared_from_this(), task_item);
                subscriber_(std::static_pointer_cast<Event>(event));
W
wxyu 已提交
229
            }
X
xj.lin 已提交
230 231 232 233
        }
    }
}

S
starlord 已提交
234 235
}  // namespace scheduler
}  // namespace milvus