Fix AWS CSI scale-up from zero

Author: Debasish-87

cluster-autoscaler/cloudprovider/aws/aws_cloud_provider.go

@@ -426,6 +426,9 @@ func (ng *AwsNodeGroup) TemplateNodeInfo() (*framework.NodeInfo, error) {
 	}
 
 	nodeInfo := framework.NewNodeInfo(node, nil, framework.NewPodInfo(cloudprovider.BuildKubeProxy(ng.asg.Name), nil))
+	if csiNode := ng.awsManager.buildCSINodeFromTemplate(template, node.Name); csiNode != nil {
+		nodeInfo.SetCSINode(csiNode)
+	}
 	return nodeInfo, nil
 }
 

cluster-autoscaler/cloudprovider/aws/aws_cloud_provider_test.go

@@ -872,3 +872,50 @@ func TestDeleteNodesWithPlaceholderAndStaleCache(t *testing.T) {
 	a.AssertNumberOfCalls(t, "TerminateInstanceInAutoScalingGroup", 2)
 
 }
+
+func TestAwsNodeGroupTemplateNodeInfoSetsCSINode(t *testing.T) {
+	const instanceTypeName = "m5.large"
+
+	manager := &AwsManager{
+		instanceTypes: map[string]*InstanceType{
+			instanceTypeName: {
+				InstanceType:   instanceTypeName,
+				VCPU:           2,
+				MemoryMb:       8192,
+				Architecture:   "amd64",
+				EBSVolumeLimit: 39,
+			},
+		},
+	}
+
+	origGetInstanceTypeFunc := getInstanceTypeForAsg
+	defer func() { getInstanceTypeForAsg = origGetInstanceTypeFunc }()
+	getInstanceTypeForAsg = func(_ *asgCache, _ *asg) (string, error) {
+		return instanceTypeName, nil
+	}
+
+	ng := &AwsNodeGroup{
+		awsManager: manager,
+		asg: &asg{
+			AwsRef:            AwsRef{Name: "test-asg"},
+			AvailabilityZones: []string{"us-east-1a"},
+			minSize:           0,
+			maxSize:           5,
+			curSize:           0,
+		},
+	}
+
+	nodeInfo, err := ng.TemplateNodeInfo()
+	assert.NoError(t, err)
+	assert.NotNil(t, nodeInfo)
+	assert.NotNil(t, nodeInfo.CSINode)
+	assert.Equal(t, nodeInfo.Node().Name, nodeInfo.CSINode.Name)
+	assert.Len(t, nodeInfo.CSINode.Spec.Drivers, 1)
+
+	driver := nodeInfo.CSINode.Spec.Drivers[0]
+	assert.Equal(t, "ebs.csi.aws.com", driver.Name)
+	assert.Equal(t, nodeInfo.Node().Name, driver.NodeID)
+	assert.NotNil(t, driver.Allocatable)
+	assert.NotNil(t, driver.Allocatable.Count)
+	assert.Equal(t, int32(39), *driver.Allocatable.Count)
+}

cluster-autoscaler/cloudprovider/aws/aws_manager.go

