Ensure GPU AMI is used when using multiple instances (#3924)

* Ensure GPU AMI is used when using multiple instances

* Add tests for instance type selection

pkg/eks/api.go

@@ -330,7 +330,7 @@ func ResolveAMI(provider api.ClusterProvider, version string, np api.NodePool) e
 		return errors.Errorf("invalid AMI value: %q", ng.AMI)
 	}
 
-	instanceType := selectInstanceType(np)
+	instanceType := SelectInstanceType(np)
 	id, err := resolver.Resolve(provider.Region(), version, instanceType, ng.AMIFamily)
 	if err != nil {
 		return errors.Wrap(err, "unable to determine AMI to use")
@@ -342,18 +342,30 @@ func ResolveAMI(provider api.ClusterProvider, version string, np api.NodePool) e
 	return nil
 }
 
-// selectInstanceType determines which instanceType is relevant for selecting an AMI
+// SelectInstanceType determines which instanceType is relevant for selecting an AMI
 // If the nodegroup has mixed instances it will prefer a GPU instance type over a general class one
 // This is to make sure that the AMI that is selected later is valid for all the types
-func selectInstanceType(np api.NodePool) string {
-	if ng, ok := np.(*api.NodeGroup); ok && api.HasMixedInstances(ng) {
-		for _, instanceType := range ng.InstancesDistribution.InstanceTypes {
+func SelectInstanceType(np api.NodePool) string {
+	var instanceTypes []string
+	switch ng := np.(type) {
+	case *api.NodeGroup:
+		if ng.InstancesDistribution != nil {
+			instanceTypes = ng.InstancesDistribution.InstanceTypes
+		}
+	case *api.ManagedNodeGroup:
+		instanceTypes = ng.InstanceTypes
+	}
+
+	hasMixedInstances := len(instanceTypes) > 0
+	if hasMixedInstances {
+		for _, instanceType := range instanceTypes {
 			if utils.IsGPUInstanceType(instanceType) {
 				return instanceType
 			}
 		}
-		return ng.InstancesDistribution.InstanceTypes[0]
+		return instanceTypes[0]
 	}
+
 	return np.BaseNodeGroup().InstanceType
 }
 

pkg/eks/instance_selection_test.go

@@ -0,0 +1,78 @@
+package eks_test
+
+import (
+	. "github.com/onsi/ginkgo/extensions/table"
+	. "github.com/onsi/gomega"
+	api "github.com/weaveworks/eksctl/pkg/apis/eksctl.io/v1alpha5"
+	"github.com/weaveworks/eksctl/pkg/eks"
+)
+
+type instanceTypeCase struct {
+	instanceTypes []string
+	instanceType  string
+
+	expectedInstanceType string
+}
+
+var _ = DescribeTable("Instance type selection", func(t instanceTypeCase) {
+	var (
+		ng  *api.NodeGroup
+		mng *api.ManagedNodeGroup
+	)
+
+	if len(t.instanceTypes) > 0 {
+		ng = &api.NodeGroup{
+			InstancesDistribution: &api.NodeGroupInstancesDistribution{
+				InstanceTypes: t.instanceTypes,
+			},
+		}
+		mng = &api.ManagedNodeGroup{
+			InstanceTypes: t.instanceTypes,
+		}
+	} else {
+		ngBase := &api.NodeGroupBase{
+			InstanceType: t.instanceType,
+		}
+		ng = &api.NodeGroup{
+			NodeGroupBase: ngBase,
+		}
+		mng = &api.ManagedNodeGroup{
+			NodeGroupBase: ngBase,
+		}
+	}
+
+	for _, np := range []api.NodePool{ng, mng} {
+		instanceType := eks.SelectInstanceType(np)
+		Expect(instanceType).To(Equal(t.expectedInstanceType))
+	}
+},
+	Entry("all ARM instances", instanceTypeCase{
+		instanceTypes: []string{"t4g.xlarge", "m6g.xlarge", "r6g.xlarge"},
+
+		expectedInstanceType: "t4g.xlarge",
+	}),
+
+	Entry("one GPU instance", instanceTypeCase{
+		instanceTypes: []string{"t2.medium", "t4.large", "g3.xlarge"},
+
+		expectedInstanceType: "g3.xlarge",
+	}),
+
+	Entry("all GPU instances", instanceTypeCase{
+		instanceTypes: []string{"p2.large", "p3.large", "g3.large"},
+
+		expectedInstanceType: "p2.large",
+	}),
+
+	Entry("single instance type", instanceTypeCase{
+		instanceType: "t4.large",
+
+		expectedInstanceType: "t4.large",
+	}),
+
+	Entry("single GPU instance type", instanceTypeCase{
+		instanceType: "t4.large",
+
+		expectedInstanceType: "t4.large",
+	}),
+)