From 6f5d66f533b51f88184bc0cb08560de7c7d867ab Mon Sep 17 00:00:00 2001
From: zouap <zouap@pcl.ac.cn>
Date: Mon, 13 Jun 2022 14:20:30 +0800
Subject: [PATCH] =?UTF-8?q?=E6=8F=90=E4=BA=A4=E4=BB=A3=E7=A0=81=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: zouap <zouap@pcl.ac.cn>
---
 routers/repo/ai_model_manage.go |  7 ++++---
 web_src/js/components/Model.vue | 16 ++++++++++++++--
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/routers/repo/ai_model_manage.go b/routers/repo/ai_model_manage.go
index 67eee5f3b..334197737 100644
--- a/routers/repo/ai_model_manage.go
+++ b/routers/repo/ai_model_manage.go
@@ -28,7 +28,7 @@ const (
 	MODEL_NOT_LATEST       = 0
 )
 
-func saveModelByParameters(jobId string, versionName string, name string, version string, label string, description string, ctx *context.Context) error {
+func saveModelByParameters(jobId string, versionName string, name string, version string, label string, description string, engine int, ctx *context.Context) error {
 	aiTask, err := models.GetCloudbrainByJobIDAndVersionName(jobId, versionName)
 	if err != nil {
 		aiTask, err = models.GetRepoCloudBrainByJobID(ctx.Repo.Repository.ID, jobId)
@@ -153,8 +153,9 @@ func SaveModel(ctx *context.Context) {
 	version := ctx.Query("Version")
 	label := ctx.Query("Label")
 	description := ctx.Query("Description")
+	engine := ctx.QueryInt("Engine")
 	trainTaskCreate := ctx.QueryBool("trainTaskCreate")
-
+	log.Info("engine=" + fmt.Sprint(engine))
 	if !trainTaskCreate {
 		if !ctx.Repo.CanWrite(models.UnitTypeModelManage) {
 			//ctx.NotFound(ctx.Req.URL.RequestURI(), nil)
@@ -173,7 +174,7 @@ func SaveModel(ctx *context.Context) {
 		return
 	}
 
-	err := saveModelByParameters(JobId, VersionName, name, version, label, description, ctx)
+	err := saveModelByParameters(JobId, VersionName, name, version, label, description, engine, ctx)
 
 	if err != nil {
 		log.Info("save model error." + err.Error())
diff --git a/web_src/js/components/Model.vue b/web_src/js/components/Model.vue
index b67ddf3f3..c3cf04d11 100644
--- a/web_src/js/components/Model.vue
+++ b/web_src/js/components/Model.vue
@@ -168,7 +168,7 @@ export default {
                 tableData= res.data
                 for(let i=0;i<tableData.length;i++){
                     TrainTaskInfo = JSON.parse(tableData[i].TrainTaskInfo)
-                    tableData[i].EngineName = TrainTaskInfo.EngineName.split('-')[0]
+                    tableData[i].EngineName = getEngineName(tableData[i])
                     tableData[i].ComputeResource = TrainTaskInfo.ComputeResource
                     tableData[i].cName=tableData[i].Name
                     tableData[i].Name=''
@@ -347,6 +347,18 @@ export default {
         })
         .modal('show')
     },
+    getEngineName(model){
+        if(model.Engine == 0){
+            return "Pytorch";
+        }else  if(model.Engine == 1 || model.Engine == 121){
+            return "TensorFlow";
+        }else  if(model.Engine == 2 || model.Engine == 122){
+            return "MindSpore";
+        }else{
+            return "Other"
+        }
+
+    },
     getModelList(){
         try {
             this.$refs.table.store.states.lazyTreeNodeMap = {}
@@ -360,7 +372,7 @@ export default {
                 for(let i=0;i<this.tableData.length;i++){
                     TrainTaskInfo = JSON.parse(this.tableData[i].TrainTaskInfo)
                     this.tableData[i].cName=this.tableData[i].Name
-                    this.tableData[i].EngineName = TrainTaskInfo.EngineName.split('-')[0]
+                    this.tableData[i].EngineName = getEngineName(this.tableData[i])
                     this.tableData[i].ComputeResource = TrainTaskInfo.ComputeResource
                     this.tableData[i].hasChildren = res.data.data[i].VersionCount===1 ? false : true
                 }