From 3d3c87767850bd09efbe52c45d1009a45be56a8a Mon Sep 17 00:00:00 2001
From: liuzx <liuzx@pcl.ac.cn>
Date: Thu, 28 Jul 2022 10:30:04 +0800
Subject: [PATCH] fix-2595

---
 models/cloudbrain.go       | 24 +++++++++++++++++++++++-
 routers/repo/cloudbrain.go | 20 ++++++++++++++------
 2 files changed, 37 insertions(+), 7 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index f74fdd825..9dc5f0285 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,7 +2182,17 @@ func InitCloudbrainOneResourceSpecMap() {
 			CloudbrainBenchmarkResourceSpecsMap[spec.Id] = spec
 		}
 	}
-	SpecsMapInitFlag = true
+	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
+	}
 }
 
 func InitCloudbrainOneGpuInfoMap() {
@@ -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..bd8b05088 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -3092,19 +3092,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")