diff --git a/models/cloudbrain.go b/models/cloudbrain.go index f74fdd825..5091a8762 100755 --- a/models/cloudbrain.go +++ b/models/cloudbrain.go @@ -2140,11 +2140,13 @@ var ( CloudbrainTrainResourceSpecsMap map[int]*ResourceSpec CloudbrainInferenceResourceSpecsMap map[int]*ResourceSpec CloudbrainBenchmarkResourceSpecsMap map[int]*ResourceSpec + CloudbrainSpecialResourceSpecsMap map[int]*ResourceSpec GpuInfosMapInitFlag = false CloudbrainDebugGpuInfosMap map[string]*GpuInfo CloudbrainTrainGpuInfosMap map[string]*GpuInfo CloudbrainInferenceGpuInfosMap map[string]*GpuInfo CloudbrainBenchmarkGpuInfosMap map[string]*GpuInfo + CloudbrainSpecialGpuInfosMap map[string]*GpuInfo ) func InitCloudbrainOneResourceSpecMap() { @@ -2180,6 +2182,16 @@ func InitCloudbrainOneResourceSpecMap() { CloudbrainBenchmarkResourceSpecsMap[spec.Id] = spec } } + if CloudbrainSpecialResourceSpecsMap == nil || len(CloudbrainSpecialResourceSpecsMap) == 0 { + t := SpecialPools{} + json.Unmarshal([]byte(setting.SpecialPools), &t) + for _, pool := range t.Pools { + CloudbrainSpecialResourceSpecsMap = make(map[int]*ResourceSpec, len(pool.ResourceSpec)) + for _, spec := range pool.ResourceSpec { + CloudbrainSpecialResourceSpecsMap[spec.Id] = spec + } + } + } SpecsMapInitFlag = true } @@ -2216,6 +2228,16 @@ func InitCloudbrainOneGpuInfoMap() { CloudbrainBenchmarkGpuInfosMap[GpuInfo.Queue] = GpuInfo } } + if CloudbrainSpecialGpuInfosMap == nil || len(CloudbrainSpecialGpuInfosMap) == 0 { + t := SpecialPools{} + json.Unmarshal([]byte(setting.SpecialPools), &t) + for _, pool := range t.Pools { + CloudbrainSpecialGpuInfosMap = make(map[string]*GpuInfo, len(pool.Pool)) + for _, GpuInfo := range pool.Pool { + CloudbrainSpecialGpuInfosMap[GpuInfo.Queue] = GpuInfo + } + } + } GpuInfosMapInitFlag = true } func GetNewestJobsByAiCenter() ([]int64, error) { diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go index aa77e8d5c..31d0bb8d0 100755 --- a/routers/repo/cloudbrain.go +++ b/routers/repo/cloudbrain.go @@ -2980,10 +2980,18 @@ func GetCloudbrainAiCenter(task models.Cloudbrain, ctx *context.Context) string } else if task.Type == models.TypeCloudBrainTwo { return ctx.Tr("repo.cloudbrain2") } else if task.Type == models.TypeC2Net { - return task.AiCenter + return getCutStringAiCenterByAiCenter(task.AiCenter) } return "" } +func getCutStringAiCenterByAiCenter(aiCenter string) string { + if aiCenter == "" { + return "" + } + index := strings.LastIndex(aiCenter, "+") + return aiCenter[index+1:] + +} func GetCloudbrainCluster(task models.Cloudbrain, ctx *context.Context) string { if task.Type == models.TypeCloudBrainOne || task.Type == models.TypeCloudBrainTwo { return ctx.Tr("cloudbrain.resource_cluster_openi") @@ -3092,19 +3100,27 @@ func GetCloudbrainFlavorName(task models.Cloudbrain) (string, error) { } func getCloudBrainOneResourceSpec(task models.Cloudbrain) (*models.ResourceSpec, *models.GpuInfo, error) { - GpuQueueDefault := "openidebug" + gpuQueueDefault := "openidebug" if task.GpuQueue != "" { - GpuQueueDefault = task.GpuQueue + gpuQueueDefault = task.GpuQueue } if task.ResourceSpecId >= 0 { if task.JobType == string(models.JobTypeTrain) { - return models.CloudbrainTrainResourceSpecsMap[task.ResourceSpecId], models.CloudbrainTrainGpuInfosMap[GpuQueueDefault], nil + if models.CloudbrainTrainResourceSpecsMap[task.ResourceSpecId] != nil { + return models.CloudbrainTrainResourceSpecsMap[task.ResourceSpecId], models.CloudbrainTrainGpuInfosMap[gpuQueueDefault], nil + } else { + return models.CloudbrainSpecialResourceSpecsMap[task.ResourceSpecId], models.CloudbrainSpecialGpuInfosMap[gpuQueueDefault], nil + } } else if task.JobType == string(models.JobTypeDebug) { - return models.CloudbrainDebugResourceSpecsMap[task.ResourceSpecId], models.CloudbrainDebugGpuInfosMap[GpuQueueDefault], nil + if models.CloudbrainDebugResourceSpecsMap[task.ResourceSpecId] != nil { + return models.CloudbrainDebugResourceSpecsMap[task.ResourceSpecId], models.CloudbrainDebugGpuInfosMap[gpuQueueDefault], nil + } else { + return models.CloudbrainSpecialResourceSpecsMap[task.ResourceSpecId], models.CloudbrainSpecialGpuInfosMap[gpuQueueDefault], nil + } } else if task.JobType == string(models.JobTypeInference) { - return models.CloudbrainInferenceResourceSpecsMap[task.ResourceSpecId], models.CloudbrainInferenceGpuInfosMap[GpuQueueDefault], nil + return models.CloudbrainInferenceResourceSpecsMap[task.ResourceSpecId], models.CloudbrainInferenceGpuInfosMap[gpuQueueDefault], nil } else if task.JobType == string(models.JobTypeBenchmark) || task.JobType == string(models.JobTypeSnn4imagenet) || task.JobType == string(models.JobTypeBrainScore) { - return models.CloudbrainBenchmarkResourceSpecsMap[task.ResourceSpecId], models.CloudbrainBenchmarkGpuInfosMap[GpuQueueDefault], nil + return models.CloudbrainBenchmarkResourceSpecsMap[task.ResourceSpecId], models.CloudbrainBenchmarkGpuInfosMap[gpuQueueDefault], nil } } else { err := errors.New("ResourceSpecId is null")