feat: check resource quota in webhook

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

Author: DSFans2014

pkg/scheduler/webhook.go

@@ -29,6 +29,7 @@ import (
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 
 	"github.com/Project-HAMi/HAMi/pkg/device"
+	"github.com/Project-HAMi/HAMi/pkg/device/nvidia"
 	"github.com/Project-HAMi/HAMi/pkg/scheduler/config"
 )
 
@@ -96,10 +97,62 @@ func (h *webhook) Handle(_ context.Context, req admission.Request) admission.Res
 			return admission.Denied("pod has node assigned")
 		}
 	}
+	if !fitResourceQuota(pod) {
+		return admission.Denied("exceeding resource quota")
+	}
 	marshaledPod, err := json.Marshal(pod)
 	if err != nil {
 		klog.Errorf(template+" - Failed to marshal pod, error: %v", pod.Namespace, pod.Name, pod.UID, err)
 		return admission.Errored(http.StatusInternalServerError, err)
 	}
 	return admission.PatchResponseFromRaw(req.Object.Raw, marshaledPod)
 }
+
+func fitResourceQuota(pod *corev1.Pod) bool {
+	for deviceName, dev := range device.GetDevices() {
+		// Only supports NVIDIA
+		if deviceName != nvidia.NvidiaGPUDevice {
+			continue
+		}
+		memoryFactor := nvidia.MemoryFactor
+		resourceNames := dev.GetResourceNames()
+		resourceName := corev1.ResourceName(corev1.ResourceName(resourceNames.ResourceCountName))
+		memResourceName := corev1.ResourceName(corev1.ResourceName(resourceNames.ResourceMemoryName))
+		coreResourceName := corev1.ResourceName(corev1.ResourceName(resourceNames.ResourceCoreName))
+		var memoryReq int64 = 0
+		var coresReq int64 = 0
+		getRequest := func(ctr *corev1.Container, resName corev1.ResourceName) (int64, bool) {
+			v, ok := ctr.Resources.Limits[resName]
+			if !ok {
+				v, ok = ctr.Resources.Requests[resName]
+			}
+			if ok {
+				if n, ok := v.AsInt64(); ok {
+					return n, true
+				}
+			}
+			return 0, false
+		}
+		for _, ctr := range pod.Spec.Containers {
+			req, ok := getRequest(&ctr, resourceName)
+			if ok && req == 1 {
+				if memReq, ok := getRequest(&ctr, memResourceName); ok {
+					memoryReq += memReq
+				}
+				if coreReq, ok := getRequest(&ctr, coreResourceName); ok {
+					coresReq += coreReq
+				}
+			}
+		}
+		if memoryFactor > 1 {
+			oriMemReq := memoryReq
+			memoryReq = memoryReq * int64(memoryFactor)
+			klog.V(5).Infof("Adjusting memory request for quota check: oriMemReq %d, memoryReq %d, factor %d", oriMemReq, memoryReq, memoryFactor)
+		}
+		if !device.GetLocalCache().FitQuota(pod.Namespace, memoryReq, memoryFactor, coresReq, deviceName) {
+			klog.Infof(template+" - Denying admission", pod.Namespace, pod.Name, pod.UID)
+			return false
+		}
+	}
+	return true
+}

pkg/scheduler/webhook_test.go

@@ -29,6 +29,7 @@ import (
 	"k8s.io/klog/v2"
 	"sigs.k8s.io/controller-runtime/pkg/webhook/admission"
 
+	"github.com/Project-HAMi/HAMi/pkg/device"
 	"github.com/Project-HAMi/HAMi/pkg/device/nvidia"
 	"github.com/Project-HAMi/HAMi/pkg/scheduler/config"
 )
@@ -240,3 +241,158 @@ func TestPodHasDifferentScheduler(t *testing.T) {
 		t.Errorf("Expected allowed response for pod with different scheduler, but got: %v", resp)
 	}
 }
