Added SubGroups to PodGroup CRD and SubGroupInfo to PodGroupInfo (#342)

Added SubGroups to PodGroup CRD and SubGroupInfo to PodGroupInfo

Author: romanbaron

deployments/crds/internal/scheduling.run.ai_podgroups.yaml

@@ -87,6 +87,25 @@ spec:
                   -1 and 1
                 format: int32
                 type: integer
+              subGroups:
+                description: SubGroups defines finer-grained subsets of pods within
+                  the PodGroup with individual scheduling constraints
+                items:
+                  properties:
+                    minMember:
+                      description: |-
+                        MinMember defines the minimal number of members/tasks to run this SubGroup;
+                        if there's not enough resources to start all tasks, the scheduler will not start anyone.
+                      format: int32
+                      type: integer
+                    name:
+                      description: Name uniquely identifies the SubGroup within the
+                        parent PodGroup.
+                      type: string
+                  required:
+                  - name
+                  type: object
+                type: array
               topologyConstraint:
                 description: TopologyConstraint defines the topology constraints for
                   this PodGroup

pkg/apis/scheduling/v2alpha2/podgroup_types.go

@@ -66,6 +66,18 @@ type PodGroupSpec struct {
 
 	// TopologyConstraint defines the topology constraints for this PodGroup
 	TopologyConstraint TopologyConstraint `json:"topologyConstraint,omitempty"`
+
+	// SubGroups defines finer-grained subsets of pods within the PodGroup with individual scheduling constraints
+	SubGroups []SubGroup `json:"subGroups,omitempty"`
+}
+
+type SubGroup struct {
+	// Name uniquely identifies the SubGroup within the parent PodGroup.
+	Name string `json:"name"`
+
+	// MinMember defines the minimal number of members/tasks to run this SubGroup;
+	// if there's not enough resources to start all tasks, the scheduler will not start anyone.
+	MinMember int32 `json:"minMember,omitempty"`
 }
 
 // PodGroupStatus defines the observed state of PodGroup

pkg/apis/scheduling/v2alpha2/zz_generated.deepcopy.go

@@ -42,6 +42,7 @@ const (
 	GPUGroup                           = "runai-gpu-group"
 	ReceivedResourceTypeAnnotationName = "received-resource-type"
 	WholeGpuIndicator                  = "-2"
+	SubGroupLabelKey                   = "kai.scheduler/subgroup-name"
 )
 
 type ResourceRequestType string
@@ -71,6 +72,8 @@ type PodInfo struct {
 	Name      string
 	Namespace string
 
+	SubGroupName string
+
 	ResourceRequestType  ResourceRequestType
 	ResourceReceivedType ResourceReceivedType
 
@@ -175,6 +178,7 @@ func NewTaskInfoWithBindRequest(pod *v1.Pod, bindRequest *bindrequest_info.BindR
 		Job:                            getPodGroupID(pod),
 		Name:                           pod.Name,
 		Namespace:                      pod.Namespace,
+		SubGroupName:                   pod.Labels[SubGroupLabelKey],
 		NodeName:                       nodeName,
 		Status:                         getTaskStatus(pod, bindRequest),
 		IsVirtualStatus:                false,

pkg/scheduler/api/pod_info/pod_info.go

@@ -86,6 +86,7 @@ type PodGroupInfo struct {
 	LastStartTimestamp *time.Time
 	PodGroup           *enginev2alpha2.PodGroup
 	PodGroupUID        types.UID
+	SubGroups          map[string]*SubGroupInfo
 
 	StalenessInfo
 
@@ -115,6 +116,8 @@ func NewPodGroupInfo(uid common_info.PodGroupID, tasks ...*pod_info.PodInfo) *Po
 			Stale:     false,
 		},
 
+		SubGroups: map[string]*SubGroupInfo{},
+
 		LastStartTimestamp:   nil,
 		activeAllocatedCount: ptr.To(0),
 	}
@@ -146,6 +149,11 @@ func (pgi *PodGroupInfo) SetPodGroup(pg *enginev2alpha2.PodGroup) {
 	pgi.PodGroup = pg
 	pgi.PodGroupUID = pg.UID
 
+	for _, sg := range pg.Spec.SubGroups {
+		subGroupInfo := fromSubGroup(&sg)
+		pgi.SubGroups[subGroupInfo.Name] = subGroupInfo
+	}
+
 	if pg.Annotations[commonconstants.StalePodgroupTimeStamp] != "" {
 		staleTimeStamp, err := time.Parse(time.RFC3339, pg.Annotations[commonconstants.StalePodgroupTimeStamp])
 		if err != nil {
@@ -187,6 +195,10 @@ func (pgi *PodGroupInfo) addTaskIndex(ti *pod_info.PodInfo) {
 
 func (pgi *PodGroupInfo) AddTaskInfo(ti *pod_info.PodInfo) {
 	pgi.PodInfos[ti.UID] = ti
+	subGroup, found := pgi.SubGroups[ti.SubGroupName]
+	if found {
+		subGroup.assignTask(ti)
+	}
 	pgi.addTaskIndex(ti)
 
 	if pod_status.AllocatedStatus(ti.Status) {
@@ -397,6 +409,7 @@ func (pgi *PodGroupInfo) CloneWithTasks(tasks []*pod_info.PodInfo) *PodGroupInfo
 
 		PodGroup:    pgi.PodGroup,
 		PodGroupUID: pgi.PodGroupUID,
+		SubGroups:   map[string]*SubGroupInfo{},
 
 		PodStatusIndex:       map[pod_status.PodStatus]pod_info.PodsMap{},
 		PodInfos:             pod_info.PodsMap{},
@@ -405,6 +418,10 @@ func (pgi *PodGroupInfo) CloneWithTasks(tasks []*pod_info.PodInfo) *PodGroupInfo
 
 	pgi.CreationTimestamp.DeepCopyInto(&info.CreationTimestamp)
 
+	for _, subGroup := range pgi.SubGroups {
+		info.SubGroups[subGroup.Name] = newSubGroupInfo(subGroup.Name, subGroup.MinAvailable)
+	}
+
 	for _, task := range tasks {
 		info.AddTaskInfo(task.Clone())
 	}

pkg/scheduler/api/podgroup_info/job_info.go

@@ -74,6 +74,7 @@ func TestAddTaskInfo(t *testing.T) {
 			expected: &PodGroupInfo{
 				UID:       case01_uid,
 				Allocated: common_info.BuildResource("4000m", "4G"),
+				SubGroups: map[string]*SubGroupInfo{},
 				PodInfos: pod_info.PodsMap{
 					case01_task1.UID: case01_task1,
 					case01_task2.UID: case01_task2,
@@ -155,6 +156,7 @@ func TestDeleteTaskInfo(t *testing.T) {
 			expected: &PodGroupInfo{
 				UID:       case01_uid,
 				Allocated: common_info.BuildResource("3000m", "3G"),
+				SubGroups: map[string]*SubGroupInfo{},
 				PodInfos: pod_info.PodsMap{
 					case01_task1.UID: case01_task1,
 					case01_task2.UID: case01_task2,
@@ -177,6 +179,7 @@ func TestDeleteTaskInfo(t *testing.T) {
 			expected: &PodGroupInfo{
 				UID:       case02_uid,
 				Allocated: common_info.BuildResource("3000m", "3G"),
+				SubGroups: map[string]*SubGroupInfo{},
 				PodInfos: pod_info.PodsMap{
 					case02_task1.UID: case02_task1,
 					case02_task2.UID: case02_task2,

pkg/scheduler/api/podgroup_info/job_info_test.go

@@ -0,0 +1,31 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package podgroup_info
+
+import (
+	"github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
+	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_info"
+)
+
+type SubGroupInfo struct {
+	Name         string
+	MinAvailable int32
+	PodInfos     pod_info.PodsMap
+}
+
+func newSubGroupInfo(name string, minAvailable int32) *SubGroupInfo {
+	return &SubGroupInfo{
+		Name:         name,
+		MinAvailable: minAvailable,
+		PodInfos:     pod_info.PodsMap{},
+	}
+}
+
+func fromSubGroup(subGroup *v2alpha2.SubGroup) *SubGroupInfo {
+	return newSubGroupInfo(subGroup.Name, subGroup.MinMember)
+}
+
+func (sgi *SubGroupInfo) assignTask(ti *pod_info.PodInfo) {
+	sgi.PodInfos[ti.UID] = ti
+}

pkg/scheduler/api/podgroup_info/subgroup_info.go

@@ -0,0 +1,75 @@
+// Copyright 2025 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package podgroup_info
+
+import (
+	"testing"
+
+	"github.com/NVIDIA/KAI-scheduler/pkg/apis/scheduling/v2alpha2"
+	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_info"
+	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_status"
+)
+
+func TestNewSubGroupInfo(t *testing.T) {
+	name := "my-subgroup"
+	minAvailable := int32(4)
+	sgi := newSubGroupInfo(name, minAvailable)
+
+	if sgi.Name != name {
+		t.Errorf("Expected Name %s, got %s", name, sgi.Name)
+	}
+	if sgi.MinAvailable != minAvailable {
+		t.Errorf("Expected MinAvailable %d, got %d", minAvailable, sgi.MinAvailable)
+	}
+	if len(sgi.PodInfos) != 0 {
+		t.Errorf("Expected empty PodInfos, got %d items", len(sgi.PodInfos))
+	}
+}
+
+func TestFromSubGroup(t *testing.T) {
+	subGroup := &v2alpha2.SubGroup{
+		Name:      "test-subgroup",
+		MinMember: 3,
+	}
+
+	sgi := fromSubGroup(subGroup)
+	if sgi.Name != subGroup.Name {
+		t.Errorf("Expected name %s, got %s", subGroup.Name, sgi.Name)
+	}
+	if sgi.MinAvailable != subGroup.MinMember {
+		t.Errorf("Expected MinAvailable %d, got %d", subGroup.MinMember, sgi.MinAvailable)
+	}
+	if len(sgi.PodInfos) != 0 {
+		t.Errorf("Expected empty PodInfos, got %d items", len(sgi.PodInfos))
+	}
+}
+
+func TestAddTaskInfoToSubGroup(t *testing.T) {
+	sgi := newSubGroupInfo("test", 1)
+	podInfo := &pod_info.PodInfo{
+		UID:    "pod-1",
+		Status: pod_status.Pending,
+	}
+
+	sgi.assignTask(podInfo)
+	if len(sgi.PodInfos) != 1 {
+		t.Errorf("Expected 1 pod info object, got %d", len(sgi.PodInfos))
+	}
+	if sgi.PodInfos[podInfo.UID] != podInfo {
+		t.Error("Pod info not properly stored in map")
+	}
+
+	podInfo2 := &pod_info.PodInfo{
+		UID:    "pod-2",
+		Status: pod_status.Pending,
+	}
+	sgi.assignTask(podInfo2)
+	if len(sgi.PodInfos) != 2 {
+		t.Errorf("Expected 2 pod info objects, got %d", len(sgi.PodInfos))
+	}
+
+	if sgi.PodInfos[podInfo2.UID] != podInfo2 {
+		t.Error("Pod info 2 not properly stored in map")
+	}
+}