fix: gpu resource device count calculation (#107)

Author: enoodle

pkg/scheduler/api/resource_info/gpu_resource_requirment.go

@@ -40,14 +40,16 @@ func NewGpuResourceRequirement() *GpuResourceRequirement {
 
 func NewGpuResourceRequirementWithGpus(gpus float64, gpuMemory int64) *GpuResourceRequirement {
 	gResource := &GpuResourceRequirement{
-		count:        fractionDefaultCount,
+		count:        0,
 		portion:      gpus,
 		gpuMemory:    gpuMemory,
 		migResources: make(map[v1.ResourceName]int64),
 	}
 	if gpus >= wholeGpuPortion {
 		gResource.count = int64(gpus)
 		gResource.portion = wholeGpuPortion
+	} else if gpus > 0 || gpuMemory > 0 { // Fraction
+		gResource.count = fractionDefaultCount
 	}
 	return gResource
 }

pkg/scheduler/api/resource_info/gpu_resource_requirment_test.go

@@ -150,6 +150,23 @@ var _ = Describe("GpuResourceRequirement mechanism", func() {
 			Expect(gpuResource1.LessEqual(gpuResource2)).To(BeTrue())
 		})
 	})
+
+	Context("GetNumOfGpuDevices", func() {
+		It("should return 0 for 0 GPUs", func() {
+			gpuResource := NewGpuResourceRequirementWithGpus(0, 0)
+			Expect(gpuResource.GetNumOfGpuDevices()).To(Equal(int64(0)))
+		})
+
+		It("should return 1 for 0.5 GPUs", func() {
+			gpuResource := NewGpuResourceRequirementWithGpus(0.5, 0)
+			Expect(gpuResource.GetNumOfGpuDevices()).To(Equal(int64(1)))
+		})
+
+		It("should return 2 for 2 GPUs", func() {
+			gpuResource := NewGpuResourceRequirementWithGpus(2, 0)
+			Expect(gpuResource.GetNumOfGpuDevices()).To(Equal(int64(2)))
+		})
+	})
 })
 
 func newGpuResourceRequirementWithValues(

pkg/scheduler/plugins/proportion/capacity_policy/max_allowed_check_test.go

@@ -6,8 +6,14 @@ package capacity_policy
 import (
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 
+	"reflect"
+
+	"github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
+	"github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
 	commonconstants "github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/common_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/podgroup_info"
@@ -205,10 +211,11 @@ var _ = Describe("Max Allowed Policy Check", func() {
 	Describe("resultsOverLimit", func() {
 		Context("resultsOverLimit tests", func() {
 			tests := map[string]struct {
-				queues         map[common_info.QueueID]*rs.QueueAttributes
-				job            *podgroup_info.PodGroupInfo
-				requestedShare rs.ResourceQuantities
-				expectedResult bool
+				queues          map[common_info.QueueID]*rs.QueueAttributes
+				job             *podgroup_info.PodGroupInfo
+				requestedShare  rs.ResourceQuantities
+				expectedResult  bool
+				expectedDetails *v2alpha2.QuotaDetails // optional, only compare if set
 			}{
 				"single queue - unlimited queue": {
 					queues: map[common_info.QueueID]*rs.QueueAttributes{
@@ -288,6 +295,39 @@ var _ = Describe("Max Allowed Policy Check", func() {
 					},
 					expectedResult: false,
 				},
+				"single queue - limited queue - results above limit 0 allocated": {
+					queues: map[common_info.QueueID]*rs.QueueAttributes{
+						"queue1": {
+							UID:               "queue1",
+							Name:              "queue1",
+							ParentQueue:       "",
+							ChildQueues:       nil,
+							CreationTimestamp: metav1.Time{},
+							QueueResourceShare: rs.QueueResourceShare{
+								GPU: rs.ResourceShare{
+									MaxAllowed: 3,
+									Allocated:  0,
+								},
+							},
+						},
+					},
+					job: &podgroup_info.PodGroupInfo{
+						Name:      "job-a",
+						Namespace: "team-a",
+						Queue:     "queue1",
+					},
+					requestedShare: rs.ResourceQuantities{
+						rs.GpuResource: 4,
+					},
+					expectedResult: false,
+					expectedDetails: &v2alpha2.QuotaDetails{
+						QueueAllocatedResources: v1.ResourceList{
+							v1.ResourceCPU:        *resource.NewMilliQuantity(0, resource.DecimalSI),
+							v1.ResourceMemory:     *resource.NewQuantity(0, resource.DecimalSI),
+							constants.GpuResource: *resource.NewQuantity(0, resource.DecimalSI),
+						},
+					},
+				},
 				"multiple queues - limited queues - results below limit": {
 					queues: map[common_info.QueueID]*rs.QueueAttributes{
 						"top-queue": {
@@ -401,6 +441,18 @@ var _ = Describe("Max Allowed Policy Check", func() {
 					capacityPolicy := New(testData.queues, true)
 					result := capacityPolicy.resultsOverLimit(testData.requestedShare, testData.job)
 					Expect(result.IsSchedulable).To(Equal(testData.expectedResult))
+					if testData.expectedDetails != nil {
+						expectedValues := reflect.ValueOf(*testData.expectedDetails)
+						resultValues := reflect.ValueOf(*result.Details.QueueDetails)
+						for i := 0; i < expectedValues.NumField(); i++ {
+							xField := expectedValues.Field(i).Interface()
+							yField := resultValues.Field(i).Interface()
+							zero := reflect.Zero(expectedValues.Field(i).Type()).Interface()
+							if !reflect.DeepEqual(xField, zero) {
+								Expect(xField).To(Equal(yField))
+							}
+						}
+					}
 				})
 			}