Allow keeping nodes for which scheduled rebalancer's drain phase has …

…timed out (#182)
castai · Jul 11, 2023 · cdbc10c · cdbc10c
1 parent 3700b51
commit cdbc10c
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 2 deletions.
diff --git a/castai/resource_rebalancing_schedule.go b/castai/resource_rebalancing_schedule.go
@@ -95,6 +95,11 @@ func resourceRebalancingSchedule() *schema.Resource {
 							ValidateDiagFunc: validation.ToDiagFunc(validation.IntAtLeast(1)),
 							Description:      "Minimum number of nodes that should be kept in the cluster after rebalancing.",
 						},
+						"keep_drain_timeout_nodes": {
+							Type:        schema.TypeBool,
+							Optional:    true,
+							Description: "Defines whether the nodes that failed to get drained until a predefined timeout, will be kept with a rebalancing.cast.ai/status=drain-failed annotation instead of forcefully drained.",
+						},
 						"execution_conditions": {
 							Type:     schema.TypeList,
 							MaxItems: 1,
@@ -243,6 +248,8 @@ func stateToSchedule(d *schema.ResourceData) (*sdk.ScheduledrebalancingV1Rebalan
 			return nil, fmt.Errorf("parsing selector: %w", err)
 		}
 
+		keepDrainTimeoutNodes := readOptionalValue[bool](launchConfigurationData, "keep_drain_timeout_nodes")
+
 		var executionConditions *sdk.ScheduledrebalancingV1ExecutionConditions
 		executionConditionsData := launchConfigurationData["execution_conditions"].([]any)
 		if len(executionConditionsData) != 0 {
@@ -256,8 +263,9 @@ func stateToSchedule(d *schema.ResourceData) (*sdk.ScheduledrebalancingV1Rebalan
 			NodeTtlSeconds:   readOptionalNumber[int, int32](launchConfigurationData, "node_ttl_seconds"),
 			NumTargetedNodes: readOptionalNumber[int, int32](launchConfigurationData, "num_targeted_nodes"),
 			RebalancingOptions: &sdk.ScheduledrebalancingV1RebalancingOptions{
-				MinNodes:            readOptionalNumber[int, int32](launchConfigurationData, "rebalancing_min_nodes"),
-				ExecutionConditions: executionConditions,
+				MinNodes:              readOptionalNumber[int, int32](launchConfigurationData, "rebalancing_min_nodes"),
+				KeepDrainTimeoutNodes: keepDrainTimeoutNodes,
+				ExecutionConditions:   executionConditions,
 			},
 			Selector: selector,
 		}
@@ -286,6 +294,7 @@ func scheduleToState(schedule *sdk.ScheduledrebalancingV1RebalancingSchedule, d
 
 	if schedule.LaunchConfiguration.RebalancingOptions != nil {
 		launchConfig["rebalancing_min_nodes"] = schedule.LaunchConfiguration.RebalancingOptions.MinNodes
+		launchConfig["keep_drain_timeout_nodes"] = schedule.LaunchConfiguration.RebalancingOptions.KeepDrainTimeoutNodes
 
 		executionConditions := schedule.LaunchConfiguration.RebalancingOptions.ExecutionConditions
 		if executionConditions != nil {

diff --git a/castai/resource_rebalancing_schedule_test.go b/castai/resource_rebalancing_schedule_test.go
@@ -82,6 +82,7 @@ resource "castai_rebalancing_schedule" "test" {
 		node_ttl_seconds = 10
 		num_targeted_nodes = 3
 		rebalancing_min_nodes = 2
+		keep_drain_timeout_nodes = true
 		selector = jsonencode({
 			nodeSelectorTerms = [{
 				matchExpressions = [

diff --git a/castai/sdk/api.gen.go b/castai/sdk/api.gen.go
diff --git a/docs/resources/rebalancing_schedule.md b/docs/resources/rebalancing_schedule.md
diff --git a/examples/eks/eks_cluster_autoscaler_polices/castai.tf b/examples/eks/eks_cluster_autoscaler_polices/castai.tf
@@ -200,6 +200,7 @@ resource "castai_rebalancing_schedule" "spots" {
 		node_ttl_seconds = 300
 		num_targeted_nodes = 3
 		rebalancing_min_nodes = 2
+		keep_drain_timeout_nodes = false
 		selector = jsonencode({
 			nodeSelectorTerms = [{
 				matchExpressions = [

diff --git a/examples/resources/castai_rebalancing_schedule/resource.tf b/examples/resources/castai_rebalancing_schedule/resource.tf
@@ -11,6 +11,7 @@ resource "castai_rebalancing_schedule" "spots" {
 		node_ttl_seconds = 300
 		num_targeted_nodes = 3
 		rebalancing_min_nodes = 2
+		keep_drain_timeout_nodes = false
 		selector = jsonencode({
 			nodeSelectorTerms = [{
 				matchExpressions = [