fix(preempt): keep nominated gang tasks eligible and retry stranded gang placements

pkg/scheduler/actions/preempt/preempt.go

@@ -28,6 +28,7 @@ import (
 	"fmt"
 	"math"
 	"math/rand"
+	"slices"
 	"sort"
 	"sync"
 	"sync/atomic"
@@ -36,6 +37,7 @@ import (
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	utilerrors "k8s.io/apimachinery/pkg/util/errors"
+	"k8s.io/apimachinery/pkg/util/sets"
 	"k8s.io/client-go/util/workqueue"
 	"k8s.io/klog/v2"
 	k8sutil "k8s.io/kubernetes/pkg/scheduler/util"
@@ -63,6 +65,11 @@ const (
 	MinCandidateNodesPercentageKey = "minCandidateNodesPercentage"
 	MinCandidateNodesAbsoluteKey   = "minCandidateNodesAbsolute"
 	MaxCandidateNodesAbsoluteKey   = "maxCandidateNodesAbsolute"
+
+	// GangPlacementRetriesKey bounds how many times a starving job's preemption
+	// transaction is retried with the previously chosen nodes excluded when the
+	// job could not be pipelined as a whole. 0 disables retries.
+	GangPlacementRetriesKey = "gangPlacementRetries"
 )
 
 type Action struct {
@@ -78,6 +85,7 @@ type Action struct {
 	minCandidateNodesPercentage   int
 	minCandidateNodesAbsolute     int
 	maxCandidateNodesAbsolute     int
+	gangPlacementRetries          int
 }
 
 func New() *Action {
@@ -89,6 +97,7 @@ func New() *Action {
 		minCandidateNodesPercentage:      10,
 		minCandidateNodesAbsolute:        1,
 		maxCandidateNodesAbsolute:        100,
+		gangPlacementRetries:             2,
 	}
 }
 
@@ -107,6 +116,7 @@ func (pmpt *Action) parseArguments(ssn *framework.Session) {
 	arguments.GetInt(&pmpt.minCandidateNodesPercentage, MinCandidateNodesPercentageKey)
 	arguments.GetInt(&pmpt.minCandidateNodesAbsolute, MinCandidateNodesAbsoluteKey)
 	arguments.GetInt(&pmpt.maxCandidateNodesAbsolute, MaxCandidateNodesAbsoluteKey)
+	arguments.GetInt(&pmpt.gangPlacementRetries, GangPlacementRetriesKey)
 	pmpt.ssn = ssn
 }
 
@@ -157,16 +167,13 @@ func (pmpt *Action) Execute(ssn *framework.Session) {
 		}
 		preemptorsMap[job.Queue].Push(job)
 		underRequest = append(underRequest, job)
-		preemptorTasks[job.UID] = util.NewPriorityQueue(ssn.TaskOrderFn)
-		for _, task := range job.TaskStatusIndex[api.Pending] {
-			if task.SchGated {
-				continue
-			}
-			preemptorTasks[job.UID].Push(task)
-		}
+		preemptorTasks[job.UID] = pendingPreemptorTasks(ssn, job)
 	}
 
-	ph := util.NewPredicateHelper()
+	// The predicate error cache is keyed by task role within a job, so helpers
+	// are held per job: sharing one across jobs buys nothing, and a job's
+	// helper is replaced when a placement attempt is abandoned.
+	predicateHelpers := map[api.JobID]util.PredicateHelper{}
 	// Preemption between Jobs within Queue.
 	for _, queue := range queues {
 		for {
@@ -180,53 +187,9 @@ func (pmpt *Action) Execute(ssn *framework.Session) {
 
 			preemptorJob := preemptors.Pop().(*api.JobInfo)
 
-			stmt := framework.NewStatement(ssn)
-			var assigned bool
-			var err error
-			for {
-				// If job is not request more resource, then stop preempting.
-				if !ssn.JobStarving(preemptorJob) {
-					break
-				}
-
-				// If not preemptor tasks, next job.
-				if preemptorTasks[preemptorJob.UID].Empty() {
-					klog.V(3).Infof("No preemptor task in job <%s/%s>.",
-						preemptorJob.Namespace, preemptorJob.Name)
-					break
-				}
-
-				preemptor := preemptorTasks[preemptorJob.UID].Pop().(*api.TaskInfo)
-
-				assigned, err = pmpt.preempt(ssn, stmt, preemptor, func(task *api.TaskInfo) bool {
-					// Ignore non running task.
-					if !api.PreemptableStatus(task.Status) {
-						return false
-					}
-					// BestEffort pod is not supported to preempt unBestEffort pod.
-					if preemptor.BestEffort && !task.BestEffort {
-						return false
-					}
-					if !task.Preemptable {
-						return false
-					}
-					job, found := ssn.Jobs[task.Job]
-					if !found {
-						return false
-					}
-					// Preempt other jobs within queue
-					return job.Queue == preemptorJob.Queue && preemptor.Job != task.Job
-				}, ph)
-				if err != nil {
-					klog.V(3).Infof("Preemptor <%s/%s> failed to preempt Task , err: %s", preemptor.Namespace, preemptor.Name, err)
-				}
-			}
-
 			// Commit changes only if job is pipelined, otherwise try next job.
-			if ssn.JobPipelined(preemptorJob) {
-				stmt.Commit()
-			} else {
-				stmt.Discard()
+			assigned, committed := pmpt.preemptForJob(ssn, preemptorJob, preemptorTasks, predicateHelpers)
+			if !committed {
 				continue
 			}
 
@@ -238,14 +201,7 @@ func (pmpt *Action) Execute(ssn *framework.Session) {
 		// Preemption between Task within Job.
 		for _, job := range underRequest {
 			// Fix: preemptor numbers lose when in same job
-			preemptorTasks[job.UID] = util.NewPriorityQueue(ssn.TaskOrderFn)
-			for _, task := range job.TaskStatusIndex[api.Pending] {
-				// Again, skip scheduling gated tasks
-				if task.SchGated {
-					continue
-				}
-				preemptorTasks[job.UID].Push(task)
-			}
+			preemptorTasks[job.UID] = pendingPreemptorTasks(ssn, job)
 			for {
 				if _, found := preemptorTasks[job.UID]; !found {
 					break
@@ -274,7 +230,7 @@ func (pmpt *Action) Execute(ssn *framework.Session) {
 
 					// Preempt tasks within job.
 					return preemptor.Job == task.Job
-				}, ph)
+				}, jobPredicateHelper(predicateHelpers, job.UID), nil)
 				if err != nil {
 					klog.V(3).Infof("Preemptor <%s/%s> failed to preempt Task , err: %s", preemptor.Namespace, preemptor.Name, err)
 				}
@@ -291,24 +247,186 @@ func (pmpt *Action) Execute(ssn *framework.Session) {
 
 func (pmpt *Action) UnInitialize() {}
 
+// pendingPreemptorTasks queues the job's pending, non-gated tasks in task order.
+func pendingPreemptorTasks(ssn *framework.Session, job *api.JobInfo) *util.PriorityQueue {
+	tasks := util.NewPriorityQueue(ssn.TaskOrderFn)
+	for _, task := range job.TaskStatusIndex[api.Pending] {
+		if task.SchGated {
+			continue
+		}
+		tasks.Push(task)
+	}
+	return tasks
+}
+
+// jobPredicateHelper returns the job's predicate helper, creating it on first use.
+func jobPredicateHelper(helpers map[api.JobID]util.PredicateHelper, jobID api.JobID) util.PredicateHelper {
+	if helper, found := helpers[jobID]; found {
+		return helper
+	}
+	helper := util.NewPredicateHelper()
+	helpers[jobID] = helper
+	return helper
+}
+
+// pipelinedNodes returns the nodes hosting the job's pipelined tasks, skipping
+// tasks listed in ignore (those pipelined by an earlier, committed statement).
+func pipelinedNodes(job *api.JobInfo, ignore sets.Set[api.TaskID]) sets.Set[string] {
+	nodes := sets.New[string]()
+	for uid, task := range job.TaskStatusIndex[api.Pipelined] {
+		if !ignore.Has(uid) {
+			nodes.Insert(task.NodeName)
+		}
+	}
+	return nodes
+}
+
+// preemptForJob runs preemption transactions for a starving job until one
+// leaves the job pipelined and is committed, or the retry budget is spent.
+//
+// Tasks are placed one at a time, so the node chosen for an early task can
+// leave later tasks with no candidate at all — e.g. a gang bound together by
+// a topology pod affinity whose first task lands in a domain that has no
+// other preemptable node. A single pass would discard the transaction every
+// session and the job would starve indefinitely while a viable placement
+// existed elsewhere. Each retry excludes the nodes the previous attempt
+// pipelined onto, steering the whole gang into a different region of the
+// cluster.
+//
+// Each retry replaces the job's predicate helper: its error cache is keyed
+// by task role, so failures recorded for one sibling under the abandoned
+// placement would otherwise be replayed against the others — in this attempt
+// and in every later use of the helper for this job.
+//
+// Returns whether the last task attempt was assigned (the caller re-queues
+// the job to keep preempting for its remaining tasks) and whether the
+// transaction was committed.
+func (pmpt *Action) preemptForJob(
+	ssn *framework.Session,
+	job *api.JobInfo,
+	preemptorTasks map[api.JobID]*util.PriorityQueue,
+	predicateHelpers map[api.JobID]util.PredicateHelper,
+) (assigned, committed bool) {
+	committedTasks := sets.KeySet(job.TaskStatusIndex[api.Pipelined])
+	excludedNodes := sets.New[string]()
+	for attempt := 0; ; attempt++ {
+		stmt := framework.NewStatement(ssn)
+		assigned = pmpt.preemptJobTasks(ssn, stmt, job, preemptorTasks[job.UID], jobPredicateHelper(predicateHelpers, job.UID), excludedNodes)
+		if ssn.JobPipelined(job) {
+			stmt.Commit()
+			return assigned, true
+		}
+
+		chosen := pipelinedNodes(job, committedTasks)
+		stmt.Discard()
+
+		if attempt >= pmpt.gangPlacementRetries || chosen.Len() == 0 {
+			if attempt > 0 {
+				clearLastTxContexts(job)
+			}
+			return false, false
+		}
+
+		klog.V(3).Infof("Job <%s/%s> not pipelined after placing tasks on %v, retrying preemption without those nodes",
+			job.Namespace, job.Name, sets.List(chosen))
+		excludedNodes = excludedNodes.Union(chosen)
+		predicateHelpers[job.UID] = util.NewPredicateHelper()
+		preemptorTasks[job.UID] = pendingPreemptorTasks(ssn, job)
+	}
+}
+
+// clearLastTxContexts drops the transaction context Discard left on the job's
+// pending tasks. After a retried placement fails for good, that context is
+// the last attempt's, made with the better nodes excluded; publishing it as
+// the tasks' scheduling reason and nominated node would steer allocate's
+// next session toward the worst placement tried.
+func clearLastTxContexts(job *api.JobInfo) {
+	for _, task := range job.TaskStatusIndex[api.Pending] {
+		task.ClearLastTxContext()
+	}
+}
+
+// preemptJobTasks preempts for the job's queued tasks in order within a single
+// statement, stopping once the job is no longer starving or no tasks remain.
+// Returns whether the last attempted task was assigned.
+func (pmpt *Action) preemptJobTasks(
+	ssn *framework.Session,
+	stmt *framework.Statement,
+	job *api.JobInfo,
+	tasks *util.PriorityQueue,
+	predicateHelper util.PredicateHelper,
+	excludedNodes sets.Set[string],
+) bool {
+	var assigned bool
+	for {
+		// If job is not request more resource, then stop preempting.
+		if !ssn.JobStarving(job) {
+			break
+		}
+
+		// If not preemptor tasks, next job.
+		if tasks.Empty() {
+			klog.V(3).Infof("No preemptor task in job <%s/%s>.", job.Namespace, job.Name)
+			break
+		}
+
+		preemptor := tasks.Pop().(*api.TaskInfo)
+
+		var err error
+		assigned, err = pmpt.preempt(ssn, stmt, preemptor, func(task *api.TaskInfo) bool {
+			// Ignore non running task.
+			if !api.PreemptableStatus(task.Status) {
+				return false
+			}
+			// BestEffort pod is not supported to preempt unBestEffort pod.
+			if preemptor.BestEffort && !task.BestEffort {
+				return false
+			}
+			if !task.Preemptable {
+				return false
+			}
+			victimJob, found := ssn.Jobs[task.Job]
+			if !found {
+				return false
+			}
+			// Preempt other jobs within queue
+			return victimJob.Queue == job.Queue && preemptor.Job != task.Job
+		}, predicateHelper, excludedNodes)
+		if err != nil {
+			klog.V(3).Infof("Preemptor <%s/%s> failed to preempt Task , err: %s", preemptor.Namespace, preemptor.Name, err)
+		}
+	}
+	return assigned
+}
+
 func (pmpt *Action) preempt(
 	ssn *framework.Session,
 	stmt *framework.Statement,
 	preemptor *api.TaskInfo,
 	filter func(*api.TaskInfo) bool,
 	predicateHelper util.PredicateHelper,
+	excludedNodes sets.Set[string],
 ) (bool, error) {
+	// Eligibility predicates the nominated node, which needs the task's cycle state.
+	if err := ssn.PrePredicateFn(preemptor); err != nil {
+		return false, fmt.Errorf("PrePredicate for task %s/%s failed for: %v", preemptor.Namespace, preemptor.Name, err)
+	}
+
 	// Check whether the task is eligible to preempt others, e.g., check preemptionPolicy is `Never` or not
 	if err := pmpt.taskEligibleToPreempt(preemptor); err != nil {
 		return false, err
 	}
 
-	if err := ssn.PrePredicateFn(preemptor); err != nil {
-		return false, fmt.Errorf("PrePredicate for task %s/%s failed for: %v", preemptor.Namespace, preemptor.Name, err)
-	}
-
 	// we should filter out those nodes that are UnschedulableAndUnresolvable status got in allocate action
 	allNodes := ssn.FilterOutUnschedulableAndUnresolvableNodesForTask(preemptor)
+	if excludedNodes.Len() > 0 {
+		// Drop the nodes abandoned gang placement attempts chose before
+		// PredicateNodes samples its candidates, so the exclusion cannot
+		// shrink an already truncated sample.
+		allNodes = slices.DeleteFunc(slices.Clone(allNodes), func(n *api.NodeInfo) bool {
+			return excludedNodes.Has(n.Name)
+		})
+	}
 	predicateNodes, fitErrors := predicateHelper.PredicateNodes(preemptor, allNodes, ssn.PredicateForPreemptAction, pmpt.enablePredicateErrorCache, ssn.NodesInShard)
 
 	// When predicate filtering returns no candidates the cluster is effectively
@@ -452,20 +570,23 @@ func (pmpt *Action) taskEligibleToPreempt(preemptor *api.TaskInfo) error {
 			return fmt.Errorf("not eligible due to the pod's nominated node is not found in the session")
 		}
 
-		err := pmpt.ssn.PredicateFn(preemptor, nodeInfo)
-		if err == nil {
-			return fmt.Errorf("not eligible due to the pod's nominated node is already schedulable, which should not happen as preemption means no node is schedulable")
-		}
-
-		fitError, ok := err.(*api.FitError)
-		if !ok {
-			return fmt.Errorf("not eligible due to the predicate returned a non-FitError error, the error is: %v", err)
-		}
+		// Diverges from upstream, which rejects the task when its nominated
+		// node passes predicates. Predicates ignore resource fit, so a passing
+		// nominated node does not mean the task can run there; that is
+		// allocate's call. Preempt stays eligible so the task can join its
+		// gang siblings' transaction (pipelineOnFittingNode keeps it from
+		// evicting anyone if the node really does fit).
+		if err := pmpt.ssn.PredicateFn(preemptor, nodeInfo); err != nil {
+			fitError, ok := err.(*api.FitError)
+			if !ok {
+				return fmt.Errorf("not eligible due to the predicate returned a non-FitError error, the error is: %v", err)
+			}
 
-		// If the pod's nominated node is considered as UnschedulableAndUnresolvable by the predicate,
-		// then the pod should be considered for preempting again.
-		if fitError.Status.ContainsUnschedulableAndUnresolvable() {
-			return nil
+			// If the pod's nominated node is considered as UnschedulableAndUnresolvable by the predicate,
+			// then the pod should be considered for preempting again.
+			if fitError.Status.ContainsUnschedulableAndUnresolvable() {
+				return nil
+			}
 		}
 
 		preemptorPodPriority := PodPriority(preemptor.Pod)