SearchTask.cpp 14.2 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/task/SearchTask.h"
S
starlord 已提交
19
#include "db/engine/EngineFactory.h"
S
starlord 已提交
20 21
#include "metrics/Metrics.h"
#include "scheduler/job/SearchJob.h"
S
starlord 已提交
22
#include "utils/Log.h"
S
starlord 已提交
23
#include "utils/TimeRecorder.h"
24

W
wxyu 已提交
25
#include <src/scheduler/SchedInst.h>
S
starlord 已提交
26
#include <algorithm>
S
starlord 已提交
27
#include <string>
28
#include <thread>
S
starlord 已提交
29
#include <utility>
30 31

namespace milvus {
W
wxyu 已提交
32
namespace scheduler {
33 34 35 36

static constexpr size_t PARALLEL_REDUCE_THRESHOLD = 10000;
static constexpr size_t PARALLEL_REDUCE_BATCH = 1000;

37
// TODO(wxyu): remove unused code
S
starlord 已提交
38 39
// bool
// NeedParallelReduce(uint64_t nq, uint64_t topk) {
40 41 42 43 44 45 46 47 48 49
//    server::ServerConfig &config = server::ServerConfig::GetInstance();
//    server::ConfigNode &db_config = config.GetConfig(server::CONFIG_DB);
//    bool need_parallel = db_config.GetBoolValue(server::CONFIG_DB_PARALLEL_REDUCE, false);
//    if (!need_parallel) {
//        return false;
//    }
//
//    return nq * topk >= PARALLEL_REDUCE_THRESHOLD;
//}
//
S
starlord 已提交
50 51
// void
// ParallelReduce(std::function<void(size_t, size_t)> &reduce_function, size_t max_index) {
52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78
//    size_t reduce_batch = PARALLEL_REDUCE_BATCH;
//
//    auto thread_count = std::thread::hardware_concurrency() - 1; //not all core do this work
//    if (thread_count > 0) {
//        reduce_batch = max_index / thread_count + 1;
//    }
//    ENGINE_LOG_DEBUG << "use " << thread_count <<
//                     " thread parallelly do reduce, each thread process " << reduce_batch << " vectors";
//
//    std::vector<std::shared_ptr<std::thread> > thread_array;
//    size_t from_index = 0;
//    while (from_index < max_index) {
//        size_t to_index = from_index + reduce_batch;
//        if (to_index > max_index) {
//            to_index = max_index;
//        }
//
//        auto reduce_thread = std::make_shared<std::thread>(reduce_function, from_index, to_index);
//        thread_array.push_back(reduce_thread);
//
//        from_index = to_index;
//    }
//
//    for (auto &thread_ptr : thread_array) {
//        thread_ptr->join();
//    }
//}
79 80 81

void
CollectFileMetrics(int file_type, size_t file_size) {
Y
yudong.cai 已提交
82
    server::MetricsBase& inst = server::Metrics::GetInstance();
83
    switch (file_type) {
W
wxyu 已提交
84 85
        case TableFileSchema::RAW:
        case TableFileSchema::TO_INDEX: {
Y
yudong.cai 已提交
86 87 88
            inst.RawFileSizeHistogramObserve(file_size);
            inst.RawFileSizeTotalIncrement(file_size);
            inst.RawFileSizeGaugeSet(file_size);
89 90 91
            break;
        }
        default: {
Y
yudong.cai 已提交
92 93 94
            inst.IndexFileSizeHistogramObserve(file_size);
            inst.IndexFileSizeTotalIncrement(file_size);
            inst.IndexFileSizeGaugeSet(file_size);
95 96 97 98 99
            break;
        }
    }
}

W
wxyu 已提交
100 101
XSearchTask::XSearchTask(TableFileSchemaPtr file, TaskLabelPtr label)
    : Task(TaskType::SearchTask, std::move(label)), file_(file) {
W
wxyu 已提交
102
    if (file_) {
103 104 105
        if (file_->metric_type_ != static_cast<int>(MetricType::L2)) {
            metric_l2 = false;
        }
S
starlord 已提交
106 107
        index_engine_ = EngineFactory::Build(file_->dimension_, file_->location_, (EngineType)file_->engine_type_,
                                             (MetricType)file_->metric_type_, file_->nlist_);
W
wxyu 已提交
108
    }
W
wxyu 已提交
109 110
}

111 112
void
XSearchTask::Load(LoadType type, uint8_t device_id) {
S
starlord 已提交
113
    TimeRecorder rc("");
W
wxyu 已提交
114 115
    Status stat = Status::OK();
    std::string error_msg;
116
    std::string type_str;
117 118

    try {
W
wxyu 已提交
119
        if (type == LoadType::DISK2CPU) {
W
wxyu 已提交
120
            stat = index_engine_->Load();
121
            type_str = "DISK2CPU";
W
wxyu 已提交
122
        } else if (type == LoadType::CPU2GPU) {
W
wxyu 已提交
123 124 125 126 127
            bool hybrid = false;
            if (index_engine_->IndexEngineType() == engine::EngineType::FAISS_IVFSQ8H) {
                hybrid = true;
            }
            stat = index_engine_->CopyToGpu(device_id, hybrid);
128
            type_str = "CPU2GPU";
W
wxyu 已提交
129
        } else if (type == LoadType::GPU2CPU) {
W
wxyu 已提交
130
            stat = index_engine_->CopyToCpu();
131
            type_str = "GPU2CPU";
W
wxyu 已提交
132
        } else {
W
wxyu 已提交
133 134
            error_msg = "Wrong load type";
            stat = Status(SERVER_UNEXPECTED_ERROR, error_msg);
W
wxyu 已提交
135
        }
S
starlord 已提交
136 137
    } catch (std::exception& ex) {
        // typical error: out of disk space or permition denied
W
wxyu 已提交
138 139 140 141 142
        error_msg = "Failed to load index file: " + std::string(ex.what());
        stat = Status(SERVER_UNEXPECTED_ERROR, error_msg);
    }

    if (!stat.ok()) {
143 144 145 146 147 148 149 150
        Status s;
        if (stat.ToString().find("out of memory") != std::string::npos) {
            error_msg = "out of memory: " + type_str;
            s = Status(SERVER_OUT_OF_MEMORY, error_msg);
        } else {
            error_msg = "Failed to load index file: " + type_str;
            s = Status(SERVER_UNEXPECTED_ERROR, error_msg);
        }
151

S
starlord 已提交
152
        if (auto job = job_.lock()) {
W
wxyu 已提交
153 154 155
            auto search_job = std::static_pointer_cast<scheduler::SearchJob>(job);
            search_job->SearchDone(file_->id_);
            search_job->GetStatus() = s;
156 157 158 159 160
        }

        return;
    }

W
wxyu 已提交
161
    size_t file_size = index_engine_->PhysicalSize();
162

J
JinHai-CN 已提交
163 164
    std::string info = "Load file id:" + std::to_string(file_->id_) +
                       " file type:" + std::to_string(file_->file_type_) + " size:" + std::to_string(file_size) +
S
starlord 已提交
165
                       " bytes from location: " + file_->location_ + " totally cost";
166
    double span = rc.ElapseFromBegin(info);
S
starlord 已提交
167 168 169
    //    for (auto &context : search_contexts_) {
    //        context->AccumLoadCost(span);
    //    }
170 171 172

    CollectFileMetrics(file_->file_type_, file_size);

S
starlord 已提交
173
    // step 2: return search task for later execution
174 175
    index_id_ = file_->id_;
    index_type_ = file_->file_type_;
S
starlord 已提交
176
    //    search_contexts_.swap(search_contexts_);
177 178 179 180 181 182 183 184
}

void
XSearchTask::Execute() {
    if (index_engine_ == nullptr) {
        return;
    }

S
starlord 已提交
185 186
    //    ENGINE_LOG_DEBUG << "Searching in file id:" << index_id_ << " with "
    //                     << search_contexts_.size() << " tasks";
187

S
starlord 已提交
188
    TimeRecorder rc("DoSearch file id:" + std::to_string(index_id_));
189

Y
Yu Kun 已提交
190
    server::CollectDurationMetrics metrics(index_type_);
191

S
starlord 已提交
192
    std::vector<int64_t> output_ids;
J
jinhai 已提交
193
    std::vector<float> output_distance;
W
wxyu 已提交
194 195 196

    if (auto job = job_.lock()) {
        auto search_job = std::static_pointer_cast<scheduler::SearchJob>(job);
S
starlord 已提交
197
        // step 1: allocate memory
W
wxyu 已提交
198 199 200
        uint64_t nq = search_job->nq();
        uint64_t topk = search_job->topk();
        uint64_t nprobe = search_job->nprobe();
S
starlord 已提交
201
        const float* vectors = search_job->vectors();
Y
yudong.cai 已提交
202 203 204

        output_ids.resize(topk * nq);
        output_distance.resize(topk * nq);
S
starlord 已提交
205 206
        std::string hdr =
            "job " + std::to_string(search_job->id()) + " nq " + std::to_string(nq) + " topk " + std::to_string(topk);
207 208

        try {
S
starlord 已提交
209
            // step 2: search
W
wxyu 已提交
210 211 212 213 214 215
            bool hybrid = false;
            if (index_engine_->IndexEngineType() == engine::EngineType::FAISS_IVFSQ8H &&
                ResMgrInst::GetInstance()->GetResource(path().Last())->type() == ResourceType::CPU) {
                hybrid = true;
            }
            index_engine_->Search(nq, vectors, topk, nprobe, output_distance.data(), output_ids.data(), hybrid);
216

Y
yudong.cai 已提交
217
            double span = rc.RecordSection(hdr + ", do search");
S
starlord 已提交
218
            //            search_job->AccumSearchCost(span);
219

Y
yudong.cai 已提交
220
            // step 3: pick up topk result
Y
yudong.cai 已提交
221
            auto spec_k = index_engine_->Count() < topk ? index_engine_->Count() : topk;
J
jinhai 已提交
222 223 224
            {
                std::unique_lock<std::mutex> lock(search_job->mutex());
                XSearchTask::MergeTopkToResultSet(output_ids, output_distance, spec_k, nq, topk, metric_l2,
225
                                              search_job->GetResult());
J
jinhai 已提交
226
            }
227

Y
yudong.cai 已提交
228
            span = rc.RecordSection(hdr + ", reduce topk");
S
starlord 已提交
229 230
            //            search_job->AccumReduceCost(span);
        } catch (std::exception& ex) {
231
            ENGINE_LOG_ERROR << "SearchTask encounter exception: " << ex.what();
S
starlord 已提交
232
            //            search_job->IndexSearchDone(index_id_);//mark as done avoid dead lock, even search failed
233 234
        }

Y
yudong.cai 已提交
235
        // step 4: notify to send result to client
W
wxyu 已提交
236
        search_job->SearchDone(index_id_);
237 238 239
    }

    rc.ElapseFromBegin("totally cost");
240 241 242

    // release index in resource
    index_engine_ = nullptr;
243 244
}

Y
yudong.cai 已提交
245
void
246 247
XSearchTask::MergeTopkToResultSet(const std::vector<int64_t>& input_ids, const std::vector<float>& input_distance,
                                  uint64_t input_k, uint64_t nq, uint64_t topk, bool ascending,
Y
yudong.cai 已提交
248
                                  scheduler::ResultSet& result) {
Y
yudong.cai 已提交
249
    if (result.empty()) {
Y
yudong.cai 已提交
250 251 252 253 254
        result.resize(nq);
    }

    for (uint64_t i = 0; i < nq; i++) {
        scheduler::Id2DistVec result_buf;
255
        auto& result_i = result[i];
Y
yudong.cai 已提交
256 257 258

        if (result[i].empty()) {
            result_buf.resize(input_k, scheduler::IdDistPair(-1, 0.0));
259
            uint64_t input_k_multi_i = topk * i;
Y
yudong.cai 已提交
260 261
            for (auto k = 0; k < input_k; ++k) {
                uint64_t idx = input_k_multi_i + k;
262
                auto& result_buf_item = result_buf[k];
Y
yudong.cai 已提交
263 264
                result_buf_item.first = input_ids[idx];
                result_buf_item.second = input_distance[idx];
265
            }
Y
yudong.cai 已提交
266 267 268 269
        } else {
            size_t tar_size = result_i.size();
            uint64_t output_k = std::min(topk, input_k + tar_size);
            result_buf.resize(output_k, scheduler::IdDistPair(-1, 0.0));
Y
yudong.cai 已提交
270 271
            size_t buf_k = 0, src_k = 0, tar_k = 0;
            uint64_t src_idx;
272
            uint64_t input_k_multi_i = topk * i;
Y
yudong.cai 已提交
273 274
            while (buf_k < output_k && src_k < input_k && tar_k < tar_size) {
                src_idx = input_k_multi_i + src_k;
275
                auto& result_buf_item = result_buf[buf_k];
Y
yudong.cai 已提交
276 277
                auto& result_item = result_i[tar_k];
                if ((ascending && input_distance[src_idx] < result_item.second) ||
S
starlord 已提交
278
                    (!ascending && input_distance[src_idx] > result_item.second)) {
Y
yudong.cai 已提交
279 280 281 282 283 284 285 286
                    result_buf_item.first = input_ids[src_idx];
                    result_buf_item.second = input_distance[src_idx];
                    src_k++;
                } else {
                    result_buf_item = result_item;
                    tar_k++;
                }
                buf_k++;
287 288
            }

Y
yudong.cai 已提交
289
            if (buf_k < output_k) {
Y
yudong.cai 已提交
290 291 292
                if (src_k < input_k) {
                    while (buf_k < output_k && src_k < input_k) {
                        src_idx = input_k_multi_i + src_k;
293
                        auto& result_buf_item = result_buf[buf_k];
Y
yudong.cai 已提交
294 295 296 297 298 299 300
                        result_buf_item.first = input_ids[src_idx];
                        result_buf_item.second = input_distance[src_idx];
                        src_k++;
                        buf_k++;
                    }
                } else {
                    while (buf_k < output_k && tar_k < tar_size) {
Y
yudong.cai 已提交
301
                        result_buf[buf_k] = result_i[tar_k];
Y
yudong.cai 已提交
302 303 304 305
                        tar_k++;
                        buf_k++;
                    }
                }
306 307
            }
        }
Y
yudong.cai 已提交
308 309

        result_i.swap(result_buf);
310
    }
Y
yudong.cai 已提交
311
}
312

S
starlord 已提交
313 314
// void
// XSearchTask::MergeTopkArray(std::vector<int64_t>& tar_ids, std::vector<float>& tar_distance, uint64_t& tar_input_k,
315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377
//                            const std::vector<int64_t>& src_ids, const std::vector<float>& src_distance,
//                            uint64_t src_input_k, uint64_t nq, uint64_t topk, bool ascending) {
//    if (src_ids.empty() || src_distance.empty()) {
//        return;
//    }
//
//    uint64_t output_k = std::min(topk, tar_input_k + src_input_k);
//    std::vector<int64_t> id_buf(nq * output_k, -1);
//    std::vector<float> dist_buf(nq * output_k, 0.0);
//
//    uint64_t buf_k, src_k, tar_k;
//    uint64_t src_idx, tar_idx, buf_idx;
//    uint64_t src_input_k_multi_i, tar_input_k_multi_i, buf_k_multi_i;
//
//    for (uint64_t i = 0; i < nq; i++) {
//        src_input_k_multi_i = src_input_k * i;
//        tar_input_k_multi_i = tar_input_k * i;
//        buf_k_multi_i = output_k * i;
//        buf_k = src_k = tar_k = 0;
//        while (buf_k < output_k && src_k < src_input_k && tar_k < tar_input_k) {
//            src_idx = src_input_k_multi_i + src_k;
//            tar_idx = tar_input_k_multi_i + tar_k;
//            buf_idx = buf_k_multi_i + buf_k;
//            if ((ascending && src_distance[src_idx] < tar_distance[tar_idx]) ||
//                (!ascending && src_distance[src_idx] > tar_distance[tar_idx])) {
//                id_buf[buf_idx] = src_ids[src_idx];
//                dist_buf[buf_idx] = src_distance[src_idx];
//                src_k++;
//            } else {
//                id_buf[buf_idx] = tar_ids[tar_idx];
//                dist_buf[buf_idx] = tar_distance[tar_idx];
//                tar_k++;
//            }
//            buf_k++;
//        }
//
//        if (buf_k < output_k) {
//            if (src_k < src_input_k) {
//                while (buf_k < output_k && src_k < src_input_k) {
//                    src_idx = src_input_k_multi_i + src_k;
//                    buf_idx = buf_k_multi_i + buf_k;
//                    id_buf[buf_idx] = src_ids[src_idx];
//                    dist_buf[buf_idx] = src_distance[src_idx];
//                    src_k++;
//                    buf_k++;
//                }
//            } else {
//                while (buf_k < output_k && tar_k < tar_input_k) {
//                    tar_idx = tar_input_k_multi_i + tar_k;
//                    buf_idx = buf_k_multi_i + buf_k;
//                    id_buf[buf_idx] = tar_ids[tar_idx];
//                    dist_buf[buf_idx] = tar_distance[tar_idx];
//                    tar_k++;
//                    buf_k++;
//                }
//            }
//        }
//    }
//
//    tar_ids.swap(id_buf);
//    tar_distance.swap(dist_buf);
//    tar_input_k = output_k;
//}
378

S
starlord 已提交
379 380
}  // namespace scheduler
}  // namespace milvus