diff --git a/.changes/unreleased/added-20260729-142722.yaml b/.changes/unreleased/added-20260729-142722.yaml new file mode 100644 index 000000000..9c12dfcf1 --- /dev/null +++ b/.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 diff --git a/pkg/scheduler/actions/utils/job_order_by_queue.go b/pkg/scheduler/actions/utils/job_order_by_queue.go index d8ca3c123..ce98df949 100644 --- a/pkg/scheduler/actions/utils/job_order_by_queue.go +++ b/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), diff --git a/pkg/scheduler/api/podgroup_info/job_info.go b/pkg/scheduler/api/podgroup_info/job_info.go index 9eefbd41b..0d098cefd 100644 --- a/pkg/scheduler/api/podgroup_info/job_info.go +++ b/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 { diff --git a/pkg/scheduler/framework/session.go b/pkg/scheduler/framework/session.go index 099a05548..37eb64def 100644 --- a/pkg/scheduler/framework/session.go +++ b/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 diff --git a/pkg/scheduler/framework/session_plugins.go b/pkg/scheduler/framework/session_plugins.go index 63c639e90..ba3e1e930 100644 --- a/pkg/scheduler/framework/session_plugins.go +++ b/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 { diff --git a/pkg/scheduler/framework/session_plugins_test.go b/pkg/scheduler/framework/session_plugins_test.go index 73f3c598f..3d735d52f 100644 --- a/pkg/scheduler/framework/session_plugins_test.go +++ b/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) + } + } + } +} diff --git a/pkg/scheduler/plugins/factory.go b/pkg/scheduler/plugins/factory.go index 23b0a4486..6fcadcc55 100644 --- a/pkg/scheduler/plugins/factory.go +++ b/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) diff --git a/pkg/scheduler/plugins/gpujoborder/gpujoborder.go b/pkg/scheduler/plugins/gpujoborder/gpujoborder.go new file mode 100644 index 000000000..3f536baed --- /dev/null +++ b/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) {} diff --git a/pkg/scheduler/plugins/gpujoborder/gpujoborder_test.go b/pkg/scheduler/plugins/gpujoborder/gpujoborder_test.go new file mode 100644 index 000000000..bcccd9f27 --- /dev/null +++ b/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) + } +}