From 71f7efaa5172425a088479f67c75bc44694777f7 Mon Sep 17 00:00:00 2001
From: zylxjtu <zhang.yuanliang@hotmail.com>
Date: Thu, 20 Feb 2025 18:57:26 +0000
Subject: [PATCH] Replace the hardcoded SKU

---
 e2e/aks_model.go   | 2 +-
 e2e/kube.go        | 3 +++
 e2e/node_config.go | 4 ++--
 e2e/vmss.go        | 2 +-
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/e2e/aks_model.go b/e2e/aks_model.go
index 0058b2d4c8a..b65f3719bf7 100644
--- a/e2e/aks_model.go
+++ b/e2e/aks_model.go
@@ -46,7 +46,7 @@ func getBaseClusterModel(clusterName string) *armcontainerservice.ManagedCluster
 				{
 					Name:         to.Ptr("nodepool1"),
 					Count:        to.Ptr[int32](1),
-					VMSize:       to.Ptr("standard_d2ds_v5"),
+					VMSize:       to.Ptr(config.Config.DefaultVMSKU),
 					MaxPods:      to.Ptr[int32](110),
 					OSType:       to.Ptr(armcontainerservice.OSTypeLinux),
 					Type:         to.Ptr(armcontainerservice.AgentPoolTypeVirtualMachineScaleSets),
diff --git a/e2e/kube.go b/e2e/kube.go
index fa6536ffe78..7b7acdbbfb9 100644
--- a/e2e/kube.go
+++ b/e2e/kube.go
@@ -619,6 +619,9 @@ func nvidiaDevicePluginDaemonSet() *appsv1.DaemonSet {
 					},
 				},
 				Spec: corev1.PodSpec{
+					NodeSelector: map[string]string{
+						"kubernetes.io/os": "linux",
+					},
 					Tolerations: []corev1.Toleration{
 						{
 							Key:      "sku",
diff --git a/e2e/node_config.go b/e2e/node_config.go
index 32b95655f46..01b12d9b7d7 100644
--- a/e2e/node_config.go
+++ b/e2e/node_config.go
@@ -60,7 +60,7 @@ func nbcToAKSNodeConfigV1(nbc *datamodel.NodeBootstrappingConfiguration) *aksnod
 		Version:            "v0",
 		DisableCustomData:  false,
 		LinuxAdminUsername: "azureuser",
-		VmSize:             "Standard_D2ds_v5",
+		VmSize:             config.Config.DefaultVMSKU,
 		ClusterConfig: &aksnodeconfigv1.ClusterConfig{
 			Location:      nbc.ContainerService.Location,
 			ResourceGroup: nbc.ResourceGroupName,
@@ -347,7 +347,7 @@ func baseTemplateLinux(t *testing.T, location string, k8sVersion string, arch st
 		},
 		AgentPoolProfile: &datamodel.AgentPoolProfile{
 			Name:                "nodepool2",
-			VMSize:              "Standard_D2ds_v5",
+			VMSize:              config.Config.DefaultVMSKU,
 			KubeletDiskType:     "",
 			WorkloadRuntime:     "",
 			DNSPrefix:           "",
diff --git a/e2e/vmss.go b/e2e/vmss.go
index 31d6fe17c7a..e99de0ceb52 100644
--- a/e2e/vmss.go
+++ b/e2e/vmss.go
@@ -427,7 +427,7 @@ func getBaseVMSSModel(s *Scenario, customData, cseCmd string) armcompute.Virtual
 	model := armcompute.VirtualMachineScaleSet{
 		Location: to.Ptr(config.Config.Location),
 		SKU: &armcompute.SKU{
-			Name:     to.Ptr("Standard_D2ds_v5"),
+			Name:     to.Ptr(config.Config.DefaultVMSKU),
 			Capacity: to.Ptr[int64](1),
 		},
 		Properties: &armcompute.VirtualMachineScaleSetProperties{