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

.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,10 @@ func (ssn *Session) AddJobOrderFn(jof common_info.CompareFn) {
 	ssn.JobOrderFns = append(ssn.JobOrderFns, jof)
 }
 
+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 +272,36 @@ 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)
+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/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) {}

pkg/scheduler/plugins/gpujoborder/gpujoborder_test.go

@@ -0,0 +1,176 @@
+// Copyright 2026 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package gpujoborder
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/common_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/api/pod_status"
+	"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/framework"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/plugins/elastic"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/scheduler_util"
+)
+
+func makeGPUPodGroup(uid string, priority int32, gpuCount float64, vm *resource_info.ResourceVectorMap) *podgroup_info.PodGroupInfo {
+	task := &pod_info.PodInfo{
+		UID:          common_info.PodID(uid + "-task"),
+		ResReqVector: resource_info.NewResourceVectorWithValues(0, 0, gpuCount, vm),
+		Status:       pod_status.Running,
+	}
+	pg := podgroup_info.NewPodGroupInfoWithVectorMap(common_info.PodGroupID(uid), vm, task)
+	pg.Priority = priority
+	return pg
+}
+
+func newPlugin(t *testing.T, mode string) *gpuJobOrderPlugin {
+	t.Helper()
+	args := framework.PluginArguments{}
+	if mode != "" {
+		args["mode"] = mode
+	}
+	rp, ok := New(args).(*gpuJobOrderPlugin)
+	if !ok {
+		t.Fatalf("New() did not return *gpuJobOrderPlugin")
+	}
+	return rp
+}
+
+func TestVictimOrderFn_PriorityDiffers_Defers(t *testing.T) {
+	vm := resource_info.NewResourceVectorMap()
+	a := makeGPUPodGroup("a", 50, 1, vm)
+	b := makeGPUPodGroup("b", 10, 1, vm)
+	rp := newPlugin(t, "")
+	if got := rp.VictimOrderFn(a, b); got != 0 {
+		t.Errorf("expected 0 when priorities differ, got %d", got)
+	}
+}
+
+func TestVictimOrderFn_SamePriority_EvictLargerFirst_Default(t *testing.T) {
+	vm := resource_info.NewResourceVectorMap()
+	small := makeGPUPodGroup("small", 10, 1, vm)
+	large := makeGPUPodGroup("large", 10, 2, vm)
+	rp := newPlugin(t, "")
+	if got := rp.VictimOrderFn(large, small); got != -1 {
+		t.Errorf("expected -1 (larger job preferred as victim), got %d", got)
+	}
+	if got := rp.VictimOrderFn(small, large); got != 1 {
+		t.Errorf("expected 1, got %d", got)
+	}
+}
+
+func TestVictimOrderFn_SamePriority_SameGPU_FallsThrough(t *testing.T) {
+	vm := resource_info.NewResourceVectorMap()
+	a := makeGPUPodGroup("a", 10, 1, vm)
+	b := makeGPUPodGroup("b", 10, 1, vm)
+	rp := newPlugin(t, "")
+	if got := rp.VictimOrderFn(a, b); got != 0 {
+		t.Errorf("expected 0 (equal GPU falls through), got %d", got)
+	}
+}
+
+func TestVictimOrderFn_EvictSmallerFirstMode(t *testing.T) {
+	vm := resource_info.NewResourceVectorMap()
+	small := makeGPUPodGroup("small", 10, 1, vm)
+	large := makeGPUPodGroup("large", 10, 2, vm)
+	rp := newPlugin(t, "evict-smaller-first")
+	if got := rp.VictimOrderFn(small, large); got != -1 {
+		t.Errorf("expected -1 (smaller job preferred as victim in evict-smaller-first mode), got %d", got)
+	}
+	if got := rp.VictimOrderFn(large, small); got != 1 {
+		t.Errorf("expected 1, got %d", got)
+	}
+}
+
+func TestNew_UnrecognizedMode_FallsBackToEvictLargerFirst(t *testing.T) {
+	vm := resource_info.NewResourceVectorMap()
+	small := makeGPUPodGroup("small", 10, 1, vm)
+	large := makeGPUPodGroup("large", 10, 2, vm)
+	rp := newPlugin(t, "totally-not-a-real-mode")
+	if got := rp.VictimOrderFn(large, small); got != -1 {
+		t.Errorf("expected fallback to evict-larger-first behavior, got %d", got)
+	}
+}
+
+func TestGpujoborder_DoesNotAffect_PendingJobOrdering(t *testing.T) {
+	rp := newPlugin(t, "")
+
+	vm := resource_info.NewResourceVectorMap()
+	small := makeGPUPodGroup("small", 10, 1, vm)
+	large := makeGPUPodGroup("large", 10, 2, vm)
+
+	ssnWithout := &framework.Session{}
+	baselineResult := ssnWithout.JobOrderFn(small, large)
+
+	ssnWith := &framework.Session{}
+	rp.OnSessionOpen(ssnWith)
+	withPluginResult := ssnWith.JobOrderFn(small, large)
+
+	if len(ssnWith.JobOrderFns) != 0 {
+		t.Errorf("expected gpujoborder to register 0 JobOrderFns, got %d", len(ssnWith.JobOrderFns))
+	}
+	if len(ssnWith.VictimOrderFns) != 1 {
+		t.Errorf("expected gpujoborder to register exactly 1 VictimOrderFn, got %d", len(ssnWith.VictimOrderFns))
+	}
+	if baselineResult != withPluginResult {
+		t.Errorf("pending-job ordering changed after registering gpujoborder: without=%v, with=%v",
+			baselineResult, withPluginResult)
+	}
+}
+
+func makeElasticJob(uid string, priority int32, minAvailable int32, allocatedTasks int,
+	gpuPerTask float64, vm *resource_info.ResourceVectorMap) *podgroup_info.PodGroupInfo {
+
+	pg := podgroup_info.NewPodGroupInfoWithVectorMap(common_info.PodGroupID(uid), vm)
+	pg.Priority = priority
+
+	root := subgroup_info.NewSubGroupSet(subgroup_info.RootSubGroupSetName, nil)
+	root.AddPodSet(subgroup_info.NewPodSet(podgroup_info.DefaultSubGroup, minAvailable, nil))
+	pg.RootSubGroupSet = root
+	pg.PodSets = root.GetDescendantPodSets()
+
+	for i := 0; i < allocatedTasks; i++ {
+		task := &pod_info.PodInfo{
+			UID:          common_info.PodID(fmt.Sprintf("%s-task-%d", uid, i)),
+			ResReqVector: resource_info.NewResourceVectorWithValues(0, 0, gpuPerTask, vm),
+			Status:       pod_status.Running,
+		}
+		pg.AddTaskInfo(task)
+	}
+	return pg
+}
+
+func TestElasticProtection_OutranksGpujoborder(t *testing.T) {
+	rp := newPlugin(t, "")
+
+	vm := resource_info.NewResourceVectorMap()
+	atMinLarge := makeElasticJob("at-min-large", 10, 2, 2, 1.0, vm)
+	aboveMinSmall := makeElasticJob("above-min-small", 10, 1, 2, 0.5, vm)
+
+	ssn := &framework.Session{}
+	ssn.AddJobOrderFn(elastic.JobOrderFn)
+	ssn.AddVictimOrderFn(rp.VictimOrderFn)
+
+	victimLessFn := func(l, r interface{}) bool {
+		return ssn.VictimOrderFn(l, r)
+	}
+	pq := scheduler_util.NewPriorityQueue(victimLessFn, scheduler_util.QueueCapacityInfinite)
+	pq.Push(atMinLarge)
+	pq.Push(aboveMinSmall)
+
+	popped := pq.Pop().(*podgroup_info.PodGroupInfo)
+
+	t.Logf("Popped as victim: %q (GPUs=%v)", popped.UID, popped.GetAliveTasksRequestedGPUs())
+
+	if popped.UID != aboveMinSmall.UID {
+		t.Errorf("expected elastic's above-min job to be evicted first (protecting the at-min job), "+
+			"but got %q evicted instead -- gpujoborder's raw GPU-size comparison incorrectly "+
+			"outranked elastic's protection", popped.UID)
+	}
+}