https://github.com/NVIDIA/KAI-Scheduler/pull (#138)

/1381- Add caching to GetTasksToAllocate and GetTasksToAllocateInitResource 
2- GetActivelyRunningTasksCount returns the count instead of the pod
3- save the amount of ActiveAllocated pods in podgroup info instead of calculating them on request

Author: davidLif

pkg/scheduler/actions/common/allocate.go

@@ -131,7 +131,7 @@ func handleFailedTaskAllocation(job *podgroup_info.PodGroupInfo, unschedulableTa
 		allocationError.SetError(common_info.DefaultPodError)
 	}
 
-	numRunningTasks := int32(len(job.GetActivelyRunningTasks()))
+	numRunningTasks := job.GetActivelyRunningTasksCount()
 	if job.MinAvailable > 1 && numRunningTasks < job.MinAvailable {
 		job.SetJobFitError(
 			podgroup_info.PodSchedulingErrors,

pkg/scheduler/actions/consolidation/consolidation.go

@@ -147,7 +147,7 @@ func buildPreemptibleFilterFunc(preemptor *podgroup_info.PodGroupInfo, maxPreemp
 			return false
 		}
 
-		if len(job.GetActiveAllocatedTasks()) == 0 {
+		if job.GetActiveAllocatedTasksCount() == 0 {
 			return false
 		}
 

pkg/scheduler/actions/utils/job_order_by_queue_test.go

@@ -564,7 +564,9 @@ func TestJobsOrderByQueues_PushJob(t *testing.T) {
 			jobsOrder.PushJob(tt.args.job)
 
 			for _, expectedJob := range tt.expected.expectedJobsList {
+				_ = expectedJob.GetActiveAllocatedTasksCount()
 				actualJob := jobsOrder.PopNextJob()
+				_ = actualJob.GetActiveAllocatedTasksCount()
 				assert.Equal(t, expectedJob, actualJob)
 			}
 		})

pkg/scheduler/api/podgroup_info/allocation_info.go

@@ -27,6 +27,10 @@ func HasTasksToAllocate(podGroupInfo *PodGroupInfo, isRealAllocation bool) bool
 func GetTasksToAllocate(
 	podGroupInfo *PodGroupInfo, taskOrderFn common_info.LessFn, isRealAllocation bool,
 ) []*pod_info.PodInfo {
+	if podGroupInfo.tasksToAllocate != nil {
+		return podGroupInfo.tasksToAllocate
+	}
+
 	taskPriorityQueue := getTasksToAllocateQueue(podGroupInfo, taskOrderFn, isRealAllocation)
 	maxNumOfTasksToAllocate := getNumOfTasksToAllocate(podGroupInfo, taskPriorityQueue.Len())
 
@@ -36,6 +40,7 @@ func GetTasksToAllocate(
 		tasksToAllocate = append(tasksToAllocate, nextPod)
 	}
 
+	podGroupInfo.tasksToAllocate = tasksToAllocate
 	return tasksToAllocate
 }
 
@@ -77,18 +82,21 @@ func GetJobsToAllocateInitResource(
 func GetTasksToAllocateInitResource(
 	podGroupInfo *PodGroupInfo, taskOrderFn common_info.LessFn, isRealAllocation bool,
 ) *resource_info.Resource {
-	tasksTotalRequestedResource := resource_info.EmptyResource()
-
 	if podGroupInfo == nil {
-		return tasksTotalRequestedResource
+		return resource_info.EmptyResource()
+	}
+	if podGroupInfo.tasksToAllocateInitResource != nil {
+		return podGroupInfo.tasksToAllocateInitResource
 	}
 
+	tasksTotalRequestedResource := resource_info.EmptyResource()
 	for _, task := range GetTasksToAllocate(podGroupInfo, taskOrderFn, isRealAllocation) {
 		if task.ShouldAllocate(isRealAllocation) {
 			tasksTotalRequestedResource.AddResourceRequirements(task.ResReq)
 		}
 	}
 
+	podGroupInfo.tasksToAllocateInitResource = tasksTotalRequestedResource
 	return tasksTotalRequestedResource
 }
 

pkg/scheduler/api/podgroup_info/job_info.go

@@ -14,6 +14,7 @@ import (
 	policyv1 "k8s.io/api/policy/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/utils/ptr"
 
 	enginev2alpha2 "github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
 	commonconstants "github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
@@ -61,8 +62,7 @@ type PodGroupInfo struct {
 	NodesFitErrors map[common_info.PodID]*common_info.FitErrors
 
 	// All tasks of the Job.
-	PodStatusIndex map[pod_status.PodStatus]pod_info.PodsMap
-	PodInfos       pod_info.PodsMap
+	PodInfos pod_info.PodsMap
 
 	Allocated *resource_info.Resource
 
@@ -76,6 +76,12 @@ type PodGroupInfo struct {
 	StalenessInfo
 
 	schedulingConstraintsSignature common_info.SchedulingConstraintsSignature
+
+	// inner cache
+	tasksToAllocate             []*pod_info.PodInfo
+	tasksToAllocateInitResource *resource_info.Resource
+	PodStatusIndex              map[pod_status.PodStatus]pod_info.PodsMap
+	activeAllocatedCount        *int
 }
 
 func NewPodGroupInfo(uid common_info.PodGroupID, tasks ...*pod_info.PodInfo) *PodGroupInfo {
@@ -94,6 +100,7 @@ func NewPodGroupInfo(uid common_info.PodGroupID, tasks ...*pod_info.PodInfo) *Po
 			TimeStamp: nil,
 			Stale:     false,
 		},
+		activeAllocatedCount: ptr.To(0),
 	}
 
 	for _, task := range tasks {
@@ -155,6 +162,11 @@ func (podGroupInfo *PodGroupInfo) addTaskIndex(ti *pod_info.PodInfo) {
 	}
 
 	podGroupInfo.PodStatusIndex[ti.Status][ti.UID] = ti
+	if pod_status.IsActiveAllocatedStatus(ti.Status) {
+		podGroupInfo.activeAllocatedCount = ptr.To(*podGroupInfo.activeAllocatedCount + 1)
+	}
+
+	podGroupInfo.invalidateTasksCache()
 }
 
 func (podGroupInfo *PodGroupInfo) AddTaskInfo(ti *pod_info.PodInfo) {
@@ -182,31 +194,44 @@ func (podGroupInfo *PodGroupInfo) UpdateTaskStatus(task *pod_info.PodInfo, statu
 func (podGroupInfo *PodGroupInfo) deleteTaskIndex(ti *pod_info.PodInfo) {
 	if tasks, found := podGroupInfo.PodStatusIndex[ti.Status]; found {
 		delete(tasks, ti.UID)
+		if pod_status.IsActiveAllocatedStatus(ti.Status) {
+			podGroupInfo.activeAllocatedCount = ptr.To(*podGroupInfo.activeAllocatedCount - 1)
+		}
 
 		if len(tasks) == 0 {
 			delete(podGroupInfo.PodStatusIndex, ti.Status)
 		}
+
+		podGroupInfo.invalidateTasksCache()
 	}
 }
 
-func (podGroupInfo *PodGroupInfo) GetActiveAllocatedTasks() []*pod_info.PodInfo {
-	var tasksToAllocate []*pod_info.PodInfo
-	for _, task := range podGroupInfo.PodInfos {
-		if pod_status.IsActiveAllocatedStatus(task.Status) {
-			tasksToAllocate = append(tasksToAllocate, task)
+func (podGroupInfo *PodGroupInfo) invalidateTasksCache() {
+	podGroupInfo.tasksToAllocate = nil
+	podGroupInfo.tasksToAllocateInitResource = nil
+}
+
+func (podGroupInfo *PodGroupInfo) GetActiveAllocatedTasksCount() int {
+	if podGroupInfo.activeAllocatedCount == nil {
+		var taskCount int
+		for _, task := range podGroupInfo.PodInfos {
+			if pod_status.IsActiveAllocatedStatus(task.Status) {
+				taskCount++
+			}
 		}
+		podGroupInfo.activeAllocatedCount = ptr.To(taskCount)
 	}
-	return tasksToAllocate
+	return *podGroupInfo.activeAllocatedCount
 }
 
-func (podGroupInfo *PodGroupInfo) GetActivelyRunningTasks() []*pod_info.PodInfo {
-	var tasks []*pod_info.PodInfo
+func (podGroupInfo *PodGroupInfo) GetActivelyRunningTasksCount() int32 {
+	tasksCount := int32(0)
 	for _, task := range podGroupInfo.PodInfos {
 		if pod_status.IsActiveUsedStatus(task.Status) {
-			tasks = append(tasks, task)
+			tasksCount += 1
 		}
 	}
-	return tasks
+	return tasksCount
 }
 
 func (podGroupInfo *PodGroupInfo) DeleteTaskInfo(ti *pod_info.PodInfo) error {
@@ -342,8 +367,9 @@ func (podGroupInfo *PodGroupInfo) CloneWithTasks(tasks []*pod_info.PodInfo) *Pod
 		PodGroup:    podGroupInfo.PodGroup,
 		PodGroupUID: podGroupInfo.PodGroupUID,
 
-		PodStatusIndex: map[pod_status.PodStatus]pod_info.PodsMap{},
-		PodInfos:       pod_info.PodsMap{},
+		PodStatusIndex:       map[pod_status.PodStatus]pod_info.PodsMap{},
+		PodInfos:             pod_info.PodsMap{},
+		activeAllocatedCount: ptr.To(0),
 	}
 
 	podGroupInfo.CreationTimestamp.DeepCopyInto(&info.CreationTimestamp)

pkg/scheduler/api/podgroup_info/job_info_test.go

@@ -9,6 +9,7 @@ import (
 
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/utils/ptr"
 
 	"github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
 	commonconstants "github.com/NVIDIA/KAI-scheduler/pkg/common/constants"
@@ -75,8 +76,9 @@ func TestAddTaskInfo(t *testing.T) {
 						case01_task4.UID: case01_task4,
 					},
 				},
-				JobFitErrors:   make(v2alpha2.UnschedulableExplanations, 0),
-				NodesFitErrors: map[common_info.PodID]*common_info.FitErrors{},
+				activeAllocatedCount: ptr.To(3),
+				JobFitErrors:         make(v2alpha2.UnschedulableExplanations, 0),
+				NodesFitErrors:       map[common_info.PodID]*common_info.FitErrors{},
 			},
 		},
 	}
@@ -150,8 +152,9 @@ func TestDeleteTaskInfo(t *testing.T) {
 					pod_status.Pending: {case01_task1.UID: case01_task1},
 					pod_status.Running: {case01_task3.UID: case01_task3},
 				},
-				JobFitErrors:   make(v2alpha2.UnschedulableExplanations, 0),
-				NodesFitErrors: map[common_info.PodID]*common_info.FitErrors{},
+				activeAllocatedCount: ptr.To(1),
+				JobFitErrors:         make(v2alpha2.UnschedulableExplanations, 0),
+				NodesFitErrors:       map[common_info.PodID]*common_info.FitErrors{},
 			},
 		},
 		{
@@ -175,8 +178,9 @@ func TestDeleteTaskInfo(t *testing.T) {
 						case02_task3.UID: case02_task3,
 					},
 				},
-				JobFitErrors:   make(v2alpha2.UnschedulableExplanations, 0),
-				NodesFitErrors: map[common_info.PodID]*common_info.FitErrors{},
+				activeAllocatedCount: ptr.To(1),
+				JobFitErrors:         make(v2alpha2.UnschedulableExplanations, 0),
+				NodesFitErrors:       map[common_info.PodID]*common_info.FitErrors{},
 			},
 		},
 	}

pkg/scheduler/cache/status_updater/default_status_updater.go

@@ -217,7 +217,7 @@ func (su *defaultStatusUpdater) recordJobNotReadyEvent(job *podgroup_info.PodGro
 func (su *defaultStatusUpdater) markPodGroupUnschedulable(job *podgroup_info.PodGroupInfo, message string) bool {
 	su.recorder.Event(job.PodGroup, v1.EventTypeNormal, enginev2alpha2.PodGroupReasonUnschedulable, message)
 
-	if len(job.GetActiveAllocatedTasks()) > 0 {
+	if job.GetActiveAllocatedTasksCount() > 0 {
 		// Don't update podgroup condition if there are any allocated pods (RUN-20673)
 		return false
 	}

pkg/scheduler/plugins/elastic/elastic.go

@@ -29,8 +29,8 @@ func JobOrderFn(l, r interface{}) int {
 	lv := l.(*podgroup_info.PodGroupInfo)
 	rv := r.(*podgroup_info.PodGroupInfo)
 
-	lvAllocatedCount := int32(len(lv.GetActiveAllocatedTasks()))
-	rvAllocatedCount := int32(len(rv.GetActiveAllocatedTasks()))
+	lvAllocatedCount := int32(lv.GetActiveAllocatedTasksCount())
+	rvAllocatedCount := int32(rv.GetActiveAllocatedTasksCount())
 
 	if lvAllocatedCount < lv.MinAvailable && rvAllocatedCount >= rv.MinAvailable {
 		return -1