diff --git a/models/cloudbrain.go b/models/cloudbrain.go index 68841544e..51e122bb6 100755 --- a/models/cloudbrain.go +++ b/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 } diff --git a/modules/redis/redis_key/cloudbrain_redis_key.go b/modules/redis/redis_key/cloudbrain_redis_key.go index 4c5d05dfc..a442d89f4 100644 --- a/modules/redis/redis_key/cloudbrain_redis_key.go +++ b/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") +} diff --git a/modules/setting/setting.go b/modules/setting/setting.go index 09e7259f2..44dff2779 100755 --- a/modules/setting/setting.go +++ b/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("") diff --git a/routers/repo/aisafety.go b/routers/repo/aisafety.go index 55f25dba6..27e123fac 100644 --- a/routers/repo/aisafety.go +++ b/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 { diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go index 64075e035..f220802bf 100755 --- a/routers/repo/cloudbrain.go +++ b/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 { diff --git a/routers/repo/grampus.go b/routers/repo/grampus.go index 74e2f06d6..339895b02 100755 --- a/routers/repo/grampus.go +++ b/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"]) diff --git a/routers/repo/modelarts.go b/routers/repo/modelarts.go index 86fbd5bdb..48536ae7e 100755 --- a/routers/repo/modelarts.go +++ b/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)) diff --git a/services/cloudbrain/cloudbrainTask/count.go b/services/cloudbrain/cloudbrainTask/count.go index 0010164ae..186d52fb9 100644 --- a/services/cloudbrain/cloudbrainTask/count.go +++ b/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.") diff --git a/services/cloudbrain/lock.go b/services/cloudbrain/lock.go new file mode 100644 index 000000000..c9d6080f1 --- /dev/null +++ b/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, "" +} diff --git a/services/lock/cloudbrain_name_lock.go b/services/lock/cloudbrain_name_lock.go new file mode 100644 index 000000000..bed388e01 --- /dev/null +++ b/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() +} diff --git a/services/lock/cloudbrain_uniqueness_lock.go b/services/lock/cloudbrain_uniqueness_lock.go new file mode 100644 index 000000000..8297e8b03 --- /dev/null +++ b/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() +} diff --git a/services/lock/lock.go b/services/lock/lock.go new file mode 100644 index 000000000..308cebd2a --- /dev/null +++ b/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 +} diff --git a/services/lock/lock_operator.go b/services/lock/lock_operator.go new file mode 100644 index 000000000..cd2002f1c --- /dev/null +++ b/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 +}