feat: add preScore

Signed-off-by: LY-today <[email protected]>
koordinator-sh · Jan 3, 2025 · 1b77ba5 · 1b77ba5
1 parent 507a5b5
commit 1b77ba5
Show file tree

Hide file tree

Showing 4 changed files with 82 additions and 9 deletions.
diff --git a/pkg/scheduler/plugins/noderesourcefitplus/node_resources_fit_plus.go b/pkg/scheduler/plugins/noderesourcefitplus/node_resources_fit_plus.go
@@ -34,6 +34,8 @@ import (
 const (
 	// Name is plugin name
 	Name = "NodeResourcesFitPlus"
+
+	preScoreStateKey = "PreScore" + Name
 )
 
 var (
@@ -63,6 +65,11 @@ func (s *Plugin) Name() string {
 	return Name
 }
 
+func (s *Plugin) PreScore(ctx context.Context, cycleState *framework.CycleState, pod *v1.Pod, nodes []*v1.Node) *framework.Status {
+	cycleState.Write(preScoreStateKey, computePodResourceRequest(pod))
+	return nil
+}
+
 func (s *Plugin) Score(ctx context.Context, state *framework.CycleState, p *v1.Pod, nodeName string) (int64, *framework.Status) {
 
 	nodeInfo, err := s.handle.SnapshotSharedLister().NodeInfos().Get(nodeName)
@@ -73,7 +80,12 @@ func (s *Plugin) Score(ctx context.Context, state *framework.CycleState, p *v1.P
 	var nodeScore int64
 	var weightSum int64
 
-	podRequest, _ := fitsRequest(computePodResourceRequest(p).Resource, nodeInfo)
+	scoreState, err := getPreScoreState(state)
+	if err != nil {
+		return 0, framework.NewStatus(framework.Error, fmt.Sprintf("get State node %q from PreScore: %v", nodeName, err))
+	}
+
+	podRequest, _ := fitsRequest(scoreState.Resource, nodeInfo)
 
 	for _, requestSourceName := range podRequest {
 		v, ok := s.args.Resources[requestSourceName]
@@ -115,14 +127,19 @@ func (p *Plugin) ScoreExtensions() framework.ScoreExtensions {
 	return nil
 }
 
-type preFilterState struct {
+type preScoreState struct {
 	framework.Resource
 }
 
-func computePodResourceRequest(pod *v1.Pod) *preFilterState {
+// Clone the prefilter state.
+func (s *preScoreState) Clone() framework.StateData {
+	return s
+}
+
+func computePodResourceRequest(pod *v1.Pod) *preScoreState {
 	// pod hasn't scheduled yet so we don't need to worry about InPlacePodVerticalScalingEnabled
 	reqs := resource.PodRequests(pod, resource.PodResourcesOptions{})
-	result := &preFilterState{}
+	result := &preScoreState{}
 	result.SetMaxResource(reqs)
 	return result
 }
@@ -169,3 +186,17 @@ func fitsRequest(podRequest framework.Resource, nodeInfo *framework.NodeInfo) ([
 
 	return podRequestResource, nodeRequestResource
 }
+
+func getPreScoreState(cycleState *framework.CycleState) (*preScoreState, error) {
+	c, err := cycleState.Read(preScoreStateKey)
+	if err != nil {
+		// preFilterState doesn't exist, likely PreFilter wasn't invoked.
+		return nil, fmt.Errorf("error reading %q from cycleState: %w", preScoreStateKey, err)
+	}
+
+	s, ok := c.(*preScoreState)
+	if !ok {
+		return nil, fmt.Errorf("%+v  convert to NodeResourcesFit.preFilterState error", c)
+	}
+	return s, nil
+}
diff --git a/pkg/scheduler/plugins/noderesourcefitplus/node_resources_fit_plus_test.go b/pkg/scheduler/plugins/noderesourcefitplus/node_resources_fit_plus_test.go
@@ -251,6 +251,12 @@ func TestPlugin_Score(t *testing.T) {
 			},
 		},
 	}
+
+	status := p.(framework.PreScorePlugin).PreScore(context.TODO(), cycleState, pod, nodes)
+	if status != nil {
+		t.Fatal("PreScore run err")
+	}
+
 	scoreNode1, _ := p.(*Plugin).Score(context.TODO(), cycleState, pod, "testNode1")
 	scoreNode2, _ := p.(*Plugin).Score(context.TODO(), cycleState, pod, "testNode2")
 	if scoreNode1 <= scoreNode2 {

diff --git a/pkg/scheduler/plugins/scarceresourceavoidance/scarce_resource_avoidance.go b/pkg/scheduler/plugins/scarceresourceavoidance/scarce_resource_avoidance.go
@@ -31,6 +31,8 @@ import (
 const (
 	// Name is plugin name
 	Name = "ScarceResourceAvoidance"
+
+	preScoreStateKey = "PreScore" + Name
 )
 
 var (
@@ -60,14 +62,22 @@ func (s *Plugin) Name() string {
 	return Name
 }
 
+func (s *Plugin) PreScore(ctx context.Context, cycleState *framework.CycleState, pod *v1.Pod, nodes []*v1.Node) *framework.Status {
+	cycleState.Write(preScoreStateKey, computePodResourceRequest(pod))
+	return nil
+}
+
 func (s *Plugin) Score(ctx context.Context, state *framework.CycleState, p *v1.Pod, nodeName string) (int64, *framework.Status) {
 	nodeInfo, err := s.handle.SnapshotSharedLister().NodeInfos().Get(nodeName)
 	if err != nil {
 		return 0, framework.NewStatus(framework.Error, fmt.Sprintf("getting node %q from Snapshot: %v", nodeName, err))
 	}
 
-	podRequest := computePodResourceRequest(p)
-	podRequestResource, nodeAllocatableResource := fitsRequest(podRequest.Resource, nodeInfo)
+	scoreState, err := getPreScoreState(state)
+	if err != nil {
+		return 0, framework.NewStatus(framework.Error, fmt.Sprintf("get State node %q from PreScore: %v", nodeName, err))
+	}
+	podRequestResource, nodeAllocatableResource := fitsRequest(scoreState.Resource, nodeInfo)
 	diffNames := difference(nodeAllocatableResource, podRequestResource)
 	intersectNames := intersection(diffNames, s.args.Resources)
 
@@ -116,14 +126,19 @@ func (p *Plugin) ScoreExtensions() framework.ScoreExtensions {
 	return nil
 }
 
-type preFilterState struct {
+type preScoreState struct {
 	framework.Resource
 }
 
-func computePodResourceRequest(pod *v1.Pod) *preFilterState {
+// Clone the prefilter state.
+func (s *preScoreState) Clone() framework.StateData {
+	return s
+}
+
+func computePodResourceRequest(pod *v1.Pod) *preScoreState {
 	// pod hasn't scheduled yet so we don't need to worry about InPlacePodVerticalScalingEnabled
 	reqs := resource.PodRequests(pod, resource.PodResourcesOptions{})
-	result := &preFilterState{}
+	result := &preScoreState{}
 	result.SetMaxResource(reqs)
 	return result
 }
@@ -170,6 +185,21 @@ func fitsRequest(podRequest framework.Resource, nodeInfo *framework.NodeInfo) ([
 
 	return podRequestResource, nodeRequestResource
 }
+
 func resourceTypesScore(requestsSourcesNum, allocatablesSourcesNum int64) int64 {
 	return (allocatablesSourcesNum - requestsSourcesNum) * framework.MaxNodeScore / allocatablesSourcesNum
 }
+
+func getPreScoreState(cycleState *framework.CycleState) (*preScoreState, error) {
+	c, err := cycleState.Read(preScoreStateKey)
+	if err != nil {
+		// preFilterState doesn't exist, likely PreFilter wasn't invoked.
+		return nil, fmt.Errorf("error reading %q from cycleState: %w", preScoreStateKey, err)
+	}
+
+	s, ok := c.(*preScoreState)
+	if !ok {
+		return nil, fmt.Errorf("%+v  convert to NodeResourcesFit.preFilterState error", c)
+	}
+	return s, nil
+}
diff --git a/pkg/scheduler/plugins/scarceresourceavoidance/scarce_resource_avoidance_test.go b/pkg/scheduler/plugins/scarceresourceavoidance/scarce_resource_avoidance_test.go
@@ -214,6 +214,12 @@ func TestPlugin_Score(t *testing.T) {
 			},
 		},
 	}
+
+	status := p.(framework.PreScorePlugin).PreScore(context.TODO(), cycleState, pod, nodes)
+	if status != nil {
+		t.Fatal("PreScore run err")
+	}
+
 	scoreNode1, _ := p.(*Plugin).Score(context.TODO(), cycleState, pod, "testNode1")
 	scoreNode2, _ := p.(*Plugin).Score(context.TODO(), cycleState, pod, "testNode2")
 	if scoreNode1 >= scoreNode2 {