FIX: 修复用户空间的单节点内存不显示的问题

PKUHPC · Dec 5, 2024 · c96d88b · c96d88b
1 parent a103247
commit c96d88b
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 14 deletions.
diff --git 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,16 @@ 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
+			totalMemInt, err = utils.ConvertMemory(totalMemsTmp)
+			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
@@ -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
+}