From 0b542f8421110c44161fed84cf2e1a9697406b5c Mon Sep 17 00:00:00 2001
From: zhangxiaokun <zhang.xiaokun@huawei.com>
Date: Thu, 24 Dec 2020 15:40:42 +0800
Subject: [PATCH] Return when Task sink failed

---
 ge/common/helper/model_helper.cc                 | 16 +++++-----------
 ge/graph/load/new_model_manager/davinci_model.cc |  4 ++--
 inc/framework/common/helper/model_helper.h       |  2 +-
 3 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/ge/common/helper/model_helper.cc b/ge/common/helper/model_helper.cc
index aacef88c..d89c326f 100644
--- a/ge/common/helper/model_helper.cc
+++ b/ge/common/helper/model_helper.cc
@@ -122,7 +122,7 @@ Status ModelHelper::SaveModelTbeKernel(std::shared_ptr<OmFileSaveHelper> &om_fil
   if (tbe_kernel_store.DataSize() > 0) {
     GE_CHK_STATUS_RET(
         SaveModelPartition(om_file_save_helper, ModelPartitionType::TBE_KERNELS,
-                           ge_model->GetTBEKernelStore().Data(), ge_model->GetTBEKernelStore().DataSize(), 
+                           ge_model->GetTBEKernelStore().Data(), ge_model->GetTBEKernelStore().DataSize(),
                            model_index), "Add tbe kernel partition failed");
   }
   // no need to check value, DATA->NetOutput
@@ -306,7 +306,7 @@ FMK_FUNC_HOST_VISIBILITY FMK_FUNC_DEV_VISIBILITY Status ModelHelper::SaveToOmRoo
       model_names.emplace_back(item.first);
     }
   }
-  
+
   vector<ge::Buffer> model_buffers(model_names.size());
   vector<ge::Buffer> task_buffers(model_names.size());
 
@@ -611,12 +611,11 @@ Status ModelHelper::LoadModelData(OmFileLoadHelper &om_load_helper) {
     return INTERNAL_ERROR;
   }
 
-  SetModelToGeModel(model);
-
+  SetModelToGeModel(model_, model);
   return SUCCESS;
 }
 
-void ModelHelper::SetModelToGeModel(ge::Model &model) {
+void ModelHelper::SetModelToGeModel(GeModelPtr &ge_model, Model &model) {
   model_->SetGraph(model.GetGraph());
   model_->SetName(model.GetName());
   model_->SetVersion(model.GetVersion());
@@ -636,12 +635,7 @@ Status ModelHelper::LoadModelData(OmFileLoadHelper &om_load_helper, GeModelPtr &
     return INTERNAL_ERROR;
   }
 
-  cur_model->SetGraph(model.GetGraph());
-  cur_model->SetName(model.GetName());
-  cur_model->SetVersion(model.GetVersion());
-  cur_model->SetPlatformVersion(model.GetPlatformVersion());
-  cur_model->SetAttr(model.MutableAttrMap());
-
+  SetModelToGeModel(cur_model, model);
   return SUCCESS;
 }
 
diff --git a/ge/graph/load/new_model_manager/davinci_model.cc b/ge/graph/load/new_model_manager/davinci_model.cc
index fd842a4b..044531f6 100755
--- a/ge/graph/load/new_model_manager/davinci_model.cc
+++ b/ge/graph/load/new_model_manager/davinci_model.cc
@@ -682,7 +682,7 @@ Status DavinciModel::Init(void *dev_ptr, size_t mem_size, void *weight_ptr, size
 
   SetDataDumperArgs(compute_graph);
   GE_TIMESTAMP_START(DoTaskSink);
-  auto ret = DoTaskSink();
+  GE_CHK_STATUS_RET(DoTaskSink(), "Task sink failed");
   GE_TIMESTAMP_END(DoTaskSink, "GraphLoader::DoTaskSink");
 
   auto all_dump_model = GetDumpProperties().GetAllDumpModel();
@@ -723,7 +723,7 @@ Status DavinciModel::Init(void *dev_ptr, size_t mem_size, void *weight_ptr, size
   }
 
   Shrink();
-  return ret;
+  return SUCCESS;
 }
 
 Status DavinciModel::ReportProfilingData() {
diff --git a/inc/framework/common/helper/model_helper.h b/inc/framework/common/helper/model_helper.h
index 7867e63d..bc0444bc 100644
--- a/inc/framework/common/helper/model_helper.h
+++ b/inc/framework/common/helper/model_helper.h
@@ -69,7 +69,7 @@ class ModelHelper {
   Status GenerateGeModel(OmFileLoadHelper &om_load_helper);
   Status GenerateGeRootModel(OmFileLoadHelper &om_load_helper);
   Status LoadModelData(OmFileLoadHelper &om_load_helper);
-  void SetModelToGeModel(ge::Model &model);
+  void SetModelToGeModel(GeModelPtr &ge_model, Model &model);
   Status LoadModelData(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index);
   Status LoadWeights(OmFileLoadHelper &om_load_helper);
   Status LoadWeights(OmFileLoadHelper &om_load_helper, GeModelPtr &cur_model, size_t mode_index);