Apply CPU startup boost in admission controller if its set

Author: kamarabbas99

vertical-pod-autoscaler/pkg/admission-controller/main.go

@@ -25,6 +25,7 @@ import (
 	"time"
 
 	"github.com/spf13/pflag"
+	"k8s.io/apimachinery/pkg/api/resource"
 	"k8s.io/client-go/informers"
 	kube_client "k8s.io/client-go/kubernetes"
 	typedadmregv1 "k8s.io/client-go/kubernetes/typed/admissionregistration/v1"
@@ -78,6 +79,7 @@ var (
 	registerWebhook      = flag.Bool("register-webhook", true, "If set to true, admission webhook object will be created on start up to register with the API server.")
 	webhookLabels        = flag.String("webhook-labels", "", "Comma separated list of labels to add to the webhook object. Format: key1:value1,key2:value2")
 	registerByURL        = flag.Bool("register-by-url", false, "If set to true, admission webhook will be registered by URL (webhookAddress:webhookPort) instead of by service name")
+	maxAllowedCpu        = flag.String("container-recommendation-max-allowed-cpu", "", "Maximum amount of CPU that will be recommended for a container.")
 )
 
 func main() {
@@ -93,6 +95,13 @@ func main() {
 		klog.FlushAndExit(klog.ExitFlushTimeout, 1)
 	}
 
+	if *maxAllowedCpu != "" {
+		if _, err := resource.ParseQuantity(*maxAllowedCpu); err != nil {
+			klog.ErrorS(err, "Failed to parse maxAllowedCpu")
+			klog.FlushAndExit(klog.ExitFlushTimeout, 1)
+		}
+	}
+
 	healthCheck := metrics.NewHealthCheck(time.Minute)
 	metrics_admission.Register()
 	server.Initialize(&commonFlags.EnableProfiling, healthCheck, address)
@@ -145,7 +154,7 @@ func main() {
 		hostname,
 	)
 
-	calculators := []patch.Calculator{patch.NewResourceUpdatesCalculator(recommendationProvider), patch.NewObservedContainersCalculator()}
+	calculators := []patch.Calculator{patch.NewResourceUpdatesCalculator(recommendationProvider, *maxAllowedCpu), patch.NewObservedContainersCalculator()}
 	as := logic.NewAdmissionServer(podPreprocessor, vpaPreprocessor, limitRangeCalculator, vpaMatcher, calculators)
 	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
 		as.Serve(w, r)

vertical-pod-autoscaler/pkg/admission-controller/resource/pod/patch/resource_updates.go

@@ -21,10 +21,14 @@ import (
 	"strings"
 
 	core "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	"k8s.io/klog/v2"
 
 	resource_admission "k8s.io/autoscaler/vertical-pod-autoscaler/pkg/admission-controller/resource"
 	"k8s.io/autoscaler/vertical-pod-autoscaler/pkg/admission-controller/resource/pod/recommendation"
 	vpa_types "k8s.io/autoscaler/vertical-pod-autoscaler/pkg/apis/autoscaling.k8s.io/v1"
+	"k8s.io/autoscaler/vertical-pod-autoscaler/pkg/features"
+	"k8s.io/autoscaler/vertical-pod-autoscaler/pkg/utils/annotations"
 	resourcehelpers "k8s.io/autoscaler/vertical-pod-autoscaler/pkg/utils/resources"
 	vpa_api_util "k8s.io/autoscaler/vertical-pod-autoscaler/pkg/utils/vpa"
 )
@@ -37,13 +41,19 @@ const (
 
 type resourcesUpdatesPatchCalculator struct {
 	recommendationProvider recommendation.Provider
+	maxAllowedCpu          resource.Quantity
 }
 
 // NewResourceUpdatesCalculator returns a calculator for
 // resource update patches.
-func NewResourceUpdatesCalculator(recommendationProvider recommendation.Provider) Calculator {
+func NewResourceUpdatesCalculator(recommendationProvider recommendation.Provider, maxAllowedCpu string) Calculator {
+	var maxAllowedCpuQuantity resource.Quantity
+	if maxAllowedCpu != "" {
+		maxAllowedCpuQuantity = resource.MustParse(maxAllowedCpu)
+	}
 	return &resourcesUpdatesPatchCalculator{
 		recommendationProvider: recommendationProvider,
+		maxAllowedCpu:          maxAllowedCpuQuantity,
 	}
 }
 
@@ -52,11 +62,22 @@ func (*resourcesUpdatesPatchCalculator) PatchResourceTarget() PatchResourceTarge
 }
 
 func (c *resourcesUpdatesPatchCalculator) CalculatePatches(pod *core.Pod, vpa *vpa_types.VerticalPodAutoscaler) ([]resource_admission.PatchRecord, error) {
+	klog.Infof("Calculating patches for pod %s/%s with VPA %s", pod.Namespace, pod.Name, vpa.Name)
 	result := []resource_admission.PatchRecord{}
 
 	containersResources, annotationsPerContainer, err := c.recommendationProvider.GetContainersResourcesForPod(pod, vpa)
 	if err != nil {
-		return []resource_admission.PatchRecord{}, fmt.Errorf("failed to calculate resource patch for pod %s/%s: %v", pod.Namespace, pod.Name, err)
+		return nil, fmt.Errorf("failed to calculate resource patch for pod %s/%s: %v", pod.Namespace, pod.Name, err)
+	}
+
+	if vpa_api_util.GetUpdateMode(vpa) == vpa_types.UpdateModeOff {
+		// If update mode is "Off", we don't want to apply any recommendations,
+		// but we still want to apply startup boost.
+		for i := range containersResources {
+			containersResources[i].Requests = nil
+			containersResources[i].Limits = nil
+		}
+		annotationsPerContainer = vpa_api_util.ContainerToAnnotationsMap{}
 	}
 
 	if annotationsPerContainer == nil {
@@ -65,9 +86,44 @@ func (c *resourcesUpdatesPatchCalculator) CalculatePatches(pod *core.Pod, vpa *v
 
 	updatesAnnotation := []string{}
 	for i, containerResources := range containersResources {
+		// Apply startup boost if configured
+		if features.Enabled(features.CPUStartupBoost) {
+			policy := vpa_api_util.GetContainerResourcePolicy(pod.Spec.Containers[i].Name, vpa.Spec.ResourcePolicy)
+			if policy != nil && policy.Mode != nil && *policy.Mode == vpa_types.ContainerScalingModeOff {
+				klog.V(4).Infof("Not applying startup boost for container %s since its scaling mode is Off", pod.Spec.Containers[i].Name)
+				continue
+			} else {
+				boost, err := getStartupBoost(&pod.Spec.Containers[i], vpa)
+				if err != nil {
+					return nil, err
+				}
+				if boost != nil {
+					if !c.maxAllowedCpu.IsZero() && boost.Cmp(c.maxAllowedCpu) > 0 {
+						cappedBoost := c.maxAllowedCpu
+						boost = &cappedBoost
+					}
+					if containerResources.Requests == nil {
+						containerResources.Requests = core.ResourceList{}
+					}
+					containerResources.Requests[core.ResourceCPU] = *boost
+					if containerResources.Limits == nil {
+						containerResources.Limits = core.ResourceList{}
+					}
+					containerResources.Limits[core.ResourceCPU] = *boost
+					originalResources, err := annotations.GetOriginalResourcesAnnotationValue(&pod.Spec.Containers[i])
+					if err != nil {
+						return nil, err
+					}
+					result = append(result, GetAddAnnotationPatch(annotations.StartupCPUBoostAnnotation, originalResources))
+				}
+			}
+		}
+
 		newPatches, newUpdatesAnnotation := getContainerPatch(pod, i, annotationsPerContainer, containerResources)
-		result = append(result, newPatches...)
-		updatesAnnotation = append(updatesAnnotation, newUpdatesAnnotation)
+		if len(newPatches) > 0 {
+			result = append(result, newPatches...)
+			updatesAnnotation = append(updatesAnnotation, newUpdatesAnnotation)
+		}
 	}
 
 	if len(updatesAnnotation) > 0 {
@@ -77,6 +133,48 @@ func (c *resourcesUpdatesPatchCalculator) CalculatePatches(pod *core.Pod, vpa *v
 	return result, nil
 }
 
+func getStartupBoost(container *core.Container, vpa *vpa_types.VerticalPodAutoscaler) (*resource.Quantity, error) {
+	policy := vpa_api_util.GetContainerResourcePolicy(container.Name, vpa.Spec.ResourcePolicy)
+	startupBoost := vpa.Spec.StartupBoost
+	if policy != nil && policy.StartupBoost != nil {
+		startupBoost = policy.StartupBoost
+	}
+	if startupBoost == nil {
+		return nil, nil
+	}
+
+	cpuRequest := container.Resources.Requests[core.ResourceCPU]
+	boostType := startupBoost.CPU.Type
+	if boostType == "" {
+		boostType = vpa_types.FactorStartupBoostType
+	}
+
+	switch boostType {
+	case vpa_types.FactorStartupBoostType:
+		if startupBoost.CPU.Factor == nil {
+			return nil, fmt.Errorf("startupBoost.CPU.Factor is required when Type is Factor or not specified")
+		}
+		factor := *startupBoost.CPU.Factor
+		if factor < 1 {
+			return nil, fmt.Errorf("boost factor must be >= 1")
+		}
+		boostedCPU := cpuRequest.MilliValue()
+		boostedCPU = int64(float64(boostedCPU) * float64(factor))
+		return resource.NewMilliQuantity(boostedCPU, resource.DecimalSI), nil
+	case vpa_types.QuantityStartupBoostType:
+		if startupBoost.CPU.Quantity == nil {
+			return nil, fmt.Errorf("startupBoost.CPU.Quantity is required when Type is Quantity")
+		}
+		quantity := *startupBoost.CPU.Quantity
+		if quantity.Cmp(cpuRequest) < 0 {
+			return nil, fmt.Errorf("boost quantity %s is less than container's request %s", quantity.String(), cpuRequest.String())
+		}
+		return &quantity, nil
+	default:
+		return nil, fmt.Errorf("unsupported startup boost type: %s", startupBoost.CPU.Type)
+	}
+}
+
 func getContainerPatch(pod *core.Pod, i int, annotationsPerContainer vpa_api_util.ContainerToAnnotationsMap, containerResources vpa_api_util.ContainerResources) ([]resource_admission.PatchRecord, string) {
 	var patches []resource_admission.PatchRecord
 	// Add empty resources object if missing.