Merge pull request #19 from castai/KUBE-171

KUBE-171: added a patch node request to set the required label/taint
castai · Dec 8, 2023 · d514924 · d514924
2 parents 01eaeb5 + e9c9ae9
commit d514924
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/handler/handler.go b/handler/handler.go
@@ -21,8 +21,16 @@ import (
 const CastNodeIDLabel = "provisioner.cast.ai/node-id"
 
 const (
-	cloudEventInterrupted             = "interrupted"
-	cloudEventRebalanceRecommendation = "rebalanceRecommendation"
+	taintNodeDraining = "autoscaling.cast.ai/draining"
+	taintNodeDrainingEffect = "NoSchedule"
+
+	labelNodeDraining                  = "autoscaling.cast.ai/draining"
+	valueNodeDrainingReasonInterrupted = "spot-interruption"
+
+	cloudEventInterrupted              = "interrupted"
+	cloudEventRebalanceRecommendation  = "rebalanceRecommendation"
+
+	valueTrue = "true"
 )
 
 type MetadataChecker interface {
@@ -144,6 +152,12 @@ func (g *SpotHandler) taintNode(ctx context.Context, node *v1.Node) error {
 
 	err := g.patchNode(ctx, node, func(n *v1.Node) error {
 		n.Spec.Unschedulable = true
+		n.Labels[labelNodeDraining] = valueNodeDrainingReasonInterrupted
+		n.Spec.Taints = append(n.Spec.Taints, v1.Taint{
+			Key:    taintNodeDraining,
+			Value:  valueTrue,
+			Effect: taintNodeDrainingEffect,
+		})
 		return nil
 	})
 	if err != nil {

diff --git a/handler/handler_test.go b/handler/handler_test.go
@@ -73,6 +73,12 @@ func TestRunLoop(t *testing.T) {
 
 		node, _ = fakeApi.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
 		r.Equal(true, node.Spec.Unschedulable)
+		r.Equal(valueNodeDrainingReasonInterrupted, node.Labels[labelNodeDraining])
+		r.Contains(node.Spec.Taints, v1.Taint{
+			Key:    taintNodeDraining,
+			Value:  valueTrue,
+			Effect: taintNodeDrainingEffect,
+		})
 	})
 
 	t.Run("keep checking interruption on context canceled", func(t *testing.T) {
@@ -110,6 +116,12 @@ func TestRunLoop(t *testing.T) {
 
 		node, _ = fakeApi.CoreV1().Nodes().Get(context.Background(), nodeName, metav1.GetOptions{})
 		r.Equal(true, node.Spec.Unschedulable)
+		r.Equal(valueNodeDrainingReasonInterrupted, node.Labels[labelNodeDraining])
+		r.Contains(node.Spec.Taints, v1.Taint{
+			Key:    taintNodeDraining,
+			Value:  valueTrue,
+			Effect: taintNodeDrainingEffect,
+		})
 	})
 
 	t.Run("handle mock interruption retries", func(t *testing.T) {