TAS: assign pods based on ranks for Job

pkg/controller/tas/topology_ungater.go

@@ -20,9 +20,12 @@ import (
 	"context"
 	"errors"
 	"slices"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/go-logr/logr"
+	batchv1 "k8s.io/api/batch/v1"
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/util/workqueue"
@@ -73,6 +76,11 @@ type podWithDomain struct {
 	domainID utiltas.TopologyDomainID
 }
 
+type domainWithCount struct {
+	domainID utiltas.TopologyDomainID
+	count    int
+}
+
 var _ reconcile.Reconciler = (*topologyUngater)(nil)
 var _ predicate.Predicate = (*topologyUngater)(nil)
 
@@ -293,6 +301,56 @@ func podsToUngateInfo(
 }
 
 func assignGatedPodsToDomains(
+	log logr.Logger,
+	psa *kueue.PodSetAssignment,
+	pods []*corev1.Pod) []podWithDomain {
+	if rankToGatedPod, ok := readRanksIfAvailable(log, psa, pods); ok {
+		return assignGatedPodsToDomainsByRanks(psa, rankToGatedPod)
+	}
+	return assignGatedPodsToDomainsGreedy(log, psa, pods)
+}
+
+func assignGatedPodsToDomainsByRanks(
+	psa *kueue.PodSetAssignment,
+	rankToGatedPod map[int]*corev1.Pod) []podWithDomain {
+	toUngate := make([]podWithDomain, 0)
+	sortedDomains := sortDomains(psa)
+	totalCount := 0
+	for i := range sortedDomains {
+		totalCount += sortedDomains[i].count
+	}
+	rankToDomainID := make([]utiltas.TopologyDomainID, totalCount)
+	index := 0
+	for _, domain := range sortedDomains {
+		for s := range domain.count {
+			rankToDomainID[index+s] = domain.domainID
+		}
+		index += domain.count
+	}
+	for rank, pod := range rankToGatedPod {
+		toUngate = append(toUngate, podWithDomain{
+			pod:      pod,
+			domainID: rankToDomainID[rank],
+		})
+	}
+	return toUngate
+}
+
+func sortDomains(psa *kueue.PodSetAssignment) []domainWithCount {
+	sortableDomains := make([]domainWithCount, len(psa.TopologyAssignment.Domains))
+	for i, domain := range psa.TopologyAssignment.Domains {
+		sortableDomains[i] = domainWithCount{
+			domainID: utiltas.DomainID(domain.Values),
+			count:    int(domain.Count),
+		}
+	}
+	slices.SortFunc(sortableDomains, func(a, b domainWithCount) int {
+		return strings.Compare(string(a.domainID), string(b.domainID))
+	})
+	return sortableDomains
+}
+
+func assignGatedPodsToDomainsGreedy(
 	log logr.Logger,
 	psa *kueue.PodSetAssignment,
 	pods []*corev1.Pod) []podWithDomain {
@@ -335,6 +393,44 @@ func assignGatedPodsToDomains(
 	return toUngate
 }
 
+func readRanksIfAvailable(log logr.Logger,
+	psa *kueue.PodSetAssignment,
+	pods []*corev1.Pod) (map[int]*corev1.Pod, bool) {
+	result := make(map[int]*corev1.Pod, 0)
+	count := int(*psa.Count)
+	for _, pod := range pods {
+		rank := readIntFromLabel(log, pod, batchv1.JobCompletionIndexAnnotation)
+		if rank == nil {
+			// the Pod has no rank information - ranks cannot be used
+			return nil, false
+		}
+		if _, found := result[*rank]; found {
+			// there is a conflict in ranks, they cannot be used
+			return nil, false
+		}
+		if *rank >= count {
+			// the rank exceeds parallelism, this scenario is not supported by
+			// the rank-based ordering of pods.
+			return nil, false
+		}
+		result[*rank] = pod
+	}
+	return result, true
+}
+
+func readIntFromLabel(log logr.Logger, pod *corev1.Pod, labelKey string) *int {
+	v, found := pod.Labels[labelKey]
+	if !found {
+		return nil
+	}
+	i, err := strconv.Atoi(v)
+	if err != nil {
+		log.Error(err, "failed to parse index annotation", "value", v)
+		return nil
+	}
+	return ptr.To(i)
+}
+
 func isAdmittedByTAS(w *kueue.Workload) bool {
 	return w.Status.Admission != nil && workload.IsAdmitted(w) &&
 		slices.ContainsFunc(w.Status.Admission.PodSetAssignments,