Merge pull request #1279 from kamilkrampa/gcp-spot-instances

Support GCP Spot VMs

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"
@@ -1256,12 +1257,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" {
-			usageType = "preemptible"
-		} else {
-			usageType = "ondemand"
-		}
+		usageType := getUsageType(gcp.Labels)
 		log.Debugf("GPU of type: \"%s\" found", t)
 		return t + "," + usageType
 	}
@@ -1308,13 +1304,7 @@ func (gcp *gcpKey) Features() string {
 
 	r, _ := util.GetRegion(gcp.Labels)
 	region := strings.ToLower(r)
-	var usageType string
-
-	if t, ok := gcp.Labels["cloud.google.com/gke-preemptible"]; ok && t == "true" {
-		usageType = "preemptible"
-	} else {
-		usageType = "ondemand"
-	}
+	usageType := getUsageType(gcp.Labels)
 
 	if _, ok := gcp.Labels[GKE_GPU_TAG]; ok {
 		return region + "," + instanceType + "," + usageType + "," + "gpu"
@@ -1409,3 +1399,13 @@ func parseGCPProjectID(id string) string {
 	// Return empty string if an account could not be parsed from provided string
 	return ""
 }
+
+func getUsageType(labels map[string]string) string {
+	if t, ok := labels["cloud.google.com/gke-preemptible"]; ok && t == "true" {
+		return "preemptible"
+	} else if t, ok := labels["cloud.google.com/gke-spot"]; ok && t == "true" {
+		// https://cloud.google.com/kubernetes-engine/docs/concepts/spot-vms
+		return "preemptible"
+	}
+	return "ondemand"
+}

pkg/cloud/gcpprovider_test.go

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