Add the missing labels/taints for VMs pool

Author: wenxuan0923

cluster-autoscaler/cloudprovider/azure/azure_template.go

@@ -89,7 +89,7 @@ const (
 // VMPoolNodeTemplate holds properties for node from VMPool
 type VMPoolNodeTemplate struct {
 	AgentPoolName string
-	Taints        []string
+	Taints        []apiv1.Taint
 	Labels        map[string]*string
 	OSDiskType    *armcontainerservice.OSDiskType
 }
@@ -155,23 +155,33 @@ func buildNodeTemplateFromVMSS(vmss compute.VirtualMachineScaleSet, inputLabels
 	}, nil
 }
 
-func buildNodeTemplateFromVMPool(vmsPool armcontainerservice.AgentPool, location string, skuName string) (NodeTemplate, error) {
+func buildNodeTemplateFromVMPool(vmsPool armcontainerservice.AgentPool, location string, skuName string, labelsFromSpec map[string]string, taintsFromSpec string) (NodeTemplate, error) {
 	if vmsPool.Properties == nil {
 		return NodeTemplate{}, fmt.Errorf("vmsPool %s has nil properties", to.String(vmsPool.Name))
 	}
-	var labels map[string]*string
-	if vmsPool.Properties.NodeLabels != nil {
-		labels = vmsPool.Properties.NodeLabels
+	// labels from the agentpool
+	labels := vmsPool.Properties.NodeLabels
+	// labels from spec
+	for k, v := range labelsFromSpec {
+		if labels == nil {
+			labels = make(map[string]*string)
+		}
+		labels[k] = to.StringPtr(v)
 	}
 
-	var taints []string
-	if vmsPool.Properties.NodeTaints != nil {
-		for _, taint := range vmsPool.Properties.NodeTaints {
-			if taint != nil {
-				taints = append(taints, *taint)
-			}
+	// taints from the agentpool
+	taintsList := []string{}
+	for _, taint := range vmsPool.Properties.NodeTaints {
+		if to.String(taint) != "" {
+			taintsList = append(taintsList, to.String(taint))
 		}
 	}
+	// taints from spec
+	if taintsFromSpec != "" {
+		taintsList = append(taintsList, taintsFromSpec)
+	}
+	taintsStr := strings.Join(taintsList, ",")
+	taints := extractTaintsFromSpecString(taintsStr)
 
 	var zones []string
 	if vmsPool.Properties.AvailabilityZones != nil {
@@ -284,8 +294,7 @@ func processVMPoolTemplate(template NodeTemplate, nodeName string, node apiv1.No
 		}
 	}
 	node.Labels = cloudprovider.JoinStringMaps(node.Labels, labels)
-	taints := buildNodeTaintsForVMPool(template.VMPoolNodeTemplate.Taints)
-	node.Spec.Taints = taints
+	node.Spec.Taints = template.VMPoolNodeTemplate.Taints
 	return node
 }
 
@@ -459,14 +468,19 @@ func extractTaintsFromTags(tags map[string]*string) []apiv1.Taint {
 	return taints
 }
 
-// extractTaintsFromSpecString is for VMSS nodepool taints
+// extractTaintsFromSpecString is for nodepool taints
 // Example of a valid taints string, is the same argument to kubelet's `--register-with-taints`
 // "dedicated=foo:NoSchedule,group=bar:NoExecute,app=fizz:PreferNoSchedule"
 func extractTaintsFromSpecString(taintsString string) []apiv1.Taint {
 	taints := make([]apiv1.Taint, 0)
+	dedupMap := make(map[string]interface{})
 	// First split the taints at the separator
 	splits := strings.Split(taintsString, ",")
 	for _, split := range splits {
+		if dedupMap[split] != nil {
+			continue
+		}
+		dedupMap[split] = struct{}{}
 		valid, taint := constructTaintFromString(split)
 		if valid {
 			taints = append(taints, taint)

cluster-autoscaler/cloudprovider/azure/azure_template_test.go

@@ -21,6 +21,7 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/resourcemanager/containerservice/armcontainerservice/v5"
 	"github.com/Azure/azure-sdk-for-go/services/compute/mgmt/2022-08-01/compute"
 	"github.com/Azure/go-autorest/autorest"
 	"github.com/Azure/go-autorest/autorest/to"
@@ -137,6 +138,11 @@ func TestExtractTaintsFromSpecString(t *testing.T) {
 			Value:  "fizz",
 			Effect: apiv1.TaintEffectPreferNoSchedule,
 		},
+		{
+			Key:    "dedicated", // duplicate key, should be ignored
+			Value:  "foo",
+			Effect: apiv1.TaintEffectNoSchedule,
+		},
 	}
 
 	taints := extractTaintsFromSpecString(strings.Join(taintsString, ","))
@@ -222,6 +228,61 @@ func TestEmptyTopologyFromScaleSet(t *testing.T) {
 	assert.True(t, ok)
 	assert.Equal(t, expectedAzureDiskTopology, azureDiskTopology)
 }
+func TestBuildNodeTemplateFromVMPool(t *testing.T) {
+	agentPoolName := "testpool"
+	location := "eastus"
+	skuName := "Standard_DS2_v2"
+	labelKey := "foo"
+	labelVal := "bar"
+	taintStr := "dedicated=foo:NoSchedule,boo=fizz:PreferNoSchedule,group=bar:NoExecute"
+
+	osType := armcontainerservice.OSTypeLinux
+	osDiskType := armcontainerservice.OSDiskTypeEphemeral
+	zone1 := "1"
+	zone2 := "2"
+
+	vmpool := armcontainerservice.AgentPool{
+		Name: to.StringPtr(agentPoolName),
+		Properties: &armcontainerservice.ManagedClusterAgentPoolProfileProperties{
+			NodeLabels: map[string]*string{
+				"existing":   to.StringPtr("label"),
+				"department": to.StringPtr("engineering"),
+			},
+			NodeTaints:        []*string{to.StringPtr("group=bar:NoExecute")},
+			OSType:            &osType,
+			OSDiskType:        &osDiskType,
+			AvailabilityZones: []*string{&zone1, &zone2},
+		},
+	}
+
+	labelsFromSpec := map[string]string{labelKey: labelVal}
+	taintsFromSpec := taintStr
+
+	template, err := buildNodeTemplateFromVMPool(vmpool, location, skuName, labelsFromSpec, taintsFromSpec)
+	assert.NoError(t, err)
+	assert.Equal(t, skuName, template.SkuName)
+	assert.Equal(t, location, template.Location)
+	assert.ElementsMatch(t, []string{zone1, zone2}, template.Zones)
+	assert.Equal(t, "linux", template.InstanceOS)
+	assert.NotNil(t, template.VMPoolNodeTemplate)
+	assert.Equal(t, agentPoolName, template.VMPoolNodeTemplate.AgentPoolName)
+	assert.Equal(t, &osDiskType, template.VMPoolNodeTemplate.OSDiskType)
+	// Labels: should include both from NodeLabels and labelsFromSpec
+	assert.Contains(t, template.VMPoolNodeTemplate.Labels, "existing")
+	assert.Equal(t, "label", *template.VMPoolNodeTemplate.Labels["existing"])
+	assert.Contains(t, template.VMPoolNodeTemplate.Labels, "department")
+	assert.Equal(t, "engineering", *template.VMPoolNodeTemplate.Labels["department"])
+	assert.Contains(t, template.VMPoolNodeTemplate.Labels, labelKey)
+	assert.Equal(t, labelVal, *template.VMPoolNodeTemplate.Labels[labelKey])
+	// Taints: should include both from NodeTaints and taintsFromSpec
+	taintSet := makeTaintSet(template.VMPoolNodeTemplate.Taints)
+	expectedTaints := []apiv1.Taint{
+		{Key: "group", Value: "bar", Effect: apiv1.TaintEffectNoExecute},
+		{Key: "dedicated", Value: "foo", Effect: apiv1.TaintEffectNoSchedule},
+		{Key: "boo", Value: "fizz", Effect: apiv1.TaintEffectPreferNoSchedule},
+	}
+	assert.Equal(t, makeTaintSet(expectedTaints), taintSet)
+}
 
 func makeTaintSet(taints []apiv1.Taint) map[apiv1.Taint]bool {
 	set := make(map[apiv1.Taint]bool)

cluster-autoscaler/cloudprovider/azure/azure_vms_pool.go

@@ -460,7 +460,9 @@ func (vmPool *VMPool) TemplateNodeInfo() (*schedulerframework.NodeInfo, error) {
 		return nil, err
 	}
 
-	template, err := buildNodeTemplateFromVMPool(ap, vmPool.manager.config.Location, vmPool.sku)
+	inputLabels := map[string]string{}
+	inputTaints := ""
+	template, err := buildNodeTemplateFromVMPool(ap, vmPool.manager.config.Location, vmPool.sku, inputLabels, inputTaints)
 	if err != nil {
 		return nil, err
 	}