feat: Enable predictive scaling via scaling policy

Author: zrudzionis

castai/resource_workload_scaling_policy.go

@@ -36,6 +36,7 @@ const (
 	FieldLimitStrategyType                         = "type"
 	FieldLimitStrategyMultiplier                   = "multiplier"
 	FieldConfidence                                = "confidence"
+	FieldPredictiveScaling                         = "predictive_scaling"
 	FieldConfidenceThreshold                       = "threshold"
 	DeprecatedFieldApplyThreshold                  = "apply_threshold"
 	FieldApplyThresholdStrategy                    = "apply_threshold_strategy"
@@ -257,6 +258,16 @@ It can be either:
 					},
 				},
 			},
+			FieldPredictiveScaling: {
+				Type:     schema.TypeList,
+				Optional: true,
+				MaxItems: 1,
+				Elem: &schema.Resource{
+					Schema: map[string]*schema.Schema{
+						FieldCPU: getPredictiveScalingResourceSchema(),
+					},
+				},
+			},
 		},
 		Timeouts: &schema.ResourceTimeout{
 			Create: schema.DefaultTimeout(15 * time.Second),
@@ -267,6 +278,24 @@ It can be either:
 	}
 }
 
+func getPredictiveScalingResourceSchema() *schema.Schema {
+	return &schema.Schema{
+		Type:        schema.TypeList,
+		Optional:    true,
+		MaxItems:    1,
+		Description: "Defines predictive scaling resource configuration.",
+		Elem: &schema.Resource{
+			Schema: map[string]*schema.Schema{
+				FieldEnabled: {
+					Type:        schema.TypeBool,
+					Required:    true,
+					Description: "Defines if predictive scaling is enabled for resource.",
+				},
+			},
+		},
+	}
+}
+
 func k8sLabelExpressionsSchema() *schema.Schema {
 	return &schema.Schema{
 		Type:        schema.TypeList,
@@ -513,6 +542,12 @@ func resourceWorkloadScalingPolicyCreate(ctx context.Context, d *schema.Resource
 
 	req.RecommendationPolicies.AntiAffinity = toAntiAffinity(toSection(d, "anti_affinity"))
 
+	ps, err := toPredictiveScaling(toSection(d, FieldPredictiveScaling))
+	if err != nil {
+		return diag.FromErr(err)
+	}
+	req.RecommendationPolicies.PredictiveScaling = ps
+
 	ar, err := toAssignmentRules(toSection(d, FieldAssignmentRules))
 	if err != nil {
 		return diag.FromErr(err)
@@ -593,6 +628,9 @@ func resourceWorkloadScalingPolicyRead(ctx context.Context, d *schema.ResourceDa
 	if err := d.Set("anti_affinity", toAntiAffinityMap(sp.RecommendationPolicies.AntiAffinity)); err != nil {
 		return diag.FromErr(fmt.Errorf("setting anti-affinity: %w", err))
 	}
+	if err := d.Set("predictive_scaling", toPredictiveScalingMap(sp.RecommendationPolicies.PredictiveScaling)); err != nil {
+		return diag.FromErr(fmt.Errorf("setting predictive scaling: %w", err))
+	}
 
 	if err := d.Set(FieldAssignmentRules, toAssignmentRulesMap(getResourceFrom(d, FieldAssignmentRules), sp.AssignmentRules)); err != nil {
 		return diag.FromErr(fmt.Errorf("setting assignment rules: %w", err))
@@ -640,20 +678,25 @@ func resourceWorkloadScalingPolicyUpdate(ctx context.Context, d *schema.Resource
 	if err != nil {
 		return diag.FromErr(err)
 	}
+	predictiveScaling, err := toPredictiveScaling(toSection(d, FieldPredictiveScaling))
+	if err != nil {
+		return diag.FromErr(err)
+	}
 
 	req := sdk.WorkloadOptimizationAPIUpdateWorkloadScalingPolicyJSONRequestBody{
 		Name:            d.Get("name").(string),
 		ApplyType:       sdk.WorkloadoptimizationV1ApplyType(d.Get("apply_type").(string)),
 		AssignmentRules: ar,
 		RecommendationPolicies: sdk.WorkloadoptimizationV1RecommendationPolicies{
-			ManagementOption: sdk.WorkloadoptimizationV1ManagementOption(d.Get("management_option").(string)),
-			Cpu:              cpu,
-			Memory:           memory,
-			Startup:          toStartup(toSection(d, "startup")),
-			Downscaling:      toDownscaling(toSection(d, "downscaling")),
-			MemoryEvent:      toMemoryEvent(toSection(d, "memory_event")),
-			AntiAffinity:     toAntiAffinity(toSection(d, "anti_affinity")),
-			Confidence:       toConfidence(toSection(d, FieldConfidence)),
+			ManagementOption:  sdk.WorkloadoptimizationV1ManagementOption(d.Get("management_option").(string)),
+			Cpu:               cpu,
+			Memory:            memory,
+			Startup:           toStartup(toSection(d, "startup")),
+			Downscaling:       toDownscaling(toSection(d, "downscaling")),
+			MemoryEvent:       toMemoryEvent(toSection(d, "memory_event")),
+			AntiAffinity:      toAntiAffinity(toSection(d, "anti_affinity")),
+			Confidence:        toConfidence(toSection(d, FieldConfidence)),
+			PredictiveScaling: predictiveScaling,
 		},
 	}
 
@@ -1161,6 +1204,57 @@ func toAntiAffinityMap(s *sdk.WorkloadoptimizationV1AntiAffinitySettings) []map[
 	return []map[string]any{m}
 }
 
+func toPredictiveScalingMap(s *sdk.WorkloadoptimizationV1PredictiveScalingSettings) []map[string]any {
+	if s == nil {
+		return nil
+	}
+
+	return []map[string]any{
+		{
+			FieldCPU: toPredictiveScalingResourceMap(s.Cpu),
+		},
+	}
+}
+
+func toPredictiveScaling(m map[string]any) (*sdk.WorkloadoptimizationV1PredictiveScalingSettings, error) {
+	if len(m) == 0 {
+		return nil, nil
+	}
+
+	cpuCfg := getFirstElem(m, FieldCPU)
+	cpuResource, err := toPredictiveScalingResource(cpuCfg)
+	if err != nil {
+		return nil, fmt.Errorf("field %q: field %q: %w", FieldPredictiveScaling, FieldCPU, err)
+	}
+
+	return &sdk.WorkloadoptimizationV1PredictiveScalingSettings{
+		Cpu: *cpuResource,
+	}, nil
+}
+
+func toPredictiveScalingResourceMap(s sdk.WorkloadoptimizationV1PredictiveScaling) []map[string]any {
+	return []map[string]any{
+		{
+			FieldEnabled: s.Enabled,
+		},
+	}
+}
+
+func toPredictiveScalingResource(m map[string]any) (*sdk.WorkloadoptimizationV1PredictiveScaling, error) {
+	if len(m) == 0 {
+		return nil, nil
+	}
+
+	enabled, err := mustGetValue[bool](m, FieldEnabled)
+	if err != nil {
+		return nil, err
+	}
+
+	return &sdk.WorkloadoptimizationV1PredictiveScaling{
+		Enabled: *enabled,
+	}, nil
+}
+
 func getWorkloadScalingPolicyByName(ctx context.Context, client sdk.ClientWithResponsesInterface, clusterID, name string) (*sdk.WorkloadoptimizationV1WorkloadScalingPolicy, error) {
 	list, err := client.WorkloadOptimizationAPIListWorkloadScalingPoliciesWithResponse(ctx, clusterID)
 	if checkErr := sdk.CheckOKResponse(list, err); checkErr != nil {

castai/resource_workload_scaling_policy_test.go

@@ -105,6 +105,7 @@ func TestAccResourceWorkloadScalingPolicy(t *testing.T) {
 					resource.TestCheckResourceAttr(resourceName, "assignment_rules.0.rules.1.workload.0.gvk.0", "DaemonSet"),
 					resource.TestCheckResourceAttr(resourceName, "assignment_rules.0.rules.1.workload.0.labels_expressions.0.key", "helm.sh/chart"),
 					resource.TestCheckResourceAttr(resourceName, "assignment_rules.0.rules.1.workload.0.labels_expressions.0.operator", "DoesNotExist"),
+					resource.TestCheckResourceAttr(resourceName, "predictive_scaling.0.cpu.0.enabled", "true"),
 				),
 			},
 		},
@@ -213,6 +214,11 @@ func scalingPolicyConfigUpdated(clusterName, projectID, name string) string {
 				}
 			}
 		}
+		predictive_scaling {
+			cpu {
+				enabled = true
+			}
+		}
 		cpu {
 			function 		= "QUANTILE"
 			overhead 		= 0.15