From 5fb14f452c1d2a84a2e68e1ca7e9f33bbd57d339 Mon Sep 17 00:00:00 2001
From: liuzx <liuzx@pcl.ac.cn>
Date: Fri, 22 Jul 2022 18:06:44 +0800
Subject: [PATCH] GetCloudbrainCardNumAndType

---
 models/cloudbrain.go | 80 +++++++++++++++++++++++++++++++++++++-------
 1 file changed, 67 insertions(+), 13 deletions(-)

diff --git a/models/cloudbrain.go b/models/cloudbrain.go
index 1b877c568..13649fe93 100755
--- a/models/cloudbrain.go
+++ b/models/cloudbrain.go
@@ -2022,6 +2022,9 @@ var (
 	CloudbrainBenchmarkResourceSpecsMap map[int]*ResourceSpec
 	ModelArtsDebugResourceSpecsMap      map[string]*FlavorInfo
 	ModelArtsTrainResourceSpecsMap      map[string]*FlavorInfo
+	GpuInfosMapInitFlag                 = false
+	CloudbrainDebugGpuInfosMap          map[string]*GpuInfo
+	CloudbrainTrainGpuInfosMap          map[string]*GpuInfo
 )
 
 type ModelArtsFlavor struct {
@@ -2080,6 +2083,26 @@ func InitResourceSpecMap() {
 	SpecsMapInitFlag = true
 }
 
+func InitGpuInfoMap() {
+	if CloudbrainDebugGpuInfosMap == nil || len(CloudbrainDebugGpuInfosMap) == 0 {
+		t := GpuInfos{}
+		json.Unmarshal([]byte(setting.GpuTypes), &t)
+		CloudbrainDebugGpuInfosMap = make(map[int]*GpuInfo, len(t.GpuInfo))
+		for _, GpuInfo := range t.GpuInfo {
+			CloudbrainDebugGpuInfosMap[GpuInfo.Queue] = GpuInfo
+		}
+	}
+	if CloudbrainTrainGpuInfosMap == nil || len(CloudbrainTrainGpuInfosMap) == 0 {
+		t := GpuInfos{}
+		json.Unmarshal([]byte(setting.TrainGpuTypes), &t)
+		CloudbrainTrainGpuInfosMap = make(map[int]*GpuInfo, len(t.GpuInfo))
+		for _, GpuInfo := range t.GpuInfo {
+			CloudbrainTrainGpuInfosMap[GpuInfo.Queue] = GpuInfo
+		}
+	}
+	GpuInfosMapInitFlag = true
+}
+
 type ResourceAndFlavor struct {
 	ResourceSpec *ResourceSpec
 	FlavorInfo   *FlavorInfo
@@ -2121,21 +2144,52 @@ func GetCloudbrainResourceSpec(jobType string, clusterType int, resourceSpecId i
 	return nil
 
 }
-func GetCloudbrainTaskCardNum(task Cloudbrain) int {
-	spec := GetCloudbrainResourceSpec(task.JobType, task.Type, task.ResourceSpecId, task.FlavorCode)
-	if spec == nil {
-		return 1
+
+// func GetCloudbrainTaskCardNum(task Cloudbrain) int {
+// 	spec := GetCloudbrainResourceSpec(task.JobType, task.Type, task.ResourceSpecId, task.FlavorCode)
+// 	if spec == nil {
+// 		return 0
+// 	}
+// 	if task.Type == TypeCloudBrainOne {
+// 		if spec.ResourceSpec == nil {
+// 			return 0
+// 		}
+// 		return spec.ResourceSpec.GpuNum
+// 	} else if task.Type == TypeCloudBrainTwo || task.Type == TypeC2Net {
+// 		return getCardnumByFlavorname(task.FlavorName)
+// 	}
+// 	return 0
+// }
+
+func getCardNumAndTypeByFlavorname(FlavorName string) (int, string) {
+	var beginIndex = strings.Index(FlavorName, ":")
+	var lastIndex = strings.LastIndex(FlavorName, ":")
+	var endIndex = strings.Index(FlavorName, "*")
+	cardNum, err := strconv.Atoi(strings.TrimSpace(FlavorName[beginIndex:endIndex]))
+	if err != nil {
+		log.Error("strconv.Atoi failed: %v", err)
+		return 0, ""
+	}
+	cardType := strings.TrimSpace(FlavorName[endIndex:lastIndex])
+	return cardNum, cardType
+}
+
+func GetCloudbrainCardNumAndType(task Cloudbrain) (int, string) {
+	if !SpecsMapInitFlag {
+		InitResourceSpecMap()
+	}
+	if !GpuInfosMapInitFlag {
+		InitGpuInfoMap()
 	}
 	if task.Type == TypeCloudBrainOne {
-		if spec.ResourceSpec == nil {
-			return 1
-		}
-		return spec.ResourceSpec.GpuNum
-	} else if task.Type == TypeCloudBrainTwo {
-		if spec.FlavorInfo == nil {
-			return 1
+		switch task.JobType {
+		case string(JobTypeDebug):
+			return CloudbrainDebugResourceSpecsMap[task.ResourceSpecId].GpuNum, CloudbrainDebugGpuInfosMap[task.GpuQueue].Value
+		case string(JobTypeTrain):
+			return CloudbrainTrainResourceSpecsMap[task.ResourceSpecId].GpuNum, CloudbrainTrainGpuInfosMap[task.GpuQueue].Value
 		}
-		return spec.FlavorInfo.UnitPrice
+	} else if task.Type == TypeCloudBrainTwo || task.Type == TypeC2Net {
+		return getCardNumAndTypeByFlavorname(task.FlavorName)
 	}
-	return 1
+	return 0, ""
 }