feat: add Node Template option to configure user-managed GPU drivers (#639)

* feat: add Node Template option to configure user-managed GPU drivers

* fix: only send user_managed_gpu_drivers when it is set to true

* fix(tests): handle autoscaler policy version conflicts with wait.Backoff retry

 When node template and autoscaler policies are updated concurrently
 in the same terraform apply, the API rejects updates with:
"node template has changed since the policies have been retrieved, refetch the policies and perform the update again"

* fix: remove additional try of updatePolicies()
fix: change constant name

Author: kasia-kujawa

castai/resource_autoscaler.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"log"
 	"net/http"
+	"strings"
 	"time"
 
 	jsonpatch "github.com/evanphx/json-patch"
@@ -17,6 +18,7 @@ import (
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
 	"github.com/mitchellh/mapstructure"
+	"k8s.io/apimachinery/pkg/util/wait"
 
 	"github.com/castai/terraform-provider-castai/castai/sdk"
 	"github.com/castai/terraform-provider-castai/castai/types"
@@ -805,23 +807,59 @@ func updateAutoscalerPolicies(ctx context.Context, data *schema.ResourceData, me
 		return nil
 	}
 
-	policies, err := getChangedPolicies(ctx, data, meta, clusterId)
-	if err != nil {
-		return err
+	// Define the update operation that will be executed with retry logic
+	updatePolicies := func() error {
+		policies, err := getChangedPolicies(ctx, data, meta, clusterId)
+		if err != nil {
+			return err
+		}
+
+		if policies == nil {
+			log.Printf("[DEBUG] changed policies json not calculated. Skipping autoscaler policies changes")
+			return nil
+		}
+
+		changedPoliciesJSON := string(policies)
+		if changedPoliciesJSON == "" {
+			log.Printf("[DEBUG] changed policies json not found. Skipping autoscaler policies changes")
+			return nil
+		}
+
+		return upsertPolicies(ctx, meta, clusterId, changedPoliciesJSON)
 	}
 
-	if policies == nil {
-		log.Printf("[DEBUG] changed policies json not calculated. Skipping autoscaler policies changes")
-		return nil
+	// Exponential backoff configuration
+	backoff := wait.Backoff{
+		Duration: 100 * time.Millisecond,
+		Factor:   2.0,
+		Jitter:   0.1,
+		Steps:    5,
+		Cap:      2 * time.Second,
 	}
 
-	changedPoliciesJSON := string(policies)
-	if changedPoliciesJSON == "" {
-		log.Printf("[DEBUG] changed policies json not found. Skipping autoscaler policies changes")
-		return nil
+	retryErr := wait.ExponentialBackoffWithContext(ctx, backoff, func(ctx context.Context) (done bool, err error) {
+		err = updatePolicies()
+		if err == nil {
+			return true, nil // Success - stop retrying
+		}
+
+		// Check if error is retryable
+		if !isNodeTemplateVersionConflict(err) {
+			return false, err // Non-retryable error - stop with error
+		}
+
+		log.Printf("[DEBUG] Retry failed with version conflict: %v", err)
+		return false, nil // Retryable error - continue retrying
+	})
+
+	if retryErr != nil {
+		if wait.Interrupted(retryErr) {
+			return fmt.Errorf("timeout waiting for autoscaler policy update after version conflicts: %w", retryErr)
+		}
+		return retryErr
 	}
 
-	return upsertPolicies(ctx, meta, clusterId, changedPoliciesJSON)
+	return nil
 }
 
 func upsertPolicies(ctx context.Context, meta interface{}, clusterId string, changedPoliciesJSON string) error {
@@ -835,6 +873,15 @@ func upsertPolicies(ctx context.Context, meta interface{}, clusterId string, cha
 	return nil
 }
 
+// isNodeTemplateVersionConflict checks if the error is due to version mismatch
+func isNodeTemplateVersionConflict(err error) bool {
+	if err == nil {
+		return false
+	}
+	errMsg := err.Error()
+	return strings.Contains(errMsg, "template has changed") || strings.Contains(errMsg, "refetch the policies")
+}
+
 func readAutoscalerPolicies(ctx context.Context, data *schema.ResourceData, meta interface{}) error {
 	log.Printf("[INFO] AUTOSCALER policies get call start")
 	defer log.Printf("[INFO] AUTOSCALER policies get call end")

castai/resource_node_template.go

@@ -90,6 +90,7 @@ const (
 	FieldNodeTemplateEdgeLocationIDs                          = "edge_location_ids"
 	FieldNodeTemplatePriceAdjustmentConfiguration             = "price_adjustment_configuration"
 	FieldNodeTemplateInstanceTypeAdjustments                  = "instance_type_adjustments"
+	FieldNodeTemplateUserManagedGPUDrivers                    = "user_managed_gpu_drivers"
 )
 
 const (
@@ -710,6 +711,12 @@ func resourceNodeTemplate() *schema.Resource {
 								},
 							},
 						},
+						FieldNodeTemplateUserManagedGPUDrivers: {
+							Type:        schema.TypeBool,
+							Optional:    true,
+							Default:     nil,
+							Description: "Enable/disable user-managed GPU drivers (for GKE clusters only).",
+						},
 					},
 				},
 			},
@@ -871,6 +878,10 @@ func flattenGpuSettings(g *sdk.NodetemplatesV1GPU) ([]map[string]any, error) {
 
 	out := make(map[string]any)
 
+	if g.UserManagedGpuDrivers != nil {
+		out[FieldNodeTemplateUserManagedGPUDrivers] = g.UserManagedGpuDrivers
+	}
+
 	if g.EnableTimeSharing != nil {
 		out[FieldNodeTemplateEnableTimeSharing] = g.EnableTimeSharing
 	}
@@ -1170,6 +1181,7 @@ func updateNodeTemplate(ctx context.Context, d *schema.ResourceData, meta any, s
 		FieldNodeTemplateClmEnabled,
 		FieldNodeTemplateEdgeLocationIDs,
 		FieldNodeTemplatePriceAdjustmentConfiguration,
+		FieldNodeTemplateUserManagedGPUDrivers,
 	) {
 		log.Printf("[INFO] Nothing to update in node template")
 		return nil
@@ -1680,6 +1692,11 @@ func toTemplateGpu(obj map[string]any) *sdk.NodetemplatesV1GPU {
 		return nil
 	}
 
+	var userManagedGPUDrivers bool
+	if v, ok := obj[FieldNodeTemplateUserManagedGPUDrivers].(bool); ok {
+		userManagedGPUDrivers = v
+	}
+
 	var defaultSharedClientsPerGpu int32
 	if v, ok := obj[FieldNodeTemplateDefaultSharedClientsPerGpu].(int); ok {
 		defaultSharedClientsPerGpu = int32(v)
@@ -1711,16 +1728,29 @@ func toTemplateGpu(obj map[string]any) *sdk.NodetemplatesV1GPU {
 	// terraform treats nil values as zero values
 	// this condition checks whether the whole gpu configuration is deleted
 	// and gpu configuration should be set to nil
-	if defaultSharedClientsPerGpu == 0 &&
-		!enableTimeSharing &&
-		len(sharingConfig) == 0 {
+	if defaultSharedClientsPerGpu == 0 && !enableTimeSharing && len(sharingConfig) == 0 && !userManagedGPUDrivers {
 		return nil
 	}
-	return &sdk.NodetemplatesV1GPU{
-		DefaultSharedClientsPerGpu: &defaultSharedClientsPerGpu,
-		EnableTimeSharing:          &enableTimeSharing,
-		SharingConfiguration:       &sharingConfig,
+
+	result := &sdk.NodetemplatesV1GPU{
+		EnableTimeSharing:    &enableTimeSharing,
+		SharingConfiguration: &sharingConfig,
+	}
+
+	// Only set DefaultSharedClientsPerGpu if it's non-zero to avoid API validation errors
+	// as terraform treats nil values as zero values
+	// API requires it to be in range (0, 48] if present
+	if defaultSharedClientsPerGpu > 0 {
+		result.DefaultSharedClientsPerGpu = &defaultSharedClientsPerGpu
 	}
+
+	// Only set UserManagedGpuDrivers if explicitly true
+	//as this is optional and as terraform treats nil values as zero values
+	if userManagedGPUDrivers {
+		result.UserManagedGpuDrivers = &userManagedGPUDrivers
+	}
+
+	return result
 }
 
 func toTemplateConstraintsInstanceFamilies(o map[string]any) *sdk.NodetemplatesV1TemplateConstraintsInstanceFamilyConstraints {