Change var names and add modify map struct.

Signed-off-by: ZhengW22 <[email protected]>

Author: ZhengW22

pkg/scheduler/nodes.go

@@ -36,7 +36,7 @@ import (
 	"github.com/Project-HAMi/HAMi/pkg/device/nvidia"
 )
 
-var vendorUUIDMap = map[string][]string{
+var vendorNoUseAnnoKeyMap = map[string][]string{
 	nvidia.GPUNoUseUUID:        {nvidia.NvidiaGPUDevice},
 	cambricon.MLUNoUseUUID:     {cambricon.CambriconMLUDevice},
 	hygon.DCUNoUseUUID:         {hygon.HygonDCUDevice},
@@ -94,35 +94,40 @@ func (m *nodeManager) addNode(nodeID string, nodeInfo *device.NodeInfo) {
 	m.nodes[nodeID].Devices = rmDeviceByNodeAnnotation(m.nodes[nodeID])
 }
 
-func rmDeviceByNodeAnnotation(nodeInfo *device.NodeInfo) []device.DeviceInfo {
-	disableGPUUUIDVendorMap := make(map[string][]string)
+func rmDeviceByNodeAnnotation(nodeInfo *util.NodeInfo) []util.DeviceInfo {
+	vendorWithDisableGPUUUIDMap := make(map[string]map[string]bool)
 	if nodeInfo.Node != nil && nodeInfo.Node.Annotations != nil {
-		for annoKey, vendor := range vendorUUIDMap {
-			klog.V(5).Infof("Current annokey is %s, and vendor is %v", annoKey, vendor)
+		for annoKey, vendors := range vendorNoUseAnnoKeyMap {
+			klog.V(5).Infof("Current annokey is %s, and vendor is %v", annoKey, vendors)
 			if value, ok := nodeInfo.Node.Annotations[annoKey]; ok {
 				disableGPUUUIDList := strings.Split(value, ",")
 				klog.V(5).Infof("Disable gpu uuid list is: %v", disableGPUUUIDList)
 				for _, disableGPUUUID := range disableGPUUUIDList {
 					if id := strings.TrimSpace(disableGPUUUID); id != "" {
-						disableGPUUUIDVendorMap[id] = vendor
+						for _, vendor := range vendors {
+							if vendorWithDisableGPUUUIDMap[vendor] == nil {
+								newVendorMap := make(map[string]bool)
+								newVendorMap[disableGPUUUID] = true
+								vendorWithDisableGPUUUIDMap[vendor] = newVendorMap
+							} else {
+								vendorWithDisableGPUUUIDMap[vendor][disableGPUUUID] = true
+							}
+						}
 					}
 				}
 			}
 		}
 	}
-	if len(disableGPUUUIDVendorMap) == 0 {
+	if len(vendorWithDisableGPUUUIDMap) == 0 {
 		return nodeInfo.Devices
 	}
 	tmp := make([]device.DeviceInfo, 0, len(nodeInfo.Devices))
 	for _, d := range nodeInfo.Devices {
 		removeFlag := false
-		if vendorList, ok := disableGPUUUIDVendorMap[d.ID]; ok {
-			for _, vendor := range vendorList {
-				if vendor == d.DeviceVendor {
-					klog.V(5).Infof("Disable gpu uuid is : %s", d.ID)
-					removeFlag = true
-					break
-				}
+		if disableGPUUUIDMap, ok := vendorWithDisableGPUUUIDMap[d.DeviceVendor]; ok {
+			if ok := disableGPUUUIDMap[d.ID]; ok {
+				klog.V(5).Infof("Disable gpu uuid is : %s", d.ID)
+				removeFlag = true
 			}
 		}
 		if !removeFlag {

pkg/scheduler/nodes_test.go

@@ -19,6 +19,7 @@ package scheduler
 import (
 	"fmt"
 	"reflect"
+	"strings"
 	"testing"
 
 	"github.com/Project-HAMi/HAMi/pkg/device"
@@ -30,6 +31,7 @@ import (
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
 	"github.com/Project-HAMi/HAMi/pkg/device"
+	"github.com/Project-HAMi/HAMi/pkg/device/metax"
 	"github.com/Project-HAMi/HAMi/pkg/device/nvidia"
 	"github.com/Project-HAMi/HAMi/pkg/util"
 )
@@ -332,7 +334,7 @@ func Test_rmNodeDevices(t *testing.T) {
 
 func Test_rmDeviceByNodeAnnotation(t *testing.T) {
 	id1 := "60151478-4709-4242-a8c1-a944252d194b"
-	id2 := "60151478-4709-4242-a8c1-a944252d194d"
+	id2 := "33c00a52-72ab-4b61-a7ce-43107588835b"
 	type args struct {
 		nodeInfo *device.NodeInfo
 	}
@@ -342,7 +344,7 @@ func Test_rmDeviceByNodeAnnotation(t *testing.T) {
 		want []device.DeviceInfo
 	}{
 		{
-			name: "Test remove device",
+			name: "Test remove one device",
 			args: args{
 				nodeInfo: &util.NodeInfo{
 					Node:    &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nvidia.GPUNoUseUUID: id1}}},
@@ -351,6 +353,26 @@ func Test_rmDeviceByNodeAnnotation(t *testing.T) {
 			},
 			want: []device.DeviceInfo{},
 		},
+		{
+			name: "Test remove two devices",
+			args: args{
+				nodeInfo: &util.NodeInfo{
+					Node:    &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nvidia.GPUNoUseUUID: strings.Join([]string{id1, id2}, ",")}}},
+					Devices: []util.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}, {DeviceVendor: nvidia.NvidiaGPUDevice, ID: id2}},
+				},
+			},
+			want: []util.DeviceInfo{},
+		},
+		{
+			name: "Test remove one device and keep one device",
+			args: args{
+				nodeInfo: &util.NodeInfo{
+					Node:    &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{nvidia.GPUNoUseUUID: strings.Join([]string{id2}, ",")}}},
+					Devices: []util.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}, {DeviceVendor: nvidia.NvidiaGPUDevice, ID: id2}},
+				},
+			},
+			want: []util.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}},
+		},
 		{
 			name: "Test no removing device, case1",
 			args: args{
@@ -371,6 +393,26 @@ func Test_rmDeviceByNodeAnnotation(t *testing.T) {
 			},
 			want: []util.DeviceInfo{{DeviceVendor: nvidia.NvidiaGPUDevice, ID: id1}},
 		},
+		{
+			name: "Test removing metax device, case1",
+			args: args{
+				nodeInfo: &util.NodeInfo{
+					Node:    &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{metax.MetaxNoUseUUID: id1}}},
+					Devices: []util.DeviceInfo{{DeviceVendor: metax.MetaxGPUDevice, ID: id1}},
+				},
+			},
+			want: []util.DeviceInfo{},
+		},
+		{
+			name: "Test removing metax device, case2",
+			args: args{
+				nodeInfo: &util.NodeInfo{
+					Node:    &corev1.Node{ObjectMeta: metav1.ObjectMeta{Annotations: map[string]string{metax.MetaxNoUseUUID: id1}}},
+					Devices: []util.DeviceInfo{{DeviceVendor: metax.MetaxSGPUDevice, ID: id1}},
+				},
+			},
+			want: []util.DeviceInfo{},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {