Merge gpu branch into dev

ecs-cli/modules/cli/cluster/cluster_app.go

@@ -17,7 +17,6 @@ import (
 	"bufio"
 	"fmt"
 	"os"
-	"regexp"
 	"strconv"
 	"strings"
 
@@ -72,10 +71,6 @@ const (
 	ParameterKeySpotPrice                = "SpotPrice"
 )
 
-const (
-	defaultARM64InstanceType = "a1.medium"
-)
-
 var flagNamesToStackParameterKeys map[string]string
 var requiredParameters []string = []string{ParameterKeyCluster}
 
@@ -317,20 +312,13 @@ func createCluster(context *cli.Context, awsClients *AWSClients, commandConfig *
 	}
 
 	if launchType == config.LaunchTypeEC2 {
-		architecture, err := determineArchitecture(cfnParams)
-		if err != nil {
-			return err
-		}
-
 		// Check if image id was supplied, else populate
 		_, err = cfnParams.GetParameter(ParameterKeyAmiId)
 		if err == cloudformation.ParameterNotFoundError {
-			amiMetadata, err := metadataClient.GetRecommendedECSLinuxAMI(architecture)
+			err := populateAMIID(cfnParams, metadataClient)
 			if err != nil {
 				return err
 			}
-			logrus.Infof("Using recommended %s AMI with ECS Agent %s and %s", amiMetadata.OsName, amiMetadata.AgentVersion, amiMetadata.RuntimeVersion)
-			cfnParams.Add(ParameterKeyAmiId, amiMetadata.ImageID)
 		} else if err != nil {
 			return err
 		}
@@ -393,26 +381,33 @@ func canEnableContainerInstanceTagging(client ecsclient.ECSClient) (bool, error)
 	return false, nil
 }
 
-func determineArchitecture(cfnParams *cloudformation.CfnStackParams) (string, error) {
-	architecture := amimetadata.ArchitectureTypeX86
+func retrieveInstanceType(cfnParams *cloudformation.CfnStackParams) (string, error) {
+	param, err := cfnParams.GetParameter(ParameterKeyInstanceType)
 
-	// a1 instances get the Arm based ECS AMI
-	instanceTypeParam, err := cfnParams.GetParameter(ParameterKeyInstanceType)
 	if err == cloudformation.ParameterNotFoundError {
-		logrus.Infof("Defaulting instance type to t2.micro")
-	} else if err != nil {
+		logrus.Infof("Defaulting instance type to %s", cloudformation.DefaultECSInstanceType)
+		return cloudformation.DefaultECSInstanceType, nil
+	}
+	if err != nil {
 		return "", err
-	} else {
-		instanceType := aws.StringValue(instanceTypeParam.ParameterValue)
-		// This regex matches all current a1 instances, and should work for any future additions as well
-		r := regexp.MustCompile("a1\\.(medium|\\d*x?large)")
-		if r.MatchString(instanceType) {
-			logrus.Infof("Using Arm ecs-optimized AMI because instance type was %s", instanceType)
-			architecture = amimetadata.ArchitectureTypeARM64
-		}
+	}
+	return aws.StringValue(param.ParameterValue), nil
+}
+
+func populateAMIID(cfnParams *cloudformation.CfnStackParams, client amimetadata.Client) error {
+	instanceType, err := retrieveInstanceType(cfnParams)
+	if err != nil {
+		return err
 	}
 
-	return architecture, nil
+	amiMetadata, err := client.GetRecommendedECSLinuxAMI(instanceType)
+	if err != nil {
+		return err
+	}
+	logrus.Infof("Using recommended %s AMI with ECS Agent %s and %s",
+		amiMetadata.OsName, amiMetadata.AgentVersion, amiMetadata.RuntimeVersion)
+	cfnParams.Add(ParameterKeyAmiId, amiMetadata.ImageID)
+	return nil
 }
 
 // unfortunately go SDK lacks a unified Tag type

ecs-cli/modules/cli/cluster/cluster_app_test.go

@@ -144,7 +144,7 @@ func TestClusterUpWithForce(t *testing.T) {
 	)
 
 	gomock.InOrder(
-		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("x86").Return(amiMetadata(amiID), nil),
+		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("t2.micro").Return(amiMetadata(amiID), nil),
 	)
 
 	gomock.InOrder(
@@ -179,7 +179,7 @@ func TestClusterUpWithoutPublicIP(t *testing.T) {
 	)
 
 	gomock.InOrder(
-		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("x86").Return(amiMetadata(amiID), nil),
+		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("t2.micro").Return(amiMetadata(amiID), nil),
 	)
 
 	gomock.InOrder(
@@ -232,7 +232,7 @@ func TestClusterUpWithUserData(t *testing.T) {
 	)
 
 	gomock.InOrder(
-		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("x86").Return(amiMetadata(amiID), nil),
+		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("t2.micro").Return(amiMetadata(amiID), nil),
 	)
 
 	gomock.InOrder(
@@ -281,7 +281,7 @@ func TestClusterUpWithSpotPrice(t *testing.T) {
 	)
 
 	gomock.InOrder(
-		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("x86").Return(amiMetadata(amiID), nil),
+		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("t2.micro").Return(amiMetadata(amiID), nil),
 	)
 
 	gomock.InOrder(
@@ -976,7 +976,7 @@ func TestClusterUpARM64(t *testing.T) {
 	)
 
 	gomock.InOrder(
-		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("arm64").Return(amiMetadata(armAMIID), nil),
+		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("a1.medium").Return(amiMetadata(armAMIID), nil),
 	)
 
 	gomock.InOrder(
@@ -1050,7 +1050,7 @@ func TestClusterUpWithTags(t *testing.T) {
 		}),
 	)
 	gomock.InOrder(
-		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("x86").Return(amiMetadata(amiID), nil),
+		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("t2.micro").Return(amiMetadata(amiID), nil),
 	)
 	gomock.InOrder(
 		mockCloudformation.EXPECT().ValidateStackExists(stackName).Return(errors.New("error")),
@@ -1131,7 +1131,7 @@ func TestClusterUpWithTagsContainerInstanceTaggingEnabled(t *testing.T) {
 		}),
 	)
 	gomock.InOrder(
-		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("x86").Return(amiMetadata(amiID), nil),
+		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("t2.micro").Return(amiMetadata(amiID), nil),
 	)
 	gomock.InOrder(
 		mockCloudformation.EXPECT().ValidateStackExists(stackName).Return(errors.New("error")),
@@ -1165,33 +1165,6 @@ func TestClusterUpWithTagsContainerInstanceTaggingEnabled(t *testing.T) {
 	assert.Equal(t, userdataMock.tags, expectedECSTags, "Expected tags to match")
 }
 
-func TestDetermineArchitecture(t *testing.T) {
-	var testCases = []struct {
-		in  string
-		out string
-	}{
-		{"a1.medium", "arm64"},
-		{"a1.large", "arm64"},
-		{"a1.xlarge", "arm64"},
-		{"a1.2xlarge", "arm64"},
-		{"a1.4xlarge", "arm64"},
-		{"t2.medium", "x86"},
-		{"c5.large", "x86"},
-		{"i3.metal", "x86"},
-		{"t3.micro", "x86"},
-	}
-
-	for _, tt := range testCases {
-		t.Run(tt.in, func(t *testing.T) {
-			cfnParams := cloudformation.NewCfnStackParams(requiredParameters)
-			cfnParams.Add(ParameterKeyInstanceType, tt.in)
-			arch, err := determineArchitecture(cfnParams)
-			assert.NoError(t, err, "Unexpected error determining architecture")
-			assert.Equal(t, tt.out, arch, "Expected architecture to match")
-		})
-	}
-}
-
 ///////////////////
 // Cluster Down //
 //////////////////
@@ -1397,7 +1370,7 @@ func mocksForSuccessfulClusterUp(mockECS *mock_ecs.MockECSClient, mockCloudforma
 		mockECS.EXPECT().CreateCluster(clusterName, gomock.Any()).Return(clusterName, nil),
 	)
 	gomock.InOrder(
-		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("x86").Return(amiMetadata(amiID), nil),
+		mockSSM.EXPECT().GetRecommendedECSLinuxAMI("t2.micro").Return(amiMetadata(amiID), nil),
 	)
 	gomock.InOrder(
 		mockCloudformation.EXPECT().ValidateStackExists(stackName).Return(errors.New("error")),

ecs-cli/modules/clients/aws/amimetadata/client.go

@@ -11,30 +11,36 @@
 // express or implied. See the License for the specific language governing
 // permissions and limitations under the License.
 
+// Package amimetadata provides AMI metadata given an instance type.
 package amimetadata
 
 import (
 	"encoding/json"
-
 	"github.com/aws/amazon-ecs-cli/ecs-cli/modules/clients"
 	"github.com/aws/amazon-ecs-cli/ecs-cli/modules/config"
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
 	"github.com/aws/aws-sdk-go/service/ssm"
 	"github.com/aws/aws-sdk-go/service/ssm/ssmiface"
 	"github.com/pkg/errors"
+	"github.com/sirupsen/logrus"
+	"regexp"
+	"strings"
 )
 
+// SSM parameter names to retrieve ECS optimized AMI.
+// See: https://docs.aws.amazon.com/AmazonECS/latest/developerguide/retrieve-ecs-optimized_AMI.html
 const (
-	amazonLinux2X86RecommendedParameterName   = "/aws/service/ecs/optimized-ami/amazon-linux-2/recommended"
-	amazonLinux2ARM64RecommendedParameterName = "/aws/service/ecs/optimized-ami/amazon-linux-2/arm64/recommended"
-)
-
-const (
-	ArchitectureTypeARM64 = "arm64"
-	ArchitectureTypeX86   = "x86"
+	amazonLinux2X86RecommendedParameterName    = "/aws/service/ecs/optimized-ami/amazon-linux-2/recommended"
+	amazonLinux2ARM64RecommendedParameterName  = "/aws/service/ecs/optimized-ami/amazon-linux-2/arm64/recommended"
+	amazonLinux2X86GPURecommendedParameterName = "/aws/service/ecs/optimized-ami/amazon-linux-2/gpu/recommended"
 )
 
+// AMIMetadata is returned through ssm:GetParameters and can be used to retrieve the ImageId
+// while launching instances.
+//
+// See: https://docs.aws.amazon.com/AmazonECS/latest/developerguide/retrieve-ecs-optimized_AMI.html
+// See: https://docs.aws.amazon.com/AWSCloudFormation/latest/UserGuide/aws-properties-as-launchconfig.html#cfn-as-launchconfig-imageid
 type AMIMetadata struct {
 	ImageID        string `json:"image_id"`
 	OsName         string `json:"os"`
@@ -47,13 +53,13 @@ type Client interface {
 	GetRecommendedECSLinuxAMI(string) (*AMIMetadata, error)
 }
 
-// ssmClient implements Client
+// metadataClient implements Client.
 type metadataClient struct {
 	client ssmiface.SSMAPI
 	region string
 }
 
-// NewSSMClient creates an instance of Client.
+// NewMetadataClient creates an instance of Client.
 func NewMetadataClient(commandConfig *config.CommandConfig) Client {
 	client := ssm.New(commandConfig.Session)
 	client.Handlers.Build.PushBackNamed(clients.CustomUserAgentHandler())
@@ -63,20 +69,31 @@ func NewMetadataClient(commandConfig *config.CommandConfig) Client {
 	}
 }
 
-func (c *metadataClient) GetRecommendedECSLinuxAMI(architecture string) (*AMIMetadata, error) {
-	ssmParam := amazonLinux2X86RecommendedParameterName
-	if architecture == ArchitectureTypeARM64 {
-		ssmParam = amazonLinux2ARM64RecommendedParameterName
+// GetRecommendedECSLinuxAMI returns the recommended Amazon ECS-Optimized AMI Metadata given the instance type.
+func (c *metadataClient) GetRecommendedECSLinuxAMI(instanceType string) (*AMIMetadata, error) {
+	if isARM64Instance(instanceType) {
+		logrus.Infof("Using Arm ecs-optimized AMI because instance type was %s", instanceType)
+		return c.parameterValueFor(amazonLinux2ARM64RecommendedParameterName)
+	}
+	if isGPUInstance(instanceType) {
+		logrus.Infof("Using GPU ecs-optimized AMI because instance type was %s", instanceType)
+		return c.parameterValueFor(amazonLinux2X86GPURecommendedParameterName)
 	}
+	return c.parameterValueFor(amazonLinux2X86RecommendedParameterName)
+}
 
+func (c *metadataClient) parameterValueFor(ssmParamName string) (*AMIMetadata, error) {
 	response, err := c.client.GetParameter(&ssm.GetParameterInput{
-		Name: aws.String(ssmParam),
+		Name: aws.String(ssmParamName),
 	})
 	if err != nil {
 		if aerr, ok := err.(awserr.Error); ok {
 			if aerr.Code() == ssm.ErrCodeParameterNotFound {
-				// Added for arm AMIs which are only supported in some regions
-				return nil, errors.Wrapf(err, "Could not find Recommended Amazon Linux 2 AMI in %s with architecture %s; the AMI may not be supported in this region", c.region, architecture)
+				// Added for AMIs which are only supported in some regions
+				return nil, errors.Wrapf(err,
+					"Could not find Recommended Amazon Linux 2 AMI %s in %s; the AMI may not be supported in this region",
+					ssmParamName,
+					c.region)
 			}
 		}
 		return nil, err
@@ -85,3 +102,24 @@ func (c *metadataClient) GetRecommendedECSLinuxAMI(architecture string) (*AMIMet
 	err = json.Unmarshal([]byte(aws.StringValue(response.Parameter.Value)), metadata)
 	return metadata, err
 }
+
+func isARM64Instance(instanceType string) bool {
+	r := regexp.MustCompile("a1\\.(medium|\\d*x?large)")
+	if r.MatchString(instanceType) {
+		return true
+	}
+	return false
+}
+
+func isGPUInstance(instanceType string) bool {
+	if strings.HasPrefix(instanceType, "p2.") {
+		return true
+	}
+	if strings.HasPrefix(instanceType, "p3.") {
+		return true
+	}
+	if strings.HasPrefix(instanceType, "p3dn.") {
+		return true
+	}
+	return false
+}