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/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/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 (
+	ModePreferLarger  = "prefer-larger"
+	ModePreferSmaller = "prefer-smaller"
+)
+
+type gpuJobOrderPlugin struct {
+	mode string
+}
+
+func New(arguments framework.PluginArguments) framework.Plugin {
+	mode := arguments.GetString("mode", ModePreferLarger)
+	if mode != ModePreferLarger && mode != ModePreferSmaller {
+		log.InfraLogger.Warningf("gpujoborder: unrecognized mode %q, defaulting to prefer-larger", mode)
+		mode = ModePreferLarger
+	}
+	return &gpuJobOrderPlugin{mode: mode}
+}
+
+func (rp *gpuJobOrderPlugin) Name() string {
+	return "gpujoborder"
+}
+
+func (rp *gpuJobOrderPlugin) OnSessionOpen(ssn *framework.Session) {
+	ssn.AddJobOrderFn(rp.JobOrderFn)
+}
+
+func (rp *gpuJobOrderPlugin) JobOrderFn(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 ModePreferSmaller:
+		if lGPU < rGPU {
+			return -1
+		}
+		if lGPU > rGPU {
+			return 1
+		}
+	default: // ModePreferLarger
+		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,95 @@
+// Copyright 2026 NVIDIA CORPORATION
+// SPDX-License-Identifier: Apache-2.0
+
+package gpujoborder
+
+import (
+	"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/resource_info"
+	"github.com/kai-scheduler/KAI-scheduler/pkg/scheduler/framework"
+)
+
+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 TestJobOrderFn_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.JobOrderFn(a, b); got != 0 {
+		t.Errorf("expected 0 when priorities differ, got %d", got)
+	}
+}
+
+func TestJobOrderFn_SamePriority_PrefersLarger_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.JobOrderFn(large, small); got != -1 {
+		t.Errorf("expected -1 (larger job preferred as victim), got %d", got)
+	}
+	if got := rp.JobOrderFn(small, large); got != 1 {
+		t.Errorf("expected 1, got %d", got)
+	}
+}
+
+func TestJobOrderFn_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.JobOrderFn(a, b); got != 0 {
+		t.Errorf("expected 0 (equal GPU falls through), got %d", got)
+	}
+}
+
+func TestJobOrderFn_PreferSmallerMode(t *testing.T) {
+	vm := resource_info.NewResourceVectorMap()
+	small := makeGPUPodGroup("small", 10, 1, vm)
+	large := makeGPUPodGroup("large", 10, 2, vm)
+	rp := newPlugin(t, "prefer-smaller")
+	if got := rp.JobOrderFn(small, large); got != -1 {
+		t.Errorf("expected -1 (smaller job preferred as victim in prefer-smaller mode), got %d", got)
+	}
+	if got := rp.JobOrderFn(large, small); got != 1 {
+		t.Errorf("expected 1, got %d", got)
+	}
+}
+
+func TestNew_UnrecognizedMode_FallsBackToPreferLarger(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.JobOrderFn(large, small); got != -1 {
+		t.Errorf("expected fallback to prefer-larger behavior, got %d", got)
+	}
+}