Did a small refactoring in allocation_info.go (#348)

Author: romanbaron

pkg/scheduler/api/podgroup_info/allocation_info.go

@@ -9,7 +9,6 @@ import (
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/common_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/common_info/resources"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_info"
-	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_status"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/resource_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/log"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/scheduler_util"
@@ -67,18 +66,6 @@ func GetTasksToAllocateRequestedGPUs(
 	return tasksTotalRequestedGPUs, tasksTotalRequestedGpuMemory
 }
 
-func GetJobsToAllocateInitResource(
-	podGroupInfos []*PodGroupInfo, taskOrderFn common_info.LessFn, isRealAllocation bool,
-) *resource_info.Resource {
-	tasksTotalRequestedResource := resource_info.EmptyResource()
-	for _, podGroupInfo := range podGroupInfos {
-		pgInitResource := GetTasksToAllocateInitResource(podGroupInfo, taskOrderFn, isRealAllocation)
-		tasksTotalRequestedResource.Add(pgInitResource)
-	}
-
-	return tasksTotalRequestedResource
-}
-
 func GetTasksToAllocateInitResource(
 	podGroupInfo *PodGroupInfo, taskOrderFn common_info.LessFn, isRealAllocation bool,
 ) *resource_info.Resource {
@@ -113,7 +100,7 @@ func getTasksToAllocateQueue(
 }
 
 func getNumOfTasksToAllocate(podGroupInfo *PodGroupInfo, numOfTasksWaitingAllocation int) int {
-	allocatedTasks := int32(getNumOfAllocatedTasks(podGroupInfo))
+	allocatedTasks := int32(podGroupInfo.GetActiveAllocatedTasksCount())
 
 	var maxTasksToAllocate int32
 	if allocatedTasks >= podGroupInfo.MinAvailable {
@@ -124,13 +111,3 @@ func getNumOfTasksToAllocate(podGroupInfo *PodGroupInfo, numOfTasksWaitingAlloca
 
 	return int(math.Min(float64(maxTasksToAllocate), float64(numOfTasksWaitingAllocation)))
 }
-
-func getNumOfAllocatedTasks(podGroupInfo *PodGroupInfo) int {
-	allocatedTasks := 0
-	for _, task := range podGroupInfo.PodInfos {
-		if pod_status.IsActiveAllocatedStatus(task.Status) {
-			allocatedTasks += 1
-		}
-	}
-	return allocatedTasks
-}

pkg/scheduler/api/podgroup_info/allocation_info_test.go

@@ -123,11 +123,10 @@ func Test_getMaxNumOfTasksToAllocate(t *testing.T) {
 			pg.MinAvailable = tt.args.minAvailable
 			for i, pod := range tt.args.pods {
 				pi := pod_info.NewTaskInfo(pod)
-				pg.AddTaskInfo(pi)
-
 				if tt.args.overridingStatus != nil {
 					pi.Status = tt.args.overridingStatus[i]
 				}
+				pg.AddTaskInfo(pi)
 			}
 
 			if got := getNumOfTasksToAllocate(pg, tt.args.podsWantingToAllocate); got != tt.want {
@@ -222,7 +221,7 @@ func Test_getNumOfAllocatedTasks(t *testing.T) {
 				}
 			}
 
-			if got := getNumOfAllocatedTasks(pg); got != tt.want {
+			if got := pg.GetActiveAllocatedTasksCount(); got != tt.want {
 				t.Errorf("getNumOfAllocatedTasks() = %v, want %v", got, tt.want)
 			}
 		})