From 81bdb8b864c4fe015611383fb642e69840a48833 Mon Sep 17 00:00:00 2001
From: zouap <zouap@pcl.ac.cn>
Date: Fri, 31 Dec 2021 17:49:05 +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>
---
 models/cloudbrain.go            | 10 ++++++----
 routers/repo/ai_model_manage.go | 17 +++++++++++++++++
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index 5bf066413..c3707c9e2 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -19,8 +19,8 @@ type JobType string
 type ModelArtsJobStatus string
 
 const (
-	NPUResource  = "NPU"
-	GPUResource  = "CPU/GPU"
+	NPUResource = "NPU"
+	GPUResource = "CPU/GPU"
 
 	JobWaiting   CloudbrainStatus = "WAITING"
 	JobStopped   CloudbrainStatus = "STOPPED"
@@ -210,7 +210,7 @@ type CloudbrainsOptions struct {
 	JobType         string
 	VersionName     string
 	IsLatestVersion string
-	JobTypeNot		bool
+	JobTypeNot      bool
 }
 
 type TaskPod struct {
@@ -975,7 +975,9 @@ func QueryModelTrainJobList(repoId int64) ([]*CloudbrainInfo, int, error) {
 	cond = cond.And(
 		builder.Eq{"Status": "COMPLETED"},
 	)
-
+	cond = cond.And(
+		builder.Eq{"job_type": "TRAIN"},
+	)
 	cloudbrains := make([]*CloudbrainInfo, 0)
 	if err := sess.Select("job_id,job_name").Table(&Cloudbrain{}).Where(cond).OrderBy("created_unix DESC").
 		Find(&cloudbrains); err != nil {
diff --git a/routers/repo/ai_model_manage.go b/routers/repo/ai_model_manage.go
index c1b5a099e..87c5dc001 100644
--- a/routers/repo/ai_model_manage.go
+++ b/routers/repo/ai_model_manage.go
@@ -474,6 +474,23 @@ func ShowOneVersionOtherModel(ctx *context.Context) {
 
 func ShowModelTemplate(ctx *context.Context) {
 	ctx.Data["isModelManage"] = true
+	repoId := ctx.Repo.Repository.ID
+	Type := -1
+	_, count, _ := models.QueryModel(&models.AiModelQueryOptions{
+		ListOptions: models.ListOptions{
+			Page:     1,
+			PageSize: 2,
+		},
+		RepoID: repoId,
+		Type:   Type,
+		New:    MODEL_LATEST,
+	})
+	ctx.Data["MODEL_COUNT"] = count
+
+	_, trainCount, _ := models.QueryModelTrainJobList(repoId)
+	log.Info("query train count=" + fmt.Sprint(trainCount))
+
+	ctx.Data["TRAIN_COUNT"] = trainCount
 	ctx.HTML(200, tplModelManageIndex)
 }