Merge pull request #1647 from biancaburtoiu/biancaburtoiu/set-up-karpenter-capacityType-label

Set up Karpenter capacity type label

pkg/cloud/awsprovider.go

@@ -667,10 +667,13 @@ func (k *awsKey) Features() string {
 // If the instance is a spot instance, it will return PreemptibleType
 // Otherwise returns an empty string
 func (k *awsKey) getUsageType(labels map[string]string) string {
-	if label, ok := labels[EKSCapacityTypeLabel]; ok && label == EKSCapacitySpotTypeValue {
+	if eksLabel, ok := labels[EKSCapacityTypeLabel]; ok && eksLabel == EKSCapacitySpotTypeValue {
 		// We currently write out spot instances as "preemptible" in the pricing data, so these need to match
 		return PreemptibleType
 	}
+	if kLabel, ok := labels[KarpenterCapacityTypeLabel]; ok && kLabel == KarpenterCapacitySpotTypeValue {
+		return PreemptibleType
+	}
 	return ""
 }
 

pkg/cloud/awsprovider_test.go

@@ -25,7 +25,7 @@ func Test_awsKey_getUsageType(t *testing.T) {
 			want: "",
 		},
 		{
-			name: "label with a capacityType set to empty string should return empty string",
+			name: "EKS label with a capacityType set to empty string should return empty string",
 			args: args{
 				labels: map[string]string{
 					EKSCapacityTypeLabel: "",
@@ -34,7 +34,7 @@ func Test_awsKey_getUsageType(t *testing.T) {
 			want: "",
 		},
 		{
-			name: "label with capacityType set to a random value should return empty string",
+			name: "EKS label with capacityType set to a random value should return empty string",
 			args: args{
 				labels: map[string]string{
 					EKSCapacityTypeLabel: "TEST_ME",
@@ -43,7 +43,7 @@ func Test_awsKey_getUsageType(t *testing.T) {
 			want: "",
 		},
 		{
-			name: "label with capacityType set to spot should return spot",
+			name: "EKS label with capacityType set to spot should return spot",
 			args: args{
 				labels: map[string]string{
 					EKSCapacityTypeLabel: EKSCapacitySpotTypeValue,
@@ -51,6 +51,33 @@ func Test_awsKey_getUsageType(t *testing.T) {
 			},
 			want: PreemptibleType,
 		},
+		{
+			name: "Karpenter label with a capacityType set to empty string should return empty string",
+			args: args{
+				labels: map[string]string{
+					KarpenterCapacityTypeLabel: "",
+				},
+			},
+			want: "",
+		},
+		{
+			name: "Karpenter label with capacityType set to a random value should return empty string",
+			args: args{
+				labels: map[string]string{
+					KarpenterCapacityTypeLabel: "TEST_ME",
+				},
+			},
+			want: "",
+		},
+		{
+			name: "Karpenter label with capacityType set to spot should return spot",
+			args: args{
+				labels: map[string]string{
+					KarpenterCapacityTypeLabel: KarpenterCapacitySpotTypeValue,
+				},
+			},
+			want: PreemptibleType,
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {

pkg/cloud/gcpprovider.go

@@ -42,6 +42,9 @@ const (
 	GCPMonthlyBasicDiskCost = 0.04
 	GCPMonthlySSDDiskCost   = 0.17
 	GCPMonthlyGP2DiskCost   = 0.1
+
+	GKEPreemptibleLabel = "cloud.google.com/gke-preemptible"
+	GKESpotLabel        = "cloud.google.com/gke-spot"
 )
 
 // List obtained by installing the `gcloud` CLI tool,
@@ -1563,11 +1566,13 @@ func parseGCPProjectID(id string) string {
 }
 
 func getUsageType(labels map[string]string) string {
-	if t, ok := labels["cloud.google.com/gke-preemptible"]; ok && t == "true" {
+	if t, ok := labels[GKEPreemptibleLabel]; ok && t == "true" {
 		return "preemptible"
-	} else if t, ok := labels["cloud.google.com/gke-spot"]; ok && t == "true" {
+	} else if t, ok := labels[GKESpotLabel]; ok && t == "true" {
 		// https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms
 		return "preemptible"
+	} else if t, ok := labels[KarpenterCapacityTypeLabel]; ok && t == KarpenterCapacitySpotTypeValue {
+		return "preemptible"
 	}
 	return "ondemand"
 }

pkg/cloud/gcpprovider_test.go

@@ -73,13 +73,19 @@ func TestGetUsageType(t *testing.T) {
 	}{
 		{
 			input: map[string]string{
-				"cloud.google.com/gke-preemptible": "true",
+				GKEPreemptibleLabel: "true",
 			},
 			expected: "preemptible",
 		},
 		{
 			input: map[string]string{
-				"cloud.google.com/gke-spot": "true",
+				GKESpotLabel: "true",
+			},
+			expected: "preemptible",
+		},
+		{
+			input: map[string]string{
+				KarpenterCapacityTypeLabel: KarpenterCapacitySpotTypeValue,
 			},
 			expected: "preemptible",
 		},

pkg/cloud/provider.go

@@ -32,6 +32,9 @@ const authSecretPath = "/var/secrets/service-key.json"
 const storageConfigSecretPath = "/var/azure-storage-config/azure-storage-config.json"
 const defaultShareTenancyCost = "true"
 
+const KarpenterCapacityTypeLabel = "karpenter.sh/capacity-type"
+const KarpenterCapacitySpotTypeValue = "spot"
+
 var createTableStatements = []string{
 	`CREATE TABLE IF NOT EXISTS names (
 		cluster_id VARCHAR(255) NOT NULL,