CA: GCE: implement GetMachineFamily, fix IsCustomMachine

cluster-autoscaler/cloudprovider/gce/gce_price_model.go

@@ -19,7 +19,6 @@ package gce
 import (
 	"math"
 	"strconv"
-	"strings"
 	"time"
 
 	apiv1 "k8s.io/api/core/v1"
@@ -142,10 +141,10 @@ func (model *GcePriceModel) getPreemptibleDiscount(node *apiv1.Node) float64 {
 	if !found {
 		return 1.0
 	}
-	instanceFamily := getInstanceFamily(instanceType)
+	instanceFamily, _ := GetMachineFamily(instanceType)
 
 	discountMap := model.PriceInfo.PredefinedPreemptibleDiscount()
-	if isInstanceCustom(instanceType) {
+	if IsCustomMachine(instanceType) {
 		discountMap = model.PriceInfo.CustomPreemptibleDiscount()
 	}
 
@@ -171,8 +170,8 @@ func (model *GcePriceModel) getBasePrice(resources apiv1.ResourceList, instanceT
 		return 0
 	}
 	hours := getHours(startTime, endTime)
-	instanceFamily := getInstanceFamily(instanceType)
-	isCustom := isInstanceCustom(instanceType)
+	instanceFamily, _ := GetMachineFamily(instanceType)
+	isCustom := IsCustomMachine(instanceType)
 	price := 0.0
 
 	cpu := resources[apiv1.ResourceCPU]
@@ -224,14 +223,6 @@ func getHours(startTime time.Time, endTime time.Time) float64 {
 	return hours
 }
 
-func getInstanceFamily(instanceType string) string {
-	return strings.Split(instanceType, "-")[0]
-}
-
-func isInstanceCustom(instanceType string) bool {
-	return strings.Contains(instanceType, "custom")
-}
-
 // hasPreemptiblePricing returns whether we should use preemptible pricing for a node, based on labels. Spot VMs have
 // dynamic pricing, which is different than the static pricing for Preemptible VMs we use here. However it should be close
 // enough in practice and we really only look at prices in comparison with each other. Spot VMs will always be cheaper

cluster-autoscaler/cloudprovider/gce/machine_types.go

@@ -18,6 +18,7 @@ package gce
 
 import (
 	"fmt"
+	"strconv"
 	"strings"
 
 	"k8s.io/autoscaler/cluster-autoscaler/utils/units"
@@ -37,8 +38,29 @@ type MachineType struct {
 }
 
 // IsCustomMachine checks if a machine type is custom or predefined.
-func IsCustomMachine(name string) bool {
-	return strings.HasPrefix(name, "custom-")
+func IsCustomMachine(machineType string) bool {
+	// Custom types are in the form "<family>-custom-<num_cpu>-<memory_mb>", with the "<family>-" prefix being optional for the N1
+	// family. Examples: n2-custom-48-184320, custom-48-184320 (equivalent to n1-custom-48-184320).
+	// Docs: https://cloud.google.com/compute/docs/instances/creating-instance-with-custom-machine-type#gcloud.
+	parts := strings.Split(machineType, "-")
+	if len(parts) < 2 {
+		return false
+	}
+	return parts[0] == "custom" || parts[1] == "custom"
+}
+
+// GetMachineFamily gets the machine family from the machine type. Machine family is determined as everything before the first
+// dash character, unless this first part is "custom", which actually means "n1":
+// https://cloud.google.com/compute/docs/instances/creating-instance-with-custom-machine-type#gcloud.
+func GetMachineFamily(machineType string) (string, error) {
+	parts := strings.Split(machineType, "-")
+	if len(parts) < 2 {
+		return "", fmt.Errorf("unable to parse machine type %q", machineType)
+	}
+	if parts[0] == "custom" {
+		return "n1", nil
+	}
+	return parts[0], nil
 }
 
 // NewMachineTypeFromAPI creates a MachineType object based on machine type representation
@@ -57,20 +79,34 @@ func NewMachineTypeFromAPI(name string, mt *gce_api.MachineType) (MachineType, e
 // NewCustomMachineType creates a MachineType object describing a custom GCE machine.
 // CPU and Memory are based on parsing custom machine name.
 func NewCustomMachineType(name string) (MachineType, error) {
-	// example custom-2-2816
-	var cpu, mem int64
-	var count int
-	count, err := fmt.Sscanf(name, "custom-%d-%d", &cpu, &mem)
+	if !IsCustomMachine(name) {
+		return MachineType{}, fmt.Errorf("%q is not a valid custom machine type", name)
+	}
+
+	parts := strings.Split(name, "-")
+	var cpuPart, memPart string
+	if len(parts) == 3 {
+		cpuPart = parts[1]
+		memPart = parts[2]
+	} else if len(parts) == 4 {
+		cpuPart = parts[2]
+		memPart = parts[3]
+	} else {
+		return MachineType{}, fmt.Errorf("unable to parse custom machine type %q", name)
+	}
+
+	cpu, err := strconv.ParseInt(cpuPart, 10, 64)
 	if err != nil {
-		return MachineType{}, err
+		return MachineType{}, fmt.Errorf("unable to parse CPU %q from machine type %q: %v", cpuPart, name, err)
 	}
-	if count != 2 {
-		return MachineType{}, fmt.Errorf("failed to parse all params in %s", name)
+	memBytes, err := strconv.ParseInt(memPart, 10, 64)
+	if err != nil {
+		return MachineType{}, fmt.Errorf("unable to parse memory %q from machine type %q: %v", memPart, name, err)
 	}
-	mem = mem * units.MiB
+
 	return MachineType{
 		Name:   name,
 		CPU:    cpu,
-		Memory: mem,
+		Memory: memBytes * units.MiB,
 	}, nil
 }

cluster-autoscaler/cloudprovider/gce/machine_types_test.go

@@ -21,6 +21,8 @@ import (
 
 	"k8s.io/autoscaler/cluster-autoscaler/utils/units"
 
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-cmp/cmp/cmpopts"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -37,6 +39,12 @@ func TestNewCustomMachineType(t *testing.T) {
 			expectCPU:    2,
 			expectMemory: 2816 * units.MiB,
 		},
+		{
+			name:         "n2-custom-2-2816",
+			expectCustom: true,
+			expectCPU:    2,
+			expectMemory: 2816 * units.MiB,
+		},
 		{
 			name: "other-a2-2816",
 		},
@@ -62,3 +70,42 @@ func TestNewCustomMachineType(t *testing.T) {
 		})
 	}
 }
+
+func TestGetMachineFamily(t *testing.T) {
+	for tn, tc := range map[string]struct {
+		machineType string
+		wantFamily  string
+		wantErr     error
+	}{
+		"predefined machine type": {
+			machineType: "n1-standard-8",
+			wantFamily:  "n1",
+		},
+		"predefined short machine type": {
+			machineType: "e2-small",
+			wantFamily:  "e2",
+		},
+		"custom machine type with family prefix": {
+			machineType: "n2-custom-2-2816",
+			wantFamily:  "n2",
+		},
+		"custom machine type without family prefix": {
+			machineType: "custom-2-2816",
+			wantFamily:  "n1",
+		},
+		"invalid machine type": {
+			machineType: "nodashes",
+			wantErr:     cmpopts.AnyError,
+		},
+	} {
+		t.Run(tn, func(t *testing.T) {
+			gotFamily, gotErr := GetMachineFamily(tc.machineType)
+			if diff := cmp.Diff(tc.wantFamily, gotFamily); diff != "" {
+				t.Errorf("GetMachineFamily(%q): diff (-want +got):\n%s", tc.machineType, diff)
+			}
+			if diff := cmp.Diff(tc.wantErr, gotErr, cmpopts.EquateErrors()); diff != "" {
+				t.Errorf("GetMachineFamily(%q): err diff (-want +got):\n%s", tc.machineType, diff)
+			}
+		})
+	}
+}