@@ -27,17 +27,18 @@ import (
 	"strings"
 	"time"
 
-	apiv1 "k8s.io/api/core/v1"
-	"k8s.io/apimachinery/pkg/api/resource"
-	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
-	"k8s.io/klog/v2"
-
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/autoscaling"
 	autoscalingtypes "github.com/aws/aws-sdk-go-v2/service/autoscaling/types"
 	"github.com/aws/aws-sdk-go-v2/service/ec2"
 	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/aws/aws-sdk-go-v2/service/eks"
+	apiv1 "k8s.io/api/core/v1"
+	storagev1 "k8s.io/api/storage/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/klog/v2"
+	"k8s.io/utils/ptr"
 	"sigs.k8s.io/cluster-autoscaler/pkg/cloudprovider"
 	"sigs.k8s.io/cluster-autoscaler/pkg/config"
 	"sigs.k8s.io/cluster-autoscaler/pkg/utils/gpu"
@@ -341,6 +342,33 @@ func (m *AwsManager) buildNodeFromTemplate(asg *asg, template *asgTemplate) (*ap
 	return &node, nil
 }
 
+func (m *AwsManager) buildCSINodeFromTemplate(template *asgTemplate, nodeName string) *storagev1.CSINode {
+	if template == nil || template.InstanceType == nil {
+		return nil
+	}
+
+	driver := storagev1.CSINodeDriver{
+		Name:   "ebs.csi.aws.com",
+		NodeID: nodeName,
+	}
+	// Prefer a CSINode without Allocatable over nil when the attachment limit is
+	// unknown. Returning nil blocks CSI-aware scale-from-zero simulation.
+	if template.InstanceType.EBSVolumeLimit > 0 {
+		driver.Allocatable = &storagev1.VolumeNodeResources{
+			Count: ptr.To(int32(template.InstanceType.EBSVolumeLimit)),
+		}
+	}
+
+	return &storagev1.CSINode{
+		ObjectMeta: metav1.ObjectMeta{
+			Name: nodeName,
+		},
+		Spec: storagev1.CSINodeSpec{
+			Drivers: []storagev1.CSINodeDriver{driver},
+		},
+	}
+}
+
 func joinNodeLabelsChoosingUserValuesOverAPIValues(extractedLabels map[string]string, mngLabels map[string]string) map[string]string {
 	result := make(map[string]string)
 

cluster-autoscaler/cloudprovider/aws/aws_manager_test.go

@@ -951,3 +951,66 @@ func TestParseASGAutoDiscoverySpecs(t *testing.T) {
 		})
 	}
 }
+
+func TestBuildCSINodeFromTemplate(t *testing.T) {
+	manager := &AwsManager{}
+
+	template := &asgTemplate{
+		InstanceType: &InstanceType{
+			InstanceType:   "m5.large",
+			EBSVolumeLimit: 39,
+		},
+	}
+
+	got := manager.buildCSINodeFromTemplate(template, "template-node")
+
+	require.NotNil(t, got)
+	require.Equal(t, "template-node", got.Name)
+	require.Len(t, got.Spec.Drivers, 1)
+
+	driver := got.Spec.Drivers[0]
+	require.Equal(t, "ebs.csi.aws.com", driver.Name)
+	require.Equal(t, "template-node", driver.NodeID)
+	require.NotNil(t, driver.Allocatable)
+	require.Equal(t, int32(39), *driver.Allocatable.Count)
+}
+
+func TestBuildCSINodeFromTemplate_NoVolumeLimit(t *testing.T) {
+	manager := &AwsManager{}
+
+	template := &asgTemplate{
+		InstanceType: &InstanceType{
+			InstanceType:   "m5.large",
+			EBSVolumeLimit: 0,
+		},
+	}
+
+	got := manager.buildCSINodeFromTemplate(template, "template-node")
+
+	require.NotNil(t, got)
+	require.Equal(t, "template-node", got.Name)
+	require.Len(t, got.Spec.Drivers, 1)
+
+	driver := got.Spec.Drivers[0]
+	require.Equal(t, "ebs.csi.aws.com", driver.Name)
+	require.Equal(t, "template-node", driver.NodeID)
+	require.Nil(t, driver.Allocatable)
+}
+
+func TestBuildCSINodeFromTemplate_NilTemplate(t *testing.T) {
+	manager := &AwsManager{}
+
+	got := manager.buildCSINodeFromTemplate(nil, "template-node")
+
+	require.Nil(t, got)
+}
+
+func TestBuildCSINodeFromTemplate_NilInstanceType(t *testing.T) {
+	manager := &AwsManager{}
+
+	template := &asgTemplate{}
+
+	got := manager.buildCSINodeFromTemplate(template, "template-node")
+
+	require.Nil(t, got)
+}

cluster-autoscaler/cloudprovider/aws/aws_util.go

