Merge pull request kubernetes#5103 from towca/jtuznik/custom-fix

CA: GCE: fix custom machine type parsing
navinjoy · Aug 18, 2022 · 221e181 · 221e181
2 parents 3429502 + 1d11d49
commit 221e181
Show file tree

Hide file tree

Showing 2 changed files with 107 additions and 15 deletions.
diff --git a/cluster-autoscaler/cloudprovider/gce/machine_types.go b/cluster-autoscaler/cloudprovider/gce/machine_types.go
@@ -83,19 +83,23 @@ func NewCustomMachineType(name string) (MachineType, error) {
 		return MachineType{}, fmt.Errorf("%q is not a valid custom machine type", name)
 	}
 
+	// Identify the "custom" part of the name, assume the next part is the CPU count, and the one after that is the memory amount.
+	// This should work if the type name has a "custom-*-*" infix, regardless of the rest of the 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 {
+	customPartIndex := -1
+	for i, part := range parts {
+		if part == "custom" {
+			customPartIndex = i
+			break
+		}
+	}
+	if customPartIndex == -1 || customPartIndex+2 >= len(parts) {
 		return MachineType{}, fmt.Errorf("unable to parse custom machine type %q", name)
 	}
+	cpuPart := parts[customPartIndex+1]
+	memPart := parts[customPartIndex+2]
 
-	cpu, err := strconv.ParseInt(cpuPart, 10, 64)
+	cpu, err := parseCustomCpu(name, cpuPart)
 	if err != nil {
 		return MachineType{}, fmt.Errorf("unable to parse CPU %q from machine type %q: %v", cpuPart, name, err)
 	}
@@ -110,3 +114,19 @@ func NewCustomMachineType(name string) (MachineType, error) {
 		Memory: memBytes * units.MiB,
 	}, nil
 }
+
+func parseCustomCpu(machineType string, cpuPart string) (int64, error) {
+	// We need to identify the family because some e2 machine types have special names.
+	family, err := GetMachineFamily(machineType)
+	if err != nil {
+		return 0, fmt.Errorf("unable to get family while parsing custom machine type %q: %v", machineType, err)
+	}
+
+	// There are e2-custom-micro-*, e2-custom-small-*, e2-custom-medium-* custom machine types in the e2 family. They all have 2 guestCpus
+	// in the API.
+	if family == "e2" && (cpuPart == "micro" || cpuPart == "small" || cpuPart == "medium") {
+		return 2, nil
+	}
+
+	return strconv.ParseInt(cpuPart, 10, 64)
+}
diff --git a/cluster-autoscaler/cloudprovider/gce/machine_types_test.go b/cluster-autoscaler/cloudprovider/gce/machine_types_test.go
@@ -30,6 +30,7 @@ func TestNewCustomMachineType(t *testing.T) {
 	testCases := []struct {
 		name         string
 		expectCustom bool
+		expectErr    bool
 		expectCPU    int64
 		expectMemory int64
 	}{
@@ -39,33 +40,84 @@ func TestNewCustomMachineType(t *testing.T) {
 			expectCPU:    2,
 			expectMemory: 2816 * units.MiB,
 		},
+		{
+			name:         "custom-2-2816-ext",
+			expectCustom: true,
+			expectCPU:    2,
+			expectMemory: 2816 * units.MiB,
+		},
 		{
 			name:         "n2-custom-2-2816",
 			expectCustom: true,
 			expectCPU:    2,
 			expectMemory: 2816 * units.MiB,
 		},
 		{
-			name: "other-a2-2816",
+			name:         "n2-custom-2-2816-ext",
+			expectCustom: true,
+			expectCPU:    2,
+			expectMemory: 2816 * units.MiB,
+		},
+		{
+			name:         "e2-custom-medium-2816",
+			expectCustom: true,
+			expectCPU:    2,
+			expectMemory: 2816 * units.MiB,
+		},
+		{
+			name:         "e2-custom-micro-2048",
+			expectCustom: true,
+			expectCPU:    2,
+			expectMemory: 2048 * units.MiB,
+		},
+		{
+			name:         "e2-custom-small-2048",
+			expectCustom: true,
+			expectCPU:    2,
+			expectMemory: 2048 * units.MiB,
+		},
+		{
+			name:         "e2-custom",
+			expectCustom: true,
+			expectErr:    true,
+		},
+		{
+			name:         "e2-custom-8",
+			expectCustom: true,
+			expectErr:    true,
+		},
+		{
+			name:         "e2-custom-abc-2048",
+			expectCustom: true,
+			expectErr:    true,
+		},
+		{
+			name:         "other-a2-2816",
+			expectCustom: false,
+			expectErr:    true,
 		},
 		{
-			name: "other-2-2816",
+			name:         "other-2-2816",
+			expectCustom: false,
+			expectErr:    true,
 		},
 		{
-			name: "n1-standard-8",
+			name:         "n1-standard-8",
+			expectCustom: false,
+			expectErr:    true,
 		},
 	}
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			assert.Equal(t, tc.expectCustom, IsCustomMachine(tc.name))
 			m, err := NewCustomMachineType(tc.name)
-			if tc.expectCustom {
+			if tc.expectErr {
+				assert.Error(t, err)
+			} else {
 				assert.NoError(t, err)
 				assert.Equal(t, tc.expectCPU, m.CPU)
 				assert.Equal(t, tc.expectMemory, m.Memory)
-			} else {
-				assert.Error(t, err)
 			}
 		})
 	}
@@ -89,10 +141,30 @@ func TestGetMachineFamily(t *testing.T) {
 			machineType: "n2-custom-2-2816",
 			wantFamily:  "n2",
 		},
+		"custom machine type with family prefix and ext suffix": {
+			machineType: "n2-custom-2-2816-ext",
+			wantFamily:  "n2",
+		},
 		"custom machine type without family prefix": {
 			machineType: "custom-2-2816",
 			wantFamily:  "n1",
 		},
+		"custom machine type without family prefix with ext suffix": {
+			machineType: "custom-2-2816-ext",
+			wantFamily:  "n1",
+		},
+		"e2 custom medium type": {
+			machineType: "e2-custom-medium-2816",
+			wantFamily:  "e2",
+		},
+		"e2 custom small type": {
+			machineType: "e2-custom-small-2048",
+			wantFamily:  "e2",
+		},
+		"e2 custom micro type": {
+			machineType: "e2-custom-micro-2048",
+			wantFamily:  "e2",
+		},
 		"invalid machine type": {
 			machineType: "nodashes",
 			wantErr:     cmpopts.AnyError,