feat: [SKU modularization] remove sku_config from v1alpha1 and implement skuHandler interface

api/v1alpha1/workspace_validation.go

@@ -6,6 +6,7 @@
 import (
 	"context"
 	"fmt"
+	"os"
 	"reflect"
 	"regexp"
 	"strconv"
@@ -168,7 +169,7 @@
 	// Currently require a preset to specified, in future we can consider defining a template
 	if r.Preset == nil {
 		errs = errs.Also(apis.ErrMissingField("Preset"))
-	} else if presetName := string(r.Preset.Name); !isValidPreset(presetName) {
+	} else if presetName := string(r.Preset.Name); !utils.IsValidPreset(presetName) {
 		errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unsupported tuning preset name %s", presetName), "presetName"))
 	}
 	return errs
@@ -295,8 +296,15 @@
 	}
 	instanceType := string(r.InstanceType)
 
-	// Check if instancetype exists in our SKUs map
-	if skuConfig, exists := SupportedGPUConfigs[instanceType]; exists {
+	skuHandler, err := utils.GetSKUHandler()
+	if err != nil {
+		errs = errs.Also(apis.ErrGeneric(fmt.Sprintf("Failed to get SKU handler: %v", err), "instanceType"))
+		return errs
+	}
+	gpuConfigs := skuHandler.GetGPUConfigs()
+
+	// Check if instancetype exists in our SKUs map for the particular cloud provider
+	if skuConfig, exists := gpuConfigs[instanceType]; exists {
 		if presetName != "" {
 			model := plugin.KaitoModelRegister.MustGet(presetName) // InferenceSpec has been validated so the name is valid.
 
@@ -350,9 +358,10 @@
 			}
 		}
 	} else {
-		// Check for other instance types pattern matches
-		if !strings.HasPrefix(instanceType, N_SERIES_PREFIX) && !strings.HasPrefix(instanceType, D_SERIES_PREFIX) {
-			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unsupported instance type %s. Supported SKUs: %s", instanceType, getSupportedSKUs()), "instanceType"))
+		provider := os.Getenv("CLOUD_PROVIDER")
+		// Check for other instance types pattern matches if cloud provider is Azure
+		if provider != consts.AzureCloudName || (!strings.HasPrefix(instanceType, N_SERIES_PREFIX) && !strings.HasPrefix(instanceType, D_SERIES_PREFIX)) {
+			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unsupported instance type %s. Supported SKUs: %s", instanceType, skuHandler.GetSupportedSKUs()), "instanceType"))
 		}
 	}
 
@@ -398,7 +407,7 @@
 	if i.Preset != nil {
 		presetName := string(i.Preset.Name)
 		// Validate preset name
-		if !isValidPreset(presetName) {
+		if !utils.IsValidPreset(presetName) {
 			errs = errs.Also(apis.ErrInvalidValue(fmt.Sprintf("Unsupported inference preset name %s", presetName), "presetName"))
 		}
 		// Validate private preset has private image specified

api/v1alpha1/workspace_validation_test.go

@@ -7,8 +7,6 @@ import (
 	"context"
 	"fmt"
 	"os"
-	"reflect"
-	"sort"
 	"strings"
 	"testing"
 
@@ -234,7 +232,7 @@ func TestResourceSpecValidateCreate(t *testing.T) {
 		{
 			name: "Insufficient total GPU memory",
 			resourceSpec: &ResourceSpec{
-				InstanceType: "Standard_NC6",
+				InstanceType: "Standard_NV6",
 				Count:        pointerToInt(1),
 			},
 			modelGPUCount:       "1",
@@ -263,7 +261,7 @@ func TestResourceSpecValidateCreate(t *testing.T) {
 		{
 			name: "Insufficient per GPU memory",
 			resourceSpec: &ResourceSpec{
-				InstanceType: "Standard_NC6",
+				InstanceType: "Standard_NV6",
 				Count:        pointerToInt(2),
 			},
 			modelGPUCount:       "1",
@@ -320,7 +318,7 @@ func TestResourceSpecValidateCreate(t *testing.T) {
 		{
 			name: "Tuning validation with single node",
 			resourceSpec: &ResourceSpec{
-				InstanceType: "Standard_NC6",
+				InstanceType: "Standard_NC6s_v3",
 				Count:        pointerToInt(1),
 			},
 			errContent:     "",
@@ -330,7 +328,7 @@ func TestResourceSpecValidateCreate(t *testing.T) {
 		{
 			name: "Tuning validation with multinode",
 			resourceSpec: &ResourceSpec{
-				InstanceType: "Standard_NC6",
+				InstanceType: "Standard_NC6s_v3",
 				Count:        pointerToInt(2),
 			},
 			errContent:     "Tuning does not currently support multinode configurations",
@@ -339,6 +337,8 @@ func TestResourceSpecValidateCreate(t *testing.T) {
 		},
 	}
 
+	os.Setenv("CLOUD_PROVIDER", consts.AzureCloudName)
+
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			if tc.validateTuning {
@@ -1438,49 +1438,3 @@ func TestDataDestinationValidateUpdate(t *testing.T) {
 		})
 	}
 }
-
-func TestGetSupportedSKUs(t *testing.T) {
-	tests := []struct {
-		name           string
-		gpuConfigs     map[string]GPUConfig
-		expectedResult []string // changed to a slice for deterministic ordering
-	}{
-		{
-			name:           "no SKUs supported",
-			gpuConfigs:     map[string]GPUConfig{},
-			expectedResult: []string{""},
-		},
-		{
-			name: "one SKU supported",
-			gpuConfigs: map[string]GPUConfig{
-				"Standard_NC6": {SKU: "Standard_NC6"},
-			},
-			expectedResult: []string{"Standard_NC6"},
-		},
-		{
-			name: "multiple SKUs supported",
-			gpuConfigs: map[string]GPUConfig{
-				"Standard_NC6":  {SKU: "Standard_NC6"},
-				"Standard_NC12": {SKU: "Standard_NC12"},
-			},
-			expectedResult: []string{"Standard_NC6", "Standard_NC12"},
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.name, func(t *testing.T) {
-			SupportedGPUConfigs = tc.gpuConfigs
-
-			resultSlice := strings.Split(getSupportedSKUs(), ", ")
-			sort.Strings(resultSlice)
-
-			// Sort the expectedResult for comparison
-			expectedResultSlice := tc.expectedResult
-			sort.Strings(expectedResultSlice)
-
-			if !reflect.DeepEqual(resultSlice, expectedResultSlice) {
-				t.Errorf("getSupportedSKUs() = %v, want %v", resultSlice, expectedResultSlice)
-			}
-		})
-	}
-}

pkg/tuning/preset-tuning.go

@@ -55,7 +55,10 @@ var (
 )
 
 func getInstanceGPUCount(sku string) int {
-	gpuConfig, exists := kaitov1alpha1.SupportedGPUConfigs[sku]
+	skuHandler, _ := utils.GetSKUHandler()
+	gpuConfigs := skuHandler.GetGPUConfigs()
+
+	gpuConfig, exists := gpuConfigs[sku]
 	if !exists {
 		return 1
 	}

pkg/tuning/preset-tuning_test.go

@@ -24,13 +24,6 @@ import (
 	"k8s.io/utils/pointer"
 )
 
-// Mocking the SupportedGPUConfigs to be used in test scenarios.
-var mockSupportedGPUConfigs = map[string]kaitov1alpha1.GPUConfig{
-	"sku1": {GPUCount: 2},
-	"sku2": {GPUCount: 4},
-	"sku3": {GPUCount: 0},
-}
-
 func normalize(s string) string {
 	return strings.Join(strings.Fields(s), " ")
 }
@@ -54,18 +47,19 @@ func saveEnv(key string) func() {
 }
 
 func TestGetInstanceGPUCount(t *testing.T) {
-	kaitov1alpha1.SupportedGPUConfigs = mockSupportedGPUConfigs
+	os.Setenv("CLOUD_PROVIDER", consts.AzureCloudName)
+
 	testcases := map[string]struct {
 		sku              string
 		expectedGPUCount int
 	}{
 		"SKU Exists With Multiple GPUs": {
-			sku:              "sku1",
-			expectedGPUCount: 2,
+			sku:              "Standard_NC24s_v3",
+			expectedGPUCount: 4,
 		},
-		"SKU Exists With Zero GPUs": {
-			sku:              "sku3",
-			expectedGPUCount: 0,
+		"SKU Exists With One GPU": {
+			sku:              "Standard_NC6s_v3",
+			expectedGPUCount: 1,
 		},
 		"SKU Does Not Exist": {
 			sku:              "sku_unknown",

pkg/utils/common-preset.go

@@ -3,6 +3,7 @@
 package utils
 
 import (
+	"github.com/azure/kaito/pkg/utils/plugin"
 	corev1 "k8s.io/api/core/v1"
 )
 
@@ -149,3 +150,7 @@
 	}
 	return volume, volumeMount
 }
+
+func IsValidPreset(preset string) bool {
+	return plugin.KaitoModelRegister.Has(preset)
+}