From 630ec097b80bc57f15e95f36d088dff04b2d3677 Mon Sep 17 00:00:00 2001 From: Theo Barber-Bany Date: Tue, 20 Aug 2019 11:46:01 +0100 Subject: [PATCH] Support new Node Label schema --- nodes/nodes.go | 45 +++++++++++++++++++++++++++++++++++++++------ nodes/nodes_test.go | 33 +++++++++++++++++++++++++++++++-- rescheduler.go | 24 ++++++++++++++++++++++-- rescheduler_test.go | 18 ++++++++++++++++++ 4 files changed, 110 insertions(+), 10 deletions(-) diff --git a/nodes/nodes.go b/nodes/nodes.go index 64f9328c..6ff5f162 100644 --- a/nodes/nodes.go +++ b/nodes/nodes.go @@ -18,6 +18,7 @@ package nodes import ( "sort" + "strings" apiv1 "k8s.io/api/core/v1" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" @@ -27,9 +28,9 @@ import ( var ( // OnDemandNodeLabel label for on-demand instances. - OnDemandNodeLabel = "node-role.kubernetes.io/worker" + OnDemandNodeLabel = "kubernetes.io/role=worker" // SpotNodeLabel label for spot instances. - SpotNodeLabel = "node-role.kubernetes.io/spot-worker" + SpotNodeLabel = "kubernetes.io/role=spot-worker" // OnDemand key for on-demand instances of NodesMap. OnDemand NodeType // Spot key for spot instances of NodesMap. @@ -157,14 +158,46 @@ func getPodCPURequests(pod *apiv1.Pod) int64 { // Determines if a node has the spotNodeLabel assigned func isSpotNode(node *apiv1.Node) bool { - _, found := node.ObjectMeta.Labels[SpotNodeLabel] - return found + splitLabel := strings.SplitN(SpotNodeLabel, "=", 2) + + // If "=" found, check for new label schema. If no "=" is found, check for + // old label schema + switch len(splitLabel) { + case 1: + _, found := node.ObjectMeta.Labels[SpotNodeLabel] + return found + case 2: + spotLabelKey := splitLabel[0] + spotLabelVal := splitLabel[1] + + val, _ := node.ObjectMeta.Labels[spotLabelKey] + if val == spotLabelVal { + return true + } + } + return false } // Determines if a node has the OnDemandNodeLabel assigned func isOnDemandNode(node *apiv1.Node) bool { - _, found := node.ObjectMeta.Labels[OnDemandNodeLabel] - return found + splitLabel := strings.SplitN(OnDemandNodeLabel, "=", 2) + + // If "=" found, check for new label schema. If no "=" is found, check for + // old label schema + switch len(splitLabel) { + case 1: + _, found := node.ObjectMeta.Labels[OnDemandNodeLabel] + return found + case 2: + onDemandLabelKey := splitLabel[0] + onDemandLabelVal := splitLabel[1] + + val, _ := node.ObjectMeta.Labels[onDemandLabelKey] + if val == onDemandLabelVal { + return true + } + } + return false } // CopyNodeInfos returns an array of copies of the NodeInfos in this array. diff --git a/nodes/nodes_test.go b/nodes/nodes_test.go index 589d1921..9d1046a9 100644 --- a/nodes/nodes_test.go +++ b/nodes/nodes_test.go @@ -29,12 +29,41 @@ import ( core "k8s.io/client-go/testing" ) +func TestIsSpotNode(t *testing.T) { + spotNode := createTestNodeWithLabel("fooSpotNode", 2000, map[string]string{"foo": "bar"}) + + SpotNodeLabel = "foo" + assert.True(t, isSpotNode(spotNode), "expected node with label 'foo' to be spot node") + + SpotNodeLabel = "foo=bar" + assert.True(t, isSpotNode(spotNode), "expected node with label 'foo' and value 'bar' to be spot node") + + SpotNodeLabel = "foo=baz" + assert.False(t, isSpotNode(spotNode), "expected node with label 'foo' and value 'bar' to not be spot node") +} + +func TestIsOnDemandNode(t *testing.T) { + onDemandNode := createTestNodeWithLabel("fooDemandNode", 2000, map[string]string{"foo": "bar"}) + + OnDemandNodeLabel = "foo" + assert.True(t, isOnDemandNode(onDemandNode), "expected node with label 'foo' to be on demand node") + + OnDemandNodeLabel = "foo=bar" + assert.True(t, isOnDemandNode(onDemandNode), "expected node with label 'foo' and value 'bar' to be on demand node") + + OnDemandNodeLabel = "foo=baz" + assert.False(t, isOnDemandNode(onDemandNode), "expected node with label 'foo' and value 'bar' to not be on demand node") +} + func TestNewNodeMap(t *testing.T) { + OnDemandNodeLabel = "kubernetes.io/role=worker" + SpotNodeLabel = "kubernetes.io/role=spot-worker" + spotLabels := map[string]string{ - SpotNodeLabel: "true", + "kubernetes.io/role": "spot-worker", } onDemandLabels := map[string]string{ - OnDemandNodeLabel: "true", + "kubernetes.io/role": "worker", } nodes := []*apiv1.Node{ diff --git a/rescheduler.go b/rescheduler.go index 7d700136..2daed8b2 100644 --- a/rescheduler.go +++ b/rescheduler.go @@ -21,6 +21,7 @@ import ( "fmt" "net/http" "os" + "strings" "time" "github.com/pusher/k8s-spot-rescheduler/metrics" @@ -96,11 +97,11 @@ func main() { // Add nodes labels as flags flags.StringVar(&nodes.OnDemandNodeLabel, "on-demand-node-label", - "node-role.kubernetes.io/worker", + "kubernetes.io/role=worker", `Name of label on nodes to be considered for draining.`) flags.StringVar(&nodes.SpotNodeLabel, "spot-node-label", - "node-role.kubernetes.io/spot-worker", + "kubernetes.io/role=spot-worker", `Name of label on nodes to be considered as targets for pods.`) flags.Parse(os.Args) @@ -110,6 +111,12 @@ func main() { os.Exit(0) } + err := validateArgs(nodes.OnDemandNodeLabel, nodes.SpotNodeLabel) + if err != nil { + fmt.Printf("Error: %s", err) + os.Exit(1) + } + glog.Infof("Running Rescheduler") // Register metrics from metrics.go @@ -419,3 +426,16 @@ func updateSpotNodeMetrics(spotNodeInfos nodes.NodeInfoArray, pdbs []*policyv1.P func podID(pod *apiv1.Pod) string { return fmt.Sprintf("%s/%s", pod.Namespace, pod.Name) } + +// Checks that the node lablels provided as arguments are in fact, sane. +func validateArgs(OnDemandNodeLabel string, SpotNodeLabel string) error { + if len(strings.Split(OnDemandNodeLabel, "=")) > 2 { + return fmt.Errorf("the on demand node label is not correctly formatted: expected '' or '=', but got %s", OnDemandNodeLabel) + } + + if len(strings.Split(SpotNodeLabel, "=")) > 2 { + return fmt.Errorf("the spot node label is not correctly formatted: expected '' or '=', but got %s", SpotNodeLabel) + } + + return nil +} diff --git a/rescheduler_test.go b/rescheduler_test.go index 1e5f1874..865516fd 100644 --- a/rescheduler_test.go +++ b/rescheduler_test.go @@ -70,6 +70,24 @@ func TestFindSpotNodeForPod(t *testing.T) { } +func TestNodeLabelValidation(t *testing.T) { + onDemandLabel := "foo.bar/role=worker" + spotLabel := "foo.bar/node-role" + + err := validateArgs(onDemandLabel, spotLabel) + assert.NoError(t, err) + + onDemandLabel = "foo.bar/broken=worker=true" + err = validateArgs(onDemandLabel, spotLabel) + assert.EqualError(t, err, "the on demand node label is not correctly formatted: expected '' or '=', but got foo.bar/broken=worker=true") + + onDemandLabel = "foo.bar/role=worker" + spotLabel = "foo.bar/node-role=spot=fail" + err = validateArgs(onDemandLabel, spotLabel) + assert.EqualError(t, err, "the spot node label is not correctly formatted: expected '' or '=', but got foo.bar/node-role=spot=fail") + +} + func TestCanDrainNode(t *testing.T) { predicateChecker := simulator.NewTestPredicateChecker()