Repository: kaito-project/kaito

Issue:
Reserved enough vram buffer for workload
**Is your feature request related to a problem? Please describe.**

Besides pre-allocated kvcache, vllm will allocate extra vram during inference. If available vram is not enough, pod will crash due to OOM.

**Describe the solution you'd like**

Make sure there is enough vram buffer (1.5GB) in the controller side. 

**Describe alternatives you've considered**

**Additional context**

Discussion:

Solution patch:
diff --git a/api/v1alpha1/inference_config_validation.go b/api/v1alpha1/inference_config_validation.go
index 9630aa9..57f927b 100644
--- a/api/v1alpha1/inference_config_validation.go
+++ b/api/v1alpha1/inference_config_validation.go
@@ -57,10 +57,9 @@ func (i *InferenceSpec) validateConfigMap(ctx context.Context, namespace, cmName
 		return apis.ErrGeneric(fmt.Sprintf("Failed to get SKU handler: %v", err), "instanceType")
 	}
 
-	gpuConfigs := skuHandler.GetGPUConfigs()
-	if skuConfig, exists := gpuConfigs[instanceType]; exists {
+	if skuConfig := skuHandler.GetGPUConfigBySKU(instanceType); skuConfig != nil {
 		// Check if this is a multi-GPU instance with less than 20GB per GPU
-		gpuMemPerGPU := skuConfig.GPUMem / skuConfig.GPUCount
+		gpuMemPerGPU := skuConfig.GPUMemGB / skuConfig.GPUCount
 		if skuConfig.GPUCount > 1 && gpuMemPerGPU < 20 {
 			// For multi-GPU instances with less than 20GB per GPU, max-model-len is required
 			maxModelLen, exists := inferenceConfig.VLLM["max-model-len"]
diff --git a/api/v1alpha1/ragengine_validation.go b/api/v1alpha1/ragengine_validation.go
index b162e4b..05c0703 100644
--- a/api/v1alpha1/ragengine_validation.go
+++ b/api/v1alpha1/ragengine_validation.go
@@ -77,9 +77,8 @@ func (r *ResourceSpec) validateRAGCreate() (errs *apis.FieldError) {
 		errs = errs.Also(apis.ErrGeneric(fmt.Sprintf("Failed to get SKU handler: %v", err), "instanceType"))
 		return errs
 	}
-	gpuConfigs := skuHandler.GetGPUConfigs()
 
-	if _, exists := gpuConfigs[instanceType]; !exists {
+	if skuConfig := skuHandler.GetGPUConfigBySKU(instanceType); skuConfig == nil {
 		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)) {
diff --git a/api/v1alpha1/workspace_validation.go b/api/v1alpha1/workspace_validation.go
index 27cb3ef..1855acb 100644
--- a/api/v1alpha1/workspace_validation.go
+++ b/api/v1alpha1/workspace_validation.go
@@ -19,6 +19,7 @@ import (
 	"k8s.io/klog/v2"
 	"knative.dev/pkg/apis"
 
+	"github.com/kaito-project/kaito/pkg/model"
 	"github.com/kaito-project/kaito/pkg/utils"
 	"github.com/kaito-project/kaito/pkg/utils/consts"
 	"github.com/kaito-project/kaito/pkg/utils/plugin"
@@ -59,9 +60,10 @@ func (w *Workspace) Validate(ctx context.Context) (errs *apis.FieldError) {
 				}
 			}
 
+			runtime := GetWorkspaceRuntimeName(w)
 			// TODO: Add Adapter Spec Validation - Including DataSource Validation for Adapter
 			errs = errs.Also(w.Resource.validateCreateWithInference(w.Inference, bypassResourceChecks).ViaField("resource"),
-				w.Inference.validateCreate(ctx, w.Namespace, w.Resource.InstanceType).ViaField("inference"))
+				w.Inference.validateCreate(ctx, w.Namespace, w.Resource.InstanceType, runtime).ViaField("inference"))
 		}
 		if w.Tuning != nil {
 			// TODO: Add validate resource based on Tuning Spec
@@ -308,17 +310,16 @@ func (r *ResourceSpec) validateCreateWithInference(inference *InferenceSpec, byp
 		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 skuConfig := skuHandler.GetGPUConfigBySKU(instanceType); skuConfig != nil {
 		if presetName != "" {
 			model := plugin.KaitoModelRegister.MustGet(presetName) // InferenceSpec has been validated so the name is valid.
 
 			machineCount := *r.Count
 			machineTotalNumGPUs := resource.NewQuantity(int64(machineCount*skuConfig.GPUCount), resource.DecimalSI)
-			machinePerGPUMemory := resource.NewQuantity(int64(skuConfig.GPUMem/skuConfig.GPUCount)*consts.GiBToBytes, resource.BinarySI) // Ensure it's per GPU
-			machineTotalGPUMem := resource.NewQuantity(int64(machineCount*skuConfig.GPUMem)*consts.GiBToBytes, resource.BinarySI)        // Total GPU memory
+			machinePerGPUMemory := resource.NewQuantity(int64(skuConfig.GPUMemGB/skuConfig.GPUCount)*consts.GiBToBytes, resource.BinarySI) // Ensure it's per GPU
+			machineTotalGPUMem := resource.NewQuantity(int64(machineCount*skuConfig.GPUMemGB)*consts.GiBToBytes, resource.BinarySI)        // Total GPU memory
 
 			modelGPUCount := resource.MustParse(model.GetInferenceParameters().GPUCountRequirement)
 			modelPerGPUMemory := resource.MustParse(model.GetInferenceParameters().PerGPUMemoryRequirement)
@@ -415,7 +416,7 @@ func (r *ResourceSpec) validateUpdate(old *ResourceSpec) (errs *apis.FieldError)
 	return errs
 }
 
-func (i *InferenceSpec) validateCreate(ctx context.Context, namespace string, instanceType string) (errs *apis.FieldError) {
+func (i *InferenceSpec) validateCreate(ctx context.Context, namespace string, instanceType string, runtime model.RuntimeName) (errs *apis.FieldError) {
 	// Check if both Preset and Template are not set
 	if i.Preset == nil && i.Template == nil {
 		errs = errs.Also(apis.ErrMissingField("Preset or Template must be specified"))
@@ -434,11 +435,19 @@ func (i *InferenceSpec) validateCreate(ctx context.Context, namespace string, in
 			// Need to return here. Otherwise, a panic will be hit when doing following checks.
 			return errs
 		}
+		modelPreset := plugin.KaitoModelRegister.MustGet(string(i.Preset.Name))
 		// Validate private preset has private image specified
-		if plugin.KaitoModelRegister.MustGet(string(i.Preset.Name)).GetInferenceParameters().ImageAccessMode == string(ModelImageAccessModePrivate) &&
+		if modelPreset.GetInferenceParameters().ImageAccessMode == string(ModelImageAccessModePrivate) &&
 			i.Preset.AccessMode != ModelImageAccessModePrivate {
 			errs = errs.Also(apis.ErrGeneric("This preset only supports private AccessMode, AccessMode must be private to continue"))
 		}
+		err := modelPreset.GetInferenceParameters().Validate(model.RuntimeContext{
+			RuntimeName: runtime,
+			UseAdapters: len(i.Adapters) > 0,
+		})
+		if err != nil {
+			errs = errs.Also(apis.ErrGeneric(fmt.Sprintf("Runtime validation: %v", err)))
+		}
 		// Additional validations for Preset
 		if i.Preset.AccessMode == ModelImageAccessModePrivate && i.Preset.Image == "" {
 			errs = errs.Also(apis.ErrGeneric("When AccessMode is private, an image must be provided in PresetOptions"))
diff --git a/api/v1beta1/inference_config_validation.go b/api/v1beta1/inference_config_validation.go
index d854947..973a36c 100644
--- a/api/v1beta1/inference_config_validation.go
+++ b/api/v1beta1/inference_config_validation.go
@@ -57,10 +57,9 @@ func (i *InferenceSpec) validateConfigMap(ctx context.Context, namespace, cmName
 		return apis.ErrGeneric(fmt.Sprintf("Failed to get SKU handler: %v", err), "instanceType")
 	}
 
-	gpuConfigs := skuHandler.GetGPUConfigs()
-	if skuConfig, exists := gpuConfigs[instanceType]; exists {
+	if skuConfig := skuHandler.GetGPUConfigBySKU(instanceType); skuConfig != nil {
 		// Check if this is a multi-GPU instance with less than 20GB per GPU
-		gpuMemPerGPU := skuConfig.GPUMem / skuConfig.GPUCount
+		gpuMemPerGPU := skuConfig.GPUMemGB / skuConfig.GPUCount
 		if skuConfig.GPUCount > 1 && gpuMemPerGPU < 20 {
 			// For multi-GPU instances with less than 20GB per GPU, max-model-len is required
 			maxModelLen, exists := inferenceConfig.VLLM["max-model-len"]
diff --git a/api/v1beta1/workspace_validation.go b/api/v1beta1/workspace_validation.go
index 268e750..05c478f 100644
--- a/api/v1beta1/workspace_validation.go
+++ b/api/v1beta1/workspace_validation.go
@@ -59,8 +59,8 @@ func (w *Workspace) Validate(ctx context.Context) (errs *apis.FieldError) {
 					bypassResourceChecks = true
 				}
 			}
-			runtime := GetWorkspaceRuntimeName(w)
 
+			runtime := GetWorkspaceRuntimeName(w)
 			// TODO: Add Adapter Spec Validation - Including DataSource Validation for Adapter
 			errs = errs.Also(w.Resource.validateCreateWithInference(w.Inference, bypassResourceChecks).ViaField("resource"),
 				w.Inference.validateCreate(ctx, w.Namespace, w.Resource.InstanceType, runtime).ViaField("inference"))
@@ -310,17 +310,16 @@ func (r *ResourceSpec) validateCreateWithInference(inference *InferenceSpec, byp
 		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 skuConfig := skuHandler.GetGPUConfigBySKU(instanceType); skuConfig != nil {
 		if presetName != "" {
 			model := plugin.KaitoModelRegister.MustGet(presetName) // InferenceSpec has been validated so the name is valid.
 
 			machineCount := *r.Count
 			machineTotalNumGPUs := resource.NewQuantity(int64(machineCount*skuConfig.GPUCount), resource.DecimalSI)
-			machinePerGPUMemory := resource.NewQuantity(int64(skuConfig.GPUMem/skuConfig.GPUCount)*consts.GiBToBytes, resource.BinarySI) // Ensure it's per GPU
-			machineTotalGPUMem := resource.NewQuantity(int64(machineCount*skuConfig.GPUMem)*consts.GiBToBytes, resource.BinarySI)        // Total GPU memory
+			machinePerGPUMemory := resource.NewQuantity(int64(skuConfig.GPUMemGB/skuConfig.GPUCount)*consts.GiBToBytes, resource.BinarySI) // Ensure it's per GPU
+			machineTotalGPUMem := resource.NewQuantity(int64(machineCount*skuConfig.GPUMemGB)*consts.GiBToBytes, resource.BinarySI)        // Total GPU memory
 
 			modelGPUCount := resource.MustParse(model.GetInferenceParameters().GPUCountRequirement)
 			modelPerGPUMemory := resource.MustParse(model.GetInferenceParameters().PerGPUMemoryRequirement)
@@ -447,7 +446,7 @@ func (i *InferenceSpec) validateCreate(ctx context.Context, namespace string, in
 			UseAdapters: len(i.Adapters) > 0,
 		})
 		if err != nil {
-			errs = errs.Also(apis.ErrGeneric(fmt.Sprintf("vLLM runtime validation: %v", err)))
+			errs = errs.Also(apis.ErrGeneric(fmt.Sprintf("Runtime validation: %v", err)))
 		}
 		// Additional validations for Preset
 		if i.Preset.AccessMode == ModelImageAccessModePrivate && i.Preset.Image == "" {
diff --git a/charts/kaito/workspace/templates/inference-params.yaml b/charts/kaito/workspace/templates/inference-params.yaml
index 590307d..29b8279 100644
--- a/charts/kaito/workspace/templates/inference-params.yaml
+++ b/charts/kaito/workspace/templates/inference-params.yaml
@@ -10,9 +10,9 @@ data:
 
     vllm:
       cpu-offload-gb: 0
-      gpu-memory-utilization: 0.95
       swap-space: 4
 
+      # gpu-memory-utilization: 0.9
       # max-seq-len-to-capture: 8192
       # num-scheduler-steps: 1
       # enable-chunked-prefill: false
diff --git a/pkg/model/interface.go b/pkg/model/interface.go
index cce7f60..c90e013 100644
--- a/pkg/model/interface.go
+++ b/pkg/model/interface.go
@@ -5,11 +5,15 @@ package model
 import (
 	"fmt"
 	"maps"
+	"math"
 	"path"
+	"strconv"
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
 
+	"github.com/kaito-project/kaito/pkg/sku"
 	"github.com/kaito-project/kaito/pkg/utils"
 )
 
@@ -27,7 +31,16 @@ const (
 	RuntimeNameHuggingfaceTransformers RuntimeName = "transformers"
 	RuntimeNameVLLM                    RuntimeName = "vllm"
 
-	ConfigfileNameVLLM = "inference_config.yaml"
+	ConfigfileNameVLLM    = "inference_config.yaml"
+	DefaultMemoryUtilVLLM = 0.9 // Default gpu memory utilization for VLLM runtime
+	UpperMemoryUtilVLLM   = 0.95
+)
+
+var (
+	// vLLM will do kvcache pre-allocation,
+	// We need to reserve enough memory for other ephemeral operations to avoid OOM.
+	// This is an empirical value.
+	ReservedNonKVCacheMemory = resource.MustParse("1.5Gi")
 )
 
 // PresetParam defines the preset inference parameters for a model.
@@ -126,8 +139,9 @@ func (v *VLLMParam) DeepCopy() VLLMParam {
 // RuntimeContext defines the runtime context for a model.
 type RuntimeContext struct {
 	RuntimeName  RuntimeName
+	GPUConfig    *sku.GPUConfig
 	ConfigVolume *corev1.VolumeMount
-	SKUNumGPUs   string
+	SKUNumGPUs   int
 	UseAdapters  bool
 }
 
@@ -160,11 +174,13 @@ func (p *PresetParam) buildVLLMInferenceCommand(rc RuntimeContext) []string {
 		p.VLLM.ModelRunParams["served-model-name"] = p.VLLM.ModelName
 	}
 	if !p.DisableTensorParallelism {
-		p.VLLM.ModelRunParams["tensor-parallel-size"] = rc.SKUNumGPUs
+		p.VLLM.ModelRunParams["tensor-parallel-size"] = strconv.Itoa(rc.SKUNumGPUs)
 	}
 	if !p.VLLM.DisallowLoRA && rc.UseAdapters {
 		p.VLLM.ModelRunParams["enable-lora"] = ""
 	}
+	gpuMemUtil := getGPUMemoryUtilForVLLM(rc.GPUConfig)
+	p.VLLM.ModelRunParams["gpu-memory-utilization"] = strconv.FormatFloat(gpuMemUtil, 'f', 2, 64)
 	if rc.ConfigVolume != nil {
 		p.VLLM.ModelRunParams["kaito-config-file"] = path.Join(rc.ConfigVolume.MountPath, ConfigfileNameVLLM)
 	}
@@ -181,3 +197,20 @@ func (p *PresetParam) Validate(rc RuntimeContext) error {
 	}
 	return nil
 }
+
+func getGPUMemoryUtilForVLLM(gpuConfig *sku.GPUConfig) float64 {
+	if gpuConfig == nil || gpuConfig.GPUMemGB <= 0 || gpuConfig.GPUCount <= 0 {
+		return DefaultMemoryUtilVLLM
+	}
+
+	gpuMem := resource.MustParse(fmt.Sprintf("%dGi", gpuConfig.GPUMemGB))
+	gpuMemPerGPU := float64(gpuMem.Value()) / float64(gpuConfig.GPUCount)
+
+	if float64(ReservedNonKVCacheMemory.Value()) >= gpuMemPerGPU {
+		// looks impossible, just prevent this case
+		return DefaultMemoryUtilVLLM
+	}
+
+	util := math.Floor((1.0-float64(ReservedNonKVCacheMemory.Value())/gpuMemPerGPU)*100) / 100
+	return math.Min(util, UpperMemoryUtilVLLM)
+}
diff --git a/pkg/ragengine/controllers/preset-rag.go b/pkg/ragengine/controllers/preset-rag.go
index 49cf6fc..26302fd 100644
--- a/pkg/ragengine/controllers/preset-rag.go
+++ b/pkg/ragengine/controllers/preset-rag.go
@@ -4,7 +4,6 @@ package controllers
 
 import (
 	"context"
-	"fmt"
 
 	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -81,18 +80,24 @@ func CreatePresetRAG(ctx context.Context, ragEngineObj *v1alpha1.RAGEngine, revi
 	var resourceReq corev1.ResourceRequirements
 
 	if ragEngineObj.Spec.Embedding.Local != nil {
-		skuNumGPUs, err := utils.GetSKUNumGPUs(ctx, kubeClient, ragEngineObj.Status.WorkerNodes,
-			ragEngineObj.Spec.Compute.InstanceType, "1")
+		var skuNumGPUs int
+		gpuConfig, err := utils.GetGPUConfigBySKU(ragEngineObj.Spec.Compute.InstanceType)
 		if err != nil {
-			return nil, fmt.Errorf("failed to get SKU num GPUs: %v", err)
+			gpuConfig, err = utils.TryGetGPUConfigFromNode(ctx, kubeClient, ragEngineObj.Status.WorkerNodes)
+			if err != nil {
+				skuNumGPUs = 1
+			}
+		}
+		if gpuConfig != nil {
+			skuNumGPUs = gpuConfig.GPUCount
 		}
 
 		resourceReq = corev1.ResourceRequirements{
 			Requests: corev1.ResourceList{
-				corev1.ResourceName(resources.CapacityNvidiaGPU): resource.MustParse(skuNumGPUs),
+				corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
 			},
 			Limits: corev1.ResourceList{
-				corev1.ResourceName(resources.CapacityNvidiaGPU): resource.MustParse(skuNumGPUs),
+				corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
 			},
 		}
 
diff --git a/pkg/sku/arc_sku_handker.go b/pkg/sku/arc_sku_handker.go
index d4392f2..5d3ef30 100644
--- a/pkg/sku/arc_sku_handker.go
+++ b/pkg/sku/arc_sku_handker.go
@@ -3,34 +3,20 @@
 
 package sku
 
-var _ CloudSKUHandler = &ArcSKUHandler{}
-
-type ArcSKUHandler struct {
-	supportedSKUs map[string]GPUConfig
-}
-
-func NewArcSKUHandler() *ArcSKUHandler {
-	return &ArcSKUHandler{
+func NewArcSKUHandler() CloudSKUHandler {
+	return &generalSKUHandler{
 		// The SKU including Azure Local official sku: https://learn.microsoft.com/en-us/azure/aks/aksarc/scale-requirements
 		supportedSKUs: map[string]GPUConfig{
-			"Standard_NK6":      {SKU: "Standard_NK6", GPUCount: 1, GPUMem: 8, GPUModel: "NVIDIA T4"},
-			"Standard_NK12":     {SKU: "Standard_NK12", GPUCount: 2, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"Standard_NC4_A2":   {SKU: "Standard_NC4_A2", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA A2"},
-			"Standard_NC8_A2":   {SKU: "Standard_NC8_A2", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA A2"},
-			"Standard_NC16_A2":  {SKU: "Standard_NC16_A2", GPUCount: 2, GPUMem: 48, GPUModel: "NVIDIA A2"},
-			"Standard_NC32_A2":  {SKU: "Standard_NC32_A2", GPUCount: 2, GPUMem: 48, GPUModel: "NVIDIA A2"},
-			"Standard_NC4_A16":  {SKU: "Standard_NC4_A16", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA A16"},
-			"Standard_NC8_A16":  {SKU: "Standard_NC8_A16", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA A16"},
-			"Standard_NC16_A16": {SKU: "Standard_NC16_A16", GPUCount: 2, GPUMem: 48, GPUModel: "NVIDIA A16"},
-			"Standard_NC32_A16": {SKU: "Standard_NC32_A16", GPUCount: 2, GPUMem: 48, GPUModel: "NVIDIA A16"},
+			"Standard_NK6":      {SKU: "Standard_NK6", GPUCount: 1, GPUMemGB: 8, GPUModel: "NVIDIA T4"},
+			"Standard_NK12":     {SKU: "Standard_NK12", GPUCount: 2, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"Standard_NC4_A2":   {SKU: "Standard_NC4_A2", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA A2"},
+			"Standard_NC8_A2":   {SKU: "Standard_NC8_A2", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA A2"},
+			"Standard_NC16_A2":  {SKU: "Standard_NC16_A2", GPUCount: 2, GPUMemGB: 48, GPUModel: "NVIDIA A2"},
+			"Standard_NC32_A2":  {SKU: "Standard_NC32_A2", GPUCount: 2, GPUMemGB: 48, GPUModel: "NVIDIA A2"},
+			"Standard_NC4_A16":  {SKU: "Standard_NC4_A16", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA A16"},
+			"Standard_NC8_A16":  {SKU: "Standard_NC8_A16", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA A16"},
+			"Standard_NC16_A16": {SKU: "Standard_NC16_A16", GPUCount: 2, GPUMemGB: 48, GPUModel: "NVIDIA A16"},
+			"Standard_NC32_A16": {SKU: "Standard_NC32_A16", GPUCount: 2, GPUMemGB: 48, GPUModel: "NVIDIA A16"},
 		},
 	}
 }
-
-func (a *ArcSKUHandler) GetSupportedSKUs() []string {
-	return GetMapKeys(a.supportedSKUs)
-}
-
-func (a *ArcSKUHandler) GetGPUConfigs() map[string]GPUConfig {
-	return a.supportedSKUs
-}
diff --git a/pkg/sku/aws_sku_handler.go b/pkg/sku/aws_sku_handler.go
index e0a149f..f009d06 100644
--- a/pkg/sku/aws_sku_handler.go
+++ b/pkg/sku/aws_sku_handler.go
@@ -3,82 +3,68 @@
 
 package sku
 
-var _ CloudSKUHandler = &AwsSKUHandler{}
-
-type AwsSKUHandler struct {
-	supportedSKUs map[string]GPUConfig
-}
-
-func NewAwsSKUHandler() *AwsSKUHandler {
-	return &AwsSKUHandler{
+func NewAwsSKUHandler() CloudSKUHandler {
+	return &generalSKUHandler{
 		// Reference: https://aws.amazon.com/ec2/instance-types/
 		supportedSKUs: map[string]GPUConfig{
-			"p2.xlarge":     {SKU: "p2.xlarge", GPUCount: 1, GPUMem: 12, GPUModel: "NVIDIA K80"},
-			"p2.8xlarge":    {SKU: "p2.8xlarge", GPUCount: 8, GPUMem: 96, GPUModel: "NVIDIA K80"},
-			"p2.16xlarge":   {SKU: "p2.16xlarge", GPUCount: 16, GPUMem: 192, GPUModel: "NVIDIA K80"},
-			"p3.2xlarge":    {SKU: "p3.2xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA V100"},
-			"p3.8xlarge":    {SKU: "p3.8xlarge", GPUCount: 4, GPUMem: 64, GPUModel: "NVIDIA V100"},
-			"p3.16xlarge":   {SKU: "p3.16xlarge", GPUCount: 8, GPUMem: 128, GPUModel: "NVIDIA V100"},
-			"p3dn.24xlarge": {SKU: "p3dn.24xlarge", GPUCount: 8, GPUMem: 256, GPUModel: "NVIDIA V100"},
-			"p4d.24xlarge":  {SKU: "p4d.24xlarge", GPUCount: 8, GPUMem: 320, GPUModel: "NVIDIA A100"},
-			"p4de.24xlarge": {SKU: "p4de.24xlarge", GPUCount: 8, GPUMem: 640, GPUModel: "NVIDIA A100"},
-			"p5.48xlarge":   {SKU: "p5.48xlarge", GPUCount: 8, GPUMem: 640, GPUModel: "NVIDIA H100"},
-			"g6.xlarge":     {SKU: "g6.xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA L4"},
-			"g6.2xlarge":    {SKU: "g6.2xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA L4"},
-			"g6.4xlarge":    {SKU: "g6.4xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA L4"},
-			"g6.8xlarge":    {SKU: "g6.8xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA L4"},
-			"g6.16xlarge":   {SKU: "g6.16xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA L4"},
-			"gr6.4xlarge":   {SKU: "gr6.4xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA L4"},
-			"gr6.8xlarge":   {SKU: "gr6.8xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA L4"},
-			"g6.12xlarge":   {SKU: "g6.12xlarge", GPUCount: 4, GPUMem: 96, GPUModel: "NVIDIA L4"},
-			"g6.24xlarge":   {SKU: "g6.24xlarge", GPUCount: 4, GPUMem: 96, GPUModel: "NVIDIA L4"},
-			"g6.48xlarge":   {SKU: "g6.48xlarge", GPUCount: 8, GPUMem: 192, GPUModel: "NVIDIA L4"},
-			"g5g.xlarge":    {SKU: "g5g.xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"g5g.2xlarge":   {SKU: "g5g.2xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"g5g.4xlarge":   {SKU: "g5g.4xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"g5g.8xlarge":   {SKU: "g5g.8xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"g5g.16xlarge":  {SKU: "g5g.16xlarge", GPUCount: 2, GPUMem: 32, GPUModel: "NVIDIA T4"},
-			"g5g.metal":     {SKU: "g5g.metal", GPUCount: 2, GPUMem: 32, GPUModel: "NVIDIA T4"},
-			"g5.xlarge":     {SKU: "g5.xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA A10G"},
-			"g5.2xlarge":    {SKU: "g5.2xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA A10G"},
-			"g5.4xlarge":    {SKU: "g5.4xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA A10G"},
-			"g5.8xlarge":    {SKU: "g5.8xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA A10G"},
-			"g5.12xlarge":   {SKU: "g5.12xlarge", GPUCount: 4, GPUMem: 96, GPUModel: "NVIDIA A10G"},
-			"g5.16xlarge":   {SKU: "g5.16xlarge", GPUCount: 1, GPUMem: 24, GPUModel: "NVIDIA A10G"},
-			"g5.24xlarge":   {SKU: "g5.24xlarge", GPUCount: 4, GPUMem: 96, GPUModel: "NVIDIA A10G"},
-			"g5.48xlarge":   {SKU: "g5.48xlarge", GPUCount: 8, GPUMem: 192, GPUModel: "NVIDIA A10G"},
-			"g4dn.xlarge":   {SKU: "g4dn.xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"g4dn.2xlarge":  {SKU: "g4dn.2xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"g4dn.4xlarge":  {SKU: "g4dn.4xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"g4dn.8xlarge":  {SKU: "g4dn.8xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"g4dn.16xlarge": {SKU: "g4dn.16xlarge", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"g4dn.12xlarge": {SKU: "g4dn.12xlarge", GPUCount: 4, GPUMem: 64, GPUModel: "NVIDIA T4"},
-			"g4dn.metal":    {SKU: "g4dn.metal", GPUCount: 8, GPUMem: 128, GPUModel: "NVIDIA T4"},
-			"g4ad.xlarge":   {SKU: "g4ad.xlarge", GPUCount: 1, GPUMem: 8, GPUModel: "AMD Radeon Pro V520"},
-			"g4ad.2xlarge":  {SKU: "g4ad.2xlarge", GPUCount: 1, GPUMem: 8, GPUModel: "AMD Radeon Pro V520"},
-			"g4ad.4xlarge":  {SKU: "g4ad.4xlarge", GPUCount: 1, GPUMem: 8, GPUModel: "AMD Radeon Pro V520"},
-			"g4ad.8xlarge":  {SKU: "g4ad.8xlarge", GPUCount: 2, GPUMem: 16, GPUModel: "AMD Radeon Pro V520"},
-			"g4ad.16xlarge": {SKU: "g4ad.16xlarge", GPUCount: 4, GPUMem: 32, GPUModel: "AMD Radeon Pro V520"},
-			"g3s.xlarge":    {SKU: "g3s.xlarge", GPUCount: 1, GPUMem: 8, GPUModel: "NVIDIA M60"},
-			"g3s.4xlarge":   {SKU: "g3s.4xlarge", GPUCount: 1, GPUMem: 8, GPUModel: "NVIDIA M60"},
-			"g3s.8xlarge":   {SKU: "g3s.8xlarge", GPUCount: 2, GPUMem: 16, GPUModel: "NVIDIA M60"},
-			"g3s.16xlarge":  {SKU: "g3s.16xlarge", GPUCount: 4, GPUMem: 32, GPUModel: "NVIDIA M60"},
+			"p2.xlarge":     {SKU: "p2.xlarge", GPUCount: 1, GPUMemGB: 12, GPUModel: "NVIDIA K80"},
+			"p2.8xlarge":    {SKU: "p2.8xlarge", GPUCount: 8, GPUMemGB: 96, GPUModel: "NVIDIA K80"},
+			"p2.16xlarge":   {SKU: "p2.16xlarge", GPUCount: 16, GPUMemGB: 192, GPUModel: "NVIDIA K80"},
+			"p3.2xlarge":    {SKU: "p3.2xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA V100"},
+			"p3.8xlarge":    {SKU: "p3.8xlarge", GPUCount: 4, GPUMemGB: 64, GPUModel: "NVIDIA V100"},
+			"p3.16xlarge":   {SKU: "p3.16xlarge", GPUCount: 8, GPUMemGB: 128, GPUModel: "NVIDIA V100"},
+			"p3dn.24xlarge": {SKU: "p3dn.24xlarge", GPUCount: 8, GPUMemGB: 256, GPUModel: "NVIDIA V100"},
+			"p4d.24xlarge":  {SKU: "p4d.24xlarge", GPUCount: 8, GPUMemGB: 320, GPUModel: "NVIDIA A100"},
+			"p4de.24xlarge": {SKU: "p4de.24xlarge", GPUCount: 8, GPUMemGB: 640, GPUModel: "NVIDIA A100"},
+			"p5.48xlarge":   {SKU: "p5.48xlarge", GPUCount: 8, GPUMemGB: 640, GPUModel: "NVIDIA H100"},
+			"g6.xlarge":     {SKU: "g6.xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA L4"},
+			"g6.2xlarge":    {SKU: "g6.2xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA L4"},
+			"g6.4xlarge":    {SKU: "g6.4xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA L4"},
+			"g6.8xlarge":    {SKU: "g6.8xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA L4"},
+			"g6.16xlarge":   {SKU: "g6.16xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA L4"},
+			"gr6.4xlarge":   {SKU: "gr6.4xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA L4"},
+			"gr6.8xlarge":   {SKU: "gr6.8xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA L4"},
+			"g6.12xlarge":   {SKU: "g6.12xlarge", GPUCount: 4, GPUMemGB: 96, GPUModel: "NVIDIA L4"},
+			"g6.24xlarge":   {SKU: "g6.24xlarge", GPUCount: 4, GPUMemGB: 96, GPUModel: "NVIDIA L4"},
+			"g6.48xlarge":   {SKU: "g6.48xlarge", GPUCount: 8, GPUMemGB: 192, GPUModel: "NVIDIA L4"},
+			"g5g.xlarge":    {SKU: "g5g.xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"g5g.2xlarge":   {SKU: "g5g.2xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"g5g.4xlarge":   {SKU: "g5g.4xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"g5g.8xlarge":   {SKU: "g5g.8xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"g5g.16xlarge":  {SKU: "g5g.16xlarge", GPUCount: 2, GPUMemGB: 32, GPUModel: "NVIDIA T4"},
+			"g5g.metal":     {SKU: "g5g.metal", GPUCount: 2, GPUMemGB: 32, GPUModel: "NVIDIA T4"},
+			"g5.xlarge":     {SKU: "g5.xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA A10G"},
+			"g5.2xlarge":    {SKU: "g5.2xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA A10G"},
+			"g5.4xlarge":    {SKU: "g5.4xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA A10G"},
+			"g5.8xlarge":    {SKU: "g5.8xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA A10G"},
+			"g5.12xlarge":   {SKU: "g5.12xlarge", GPUCount: 4, GPUMemGB: 96, GPUModel: "NVIDIA A10G"},
+			"g5.16xlarge":   {SKU: "g5.16xlarge", GPUCount: 1, GPUMemGB: 24, GPUModel: "NVIDIA A10G"},
+			"g5.24xlarge":   {SKU: "g5.24xlarge", GPUCount: 4, GPUMemGB: 96, GPUModel: "NVIDIA A10G"},
+			"g5.48xlarge":   {SKU: "g5.48xlarge", GPUCount: 8, GPUMemGB: 192, GPUModel: "NVIDIA A10G"},
+			"g4dn.xlarge":   {SKU: "g4dn.xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"g4dn.2xlarge":  {SKU: "g4dn.2xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"g4dn.4xlarge":  {SKU: "g4dn.4xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"g4dn.8xlarge":  {SKU: "g4dn.8xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"g4dn.16xlarge": {SKU: "g4dn.16xlarge", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"g4dn.12xlarge": {SKU: "g4dn.12xlarge", GPUCount: 4, GPUMemGB: 64, GPUModel: "NVIDIA T4"},
+			"g4dn.metal":    {SKU: "g4dn.metal", GPUCount: 8, GPUMemGB: 128, GPUModel: "NVIDIA T4"},
+			"g4ad.xlarge":   {SKU: "g4ad.xlarge", GPUCount: 1, GPUMemGB: 8, GPUModel: "AMD Radeon Pro V520"},
+			"g4ad.2xlarge":  {SKU: "g4ad.2xlarge", GPUCount: 1, GPUMemGB: 8, GPUModel: "AMD Radeon Pro V520"},
+			"g4ad.4xlarge":  {SKU: "g4ad.4xlarge", GPUCount: 1, GPUMemGB: 8, GPUModel: "AMD Radeon Pro V520"},
+			"g4ad.8xlarge":  {SKU: "g4ad.8xlarge", GPUCount: 2, GPUMemGB: 16, GPUModel: "AMD Radeon Pro V520"},
+			"g4ad.16xlarge": {SKU: "g4ad.16xlarge", GPUCount: 4, GPUMemGB: 32, GPUModel: "AMD Radeon Pro V520"},
+			"g3s.xlarge":    {SKU: "g3s.xlarge", GPUCount: 1, GPUMemGB: 8, GPUModel: "NVIDIA M60"},
+			"g3s.4xlarge":   {SKU: "g3s.4xlarge", GPUCount: 1, GPUMemGB: 8, GPUModel: "NVIDIA M60"},
+			"g3s.8xlarge":   {SKU: "g3s.8xlarge", GPUCount: 2, GPUMemGB: 16, GPUModel: "NVIDIA M60"},
+			"g3s.16xlarge":  {SKU: "g3s.16xlarge", GPUCount: 4, GPUMemGB: 32, GPUModel: "NVIDIA M60"},
 			//accelerator optimized
-			"trn1.2xlarge":   {SKU: "trn1.2xlarge", GPUCount: 1, GPUMem: 32, GPUModel: "AWS Trainium accelerators"},
-			"trn1.32xlarge":  {SKU: "trn1.32xlarge", GPUCount: 16, GPUMem: 512, GPUModel: "AWS Trainium accelerators"},
-			"trn1n.32xlarge": {SKU: "trn1n.32xlarge", GPUCount: 16, GPUMem: 512, GPUModel: "AWS Trainium accelerators"},
-			"inf2.xlarge":    {SKU: "inf2.xlarge", GPUCount: 1, GPUMem: 32, GPUModel: "AWS Inferentia2 accelerators"},
-			"inf2.8xlarge":   {SKU: "inf2.8xlarge", GPUCount: 1, GPUMem: 32, GPUModel: "AWS Inferentia2 accelerators"},
-			"inf2.24xlarge":  {SKU: "inf2.24xlarge", GPUCount: 6, GPUMem: 192, GPUModel: "AWS Inferentia2 accelerators"},
-			"inf2.48xlarge":  {SKU: "inf2.48xlarge", GPUCount: 12, GPUMem: 384, GPUModel: "AWS Inferentia2 accelerators"},
+			"trn1.2xlarge":   {SKU: "trn1.2xlarge", GPUCount: 1, GPUMemGB: 32, GPUModel: "AWS Trainium accelerators"},
+			"trn1.32xlarge":  {SKU: "trn1.32xlarge", GPUCount: 16, GPUMemGB: 512, GPUModel: "AWS Trainium accelerators"},
+			"trn1n.32xlarge": {SKU: "trn1n.32xlarge", GPUCount: 16, GPUMemGB: 512, GPUModel: "AWS Trainium accelerators"},
+			"inf2.xlarge":    {SKU: "inf2.xlarge", GPUCount: 1, GPUMemGB: 32, GPUModel: "AWS Inferentia2 accelerators"},
+			"inf2.8xlarge":   {SKU: "inf2.8xlarge", GPUCount: 1, GPUMemGB: 32, GPUModel: "AWS Inferentia2 accelerators"},
+			"inf2.24xlarge":  {SKU: "inf2.24xlarge", GPUCount: 6, GPUMemGB: 192, GPUModel: "AWS Inferentia2 accelerators"},
+			"inf2.48xlarge":  {SKU: "inf2.48xlarge", GPUCount: 12, GPUMemGB: 384, GPUModel: "AWS Inferentia2 accelerators"},
 		},
 	}
 }
-
-func (a *AwsSKUHandler) GetSupportedSKUs() []string {
-	return GetMapKeys(a.supportedSKUs)
-}
-
-func (a *AwsSKUHandler) GetGPUConfigs() map[string]GPUConfig {
-	return a.supportedSKUs
-}
diff --git a/pkg/sku/azure_sku_handler.go b/pkg/sku/azure_sku_handler.go
index 04d3fe6..60e63c5 100644
--- a/pkg/sku/azure_sku_handler.go
+++ b/pkg/sku/azure_sku_handler.go
@@ -3,37 +3,31 @@
 
 package sku
 
-var _ CloudSKUHandler = &AzureSKUHandler{}
-
-type AzureSKUHandler struct {
-	supportedSKUs map[string]GPUConfig
-}
-
-func NewAzureSKUHandler() *AzureSKUHandler {
-	return &AzureSKUHandler{
+func NewAzureSKUHandler() CloudSKUHandler {
+	return &generalSKUHandler{
 		supportedSKUs: map[string]GPUConfig{
-			"Standard_NC6s_v3":          {SKU: "Standard_NC6s_v3", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA V100"},
-			"Standard_NC12s_v3":         {SKU: "Standard_NC12s_v3", GPUCount: 2, GPUMem: 32, GPUModel: "NVIDIA V100"},
-			"Standard_NC24s_v3":         {SKU: "Standard_NC24s_v3", GPUCount: 4, GPUMem: 64, GPUModel: "NVIDIA V100"},
-			"Standard_NC24rs_v3":        {SKU: "Standard_NC24rs_v3", GPUCount: 4, GPUMem: 64, GPUModel: "NVIDIA V100"},
-			"Standard_NC4as_T4_v3":      {SKU: "Standard_NC4as_T4_v3", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"Standard_NC8as_T4_v3":      {SKU: "Standard_NC8as_T4_v3", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"Standard_NC16as_T4_v3":     {SKU: "Standard_NC16as_T4_v3", GPUCount: 1, GPUMem: 16, GPUModel: "NVIDIA T4"},
-			"Standard_NC64as_T4_v3":     {SKU: "Standard_NC64as_T4_v3", GPUCount: 4, GPUMem: 64, GPUModel: "NVIDIA T4"},
-			"Standard_NC24ads_A100_v4":  {SKU: "Standard_NC24ads_A100_v4", GPUCount: 1, GPUMem: 80, GPUModel: "NVIDIA A100"},
-			"Standard_NC48ads_A100_v4":  {SKU: "Standard_NC48ads_A100_v4", GPUCount: 2, GPUMem: 160, GPUModel: "NVIDIA A100"},
-			"Standard_NC96ads_A100_v4":  {SKU: "Standard_NC96ads_A100_v4", GPUCount: 4, GPUMem: 320, GPUModel: "NVIDIA A100"},
-			"Standard_ND96asr_A100_v4":  {SKU: "Standard_ND96asr_A100_v4", GPUCount: 8, GPUMem: 320, GPUModel: "NVIDIA A100"},
-			"Standard_NG32ads_V620_v1":  {SKU: "Standard_NG32ads_V620_v1", GPUCount: 1, GPUMem: 32, GPUModel: "AMD Radeon PRO V620"},
-			"Standard_NG32adms_V620_v1": {SKU: "Standard_NG32adms_V620_v1", GPUCount: 1, GPUMem: 32, GPUModel: "AMD Radeon PRO V620"},
-			"Standard_NV6":              {SKU: "Standard_NV6", GPUCount: 1, GPUMem: 8, GPUModel: "NVIDIA M60"},
-			"Standard_NV12":             {SKU: "Standard_NV12", GPUCount: 2, GPUMem: 16, GPUModel: "NVIDIA M60"},
-			"Standard_NV24":             {SKU: "Standard_NV24", GPUCount: 4, GPUMem: 32, GPUModel: "NVIDIA M60"},
-			"Standard_NV12s_v3":         {SKU: "Standard_NV12s_v3", GPUCount: 1, GPUMem: 8, GPUModel: "NVIDIA M60"},
-			"Standard_NV24s_v3":         {SKU: "Standard_NV24s_v3", GPUCount: 2, GPUMem: 16, GPUModel: "NVIDIA M60"},
-			"Standard_NV48s_v3":         {SKU: "Standard_NV48s_v3", GPUCount: 4, GPUMem: 32, GPUModel: "NVIDIA M60"},
-			"Standard_NV32as_v4":        {SKU: "Standard_NV32as_v4", GPUCount: 1, GPUMem: 16, GPUModel: "AMD Radeon Instinct MI25"},
-			"Standard_ND96amsr_A100_v4": {SKU: "Standard_ND96amsr_A100_v4", GPUCount: 8, GPUMem: 80, GPUModel: "NVIDIA A100"},
+			"Standard_NC6s_v3":          {SKU: "Standard_NC6s_v3", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA V100"},
+			"Standard_NC12s_v3":         {SKU: "Standard_NC12s_v3", GPUCount: 2, GPUMemGB: 32, GPUModel: "NVIDIA V100"},
+			"Standard_NC24s_v3":         {SKU: "Standard_NC24s_v3", GPUCount: 4, GPUMemGB: 64, GPUModel: "NVIDIA V100"},
+			"Standard_NC24rs_v3":        {SKU: "Standard_NC24rs_v3", GPUCount: 4, GPUMemGB: 64, GPUModel: "NVIDIA V100"},
+			"Standard_NC4as_T4_v3":      {SKU: "Standard_NC4as_T4_v3", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"Standard_NC8as_T4_v3":      {SKU: "Standard_NC8as_T4_v3", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"Standard_NC16as_T4_v3":     {SKU: "Standard_NC16as_T4_v3", GPUCount: 1, GPUMemGB: 16, GPUModel: "NVIDIA T4"},
+			"Standard_NC64as_T4_v3":     {SKU: "Standard_NC64as_T4_v3", GPUCount: 4, GPUMemGB: 64, GPUModel: "NVIDIA T4"},
+			"Standard_NC24ads_A100_v4":  {SKU: "Standard_NC24ads_A100_v4", GPUCount: 1, GPUMemGB: 80, GPUModel: "NVIDIA A100"},
+			"Standard_NC48ads_A100_v4":  {SKU: "Standard_NC48ads_A100_v4", GPUCount: 2, GPUMemGB: 160, GPUModel: "NVIDIA A100"},
+			"Standard_NC96ads_A100_v4":  {SKU: "Standard_NC96ads_A100_v4", GPUCount: 4, GPUMemGB: 320, GPUModel: "NVIDIA A100"},
+			"Standard_ND96asr_A100_v4":  {SKU: "Standard_ND96asr_A100_v4", GPUCount: 8, GPUMemGB: 320, GPUModel: "NVIDIA A100"},
+			"Standard_NG32ads_V620_v1":  {SKU: "Standard_NG32ads_V620_v1", GPUCount: 1, GPUMemGB: 32, GPUModel: "AMD Radeon PRO V620"},
+			"Standard_NG32adms_V620_v1": {SKU: "Standard_NG32adms_V620_v1", GPUCount: 1, GPUMemGB: 32, GPUModel: "AMD Radeon PRO V620"},
+			"Standard_NV6":              {SKU: "Standard_NV6", GPUCount: 1, GPUMemGB: 8, GPUModel: "NVIDIA M60"},
+			"Standard_NV12":             {SKU: "Standard_NV12", GPUCount: 2, GPUMemGB: 16, GPUModel: "NVIDIA M60"},
+			"Standard_NV24":             {SKU: "Standard_NV24", GPUCount: 4, GPUMemGB: 32, GPUModel: "NVIDIA M60"},
+			"Standard_NV12s_v3":         {SKU: "Standard_NV12s_v3", GPUCount: 1, GPUMemGB: 8, GPUModel: "NVIDIA M60"},
+			"Standard_NV24s_v3":         {SKU: "Standard_NV24s_v3", GPUCount: 2, GPUMemGB: 16, GPUModel: "NVIDIA M60"},
+			"Standard_NV48s_v3":         {SKU: "Standard_NV48s_v3", GPUCount: 4, GPUMemGB: 32, GPUModel: "NVIDIA M60"},
+			"Standard_NV32as_v4":        {SKU: "Standard_NV32as_v4", GPUCount: 1, GPUMemGB: 16, GPUModel: "AMD Radeon Instinct MI25"},
+			"Standard_ND96amsr_A100_v4": {SKU: "Standard_ND96amsr_A100_v4", GPUCount: 8, GPUMemGB: 80, GPUModel: "NVIDIA A100"},
 
 			// Not supporting partial gpu skus for now
 			// "Standard_NG8ads_V620_v1":   {SKU: "Standard_NG8ads_V620_v1", GPUCount: 1.0 / 4.0, GPUMem: 8, GPUModel: "AMD Radeon PRO V620"},
@@ -44,11 +38,3 @@ func NewAzureSKUHandler() *AzureSKUHandler {
 		},
 	}
 }
-
-func (a *AzureSKUHandler) GetSupportedSKUs() []string {
-	return GetMapKeys(a.supportedSKUs)
-}
-
-func (a *AzureSKUHandler) GetGPUConfigs() map[string]GPUConfig {
-	return a.supportedSKUs
-}
diff --git a/pkg/sku/cloud_sku_handler.go b/pkg/sku/cloud_sku_handler.go
index 0a6fe72..b3b2dfd 100644
--- a/pkg/sku/cloud_sku_handler.go
+++ b/pkg/sku/cloud_sku_handler.go
@@ -10,12 +10,13 @@ import (
 type CloudSKUHandler interface {
 	GetSupportedSKUs() []string
 	GetGPUConfigs() map[string]GPUConfig
+	GetGPUConfigBySKU(sku string) *GPUConfig
 }
 
 type GPUConfig struct {
 	SKU      string
 	GPUCount int
-	GPUMem   int
+	GPUMemGB int
 	GPUModel string
 }
 
@@ -31,3 +32,26 @@ func GetCloudSKUHandler(cloud string) CloudSKUHandler {
 		return nil
 	}
 }
+
+type generalSKUHandler struct {
+	supportedSKUs map[string]GPUConfig
+}
+
+func (b *generalSKUHandler) GetSupportedSKUs() []string {
+	keys := make([]string, 0, len(b.supportedSKUs))
+	for k := range b.supportedSKUs {
+		keys = append(keys, k)
+	}
+	return keys
+}
+
+func (b *generalSKUHandler) GetGPUConfigs() map[string]GPUConfig {
+	return b.supportedSKUs
+}
+
+func (b *generalSKUHandler) GetGPUConfigBySKU(sku string) *GPUConfig {
+	if config, ok := b.supportedSKUs[sku]; ok {
+		return &config
+	}
+	return nil
+}
diff --git a/pkg/sku/sku_utils.go b/pkg/sku/sku_utils.go
deleted file mode 100644
index 4198e3b..0000000
--- a/pkg/sku/sku_utils.go
+++ /dev/null
@@ -1,12 +0,0 @@
-// Copyright (c) Microsoft Corporation.
-// Licensed under the MIT license.
-
-package sku
-
-func GetMapKeys(m map[string]GPUConfig) []string {
-	keys := make([]string, 0, len(m))
-	for k := range m {
-		keys = append(keys, k)
-	}
-	return keys
-}
diff --git a/pkg/utils/common.go b/pkg/utils/common.go
index 141aaae..881752f 100644
--- a/pkg/utils/common.go
+++ b/pkg/utils/common.go
@@ -152,33 +152,32 @@ func GetSKUHandler() (sku.CloudSKUHandler, error) {
 	return skuHandler, nil
 }
 
-func GetSKUNumGPUs(ctx context.Context, kubeClient client.Client, workerNodes []string, instanceType, defaultGPUCount string) (string, error) {
+func GetGPUConfigBySKU(instanceType string) (*sku.GPUConfig, error) {
 	skuHandler, err := GetSKUHandler()
 	if err != nil {
-		return "", apis.ErrInvalidValue(fmt.Sprintf("Failed to get SKU handler: %v", err), "sku")
+		return nil, apis.ErrInvalidValue(fmt.Sprintf("Failed to get SKU handler: %v", err), "sku")
 	}
 
-	skuNumGPUs := defaultGPUCount // Default to using the provided default GPU count
+	return skuHandler.GetGPUConfigBySKU(instanceType), nil
+}
 
-	skuConfig, skuExists := skuHandler.GetGPUConfigs()[instanceType]
-	if skuExists {
-		skuNumGPUs = fmt.Sprintf("%d", skuConfig.GPUCount)
-	} else {
-		skuGPUCount, err := FetchGPUCountFromNodes(ctx, kubeClient, workerNodes)
-		if err != nil {
-			fmt.Printf("Failed to fetch GPU count from nodes: %v", err)
-		} else if skuGPUCount != "" {
-			skuNumGPUs = skuGPUCount
-		}
+func TryGetGPUConfigFromNode(ctx context.Context, kubeClient client.Client, workerNodes []string) (*sku.GPUConfig, error) {
+	skuGPUCount, err := FetchGPUCountFromNodes(ctx, kubeClient, workerNodes)
+	if err != nil || skuGPUCount == 0 {
+		return nil, fmt.Errorf("failed to fetch GPU count from nodes: %w", err)
 	}
 
-	return skuNumGPUs, nil
+	return &sku.GPUConfig{
+		SKU:      "unknown", // SKU is not available from nodes
+		GPUCount: skuGPUCount,
+		GPUModel: "unknown", // GPU model is not available from nodes
+	}, nil
 }
 
 // FetchGPUCountFromNodes retrieves the GPU count from the given node names.
-func FetchGPUCountFromNodes(ctx context.Context, kubeClient client.Client, nodeNames []string) (string, error) {
+func FetchGPUCountFromNodes(ctx context.Context, kubeClient client.Client, nodeNames []string) (int, error) {
 	if len(nodeNames) == 0 {
-		return "", fmt.Errorf("no worker nodes found in the workspace")
+		return 0, fmt.Errorf("no worker nodes found in the workspace")
 	}
 
 	var allNodes v1.NodeList
@@ -198,14 +197,14 @@ func FetchGPUCountFromNodes(ctx context.Context, kubeClient client.Client, nodeN
 	return GetPerNodeGPUCountFromNodes(&allNodes), nil
 }
 
-func GetPerNodeGPUCountFromNodes(nodeList *v1.NodeList) string {
+func GetPerNodeGPUCountFromNodes(nodeList *v1.NodeList) int {
 	for _, node := range nodeList.Items {
 		gpuCount, exists := node.Status.Capacity[consts.NvidiaGPU]
 		if exists && gpuCount.String() != "" {
-			return gpuCount.String()
+			return int(gpuCount.Value())
 		}
 	}
-	return ""
+	return 0
 }
 
 func ExtractAndValidateRepoName(image string) error {
diff --git a/pkg/workspace/inference/preset-inferences.go b/pkg/workspace/inference/preset-inferences.go
index 35e29b3..82c0cf3 100644
--- a/pkg/workspace/inference/preset-inferences.go
+++ b/pkg/workspace/inference/preset-inferences.go
@@ -154,20 +154,26 @@ func CreatePresetInference(ctx context.Context, workspaceObj *v1beta1.Workspace,
 	}
 
 	// resource requirements
-	skuNumGPUs, err := utils.GetSKUNumGPUs(ctx, kubeClient, workspaceObj.Status.WorkerNodes,
-		workspaceObj.Resource.InstanceType, inferenceParam.GPUCountRequirement)
+	var skuNumGPUs int
+	gpuConfig, err := utils.GetGPUConfigBySKU(workspaceObj.Resource.InstanceType)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get SKU num GPUs: %v", err)
+		gpuConfig, err = utils.TryGetGPUConfigFromNode(ctx, kubeClient, workspaceObj.Status.WorkerNodes)
+		if err != nil {
+			defaultNumGPU := resource.MustParse(inferenceParam.GPUCountRequirement)
+			skuNumGPUs = int(defaultNumGPU.Value())
+		}
+	}
+	if gpuConfig != nil {
+		skuNumGPUs = gpuConfig.GPUCount
 	}
 	resourceReq := corev1.ResourceRequirements{
 		Requests: corev1.ResourceList{
-			corev1.ResourceName(resources.CapacityNvidiaGPU): resource.MustParse(skuNumGPUs),
+			corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
 		},
 		Limits: corev1.ResourceList{
-			corev1.ResourceName(resources.CapacityNvidiaGPU): resource.MustParse(skuNumGPUs),
+			corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
 		},
 	}
-	skuGPUCount, _ := strconv.Atoi(skuNumGPUs)
 
 	// additional volume
 	var volumes []corev1.Volume
@@ -179,7 +185,7 @@ func CreatePresetInference(ctx context.Context, workspaceObj *v1beta1.Workspace,
 	volumeMounts = append(volumeMounts, cmVolumeMount)
 
 	// add share memory for cross process communication
-	shmVolume, shmVolumeMount := utils.ConfigSHMVolume(skuGPUCount)
+	shmVolume, shmVolumeMount := utils.ConfigSHMVolume(skuNumGPUs)
 	if shmVolume.Name != "" {
 		volumes = append(volumes, shmVolume)
 	}
@@ -196,6 +202,7 @@ func CreatePresetInference(ctx context.Context, workspaceObj *v1beta1.Workspace,
 	runtimeName := v1beta1.GetWorkspaceRuntimeName(workspaceObj)
 	commands := inferenceParam.GetInferenceCommand(pkgmodel.RuntimeContext{
 		RuntimeName:  runtimeName,
+		GPUConfig:    gpuConfig,
 		ConfigVolume: &cmVolumeMount,
 		SKUNumGPUs:   skuNumGPUs,
 		UseAdapters:  len(workspaceObj.Inference.Adapters) > 0,
diff --git a/pkg/workspace/tuning/preset-tuning.go b/pkg/workspace/tuning/preset-tuning.go
index ead8cbc..9d624e0 100644
--- a/pkg/workspace/tuning/preset-tuning.go
+++ b/pkg/workspace/tuning/preset-tuning.go
@@ -55,13 +55,10 @@ var (
 
 func getInstanceGPUCount(sku string) int {
 	skuHandler, _ := utils.GetSKUHandler()
-	gpuConfigs := skuHandler.GetGPUConfigs()
-
-	gpuConfig, exists := gpuConfigs[sku]
-	if !exists {
-		return 1
+	if gpuConfig := skuHandler.GetGPUConfigBySKU(sku); gpuConfig != nil {
+		return gpuConfig.GPUCount
 	}
-	return gpuConfig.GPUCount
+	return 1
 }
 
 func GetTuningImageInfo(ctx context.Context, workspaceObj *kaitov1beta1.Workspace, presetObj *model.PresetParam) (string, []corev1.LocalObjectReference) {
@@ -231,10 +228,17 @@ func CreatePresetTuning(ctx context.Context, workspaceObj *kaitov1beta1.Workspac
 		return nil, err
 	}
 
-	skuNumGPUs, err := utils.GetSKUNumGPUs(ctx, kubeClient, workspaceObj.Status.WorkerNodes,
-		workspaceObj.Resource.InstanceType, tuningObj.GPUCountRequirement)
+	var skuNumGPUs int
+	gpuConfig, err := utils.GetGPUConfigBySKU(workspaceObj.Resource.InstanceType)
 	if err != nil {
-		return nil, fmt.Errorf("failed to get SKU num GPUs: %v", err)
+		gpuConfig, err = utils.TryGetGPUConfigFromNode(ctx, kubeClient, workspaceObj.Status.WorkerNodes)
+		if err != nil {
+			defaultNumGPU := resource.MustParse(tuningObj.GPUCountRequirement)
+			skuNumGPUs = int(defaultNumGPU.Value())
+		}
+	}
+	if gpuConfig != nil {
+		skuNumGPUs = gpuConfig.GPUCount
 	}
 
 	commands, resourceReq := prepareTuningParameters(ctx, workspaceObj, modelCommand, tuningObj, skuNumGPUs)
@@ -375,7 +379,7 @@ func prepareModelRunParameters(ctx context.Context, tuningObj *model.PresetParam
 // and sets the GPU resources required for tuning.
 // Returns the command and resource configuration.
 func prepareTuningParameters(ctx context.Context, wObj *kaitov1beta1.Workspace, modelCommand string,
-	tuningObj *model.PresetParam, skuNumGPUs string) ([]string, corev1.ResourceRequirements) {
+	tuningObj *model.PresetParam, skuNumGPUs int) ([]string, corev1.ResourceRequirements) {
 	hfParam := tuningObj.Transformers // Only support Huggingface for now
 	if hfParam.TorchRunParams == nil {
 		hfParam.TorchRunParams = make(map[string]string)
@@ -388,10 +392,10 @@ func prepareTuningParameters(ctx context.Context, wObj *kaitov1beta1.Workspace,
 
 	resourceRequirements := corev1.ResourceRequirements{
 		Requests: corev1.ResourceList{
-			corev1.ResourceName(resources.CapacityNvidiaGPU): resource.MustParse(skuNumGPUs),
+			corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
 		},
 		Limits: corev1.ResourceList{
-			corev1.ResourceName(resources.CapacityNvidiaGPU): resource.MustParse(skuNumGPUs),
+			corev1.ResourceName(resources.CapacityNvidiaGPU): *resource.NewQuantity(int64(skuNumGPUs), resource.DecimalSI),
 		},
 	}
 



Test patch:
diff --git a/api/v1alpha1/workspace_validation_test.go b/api/v1alpha1/workspace_validation_test.go
index 91cf3d5..7b46aad 100644
--- a/api/v1alpha1/workspace_validation_test.go
+++ b/api/v1alpha1/workspace_validation_test.go
@@ -770,7 +770,7 @@ func TestInferenceSpecValidateCreate(t *testing.T) {
 					}
 				}()
 			}
-			errs := tc.inferenceSpec.validateCreate(ctx, DefaultReleaseNamespace, "Standard_NC24ads_A100_v4")
+			errs := tc.inferenceSpec.validateCreate(ctx, DefaultReleaseNamespace, "Standard_NC24ads_A100_v4", model.RuntimeNameHuggingfaceTransformers)
 			hasErrs := errs != nil
 			if hasErrs != tc.expectErrs {
 				t.Errorf("validateCreate() errors = %v, expectErrs %v", errs, tc.expectErrs)
@@ -1871,7 +1871,7 @@ other_field: value
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
 			// Validate the inference spec
-			errs := tc.inferenceSpec.validateCreate(ctx, DefaultReleaseNamespace, tc.resourceSpec.InstanceType)
+			errs := tc.inferenceSpec.validateCreate(ctx, DefaultReleaseNamespace, tc.resourceSpec.InstanceType, model.RuntimeNameHuggingfaceTransformers)
 			hasErrs := errs != nil
 
 			if hasErrs != tc.expectErrs {
diff --git a/pkg/model/interface_test.go b/pkg/model/interface_test.go
new file mode 100644
index 0000000..8770613
--- /dev/null
+++ b/pkg/model/interface_test.go
@@ -0,0 +1,80 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+package model
+
+import (
+	"testing"
+
+	"github.com/kaito-project/kaito/pkg/sku"
+)
+
+func TestGetGPUMemoryUtilForVLLM(t *testing.T) {
+	tests := []struct {
+		name      string
+		gpuConfig *sku.GPUConfig
+		expected  float64
+	}{
+		{
+			name:      "Nil GPUConfig",
+			gpuConfig: nil,
+			expected:  0.9,
+		},
+		{
+			name: "Zero GPU Memory",
+			gpuConfig: &sku.GPUConfig{
+				GPUMemGB: 0,
+				GPUCount: 1,
+			},
+			expected: 0.9,
+		},
+		{
+			name: "Zero GPU Count",
+			gpuConfig: &sku.GPUConfig{
+				GPUMemGB: 16,
+				GPUCount: 0,
+			},
+			expected: 0.9,
+		},
+		{
+			name: "V100",
+			gpuConfig: &sku.GPUConfig{
+				GPUMemGB: 16,
+				GPUCount: 1,
+			},
+			expected: 0.9,
+		},
+		{
+			name: "V100 x 2",
+			gpuConfig: &sku.GPUConfig{
+				GPUMemGB: 32,
+				GPUCount: 2,
+			},
+			expected: 0.9,
+		},
+		{
+			name: "A100",
+			gpuConfig: &sku.GPUConfig{
+				GPUMemGB: 80,
+				GPUCount: 1,
+			},
+			expected: 0.95,
+		},
+		{
+			name: "Invalid",
+			gpuConfig: &sku.GPUConfig{
+				GPUMemGB: 1,
+				GPUCount: 1,
+			},
+			expected: 0.9,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := getGPUMemoryUtilForVLLM(tt.gpuConfig)
+			if result != tt.expected {
+				t.Errorf("expected %v, got %v", tt.expected, result)
+			}
+		})
+	}
+}
diff --git a/pkg/utils/common_test.go b/pkg/utils/common_test.go
index b214150..a1920a0 100644
--- a/pkg/utils/common_test.go
+++ b/pkg/utils/common_test.go
@@ -51,7 +51,7 @@ func TestFetchGPUCountFromNodes(t *testing.T) {
 		name        string
 		nodeNames   []string
 		nodes       []runtime.Object
-		expectedGPU string
+		expectedGPU int
 		expectErr   bool
 		expectedErr string
 	}{
@@ -59,25 +59,25 @@ func TestFetchGPUCountFromNodes(t *testing.T) {
 			name:        "Single Node with GPU",
 			nodeNames:   []string{"node-1"},
 			nodes:       []runtime.Object{node1},
-			expectedGPU: "2",
+			expectedGPU: 2,
 		},
 		{
 			name:        "Multiple Nodes with GPU",
 			nodeNames:   []string{"node-1", "node-2"},
 			nodes:       []runtime.Object{node1, node2},
-			expectedGPU: "2",
+			expectedGPU: 2,
 		},
 		{
 			name:        "Node without GPU",
 			nodeNames:   []string{"node-3"},
 			nodes:       []runtime.Object{node3},
-			expectedGPU: "",
+			expectedGPU: 0,
 		},
 		{
 			name:        "No Worker Nodes",
 			nodeNames:   []string{},
 			nodes:       []runtime.Object{},
-			expectedGPU: "",
+			expectedGPU: 0,
 			expectErr:   true,
 			expectedErr: "no worker nodes found in the workspace",
 		},
diff --git a/pkg/workspace/inference/preset-inferences_test.go b/pkg/workspace/inference/preset-inferences_test.go
index 324f83d..b5a6774 100644
--- a/pkg/workspace/inference/preset-inferences_test.go
+++ b/pkg/workspace/inference/preset-inferences_test.go
@@ -46,7 +46,7 @@ func TestCreatePresetInference(t *testing.T) {
 			workload: "Deployment",
 			// No BaseCommand, TorchRunParams, TorchRunRdzvParams, or ModelRunParams
 			// So expected cmd consists of shell command and inference file
-			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --tensor-parallel-size=2 --served-model-name=mymodel --kaito-config-file=/mnt/config/inference_config.yaml",
+			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --tensor-parallel-size=2 --served-model-name=mymodel --gpu-memory-utilization=0.90 --kaito-config-file=/mnt/config/inference_config.yaml",
 			hasAdapters: false,
 		},
 
@@ -61,7 +61,7 @@ func TestCreatePresetInference(t *testing.T) {
 			workload: "Deployment",
 			// No BaseCommand, TorchRunParams, TorchRunRdzvParams, or ModelRunParams
 			// So expected cmd consists of shell command and inference file
-			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml",
+			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml --gpu-memory-utilization=0.90",
 			hasAdapters: false,
 		},
 
@@ -76,7 +76,7 @@ func TestCreatePresetInference(t *testing.T) {
 			workload: "Deployment",
 			// No BaseCommand, TorchRunParams, TorchRunRdzvParams, or ModelRunParams
 			// So expected cmd consists of shell command and inference file
-			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml",
+			expectedCmd: "/bin/sh -c python3 /workspace/vllm/inference_api.py --kaito-config-file=/mnt/config/inference_config.yaml --gpu-memory-utilization=0.90",
 			hasAdapters: false,
 		},
 
@@ -89,7 +89,7 @@ func TestCreatePresetInference(t *testing.T) {
 				c.On("Create", mock.IsType(context.TODO()), mock.IsType(&appsv1.Deployment{}), mock.Anything).Return(nil)
 			},
 			workload:       "Deployment",
-			expectedCmd:    "/bin/sh -c python3 /workspace/vllm/inference_api.py --enable-lora --tensor-parallel-size=2 --served-model-name=mymodel --kaito-config-file=/mnt/config/inference_config.yaml",
+			expectedCmd:    "/bin/sh -c python3 /workspace/vllm/inference_api.py --enable-lora --tensor-parallel-size=2 --served-model-name=mymodel --gpu-memory-utilization=0.90 --kaito-config-file=/mnt/config/inference_config.yaml",
 			hasAdapters:    true,
 			expectedVolume: "adapter-volume",
 		},
diff --git a/pkg/workspace/tuning/preset-tuning_test.go b/pkg/workspace/tuning/preset-tuning_test.go
index bffa67b..c55db9a 100644
--- a/pkg/workspace/tuning/preset-tuning_test.go
+++ b/pkg/workspace/tuning/preset-tuning_test.go
@@ -304,10 +304,10 @@ func TestPrepareTuningParameters(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			t.Setenv("CLOUD_PROVIDER", consts.AzureCloudName)
 
-			commands, resources := prepareTuningParameters(ctx, tc.workspaceObj, tc.modelCommand, tc.tuningObj, "2")
+			commands, resources := prepareTuningParameters(ctx, tc.workspaceObj, tc.modelCommand, tc.tuningObj, 2)
 			assert.Equal(t, tc.expectedCommands, commands)
-			assert.Equal(t, tc.expectedRequirements.Requests, resources.Requests)
-			assert.Equal(t, tc.expectedRequirements.Limits, resources.Limits)
+			assert.True(t, tc.expectedRequirements.Requests.Name("nvidia.com/gpu", resource.DecimalSI).Equal(*resources.Requests.Name("nvidia.com/gpu", resource.DecimalSI)))
+			assert.True(t, tc.expectedRequirements.Limits.Name("nvidia.com/gpu", resource.DecimalSI).Equal(*resources.Limits.Name("nvidia.com/gpu", resource.DecimalSI)))
 		})
 	}
 }