feat(util): use ContainerDevice Idx

Project-HAMi · Dec 26, 2024 · ba2ad3c · ba2ad3c
1 parent 311d18b
commit ba2ad3c
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 18 deletions.
diff --git a/pkg/util/util.go b/pkg/util/util.go
@@ -207,8 +207,9 @@ func UnMarshalNodeDevices(str string) ([]*DeviceInfo, error) {
 func EncodeContainerDevices(cd ContainerDevices) string {
 	tmp := ""
 	for _, val := range cd {
-		tmp += val.UUID + "," + val.Type + "," + strconv.Itoa(int(val.Usedmem)) + "," + strconv.Itoa(int(val.Usedcores)) + OneContainerMultiDeviceSplitSymbol
-	}
+        tmp += val.UUID + "," + val.Type + "," + strconv.Itoa(int(val.Usedmem)) + "," +
+            strconv.Itoa(int(val.Usedcores)) + "," + strconv.Itoa(val.Idx) + OneContainerMultiDeviceSplitSymbol
+    }
 	klog.Infof("Encoded container Devices: %s", tmp)
 	return tmp
 	//return strings.Join(cd, ",")
@@ -218,8 +219,9 @@ func EncodeContainerDeviceType(cd ContainerDevices, t string) string {
 	tmp := ""
 	for _, val := range cd {
 		if strings.Compare(val.Type, t) == 0 {
-			tmp += val.UUID + "," + val.Type + "," + strconv.Itoa(int(val.Usedmem)) + "," + strconv.Itoa(int(val.Usedcores))
-		}
+            tmp += val.UUID + "," + val.Type + "," + strconv.Itoa(int(val.Usedmem)) + "," +
+                strconv.Itoa(int(val.Usedcores)) + "," + strconv.Itoa(val.Idx)
+        }
 		tmp += OneContainerMultiDeviceSplitSymbol
 	}
 	klog.Infof("Encoded container Certain Device type: %s->%s", t, tmp)
@@ -258,20 +260,27 @@ func DecodeContainerDevices(str string) (ContainerDevices, error) {
 		return ContainerDevices{}, nil
 	}
 	for _, val := range cd {
-		if strings.Contains(val, ",") {
-			//fmt.Println("cd is ", val)
-			tmpstr := strings.Split(val, ",")
-			if len(tmpstr) < 4 {
-				return ContainerDevices{}, fmt.Errorf("pod annotation format error; information missing, please do not use nodeName field in task")
-			}
-			tmpdev.UUID = tmpstr[0]
-			tmpdev.Type = tmpstr[1]
-			devmem, _ := strconv.ParseInt(tmpstr[2], 10, 32)
-			tmpdev.Usedmem = int32(devmem)
-			devcores, _ := strconv.ParseInt(tmpstr[3], 10, 32)
-			tmpdev.Usedcores = int32(devcores)
+        if !strings.Contains(val, ",") {
+            continue
+        }
+        tmpstr := strings.Split(val, ",")
+        if len(tmpstr) < 4 {
+            return ContainerDevices{}, fmt.Errorf("pod annotation format error; information missing, please do not use nodeName field in task")
+        }
+        tmpdev.UUID = tmpstr[0]
+        tmpdev.Type = tmpstr[1]
+        devmem, _ := strconv.ParseInt(tmpstr[2], 10, 32)
+        tmpdev.Usedmem = int32(devmem)
+        devcores, _ := strconv.ParseInt(tmpstr[3], 10, 32)
+        tmpdev.Usedcores = int32(devcores)
+        if len(tmpstr) == 5 {
+            idx, err := strconv.Atoi(tmpstr[4])
+            if err != nil {
+                return ContainerDevices{}, fmt.Errorf("invalid index value in pod annotation: %v", err)
+            }
+            tmpdev.Idx = idx
+        }
 			contdev = append(contdev, tmpdev)
-		}
 	}
 	klog.V(5).Infof("Finished decoding container devices. Total devices: %d", len(contdev))
 	return contdev, nil

diff --git a/pkg/util/util_test.go b/pkg/util/util_test.go
@@ -43,7 +43,10 @@ func TestExtractMigTemplatesFromUUID(t *testing.T) {
 }
 
 func TestEmptyContainerDevicesCoding(t *testing.T) {
-	cd1 := ContainerDevices{}
+    cd1 := ContainerDevices{
+        ContainerDevice{0, "UUID1", "Type1", 1000, 30},
+        ContainerDevice{1, "UUID2", "Type1", 1000, 30},
+    }
 	s := EncodeContainerDevices(cd1)
 	fmt.Println(s)
 	cd2, _ := DecodeContainerDevices(s)