Merge pull request '不同集群的训练任务共用1个任务的数量限制' (#3577) from fix-2417 into V20230215

Reviewed-on: https://openi.pcl.ac.cn/OpenI/aiforge/pulls/3577
Reviewed-by: ychao_1983 <ychao_1983@sina.com>

models/cloudbrain.go

@@ -2315,11 +2315,10 @@ func GetWaitingCloudbrainCount(cloudbrainType int, computeResource string, jobTy
 	}
 	return sess.Count(new(Cloudbrain))
 }
-func GetNotFinalStatusTaskCount(userID int64, notFinalStatus []string, jobTypes []JobType, cloudbrainTypes []int, computeResource string) (int, error) {
+func GetNotFinalStatusTaskCount(userID int64, notFinalStatus []string, jobTypes []JobType) (int, error) {
 	count, err := x.In("status", notFinalStatus).
 		In("job_type", jobTypes).
-		In("type", cloudbrainTypes).
-		And("user_id = ? and compute_resource = ?", userID, computeResource).Count(new(Cloudbrain))
+		And("user_id = ? ", userID).Count(new(Cloudbrain))
 	return int(count), err
 }
 

modules/redis/redis_key/cloudbrain_redis_key.go

@@ -1,7 +1,13 @@
 package redis_key
 
+import "fmt"
+
 const CLOUDBRAIN_PREFIX = "cloudbrain"
 
 func CloudbrainBindingJobNameKey(repoId string, jobType string, jobName string) string {
 	return KeyJoin(CLOUDBRAIN_PREFIX, repoId, jobType, jobName, "redis_key")
 }
+
+func CloudbrainUniquenessKey(userId int64, jobType string) string {
+	return KeyJoin(CLOUDBRAIN_PREFIX, fmt.Sprint(userId), jobType, "uniqueness")
+}

modules/setting/setting.go

@@ -690,6 +690,8 @@ var (
 	IncubationSourceOrgName string
 	PaperRepoTopicName      string
 
+	CloudbrainUniquenessLockTime time.Duration
+
 	//nginx proxy
 	PROXYURL string
 	RadarMap = struct {
@@ -1506,6 +1508,8 @@ func NewContext() {
 	CullInterval = sec.Key("CULL_INTERVAL").MustString("60")
 	DebugAttachSize = sec.Key("DEBUG_ATTACH_SIZE").MustInt(20)
 
+	CloudbrainUniquenessLockTime = sec.Key("UNIQUENESS_LOCK_TIME").MustDuration(5 * time.Minute)
+
 	sec = Cfg.Section("benchmark")
 	IsBenchmarkEnabled = sec.Key("ENABLED").MustBool(false)
 	BenchmarkOwner = sec.Key("OWNER").MustString("")

routers/repo/aisafety.go

@@ -2,6 +2,7 @@ package repo
 
 import (
 	"bufio"
+	"code.gitea.io/gitea/services/lock"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -575,6 +576,20 @@ func AiSafetyCreateForPost(ctx *context.Context) {
 		tpname = tplCloudBrainModelSafetyNewGpu
 	}
 
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: string(models.JobTypeInference)}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
+		modelSafetyNewDataPrepare(ctx)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tpname, nil)
+		return
+	}
+
 	tasks, err := models.GetCloudbrainsByDisplayJobName(repo.ID, string(models.JobTypeModelSafety), displayJobName)
 	if err == nil {
 		if len(tasks) != 0 {

routers/repo/cloudbrain.go

@@ -2,6 +2,7 @@ package repo
 
 import (
 	"bufio"
+	"code.gitea.io/gitea/services/lock"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -40,8 +41,6 @@ import (
 	"code.gitea.io/gitea/modules/git"
 	"code.gitea.io/gitea/modules/log"
 	"code.gitea.io/gitea/modules/modelarts"
-	"code.gitea.io/gitea/modules/redis/redis_key"
-	"code.gitea.io/gitea/modules/redis/redis_lock"
 	"code.gitea.io/gitea/modules/setting"
 	"code.gitea.io/gitea/modules/storage"
 	"code.gitea.io/gitea/modules/util"
@@ -228,16 +227,19 @@ func cloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	if jobType == string(models.JobTypeTrain) {
 		tpl = tplCloudBrainTrainJobNew
 	}
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: form.JobType}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), jobType, displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
 		cloudBrainNewDataPrepare(ctx, jobType)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tpl, &form)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tpl, &form)
 		return
 	}
-	defer lock.UnLock()
 
 	tasks, err := models.GetCloudbrainsByDisplayJobName(repo.ID, jobType, displayJobName)
 	if err == nil {
@@ -462,15 +464,19 @@ func CloudBrainInferenceJobCreate(ctx *context.Context, form auth.CreateCloudBra
 	repo := ctx.Repo.Repository
 	tpl := tplCloudBrainInferenceJobNew
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), jobType, displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
-		cloudBrainNewDataPrepare(ctx, jobType)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tpl, &form)
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: jobType}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
+		grampusTrainJobNewDataPrepare(ctx, grampus.ProcessorTypeNPU)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tpl, &form)
 		return
 	}
