Configure kubelet labels and taints dynamicly

Pulls the values from the same tags used by the cluster-autoscaler.

See: https://github.com/kubernetes/autoscaler/blob/master/cluster-autoscaler/cloudprovider/aws/README.md#scaling-a-node-group-to-0

pkg/node/node.go

@@ -36,6 +36,8 @@ type Node struct {
 	ReservedMemory string
 	ClusterDNS     string
 	Region         string
+	Labels         []string
+	Taints         []string
 }
 
 type metadataClient interface {
@@ -67,6 +69,8 @@ func New(e ec2iface.EC2API, m metadataClient, region *string) (*Node, error) {
 			ReservedMemory: reservedMemory(instance.InstanceType),
 			ClusterDNS:     clusterDNS(instance.PrivateIpAddress),
 			Region:         *region,
+			Labels:         lables(instance.Tags),
+			Taints:         taints(instance.Tags),
 		}
 		if node.ClusterName() == "" {
 			sleepFor := b.Duration(tries)
@@ -92,6 +96,28 @@ func (n *Node) ClusterName() string {
 	return ""
 }
 
+func lables(tags []*ec2.Tag) []string {
+	var l []string
+	re := regexp.MustCompile(`k8s.io\/cluster-autoscaler\/node-template\/label\/(.*)`)
+	for _, t := range tags {
+		if matches := re.FindStringSubmatch(*t.Key); len(matches) == 2 {
+			l = append(l, matches[1]+"="+*t.Value)
+		}
+	}
+	return l
+}
+
+func taints(tags []*ec2.Tag) []string {
+	var ts []string
+	re := regexp.MustCompile(`k8s.io\/cluster-autoscaler\/node-template\/taint\/(.*)`)
+	for _, t := range tags {
+		if matches := re.FindStringSubmatch(*t.Key); len(matches) == 2 {
+			ts = append(ts, matches[1]+"="+*t.Value)
+		}
+	}
+	return ts
+}
+
 func instanceID(m metadataClient) (*string, error) {
 	result, err := m.GetMetadata("instance-id")
 	if err != nil {

pkg/node/node_test.go

@@ -18,6 +18,7 @@ package node
 
 import (
 	"errors"
+	"reflect"
 	"testing"
 
 	"github.com/errm/ekstrap/pkg/backoff"
@@ -63,6 +64,72 @@ func TestNewNode(t *testing.T) {
 	}
 }
 
+func TestNodeLabels(t *testing.T) {
+	e := &mockEC2{
+		tags: [][]*ec2.Tag{
+			{},
+			{},
+			{
+				tag("kubernetes.io/cluster/cluster-name", "owned"),
+				tag("k8s.io/cluster-autoscaler/node-template/label/node-role.kubernetes.io/spot-worker", "true"),
+				tag("k8s.io/cluster-autoscaler/node-template/label/nvidia-gpu", "K80"),
+			},
+		},
+	}
+	metadata := mockMetadata{
+		data: map[string]string{
+			"instance-id": "1234",
+		},
+	}
+	region := "us-east-1"
+	node, err := New(e, metadata, &region)
+	if err != nil {
+		t.Errorf("unexpected error: %s", err)
+	}
+
+	expected := []string{
+		"node-role.kubernetes.io/spot-worker=true",
+		"nvidia-gpu=K80",
+	}
+
+	if !reflect.DeepEqual(node.Labels, expected) {
+		t.Errorf("Expected node.Labels to be %v but was %v", expected, node.Labels)
+	}
+}
+
+func TestNodeTaints(t *testing.T) {
+	e := &mockEC2{
+		tags: [][]*ec2.Tag{
+			{},
+			{},
+			{
+				tag("kubernetes.io/cluster/cluster-name", "owned"),
+				tag("k8s.io/cluster-autoscaler/node-template/label/foo", "bar"),
+				tag("k8s.io/cluster-autoscaler/node-template/taint/dedicated", "foo:NoSchedule"),
+				tag("k8s.io/cluster-autoscaler/node-template/label/nvidia-gpu", "K80"),
+			},
+		},
+	}
+	metadata := mockMetadata{
+		data: map[string]string{
+			"instance-id": "1234",
+		},
+	}
+	region := "us-east-1"
+	node, err := New(e, metadata, &region)
+	if err != nil {
+		t.Errorf("unexpected error: %s", err)
+	}
+
+	expected := []string{
+		"dedicated=foo:NoSchedule",
+	}
+
+	if !reflect.DeepEqual(node.Taints, expected) {
+		t.Errorf("Expected node.Taints to be %v but was %v", expected, node.Taints)
+	}
+}
+
 func TestClusterDNS(t *testing.T) {
 	e := &mockEC2{
 		PrivateIPAddress: "10.1.123.4",

pkg/system/system_test.go

@@ -33,7 +33,7 @@ func TestConfigure(t *testing.T) {
 	hn := &FakeHostname{}
 	init := &FakeInit{}
 
-	i := instance("10.6.28.199", "ip-10-6-28-199.us-west-2.compute.internal", 18, "60m", "960Mi", false)
+	i := instance("10.6.28.199", "ip-10-6-28-199.us-west-2.compute.internal", 18, "60m", "960Mi", []string{}, []string{})
 	c := cluster(
 		"aws-om-cluster",
 		"https://74770F6B05F7A8FB0F02CFB5F7AF530C.yl4.us-west-2.eks.amazonaws.com",
@@ -46,8 +46,8 @@ func TestConfigure(t *testing.T) {
 		t.Errorf("unexpected error %v", err)
 	}
 
-	if len(fs.files) != 6 {
-		t.Errorf("expected 6 files, got %v", len(fs.files))
+	if len(fs.files) != 8 {
+		t.Errorf("expected 8 files, got %v", len(fs.files))
 	}
 
 	expected := `apiVersion: v1
@@ -100,7 +100,7 @@ ExecStart=/usr/bin/kubelet \
   --kubeconfig=/var/lib/kubelet/kubeconfig \
   --feature-gates=RotateKubeletServerCertificate=true \
   --anonymous-auth=false \
-  --client-ca-file=/etc/kubernetes/pki/ca.crt $KUBELET_ARGS $KUBELET_MAX_PODS $KUBELET_KUBE_RESERVED $KUBELET_EXTRA_ARGS
+  --client-ca-file=/etc/kubernetes/pki/ca.crt $KUBELET_ARGS $KUBELET_MAX_PODS $KUBELET_KUBE_RESERVED $KUBELET_NODE_LABELS $KUBELET_NODE_TAINTS $KUBELET_EXTRA_ARGS
 
 Restart=always
 StartLimitInterval=0
@@ -126,6 +126,12 @@ Environment='KUBELET_KUBE_RESERVED=--kube-reserved=cpu=60m,memory=960Mi'
 `
 	fs.Check(t, "/etc/systemd/system/kubelet.service.d/30-kube-reserved.conf", expected, 0640)
 
+	expected = `[Service]`
+	fs.Check(t, "/etc/systemd/system/kubelet.service.d/40-labels.conf", expected, 0640)
+
+	expected = `[Service]`
+	fs.Check(t, "/etc/systemd/system/kubelet.service.d/50-taints.conf", expected, 0640)
+
 	expected = `thisisthecertdata
 `
 	fs.Check(t, "/etc/kubernetes/pki/ca.crt", expected, 0640)
@@ -148,7 +154,7 @@ func TestConfigureNoReserved(t *testing.T) {
 	hn := &FakeHostname{}
 	init := &FakeInit{}
 
-	i := instance("10.6.28.199", "ip-10-6-28-199.us-west-2.compute.internal", 18, "", "", false)
+	i := instance("10.6.28.199", "ip-10-6-28-199.us-west-2.compute.internal", 18, "", "", []string{}, []string{})
 	c := cluster(
 		"aws-om-cluster",
 		"https://74770F6B05F7A8FB0F02CFB5F7AF530C.yl4.us-west-2.eks.amazonaws.com",
@@ -165,7 +171,92 @@ func TestConfigureNoReserved(t *testing.T) {
 	fs.Check(t, "/etc/systemd/system/kubelet.service.d/30-kube-reserved.conf", expected, 0640)
 }
 
-func instance(ip, dnsName string, maxPods int, reservedCPU, reservedMemory string, spot bool) *node.Node {
+func TestConfigureLabels(t *testing.T) {
+	fs := &FakeFileSystem{}
+	hn := &FakeHostname{}
+	init := &FakeInit{}
+
+	labels := []string{
+		"node-role.kubernetes.io/worker=true",
+	}
+
+	i := instance("10.6.28.199", "ip-10-6-28-199.us-west-2.compute.internal", 18, "", "", labels, []string{})
+	c := cluster(
+		"aws-om-cluster",
+		"https://74770F6B05F7A8FB0F02CFB5F7AF530C.yl4.us-west-2.eks.amazonaws.com",
+		"dGhpc2lzdGhlY2VydGRhdGE=",
+	)
+	system := System{Filesystem: fs, Hostname: hn, Init: init}
+	err := system.Configure(i, c)
+
+	if err != nil {
+		t.Errorf("unexpected error %v", err)
+	}
+
+	expected := `[Service]
+Environment='KUBELET_NODE_LABELS=--node-labels="node-role.kubernetes.io/worker=true"'
+`
+	fs.Check(t, "/etc/systemd/system/kubelet.service.d/40-labels.conf", expected, 0640)
+}
+
+func TestConfigureMultipleLabels(t *testing.T) {
+	fs := &FakeFileSystem{}
+	hn := &FakeHostname{}
+	init := &FakeInit{}
+
+	labels := []string{
+		"node-role.kubernetes.io/worker=true",
+		"gpu-type=K80",
+	}
+
+	i := instance("10.6.28.199", "ip-10-6-28-199.us-west-2.compute.internal", 18, "", "", labels, []string{})
+	c := cluster(
+		"aws-om-cluster",
+		"https://74770F6B05F7A8FB0F02CFB5F7AF530C.yl4.us-west-2.eks.amazonaws.com",
+		"dGhpc2lzdGhlY2VydGRhdGE=",
+	)
+	system := System{Filesystem: fs, Hostname: hn, Init: init}
+	err := system.Configure(i, c)
+
+	if err != nil {
+		t.Errorf("unexpected error %v", err)
+	}
+
+	expected := `[Service]
+Environment='KUBELET_NODE_LABELS=--node-labels="node-role.kubernetes.io/worker=true,gpu-type=K80"'
+`
+	fs.Check(t, "/etc/systemd/system/kubelet.service.d/40-labels.conf", expected, 0640)
+}
+
+func TestConfigureTaints(t *testing.T) {
+	fs := &FakeFileSystem{}
+	hn := &FakeHostname{}
+	init := &FakeInit{}
+
+	taints := []string{
+		"node-role.kubernetes.io/worker=true:PreferNoSchedule",
+	}
+
+	i := instance("10.6.28.199", "ip-10-6-28-199.us-west-2.compute.internal", 18, "", "", []string{}, taints)
+	c := cluster(
+		"aws-om-cluster",
+		"https://74770F6B05F7A8FB0F02CFB5F7AF530C.yl4.us-west-2.eks.amazonaws.com",
+		"dGhpc2lzdGhlY2VydGRhdGE=",
+	)
+	system := System{Filesystem: fs, Hostname: hn, Init: init}
+	err := system.Configure(i, c)
+
+	if err != nil {
+		t.Errorf("unexpected error %v", err)
+	}
+
+	expected := `[Service]
+Environment='KUBELET_NODE_TAINTS=--register-with-taints="node-role.kubernetes.io/worker=true:PreferNoSchedule"'
+`
+	fs.Check(t, "/etc/systemd/system/kubelet.service.d/50-taints.conf", expected, 0640)
+}
+
+func instance(ip, dnsName string, maxPods int, reservedCPU, reservedMemory string, labels, taints []string) *node.Node {
 	return &node.Node{
 		Instance: &ec2.Instance{
 			PrivateIpAddress: &ip,
@@ -176,6 +267,8 @@ func instance(ip, dnsName string, maxPods int, reservedCPU, reservedMemory strin
 		Region:         "us-east-1",
 		ReservedCPU:    reservedCPU,
 		ReservedMemory: reservedMemory,
+		Labels:         labels,
+		Taints:         taints,
 	}
 }
 

pkg/system/templates/etc/systemd/system/kubelet.service

@@ -22,7 +22,7 @@ ExecStart=/usr/bin/kubelet \
   --kubeconfig=/var/lib/kubelet/kubeconfig \
   --feature-gates=RotateKubeletServerCertificate=true \
   --anonymous-auth=false \
-  --client-ca-file=/etc/kubernetes/pki/ca.crt $KUBELET_ARGS $KUBELET_MAX_PODS $KUBELET_KUBE_RESERVED $KUBELET_EXTRA_ARGS
+  --client-ca-file=/etc/kubernetes/pki/ca.crt $KUBELET_ARGS $KUBELET_MAX_PODS $KUBELET_KUBE_RESERVED $KUBELET_NODE_LABELS $KUBELET_NODE_TAINTS $KUBELET_EXTRA_ARGS
 
 Restart=always
 StartLimitInterval=0

pkg/system/templates/etc/systemd/system/kubelet.service.d/40-labels.conf

@@ -0,0 +1,4 @@
+[Service]
+{{- if .Node.Labels }}
+Environment='KUBELET_NODE_LABELS=--node-labels="{{ range $index, $label := .Node.Labels }}{{ if $index }},{{ end }}{{ $label }}{{ end }}"'
+{{ end -}}

pkg/system/templates/etc/systemd/system/kubelet.service.d/50-taints.conf

@@ -0,0 +1,4 @@
+[Service]
+{{- if .Node.Taints }}
+Environment='KUBELET_NODE_TAINTS=--register-with-taints="{{ range $index, $taint := .Node.Taints }}{{ if $index }},{{ end }}{{ $taint }}{{ end }}"'
+{{ end -}}