kamilkrampa 3 лет назад
Родитель
Сommit
a5f398d504
2 измененных файлов с 51 добавлено и 3 удалено
  1. 14 3
      pkg/cloud/gcpprovider.go
  2. 37 0
      pkg/cloud/gcpprovider_test.go

+ 14 - 3
pkg/cloud/gcpprovider.go

@@ -3,7 +3,6 @@ package cloud
 import (
 	"context"
 	"fmt"
-	"github.com/kubecost/opencost/pkg/kubecost"
 	"io"
 	"io/ioutil"
 	"math"
@@ -14,6 +13,8 @@ import (
 	"sync"
 	"time"
 
+	"github.com/kubecost/opencost/pkg/kubecost"
+
 	"github.com/kubecost/opencost/pkg/clustercache"
 	"github.com/kubecost/opencost/pkg/env"
 	"github.com/kubecost/opencost/pkg/log"
@@ -1257,7 +1258,7 @@ func (gcp *gcpKey) ID() string {
 func (gcp *gcpKey) GPUType() string {
 	if t, ok := gcp.Labels[GKE_GPU_TAG]; ok {
 		var usageType string
-		if t, ok := gcp.Labels["cloud.google.com/gke-preemptible"]; ok && t == "true" {
+		if isPreemptible(gcp.Labels) {
 			usageType = "preemptible"
 		} else {
 			usageType = "ondemand"
@@ -1310,7 +1311,7 @@ func (gcp *gcpKey) Features() string {
 	region := strings.ToLower(r)
 	var usageType string
 
-	if t, ok := gcp.Labels["cloud.google.com/gke-preemptible"]; ok && t == "true" {
+	if isPreemptible(gcp.Labels) {
 		usageType = "preemptible"
 	} else {
 		usageType = "ondemand"
@@ -1409,3 +1410,13 @@ func parseGCPProjectID(id string) string {
 	// Return empty string if an account could not be parsed from provided string
 	return ""
 }
+
+func isPreemptible(labels map[string]string) bool {
+	if t, ok := labels["cloud.google.com/gke-preemptible"]; ok && t == "true" {
+		return true
+	} else if t, ok := labels["cloud.google.com/gke-spot"]; ok && t == "true" {
+		// https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms
+		return true
+	}
+	return false
+}

+ 37 - 0
pkg/cloud/gcpprovider_test.go

@@ -65,3 +65,40 @@ func TestParseGCPProjectID(t *testing.T) {
 		}
 	}
 }
+
+func TestIsPreemptible(t *testing.T) {
+	cases := []struct {
+		input    map[string]string
+		expected bool
+	}{
+		{
+			input: map[string]string{
+				"cloud.google.com/gke-preemptible": "true",
+			},
+			expected: true,
+		},
+		{
+			input: map[string]string{
+				"cloud.google.com/gke-spot": "true",
+			},
+			expected: true,
+		},
+		{
+			input: map[string]string{
+				"someotherlabel": "true",
+			},
+			expected: false,
+		},
+		{
+			input:    map[string]string{},
+			expected: false,
+		},
+	}
+
+	for _, test := range cases {
+		result := isPreemptible(test.input)
+		if result != test.expected {
+			t.Errorf("Input: %v, Expected: %t, Actual: %t", test.input, test.expected, result)
+		}
+	}
+}