Merge pull request #172 from aleksandra-malinowska/add-taints-to-node-templates

add taints to GCE node template

cluster-autoscaler/cloudprovider/gce/gce_manager.go

@@ -26,7 +26,7 @@ import (
 	"sync"
 	"time"
 
-	"gopkg.in/gcfg.v1"
+	gcfg "gopkg.in/gcfg.v1"
 
 	"github.com/golang/glog"
 	"golang.org/x/oauth2"
@@ -327,7 +327,7 @@ func (m *GceManager) buildNodeFromTemplate(mig *Mig, template *gce.InstanceTempl
 	// TODO: use proper allocatable!!
 	node.Status.Allocatable = node.Status.Capacity
 
-	// KubeEnvLabels
+	// KubeEnv labels & taints
 	if template.Properties.Metadata == nil {
 		return nil, fmt.Errorf("instance template %s has no metadata", template.Name)
 	}
@@ -336,11 +336,18 @@ func (m *GceManager) buildNodeFromTemplate(mig *Mig, template *gce.InstanceTempl
 			if item.Value == nil {
 				return nil, fmt.Errorf("no kube-env content in metadata")
 			}
+			// Extract labels
 			kubeEnvLabels, err := extractLabelsFromKubeEnv(*item.Value)
 			if err != nil {
 				return nil, err
 			}
 			node.Labels = cloudprovider.JoinStringMaps(node.Labels, kubeEnvLabels)
+			// Extract taints
+			kubeEnvTaints, err := extractTaintsFromKubeEnv(*item.Value)
+			if err != nil {
+				return nil, err
+			}
+			node.Spec.Taints = append(node.Spec.Taints, kubeEnvTaints...)
 		}
 	}
 	// GenericLabels
@@ -390,6 +397,18 @@ func parseCustomMachineType(machineType string) (cpu, mem int64, err error) {
 }
 
 func extractLabelsFromKubeEnv(kubeEnv string) (map[string]string, error) {
+	return extractFromKubeEnv(kubeEnv, "NODE_LABELS")
+}
+
+func extractTaintsFromKubeEnv(kubeEnv string) ([]apiv1.Taint, error) {
+	taintMap, err := extractFromKubeEnv(kubeEnv, "NODE_TAINTS")
+	if err != nil {
+		return nil, err
+	}
+	return buildTaints(taintMap)
+}
+
+func extractFromKubeEnv(kubeEnv, resource string) (map[string]string, error) {
 	result := make(map[string]string)
 
 	for line, env := range strings.Split(kubeEnv, "\n") {
@@ -403,15 +422,31 @@ func extractLabelsFromKubeEnv(kubeEnv string) (map[string]string, error) {
 		}
 		key := strings.Trim(items[0], " ")
 		value := strings.Trim(items[1], " \"'")
-		if key == "NODE_LABELS" {
-			for _, label := range strings.Split(value, ",") {
-				labelItems := strings.SplitN(label, "=", 2)
-				if len(labelItems) != 2 {
-					return nil, fmt.Errorf("error while parsing label: %s", label)
+		if key == resource {
+			for _, val := range strings.Split(value, ",") {
+				valItems := strings.SplitN(val, "=", 2)
+				if len(valItems) != 2 {
+					return nil, fmt.Errorf("error while parsing kube env value: %s", val)
 				}
-				result[labelItems[0]] = labelItems[1]
+				result[valItems[0]] = valItems[1]
 			}
 		}
 	}
 	return result, nil
 }
+
+func buildTaints(kubeEnvTaints map[string]string) ([]apiv1.Taint, error) {
+	taints := make([]apiv1.Taint, 0)
+	for key, value := range kubeEnvTaints {
+		values := strings.SplitN(value, ":", 2)
+		if len(values) != 2 {
+			return nil, fmt.Errorf("error while parsing node taint value and effect: %s", value)
+		}
+		taints = append(taints, apiv1.Taint{
+			Key:    key,
+			Value:  values[0],
+			Effect: apiv1.TaintEffect(values[1]),
+		})
+	}
+	return taints, nil
+}

cluster-autoscaler/cloudprovider/gce/gce_manager_test.go

@@ -20,6 +20,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	apiv1 "k8s.io/api/core/v1"
 	"k8s.io/autoscaler/cluster-autoscaler/cloudprovider"
 	kubeletapis "k8s.io/kubernetes/pkg/kubelet/apis"
 )
@@ -53,6 +54,45 @@ func TestExtractLabelsFromKubeEnv(t *testing.T) {
 	assert.Equal(t, "true", labels["cloud.google.com/gke-preemptible"])
 }
 
+func TestExtractTaintsFromKubeEnv(t *testing.T) {
+	kubeenv := "ENABLE_NODE_PROBLEM_DETECTOR: 'daemonset'\n" +
+		"NODE_LABELS: a=b,c=d,cloud.google.com/gke-nodepool=pool-3,cloud.google.com/gke-preemptible=true\n" +
+		"DNS_SERVER_IP: '10.0.0.10'\n" +
+		"NODE_TAINTS: 'dedicated=ml:NoSchedule,test=dev:PreferNoSchedule,a=b:c'\n"
+
+	expectedTaints := []apiv1.Taint{
+		{
+			Key:    "dedicated",
+			Value:  "ml",
+			Effect: apiv1.TaintEffectNoSchedule,
+		},
+		{
+			Key:    "test",
+			Value:  "dev",
+			Effect: apiv1.TaintEffectPreferNoSchedule,
+		},
+		{
+			Key:    "a",
+			Value:  "b",
+			Effect: apiv1.TaintEffect("c"),
+		},
+	}
+
+	taints, err := extractTaintsFromKubeEnv(kubeenv)
+	assert.Nil(t, err)
+	assert.Equal(t, 3, len(taints))
+	assert.Equal(t, makeTaintSet(expectedTaints), makeTaintSet(taints))
+
+}
+
+func makeTaintSet(taints []apiv1.Taint) map[apiv1.Taint]bool {
+	set := make(map[apiv1.Taint]bool)
+	for _, taint := range taints {
+		set[taint] = true
+	}
+	return set
+}
+
 func TestParseCustomMachineType(t *testing.T) {
 	cpu, mem, err := parseCustomMachineType("custom-2-2816")
 	assert.NoError(t, err)

cluster-autoscaler/core/utils.go

@@ -293,7 +293,7 @@ func sanitizeTemplateNode(node *apiv1.Node, nodeGroup string) (*apiv1.Node, erro
 	return newNode, nil
 }
 
-// Removes unregisterd nodes if needed. Returns true if anything was removed and error if such occurred.
+// Removes unregistered nodes if needed. Returns true if anything was removed and error if such occurred.
 func removeOldUnregisteredNodes(unregisteredNodes []clusterstate.UnregisteredNode, context *AutoscalingContext,
 	currentTime time.Time) (bool, error) {
 	removedAny := false