Merge pull request #2067 from k8s-infra-cherrypick-robot/cherry-pick-2064-to-release-1.14

[release-1.14] Add Attach Limit for Hyperdisk + Gen4 VMs

Author: mattcary

pkg/common/constants.go

@@ -33,3 +33,16 @@ const (
 	// Label that is set on a disk when it is used by a 'multi-zone' VolumeHandle
 	MultiZoneLabel = "goog-gke-multi-zone"
 )
+
+// doc https://cloud.google.com/compute/docs/disks/hyperdisks#max-total-disks-per-vm
+var Gen4MachineHyperdiskAttachLimitMap = []struct {
+	max   int64
+	value int64
+}{
+	{max: 4, value: 15},
+	{max: 8, value: 23},
+	{max: 16, value: 31},
+	{max: 32, value: 49},
+	{max: 64, value: 63},
+	{max: 1024, value: 127},
+}

pkg/common/utils.go

@@ -696,3 +696,17 @@ func NewLimiter(limit, burst int, emptyBucket bool) *rate.Limiter {
 
 	return limiter
 }
+
+func IsHyperdisk(diskType string) bool {
+	return strings.HasPrefix(diskType, "hyperdisk-")
+}
+
+// MapNumber is a function to map input cpu number to the Hyperdisk attach limit
+func MapNumber(num int64) int64 {
+	for _, r := range Gen4MachineHyperdiskAttachLimitMap {
+		if num <= r.max {
+			return r.value
+		}
+	}
+	return 0
+}

pkg/gce-pd-csi-driver/node.go

@@ -23,6 +23,7 @@ import (
 	"regexp"
 	"runtime"
 	"strconv"
+	"strings"
 	"time"
 
 	"google.golang.org/grpc/codes"
@@ -70,8 +71,12 @@ var _ csi.NodeServer = &GCENodeServer{}
 // node boot disk is considered an attachable disk so effective attach limit is
 // one less.
 const (
-	volumeLimitSmall     int64 = 15
-	volumeLimitBig       int64 = 127
+	volumeLimitSmall int64 = 15
+	volumeLimitBig   int64 = 127
+	// doc https://cloud.google.com/compute/docs/memory-optimized-machines#x4_disks
+	x4HyperdiskLimit int64 = 39
+	// doc https://cloud.google.com/compute/docs/accelerator-optimized-machines#a4-disks
+	a4HyperdiskLimit     int64 = 127
 	defaultLinuxFsType         = "ext4"
 	defaultWindowsFsType       = "ntfs"
 	fsTypeExt3                 = "ext3"
@@ -501,6 +506,9 @@ func (ns *GCENodeServer) NodeGetInfo(ctx context.Context, req *csi.NodeGetInfoRe
 	nodeID := common.CreateNodeID(ns.MetadataService.GetProject(), ns.MetadataService.GetZone(), ns.MetadataService.GetName())
 
 	volumeLimits, err := ns.GetVolumeLimits()
+	if err != nil {
+		klog.Errorf("GetVolumeLimits failed: %v", err.Error())
+	}
 
 	resp := &csi.NodeGetInfoResponse{
 		NodeId:             nodeID,
@@ -667,5 +675,24 @@ func (ns *GCENodeServer) GetVolumeLimits() (int64, error) {
 			return volumeLimitSmall, nil
 		}
 	}
+	gen4MachineTypesPrefix := []string{"c4a-", "c4-", "n4-"}
+	for _, gen4Prefix := range gen4MachineTypesPrefix {
+		if strings.HasPrefix(machineType, gen4Prefix) {
+			cpuString := machineType[strings.LastIndex(machineType, "-")+1:]
+			cpus, err := strconv.ParseInt(cpuString, 10, 64)
+			if err != nil {
+				return volumeLimitSmall, fmt.Errorf("invalid cpuString %s for machine type: %v", cpuString, machineType)
+			}
+			return common.MapNumber(cpus), nil
+
+		}
+		if strings.HasPrefix(machineType, "x4-") {
+			return x4HyperdiskLimit, nil
+		}
+		if strings.HasPrefix(machineType, "a4-") {
+			return a4HyperdiskLimit, nil
+		}
+	}
+
 	return volumeLimitBig, nil
 }

pkg/gce-pd-csi-driver/node_test.go

@@ -204,6 +204,7 @@ func TestNodeGetVolumeLimits(t *testing.T) {
 		name           string
 		machineType    string
 		expVolumeLimit int64
+		expectError    bool
 	}{
 		{
 			name:           "Predifined standard machine",
@@ -235,13 +236,54 @@ func TestNodeGetVolumeLimits(t *testing.T) {
 			machineType:    "e2-micro",
 			expVolumeLimit: volumeLimitSmall,
 		},
+		{
+			name:           "c4-standard-192",
+			machineType:    "c4-standard-192",
+			expVolumeLimit: 127,
+		},
+		{
+			name:           "c4-standard-48",
+			machineType:    "c4-standard-48",
+			expVolumeLimit: 63,
+		},
+		{
+			name:           "c4a-standard-4",
+			machineType:    "c4a-standard-4",
+			expVolumeLimit: 15,
+		},
+		{
+			name:           "n4-standard-16",
+			machineType:    "n4-standard-16",
+			expVolumeLimit: 31,
+		},
+		{
+			name:           "n4-highcpu-4",
+			machineType:    "n4-highcpu-4",
+			expVolumeLimit: 15,
+		},
+		{
+			name:           "invalid gen4 machine type",
+			machineType:    "n4-highcpu-4xyz",
+			expVolumeLimit: volumeLimitSmall,
+			expectError:    true,
+		},
+		{
+			name:           "x4-megamem-960-metal",
+			machineType:    "x4-megamem-960-metal",
+			expVolumeLimit: x4HyperdiskLimit,
+		},
+		{
+			name:           "a4-highgpu-8g",
+			machineType:    "a4-highgpu-8g",
+			expVolumeLimit: a4HyperdiskLimit,
+		},
 	}
 
 	for _, tc := range testCases {
 		t.Logf("Test case: %s", tc.name)
 		metadataservice.SetMachineType(tc.machineType)
 		res, err := ns.NodeGetInfo(context.Background(), req)
-		if err != nil {
+		if err != nil && !tc.expectError {
 			t.Fatalf("Failed to get node info: %v", err)
 		} else {
 			volumeLimit := res.GetMaxVolumesPerNode()