Changed SubSetNodes signature to use SubGroupInfo instead of SubGroupSet (#561)

Author: romanbaron

pkg/scheduler/actions/common/allocate.go

@@ -27,7 +27,8 @@ func AllocateJob(ssn *framework.Session, stmt *framework.Statement, nodes []*nod
 		return false
 	}
 
-	nodeSets, err := ssn.SubsetNodesFn(job, job.RootSubGroupSet, tasksToAllocate, nodes)
+	podSets := job.RootSubGroupSet.GetAllPodSets()
+	nodeSets, err := ssn.SubsetNodesFn(job, &job.RootSubGroupSet.SubGroupInfo, podSets, tasksToAllocate, nodes)
 	if err != nil {
 		log.InfraLogger.Errorf(
 			"Failed to run SubsetNodes on job <%s/%s>: %v", job.Namespace, job.Namespace, err)

pkg/scheduler/api/podgroup_info/job_info.go

@@ -116,7 +116,7 @@ func NewPodGroupInfo(uid common_info.PodGroupID, tasks ...*pod_info.PodInfo) *Po
 			Stale:     false,
 		},
 		RootSubGroupSet: defaultSubGroupSet,
-		PodSets:         defaultSubGroupSet.GetPodSets(),
+		PodSets:         defaultSubGroupSet.GetAllPodSets(),
 
 		LastStartTimestamp:   nil,
 		activeAllocatedCount: ptr.To(0),
@@ -199,7 +199,7 @@ func (pgi *PodGroupInfo) setSubGroups(podGroup *enginev2alpha2.PodGroup) error {
 		return err
 	}
 	pgi.RootSubGroupSet = rootSubGroupSet
-	podSets := rootSubGroupSet.GetPodSets()
+	podSets := rootSubGroupSet.GetAllPodSets()
 	if len(podSets) > 0 {
 		pgi.PodSets = podSets
 	} else {
@@ -480,7 +480,7 @@ func (pgi *PodGroupInfo) CloneWithTasks(tasks []*pod_info.PodInfo) *PodGroupInfo
 	pgi.CreationTimestamp.DeepCopyInto(&info.CreationTimestamp)
 
 	info.RootSubGroupSet = pgi.RootSubGroupSet.Clone()
-	info.PodSets = info.RootSubGroupSet.GetPodSets()
+	info.PodSets = info.RootSubGroupSet.GetAllPodSets()
 
 	for _, task := range tasks {
 		info.AddTaskInfo(task.Clone())

pkg/scheduler/api/podgroup_info/subgroup_info/subgroupset.go

@@ -63,13 +63,13 @@ func (sgs *SubGroupSet) Clone() *SubGroupSet {
 	return root
 }
 
-func (sgs *SubGroupSet) GetPodSets() map[string]*PodSet {
+func (sgs *SubGroupSet) GetAllPodSets() map[string]*PodSet {
 	result := make(map[string]*PodSet)
 	for _, podSet := range sgs.podSets {
 		result[podSet.GetName()] = podSet
 	}
 	for _, subGroup := range sgs.groups {
-		podSets := subGroup.GetPodSets()
+		podSets := subGroup.GetAllPodSets()
 		for name, podSet := range podSets {
 			result[name] = podSet
 		}

pkg/scheduler/api/podgroup_info/subgroup_info/subgroupset_test.go

@@ -175,7 +175,7 @@ func TestGetPodSets(t *testing.T) {
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			sg, want := tt.build()
-			got := sg.GetPodSets()
+			got := sg.GetAllPodSets()
 			if len(got) != len(want) {
 				t.Errorf("expected %d podsets, got %d", len(want), len(got))
 			}

pkg/scheduler/api/types.go

@@ -30,7 +30,8 @@ import (
 )
 
 // SubsetNodesFn is used to divide the nodes into sets
-type SubsetNodesFn func(podGroup *podgroup_info.PodGroupInfo, subGroupSet *subgroup_info.SubGroupSet, tasks []*pod_info.PodInfo, nodeSet node_info.NodeSet) ([]node_info.NodeSet, error)
+type SubsetNodesFn func(podGroup *podgroup_info.PodGroupInfo, subGroup *subgroup_info.SubGroupInfo,
+	podSets map[string]*subgroup_info.PodSet, tasks []*pod_info.PodInfo, nodeSet node_info.NodeSet) ([]node_info.NodeSet, error)
 
 // PredicateFn is used to predicate node for task.
 type PredicateFn func(*pod_info.PodInfo, *podgroup_info.PodGroupInfo, *node_info.NodeInfo) error

pkg/scheduler/framework/session_plugins.go

@@ -264,7 +264,9 @@ func (ssn *Session) SubGroupOrderFn(l, r interface{}) bool {
 	return lSubGroup.GetName() < rSubGroup.GetName()
 }
 
-func (ssn *Session) QueueOrderFn(lQ, rQ *queue_info.QueueInfo, lJob, rJob *podgroup_info.PodGroupInfo, lVictims, rVictims []*podgroup_info.PodGroupInfo) bool {
+func (ssn *Session) QueueOrderFn(lQ, rQ *queue_info.QueueInfo, lJob, rJob *podgroup_info.PodGroupInfo,
+	lVictims, rVictims []*podgroup_info.PodGroupInfo,
+) bool {
 	for _, qof := range ssn.QueueOrderFns {
 		if j := qof(lQ, rQ, lJob, rJob, lVictims, rVictims); j != 0 {
 			return j < 0
@@ -322,14 +324,17 @@ func (ssn *Session) IsTaskAllocationOnNodeOverCapacityFn(task *pod_info.PodInfo,
 	}
 }
 
-func (ssn *Session) SubsetNodesFn(podGroup *podgroup_info.PodGroupInfo, subGroupSet *subgroup_info.SubGroupSet, tasks []*pod_info.PodInfo, initNodeSet node_info.NodeSet) ([]node_info.NodeSet, error) {
+func (ssn *Session) SubsetNodesFn(
+	podGroup *podgroup_info.PodGroupInfo, subGroupInfo *subgroup_info.SubGroupInfo,
+	podSets map[string]*subgroup_info.PodSet, tasks []*pod_info.PodInfo, initNodeSet node_info.NodeSet,
+) ([]node_info.NodeSet, error) {
 	nodeSets := []node_info.NodeSet{initNodeSet}
 	for _, subsetNodesFn := range ssn.SubsetNodesFns {
 		log.InfraLogger.V(7).Infof(
 			"Running plugin func <%v> on podGroup <%s/%s>", subsetNodesFn, podGroup.Namespace, podGroup.Namespace)
 		var newNodeSets []node_info.NodeSet
 		for _, nodeSet := range nodeSets {
-			nodeSubsets, err := subsetNodesFn(podGroup, subGroupSet, tasks, nodeSet)
+			nodeSubsets, err := subsetNodesFn(podGroup, subGroupInfo, podSets, tasks, nodeSet)
 			if err != nil {
 				return nil, err
 			}

pkg/scheduler/framework/session_plugins_test.go

@@ -113,7 +113,7 @@ func TestPartitionMultiImplementation(t *testing.T) {
 		},
 	}
 
-	shardClusterSubseting := func(_ *podgroup_info.PodGroupInfo, _ *subgroup_info.SubGroupSet, _ []*pod_info.PodInfo, nodeset node_info.NodeSet) ([]node_info.NodeSet, error) {
+	shardClusterSubseting := func(_ *podgroup_info.PodGroupInfo, _ *subgroup_info.SubGroupInfo, _ map[string]*subgroup_info.PodSet, _ []*pod_info.PodInfo, nodeset node_info.NodeSet) ([]node_info.NodeSet, error) {
 		var subset1 []*node_info.NodeInfo
 		var subset2 []*node_info.NodeInfo
 		for _, node := range nodeset {
@@ -126,7 +126,7 @@ func TestPartitionMultiImplementation(t *testing.T) {
 		return []node_info.NodeSet{subset1, subset2}, nil
 	}
 
-	topologySubseting := func(_ *podgroup_info.PodGroupInfo, _ *subgroup_info.SubGroupSet, _ []*pod_info.PodInfo, nodeset node_info.NodeSet) ([]node_info.NodeSet, error) {
+	topologySubseting := func(_ *podgroup_info.PodGroupInfo, _ *subgroup_info.SubGroupInfo, _ map[string]*subgroup_info.PodSet, _ []*pod_info.PodInfo, nodeset node_info.NodeSet) ([]node_info.NodeSet, error) {
 		var subset1 []*node_info.NodeInfo
 		var subset2 []*node_info.NodeInfo
 		for _, node := range nodeset {
@@ -144,7 +144,7 @@ func TestPartitionMultiImplementation(t *testing.T) {
 	ssn.AddSubsetNodesFn(shardClusterSubseting)
 	ssn.AddSubsetNodesFn(topologySubseting)
 
-	partitions, _ := ssn.SubsetNodesFn(podgroup_info.NewPodGroupInfo("a"), nil, nil, nodes)
+	partitions, _ := ssn.SubsetNodesFn(podgroup_info.NewPodGroupInfo("a"), nil, nil, nil, nodes)
 
 	assert.Equal(t, 4, len(partitions))
 

pkg/scheduler/plugins/topology/job_filtering.go

@@ -21,14 +21,15 @@ type jobAllocationMetaData struct {
 }
 
 func (t *topologyPlugin) subSetNodesFn(
-	job *podgroup_info.PodGroupInfo, subGroupSet *subgroup_info.SubGroupSet, tasks []*pod_info.PodInfo, nodeSet node_info.NodeSet,
+	job *podgroup_info.PodGroupInfo, subGroup *subgroup_info.SubGroupInfo, podSets map[string]*subgroup_info.PodSet,
+	tasks []*pod_info.PodInfo, nodeSet node_info.NodeSet,
 ) ([]node_info.NodeSet, error) {
-	topologyTree, found := t.getJobTopology(subGroupSet)
+	topologyTree, found := t.getJobTopology(subGroup)
 	if !found {
 		job.SetJobFitError(
 			podgroup_info.PodSchedulingErrors,
 			fmt.Sprintf("Matching topology %s does not exist",
-				subGroupSet.GetTopologyConstraint().Topology),
+				subGroup.GetTopologyConstraint().Topology),
 			nil)
 		return []node_info.NodeSet{}, nil
 	}
@@ -50,7 +51,7 @@ func (t *topologyPlugin) subSetNodesFn(
 		return []node_info.NodeSet{}, nil
 	}
 
-	jobAllocatableDomains, err := t.getJobAllocatableDomains(job, subGroupSet, len(tasks), topologyTree)
+	jobAllocatableDomains, err := t.getJobAllocatableDomains(job, subGroup, podSets, len(tasks), topologyTree)
 	if err != nil {
 		return nil, err
 	}
@@ -67,11 +68,11 @@ func (t *topologyPlugin) subSetNodesFn(
 	return domainNodeSets, nil
 }
 
-func (t *topologyPlugin) getJobTopology(subGroupSet *subgroup_info.SubGroupSet) (*Info, bool) {
-	if subGroupSet.GetTopologyConstraint() == nil {
+func (t *topologyPlugin) getJobTopology(subGroup *subgroup_info.SubGroupInfo) (*Info, bool) {
+	if subGroup.GetTopologyConstraint() == nil {
 		return nil, true
 	}
-	jobTopologyName := subGroupSet.GetTopologyConstraint().Topology
+	jobTopologyName := subGroup.GetTopologyConstraint().Topology
 	if jobTopologyName == "" {
 		return nil, true
 	}
@@ -188,18 +189,19 @@ func calcNextAllocationTestPodResources(previousTestResources, maxPodResources *
 }
 
 func (t *topologyPlugin) getJobAllocatableDomains(
-	job *podgroup_info.PodGroupInfo, subGroupSet *subgroup_info.SubGroupSet, taskToAllocateCount int, topologyTree *Info,
+	job *podgroup_info.PodGroupInfo, subGroup *subgroup_info.SubGroupInfo, podSets map[string]*subgroup_info.PodSet,
+	taskToAllocateCount int, topologyTree *Info,
 ) ([]*DomainInfo, error) {
-	relevantLevels, err := t.calculateRelevantDomainLevels(subGroupSet, topologyTree)
+	relevantLevels, err := t.calculateRelevantDomainLevels(subGroup, topologyTree)
 	if err != nil {
 		return nil, err
 	}
 
 	// Validate that the domains do not clash with the chosen domain for active pods of the job
 	var relevantDomainsByLevel domainsByLevel
-	if hasActiveAllocatedTasks(subGroupSet) && hasTopologyRequiredConstraint(subGroupSet) {
-		relevantDomainsByLevel = getRelevantDomainsWithAllocatedPods(subGroupSet, topologyTree,
-			DomainLevel(subGroupSet.GetTopologyConstraint().RequiredLevel))
+	if hasActiveAllocatedTasks(podSets) && hasTopologyRequiredConstraint(subGroup) {
+		relevantDomainsByLevel = getRelevantDomainsWithAllocatedPods(podSets, topologyTree,
+			DomainLevel(subGroup.GetTopologyConstraint().RequiredLevel))
 	} else {
 		relevantDomainsByLevel = topologyTree.DomainsByLevel
 	}
@@ -223,8 +225,8 @@ func (t *topologyPlugin) getJobAllocatableDomains(
 	return domains, nil
 }
 
-func hasActiveAllocatedTasks(subGroupSet *subgroup_info.SubGroupSet) bool {
-	for _, podSet := range subGroupSet.GetPodSets() {
+func hasActiveAllocatedTasks(podSets map[string]*subgroup_info.PodSet) bool {
+	for _, podSet := range podSets {
 		if podSet.GetNumActiveAllocatedTasks() > 0 {
 			return true
 		}
@@ -233,20 +235,20 @@ func hasActiveAllocatedTasks(subGroupSet *subgroup_info.SubGroupSet) bool {
 }
 
 func getRelevantDomainsWithAllocatedPods(
-	subGroupSet *subgroup_info.SubGroupSet, topologyTree *Info, requiredLevel DomainLevel,
+	podSets map[string]*subgroup_info.PodSet, topologyTree *Info, requiredLevel DomainLevel,
 ) domainsByLevel {
 	relevantDomainsByLevel := domainsByLevel{}
 	for _, domainAtRequiredLevel := range topologyTree.DomainsByLevel[requiredLevel] {
-		if !hasActiveJobPodInDomain(subGroupSet, domainAtRequiredLevel) {
+		if !hasActiveJobPodInDomain(podSets, domainAtRequiredLevel) {
 			continue // if the domain at the top level does not have any active pods, then any domains under the subtree cannot satisfy the required constraint for both active and pending pods
 		}
 		addSubTreeToDomainMap(domainAtRequiredLevel, relevantDomainsByLevel)
 	}
 	return relevantDomainsByLevel
 }
 
-func hasActiveJobPodInDomain(subGroupSet *subgroup_info.SubGroupSet, domain *DomainInfo) bool {
-	for _, podSet := range subGroupSet.GetPodSets() {
+func hasActiveJobPodInDomain(podSets map[string]*subgroup_info.PodSet, domain *DomainInfo) bool {
+	for _, podSet := range podSets {
 		for _, pod := range podSet.GetPodInfos() {
 			if pod_status.IsActiveAllocatedStatus(pod.Status) {
 				podInDomain := domain.Nodes[pod.NodeName] != nil
@@ -269,19 +271,19 @@ func addSubTreeToDomainMap(domain *DomainInfo, domainsMap domainsByLevel) {
 	domainsMap[domain.Level][domain.ID] = domain
 }
 
-func hasTopologyRequiredConstraint(subGroupSet *subgroup_info.SubGroupSet) bool {
-	return subGroupSet.GetTopologyConstraint().RequiredLevel != ""
+func hasTopologyRequiredConstraint(subGroup *subgroup_info.SubGroupInfo) bool {
+	return subGroup.GetTopologyConstraint().RequiredLevel != ""
 }
 
 func (*topologyPlugin) calculateRelevantDomainLevels(
-	subGroupSet *subgroup_info.SubGroupSet, topologyTree *Info,
+	subGroup *subgroup_info.SubGroupInfo, topologyTree *Info,
 ) ([]DomainLevel, error) {
-	topologyConstraint := subGroupSet.GetTopologyConstraint()
+	topologyConstraint := subGroup.GetTopologyConstraint()
 	requiredPlacement := DomainLevel(topologyConstraint.RequiredLevel)
 	preferredPlacement := DomainLevel(topologyConstraint.PreferredLevel)
 	if requiredPlacement == "" && preferredPlacement == "" {
 		return nil, fmt.Errorf("no topology constraints were found for subgroup %s, with topology name %s",
-			subGroupSet.GetName(), topologyTree.Name)
+			subGroup.GetName(), topologyTree.Name)
 	}
 
 	foundRequiredLevel := false

pkg/scheduler/plugins/topology/job_filtering_test.go

@@ -390,7 +390,9 @@ func TestTopologyPlugin_subsetNodesFn(t *testing.T) {
 			}
 
 			// Call the function under test
-			subsets, err := plugin.subSetNodesFn(job, job.RootSubGroupSet, podgroup_info.GetTasksToAllocate(job, nil, nil, true), maps.Values(nodesInfoMap))
+			subsets, err := plugin.subSetNodesFn(job, &job.RootSubGroupSet.SubGroupInfo,
+				job.RootSubGroupSet.GetAllPodSets(), podgroup_info.GetTasksToAllocate(job, nil, nil, true),
+				maps.Values(nodesInfoMap))
 
 			// Check error
 			if tt.expectedError != "" {
@@ -723,7 +725,7 @@ func TestTopologyPlugin_calculateRelevantDomainLevels(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			plugin := &topologyPlugin{}
 
-			result, err := plugin.calculateRelevantDomainLevels(tt.subGroupSet, tt.topologyTree)
+			result, err := plugin.calculateRelevantDomainLevels(&tt.subGroupSet.SubGroupInfo, tt.topologyTree)
 
 			// Check error
 			if tt.expectedError != "" {
@@ -1081,7 +1083,7 @@ func TestTopologyPlugin_calcTreeAllocatable(t *testing.T) {
 			},
 		},
 		{
-			name: "no leaf domains - no allocateable domains",
+			name: "no leaf domains - no allocatable domains",
 			job: &jobs_fake.TestJobBasic{
 				Name:                "test-job",
 				RequiredCPUsPerTask: 2000, // Too much for any node
@@ -1840,7 +1842,9 @@ func TestTopologyPlugin_getJobAllocatableDomains(t *testing.T) {
 			for _, podSet := range tt.job.PodSets {
 				tt.job.RootSubGroupSet.AddPodSet(podSet)
 			}
-			result, err := plugin.getJobAllocatableDomains(tt.job, tt.job.RootSubGroupSet, len(podgroup_info.GetTasksToAllocate(tt.job, nil, nil, true)), tt.topologyTree)
+			result, err := plugin.getJobAllocatableDomains(tt.job, &tt.job.RootSubGroupSet.SubGroupInfo,
+				tt.job.RootSubGroupSet.GetAllPodSets(), len(podgroup_info.GetTasksToAllocate(tt.job, nil, nil, true)),
+				tt.topologyTree)
 
 			// Check error
 			if tt.expectedError != "" {

pkg/scheduler/test_utils/test_utils.go

@@ -18,6 +18,8 @@ import (
 	"k8s.io/client-go/informers"
 	"k8s.io/client-go/kubernetes/fake"
 
+	kueuev1alpha1 "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
+
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/actions"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/common_info"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/api/pod_status"
@@ -31,7 +33,6 @@ import (
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils/dra_fake"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils/jobs_fake"
 	"github.com/NVIDIA/KAI-scheduler/pkg/scheduler/test_utils/nodes_fake"
-	kueuev1alpha1 "sigs.k8s.io/kueue/apis/kueue/v1alpha1"
 )
 
 var SchedulerVerbosity = flag.String("vv", "", "Scheduler's verbosity")