fix(webhook) : Multi-Container & InitContainer Mutating Webhook Test

CHANGELOG.md

@@ -408,5 +408,7 @@ Add "NVIDIA_VISIBLE_DEVICES=none" to none-gpu tasks
 - Fix initialization error when using tensor parallelism on vLLM above 0.18
 - Fix multiple device typos
 - Add unit test coverage for node discovery handshake parsing with malformed or empty annotations
+- Fix webhook mutating requests for pods with initContainers and multiple application containers
+
 
 

pkg/device/nvidia/device_test.go

@@ -1467,6 +1467,9 @@ func TestFitQuota(t *testing.T) {
 	device.DevicesMap[NvidiaGPUDevice] = dev
 
 	qm := device.NewQuotaManager()
+	t.Cleanup(func() {
+		delete(qm.Quotas, "default")
+	})
 	qm.AddQuota(&corev1.ResourceQuota{
 		TypeMeta: metav1.TypeMeta{
 			APIVersion: "v1",

pkg/device/quota.go

@@ -45,7 +45,7 @@ type QuotaManager struct {
 var localCache QuotaManager
 
 func GetLocalCache() *QuotaManager {
-	return &localCache
+	return NewQuotaManager()
 }
 
 var once sync.Once
@@ -123,6 +123,9 @@ func (q *QuotaManager) AddUsage(pod *corev1.Pod, podDev PodDevices) {
 	}
 	q.mutex.Lock()
 	defer q.mutex.Unlock()
+	if q.Quotas == nil {
+		q.Quotas = make(map[string]*DeviceQuota)
+	}
 	if q.Quotas[pod.Namespace] == nil {
 		q.Quotas[pod.Namespace] = &DeviceQuota{}
 	}
@@ -232,6 +235,9 @@ func (q *QuotaManager) addQuotaLocked(quota *corev1.ResourceQuota) {
 			if !ok {
 				continue
 			}
+			if q.Quotas == nil {
+				q.Quotas = make(map[string]*DeviceQuota)
+			}
 			if q.Quotas[quota.Namespace] == nil {
 				q.Quotas[quota.Namespace] = &DeviceQuota{}
 			}

pkg/scheduler/webhook.go

@@ -70,6 +70,17 @@ func (h *webhook) Handle(_ context.Context, req admission.Request) admission.Res
 	klog.V(5).Infof(template, pod.Namespace, pod.Name, pod.UID)
 	privilegedName, hasPrivileged := privilegedContainerName(pod)
 	hasResource := false
+	for idx := range pod.Spec.InitContainers {
+		c := &pod.Spec.InitContainers[idx]
+		for _, val := range device.GetDevices() {
+			found, err := val.MutateAdmission(c, pod)
+			if err != nil {
+				klog.Errorf("validating pod failed:%s", err.Error())
+				return admission.Errored(http.StatusInternalServerError, err)
+			}
+			hasResource = hasResource || found
+		}
+	}
 	for idx := range pod.Spec.Containers {
 		c := &pod.Spec.Containers[idx]
 		for _, val := range device.GetDevices() {
@@ -139,6 +150,21 @@ func fitResourceQuota(pod *corev1.Pod) bool {
 		// container spec here. It applies its own memory factor, defaults and
 		// template rounding, which is what the scheduler later records as used,
 		// so this keeps admission and the scheduler on the same numbers.
+		var initMemoryReqMax, initCoresReqMax int64
+		for i := range pod.Spec.InitContainers {
+			req := dev.GenerateResourceRequests(&pod.Spec.InitContainers[i])
+			if req.Nums == 0 {
+				continue
+			}
+			mem := int64(req.Memreq) * int64(req.Nums)
+			cores := int64(req.Coresreq) * int64(req.Nums)
+			if mem > initMemoryReqMax {
+				initMemoryReqMax = mem
+			}
+			if cores > initCoresReqMax {
+				initCoresReqMax = cores
+			}
+		}
 		var memoryReq, coresReq int64
 		for i := range pod.Spec.Containers {
 			req := dev.GenerateResourceRequests(&pod.Spec.Containers[i])
@@ -148,6 +174,12 @@ func fitResourceQuota(pod *corev1.Pod) bool {
 			memoryReq += int64(req.Memreq) * int64(req.Nums)
 			coresReq += int64(req.Coresreq) * int64(req.Nums)
 		}
+		if initMemoryReqMax > memoryReq {
+			memoryReq = initMemoryReqMax
+		}
+		if initCoresReqMax > coresReq {
+			coresReq = initCoresReqMax
+		}
 		if memoryReq == 0 && coresReq == 0 {
 			continue
 		}