From 6dbf4df3d6e4b705457d4830999927d600d9a752 Mon Sep 17 00:00:00 2001
From: lewis <747342561@qq.com>
Date: Thu, 4 Nov 2021 11:20:18 +0800
Subject: [PATCH] opts

---
 models/cloudbrain.go           | 10 +++++++++-
 models/file_chunk.go           |  7 ++-----
 modules/modelarts/modelarts.go |  6 +++---
 routers/repo/attachment.go     |  2 +-
 routers/repo/modelarts.go      |  6 ++++--
 5 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index 563ab9d06..af5e9f169 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -29,6 +29,7 @@ const (
 	JobTypeBenchmark    JobType = "BENCHMARK"
 	JobTypeSnn4imagenet JobType = "SNN4IMAGENET"
 	JobTypeBrainScore   JobType = "BRAINSCORE"
+	JobTypeTrain        JobType = "TRAIN"
 
 	ModelArtsCreateQueue  ModelArtsJobStatus = "CREATE_QUEUING" //免费资源创建排队中
 	ModelArtsCreating     ModelArtsJobStatus = "CREATING"       //创建中
@@ -153,7 +154,8 @@ type CloudbrainsOptions struct {
 	SortType      string
 	CloudbrainIDs []int64
 	// JobStatus     CloudbrainStatus
-	Type int
+	Type    	  int
+	JobType 	  string
 }
 type TaskPod struct {
 	TaskRoleStatus struct {
@@ -846,6 +848,12 @@ func Cloudbrains(opts *CloudbrainsOptions) ([]*CloudbrainInfo, int64, error) {
 		)
 	}
 
+	if (opts.JobType) != "" {
+		cond = cond.And(
+			builder.Eq{"cloudbrain.job_type": opts.JobType},
+		)
+	}
+
 	// switch opts.JobStatus {
 	// case JobWaiting:
 	// 	cond.And(builder.Eq{"cloudbrain.status": int(JobWaiting)})
diff --git a/models/file_chunk.go b/models/file_chunk.go
index 8decb7b44..76c926dc5 100755
--- a/models/file_chunk.go
+++ b/models/file_chunk.go
@@ -14,11 +14,8 @@ const (
 )
 
 const (
-	TypeCloudBrainOne      = 0
-	TypeCloudBrainNotebook = 1
-	TypeCloudBrainTrainJob = 2
-
-	TypeCloudBrainTwo = 1
+	TypeCloudBrainOne int = iota
+	TypeCloudBrainTwo
 )
 
 type FileChunk struct {
diff --git a/modules/modelarts/modelarts.go b/modules/modelarts/modelarts.go
index 5717d044f..e1dbe9f5a 100755
--- a/modules/modelarts/modelarts.go
+++ b/modules/modelarts/modelarts.go
@@ -159,7 +159,7 @@ func GenerateTask(ctx *context.Context, jobName, uuid, description string) error
 		JobID:   jobResult.ID,
 		JobName: jobName,
 		JobType: string(models.JobTypeDebug),
-		Type:    models.TypeCloudBrainNotebook,
+		Type:    models.TypeCloudBrainTwo,
 		Uuid:    uuid,
 	})
 
@@ -207,8 +207,8 @@ func GenerateTrainJob(ctx *context.Context, req *GenerateTrainJobReq) error {
 		RepoID:      ctx.Repo.Repository.ID,
 		JobID:       strconv.FormatInt(jobResult.JobID, 10),
 		JobName:     req.JobName,
-		JobType:     string(models.JobTypeDebug),
-		Type:        models.TypeCloudBrainTrainJob,
+		JobType:     string(models.JobTypeTrain),
+		Type:        models.TypeCloudBrainTwo,
 		VersionID:   jobResult.VersionID,
 		VersionName: jobResult.VersionName,
 		Uuid:        req.Uuid,
diff --git a/routers/repo/attachment.go b/routers/repo/attachment.go
index 5bfc4f9a2..2c082a289 100755
--- a/routers/repo/attachment.go
+++ b/routers/repo/attachment.go
@@ -1020,7 +1020,7 @@ func queryDatasets(ctx *context.Context, attachs []*models.AttachmentUsername) {
 }
 
 func checkTypeCloudBrain(typeCloudBrain int) error {
-	if typeCloudBrain != models.TypeCloudBrainOne && typeCloudBrain != models.TypeCloudBrainNotebook {
+	if typeCloudBrain != models.TypeCloudBrainOne && typeCloudBrain != models.TypeCloudBrainTwo {
 		log.Error("type error:", typeCloudBrain)
 		return errors.New("type error")
 	}
diff --git a/routers/repo/modelarts.go b/routers/repo/modelarts.go
index ea8ab0c06..f33d345cd 100755
--- a/routers/repo/modelarts.go
+++ b/routers/repo/modelarts.go
@@ -287,7 +287,8 @@ func NotebookIndex(ctx *context.Context) {
 			PageSize: setting.UI.IssuePagingNum,
 		},
 		RepoID: repo.ID,
-		Type:   models.TypeCloudBrainNotebook,
+		Type:   models.TypeCloudBrainTwo,
+		JobType: string(models.JobTypeDebug),
 	})
 	if err != nil {
 		ctx.ServerError("Cloudbrain", err)
@@ -512,7 +513,8 @@ func TrainJobIndex(ctx *context.Context) {
 			PageSize: setting.UI.IssuePagingNum,
 		},
 		RepoID: repo.ID,
-		Type:   models.TypeCloudBrainTrainJob,
+		Type:   models.TypeCloudBrainTwo,
+		JobType: string(models.JobTypeTrain),
 	})
 	if err != nil {
 		ctx.ServerError("Cloudbrain", err)