+
+func TestFitResourceQuota(t *testing.T) {
+	config.SchedulerName = "hami-scheduler"
+
+	sConfig := &config.Config{
+		NvidiaConfig: nvidia.NvidiaConfig{
+			ResourceCountName:            "nvidia.com/gpu",
+			ResourceMemoryName:           "nvidia.com/gpumem",
+			ResourceMemoryPercentageName: "nvidia.com/gpumem-percentage",
+			ResourceCoreName:             "nvidia.com/gpucores",
+			DefaultMemory:                0,
+			DefaultCores:                 0,
+			DefaultGPUNum:                1,
+			MemoryFactor:                 1,
+		},
+	}
+
+	if err := config.InitDevicesWithConfig(sConfig); err != nil {
+		klog.Fatalf("Failed to initialize devices with config: %v", err)
+	}
+
+	qm := device.NewQuotaManager()
+	ns := "default"
+	memName := "nvidia.com/gpumem"
+	coreName := "nvidia.com/gpucores"
+
+	qm.Quotas[ns] = &device.DeviceQuota{
+		memName:  &device.Quota{Used: 1000, Limit: 2000},
+		coreName: &device.Quota{Used: 200, Limit: 400},
+	}
+
+	testCases := []struct {
+		name string
+		pod  *corev1.Pod
+		fit  bool
+	}{
+		{
+			name: "quota passed",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "default",
+				},
+				Spec: corev1.PodSpec{
+					SchedulerName: "hami-scheduler",
+					Containers: []corev1.Container{
+						{
+							Name: "container1",
+							SecurityContext: &corev1.SecurityContext{
+								Privileged: nil,
+							},
+							Resources: corev1.ResourceRequirements{
+								Limits: corev1.ResourceList{
+									"nvidia.com/gpu":    resource.MustParse("1"),
+									"nvidia.com/gpumem": resource.MustParse("100"),
+								},
+							},
+						},
+					},
+				},
+			},
+			fit: true,
+		},
+		{
+			name: "quota exceeded",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "default",
+				},
+				Spec: corev1.PodSpec{
+					SchedulerName: "hami-scheduler",
+					Containers: []corev1.Container{
+						{
+							Name: "container1",
+							SecurityContext: &corev1.SecurityContext{
+								Privileged: nil,
+							},
+							Resources: corev1.ResourceRequirements{
+								Limits: corev1.ResourceList{
+									"nvidia.com/gpu":    resource.MustParse("1"),
+									"nvidia.com/gpumem": resource.MustParse("1024"),
+								},
+							},
+						},
+					},
+				},
+			},
+			fit: false,
+		},
+		{
+			name: "request multiple gpus",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "default",
+				},
+				Spec: corev1.PodSpec{
+					SchedulerName: "hami-scheduler",
+					Containers: []corev1.Container{
+						{
+							Name: "container1",
+							SecurityContext: &corev1.SecurityContext{
+								Privileged: nil,
+							},
+							Resources: corev1.ResourceRequirements{
+								Limits: corev1.ResourceList{
+									"nvidia.com/gpu":    resource.MustParse("2"),
+									"nvidia.com/gpumem": resource.MustParse("1024"),
+								},
+							},
+						},
+					},
+				},
+			},
+			fit: true,
+		},
+		{
+			name: "request ascend",
+			pod: &corev1.Pod{
+				ObjectMeta: metav1.ObjectMeta{
+					Name:      "test-pod",
+					Namespace: "default",
+				},
+				Spec: corev1.PodSpec{
+					SchedulerName: "hami-scheduler",
+					Containers: []corev1.Container{
+						{
+							Name: "container1",
+							SecurityContext: &corev1.SecurityContext{
+								Privileged: nil,
+							},
+							Resources: corev1.ResourceRequirements{
+								Limits: corev1.ResourceList{
+									"huawei.com/Ascend910B":        resource.MustParse("1"),
+									"huawei.com/Ascend910B-memory": resource.MustParse("1024"),
+								},
+							},
+						},
+					},
+				},
+			},
+			fit: true,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			result := fitResourceQuota(tc.pod)
+			if tc.fit != result {
+				t.Errorf("Expected %v, but got %v", tc.fit, result)
+			}
+		})
+	}
+}