Use peek to consider multiple jobs

Author: itsomri

pkg/scheduler/actions/common/solvers/pod_scenario_builder.go

@@ -113,5 +113,7 @@ func (asb *PodAccumulatedScenarioBuilder) GetNextScenario() *solverscenario.ByNo
 		}
 	}
 
+	asb.victimsJobsQueue.ResetQueuePopsMap()
+
 	return asb.lastScenario
 }

pkg/scheduler/actions/utils/job_order_by_queue.go

@@ -29,6 +29,7 @@ type JobsOrderByQueues struct {
 	departmentIdToDepartmentMetadata map[common_info.QueueID]*departmentMetadata
 	ssn                              *framework.Session
 	jobsOrderInitOptions             JobsOrderInitOptions
+	queuePopsMap                     map[common_info.QueueID]int
 }
 
 func NewJobsOrderByQueues(ssn *framework.Session, options JobsOrderInitOptions) JobsOrderByQueues {
@@ -37,9 +38,14 @@ func NewJobsOrderByQueues(ssn *framework.Session, options JobsOrderInitOptions)
 		queueIdToQueueMetadata:           map[common_info.QueueID]*jobsQueueMetadata{},
 		departmentIdToDepartmentMetadata: map[common_info.QueueID]*departmentMetadata{},
 		jobsOrderInitOptions:             options,
+		queuePopsMap:                     map[common_info.QueueID]int{},
 	}
 }
 
+func (jobsOrder *JobsOrderByQueues) ResetQueuePopsMap() {
+	jobsOrder.queuePopsMap = map[common_info.QueueID]int{}
+}
+
 func (jobsOrder *JobsOrderByQueues) IsEmpty() bool {
 	return jobsOrder.activeDepartments == nil || jobsOrder.activeDepartments.Empty()
 }
@@ -71,6 +77,12 @@ func (jobsOrder *JobsOrderByQueues) PopNextJob() *podgroup_info.PodGroupInfo {
 	job := jobsOrder.queueIdToQueueMetadata[queue.UID].jobsInQueue.Pop().(*podgroup_info.PodGroupInfo)
 	jobsOrder.updateQueuePriorityQueue(queue, department)
 	jobsOrder.updateDepartmentPriorityQueue(department)
+
+	if _, found := jobsOrder.queuePopsMap[queue.UID]; !found {
+		jobsOrder.queuePopsMap[queue.UID] = 0
+	}
+	jobsOrder.queuePopsMap[queue.UID]++
+
 	log.InfraLogger.V(7).Infof("Popped job: %v", job.Name)
 	return job
 }
@@ -237,8 +249,18 @@ func (jobsOrder *JobsOrderByQueues) buildFuncOrderBetweenQueuesWithJobs(jobsQueu
 			return reverseOrder // When l has higher priority, return false
 		}
 
-		lJobInfosRaw := jobsQueueMetadataPerQueue[lQueue.UID].jobsInQueue.Peek(1)
-		rJobInfosRaw := jobsQueueMetadataPerQueue[rQueue.UID].jobsInQueue.Peek(1)
+		lJobsToPop := 1
+		if jobsOrder.queuePopsMap[lQueue.UID] > 0 {
+			lJobsToPop = jobsOrder.queuePopsMap[lQueue.UID]
+		}
+
+		rJobsToPop := 1
+		if jobsOrder.queuePopsMap[rQueue.UID] > 0 {
+			rJobsToPop = jobsOrder.queuePopsMap[rQueue.UID]
+		}
+
+		lJobInfosRaw := jobsQueueMetadataPerQueue[lQueue.UID].jobsInQueue.Peek(lJobsToPop)
+		rJobInfosRaw := jobsQueueMetadataPerQueue[rQueue.UID].jobsInQueue.Peek(rJobsToPop)
 
 		lJobInfos := podgroup_info.NewPodGroupInfos(lJobInfosRaw)
 		rJobInfos := podgroup_info.NewPodGroupInfos(rJobInfosRaw)
@@ -266,9 +288,17 @@ func (jobsOrder *JobsOrderByQueues) buildFuncOrderBetweenDepartmentsWithJobs(rev
 		lBestQueue := jobsOrder.departmentIdToDepartmentMetadata[lDepartment.UID].queuesPriorityQueue.Pop().(*queue_info.QueueInfo)
 		rBestQueue := jobsOrder.departmentIdToDepartmentMetadata[rDepartment.UID].queuesPriorityQueue.Pop().(*queue_info.QueueInfo)
 
-		lJobInfosRaw := jobsOrder.queueIdToQueueMetadata[lBestQueue.UID].jobsInQueue.Peek(1)
-		rJobInfosRaw := jobsOrder.queueIdToQueueMetadata[rBestQueue.UID].jobsInQueue.Peek(1)
+		lJobsToPop := 1
+		if jobsOrder.queuePopsMap[lBestQueue.UID] > 0 {
+			lJobsToPop = jobsOrder.queuePopsMap[lBestQueue.UID]
+		}
 
+		rJobsToPop := 1
+		if jobsOrder.queuePopsMap[rBestQueue.UID] > 0 {
+			rJobsToPop = jobsOrder.queuePopsMap[rBestQueue.UID]
+		}
+		lJobInfosRaw := jobsOrder.queueIdToQueueMetadata[lBestQueue.UID].jobsInQueue.Peek(lJobsToPop)
+		rJobInfosRaw := jobsOrder.queueIdToQueueMetadata[rBestQueue.UID].jobsInQueue.Peek(rJobsToPop)
 		lJobInfos := podgroup_info.NewPodGroupInfos(lJobInfosRaw)
 		rJobInfos := podgroup_info.NewPodGroupInfos(rJobInfosRaw)
 

pkg/scheduler/actions/utils/job_order_by_queue_test.go

@@ -360,6 +360,7 @@ func TestJobsOrderByQueues_PushJob(t *testing.T) {
 				departmentIdToDepartmentMetadata: tt.fields.departmentIdToDepartmentMetadata,
 				ssn:                              ssn,
 				jobsOrderInitOptions:             tt.fields.jobsOrderInitOptions,
+				queuePopsMap:                     map[common_info.QueueID]int{},
 			}
 			jobsOrder.InitializeWithJobs(tt.fields.InsertedJob)
 			jobsOrder.PushJob(tt.args.job)
@@ -457,6 +458,7 @@ func TestJobsOrderByQueues_RequeueJob(t *testing.T) {
 				departmentIdToDepartmentMetadata: tt.fields.departmentIdToDepartmentMetadata,
 				ssn:                              ssn,
 				jobsOrderInitOptions:             tt.fields.jobsOrderInitOptions,
+				queuePopsMap:                     map[common_info.QueueID]int{},
 			}
 			jobsOrder.InitializeWithJobs(tt.fields.InsertedJob)