scheduler: permit other plugin to set numaAffinity (#2182)

Signed-off-by: wangjianyu.wjy <wangjianyu.wjy@alibaba-inc.com> Co-authored-by: wangjianyu.wjy <wangjianyu.wjy@alibaba-inc.com>
koordinator-sh · Sep 2, 2024 · cb89d6d · cb89d6d
1 parent 4d89c84
commit cb89d6d
Show file tree

Hide file tree

Showing 7 changed files with 21 additions and 20 deletions.
diff --git a/pkg/scheduler/frameworkext/topologymanager/manager.go b/pkg/scheduler/frameworkext/topologymanager/manager.go
@@ -63,15 +63,16 @@ func (m *topologyManager) Admit(ctx context.Context, cycleState *framework.Cycle
 	store := s.(*Store)
 
 	policy := createNUMATopologyPolicy(policyType, numaNodes)
-
-	bestHint, admit := m.calculateAffinity(ctx, cycleState, policy, pod, nodeName, exclusivePolicy, allNUMANodeStatus)
-	klog.V(5).Infof("Best TopologyHint for (pod: %v): %v on node: %v", klog.KObj(pod), bestHint, nodeName)
-	if !admit {
-		return framework.NewStatus(framework.Unschedulable, "node(s) NUMA Topology affinity error")
+	bestHint, ok := store.GetAffinity(nodeName)
+	if !ok {
+		var admit bool
+		bestHint, admit = m.calculateAffinity(ctx, cycleState, policy, pod, nodeName, exclusivePolicy, allNUMANodeStatus)
+		klog.V(5).Infof("Best TopologyHint for (pod: %v): %v on node: %v", klog.KObj(pod), bestHint, nodeName)
+		if !admit {
+			return framework.NewStatus(framework.Unschedulable, "node(s) NUMA Topology affinity error")
+		}
+		store.SetAffinity(nodeName, bestHint)
 	}
-
-	store.SetAffinity(nodeName, bestHint)
-
 	status := m.allocateResources(ctx, cycleState, bestHint, pod, nodeName)
 	if !status.IsSuccess() {
 		return status

diff --git a/pkg/scheduler/frameworkext/topologymanager/store.go b/pkg/scheduler/frameworkext/topologymanager/store.go
@@ -56,11 +56,11 @@ func (s *Store) SetAffinity(nodeName string, affinity NUMATopologyHint) {
 	s.affinityMap.Store(nodeName, &affinity)
 }
 
-func (s *Store) GetAffinity(nodeName string) NUMATopologyHint {
+func (s *Store) GetAffinity(nodeName string) (NUMATopologyHint, bool) {
 	val, ok := s.affinityMap.Load(nodeName)
 	if !ok {
-		return NUMATopologyHint{}
+		return NUMATopologyHint{}, false
 	}
 	hint := val.(*NUMATopologyHint)
-	return *hint
+	return *hint, true
 }
diff --git a/pkg/scheduler/plugins/deviceshare/plugin.go b/pkg/scheduler/plugins/deviceshare/plugin.go
@@ -300,7 +300,7 @@ func (p *Plugin) Filter(ctx context.Context, cycleState *framework.CycleState, p
 	var affinity topologymanager.NUMATopologyHint
 	if !p.disableDeviceNUMATopologyAlignment {
 		store := topologymanager.GetStore(cycleState)
-		affinity = store.GetAffinity(nodeInfo.Node().Name)
+		affinity, _ = store.GetAffinity(nodeInfo.Node().Name)
 	}
 
 	allocator := &AutopilotAllocator{
@@ -369,7 +369,7 @@ func (p *Plugin) FilterReservation(ctx context.Context, cycleState *framework.Cy
 	var affinity topologymanager.NUMATopologyHint
 	if !p.disableDeviceNUMATopologyAlignment {
 		store := topologymanager.GetStore(cycleState)
-		affinity = store.GetAffinity(nodeInfo.Node().Name)
+		affinity, _ = store.GetAffinity(nodeInfo.Node().Name)
 	}
 
 	allocator := &AutopilotAllocator{
@@ -417,7 +417,7 @@ func (p *Plugin) Reserve(ctx context.Context, cycleState *framework.CycleState,
 	var affinity topologymanager.NUMATopologyHint
 	if !p.disableDeviceNUMATopologyAlignment {
 		store := topologymanager.GetStore(cycleState)
-		affinity = store.GetAffinity(nodeInfo.Node().Name)
+		affinity, _ = store.GetAffinity(nodeInfo.Node().Name)
 	}
 
 	allocator := &AutopilotAllocator{

diff --git a/pkg/scheduler/plugins/deviceshare/scoring.go b/pkg/scheduler/plugins/deviceshare/scoring.go
@@ -62,7 +62,7 @@ func (p *Plugin) Score(ctx context.Context, cycleState *framework.CycleState, po
 	}
 
 	store := topologymanager.GetStore(cycleState)
-	affinity := store.GetAffinity(nodeName)
+	affinity, _ := store.GetAffinity(nodeName)
 
 	allocator := &AutopilotAllocator{
 		state:      state,
@@ -130,7 +130,7 @@ func (p *Plugin) ScoreReservation(ctx context.Context, cycleState *framework.Cyc
 	}
 
 	store := topologymanager.GetStore(cycleState)
-	affinity := store.GetAffinity(nodeInfo.Node().Name)
+	affinity, _ := store.GetAffinity(nodeInfo.Node().Name)
 
 	allocator := &AutopilotAllocator{
 		state:      state,

diff --git a/pkg/scheduler/plugins/nodenumaresource/plugin.go b/pkg/scheduler/plugins/nodenumaresource/plugin.go
@@ -484,7 +484,7 @@ func (p *Plugin) FilterReservation(ctx context.Context, cycleState *framework.Cy
 	restoreState := reservationRestoreState.getNodeState(nodeName)
 
 	store := topologymanager.GetStore(cycleState)
-	affinity := store.GetAffinity(nodeName)
+	affinity, _ := store.GetAffinity(nodeName)
 	resourceOptions, err := p.getResourceOptions(state, node, pod, requestCPUBind, affinity, topologyOptions)
 	if err != nil {
 		return framework.NewStatus(framework.UnschedulableAndUnresolvable, err.Error())
@@ -541,7 +541,7 @@ func (p *Plugin) Reserve(ctx context.Context, cycleState *framework.CycleState,
 
 	// TODO: de-duplicate logic done by the Filter phase and move head the pre-process of the resource options
 	store := topologymanager.GetStore(cycleState)
-	affinity := store.GetAffinity(nodeName)
+	affinity, _ := store.GetAffinity(nodeName)
 	resourceOptions, err := p.getResourceOptions(state, node, pod, requestCPUBind, affinity, topologyOptions)
 	if err != nil {
 		return framework.AsStatus(err)

diff --git a/pkg/scheduler/plugins/nodenumaresource/plugin_test.go b/pkg/scheduler/plugins/nodenumaresource/plugin_test.go
@@ -2183,7 +2183,7 @@ func TestFilterWithNUMANodeScoring(t *testing.T) {
 			assert.NoError(t, err)
 			status = pl.Filter(context.TODO(), cycleState, tt.requestedPod, nodeInfo)
 			assert.True(t, status.IsSuccess())
-			hint := topologymanager.GetStore(cycleState).GetAffinity(tt.node.Name)
+			hint, _ := topologymanager.GetStore(cycleState).GetAffinity(tt.node.Name)
 			assert.Equal(t, tt.wantAffinity.GetBits(), hint.NUMANodeAffinity.GetBits())
 		})
 	}

diff --git a/pkg/scheduler/plugins/nodenumaresource/scoring.go b/pkg/scheduler/plugins/nodenumaresource/scoring.go
@@ -92,7 +92,7 @@ func (p *Plugin) Score(ctx context.Context, cycleState *framework.CycleState, po
 	}
 
 	store := topologymanager.GetStore(cycleState)
-	affinity := store.GetAffinity(nodeName)
+	affinity, _ := store.GetAffinity(nodeName)
 	resourceOptions, err := p.getResourceOptions(state, node, pod, requestCPUBind, affinity, topologyOptions)
 	if err != nil {
 		return 0, nil