Add default sg only if no sgs

Author: omer-dayan

pkg/scheduler/api/podgroup_info/allocation_info_test.go

@@ -135,7 +135,9 @@ func Test_getTasksFromQueue(t *testing.T) {
 func Test_getTasksPriorityQueuePerSubGroup(t *testing.T) {
 	pg := NewPodGroupInfo("test-pg")
 	sg := NewSubGroupInfo("test-sub-group", 1)
-	pg.SubGroups["test-sub-group"] = sg
+	pg.SubGroups = map[string]*SubGroupInfo{
+		"test-sub-group": sg,
+	}
 
 	pg.AddTaskInfo(simpleTask("a", "test-sub-group", pod_status.Pending))
 	m := getTasksPriorityQueuePerSubGroup(pg, tasksOrderFn, true)

pkg/scheduler/api/podgroup_info/job_info.go

@@ -110,7 +110,9 @@ func NewPodGroupInfo(uid common_info.PodGroupID, tasks ...*pod_info.PodInfo) *Po
 			Stale:     false,
 		},
 
-		SubGroups: map[string]*SubGroupInfo{},
+		SubGroups: map[string]*SubGroupInfo{
+			DefaultSubGroup: NewSubGroupInfo(DefaultSubGroup, 1),
+		},
 
 		LastStartTimestamp:   nil,
 		activeAllocatedCount: ptr.To(0),
@@ -176,10 +178,12 @@ func (pgi *PodGroupInfo) SetPodGroup(pg *enginev2alpha2.PodGroup) {
 		if _, exists := pgi.SubGroups[DefaultSubGroup]; !exists {
 			pgi.SubGroups[DefaultSubGroup] = NewSubGroupInfo(DefaultSubGroup, max(pg.Spec.MinMember, 1))
 		}
-	}
-	for _, sg := range pg.Spec.SubGroups {
-		subGroupInfo := FromSubGroup(&sg)
-		pgi.SubGroups[subGroupInfo.name] = subGroupInfo
+	} else {
+		pgi.SubGroups = map[string]*SubGroupInfo{}
+		for _, sg := range pg.Spec.SubGroups {
+			subGroupInfo := FromSubGroup(&sg)
+			pgi.SubGroups[subGroupInfo.name] = subGroupInfo
+		}
 	}
 
 	if pg.Annotations[commonconstants.StalePodgroupTimeStamp] != "" {
@@ -222,16 +226,17 @@ func (pgi *PodGroupInfo) addTaskIndex(ti *pod_info.PodInfo) {
 }
 
 func (pgi *PodGroupInfo) AddTaskInfo(ti *pod_info.PodInfo) {
-	subGroup, found := pgi.SubGroups[ti.SubGroupName]
-	if found {
-		subGroup.AssignTask(ti)
-	} else {
-		if pgi.SubGroups[DefaultSubGroup] == nil {
-			pgi.SubGroups[DefaultSubGroup] = NewSubGroupInfo(DefaultSubGroup, 0)
-		}
-		pgi.SubGroups[DefaultSubGroup].AssignTask(ti)
+	taskSubGroupName := DefaultSubGroup
+	if ti.SubGroupName != "" {
+		taskSubGroupName = ti.SubGroupName
+	}
+	subGroup, found := pgi.SubGroups[taskSubGroupName]
+	if !found {
+		log.InfraLogger.Warningf("AddTaskInfo for task <%s/%s> of podGroup: <%s/%s>: SubGroup not found <%s>", ti.Namespace, ti.Name, pgi.Namespace, pgi.Name, taskSubGroupName)
+		return
 	}
 
+	subGroup.AssignTask(ti)
 	pgi.addTaskIndex(ti)
 
 	if pod_status.AllocatedStatus(ti.Status) {

pkg/scheduler/api/podgroup_info/job_info_test.go

@@ -74,7 +74,7 @@ func TestAddTaskInfo(t *testing.T) {
 			expected: &PodGroupInfo{
 				UID:       case01_uid,
 				Allocated: common_info.BuildResource("4000m", "4G"),
-				SubGroups: map[string]*SubGroupInfo{DefaultSubGroup: NewSubGroupInfo(DefaultSubGroup, 0).WithPodInfos(pod_info.PodsMap{
+				SubGroups: map[string]*SubGroupInfo{DefaultSubGroup: NewSubGroupInfo(DefaultSubGroup, 1).WithPodInfos(pod_info.PodsMap{
 					case01_task1.UID: case01_task1,
 					case01_task2.UID: case01_task2,
 					case01_task3.UID: case01_task3,
@@ -155,7 +155,7 @@ func TestDeleteTaskInfo(t *testing.T) {
 			expected: &PodGroupInfo{
 				UID:       case01_uid,
 				Allocated: common_info.BuildResource("3000m", "3G"),
-				SubGroups: map[string]*SubGroupInfo{DefaultSubGroup: NewSubGroupInfo(DefaultSubGroup, 0).WithPodInfos(pod_info.PodsMap{
+				SubGroups: map[string]*SubGroupInfo{DefaultSubGroup: NewSubGroupInfo(DefaultSubGroup, 1).WithPodInfos(pod_info.PodsMap{
 					case01_task1.UID: case01_task1,
 					case01_task2.UID: case01_task2,
 					case01_task3.UID: case01_task3,
@@ -177,7 +177,7 @@ func TestDeleteTaskInfo(t *testing.T) {
 			expected: &PodGroupInfo{
 				UID:       case02_uid,
 				Allocated: common_info.BuildResource("3000m", "3G"),
-				SubGroups: map[string]*SubGroupInfo{DefaultSubGroup: NewSubGroupInfo(DefaultSubGroup, 0).WithPodInfos(pod_info.PodsMap{
+				SubGroups: map[string]*SubGroupInfo{DefaultSubGroup: NewSubGroupInfo(DefaultSubGroup, 1).WithPodInfos(pod_info.PodsMap{
 					case02_task1.UID: case02_task1,
 					case02_task2.UID: case02_task2,
 					case02_task3.UID: case02_task3,
@@ -993,10 +993,11 @@ func TestPodGroupInfo_IsStale(t *testing.T) {
 				pgi := NewPodGroupInfo("test-podgroup")
 
 				sg1 := NewSubGroupInfo("sg1", 1)
-				pgi.SubGroups["sg1"] = sg1
-
 				sg2 := NewSubGroupInfo("sg2", 1)
-				pgi.SubGroups["sg2"] = sg2
+				pgi.SubGroups = map[string]*SubGroupInfo{
+					"sg1": sg1,
+					"sg2": sg2,
+				}
 
 				task1 := pod_info.NewTaskInfo(&v1.Pod{
 					ObjectMeta: metav1.ObjectMeta{