From 575a5fe15c9a836f61594f526fe94069dbbdc657 Mon Sep 17 00:00:00 2001
From: qinguangrui <283713406@qq.com>
Date: Mon, 2 Dec 2024 10:40:20 +0800
Subject: [PATCH] =?UTF-8?q?FIX:=20=E4=BF=AE=E5=A4=8D=E7=94=A8=E6=88=B7?=
 =?UTF-8?q?=E7=A9=BA=E9=97=B4=E7=9A=84=E5=8D=95=E8=8A=82=E7=82=B9=E5=86=85?=
 =?UTF-8?q?=E5=AD=98=E4=B8=8D=E6=98=BE=E7=A4=BA=E7=9A=84=E9=97=AE=E9=A2=98?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 services/config/config.go | 24 ++++++++++-------------
 utils/utils.go            | 41 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 14 deletions(-)

diff --git a/services/config/config.go b/services/config/config.go
index cef7134..eba97e9 100644
--- a/services/config/config.go
+++ b/services/config/config.go
@@ -88,7 +88,6 @@ func (s *ServerConfig) GetClusterConfig(ctx context.Context, in *pb.GetClusterCo
 			totalGpus    uint32
 			comment      string
 			qos          []string
-			totalMems    int
 			totalCpus    string
 			totalMemsTmp string
 			totalNodes   string
@@ -173,20 +172,17 @@ func (s *ServerConfig) GetClusterConfig(ctx context.Context, in *pb.GetClusterCo
 				}
 			}
 
-			if strings.Contains(totalMemsTmp, "M") {
-				totalMemsInt, _ := strconv.Atoi(strings.Split(totalMemsTmp, "M")[0])
-				totalMems = totalMemsInt
-			} else if strings.Contains(totalMemsTmp, "G") {
-				totalMemsInt, _ := strconv.Atoi(strings.Split(totalMemsTmp, "G")[0])
-				totalMems = totalMemsInt * 1024
-			} else if strings.Contains(totalMemsTmp, "T") {
-				totalMemsInt, _ := strconv.Atoi(strings.Split(totalMemsTmp, "T")[0])
-				totalMems = totalMemsInt * 1024 * 1024
+			memString := strings.Split(totalMemsTmp, "M")[0]
+			totalMemInt, err = utils.ConvertMemory(memString)
+			if err != nil {
+				errInfo := &errdetails.ErrorInfo{
+					Reason: "CONVERT_MEMORY_FAILED",
+				}
+				st := status.New(codes.Internal, "convert memory error")
+				st, _ = st.WithDetails(errInfo)
+				caller.Logger.Errorf("GetClusterConfig failed: %v", st.Err())
+				return nil, st.Err()
 			}
-
-			// 将字符串转换为int
-			totalCpuInt, _ = strconv.Atoi(totalCpus)
-			totalMemInt = totalMems
 			totalNodeNumInt, _ = strconv.Atoi(totalNodes)
 		} else if err != nil && !utils.CheckSlurmStatus(output) {
 			// 获取总cpu、总内存、总节点数
diff --git a/utils/utils.go b/utils/utils.go
index d073919..8ae1c6a 100644
--- a/utils/utils.go
+++ b/utils/utils.go
@@ -597,3 +597,44 @@ func ExtractValue(input, key string) string {
 	}
 	return ""
 }
+
+func ConvertMemory(memoryString string) (int, error) {
+	// 使用正则表达式提取数字部分和单位
+	re := regexp.MustCompile(`^(\d+\.\d+)([A-Z]+)$`)
+	matches := re.FindStringSubmatch(memoryString)
+
+	if len(matches) != 3 {
+		return 0, fmt.Errorf("illegal memory format")
+	}
+
+	// 提取数字部分
+	memorySizeStr := matches[1]
+
+	// 提取单位
+	unit := matches[2]
+
+	// 将字符串转换为float64
+	memorySize, err := strconv.ParseFloat(memorySizeStr, 64)
+	if err != nil {
+		return 0, fmt.Errorf("convert err: %v", err)
+	}
+
+	// 将TB转换为MB，1TB = 1024 * 1024MB
+	const bytesInMB = 1024 * 1024
+	var memorySizeMB float64
+
+	switch unit {
+	case "P":
+		memorySizeMB = memorySize * bytesInMB * 1024
+	case "T":
+		memorySizeMB = memorySize * bytesInMB
+	case "G":
+		memorySizeMB = memorySize * bytesInMB / 1024
+	case "M":
+		memorySizeMB = memorySize
+	default:
+		return 0, fmt.Errorf("unknown memory unit")
+	}
+
+	return int(memorySizeMB), nil
+}