Merge pull request #182 from castai/spot_labels_gke_aks

feat: recognize karpetner spot label in GKE and AKS
castai · Aug 2, 2024 · a70be5a · a70be5a
2 parents 3d6f471 + 82f1f01
commit a70be5a
Show file tree

Hide file tree

Showing 4 changed files with 40 additions and 3 deletions.
diff --git a/internal/services/providers/aks/aks.go b/internal/services/providers/aks/aks.go
@@ -11,6 +11,7 @@ import (
 	"castai-agent/internal/config"
 	"castai-agent/internal/services/providers/aks/metadata"
 	"castai-agent/internal/services/providers/types"
+	"castai-agent/pkg/labels"
 )
 
 type Provider struct {
@@ -92,16 +93,27 @@ func (p *Provider) FilterSpot(_ context.Context, nodes []*corev1.Node) ([]*corev
 	var ret []*corev1.Node
 
 	for _, node := range nodes {
-		if val, ok := node.ObjectMeta.Labels["kubernetes.azure.com/scalesetpriority"]; !ok || val != "spot" {
-			continue
+		if isSpot(node) {
+			ret = append(ret, node)
 		}
 
-		ret = append(ret, node)
 	}
 
 	return ret, nil
 }
 
+func isSpot(node *corev1.Node) bool {
+	if val, ok := node.ObjectMeta.Labels["kubernetes.azure.com/scalesetpriority"]; ok && val == "spot" {
+		return true
+	}
+
+	if val, ok := node.ObjectMeta.Labels[labels.KarpenterCapacityType]; ok && val == labels.ValueKarpenterCapacityTypeSpot {
+		return true
+	}
+
+	return false
+}
+
 func (p *Provider) Name() string {
 	return Name
 }
diff --git a/internal/services/providers/aks/aks_test.go b/internal/services/providers/aks/aks_test.go
@@ -15,6 +15,7 @@ import (
 	"castai-agent/internal/castai"
 	mock_castai "castai-agent/internal/castai/mock"
 	"castai-agent/internal/services/providers/types"
+	"castai-agent/pkg/labels"
 )
 
 func TestProvider_RegisterCluster(t *testing.T) {
@@ -69,4 +70,19 @@ func TestProvider_IsSpot(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, []*v1.Node{node}, got)
 	})
+
+	t.Run("spot instance karpenter label", func(t *testing.T) {
+		p := &Provider{
+			log: logrus.New(),
+		}
+
+		node := &v1.Node{
+			ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{labels.KarpenterCapacityType: labels.ValueKarpenterCapacityTypeSpot}},
+		}
+
+		got, err := p.FilterSpot(context.Background(), []*v1.Node{node})
+
+		require.NoError(t, err)
+		require.Equal(t, []*v1.Node{node}, got)
+	})
 }
diff --git a/internal/services/providers/gke/gke.go b/internal/services/providers/gke/gke.go
@@ -112,6 +112,10 @@ func isSpot(node *corev1.Node) bool {
 		return true
 	}
 
+	if val, ok := node.Labels[labels.KarpenterCapacityType]; ok && val == labels.ValueKarpenterCapacityTypeSpot {
+		return true
+	}
+
 	return false
 }
 

diff --git a/internal/services/providers/gke/gke_test.go b/internal/services/providers/gke/gke_test.go
@@ -99,6 +99,11 @@ func TestProvider_IsSpot(t *testing.T) {
 			node:     &v1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{LabelSpot: "true"}}},
 			expected: true,
 		},
+		{
+			name:     "gke karpenter spot node",
+			node:     &v1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{labels.KarpenterCapacityType: labels.ValueKarpenterCapacityTypeSpot}}},
+			expected: true,
+		},
 		{
 			name:     "on demand node",
 			node:     &v1.Node{ObjectMeta: metav1.ObjectMeta{Labels: map[string]string{}}},