From 4c3d2dd19caeeaa8edf993be9c43d08d81c20baf Mon Sep 17 00:00:00 2001
From: mesutoezdil <mesudozdil@gmail.com>
Date: Thu, 4 Jun 2026 15:13:49 +0200
Subject: [PATCH] fix(scheduler): make NUMA sort opt-in for binpack/spread
 policies

Add NumaIgnore bool to DeviceUsageList. When the pod annotation
hami.io/topology-aware-scoring is set to "false", Less() sorts by
Score only so spread and binpack are not overridden by NUMA grouping.
Default behavior is unchanged (NUMA-first, backward compatible).

Fixes #1806

Signed-off-by: mesutoezdil <mesudozdil@gmail.com>
---
 pkg/scheduler/policy/gpu_policy.go      | 18 ++++++++++--------
 pkg/scheduler/policy/gpu_policy_test.go | 22 ++++++++++++++++++++++
 pkg/scheduler/scheduler.go              |  4 +++-
 pkg/util/types.go                       |  2 ++
 4 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/pkg/scheduler/policy/gpu_policy.go b/pkg/scheduler/policy/gpu_policy.go
index e404dde56..35c52a259 100644
--- a/pkg/scheduler/policy/gpu_policy.go
+++ b/pkg/scheduler/policy/gpu_policy.go
@@ -32,6 +32,7 @@ type DeviceListsScore struct {
 type DeviceUsageList struct {
 	DeviceLists []*DeviceListsScore
 	Policy      string
+	NumaIgnore  bool
 }
 
 func (l DeviceUsageList) Len() int {
@@ -43,17 +44,17 @@ func (l DeviceUsageList) Swap(i, j int) {
 }
 
 func (l DeviceUsageList) Less(i, j int) bool {
-	if l.Policy == util.GPUSchedulerPolicyBinpack.String() {
-		if l.DeviceLists[i].Device.Numa == l.DeviceLists[j].Device.Numa {
-			return l.DeviceLists[i].Score < l.DeviceLists[j].Score
+	di, dj := l.DeviceLists[i], l.DeviceLists[j]
+	if !l.NumaIgnore && di.Device.Numa != dj.Device.Numa {
+		if l.Policy == util.GPUSchedulerPolicyBinpack.String() {
+			return di.Device.Numa > dj.Device.Numa
 		}
-		return l.DeviceLists[i].Device.Numa > l.DeviceLists[j].Device.Numa
+		return di.Device.Numa < dj.Device.Numa
 	}
-	// default policy is spread
-	if l.DeviceLists[i].Device.Numa == l.DeviceLists[j].Device.Numa {
-		return l.DeviceLists[i].Score > l.DeviceLists[j].Score
+	if l.Policy == util.GPUSchedulerPolicyBinpack.String() {
+		return di.Score < dj.Score
 	}
-	return l.DeviceLists[i].Device.Numa < l.DeviceLists[j].Device.Numa
+	return di.Score > dj.Score
 }
 
 func (l DeviceUsageList) DeepCopy() DeviceUsageList {
@@ -67,6 +68,7 @@ func (l DeviceUsageList) DeepCopy() DeviceUsageList {
 	return DeviceUsageList{
 		DeviceLists: deviceLists,
 		Policy:      l.Policy,
+		NumaIgnore:  l.NumaIgnore,
 	}
 }
 
diff --git a/pkg/scheduler/policy/gpu_policy_test.go b/pkg/scheduler/policy/gpu_policy_test.go
index f70eb2f21..ce6efd837 100644
--- a/pkg/scheduler/policy/gpu_policy_test.go
+++ b/pkg/scheduler/policy/gpu_policy_test.go
@@ -124,6 +124,7 @@ func TestDeviceUsageList_Less(t *testing.T) {
 	tests := []struct {
 		name         string
 		policy       string
+		numaIgnore   bool
 		deviceLists  []*DeviceListsScore
 		expectedLess bool
 	}{
@@ -181,12 +182,33 @@ func TestDeviceUsageList_Less(t *testing.T) {
 			},
 			expectedLess: true,
 		},
+		{
+			name:       "Binpack NumaIgnore score primary",
+			policy:     "binpack",
+			numaIgnore: true,
+			deviceLists: []*DeviceListsScore{
+				{Device: &device.DeviceUsage{Numa: 0, Used: 10}, Score: 10},
+				{Device: &device.DeviceUsage{Numa: 1, Used: 20}, Score: 20},
+			},
+			expectedLess: true,
+		},
+		{
+			name:       "Spread NumaIgnore score primary",
+			policy:     "spread",
+			numaIgnore: true,
+			deviceLists: []*DeviceListsScore{
+				{Device: &device.DeviceUsage{Numa: 1, Used: 10}, Score: 10},
+				{Device: &device.DeviceUsage{Numa: 0, Used: 20}, Score: 20},
+			},
+			expectedLess: false,
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			l := DeviceUsageList{
 				Policy:      tt.policy,
+				NumaIgnore:  tt.numaIgnore,
 				DeviceLists: tt.deviceLists,
 			}
 			i, j := 0, 1
diff --git a/pkg/scheduler/scheduler.go b/pkg/scheduler/scheduler.go
index f1df5c87e..5cc7c22b0 100644
--- a/pkg/scheduler/scheduler.go
+++ b/pkg/scheduler/scheduler.go
@@ -531,12 +531,14 @@ func (s *Scheduler) getNodesUsage(nodes *[]string, task *corev1.Pod) (*map[strin
 		return &overallnodeMap, failedNodes, err
 	}
 
+	userGPUPolicy := util.GetGPUSchedulerPolicyByPod(device.GPUSchedulerPolicy, task)
+	numaIgnore := task != nil && task.Annotations != nil && task.Annotations[util.GPUTopologyAwareAnnotationKey] == "false"
 	for _, node := range allNodes {
 		nodeInfo := &NodeUsage{}
-		userGPUPolicy := util.GetGPUSchedulerPolicyByPod(device.GPUSchedulerPolicy, task)
 		nodeInfo.Node = node.Node
 		nodeInfo.Devices = policy.DeviceUsageList{
 			Policy:      userGPUPolicy,
+			NumaIgnore:  numaIgnore,
 			DeviceLists: make([]*policy.DeviceListsScore, 0),
 		}
 		for _, k := range node.Devices {
diff --git a/pkg/util/types.go b/pkg/util/types.go
index 780e4dfe1..f5f298e57 100644
--- a/pkg/util/types.go
+++ b/pkg/util/types.go
@@ -78,6 +78,8 @@ const (
 	NodeSchedulerPolicyAnnotationKey = "hami.io/node-scheduler-policy"
 	// GPUSchedulerPolicyAnnotationKey is user set Pod annotation to change this default GPU policy.
 	GPUSchedulerPolicyAnnotationKey = "hami.io/gpu-scheduler-policy"
+	// GPUTopologyAwareAnnotationKey disables NUMA-first sorting when set to "false".
+	GPUTopologyAwareAnnotationKey = "hami.io/topology-aware-scoring"
 )
 
 func (s SchedulerPolicyName) String() string {