Merge branch 'master' of github.com:Project-HAMi/HAMi

Project-HAMi · May 11, 2024 · 17bf9e5 · 17bf9e5
2 parents 3648648 + fd8deca
commit 17bf9e5
Show file tree

Hide file tree

Showing 4 changed files with 95 additions and 28 deletions.
diff --git a/docs/gpu-dashboard.json b/docs/gpu-dashboard.json
@@ -13,7 +13,7 @@
       }
     ]
   },
-  "description": "This dashboard is gpu metrics dashboard base on NVIDIA DCGM Exporter and 4paradigm/k8s-vgpu-scheduler",
+  "description": "This dashboard is gpu metrics dashboard base on NVIDIA DCGM Exporter and HAMi/k8s-vgpu-scheduler",
   "editable": true,
   "gnetId": 12239,
   "graphTooltip": 0,
@@ -721,7 +721,7 @@
       "timeFrom": null,
       "timeRegions": [],
       "timeShift": null,
-      "title": "4paradigm-pod显存使用量（byte）",
+      "title": "HAMi-pod显存使用量（byte）",
       "tooltip": {
         "shared": true,
         "sort": 0,
@@ -819,7 +819,7 @@
       "timeFrom": null,
       "timeRegions": [],
       "timeShift": null,
-      "title": "4paradigm-节点GPU显存使用量",
+      "title": "HAMi-节点GPU显存使用量",
       "tooltip": {
         "shared": true,
         "sort": 0,
@@ -917,7 +917,7 @@
       "timeFrom": null,
       "timeRegions": [],
       "timeShift": null,
-      "title": "4paradigm-节点GPU算力使用率",
+      "title": "HAMi-节点GPU算力使用率",
       "tooltip": {
         "shared": true,
         "sort": 0,

diff --git a/pkg/device/nvidia/device.go b/pkg/device/nvidia/device.go
@@ -166,37 +166,35 @@ func (dev *NvidiaGPUDevices) MutateAdmission(ctr *corev1.Container) (bool, error
 }
 
 func checkGPUtype(annos map[string]string, cardtype string) bool {
+	cardtype = strings.ToUpper(cardtype)
 	if inuse, ok := annos[GPUInUse]; ok {
-		if !strings.Contains(inuse, ",") {
-			if strings.Contains(strings.ToUpper(cardtype), strings.ToUpper(inuse)) {
-				return true
-			}
-		} else {
-			for _, val := range strings.Split(inuse, ",") {
-				if strings.Contains(strings.ToUpper(cardtype), strings.ToUpper(val)) {
-					return true
-				}
-			}
+		useTypes := strings.Split(inuse, ",")
+		if !ContainsSliceFunc(useTypes, func(useType string) bool {
+			return strings.Contains(cardtype, strings.ToUpper(useType))
+		}) {
+			return false
 		}
-		return false
 	}
-	if nouse, ok := annos[GPUNoUse]; ok {
-		if !strings.Contains(nouse, ",") {
-			if strings.Contains(strings.ToUpper(cardtype), strings.ToUpper(nouse)) {
-				return false
-			}
-		} else {
-			for _, val := range strings.Split(nouse, ",") {
-				if strings.Contains(strings.ToUpper(cardtype), strings.ToUpper(val)) {
-					return false
-				}
-			}
+	if unuse, ok := annos[GPUNoUse]; ok {
+		unuseTypes := strings.Split(unuse, ",")
+		if ContainsSliceFunc(unuseTypes, func(unuseType string) bool {
+			return strings.Contains(cardtype, strings.ToUpper(unuseType))
+		}) {
+			return false
 		}
-		return true
 	}
 	return true
 }
 
+func ContainsSliceFunc[S ~[]E, E any](s S, match func(E) bool) bool {
+	for _, e := range s {
+		if match(e) {
+			return true
+		}
+	}
+	return false
+}
+
 func assertNuma(annos map[string]string) bool {
 	numabind, ok := annos[NumaBind]
 	if ok {

diff --git a/pkg/device/nvidia/device_test.go b/pkg/device/nvidia/device_test.go
@@ -206,3 +206,72 @@ func Test_CheckUUID(t *testing.T) {
 		})
 	}
 }
+
+func Test_CheckType(t *testing.T) {
+	gpuDevices := &NvidiaGPUDevices{}
+	tests := []struct {
+		name string
+		args struct {
+			annos map[string]string
+			d     util.DeviceUsage
+		}
+		want bool
+	}{
+		{
+			name: "use set GPUInUse don't set GPUNoUse annotation,device match",
+			args: struct {
+				annos map[string]string
+				d     util.DeviceUsage
+			}{
+				annos: map[string]string{
+					GPUInUse: "A10",
+				},
+				d: util.DeviceUsage{
+					Type: "NVIDIA A100",
+				},
+			},
+			want: true,
+		},
+		{
+			name: "use set GPUInUse set GPUNoUse annotation,device don't match",
+			args: struct {
+				annos map[string]string
+				d     util.DeviceUsage
+			}{
+				annos: map[string]string{
+					GPUInUse: "A10",
+					GPUNoUse: "A100",
+				},
+				d: util.DeviceUsage{
+					Type: "NVIDIA A100",
+				},
+			},
+			want: false,
+		},
+		{
+			name: "use set GPUInUse set GPUNoUse annotation,device match",
+			args: struct {
+				annos map[string]string
+				d     util.DeviceUsage
+			}{
+				annos: map[string]string{
+					GPUInUse: "A10",
+					GPUNoUse: "A100",
+				},
+				d: util.DeviceUsage{
+					Type: "NVIDIA A10",
+				},
+			},
+			want: true,
+		},
+	}
+	req := util.ContainerDeviceRequest{
+		Type: NvidiaGPUDevice,
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			_, got, _ := gpuDevices.CheckType(test.args.annos, test.args.d, req)
+			assert.Equal(t, test.want, got)
+		})
+	}
+}
diff --git a/version.mk b/version.mk
@@ -9,5 +9,5 @@ NVIDIA_IMAGE=nvidia/cuda:12.3.2-devel-ubuntu22.04
 DEST_DIR=/usr/local/vgpu/
 
 VERSION = v0.0.1
-IMG_NAME ="hami"
+IMG_NAME =hami
 IMG_TAG="${IMG_NAME}:${VERSION}"