Fix: incorrect scheduling decision and calculation when using MIG (#422)

* fix mig scheduling issue

* clean up the code

---------

Co-authored-by: bmao <[email protected]>

Author: hello2mao

pkg/scheduler/api/pod_info/pod_info.go

@@ -371,10 +371,6 @@ func getTaskStatus(pod *v1.Pod, bindRequest *bindrequest_info.BindRequestInfo) p
 }
 
 func (pi *PodInfo) updatePodAdditionalFields(bindRequest *bindrequest_info.BindRequestInfo) {
-	if len(pi.Job) == 0 {
-		return
-	}
-
 	if bindRequest != nil && len(bindRequest.BindRequest.Spec.SelectedGPUGroups) > 0 {
 		pi.GPUGroups = bindRequest.BindRequest.Spec.SelectedGPUGroups
 	} else {

pkg/scheduler/cache/cluster_info/cluster_info_test.go

@@ -262,6 +262,15 @@ func TestSnapshotNodes(t *testing.T) {
 			Phase: v1core.PodRunning,
 		},
 	}
+	exampleMIGPod := examplePod.DeepCopy()
+	exampleMIGPod.Name = "mig-pod"
+	exampleMIGPod.Spec.Containers[0].Resources.Requests["nvidia.com/mig-1g.5gb"] = resource.MustParse("2")
+	exampleMIGPodWithPG := examplePod.DeepCopy()
+	exampleMIGPodWithPG.Name = "mig-pod-with-pg"
+	exampleMIGPodWithPG.Annotations = map[string]string{
+		commonconstants.PodGroupAnnotationForPod: "pg-1",
+	}
+	exampleMIGPodWithPG.Spec.Containers[0].Resources.Requests["nvidia.com/mig-1g.5gb"] = resource.MustParse("2")
 	tests := map[string]struct {
 		objs          []runtime.Object
 		resultNodes   []*node_info.NodeInfo
@@ -400,6 +409,48 @@ func TestSnapshotNodes(t *testing.T) {
 			resultPodsLen: 1,
 			nodePoolName:  "pool-a",
 		},
+		"MIG Job": {
+			objs: []runtime.Object{
+				&v1core.Node{
+					ObjectMeta: v1.ObjectMeta{
+						Name: "node-1",
+					},
+					Status: v1core.NodeStatus{
+						Allocatable: v1core.ResourceList{
+							"cpu":                   resource.MustParse("10"),
+							"nvidia.com/mig-1g.5gb": resource.MustParse("10"),
+						},
+					},
+				},
+				exampleMIGPod,
+				exampleMIGPodWithPG,
+			},
+			resultNodes: []*node_info.NodeInfo{
+				{
+					Name: "node-1",
+					Idle: resource_info.ResourceFromResourceList(
+						v1core.ResourceList{
+							"cpu":                   resource.MustParse("6"),
+							"nvidia.com/mig-1g.5gb": resource.MustParse("6"),
+						},
+					),
+					Used: resource_info.ResourceFromResourceList(
+						v1core.ResourceList{
+							"cpu":                   resource.MustParse("4"),
+							"memory":                resource.MustParse("0"),
+							"nvidia.com/mig-1g.5gb": resource.MustParse("4"),
+						},
+					),
+					Releasing: resource_info.ResourceFromResourceList(
+						v1core.ResourceList{
+							"cpu":    resource.MustParse("0"),
+							"memory": resource.MustParse("0"),
+						},
+					),
+				},
+			},
+			resultPodsLen: 2,
+		},
 	}
 
 	for name, test := range tests {

pkg/scheduler/plugins/proportion/proportion.go

@@ -26,7 +26,6 @@ import (
 	commonconstants "github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/common_info"
-	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/common_info/resources"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/node_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_status"
@@ -247,18 +246,6 @@ func getNodeResources(ssn *framework.Session, node *node_info.NodeInfo) rs.Resou
 		nodeResource.Add(rs.NewResourceQuantities(node.Allocatable.Cpu(), node.Allocatable.Memory(), 0))
 	} else {
 		nodeResource.Add(utils.QuantifyResource(node.Allocatable))
-		migEnabledGpus := 0
-		for resource, qty := range node.Node.Status.Allocatable {
-			if resource_info.IsMigResource(resource) {
-				gpu, _, err := resources.ExtractGpuAndMemoryFromMigResourceName(string(resource))
-				if err != nil {
-					log.InfraLogger.Errorf("Failed to extract gpu and memory from mig resource %v: %v", resource, err)
-					continue
-				}
-				migEnabledGpus += int(qty.Value()) * gpu
-			}
-		}
-		nodeResource[rs.GpuResource] += float64(migEnabledGpus)
 	}
 
 	// Subtract resources of non-related pods

pkg/scheduler/plugins/proportion/proportion_test.go

@@ -565,6 +565,28 @@ var _ = Describe("Set Fair Share in Proportion", func() {
 					rs.GpuResource:    2,
 				},
 			},
+			{
+				name:           "mig gpu node",
+				isRestrictNode: true,
+				node: &node_info.NodeInfo{
+					Name: "n1",
+					Node: &v1.Node{
+						ObjectMeta: metav1.ObjectMeta{
+							Labels: map[string]string{
+								"node-role.kubernetes.io/gpu-worker": "true",
+							},
+						},
+					},
+					Allocatable: resource_info.ResourceFromResourceList(
+						common_info.BuildResourceListWithMig("8000m", "10G", "nvidia.com/mig-1g.5gb"),
+					),
+				},
+				want: rs.ResourceQuantities{
+					rs.CpuResource:    8000,
+					rs.MemoryResource: 10000000000,
+					rs.GpuResource:    1,
+				},
+			},
 			{
 				name:           "ignore extra resources",
 				isRestrictNode: true,