Adjust Attach Limits for Gen3 Machines + Add Labels Override

Author: sunnylovestiramisu

pkg/common/constants.go

@@ -60,6 +60,9 @@ const (
 	// Node label for Data Cache (only applicable to GKE nodes)
 	NodeLabelPrefix         = "cloud.google.com/%s"
 	DataCacheLssdCountLabel = "gke-data-cache-disk"
+	// Node label for attach limit override
+	NodeRestrictionLabelPrefix = "node-restriction.kubernetes.io/%s"
+	AttachLimitOverrideLabel   = "gke-volume-attach-limit-override"
 )
 
 // doc https://cloud.google.com/compute/docs/disks/hyperdisks#max-total-disks-per-vm

pkg/common/utils.go

@@ -24,6 +24,7 @@ import (
 	"net/http"
 	"regexp"
 	"slices"
+	"strconv"
 	"strings"
 	"time"
 
@@ -772,6 +773,25 @@ func MapNumber(num int64) int64 {
 	return 0
 }
 
+func ExtractCPUFromMachineType(input string) (int64, error) {
+	// Regex to find the number at the end of the string,
+	// it allows optional -lssd suffix.
+	re := regexp.MustCompile(`(\d+)(?:-lssd|-metal)?$`)
+
+	match := re.FindStringSubmatch(input)
+	if len(match) < 2 {
+		return 0, fmt.Errorf("no number found at the end of the input string: %s", input)
+	}
+
+	numberStr := match[1]
+	number, err := strconv.ParseInt(numberStr, 10, 64)
+	if err != nil {
+		return 0, fmt.Errorf("failed to convert string '%s' to integer: %w", numberStr, err)
+	}
+
+	return number, nil
+}
+
 func DiskTypeLabelKey(diskType string) string {
 	return fmt.Sprintf("%s/%s", DiskTypeKeyPrefix, diskType)
 }

pkg/common/utils_test.go

@@ -2149,3 +2149,47 @@ func TestGetMinIopsThroughput(t *testing.T) {
 		})
 	}
 }
