From 299b02ef8b69a0c1a6c224815c5cb4bc5cf5d53d Mon Sep 17 00:00:00 2001
From: yuyuanshifu <747342561@qq.com>
Date: Thu, 7 Jan 2021 11:15:34 +0800
Subject: [PATCH] add benchmark enabled

---
 models/cloudbrain.go             | 5 +++++
 modules/auth/cloudbrain.go       | 1 +
 modules/cloudbrain/cloudbrain.go | 3 ++-
 routers/repo/cloudbrain.go       | 5 +++--
 4 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index da6baa169..c6d4e30c9 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -13,6 +13,7 @@ import (
 )
 
 type CloudbrainStatus string
+type JobType string
 
 const (
 	JobWaiting   CloudbrainStatus = "WAITING"
@@ -20,11 +21,15 @@ const (
 	JobSucceeded CloudbrainStatus = "SUCCEEDED"
 	JobFailed    CloudbrainStatus = "FAILED"
 	JobRunning   CloudbrainStatus = "RUNNING"
+
+	JobTypeDebug		JobType = "DEBUG"
+	JobTypeBenchmark	JobType = "BENCHMARK"
 )
 
 type Cloudbrain struct {
 	ID          int64  `xorm:"pk autoincr"`
 	JobID       string `xorm:"INDEX NOT NULL"`
+	JobType		string `xorm:"INDEX NOT NULL"`
 	JobName     string			   `xorm:"INDEX"`
 	Status      string             `xorm:"INDEX"`
 	UserID      int64              `xorm:"INDEX"`
diff --git a/modules/auth/cloudbrain.go b/modules/auth/cloudbrain.go
index b7b59ce63..ce2733afb 100755
--- a/modules/auth/cloudbrain.go
+++ b/modules/auth/cloudbrain.go
@@ -11,6 +11,7 @@ type CreateCloudBrainForm struct {
 	Image   string `form:"image" binding:"Required"`
 	Command string `form:"command" binding:"Required"`
 	Attachment string `form:"attachment" binding:"Required"`
+	JobType string `form:"job_type" binding:"Required"`
 }
 
 type CommitImageCloudBrainForm struct {
diff --git a/modules/cloudbrain/cloudbrain.go b/modules/cloudbrain/cloudbrain.go
index 264efd6b4..dccbee5c6 100755
--- a/modules/cloudbrain/cloudbrain.go
+++ b/modules/cloudbrain/cloudbrain.go
@@ -22,7 +22,7 @@ const (
 	Success = "S000"
 )
 
-func GenerateTask(ctx *context.Context, jobName, image, command, uuid, codePath, modelPath, benchmarkPath string) error {
+func GenerateTask(ctx *context.Context, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, jobType string) error {
 	dataActualPath := setting.Attachment.Minio.RealPath +
 		setting.Attachment.Minio.Bucket + "/" +
 		setting.Attachment.Minio.BasePath +
@@ -97,6 +97,7 @@ func GenerateTask(ctx *context.Context, jobName, image, command, uuid, codePath,
 		JobID:   jobID,
 		JobName: jobName,
 		SubTaskName: SubTaskName,
+		JobType: jobType,
 	})
 
 	if err != nil {
diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index ed53e86d0..8b1e3c384 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -122,6 +122,7 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	image := form.Image
 	command := form.Command
 	uuid := form.Attachment
+	jobType := form.JobType
 	codePath := setting.JobPath + jobName + cloudbrain.CodeMountPath
 	repo := ctx.Repo.Repository
 	downloadCode(repo, codePath)
@@ -134,11 +135,11 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	}
 
 	benchmarkPath := setting.JobPath + jobName + cloudbrain.BenchMarkMountPath
-	if setting.IsBenchmarkEnabled {
+	if setting.IsBenchmarkEnabled && jobType == string(models.JobTypeBenchmark) {
 		downloadBenchmarkCode(repo, jobName, benchmarkPath)
 	}
 
-	err = cloudbrain.GenerateTask(ctx, jobName, image, command, uuid, codePath, modelPath, benchmarkPath)
+	err = cloudbrain.GenerateTask(ctx, jobName, image, command, uuid, codePath, modelPath, benchmarkPath, jobType)
 	if err != nil {
 		ctx.RenderWithErr(err.Error(), tplCloudBrainNew, &form)
 		return