-	defer lock.UnLock()
 
 	ckptUrl := setting.Attachment.Minio.RealPath + form.TrainUrl + form.CkptName
 	log.Info("ckpt url:" + ckptUrl)
@@ -647,7 +653,24 @@ func CloudBrainRestart(ctx *context.Context) {
 	var errorMsg = ""
 	var status = string(models.JobWaiting)
 	task := ctx.Cloudbrain
+
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainRestart(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{JobType: task.JobType}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
+		resultCode = "-1"
+		errorMsg = ctx.Tr(errMsg)
+	}
 	for {
+		if errorMsg != "" {
+			break
+		}
+
 		if task.Status != string(models.JobStopped) && task.Status != string(models.JobSucceeded) && task.Status != string(models.JobFailed) {
 			log.Error("the job(%s) is not stopped", task.JobName, ctx.Data["MsgID"])
 			resultCode = "-1"
@@ -2387,15 +2410,19 @@ func BenchMarkAlgorithmCreate(ctx *context.Context, form auth.CreateCloudBrainFo
 	ctx.Data["benchmarkTypeID"] = benchmarkTypeID
 	ctx.Data["benchmark_child_types_id_hidden"] = benchmarkChildTypeID
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), form.JobType, displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: form.JobType}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
 		cloudBrainNewDataPrepare(ctx, jobType)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tplCloudBrainBenchmarkNew, &form)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tplCloudBrainBenchmarkNew, &form)
 		return
 	}
-	defer lock.UnLock()
 
 	tasks, err := models.GetCloudbrainsByDisplayJobName(repo.ID, jobType, displayJobName)
 	if err == nil {
@@ -2576,15 +2603,19 @@ func ModelBenchmarkCreate(ctx *context.Context, form auth.CreateCloudBrainForm)
 	tpl := tplCloudBrainBenchmarkNew
 	command := cloudbrain.GetCloudbrainDebugCommand()
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), jobType, displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
-		cloudBrainNewDataPrepare(ctx, jobType)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tpl, &form)
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: jobType}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
+		grampusTrainJobNewDataPrepare(ctx, grampus.ProcessorTypeNPU)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tpl, &form)
 		return
 	}
