Merge pull request #5974 from azylinski/add-estimator-PostPackingFunc

Add PostPackingFunc to be run at the end of the estimation logic
kubernetes · Jul 26, 2023 · 63eab4e · 63eab4e
2 parents 8c749fc + 21229d3
commit 63eab4e
Show file tree

Hide file tree

Showing 5 changed files with 26 additions and 14 deletions.
diff --git a/cluster-autoscaler/core/autoscaler.go b/cluster-autoscaler/core/autoscaler.go
@@ -127,6 +127,7 @@ func initializeDefaultOptions(opts *AutoscalerOptions) error {
 			opts.EstimatorName,
 			estimator.NewThresholdBasedEstimationLimiter(thresholds),
 			estimator.NewDecreasingPodOrderer(),
+			/* EstimationAnalyserFunc */ nil,
 		)
 		if err != nil {
 			return err

diff --git a/cluster-autoscaler/core/test/common.go b/cluster-autoscaler/core/test/common.go
@@ -212,6 +212,7 @@ func NewScaleTestAutoscalingContext(
 		options.EstimatorName,
 		estimator.NewThresholdBasedEstimationLimiter(nil),
 		estimator.NewDecreasingPodOrderer(),
+		/* EstimationAnalyserFunc */ nil,
 	)
 	predicateChecker, err := predicatechecker.NewTestPredicateChecker()
 	if err != nil {

diff --git a/cluster-autoscaler/estimator/binpacking_estimator.go b/cluster-autoscaler/estimator/binpacking_estimator.go
@@ -30,11 +30,12 @@ import (
 
 // BinpackingNodeEstimator estimates the number of needed nodes to handle the given amount of pods.
 type BinpackingNodeEstimator struct {
-	predicateChecker predicatechecker.PredicateChecker
-	clusterSnapshot  clustersnapshot.ClusterSnapshot
-	limiter          EstimationLimiter
-	podOrderer       EstimationPodOrderer
-	context          EstimationContext
+	predicateChecker       predicatechecker.PredicateChecker
+	clusterSnapshot        clustersnapshot.ClusterSnapshot
+	limiter                EstimationLimiter
+	podOrderer             EstimationPodOrderer
+	context                EstimationContext
+	estimationAnalyserFunc EstimationAnalyserFunc // optional
 }
 
 // NewBinpackingNodeEstimator builds a new BinpackingNodeEstimator.
@@ -44,13 +45,15 @@ func NewBinpackingNodeEstimator(
 	limiter EstimationLimiter,
 	podOrderer EstimationPodOrderer,
 	context EstimationContext,
+	estimationAnalyserFunc EstimationAnalyserFunc,
 ) *BinpackingNodeEstimator {
 	return &BinpackingNodeEstimator{
-		predicateChecker: predicateChecker,
-		clusterSnapshot:  clusterSnapshot,
-		limiter:          limiter,
-		podOrderer:       podOrderer,
-		context:          context,
+		predicateChecker:       predicateChecker,
+		clusterSnapshot:        clusterSnapshot,
+		limiter:                limiter,
+		podOrderer:             podOrderer,
+		context:                context,
+		estimationAnalyserFunc: estimationAnalyserFunc,
 	}
 }
 
@@ -141,6 +144,11 @@ func (e *BinpackingNodeEstimator) Estimate(
 			scheduledPods = append(scheduledPods, pod)
 		}
 	}
+
+	if e.estimationAnalyserFunc != nil {
+		e.estimationAnalyserFunc(e.clusterSnapshot, nodeGroup, newNodesWithPods)
+	}
+
 	return len(newNodesWithPods), scheduledPods
 }
 

diff --git a/cluster-autoscaler/estimator/binpacking_estimator_test.go b/cluster-autoscaler/estimator/binpacking_estimator_test.go
@@ -187,8 +187,7 @@ func TestBinpackingEstimate(t *testing.T) {
 			assert.NoError(t, err)
 			limiter := NewThresholdBasedEstimationLimiter([]Threshold{NewStaticThreshold(tc.maxNodes, time.Duration(0))})
 			processor := NewDecreasingPodOrderer()
-			estimator := NewBinpackingNodeEstimator(predicateChecker, clusterSnapshot, limiter, processor, nil)
-
+			estimator := NewBinpackingNodeEstimator(predicateChecker, clusterSnapshot, limiter, processor, nil /* EstimationContext */, nil /* EstimationAnalyserFunc */)
 			node := makeNode(tc.millicores, tc.memory, "template", "zone-mars")
 			nodeInfo := schedulerframework.NewNodeInfo()
 			nodeInfo.SetNode(node)

diff --git a/cluster-autoscaler/estimator/estimator.go b/cluster-autoscaler/estimator/estimator.go
@@ -45,15 +45,18 @@ type Estimator interface {
 // EstimatorBuilder creates a new estimator object.
 type EstimatorBuilder func(predicatechecker.PredicateChecker, clustersnapshot.ClusterSnapshot, EstimationContext) Estimator
 
+// EstimationAnalyserFunc to be run at the end of the estimation logic.
+type EstimationAnalyserFunc func(clustersnapshot.ClusterSnapshot, cloudprovider.NodeGroup, map[string]bool)
+
 // NewEstimatorBuilder creates a new estimator object from flag.
-func NewEstimatorBuilder(name string, limiter EstimationLimiter, orderer EstimationPodOrderer) (EstimatorBuilder, error) {
+func NewEstimatorBuilder(name string, limiter EstimationLimiter, orderer EstimationPodOrderer, estimationAnalyserFunc EstimationAnalyserFunc) (EstimatorBuilder, error) {
 	switch name {
 	case BinpackingEstimatorName:
 		return func(
 			predicateChecker predicatechecker.PredicateChecker,
 			clusterSnapshot clustersnapshot.ClusterSnapshot,
 			context EstimationContext) Estimator {
-			return NewBinpackingNodeEstimator(predicateChecker, clusterSnapshot, limiter, orderer, context)
+			return NewBinpackingNodeEstimator(predicateChecker, clusterSnapshot, limiter, orderer, context, estimationAnalyserFunc)
 		}, nil
 	}
 	return nil, fmt.Errorf("unknown estimator: %s", name)