@@ -75,6 +75,9 @@ func transformInstanceType(rawInstanceType *ec2types.InstanceTypeInfo) *Instance
 	if rawInstanceType.ProcessorInfo != nil && len(rawInstanceType.ProcessorInfo.SupportedArchitectures) > 0 {
 		instanceType.Architecture = interpretEc2SupportedArchitecure(string(rawInstanceType.ProcessorInfo.SupportedArchitectures[0]))
 	}
+	if rawInstanceType.EbsInfo != nil && rawInstanceType.EbsInfo.MaximumEbsAttachments != nil {
+		instanceType.EBSVolumeLimit = int64(*rawInstanceType.EbsInfo.MaximumEbsAttachments)
+	}
 	return instanceType
 }
 

cluster-autoscaler/cloudprovider/aws/aws_util_test.go

@@ -41,6 +41,9 @@ func TestInstanceTypeTransform(t *testing.T) {
 		MemoryInfo: &ec2types.MemoryInfo{
 			SizeInMiB: aws.Int64(7680),
 		},
+		EbsInfo: &ec2types.EbsInfo{
+			MaximumEbsAttachments: aws.Int32(39),
+		},
 	}
 
 	instanceType := transformInstanceType(&rawInstanceType)
@@ -50,6 +53,27 @@ func TestInstanceTypeTransform(t *testing.T) {
 	assert.Equal(t, int64(7680), instanceType.MemoryMb)
 	assert.Equal(t, int64(0), instanceType.GPU)
 	assert.Equal(t, "amd64", instanceType.Architecture)
+	assert.Equal(t, int64(39), instanceType.EBSVolumeLimit)
+}
+
+func TestInstanceTypeTransform_NilEbsInfo(t *testing.T) {
+	rawInstanceType := ec2types.InstanceTypeInfo{
+		InstanceType: ec2types.InstanceType("c4.xlarge"),
+		ProcessorInfo: &ec2types.ProcessorInfo{
+			SupportedArchitectures: []ec2types.ArchitectureType{ec2types.ArchitectureTypeX8664},
+		},
+		VCpuInfo: &ec2types.VCpuInfo{
+			DefaultVCpus: aws.Int32(4),
+		},
+		MemoryInfo: &ec2types.MemoryInfo{
+			SizeInMiB: aws.Int64(7680),
+		},
+	}
+
+	instanceType := transformInstanceType(&rawInstanceType)
+
+	assert.Equal(t, "c4.xlarge", instanceType.InstanceType)
+	assert.Equal(t, int64(0), instanceType.EBSVolumeLimit)
 }
 
 func TestInterpretEc2SupportedArchitecure(t *testing.T) {

cluster-autoscaler/cloudprovider/aws/ec2_instance_types.go

@@ -20,11 +20,12 @@ package aws
 
 // InstanceType is spec of EC2 instance
 type InstanceType struct {
-	InstanceType string
-	VCPU         int64
-	MemoryMb     int64
-	GPU          int64
-	Architecture string
+	InstanceType   string
+	VCPU           int64
+	MemoryMb       int64
+	GPU            int64
+	Architecture   string
+	EBSVolumeLimit int64
 }
 
 // StaticListLastUpdateTime is a string declaring the last time the static list was updated.

cluster-autoscaler/cloudprovider/aws/ec2_instance_types/gen.go

@@ -59,6 +59,7 @@ type InstanceType struct {
 	MemoryMb     int64
 	GPU          int64
 	Architecture string
+	EBSVolumeLimit int64
 }
 
 // StaticListLastUpdateTime is a string declaring the last time the static list was updated.
@@ -73,6 +74,7 @@ var InstanceTypes = map[string]*InstanceType{
 		MemoryMb:     {{ .MemoryMb }},
 		GPU:          {{ .GPU }},
 		Architecture: "{{ .Architecture }}",
+		EBSVolumeLimit:  {{ .EBSVolumeLimit }},
 	},
 {{- end }}
 }