Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions .changes/unreleased/added-20260729-142722.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
kind: Added
body: |-
Add gpujoborder plugin for configurable GPU-based JobOrderFn tiebreak on priority ties
2 changes: 1 addition & 1 deletion pkg/scheduler/actions/utils/job_order_by_queue.go
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
16 changes: 10 additions & 6 deletions pkg/scheduler/api/podgroup_info/job_info.go
Original file line number Diff line number Diff line change
Expand Up @@ -93,6 +93,7 @@ type PodGroupInfo struct {
tasksToAllocateInitResourceVector resource_info.ResourceVector
PodStatusIndex map[pod_status.PodStatus]pod_info.PodsMap
activeAllocatedCount *int
aliveTasksRequestedGPUs *float64
Comment thread
CoolingCube marked this conversation as resolved.
}

func NewPodGroupInfo(uid common_info.PodGroupID, tasks ...*pod_info.PodInfo) *PodGroupInfo {
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -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 {
Expand Down
2 changes: 2 additions & 0 deletions pkg/scheduler/framework/session.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
51 changes: 44 additions & 7 deletions pkg/scheduler/framework/session_plugins.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
Expand Down Expand Up @@ -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 {
Comment thread
CoolingCube marked this conversation as resolved.
for _, jof := range ssn.JobOrderFns {
if j := jof(l, r); j != 0 {
return j > 0
}
}
for _, vof := range ssn.VictimOrderFns {
Comment thread
CoolingCube marked this conversation as resolved.
if v := vof(l, r); v != 0 {
return v < 0
}
}
return !jobOrderCreationFallback(l, r)
}

func (ssn *Session) TaskOrderFn(l, r interface{}) bool {
Expand Down
42 changes: 42 additions & 0 deletions pkg/scheduler/framework/session_plugins_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
)

Expand Down Expand Up @@ -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)
}
}
}
}
2 changes: 2 additions & 0 deletions pkg/scheduler/plugins/factory.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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)
Expand Down
68 changes: 68 additions & 0 deletions pkg/scheduler/plugins/gpujoborder/gpujoborder.go
Original file line number Diff line number Diff line change
@@ -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) {}
Loading
Loading