diff --git a/modules/cloudbrain/cloudbrain.go b/modules/cloudbrain/cloudbrain.go index f9c7c5b5c..a2e71f8b2 100755 --- a/modules/cloudbrain/cloudbrain.go +++ b/modules/cloudbrain/cloudbrain.go @@ -259,21 +259,7 @@ func GenerateTask(req GenerateCloudBrainTaskReq) error { } //如果没有匹配到spec信息,尝试从专属资源池获取 if resourceSpec == nil && SpecialPools != nil { - for _, specialPool := range SpecialPools.Pools { - if resourceSpec != nil { - break - } - 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 - } - } - } - } - } + resourceSpec = geMatchResourceSpec(req.JobType, req.GpuQueue, req.ResourceSpecId) } if resourceSpec == nil { @@ -464,6 +450,11 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain, newID *string) e } } + //如果没有匹配到spec信息,尝试从专属资源池获取 + if resourceSpec == nil && SpecialPools != nil { + resourceSpec = geMatchResourceSpec(task.JobType, task.GpuQueue, task.ResourceSpecId) + } + if resourceSpec == nil { log.Error("no such resourceSpecId(%d)", task.ResourceSpecId, ctx.Data["MsgID"]) return errors.New("no such resourceSpec") @@ -602,6 +593,23 @@ func RestartTask(ctx *context.Context, task *models.Cloudbrain, newID *string) e return nil } +func geMatchResourceSpec(jobType string, gpuQueue string, resourceSpecId int) *models.ResourceSpec { + + for _, specialPool := range SpecialPools.Pools { + + if specialPool.ResourceSpec != nil { + if IsElementExist(specialPool.JobType, jobType) && IsQueueInSpecialtPool(specialPool.Pool, gpuQueue) { + for _, spec := range specialPool.ResourceSpec { + if resourceSpecId == spec.Id { + return spec + } + } + } + } + } + return nil +} + func DelCloudBrainJob(jobId string) string { task, err := models.GetCloudbrainByJobID(jobId) if err != nil { diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go index e8ea71ab7..5276b233e 100755 --- a/routers/repo/cloudbrain.go +++ b/routers/repo/cloudbrain.go @@ -134,7 +134,7 @@ func cloudBrainNewDataPrepare(ctx *context.Context) error { if inferenceGpuInfos == nil { json.Unmarshal([]byte(setting.InferenceGpuTypes), &inferenceGpuInfos) } - ctx.Data["train_gpu_types"] = trainGpuInfos.GpuInfo + ctx.Data["inference_gpu_types"] = inferenceGpuInfos.GpuInfo if benchmarkGpuInfos == nil { json.Unmarshal([]byte(setting.BenchmarkGpuTypes), &benchmarkGpuInfos) @@ -558,7 +558,18 @@ func CloudBrainRestart(ctx *context.Context) { for _, resourceType := range gpuInfos.GpuInfo { if resourceType.Queue == task.GpuQueue { hasSameResource = true - continue + break + } + } + if !hasSameResource && cloudbrain.SpecialPools != nil { + + for _, specialPool := range cloudbrain.SpecialPools.Pools { + cloudbrain.IsElementExist(specialPool.JobType, string(models.JobTypeDebug)) + for _, pool := range specialPool.Pool { + if pool.Queue == task.GpuQueue { + hasSameResource = true + } + } } }