feat(scheduler): correct resource accounting for init containers

pkg/scheduler/scheduler.go

@@ -738,12 +738,32 @@ ReleaseNodeLocks:
 func (s *Scheduler) Filter(args extenderv1.ExtenderArgs) (*extenderv1.ExtenderFilterResult, error) {
 	klog.InfoS("Starting schedule filter process", "pod", args.Pod.Name, "uuid", args.Pod.UID, "namespace", args.Pod.Namespace)
 	resourceReqs := device.Resourcereqs(args.Pod)
-	resourceReqTotal := 0
-	for _, n := range resourceReqs {
-		for _, k := range n {
-			resourceReqTotal += int(k.Nums)
+
+	initReqTotal := 0
+	appReqTotal := 0
+	numInitContainers := len(args.Pod.Spec.InitContainers)
+
+	for i := range numInitContainers {
+		currentInitReq := 0
+		for _, k := range resourceReqs[i] {
+			currentInitReq += int(k.Nums)
+		}
+		if currentInitReq > initReqTotal {
+			initReqTotal = currentInitReq
+		}
+	}
+
+	// 2. Calculate the SUM of requests among Regular Containers
+	for i := numInitContainers; i < len(resourceReqs); i++ {
+		for _, k := range resourceReqs[i] {
+			appReqTotal += int(k.Nums)
 		}
 	}
+
+	// 3. The effective total request is the MAX of (InitContainers, RegularContainers)
+	resourceReqTotal := appReqTotal
+	resourceReqTotal = max(resourceReqTotal, initReqTotal)
+
 	if resourceReqTotal == 0 {
 		klog.V(1).InfoS("Pod does not request any resources",
 			"pod", args.Pod.Name)

pkg/scheduler/scheduler_test.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"maps"
+	"slices"
 	"strings"
 	"sync/atomic"
 	"testing"
@@ -459,11 +460,11 @@ func Test_Filter(t *testing.T) {
 	}
 
 	tests := []struct {
-		name                      string
-		args                      extenderv1.ExtenderArgs
-		want                      *extenderv1.ExtenderFilterResult
-		wantPodAnnotationDeviceID string
-		wantErr                   error
+		name                       string
+		args                       extenderv1.ExtenderArgs
+		want                       *extenderv1.ExtenderFilterResult
+		wantPodAnnotationDeviceIDs []string
+		wantErr                    error
 	}{
 		{
 			name: "node use binpack gpu use binpack policy",
@@ -500,7 +501,56 @@ func Test_Filter(t *testing.T) {
 			want: &extenderv1.ExtenderFilterResult{
 				NodeNames: &[]string{"node2"},
 			},
-			wantPodAnnotationDeviceID: "device4",
+			wantPodAnnotationDeviceIDs: []string{"device4"},
+		},
+		{
+			name: "pod with init containers fits correctly using max resource logic (Binpack)",
+			args: extenderv1.ExtenderArgs{
+				Pod: &corev1.Pod{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test-init-containers",
+						UID:  "test-init-uid",
+						Annotations: map[string]string{
+							util.GPUSchedulerPolicyAnnotationKey:  util.GPUSchedulerPolicyBinpack.String(),
+							util.NodeSchedulerPolicyAnnotationKey: util.NodeSchedulerPolicyBinpack.String(),
+						},
+					},
+					Spec: corev1.PodSpec{
+						InitContainers: []corev1.Container{
+							{
+								Name:  "init-1",
+								Image: "busybox",
+								Resources: corev1.ResourceRequirements{
+									Limits: corev1.ResourceList{
+										"hami.io/gpu":      *resource.NewQuantity(1, resource.BinarySI),
+										"hami.io/gpucores": *resource.NewQuantity(20, resource.BinarySI),
+										"hami.io/gpumem":   *resource.NewQuantity(5000, resource.BinarySI),
+									},
+								},
+							},
+						},
+						Containers: []corev1.Container{
+							{
+								Name:  "app-1",
+								Image: "chrstnhntschl/gpu_burn",
+								Resources: corev1.ResourceRequirements{
+									Limits: corev1.ResourceList{
+										"hami.io/gpu":      *resource.NewQuantity(1, resource.BinarySI),
+										"hami.io/gpucores": *resource.NewQuantity(20, resource.BinarySI),
+										"hami.io/gpumem":   *resource.NewQuantity(4000, resource.BinarySI),
+									},
+								},
+							},
+						},
+					},
+				},
+				NodeNames: &[]string{"node1", "node2"},
+			},
+			wantErr: nil,
+			want: &extenderv1.ExtenderFilterResult{
+				NodeNames: &[]string{"node2"},
+			},
+			wantPodAnnotationDeviceIDs: []string{"device3"},
 		},
 		{
 			name: "node use binpack gpu use spread policy",
@@ -537,7 +587,7 @@ func Test_Filter(t *testing.T) {
 			want: &extenderv1.ExtenderFilterResult{
 				NodeNames: &[]string{"node2"},
 			},
-			wantPodAnnotationDeviceID: "device3",
+			wantPodAnnotationDeviceIDs: []string{"device3", "device4"}, // Both are acceptable due to tie
 		},
 		{
 			name: "node use spread gpu use binpack policy",
@@ -574,7 +624,7 @@ func Test_Filter(t *testing.T) {
 			want: &extenderv1.ExtenderFilterResult{
 				NodeNames: &[]string{"node1"},
 			},
-			wantPodAnnotationDeviceID: "device1",
+			wantPodAnnotationDeviceIDs: []string{"device1"},
 		},
 		{
 			name: "node use spread gpu use spread policy",
@@ -611,7 +661,7 @@ func Test_Filter(t *testing.T) {
 			want: &extenderv1.ExtenderFilterResult{
 				NodeNames: &[]string{"node1"},
 			},
-			wantPodAnnotationDeviceID: "device2",
+			wantPodAnnotationDeviceIDs: []string{"device2"},
 		},
 	}
 
@@ -624,7 +674,12 @@ func Test_Filter(t *testing.T) {
 			assert.DeepEqual(t, test.want, got)
 			getPod, _ := client.KubeClient.CoreV1().Pods(test.args.Pod.Namespace).Get(context.Background(), test.args.Pod.Name, metav1.GetOptions{})
 			podDevices, _ := device.DecodePodDevices(device.SupportDevices, getPod.Annotations)
-			assert.DeepEqual(t, test.wantPodAnnotationDeviceID, podDevices["NVIDIA"][0][0].UUID)
+
+			actualUUID := podDevices["NVIDIA"][0][0].UUID
+
+			if !slices.Contains(test.wantPodAnnotationDeviceIDs, actualUUID) {
+				t.Errorf("expected one of %v, got %s", test.wantPodAnnotationDeviceIDs, actualUUID)
+			}
 		})
 	}
 }