Fix the multi-gpus building index enable problem (#3845)

Signed-off-by: cqy <yaya645@126.com>
Signed-off-by: shengjun.li <shengjun.li@zilliz.com>

core/src/db/engine/ExecutionEngine.h

@@ -46,7 +46,7 @@ class ExecutionEngine {
     Search(ExecutionEngineContext& context) = 0;
 
     virtual Status
-    BuildIndex() = 0;
+    BuildIndex(uint64_t device_id) = 0;
 };
 
 using ExecutionEnginePtr = std::shared_ptr<ExecutionEngine>;

core/src/db/engine/ExecutionEngineImpl.cpp

@@ -602,7 +602,7 @@ ExecutionEngineImpl::ProcessRangeQuery(const std::unordered_map<std::string, Dat
 }
 
 Status
-ExecutionEngineImpl::BuildIndex() {
+ExecutionEngineImpl::BuildIndex(uint64_t device_id) {
     TimeRecorderAuto rc("ExecutionEngineImpl::BuildIndex");
 
     SegmentPtr segment_ptr;
@@ -616,6 +616,8 @@ ExecutionEngineImpl::BuildIndex() {
     context.prev_partition = snapshot->GetResource<snapshot::Partition>(segment->GetPartitionId());
     auto build_op = std::make_shared<snapshot::ChangeSegmentFileOperation>(context, snapshot);
 
+    gpu_num_ = device_id;
+
     for (auto& field_name : target_fields_) {
         // create snapshot segment files
         CollectionIndex index_info;

core/src/db/engine/ExecutionEngineImpl.h

@@ -38,7 +38,7 @@ class ExecutionEngineImpl : public ExecutionEngine {
     Search(ExecutionEngineContext& context) override;
 
     Status
-    BuildIndex() override;
+    BuildIndex(uint64_t device_id) override;
 
  private:
     Status

core/src/scheduler/task/BuildIndexTask.cpp

@@ -58,6 +58,7 @@ BuildIndexTask::OnLoad(milvus::scheduler::LoadType type, uint8_t device_id) {
             // no need to copy flat to gpu,
             //            stat = execution_engine_->CopyToGpu(device_id);
             //            type_str = "CPU2GPU:" + std::to_string(device_id);
+            gpu_device_id = device_id;
         } else {
             error_msg = "Wrong load type";
             stat = Status(SERVER_UNEXPECTED_ERROR, error_msg);
@@ -100,7 +101,7 @@ BuildIndexTask::OnExecute() {
 
     Status status;
     try {
-        status = execution_engine_->BuildIndex();
+        status = execution_engine_->BuildIndex(gpu_device_id);
     } catch (std::exception& e) {
         status = Status(DB_ERROR, e.what());
     }

core/src/scheduler/task/BuildIndexTask.h

@@ -58,6 +58,9 @@ class BuildIndexTask : public Task {
     engine::TargetFields target_fields_;
 
     engine::ExecutionEnginePtr execution_engine_;
+
+ private:
+    int64_t gpu_device_id = 0;
 };
 
 }  // namespace scheduler