feat: Add support for GPU sharingStrategy in Node Template

Author: kasia-kujawa

castai/resource_node_template.go

@@ -86,6 +86,7 @@ const (
 	FieldNodeTemplateSharingConfiguration                     = "sharing_configuration"
 	FieldNodeTemplateSharedClientsPerGpu                      = "shared_clients_per_gpu"
 	FieldNodeTemplateSharedGpuName                            = "gpu_name"
+	FieldNodeTemplateSharingStrategy                          = "sharing_strategy"
 	FieldNodeTemplateClmEnabled                               = "clm_enabled"
 	FieldNodeTemplateEdgeLocationIDs                          = "edge_location_ids"
 	FieldNodeTemplatePriceAdjustmentConfiguration             = "price_adjustment_configuration"
@@ -679,11 +680,18 @@ func resourceNodeTemplate() *schema.Resource {
 				DiffSuppressFunc: compareLists,
 				Elem: &schema.Resource{
 					Schema: map[string]*schema.Schema{
+						FieldNodeTemplateSharingStrategy: {
+							Type:             schema.TypeString,
+							Optional:         true,
+							ValidateDiagFunc: validation.ToDiagFunc(validation.StringInSlice([]string{"time-slicing", "mps"}, false)),
+							Description:      "GPU sharing strategy. Supported values: `time-slicing`, `mps`.",
+						},
 						FieldNodeTemplateEnableTimeSharing: {
 							Type:        schema.TypeBool,
 							Optional:    true,
 							Default:     nil,
-							Description: "Enable/disable GPU time-sharing.",
+							Deprecated:  "Use sharing_strategy instead.",
+							Description: "Enable/disable GPU time-sharing. Deprecated: use sharing_strategy = \"time-slicing\" instead.",
 						},
 						FieldNodeTemplateDefaultSharedClientsPerGpu: {
 							Type:        schema.TypeInt,
@@ -882,6 +890,10 @@ func flattenGpuSettings(g *sdk.NodetemplatesV1GPU) ([]map[string]any, error) {
 		out[FieldNodeTemplateUserManagedGPUDrivers] = g.UserManagedGpuDrivers
 	}
 
+	if g.SharingStrategy != nil {
+		out[FieldNodeTemplateSharingStrategy] = gpuSharingStrategyToTerraform(*g.SharingStrategy)
+	}
+
 	if g.EnableTimeSharing != nil {
 		out[FieldNodeTemplateEnableTimeSharing] = g.EnableTimeSharing
 	}
@@ -1175,6 +1187,7 @@ func updateNodeTemplate(ctx context.Context, d *schema.ResourceData, meta any, s
 		FieldNodeTemplateGpu,
 		FieldNodeTemplateDefaultSharedClientsPerGpu,
 		FieldNodeTemplateEnableTimeSharing,
+		FieldNodeTemplateSharingStrategy,
 		FieldNodeTemplateSharingConfiguration,
 		FieldNodeTemplateSharedGpuName,
 		FieldNodeTemplateSharedClientsPerGpu,
@@ -1702,6 +1715,12 @@ func toTemplateGpu(obj map[string]any) *sdk.NodetemplatesV1GPU {
 		defaultSharedClientsPerGpu = int32(v)
 	}
 
+	var sharingStrategy *sdk.NodetemplatesV1GPUSharingStrategy
+	if v, ok := obj[FieldNodeTemplateSharingStrategy].(string); ok && v != "" {
+		s := gpuSharingStrategyToAPI(v)
+		sharingStrategy = &s
+	}
+
 	var enableTimeSharing bool
 	if v, ok := obj[FieldNodeTemplateEnableTimeSharing].(bool); ok {
 		enableTimeSharing = v
@@ -1728,13 +1747,14 @@ 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 && !userManagedGPUDrivers {
+	if defaultSharedClientsPerGpu == 0 && !enableTimeSharing && sharingStrategy == nil && len(sharingConfig) == 0 && !userManagedGPUDrivers {
 		return nil
 	}
 
 	result := &sdk.NodetemplatesV1GPU{
 		EnableTimeSharing:    &enableTimeSharing,
 		SharingConfiguration: &sharingConfig,
+		SharingStrategy:      sharingStrategy,
 	}
 
 	// Only set DefaultSharedClientsPerGpu if it's non-zero to avoid API validation errors
@@ -1925,3 +1945,27 @@ func compareLists(key, oldValue, newValue string, d *schema.ResourceData) bool {
 	}
 	return false
 }
+
+// gpuSharingStrategyToAPI converts a terraform-friendly strategy string to the API enum value.
+func gpuSharingStrategyToAPI(s string) sdk.NodetemplatesV1GPUSharingStrategy {
+	switch s {
+	case "mps":
+		return sdk.GPUSHARINGSTRATEGYMPS
+	case "time-slicing":
+		return sdk.GPUSHARINGSTRATEGYTIMESLICING
+	default:
+		return sdk.GPUSHARINGSTRATEGYUNSPECIFIED
+	}
+}
+
+// gpuSharingStrategyToTerraform converts the API enum value to a terraform-friendly string.
+func gpuSharingStrategyToTerraform(s sdk.NodetemplatesV1GPUSharingStrategy) string {
+	switch s {
+	case sdk.GPUSHARINGSTRATEGYMPS:
+		return "mps"
+	case sdk.GPUSHARINGSTRATEGYTIMESLICING:
+		return "time-slicing"
+	default:
+		return ""
+	}
+}

castai/resource_node_template_test.go

@@ -273,6 +273,7 @@ edge_location_ids.1 = b2c3d4e5-f6a7-8901-bcde-f12345678901
 gpu.# = 1
 gpu.0.default_shared_clients_per_gpu = 10
 gpu.0.enable_time_sharing = true
+gpu.0.sharing_strategy = 
 gpu.0.user_managed_gpu_drivers = true
 gpu.0.sharing_configuration.# = 1
 gpu.0.sharing_configuration.0.gpu_name = A100
@@ -1334,6 +1335,45 @@ func Test_toTemplateGpu(t *testing.T) {
 				// UserManagedGpuDrivers not set (nil) - only sent when explicitly true (GKE-only)
 			},
 		},
+		{
+			name: "sharing_strategy mps",
+			input: map[string]any{
+				FieldNodeTemplateSharingStrategy:            "mps",
+				FieldNodeTemplateDefaultSharedClientsPerGpu: 4,
+			},
+			want: &sdk.NodetemplatesV1GPU{
+				DefaultSharedClientsPerGpu: lo.ToPtr(int32(4)),
+				EnableTimeSharing:          lo.ToPtr(false),
+				SharingConfiguration:       &map[string]sdk.NodetemplatesV1SharedGPU{},
+				SharingStrategy:            lo.ToPtr(sdk.GPUSHARINGSTRATEGYMPS),
+			},
+		},
+		{
+			name: "sharing_strategy time-slicing",
+			input: map[string]any{
+				FieldNodeTemplateSharingStrategy:            "time-slicing",
+				FieldNodeTemplateDefaultSharedClientsPerGpu: 8,
+			},
+			want: &sdk.NodetemplatesV1GPU{
+				DefaultSharedClientsPerGpu: lo.ToPtr(int32(8)),
+				EnableTimeSharing:          lo.ToPtr(false),
+				SharingConfiguration:       &map[string]sdk.NodetemplatesV1SharedGPU{},
+				SharingStrategy:            lo.ToPtr(sdk.GPUSHARINGSTRATEGYTIMESLICING),
+			},
+		},
+		{
+			name: "sharing_strategy mps with user_managed_gpu_drivers",
+			input: map[string]any{
+				FieldNodeTemplateSharingStrategy:       "mps",
+				FieldNodeTemplateUserManagedGPUDrivers: true,
+			},
+			want: &sdk.NodetemplatesV1GPU{
+				EnableTimeSharing:    lo.ToPtr(false),
+				SharingConfiguration: &map[string]sdk.NodetemplatesV1SharedGPU{},
+				SharingStrategy:      lo.ToPtr(sdk.GPUSHARINGSTRATEGYMPS),
+				UserManagedGpuDrivers: lo.ToPtr(true),
+			},
+		},
 	}
 
 	for _, tt := range tests {
@@ -1356,6 +1396,7 @@ func Test_toTemplateGpu(t *testing.T) {
 			}
 
 			require.Equal(t, tt.want.EnableTimeSharing, got.EnableTimeSharing)
+			require.Equal(t, tt.want.SharingStrategy, got.SharingStrategy)
 			require.Equal(t, tt.want.UserManagedGpuDrivers, got.UserManagedGpuDrivers)
 
 			if tt.want.SharingConfiguration != nil && got.SharingConfiguration != nil {
@@ -1471,6 +1512,34 @@ func Test_flattenGpuSettings(t *testing.T) {
 			},
 			wantErr: false,
 		},
+		{
+			name: "sharing_strategy mps",
+			input: &sdk.NodetemplatesV1GPU{
+				SharingStrategy:            lo.ToPtr(sdk.GPUSHARINGSTRATEGYMPS),
+				DefaultSharedClientsPerGpu: lo.ToPtr(int32(4)),
+			},
+			want: []map[string]any{
+				{
+					FieldNodeTemplateSharingStrategy:            "mps",
+					FieldNodeTemplateDefaultSharedClientsPerGpu: lo.ToPtr(int32(4)),
+				},
+			},
+			wantErr: false,
+		},
+		{
+			name: "sharing_strategy time-slicing",
+			input: &sdk.NodetemplatesV1GPU{
+				SharingStrategy:            lo.ToPtr(sdk.GPUSHARINGSTRATEGYTIMESLICING),
+				DefaultSharedClientsPerGpu: lo.ToPtr(int32(8)),
+			},
+			want: []map[string]any{
+				{
+					FieldNodeTemplateSharingStrategy:            "time-slicing",
+					FieldNodeTemplateDefaultSharedClientsPerGpu: lo.ToPtr(int32(8)),
+				},
+			},
+			wantErr: false,
+		},
 	}
 
 	for _, tt := range tests {
@@ -1491,6 +1560,7 @@ func Test_flattenGpuSettings(t *testing.T) {
 			// Compare all fields except sharing_configuration
 			require.Equal(t, tt.want[0][FieldNodeTemplateUserManagedGPUDrivers], got[0][FieldNodeTemplateUserManagedGPUDrivers])
 			require.Equal(t, tt.want[0][FieldNodeTemplateEnableTimeSharing], got[0][FieldNodeTemplateEnableTimeSharing])
+			require.Equal(t, tt.want[0][FieldNodeTemplateSharingStrategy], got[0][FieldNodeTemplateSharingStrategy])
 			require.Equal(t, tt.want[0][FieldNodeTemplateDefaultSharedClientsPerGpu], got[0][FieldNodeTemplateDefaultSharedClientsPerGpu])
 
 			// Compare sharing_configuration with ElementsMatch (order doesn't matter)