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)})

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 {

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,

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")
 	}

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)