-	defer lock.UnLock()
 
 	tasks, err := models.GetCloudbrainsByDisplayJobName(repo.ID, jobType, displayJobName)
 	if err == nil {

routers/repo/grampus.go

@@ -1,6 +1,7 @@
 package repo
 
 import (
+	"code.gitea.io/gitea/services/lock"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -27,8 +28,6 @@ import (
 	"code.gitea.io/gitea/modules/grampus"
 	"code.gitea.io/gitea/modules/modelarts"
 	"code.gitea.io/gitea/modules/notification"
-	"code.gitea.io/gitea/modules/redis/redis_key"
-	"code.gitea.io/gitea/modules/redis/redis_lock"
 	"code.gitea.io/gitea/modules/timeutil"
 	"code.gitea.io/gitea/modules/util"
 	"github.com/unknwon/com"
@@ -131,13 +130,17 @@ func GrampusNotebookCreate(ctx *context.Context, form auth.CreateGrampusNotebook
 		codeStoragePath = setting.CBCodePathPrefix + jobName + cloudbrain.CodeMountPath + "/"
 	}
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), string(models.JobTypeDebug), displayJobName))
-	defer lock.UnLock()
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: string(models.JobTypeDebug)}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
 		grampusNotebookNewDataPrepare(ctx, processType)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tpl, &form)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tpl, &form)
 		return
 	}
 
@@ -514,19 +517,23 @@ func grampusTrainJobGpuCreate(ctx *context.Context, form auth.CreateGrampusTrain
 	image := strings.TrimSpace(form.Image)
 	tpl := tplGrampusTrainJobGPUNew
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), string(models.JobTypeTrain), displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
+	if !jobNamePattern.MatchString(displayJobName) {
 		grampusTrainJobNewDataPrepare(ctx, grampus.ProcessorTypeGPU)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tplGrampusTrainJobGPUNew, &form)
+		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_jobname_err"), tpl, &form)
 		return
 	}
-	defer lock.UnLock()
 
-	if !jobNamePattern.MatchString(displayJobName) {
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: string(models.JobTypeTrain)}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
 		grampusTrainJobNewDataPrepare(ctx, grampus.ProcessorTypeGPU)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_jobname_err"), tpl, &form)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tpl, &form)
 		return
 	}
 
@@ -762,19 +769,23 @@ func grampusTrainJobNpuCreate(ctx *context.Context, form auth.CreateGrampusTrain
 	engineName := form.EngineName
 	tpl := tplGrampusTrainJobNPUNew
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), string(models.JobTypeTrain), displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
+	if !jobNamePattern.MatchString(displayJobName) {
 		grampusTrainJobNewDataPrepare(ctx, grampus.ProcessorTypeNPU)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tplGrampusTrainJobNPUNew, &form)
+		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_jobname_err"), tpl, &form)
 		return
 	}
-	defer lock.UnLock()
 
-	if !jobNamePattern.MatchString(displayJobName) {
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: string(models.JobTypeTrain)}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
 		grampusTrainJobNewDataPrepare(ctx, grampus.ProcessorTypeNPU)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_jobname_err"), tpl, &form)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tpl, &form)
 		return
 	}
 
@@ -1663,7 +1674,22 @@ func GrampusNotebookRestart(ctx *context.Context) {
 		return
 	}
 
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainRestart(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{JobType: task.JobType}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
+		errorMsg = ctx.Tr(errMsg)
+	}
+
 	for {
+		if errorMsg != "" {
+			break
+		}
 
 		if task.Status != models.GrampusStatusStopped && task.Status != models.GrampusStatusSucceeded && task.Status != models.GrampusStatusFailed {
 			log.Error("the job(%s) is not stopped", task.JobName, ctx.Data["MsgID"])

routers/repo/modelarts.go

@@ -2,6 +2,7 @@ package repo
 
 import (
 	"archive/zip"
+	"code.gitea.io/gitea/services/lock"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -35,8 +36,6 @@ import (
 	"code.gitea.io/gitea/modules/modelarts"
 	"code.gitea.io/gitea/modules/notification"
 	"code.gitea.io/gitea/modules/obs"
-	"code.gitea.io/gitea/modules/redis/redis_key"
-	"code.gitea.io/gitea/modules/redis/redis_lock"
 	"code.gitea.io/gitea/modules/setting"
 	"code.gitea.io/gitea/modules/storage"
 	"code.gitea.io/gitea/modules/timeutil"
@@ -176,15 +175,19 @@ func Notebook2Create(ctx *context.Context, form auth.CreateModelArtsNotebookForm
 	imageId := form.ImageId
 	repo := ctx.Repo.Repository
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), string(models.JobTypeDebug), displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: string(models.JobTypeDebug)}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
 		notebookNewDataPrepare(ctx)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tplModelArtsNotebookNew, &form)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tplModelArtsNotebookNew, &form)
 		return
 	}