+
+func TestExtractCPUFromMachineType(t *testing.T) {
+	testcases := []struct {
+		name         string
+		input        string
+		expectOutput int64
+		expectErr    bool
+	}{
+		{
+			name:         "c3-highmem-176",
+			input:        "c3-highmem-176",
+			expectOutput: 176,
+		},
+		{
+			name:         "c3-standard-8-lssd",
+			input:        "c3-standard-8-lssd",
+			expectOutput: 8,
+		},
+		{
+			name:         "c3-standard-192-metal",
+			input:        "c3-standard-192-metal",
+			expectOutput: 192,
+		},
+		{
+			name:         "invalid input",
+			input:        "something-not-valid",
+			expectOutput: 0,
+			expectErr:    true,
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(tc.name, func(t *testing.T) {
+			output, err := ExtractCPUFromMachineType(tc.input)
+			if output != tc.expectOutput {
+				t.Errorf("ExtractCPUFromMachineType: got %v, want %v", output, tc.expectOutput)
+			}
+
+			if gotErr := err != nil; gotErr != tc.expectErr {
+				t.Fatalf("ExtractCPUFromMachineType(%+v) = %v; expectedErr: %v", tc.input, err, tc.expectErr)
+			}
+		})
+	}
+}

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

@@ -31,6 +31,8 @@ import (
 
 	csi "github.com/container-storage-interface/spec/lib/go/csi"
 
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
 	"k8s.io/klog/v2"
 	"k8s.io/mount-utils"
 
@@ -101,7 +103,12 @@ const (
 	// 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
+	a4HyperdiskLimit int64 = 127
+	// doc https://cloud.google.com/compute/docs/storage-optimized-machines#z3_disks
+	// doc https://cloud.google.com/compute/docs/accelerator-optimized-machines#a3-disks
+	gen3HyperdiskLimit int64 = 31
+	// doc https://cloud.google.com/compute/docs/compute-optimized-machines#h3_disks
+	h3HyperdiskLimit     int64 = 7 // Use limit for Hyperdisk Balanced
 	defaultLinuxFsType         = "ext4"
 	defaultWindowsFsType       = "ntfs"
 	fsTypeExt3                 = "ext3"
@@ -571,7 +578,7 @@ 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()
+	volumeLimits, err := ns.GetVolumeLimits(ctx)
 	if err != nil {
 		klog.Errorf("GetVolumeLimits failed: %v", err.Error())
 	}
@@ -731,7 +738,7 @@ func (ns *GCENodeServer) NodeExpandVolume(ctx context.Context, req *csi.NodeExpa
 	}, nil
 }
 
-func (ns *GCENodeServer) GetVolumeLimits() (int64, error) {
+func (ns *GCENodeServer) GetVolumeLimits(ctx context.Context) (int64, error) {
 	// Machine-type format: n1-type-CPUS or custom-CPUS-RAM or f1/g1-type
 	machineType := ns.MetadataService.GetMachineType()
 
@@ -741,6 +748,22 @@ func (ns *GCENodeServer) GetVolumeLimits() (int64, error) {
 			return volumeLimitSmall, nil
 		}
 	}
+
+	// Get attach limit override from label
+	attachLimitOverride, err := GetAttachLimitsOverrideFromNodeLabel(ctx, ns.MetadataService.GetName())
+	if err == nil && attachLimitOverride > 0 && attachLimitOverride < 128 {
+		return attachLimitOverride, nil
+	} else {
+		// If there is an error or the range is not valid, still proceed to get defaults for the machine type
+		if err != nil {
+			klog.Warningf("using default value due to err getting node-restriction.kubernetes.io/gke-volume-attach-limit-override: %v", err)
+		}
+		if attachLimitOverride != 0 {
+			klog.Warningf("using default value due to invalid node-restriction.kubernetes.io/gke-volume-attach-limit-override: %d", attachLimitOverride)
+		}
+	}
+
+	// Process gen4 machine attach limits
 	gen4MachineTypesPrefix := []string{"c4a-", "c4-", "n4-"}
 	for _, gen4Prefix := range gen4MachineTypesPrefix {
 		if strings.HasPrefix(machineType, gen4Prefix) {
@@ -760,5 +783,59 @@ func (ns *GCENodeServer) GetVolumeLimits() (int64, error) {
 		}
 	}
 
+	// Process gen3 machine attach limits
+	gen3MachineTypesPrefix := []string{"c3-", "c3d-"}
+	for _, gen3Prefix := range gen3MachineTypesPrefix {
+		if strings.HasPrefix(machineType, gen3Prefix) {
+			cpus, err := common.ExtractCPUFromMachineType(machineType)
+			if err != nil {
+				return volumeLimitSmall, err
+			}
+			if cpus <= 8 || strings.Contains(machineType, "metal") {
+				return volumeLimitSmall, nil
+			}
+			return gen3HyperdiskLimit, nil
+
+		}
+		if strings.HasPrefix(machineType, "z3-") {
+			return gen3HyperdiskLimit, nil
+		}
+		if strings.HasPrefix(machineType, "h3-") {
+			return h3HyperdiskLimit, nil
+		}
+		if strings.HasPrefix(machineType, "a3-") {
+			if machineType == "a3-ultragpu-8g" {
+				return volumeLimitBig, nil
+			} else {
+				return gen3HyperdiskLimit, nil
+			}
+		}
+
+	}
+
 	return volumeLimitBig, nil
 }
+
+func GetAttachLimitsOverrideFromNodeLabel(ctx context.Context, nodeName string) (int64, error) {
+	cfg, err := rest.InClusterConfig()
+	if err != nil {
+		return 0, err
+	}
+	kubeClient, err := kubernetes.NewForConfig(cfg)
+	if err != nil {
+		return 0, err
+	}
+	node, err := getNodeWithRetry(ctx, kubeClient, nodeName)
+	if err != nil {
+		return 0, err
+	}
+	if val, found := node.GetLabels()[fmt.Sprintf(common.NodeRestrictionLabelPrefix, common.AttachLimitOverrideLabel)]; found {
+		attachLimitOverrideForNode, err := strconv.ParseInt(val, 10, 64)
+		if err != nil {
+			return 0, fmt.Errorf("error getting attach limit override from node label: %v", err)
+		}
+		klog.V(4).Infof("attach limit override for the node: %v", attachLimitOverrideForNode)
+		return attachLimitOverrideForNode, nil
+	}
+	return 0, nil
+}

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

@@ -309,6 +309,41 @@ func TestNodeGetVolumeLimits(t *testing.T) {
 			machineType:    "a4-highgpu-8g",
 			expVolumeLimit: a4HyperdiskLimit,
 		},
+		{
+			name:           "z3-highmem-176",
+			machineType:    "z3-highmem-176",
+			expVolumeLimit: gen3HyperdiskLimit,
+		},
+		{
+			name:           "h3-standard-88",
+			machineType:    "h3-standard-88",
+			expVolumeLimit: h3HyperdiskLimit,
+		},
+		{
+			name:           "a3-ultragpu-8g",
+			machineType:    "a3-ultragpu-8g",
+			expVolumeLimit: volumeLimitBig,
+		},
+		{
+			name:           "a3-megagpu-8g",
+			machineType:    "a3-megagpu-8g",
+			expVolumeLimit: gen3HyperdiskLimit, // 31
+		},
+		{
+			name:           "c3d-highmem-8-lssd",
+			machineType:    "c3d-highmem-8-lssd",
+			expVolumeLimit: volumeLimitSmall, // 15
+		},
+		{
+			name:           "c3-standard-192-metal",
+			machineType:    "c3-standard-192-metal",
+			expVolumeLimit: volumeLimitSmall, // 15
+		},
+		{
+			name:           "c3-standard-176",
+			machineType:    "c3-standard-176",
+			expVolumeLimit: gen3HyperdiskLimit, // 31
+		},
 	}
 
 	for _, tc := range testCases {