fix bug: task count

models/cloudbrain.go

@@ -1194,8 +1194,8 @@ func GetCloudBrainUnStoppedJob() ([]*Cloudbrain, error) {
 		Find(&cloudbrains)
 }
 
-func GetCloudbrainCountByUserID(userID int64) (int, error) {
-	count, err := x.In("status", JobWaiting, JobRunning).And("job_type = ? and user_id = ? and type = ?", JobTypeDebug, userID, TypeCloudBrainOne).Count(new(Cloudbrain))
+func GetCloudbrainCountByUserID(userID int64, jobType string) (int, error) {
+	count, err := x.In("status", JobWaiting, JobRunning).And("job_type = ? and user_id = ? and type = ?", jobType, userID, TypeCloudBrainOne).Count(new(Cloudbrain))
 	return int(count), err
 }
 

modules/cloudbrain/cloudbrain.go

@@ -14,7 +14,10 @@ import (
 )
 
 const (
-	Command               = `pip3 install jupyterlab==2.2.5 -i https://pypi.tuna.tsinghua.edu.cn/simple;service ssh stop;jupyter lab --no-browser --ip=0.0.0.0 --allow-root --notebook-dir="/code" --port=80 --LabApp.token="" --LabApp.allow_origin="self https://cloudbrain.pcl.ac.cn"`
+	Command               = `pip3 install jupyterlab==2.2.5 -i https://pypi.tuna.tsinghua.edu.cn/simple;
+							service ssh stop;
+							jupyter lab --no-browser --ip=0.0.0.0 --allow-root --notebook-dir="/code" --port=80 --LabApp.token="" --LabApp.allow_origin="self https://cloudbrain.pcl.ac.cn"`
+	CommandBenchmark      = `echo "begin";python /code/test.py`
 	CodeMountPath         = "/code"
 	DataSetMountPath      = "/dataset"
 	ModelMountPath        = "/model"

routers/repo/cloudbrain.go

@@ -155,7 +155,6 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	ctx.Data["PageIsCloudBrain"] = true
 	jobName := form.JobName
 	image := form.Image
-	command := form.Command
 	uuid := form.Attachment
 	jobType := form.JobType
 	gpuQueue := form.GpuType
@@ -174,7 +173,14 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 		return
 	}
 
-	count, err := models.GetCloudbrainCountByUserID(ctx.User.ID)
+	var command string
+	if jobType == string(models.JobTypeBenchmark) {
+		command = cloudbrain.CommandBenchmark
+	} else {
+		command = cloudbrain.Command
+	}
+
+	count, err := models.GetCloudbrainCountByUserID(ctx.User.ID, jobType)
 	if err != nil {
 		log.Error("GetCloudbrainCountByUserID failed:%v", err, ctx.Data["MsgID"])
 		cloudBrainNewDataPrepare(ctx)
@@ -244,7 +250,12 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 		ctx.RenderWithErr(err.Error(), tplCloudBrainNew, &form)
 		return
 	}
-	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/debugjob?debugListType=all")
+
+	if jobType == string(models.JobTypeBenchmark) {
+		ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/debugjob?debugListType=all")
+	} else {
+		ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/debugjob?debugListType=all")
+	}
 }
 
 func CloudBrainRestart(ctx *context.Context) {
@@ -276,7 +287,7 @@ func CloudBrainRestart(ctx *context.Context) {
 			break
 		}
 
-		count, err := models.GetCloudbrainCountByUserID(ctx.User.ID)
+		count, err := models.GetCloudbrainCountByUserID(ctx.User.ID, string(models.JobTypeDebug))
 		if err != nil {
 			log.Error("GetCloudbrainCountByUserID failed:%v", err, ctx.Data["MsgID"])
 			resultCode = "-1"