-	defer lock.UnLock()
 
 	count, err := cloudbrainTask.GetNotFinalStatusTaskCount(ctx.User.ID, models.TypeCloudBrainTwo, string(models.JobTypeDebug))
 
@@ -494,8 +497,21 @@ func NotebookRestart(ctx *context.Context) {
 	var spec *models.Specification
 
 	task := ctx.Cloudbrain
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainRestart(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{JobType: task.JobType}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
 
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
+		errorMsg = ctx.Tr(errMsg)
+	}
 	for {
+		if errMsg != "" {
+			break
+		}
 		ctx.CheckWechatBind()
 		if ctx.Written() {
 			return
@@ -1083,15 +1099,19 @@ func TrainJobCreate(ctx *context.Context, form auth.CreateModelArtsTrainJobForm)
 		return
 	}
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), string(models.JobTypeTrain), displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: string(models.JobTypeTrain)}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
 		trainJobNewDataPrepare(ctx)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tplModelArtsTrainJobNew, &form)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tplModelArtsTrainJobNew, &form)
 		return
 	}
-	defer lock.UnLock()
 
 	count, err := cloudbrainTask.GetNotFinalStatusTaskCount(ctx.User.ID, models.TypeCloudBrainTwo, string(models.JobTypeTrain))
 
