diff --git a/pkg/node/node.go b/pkg/node/node.go index 969e3fc..b842f02 100644 --- a/pkg/node/node.go +++ b/pkg/node/node.go @@ -36,7 +36,8 @@ type Node struct { ReservedMemory string ClusterDNS string Region string - Spot bool + Labels []string + Taints []string } type metadataClient interface { @@ -61,7 +62,16 @@ func New(e ec2iface.EC2API, m metadataClient, region *string) (*Node, error) { return nil, err } instance := output.Reservations[0].Instances[0] - node := Node{Instance: instance, MaxPods: maxPods(instance.InstanceType), ReservedCPU: reservedCPU(instance.InstanceType), ReservedMemory: reservedMemory(instance.InstanceType), ClusterDNS: clusterDNS(instance.PrivateIpAddress), Region: *region, Spot: spot(instance.InstanceLifecycle)} + node := Node{ + Instance: instance, + MaxPods: maxPods(instance.InstanceType), + ReservedCPU: reservedCPU(instance.InstanceType), + 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) log.Printf("The kubernetes.io/cluster/ tag is not yet set, will try again in %s", sleepFor) @@ -86,11 +96,26 @@ func (n *Node) ClusterName() string { return "" } -func spot(lifecycleType *string) bool { - if lifecycleType != nil && *lifecycleType == ec2.InstanceLifecycleTypeSpot { - return true +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 false + return ts } func instanceID(m metadataClient) (*string, error) { diff --git a/pkg/node/node_test.go b/pkg/node/node_test.go index 3e4ceeb..62b11cb 100644 --- a/pkg/node/node_test.go +++ b/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, ®ion) + 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, ®ion) + 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", @@ -360,57 +427,6 @@ func TestMemory(t *testing.T) { } } -func TestSpot(t *testing.T) { - tests := []struct { - lifecycleType string - expected bool - }{ - { - - lifecycleType: ec2.InstanceLifecycleTypeSpot, - expected: true, - }, - { - - lifecycleType: ec2.InstanceLifecycleTypeScheduled, - expected: false, - }, - { - // OnDemand instances do not return this field - lifecycleType: "", - expected: false, - }, - { - - lifecycleType: "something-unexpected", - expected: false, - }, - } - - for _, test := range tests { - e := &mockEC2{ - tags: [][]*ec2.Tag{ - {tag("kubernetes.io/cluster/cluster-name", "owned")}, - }, - lifecycleType: test.lifecycleType, - } - metadata := mockMetadata{ - data: map[string]string{ - "instance-id": "1234", - }, - } - region := "us-west-2" - node, err := New(e, metadata, ®ion) - if err != nil { - t.Errorf("unexpected error: %s", err) - } - - if node.Spot != test.expected { - t.Errorf("expected Spot for %v to be: %v, but it was %v", test.lifecycleType, test.expected, node.Spot) - } - } -} - func tag(key, value string) *ec2.Tag { return &ec2.Tag{ Key: &key, @@ -421,10 +437,9 @@ func tag(key, value string) *ec2.Tag { type mockEC2 struct { PrivateIPAddress string ec2iface.EC2API - tags [][]*ec2.Tag - instanceType string - err error - lifecycleType string + tags [][]*ec2.Tag + instanceType string + err error } func (m *mockEC2) DescribeInstances(input *ec2.DescribeInstancesInput) (*ec2.DescribeInstancesOutput, error) { @@ -434,20 +449,15 @@ func (m *mockEC2) DescribeInstances(input *ec2.DescribeInstancesInput) (*ec2.Des var tags []*ec2.Tag //Pop the first set of tags tags, m.tags = m.tags[0], m.tags[1:] - var lifecycleType *string - if m.lifecycleType != "" { - lifecycleType = &m.lifecycleType - } if len(input.InstanceIds) > 0 { return &ec2.DescribeInstancesOutput{ Reservations: []*ec2.Reservation{{ Instances: []*ec2.Instance{ { - InstanceId: input.InstanceIds[0], - Tags: tags, - InstanceType: &m.instanceType, - PrivateIpAddress: &m.PrivateIPAddress, - InstanceLifecycle: lifecycleType, + InstanceId: input.InstanceIds[0], + Tags: tags, + InstanceType: &m.instanceType, + PrivateIpAddress: &m.PrivateIPAddress, }, }, }, diff --git a/pkg/system/system_test.go b/pkg/system/system_test.go index 84d84ee..5a49eee 100644 --- a/pkg/system/system_test.go +++ b/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) != 7 { - t.Errorf("expected 7 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_SPOT_ARGS $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,10 +126,11 @@ 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] -Environment='KUBELET_SPOT_ARGS=--node-labels="node-role.kubernetes.io/worker=true" --register-with-taints="node-role.kubernetes.io/worker=true:PreferNoSchedule"' -` - fs.Check(t, "/etc/systemd/system/kubelet.service.d/40-spot-args.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 ` @@ -153,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", @@ -170,12 +171,73 @@ func TestConfigureNoReserved(t *testing.T) { fs.Check(t, "/etc/systemd/system/kubelet.service.d/30-kube-reserved.conf", expected, 0640) } -func TestConfigureSpotInstance(t *testing.T) { +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{} - i := instance("10.6.28.199", "ip-10-6-28-199.us-west-2.compute.internal", 18, "", "", true) + 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", @@ -189,12 +251,12 @@ func TestConfigureSpotInstance(t *testing.T) { } expected := `[Service] -Environment='KUBELET_SPOT_ARGS=--node-labels="node-role.kubernetes.io/spot-worker=true"' +Environment='KUBELET_NODE_TAINTS=--register-with-taints="node-role.kubernetes.io/worker=true:PreferNoSchedule"' ` - fs.Check(t, "/etc/systemd/system/kubelet.service.d/40-spot-args.conf", expected, 0640) + fs.Check(t, "/etc/systemd/system/kubelet.service.d/50-taints.conf", expected, 0640) } -func instance(ip, dnsName string, maxPods int, reservedCPU, reservedMemory string, spot bool) *node.Node { +func instance(ip, dnsName string, maxPods int, reservedCPU, reservedMemory string, labels, taints []string) *node.Node { return &node.Node{ Instance: &ec2.Instance{ PrivateIpAddress: &ip, @@ -205,7 +267,8 @@ func instance(ip, dnsName string, maxPods int, reservedCPU, reservedMemory strin Region: "us-east-1", ReservedCPU: reservedCPU, ReservedMemory: reservedMemory, - Spot: spot, + Labels: labels, + Taints: taints, } } diff --git a/pkg/system/templates/etc/systemd/system/kubelet.service b/pkg/system/templates/etc/systemd/system/kubelet.service index 12ab085..afb4aaf 100644 --- a/pkg/system/templates/etc/systemd/system/kubelet.service +++ b/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_SPOT_ARGS $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 diff --git a/pkg/system/templates/etc/systemd/system/kubelet.service.d/40-labels.conf b/pkg/system/templates/etc/systemd/system/kubelet.service.d/40-labels.conf new file mode 100644 index 0000000..15d7e2b --- /dev/null +++ b/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 -}} diff --git a/pkg/system/templates/etc/systemd/system/kubelet.service.d/40-spot-args.conf b/pkg/system/templates/etc/systemd/system/kubelet.service.d/40-spot-args.conf deleted file mode 100644 index ab2e2c6..0000000 --- a/pkg/system/templates/etc/systemd/system/kubelet.service.d/40-spot-args.conf +++ /dev/null @@ -1,6 +0,0 @@ -[Service] -{{- if .Node.Spot }} -Environment='KUBELET_SPOT_ARGS=--node-labels="node-role.kubernetes.io/spot-worker=true"' -{{ else }} -Environment='KUBELET_SPOT_ARGS=--node-labels="node-role.kubernetes.io/worker=true" --register-with-taints="node-role.kubernetes.io/worker=true:PreferNoSchedule"' -{{ end -}} diff --git a/pkg/system/templates/etc/systemd/system/kubelet.service.d/50-taints.conf b/pkg/system/templates/etc/systemd/system/kubelet.service.d/50-taints.conf new file mode 100644 index 0000000..1cec3b7 --- /dev/null +++ b/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 -}}