From f902cc224c22e344ea718dca352dfc2b0374cbfa Mon Sep 17 00:00:00 2001
From: Mazen Selim <mazensel@amazon.com>
Date: Tue, 3 Dec 2024 23:15:37 +0000
Subject: [PATCH] Dynamically filter out invalid instance types

---
 .../internal/deployers/eksapi/nodegroup.go    | 87 +++++++++++++++++--
 1 file changed, 81 insertions(+), 6 deletions(-)

diff --git a/kubetest2/internal/deployers/eksapi/nodegroup.go b/kubetest2/internal/deployers/eksapi/nodegroup.go
index 5095c05df..12bd21fa5 100644
--- a/kubetest2/internal/deployers/eksapi/nodegroup.go
+++ b/kubetest2/internal/deployers/eksapi/nodegroup.go
@@ -6,6 +6,7 @@ import (
 	_ "embed"
 	"errors"
 	"fmt"
+	"slices"
 	"strconv"
 	"strings"
 	"time"
@@ -18,6 +19,7 @@ import (
 	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
 	"github.com/aws/aws-sdk-go-v2/service/eks"
 	ekstypes "github.com/aws/aws-sdk-go-v2/service/eks/types"
+	"github.com/aws/smithy-go"
 	"k8s.io/klog/v2"
 
 	"github.com/aws/aws-k8s-tester/kubetest2/internal/deployers/eksapi/templates"
@@ -78,9 +80,9 @@ func (m *NodegroupManager) createNodegroup(infra *Infrastructure, cluster *Clust
 				return fmt.Errorf("failed to describe AMI when populating default instance types: %s: %v", opts.AMI, err)
 			} else {
 				amiArch := out.Images[0].Architecture
-				defaultInstanceTypes, ok := defaultInstanceTypesByEC2ArchitectureValues[amiArch]
-				if !ok {
-					return fmt.Errorf("no default instance types known for AMI architecture: %v", amiArch)
+				defaultInstanceTypes, err := m.getValidDefaultInstanceTypesByEC2Arch(amiArch)
+				if err != nil {
+					return err
 				}
 				opts.InstanceTypes = defaultInstanceTypes
 				klog.V(2).Infof("Using default instance types for AMI architecture: %v: %v", amiArch, opts.InstanceTypes)
@@ -116,9 +118,9 @@ func (m *NodegroupManager) createManagedNodegroup(infra *Infrastructure, cluster
 	} else {
 		// managed nodegroups uses a t3.medium by default at the time of writing
 		// this only supports 17 pods, which can cause some flakes in the k8s e2e suite
-		defaultInstanceTypes, ok := defaultInstanceTypesByEKSAMITypes[input.AmiType]
-		if !ok {
-			return fmt.Errorf("no default instance types known for AmiType: %v", input.AmiType)
+		defaultInstanceTypes, err := m.getValidDefaultInstanceTypesByEKSAMIType(input.AmiType)
+		if err != nil {
+			return err
 		}
 		input.InstanceTypes = defaultInstanceTypes
 	}
@@ -532,3 +534,76 @@ func (m *NodegroupManager) getSubnetWithCapacity(infra *Infrastructure, opts *de
 	klog.Infof("Using subnet: %s", subnetId)
 	return subnetId, capacityReservationId, nil
 }
+
+func (m *NodegroupManager) getValidDefaultInstanceTypesByEKSAMIType(amiType ekstypes.AMITypes) ([]string, error) {
+	defaults, ok := defaultInstanceTypesByEKSAMITypes[amiType]
+	if !ok {
+		return []string{}, fmt.Errorf("no default instance types known for AmiType: %v", amiType)
+	}
+
+	return m.getValidInstanceTypesFromList(defaults), nil
+}
+
+func (m *NodegroupManager) getValidDefaultInstanceTypesByEC2Arch(arch ec2types.ArchitectureValues) ([]string, error) {
+	defaults, ok := defaultInstanceTypesByEC2ArchitectureValues[arch]
+	if !ok {
+		return []string{}, fmt.Errorf("no default instance types known for AMI architecture: %v", arch)
+	}
+
+	return m.getValidInstanceTypesFromList(defaults), nil
+}
+
+func (m *NodegroupManager) getValidInstanceTypesFromList(desiredInstanceTypes []string) []string {
+	desiredEc2InstanceTypes := convertToEc2InstanceTypes(desiredInstanceTypes)
+	_, err := m.clients.EC2().DescribeInstanceTypes(context.TODO(), &ec2.DescribeInstanceTypesInput{
+		InstanceTypes: desiredEc2InstanceTypes,
+	})
+
+	invalidInstances := []string{}
+	if err != nil {
+		var apierr smithy.APIError
+		if errors.As(err, &apierr) {
+			errorCode := apierr.ErrorCode()
+			if errorCode == "InvalidInstanceType" {
+				invalidInstances = parseInvalidInstanceTypesStrings(apierr.ErrorMessage())
+			} else {
+				klog.Infof("Could not determine invalid instance types, received error: %s", apierr.ErrorCode())
+			}
+		}
+	}
+
+	return parseValidInstanceTypes(desiredInstanceTypes, invalidInstances)
+}
+
+func convertToEc2InstanceTypes(instanceTypes []string) []ec2types.InstanceType {
+	desiredEc2InstanceTypes := make([]ec2types.InstanceType, len(instanceTypes))
+	for i, instanceType := range instanceTypes {
+		desiredEc2InstanceTypes[i] = ec2types.InstanceType(instanceType)
+	}
+	return desiredEc2InstanceTypes
+}
+
+func parseInvalidInstanceTypesStrings(errString string) []string {
+	invalidInstancesStartIndex := strings.LastIndex(errString, "[") + 1 // trim leading '['
+	invalidInstancesEndIndex := len(errString) - 1                      // trim trailing ']'
+	instanceSeparator := ", "                                           // divider between instance types in output
+
+	invalidInstancesStr := errString[invalidInstancesStartIndex:invalidInstancesEndIndex]
+	invalidInstancesArr := strings.Split(invalidInstancesStr, instanceSeparator)
+
+	return invalidInstancesArr
+}
+
+func parseValidInstanceTypes(desiredInstances []string, invalidInstances []string) []string {
+	validInstances := []string{}
+
+	for _, instanceType := range desiredInstances {
+		if !slices.Contains(invalidInstances, instanceType) {
+			validInstances = append(validInstances, instanceType)
+		} else {
+			klog.Infof("Eliminating instance type %s as an option", instanceType)
+		}
+	}
+
+	return validInstances
+}