feat(scheduler): add gpujoborder plugin for configurable JobOrderFn tiebreak (#1995)

Signed-off-by: CoolingCube <CoolingCubeInfo@proton.me>

Author: CoolingCube

.changes/unreleased/added-20260729-142722.yaml

@@ -0,0 +1,3 @@
+kind: Added
+body: |-
+  Add gpujoborder plugin for configurable GPU-based JobOrderFn tiebreak on priority ties

pkg/scheduler/actions/utils/job_order_by_queue.go

@@ -255,7 +255,7 @@ func (jo *JobsOrderByQueues) createLeafNode(queue *queue_info.QueueInfo) *queueN
 		queue: queue,
 		children: scheduler_util.NewPriorityQueue(func(l, r interface{}) bool {
 			if jo.options.VictimQueue {
-				return !jo.ssn.JobOrderFn(l, r)
+				return jo.ssn.VictimOrderFn(l, r)
 			}
 			return jo.ssn.JobOrderFn(l, r)
 		}, jo.options.MaxJobsQueueDepth),

pkg/scheduler/api/podgroup_info/job_info.go

@@ -93,6 +93,7 @@ type PodGroupInfo struct {
 	tasksToAllocateInitResourceVector resource_info.ResourceVector
 	PodStatusIndex                    map[pod_status.PodStatus]pod_info.PodsMap
 	activeAllocatedCount              *int
+	aliveTasksRequestedGPUs           *float64
 }
 
 func NewPodGroupInfo(uid common_info.PodGroupID, tasks ...*pod_info.PodInfo) *PodGroupInfo {
@@ -365,6 +366,7 @@ func (pgi *PodGroupInfo) invalidateTasksCache() {
 	pgi.allPodsMap = nil
 	pgi.tasksToAllocate = nil
 	pgi.tasksToAllocateInitResourceVector = nil
+	pgi.aliveTasksRequestedGPUs = nil
 }
 
 func (pgi *PodGroupInfo) GetActiveAllocatedTasksCount() int {
@@ -456,14 +458,16 @@ func (pgi *PodGroupInfo) GetNumGatedTasks() int {
 }
 
 func (pgi *PodGroupInfo) GetAliveTasksRequestedGPUs() float64 {
-	tasksTotalRequestedGPUs := float64(0)
-	for _, task := range pgi.GetAllPodsMap() {
-		if pod_status.IsAliveStatus(task.Status) {
-			tasksTotalRequestedGPUs += task.ResReqVector.Get(resource_info.GPUIndex)
+	if pgi.aliveTasksRequestedGPUs == nil {
+		tasksTotalRequestedGPUs := float64(0)
+		for _, task := range pgi.GetAllPodsMap() {
+			if pod_status.IsAliveStatus(task.Status) {
+				tasksTotalRequestedGPUs += task.ResReqVector.Get(resource_info.GPUIndex)
+			}
 		}
+		pgi.aliveTasksRequestedGPUs = ptr.To(tasksTotalRequestedGPUs)
 	}
-
-	return tasksTotalRequestedGPUs
+	return *pgi.aliveTasksRequestedGPUs
 }
 
 func (pgi *PodGroupInfo) GetTasksActiveAllocatedReqResourceVector() resource_info.ResourceVector {

pkg/scheduler/framework/session.go

@@ -76,6 +76,7 @@ type Session struct {
 	NodePreOrderFns                       []api.NodePreOrderFn
 	NodeOrderFns                          []api.NodeOrderFn
 	JobOrderFns                           []common_info.CompareFn
+	VictimOrderFns                        []common_info.CompareFn
 	SubGroupOrderFns                      []common_info.CompareFn
 	TaskOrderFns                          []common_info.CompareFn
 	QueueOrderFns                         []api.CompareQueueFn
@@ -434,6 +435,7 @@ func (ssn *Session) clear() {
 	ssn.NodePreOrderFns = nil
 	ssn.NodeOrderFns = nil
 	ssn.JobOrderFns = nil
+	ssn.VictimOrderFns = nil
 	ssn.SubGroupOrderFns = nil
 	ssn.TaskOrderFns = nil
 	ssn.QueueOrderFns = nil

pkg/scheduler/framework/session_plugins.go

@@ -68,6 +68,22 @@ func (ssn *Session) AddJobOrderFn(jof common_info.CompareFn) {
 	ssn.JobOrderFns = append(ssn.JobOrderFns, jof)
 }
 
+// AddVictimOrderFn registers a comparator used only to rank candidate
+// eviction victims, without affecting pending-job allocation order.
+//
+// The registered JobOrderFns chain is always checked first (inverted,
+// matching existing allocation-order semantics); comparators registered
+// here apply only as a tiebreak, when every JobOrderFn treats l and r as
+// equal. This ensures job-level protections (e.g. elastic at-min/above-min
+// status) take precedence over any victim-specific comparator.
+//
+// Sign convention: vof(l, r) < 0 means l should be evicted before r. This
+// is the direct, non-inverted sense -- unlike the JobOrderFn chain, which
+// VictimOrderFn inverts internally when using it for the eviction path.
+func (ssn *Session) AddVictimOrderFn(vof common_info.CompareFn) {
+	ssn.VictimOrderFns = append(ssn.VictimOrderFns, vof)
+}
+
 func (ssn *Session) AddTaskOrderFn(tof common_info.CompareFn) {
 	ssn.TaskOrderFns = append(ssn.TaskOrderFns, tof)
 }
@@ -268,21 +284,42 @@ func (ssn *Session) QueueAllocatedResources(queue *queue_info.QueueInfo) *resour
 	return nil
 }
 
+func jobOrderCreationFallback(l, r interface{}) bool {
+	lv := l.(*podgroup_info.PodGroupInfo)
+	rv := r.(*podgroup_info.PodGroupInfo)
+	if lv.CreationTimestamp.Equal(&rv.CreationTimestamp) {
+		return lv.UID < rv.UID
+	}
+	return lv.CreationTimestamp.Before(&rv.CreationTimestamp)
+}
+
 func (ssn *Session) JobOrderFn(l, r interface{}) bool {
 	for _, jof := range ssn.JobOrderFns {
 		if j := jof(l, r); j != 0 {
 			return j < 0
 		}
 	}
+	return jobOrderCreationFallback(l, r)
+}
 
-	// If no job order funcs, order job by CreationTimestamp first, then by UID.
-	lv := l.(*podgroup_info.PodGroupInfo)
-	rv := r.(*podgroup_info.PodGroupInfo)
-	if lv.CreationTimestamp.Equal(&rv.CreationTimestamp) {
-		return lv.UID < rv.UID
-	} else {
-		return lv.CreationTimestamp.Before(&rv.CreationTimestamp)
+// VictimOrderFn reports whether l should be evicted before r.
+//
+// The JobOrderFns chain is consulted first, inverted; VictimOrderFns
+// registered via AddVictimOrderFn apply only when every JobOrderFn treats
+// l and r as equal. With no VictimOrderFns registered, this reduces to
+// the original !JobOrderFn(l, r) behavior.
+func (ssn *Session) VictimOrderFn(l, r interface{}) bool {
+	for _, jof := range ssn.JobOrderFns {
+		if j := jof(l, r); j != 0 {
+			return j > 0
+		}
+	}
+	for _, vof := range ssn.VictimOrderFns {
+		if v := vof(l, r); v != 0 {
+			return v < 0
+		}
 	}
+	return !jobOrderCreationFallback(l, r)
 }
 
 func (ssn *Session) TaskOrderFn(l, r interface{}) bool {

pkg/scheduler/framework/session_plugins_test.go

@@ -18,10 +18,12 @@ import (
 	kaiv1 "github.com/kai-scheduler/KAI-scheduler/pkg/apis/kai/v1"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/common/constants"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/node_info"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/podgroup_info"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/podgroup_info/subgroup_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/resource_info"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/conf"
 )
 
@@ -305,3 +307,43 @@ func scenarioSearchDurationForTest(value string) metav1.Duration {
 	}
 	return metav1.Duration{Duration: duration}
 }
+
+func TestVictimOrderFn_PreservesExistingJobOrderFnBehavior(t *testing.T) {
+	priorityLikeJobOrderFn := func(l, r interface{}) int {
+		lv := l.(*podgroup_info.PodGroupInfo)
+		rv := r.(*podgroup_info.PodGroupInfo)
+		if lv.Priority > rv.Priority {
+			return -1
+		}
+		if lv.Priority < rv.Priority {
+			return 1
+		}
+		return 0
+	}
+
+	vm := resource_info.NewResourceVectorMap()
+	priorities := []int32{1, 1, 5, 5, 10, 20, 20, 100}
+
+	for i, lp := range priorities {
+		for j, rp := range priorities {
+			if i == j {
+				continue
+			}
+			ssn := &Session{JobOrderFns: []common_info.CompareFn{priorityLikeJobOrderFn}}
+			lJob := podgroup_info.NewPodGroupInfoWithVectorMap(
+				common_info.PodGroupID("l"), vm)
+			lJob.Priority = lp
+			rJob := podgroup_info.NewPodGroupInfoWithVectorMap(
+				common_info.PodGroupID("r"), vm)
+			rJob.Priority = rp
+
+			oldBehavior := !ssn.JobOrderFn(lJob, rJob)
+			newBehavior := ssn.VictimOrderFn(lJob, rJob)
+
+			if oldBehavior != newBehavior {
+				t.Errorf("mismatch for priorities l=%d r=%d: old !JobOrderFn=%v, new VictimOrderFn=%v",
+					lp, rp, oldBehavior, newBehavior)
+			}
+		}
+	}
+}

pkg/scheduler/plugins/factory.go

@@ -23,6 +23,7 @@ import (
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/framework"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/plugins/dynamicresources"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/plugins/elastic"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/plugins/gpujoborder"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/plugins/gpupack"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/plugins/gpusharingorder"
 	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/plugins/gpuspread"
@@ -51,6 +52,7 @@ func InitDefaultPlugins() {
 	// Plugins for PodGroupInfos
 	framework.RegisterPluginBuilder("predicates", predicates.New)
 	framework.RegisterPluginBuilder("priority", priority.New)
+	framework.RegisterPluginBuilder("gpujoborder", gpujoborder.New)
 	framework.RegisterPluginBuilder("nodeplacement", nodeplacement.New)
 	framework.RegisterPluginBuilder("nominatednode", nominatednode.New)
 	framework.RegisterPluginBuilder("numa", numa.New)

pkg/scheduler/plugins/gpujoborder/gpujoborder.go

@@ -0,0 +1,68 @@
+// Copyright 2026 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package gpujoborder
+
+import (
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/podgroup_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/framework"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/log"
+)
+
+const (
+	ModeEvictLargerFirst  = "evict-larger-first"
+	ModeEvictSmallerFirst = "evict-smaller-first"
+)
+
+type gpuJobOrderPlugin struct {
+	mode string
+}
+
+func New(arguments framework.PluginArguments) framework.Plugin {
+	mode := arguments.GetString("mode", ModeEvictLargerFirst)
+	if mode != ModeEvictLargerFirst && mode != ModeEvictSmallerFirst {
+		log.InfraLogger.Warningf("gpujoborder: unrecognized mode %q, defaulting to %s", mode, ModeEvictLargerFirst)
+		mode = ModeEvictLargerFirst
+	}
+	return &gpuJobOrderPlugin{mode: mode}
+}
+
+func (rp *gpuJobOrderPlugin) Name() string {
+	return "gpujoborder"
+}
+
+func (rp *gpuJobOrderPlugin) OnSessionOpen(ssn *framework.Session) {
+	ssn.AddVictimOrderFn(rp.VictimOrderFn)
+}
+
+func (rp *gpuJobOrderPlugin) VictimOrderFn(l, r interface{}) int {
+	lv := l.(*podgroup_info.PodGroupInfo)
+	rv := r.(*podgroup_info.PodGroupInfo)
+
+	if lv.Priority != rv.Priority {
+		return 0
+	}
+
+	lGPU := lv.GetAliveTasksRequestedGPUs()
+	rGPU := rv.GetAliveTasksRequestedGPUs()
+
+	switch rp.mode {
+	case ModeEvictSmallerFirst:
+		if lGPU < rGPU {
+			return -1
+		}
+		if lGPU > rGPU {
+			return 1
+		}
+	default:
+		if lGPU > rGPU {
+			return -1
+		}
+		if lGPU < rGPU {
+			return 1
+		}
+	}
+	return 0
+}
+
+func (rp *gpuJobOrderPlugin) OnSessionClose(_ *framework.Session) {}