From 9754ae34174362bc21f673c4a36e93863cf3f4f6 Mon Sep 17 00:00:00 2001
From: JesseStutler <chenzicong4@huawei.com>
Date: Tue, 24 Dec 2024 10:51:07 +0800
Subject: [PATCH] fix: if the capability cpu or memory is not specified in the
 hierarchical queue, it will be set to the corresponding value of the parent
 queue

Signed-off-by: JesseStutler <chenzicong4@huawei.com>
---
 pkg/scheduler/plugins/capacity/capacity.go    | 38 +++++++++++++------
 .../plugins/capacity/capacity_test.go         |  2 +-
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/pkg/scheduler/plugins/capacity/capacity.go b/pkg/scheduler/plugins/capacity/capacity.go
index 9c661588c0..63f7a6d91e 100644
--- a/pkg/scheduler/plugins/capacity/capacity.go
+++ b/pkg/scheduler/plugins/capacity/capacity.go
@@ -419,9 +419,7 @@ func (cp *capacityPlugin) buildHierarchicalQueueAttrs(ssn *framework.Session) bo
 			continue
 		}
 
-		attr := cp.newQueueAttr(queue)
-		cp.queueOpts[queue.UID] = attr
-		err := cp.updateParents(queue, ssn)
+		err := cp.initializeHierarchicalQueueAttrs(queue, ssn)
 		if err != nil {
 			klog.Errorf("Failed to update Queue <%s> attributes, error: %v", queue.Name, err)
 			return false
@@ -579,7 +577,7 @@ func (cp *capacityPlugin) buildHierarchicalQueueAttrs(ssn *framework.Session) bo
 	return true
 }
 
-func (cp *capacityPlugin) newQueueAttr(queue *api.QueueInfo) *queueAttr {
+func (cp *capacityPlugin) newQueueAttr(queue *api.QueueInfo, parent *api.QueueInfo) *queueAttr {
 	attr := &queueAttr{
 		queueID:  queue.UID,
 		name:     queue.Name,
@@ -593,13 +591,29 @@ func (cp *capacityPlugin) newQueueAttr(queue *api.QueueInfo) *queueAttr {
 		inqueue:   api.EmptyResource(),
 		guarantee: api.EmptyResource(),
 	}
+
+	if queue.Name == cp.rootQueue {
+		attr.realCapability = cp.totalResource
+		attr.capability = attr.realCapability
+		attr.deserved = cp.totalResource
+		return attr
+	}
+
 	if len(queue.Queue.Spec.Capability) != 0 {
 		attr.capability = api.NewResource(queue.Queue.Spec.Capability)
+		parentCapability := api.EmptyResource()
+		if parent.Queue.Name == cp.rootQueue {
+			// if current queue's parent is root, set parentCapability to totalResource (because root's realCapability is totalResource, and the attr of root queue may not have been initialized yet)
+			parentCapability = cp.totalResource
+		} else {
+			parentCapability = api.NewResource(parent.Queue.Spec.Capability)
+		}
+		// If the user does not set CPU or memory in capability, we set the value to be the same as parent (we do not consider the situation where the user sets CPU or memory<=0)
 		if attr.capability.MilliCPU <= 0 {
-			attr.capability.MilliCPU = math.MaxFloat64
+			attr.capability.MilliCPU = parentCapability.MilliCPU
 		}
 		if attr.capability.Memory <= 0 {
-			attr.capability.Memory = math.MaxFloat64
+			attr.capability.Memory = parentCapability.Memory
 		}
 	}
 
@@ -610,8 +624,9 @@ func (cp *capacityPlugin) newQueueAttr(queue *api.QueueInfo) *queueAttr {
 	return attr
 }
 
-func (cp *capacityPlugin) updateParents(queue *api.QueueInfo, ssn *framework.Session) error {
+func (cp *capacityPlugin) initializeHierarchicalQueueAttrs(queue *api.QueueInfo, ssn *framework.Session) error {
 	if queue.Name == cp.rootQueue {
+		cp.queueOpts[queue.UID] = cp.newQueueAttr(queue, nil)
 		return nil
 	}
 
@@ -624,10 +639,11 @@ func (cp *capacityPlugin) updateParents(queue *api.QueueInfo, ssn *framework.Ses
 	}
 
 	parentInfo := ssn.Queues[api.QueueID(parent)]
+	curQueueAttr := cp.newQueueAttr(queue, parentInfo)
+	cp.queueOpts[curQueueAttr.queueID] = curQueueAttr
+
 	if _, found := cp.queueOpts[parentInfo.UID]; !found {
-		parentAttr := cp.newQueueAttr(parentInfo)
-		cp.queueOpts[parentAttr.queueID] = parentAttr
-		err := cp.updateParents(parentInfo, ssn)
+		err := cp.initializeHierarchicalQueueAttrs(parentInfo, ssn)
 		if err != nil {
 			return err
 		}
@@ -653,8 +669,6 @@ func (cp *capacityPlugin) checkHierarchicalQueue(attr *queueAttr) error {
 	if attr.name == cp.rootQueue {
 		attr.guarantee = totalGuarantee
 		cp.totalGuarantee = totalGuarantee
-		attr.realCapability = cp.totalResource
-		attr.deserved = cp.totalResource
 	}
 
 	for _, childAttr := range attr.children {
diff --git a/pkg/scheduler/plugins/capacity/capacity_test.go b/pkg/scheduler/plugins/capacity/capacity_test.go
index 774e37809d..bd7211fc51 100644
--- a/pkg/scheduler/plugins/capacity/capacity_test.go
+++ b/pkg/scheduler/plugins/capacity/capacity_test.go
@@ -348,7 +348,7 @@ func Test_capacityPlugin_OnSessionOpenWithHierarchy(t *testing.T) {
 	// podgroup
 	pg1 := util.BuildPodGroup("pg1", "ns1", "q11", 1, nil, schedulingv1beta1.PodGroupInqueue)
 	// queue
-	root := buildQueueWithParents("root", "", api.BuildResourceList("8", "8Gi", []api.ScalarResource{{Name: "pods", Value: "10"}}...), api.BuildResourceList("8", "8Gi", []api.ScalarResource{{Name: "pods", Value: "10"}}...))
+	root := buildQueueWithParents("root", "", nil, nil)
 	queue1 := buildQueueWithParents("q1", "root", nil, api.BuildResourceList("4", "4Gi"))
 	queue2 := buildQueueWithParents("q2", "root", nil, api.BuildResourceList("4", "4Gi"))
 	queue11 := buildQueueWithParents("q11", "q1", nil, api.BuildResourceList("1", "1Gi"))