@@ -1461,6 +1481,20 @@ func TrainJobCreateVersion(ctx *context.Context, form auth.CreateModelArtsTrainJ
 	ctx.Data["PageIsTrainJob"] = true
 	var jobID = ctx.Params(":jobid")
 
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: form.DisplayJobName, JobType: string(models.JobTypeTrain)}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
+		trainJobNewVersionDataPrepare(ctx)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tplModelArtsTrainJobVersionNew, &form)
+		return
+	}
+
 	errStr := checkMultiNode(ctx.User.ID, form.WorkServerNumber)
 	if errStr != "" {
 		trainJobNewVersionDataPrepare(ctx)
@@ -1512,16 +1546,6 @@ func TrainJobCreateVersion(ctx *context.Context, form auth.CreateModelArtsTrainJ
 	EngineName := form.EngineName
 	isLatestVersion := modelarts.IsLatestVersion
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), string(models.JobTypeTrain), displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
-		trainJobNewVersionDataPrepare(ctx)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tplModelArtsTrainJobVersionNew, &form)
-		return
-	}
-	defer lock.UnLock()
-
 	canNewJob, _ := canUserCreateTrainJobVersion(ctx, latestTask.UserID)
 	if !canNewJob {
 		trainJobNewVersionDataPrepare(ctx)
@@ -2109,15 +2133,19 @@ func InferenceJobCreate(ctx *context.Context, form auth.CreateModelArtsInference
 		return
 	}
 
-	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(repo.ID), string(models.JobTypeInference), displayJobName))
-	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
-	if !isOk {
-		log.Error("lock processed failed:%v", err, ctx.Data["MsgID"])
+	lockOperator, errMsg := cloudbrainService.Lock4CloudbrainCreation(&lock.LockContext{Repo: ctx.Repo.Repository, Task: &models.Cloudbrain{DisplayJobName: displayJobName, JobType: string(models.JobTypeInference)}, User: ctx.User})
+	defer func() {
+		if lockOperator != nil {
+			lockOperator.Unlock()
+		}
+	}()
+
+	if errMsg != "" {
+		log.Error("lock processed failed:%s", errMsg, ctx.Data["MsgID"])
 		inferenceJobErrorNewDataPrepare(ctx, form)
-		ctx.RenderWithErr(ctx.Tr("repo.cloudbrain_samejob_err"), tplModelArtsInferenceJobNew, &form)
+		ctx.RenderWithErr(ctx.Tr(errMsg), tplModelArtsInferenceJobNew, &form)
 		return
 	}
-	defer lock.UnLock()
 
 	count, err := cloudbrainTask.GetNotFinalStatusTaskCount(ctx.User.ID, models.TypeCloudBrainTwo, string(models.JobTypeInference))
 

services/cloudbrain/cloudbrainTask/count.go

@@ -16,7 +16,7 @@ type StatusInfo struct {
 
 var CloudbrainOneNotFinalStatuses = []string{string(models.JobWaiting), string(models.JobRunning)}
 var CloudbrainTwoNotFinalStatuses = []string{string(models.ModelArtsTrainJobInit), string(models.ModelArtsTrainJobImageCreating), string(models.ModelArtsTrainJobSubmitTrying), string(models.ModelArtsTrainJobWaiting), string(models.ModelArtsTrainJobRunning), string(models.ModelArtsTrainJobScaling), string(models.ModelArtsTrainJobCheckInit), string(models.ModelArtsTrainJobCheckRunning), string(models.ModelArtsTrainJobCheckRunningCompleted)}
-var GrampusNotFinalStatuses = []string{models.GrampusStatusWaiting, models.GrampusStatusRunning}
+var GrampusNotFinalStatuses = []string{models.GrampusStatusWaiting, models.GrampusStatusRunning, models.GrampusStatusPending}
 var StatusInfoDict = map[string]StatusInfo{string(models.JobTypeDebug) + "-" + strconv.Itoa(models.TypeCloudBrainOne): {
 	CloudBrainTypes:  []int{models.TypeCloudBrainOne},
 	JobType:          []models.JobType{models.JobTypeDebug},
@@ -92,7 +92,7 @@ func GetNotFinalStatusTaskCount(uid int64, cloudbrainType int, jobType string, c
 
 	if statusInfo, ok := StatusInfoDict[key]; ok {
 
-		return models.GetNotFinalStatusTaskCount(uid, statusInfo.NotFinalStatuses, statusInfo.JobType, statusInfo.CloudBrainTypes, statusInfo.ComputeResource)
+		return models.GetNotFinalStatusTaskCount(uid, statusInfo.NotFinalStatuses, statusInfo.JobType)
 
 	} else {
 		return 0, fmt.Errorf("Can not find the status info.")

services/cloudbrain/lock.go

@@ -0,0 +1,20 @@
+package cloudbrain
+
+import "code.gitea.io/gitea/services/lock"
+
+func Lock4CloudbrainCreation(ctx *lock.LockContext) (*lock.LockChainOperator, string) {
+	op := lock.NewLockChainOperator(ctx).Add(lock.CloudbrainUniquenessLock{}).Add(lock.CloudbrainDisplayJobNameLock{})
+	errCode := op.Lock()
+	if errCode != "" {
+		return nil, errCode
+	}
+	return op, ""
+}
+func Lock4CloudbrainRestart(ctx *lock.LockContext) (*lock.LockChainOperator, string) {
+	op := lock.NewLockChainOperator(ctx).Add(lock.CloudbrainUniquenessLock{})
+	errCode := op.Lock()
+	if errCode != "" {
+		return nil, errCode
+	}
+	return op, ""
+}

services/lock/cloudbrain_name_lock.go

@@ -0,0 +1,31 @@
+package lock
+
+import (
+	"code.gitea.io/gitea/models"
+	"code.gitea.io/gitea/modules/log"
+	"code.gitea.io/gitea/modules/redis/redis_key"
+	"code.gitea.io/gitea/modules/redis/redis_lock"
+	"fmt"
+)
+
+type CloudbrainDisplayJobNameLock struct {
+}
+
+func (c CloudbrainDisplayJobNameLock) IsMatch(ctx *LockContext) bool {
+	return true
+}
+
+func (c CloudbrainDisplayJobNameLock) Lock(ctx *LockContext) string {
+	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(ctx.Repo.ID), ctx.Task.JobType, ctx.Task.DisplayJobName))
+	isOk, err := lock.Lock(models.CloudbrainKeyDuration)
+	if !isOk {
+		log.Error("CloudbrainDisplayJobNameLock lock failed:%v", err)
+		return "repo.cloudbrain_samejob_err"
+	}
+	return ""
+}
+
+func (c CloudbrainDisplayJobNameLock) Unlock(ctx *LockContext) error {
+	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainBindingJobNameKey(fmt.Sprint(ctx.Repo.ID), ctx.Task.JobType, ctx.Task.DisplayJobName))
+	return lock.UnLock()
+}

services/lock/cloudbrain_uniqueness_lock.go

@@ -0,0 +1,30 @@
+package lock
+
+import (
+	"code.gitea.io/gitea/modules/log"
+	"code.gitea.io/gitea/modules/redis/redis_key"
+	"code.gitea.io/gitea/modules/redis/redis_lock"
+	"code.gitea.io/gitea/modules/setting"
+)
+
+type CloudbrainUniquenessLock struct {
+}
+
+func (c CloudbrainUniquenessLock) IsMatch(ctx *LockContext) bool {
+	return true
+}
+
+func (c CloudbrainUniquenessLock) Lock(ctx *LockContext) string {
+	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainUniquenessKey(ctx.User.ID, ctx.Task.JobType))
+	isOk, err := lock.Lock(setting.CloudbrainUniquenessLockTime)
+	if !isOk {
+		log.Error("CloudbrainDisplayJobNameLock lock failed:%v", err)
+		return "repo.cloudbrain.morethanonejob"
+	}
+	return ""
+}
+
+func (c CloudbrainUniquenessLock) Unlock(ctx *LockContext) error {
+	lock := redis_lock.NewDistributeLock(redis_key.CloudbrainUniquenessKey(ctx.User.ID, ctx.Task.JobType))
+	return lock.UnLock()
+}

services/lock/lock.go

@@ -0,0 +1,17 @@
+package lock
+
+import (
+	"code.gitea.io/gitea/models"
+)
+
+type LockContext struct {
+	Repo *models.Repository
+	Task *models.Cloudbrain
+	User *models.User
+}
+
+type Lock interface {
+	IsMatch(ctx *LockContext) bool
+	Lock(ctx *LockContext) string
+	Unlock(ctx *LockContext) error
+}

services/lock/lock_operator.go

@@ -0,0 +1,42 @@
+package lock
+
+type LockChainOperator struct {
+	chainList  []Lock
+	lockedList []Lock
+	ctx        *LockContext
+}
+
+func NewLockChainOperator(ctx *LockContext) *LockChainOperator {
+	return &LockChainOperator{ctx: ctx}
+}
+
+func (b *LockChainOperator) Add(l Lock) *LockChainOperator {
+	b.chainList = append(b.chainList, l)
+	return b
+}
+
+func (b *LockChainOperator) Lock() string {
+	for i := 0; i < len(b.chainList); i++ {
+		l := b.chainList[i]
+		if !l.IsMatch(b.ctx) {
+			continue
+		}
+
+		if errCode := l.Lock(b.ctx); errCode != "" {
+			b.Unlock()
+			return errCode
+		}
+		b.lockedList = append(b.lockedList, l)
+	}
+	return ""
+}
+
+func (b *LockChainOperator) Unlock() error {
+	if b.chainList == nil || len(b.chainList) == 0 {
+		return nil
+	}
+	for j := len(b.lockedList) - 1; j >= 0; j-- {
+		b.lockedList[j].Unlock(b.ctx)
+	}
+	return nil
+}