Added SubGroupOrder plugin (#357)

* Added SubGroupOrder plugin

* Added prioritization consts to subgroup order plugin

Author: romanbaron

deployments/kai-scheduler/templates/services/scheduler-configmap.yaml

@@ -16,6 +16,7 @@ data:
           - name: elastic
           - name: kubeflow
           - name: ray
+          - name: subgrouporder
           - name: taskorder
           - name: nominatednode
           - name: dynamicresources

pkg/scheduler/actions/allocate/allocate.go

@@ -79,7 +79,7 @@ func (alloc *allocateAction) Execute(ssn *framework.Session) {
 func attemptToAllocateJob(ssn *framework.Session, stmt *framework.Statement, job *podgroup_info.PodGroupInfo) (allocated, pipelined bool) {
 	queue := ssn.Queues[job.Queue]
 
-	resReq := podgroup_info.GetTasksToAllocateInitResource(job, ssn.TaskOrderFn, true)
+	resReq := podgroup_info.GetTasksToAllocateInitResource(job, ssn.SubGroupOrderFn, ssn.TaskOrderFn, true)
 	log.InfraLogger.V(3).Infof("Attempting to allocate job: <%v/%v> of queue <%v>, resources: <%v>",
 		job.Namespace, job.Name, queue.Name, resReq)
 

pkg/scheduler/actions/common/action.go

@@ -94,13 +94,14 @@ func TryToVirtuallyAllocatePreemptorAndGetVictims(
 		}
 
 		resReq := podgroup_info.GetTasksToAllocateInitResource(
-			jobToAllocate, ssn.TaskOrderFn, false)
+			jobToAllocate, ssn.SubGroupOrderFn, ssn.TaskOrderFn, false)
 		log.InfraLogger.V(6).Infof("Trying to pipeline job: <%s/%s>. resources required: %v",
 			jobToAllocate.Namespace, jobToAllocate.Name, resReq)
 
 		if jobToAllocate.UID != preemptor.UID {
 			if !AllocateJob(ssn, stmt, nodes, jobToAllocate, true) {
-				tasksToAllocate := podgroup_info.GetTasksToAllocate(jobToAllocate, ssn.TaskOrderFn, false)
+				tasksToAllocate := podgroup_info.GetTasksToAllocate(jobToAllocate, ssn.SubGroupOrderFn,
+					ssn.TaskOrderFn, false)
 				newVictims = append(newVictims, tasksToAllocate...)
 			}
 			continue

pkg/scheduler/actions/common/allocate.go

@@ -17,7 +17,7 @@ import (
 
 func AllocateJob(ssn *framework.Session, stmt *framework.Statement, nodes []*node_info.NodeInfo,
 	job *podgroup_info.PodGroupInfo, isPipelineOnly bool) bool {
-	tasksToAllocate := podgroup_info.GetTasksToAllocate(job, ssn.TaskOrderFn, !isPipelineOnly)
+	tasksToAllocate := podgroup_info.GetTasksToAllocate(job, ssn.SubGroupOrderFn, ssn.TaskOrderFn, !isPipelineOnly)
 
 	result := ssn.IsJobOverQueueCapacityFn(job, tasksToAllocate)
 	if !result.IsSchedulable {

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

@@ -51,7 +51,7 @@ func (s *JobSolver) Solve(
 
 	var statement *framework.Statement
 	var pendingTasks []*pod_info.PodInfo
-	tasksToAllocate := podgroup_info.GetTasksToAllocate(pendingJob, ssn.TaskOrderFn, false)
+	tasksToAllocate := podgroup_info.GetTasksToAllocate(pendingJob, ssn.SubGroupOrderFn, ssn.TaskOrderFn, false)
 	for _, nextTaskToSolve := range tasksToAllocate {
 		nextTasksToSolve := []*pod_info.PodInfo{nextTaskToSolve}
 		pendingTasks = append(pendingTasks, nextTasksToSolve...)

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

@@ -34,7 +34,7 @@ func NewPodAccumulatedScenarioBuilder(
 
 	var scenario *solverscenario.ByNodeScenario = nil
 	recordedVictimsTasks := make(map[common_info.PodID]*pod_info.PodInfo)
-	tasksToAllocate := podgroup_info.GetTasksToAllocate(pendingJob, session.TaskOrderFn, false)
+	tasksToAllocate := podgroup_info.GetTasksToAllocate(pendingJob, session.SubGroupOrderFn, session.TaskOrderFn, false)
 	if len(tasksToAllocate) != 0 {
 		scenario = solverscenario.NewByNodeScenario(session, pendingJob, pendingJob, nil, recordedVictimsJobs)
 		for _, job := range recordedVictimsJobs {
@@ -76,7 +76,7 @@ func (asb *PodAccumulatedScenarioBuilder) addNextPotentialVictims() bool {
 	nextVictimJob := asb.victimsJobsQueue.PopNextJob()
 
 	potentialVictimTasks, jobHasMoreTasks := podgroup_info.GetTasksToEvict(
-		nextVictimJob, asb.session.TaskOrderFn,
+		nextVictimJob, asb.session.SubGroupOrderFn, asb.session.TaskOrderFn,
 	)
 
 	// Jump over recorded victims in potential victims generation

pkg/scheduler/actions/consolidation/consolidation.go

@@ -79,7 +79,7 @@ func (alloc *consolidationAction) Execute(ssn *framework.Session) {
 
 func attemptToConsolidateForPreemptor(
 	ssn *framework.Session, job *podgroup_info.PodGroupInfo) (bool, *framework.Statement) {
-	resReq := podgroup_info.GetTasksToAllocateInitResource(job, ssn.TaskOrderFn, false)
+	resReq := podgroup_info.GetTasksToAllocateInitResource(job, ssn.SubGroupOrderFn, ssn.TaskOrderFn, false)
 	log.InfraLogger.V(3).Infof(
 		"Attempting to consolidate running jobs in order to make room for job: <%s/%s>, resources: <%v>",
 		job.Namespace, job.Name, resReq)

pkg/scheduler/actions/preempt/preempt.go

@@ -99,12 +99,12 @@ func (alloc *preemptAction) Execute(ssn *framework.Session) {
 func attemptToPreemptForPreemptor(
 	ssn *framework.Session, preemptor *podgroup_info.PodGroupInfo,
 ) (bool, *framework.Statement, []string) {
-	resReq := podgroup_info.GetTasksToAllocateInitResource(preemptor, ssn.TaskOrderFn, false)
+	resReq := podgroup_info.GetTasksToAllocateInitResource(preemptor, ssn.SubGroupOrderFn, ssn.TaskOrderFn, false)
 	log.InfraLogger.V(3).Infof(
 		"Attempting to preempt for job: <%v/%v>, priority: <%v>, queue: <%v>, resources: <%v>",
 		preemptor.Namespace, preemptor.Name, preemptor.Priority, preemptor.Queue, resReq)
 
-	preemptorTasks := podgroup_info.GetTasksToAllocate(preemptor, ssn.TaskOrderFn, false)
+	preemptorTasks := podgroup_info.GetTasksToAllocate(preemptor, ssn.SubGroupOrderFn, ssn.TaskOrderFn, false)
 	if result := ssn.IsNonPreemptibleJobOverQueueQuotaFn(preemptor, preemptorTasks); !result.IsSchedulable {
 		log.InfraLogger.V(3).Infof("Job <%v/%v> would have placed the queue resources over quota",
 			preemptor.Namespace, preemptor.Name)

pkg/scheduler/actions/reclaim/reclaim.go

@@ -103,7 +103,7 @@ func (ra *reclaimAction) attemptToReclaimForSpecificJob(
 	ssn *framework.Session, reclaimer *podgroup_info.PodGroupInfo,
 ) (bool, *framework.Statement, []string) {
 	queue := ssn.Queues[reclaimer.Queue]
-	resReq := podgroup_info.GetTasksToAllocateInitResource(reclaimer, ssn.TaskOrderFn, false)
+	resReq := podgroup_info.GetTasksToAllocateInitResource(reclaimer, ssn.SubGroupOrderFn, ssn.TaskOrderFn, false)
 	log.InfraLogger.V(3).Infof("Attempting to reclaim for job: <%v/%v> of queue <%v>, resources: <%v>",
 		reclaimer.Namespace, reclaimer.Name, queue.Name, resReq)
 

pkg/scheduler/actions/utils/action.go

@@ -130,9 +130,9 @@ func GetAllPendingJobs(ssn *framework.Session) map[common_info.PodGroupID]*podgr
 
 func IsEnoughGPUsAllocatableForJob(job *podgroup_info.PodGroupInfo, ssn *framework.Session, isRealAllocation bool) bool {
 	sumOfAllAllocatableGPUs, sumOfAllAllocatableGPUsMemory := getSumOfAvailableGPUs(ssn)
-	requestedGPUs, requestedGpuMemory := podgroup_info.GetTasksToAllocateRequestedGPUs(job, ssn.TaskOrderFn,
-		isRealAllocation)
-	resReq := podgroup_info.GetTasksToAllocateInitResource(job, ssn.TaskOrderFn, isRealAllocation)
+	requestedGPUs, requestedGpuMemory := podgroup_info.GetTasksToAllocateRequestedGPUs(job, ssn.SubGroupOrderFn,
+		ssn.TaskOrderFn, isRealAllocation)
+	resReq := podgroup_info.GetTasksToAllocateInitResource(job, ssn.SubGroupOrderFn, ssn.TaskOrderFn, isRealAllocation)
 	log.InfraLogger.V(7).Infof(
 		"Task: <%v/%v> resources requires: <%v>, sumOfAllAllocatableGPUs: <%v, %v mb>",
 		job.Namespace, job.Name, resReq, sumOfAllAllocatableGPUs, sumOfAllAllocatableGPUsMemory)