diff --git a/models/cloudbrain.go b/models/cloudbrain.go index 2c23bf82d..ea1826de0 100755 --- a/models/cloudbrain.go +++ b/models/cloudbrain.go @@ -574,8 +574,8 @@ type SpecialPool struct { Type string `json:"type"` IsExclusive bool `json:"isExclusive"` Pool []*GpuInfo `json:"pool"` - JobType []string `json:"support_type"` - ResourceSpec []*ResourceSpec `json:"resource_specs"` + JobType []string `json:"jobType"` + ResourceSpec []*ResourceSpec `json:"resourceSpecs"` } type ImageInfosModelArts struct { diff --git a/modules/cloudbrain/cloudbrain.go b/modules/cloudbrain/cloudbrain.go index 09d6b9c0e..f48665cb0 100755 --- a/modules/cloudbrain/cloudbrain.go +++ b/modules/cloudbrain/cloudbrain.go @@ -223,6 +223,7 @@ func GenerateTask(req GenerateCloudBrainTaskReq) error { for _, spec := range TrainResourceSpecs.ResourceSpec { if req.ResourceSpecId == spec.Id { resourceSpec = spec + break } } } else { @@ -232,10 +233,26 @@ func GenerateTask(req GenerateCloudBrainTaskReq) error { for _, spec := range ResourceSpecs.ResourceSpec { if req.ResourceSpecId == spec.Id { resourceSpec = spec + break } } } + //如果没有匹配到spec信息,尝试从专属资源池获取 + if resourceSpec == nil && SpecialPools != nil { + for _, specialPool := range SpecialPools.Pools { + if specialPool.ResourceSpec != nil { + if IsElementExist(specialPool.JobType, req.JobType) && IsQueueInSpecialtPool(specialPool.Pool, req.GpuQueue) { + for _, spec := range specialPool.ResourceSpec { + if req.ResourceSpecId == spec.Id { + resourceSpec = spec + break + } + } + } + } + } + } if resourceSpec == nil { log.Error("no such resourceSpecId(%d)", req.ResourceSpecId, req.Ctx.Data["MsgID"]) @@ -545,3 +562,33 @@ func InitSpecialPool() { json.Unmarshal([]byte(setting.SpecialPools), &SpecialPools) } } + +func IsResourceSpecInSpecialPool(resourceSpecs []*models.ResourceSpec, resourceSpecId int) bool { + if resourceSpecs == nil || len(resourceSpecs) == 0 { + return true + } + for _, v := range resourceSpecs { + if v.Id == resourceSpecId { + return true + } + } + return false +} + +func IsQueueInSpecialtPool(pool []*models.GpuInfo, queue string) bool { + for _, v := range pool { + if v.Queue == queue { + return true + } + } + return false +} + +func IsElementExist(s []string, str string) bool { + for _, v := range s { + if v == str { + return true + } + } + return false +} diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go index 55776a2b4..17a7b0214 100755 --- a/routers/repo/cloudbrain.go +++ b/routers/repo/cloudbrain.go @@ -213,7 +213,7 @@ func cloudBrainNewDataPrepare(ctx *context.Context) error { } if len(trainGpuTypes) > 0 { - ctx.Data["train_gpu_types"] = trainGpuInfos + ctx.Data["train_gpu_types"] = trainGpuTypes } } @@ -257,6 +257,10 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) { repo := ctx.Repo.Repository tpl := tplCloudBrainNew + if jobType == string(models.JobTypeTrain) { + tpl = tplCloudBrainTrainJobNew + } + tasks, err := models.GetCloudbrainsByDisplayJobName(repo.ID, jobType, displayJobName) if err == nil { if len(tasks) != 0 { @@ -322,6 +326,13 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) { command = commandTrain } + errStr := checkCloudBrainSpecialPool(ctx, jobType, gpuQueue, resourceSpecId) + + if errStr != "" { + cloudBrainNewDataPrepare(ctx) + ctx.RenderWithErr(errStr, tpl, &form) + } + if branchName == "" { branchName = cloudbrain.DefaultBranchName } @@ -374,6 +385,42 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) { } } +/** + 检查用户传输的参数是否符合专属资源池 +*/ +func checkCloudBrainSpecialPool(ctx *context.Context, jobType string, queue string, resourceSpecId int) string { + if cloudbrain.SpecialPools != nil { + + var isInPoolOrg = false + var matchSpecialPool = false + + for _, specialPool := range cloudbrain.SpecialPools.Pools { + + if cloudbrain.IsElementExist(specialPool.JobType, jobType) && cloudbrain.IsQueueInSpecialtPool(specialPool.Pool, queue) { + if cloudbrain.IsResourceSpecInSpecialPool(specialPool.ResourceSpec, resourceSpecId) { + matchSpecialPool = true + org, _ := models.GetOrgByName(specialPool.Org) + if org != nil { + isInPoolOrg, _ = models.IsOrganizationMember(org.ID, ctx.User.ID) + if isInPoolOrg { + break //传入参数,和专属资源池匹配上了,检查通过 + } + } + } + + } + + } + //资源池有匹配上,但是用户不在相应的组织中,返回错误信息。界面已经过滤了选择,界面操作不会到这个逻辑 + if matchSpecialPool && !isInPoolOrg { + return ctx.Tr("repo.grampus.no_operate_right") + } + + } + //没有匹配到资源池或者没有设置专属资源池,检查通过; 获取和资源池完全匹配检查通过 + return "" +} + func CloudBrainRestart(ctx *context.Context) { var ID = ctx.Params(":id") var resultCode = "0"