Skip to content
This repository has been archived by the owner on Oct 15, 2020. It is now read-only.

Commit

Permalink
Merge pull request #47 from errm/dynamic-tags-and-taints
Browse files Browse the repository at this point in the history 
Dynamic tags and taints
  • Loading branch information
@errm
errm committed Nov 13, 2018
2 parents 6aa77c1 + 8491dd7 commit a5d395a
Show file tree
Hide file tree
Showing 7 changed files with 192 additions and 92 deletions.
37 changes: 31 additions & 6 deletions pkg/node/node.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,8 @@ type Node struct {
ReservedMemory string
ClusterDNS string
Region string
Spot bool
Labels []string
Taints []string
}

type metadataClient interface {
Expand All @@ -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/<name> tag is not yet set, will try again in %s", sleepFor)
Expand All @@ -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) {
Expand Down
138 changes: 74 additions & 64 deletions pkg/node/node_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ package node

import (
"errors"
"reflect"
"testing"

"github.com/errm/ekstrap/pkg/backoff"
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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, &region)
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,
Expand All @@ -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) {
Expand All @@ -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,
},
},
},
Expand Down
93 changes: 78 additions & 15 deletions pkg/system/system_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
`
Expand All @@ -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",
Expand All @@ -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",
Expand All @@ -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,
Expand All @@ -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,
}
}

Expand Down
Loading

0 comments on commit a5d395a

Please sign in to comment.