From 909dfea677cfe62b441bd89cfaa9c611eb2d9836 Mon Sep 17 00:00:00 2001
From: liuzx <liuzx@pcl.ac.cn>
Date: Tue, 26 Jul 2022 12:03:50 +0800
Subject: [PATCH] fix-2454

---
 models/cloudbrain.go       | 27 +++++++++++++++++++++++++++
 routers/repo/cloudbrain.go | 10 +++++++---
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index 4d81e65a7..200ae83ad 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -2097,10 +2097,13 @@ var (
 	SpecsMapInitFlag                    = false
 	CloudbrainDebugResourceSpecsMap     map[int]*ResourceSpec
 	CloudbrainTrainResourceSpecsMap     map[int]*ResourceSpec
+	CloudbrainInferenceResourceSpecsMap map[int]*ResourceSpec
 	CloudbrainBenchmarkResourceSpecsMap map[int]*ResourceSpec
 	GpuInfosMapInitFlag                 = false
 	CloudbrainDebugGpuInfosMap          map[string]*GpuInfo
 	CloudbrainTrainGpuInfosMap          map[string]*GpuInfo
+	CloudbrainInferenceGpuInfosMap      map[string]*GpuInfo
+	CloudbrainBenchmarkGpuInfosMap      map[string]*GpuInfo
 )
 
 func InitCloudbrainOneResourceSpecMap() {
@@ -2120,6 +2123,14 @@ func InitCloudbrainOneResourceSpecMap() {
 			CloudbrainTrainResourceSpecsMap[spec.Id] = spec
 		}
 	}
+	if CloudbrainInferenceResourceSpecsMap == nil || len(CloudbrainInferenceResourceSpecsMap) == 0 {
+		t := ResourceSpecs{}
+		json.Unmarshal([]byte(setting.InferenceResourceSpecs), &t)
+		CloudbrainInferenceResourceSpecsMap = make(map[int]*ResourceSpec, len(t.ResourceSpec))
+		for _, spec := range t.ResourceSpec {
+			CloudbrainInferenceResourceSpecsMap[spec.Id] = spec
+		}
+	}
 	if CloudbrainBenchmarkResourceSpecsMap == nil || len(CloudbrainBenchmarkResourceSpecsMap) == 0 {
 		t := ResourceSpecs{}
 		json.Unmarshal([]byte(setting.BenchmarkResourceSpecs), &t)
@@ -2148,6 +2159,22 @@ func InitCloudbrainOneGpuInfoMap() {
 			CloudbrainTrainGpuInfosMap[GpuInfo.Queue] = GpuInfo
 		}
 	}
+	if CloudbrainInferenceGpuInfosMap == nil || len(CloudbrainInferenceGpuInfosMap) == 0 {
+		t := GpuInfos{}
+		json.Unmarshal([]byte(setting.InferenceGpuTypes), &t)
+		CloudbrainInferenceGpuInfosMap = make(map[string]*GpuInfo, len(t.GpuInfo))
+		for _, GpuInfo := range t.GpuInfo {
+			CloudbrainInferenceGpuInfosMap[GpuInfo.Queue] = GpuInfo
+		}
+	}
+	if CloudbrainBenchmarkGpuInfosMap == nil || len(CloudbrainBenchmarkGpuInfosMap) == 0 {
+		t := GpuInfos{}
+		json.Unmarshal([]byte(setting.BenchmarkGpuTypes), &t)
+		CloudbrainBenchmarkGpuInfosMap = make(map[string]*GpuInfo, len(t.GpuInfo))
+		for _, GpuInfo := range t.GpuInfo {
+			CloudbrainBenchmarkGpuInfosMap[GpuInfo.Queue] = GpuInfo
+		}
+	}
 	GpuInfosMapInitFlag = true
 }
 func GetNewestJobsByAiCenter() ([]int64, error) {
diff --git a/routers/repo/cloudbrain.go b/routers/repo/cloudbrain.go
index 0d2998d90..53f465fae 100755
--- a/routers/repo/cloudbrain.go
+++ b/routers/repo/cloudbrain.go
@@ -373,7 +373,6 @@ func CloudBrainCreate(ctx *context.Context, form auth.CreateCloudBrainForm) {
 	}
 }
 
-
 func CloudBrainInferenceJobCreate(ctx *context.Context, form auth.CreateCloudBrainInferencForm) {
 	ctx.Data["PageIsCloudBrain"] = true
 	displayJobName := form.DisplayJobName
@@ -494,6 +493,7 @@ func CloudBrainInferenceJobCreate(ctx *context.Context, form auth.CreateCloudBra
 	ctx.Redirect(setting.AppSubURL + ctx.Repo.RepoLink + "/modelarts/inference-job")
 
 }
+
 /**
   检查用户传输的参数是否符合专属资源池
 */
@@ -2665,6 +2665,7 @@ func GetBenchmarkTypes(ctx *context.Context) *models.BenchmarkTypes {
 	}
 	return benchmarkTypesMap[lang]
 }
+
 func GetCloudbrainAiCenter(task models.Cloudbrain, ctx *context.Context) string {
 	if task.Type == models.TypeCloudBrainOne {
 		return ctx.Tr("repo.cloudbrain1")
@@ -2792,11 +2793,14 @@ func getCloudBrainOneResourceSpec(task models.Cloudbrain) (*models.ResourceSpec,
 			return models.CloudbrainTrainResourceSpecsMap[task.ResourceSpecId], models.CloudbrainTrainGpuInfosMap[GpuQueueDefault], nil
 		} else if task.JobType == string(models.JobTypeDebug) {
 			return models.CloudbrainDebugResourceSpecsMap[task.ResourceSpecId], models.CloudbrainDebugGpuInfosMap[GpuQueueDefault], nil
-		} else {
-			return models.CloudbrainDebugResourceSpecsMap[task.ResourceSpecId], models.CloudbrainDebugGpuInfosMap[GpuQueueDefault], nil
+		} else if task.JobType == string(models.JobTypeInference) {
+			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
 		}
 	} else {
 		err := errors.New("ResourceSpecId is null")
 		return nil, nil, err
 	}
+	return nil, nil, nil
 }