Improve AWS instance parsing

pkg/lib/aws/ec2.go

@@ -30,20 +30,91 @@ import (
 	s "github.com/cortexlabs/cortex/pkg/lib/strings"
 )
 
-// aws instance types take this form: (\w+)([0-9]+)(\w*).(\w+)
-// the first group is the instance series, e.g. "m", "t", "g", "inf", ...
-// the second group is a version number for that series, e.g. 3, 4, ...
-// the third group is optional, and is a set of single-character "flags"
+var _digitsRegex = regexp.MustCompile(`[0-9]+`)
+
+type ParsedInstanceType struct {
+	Family       string
+	Generation   int
+	Capabilities strset.Set
+	Size         string
+}
+
+// Checks weather the input is an AWS instance type
+func IsValidInstanceType(instanceType string) bool {
+	return AllInstanceTypes.Has(instanceType)
+}
+
+// Checks whether the input is an AWS instance type
+func CheckValidInstanceType(instanceType string) error {
+	if !IsValidInstanceType(instanceType) {
+		return ErrorInvalidInstanceType(instanceType)
+	}
+	return nil
+}
+
+// AWS instance types take the form of: [family][generation][capabilities].[size]
+// the first group is the instance family, e.g. "m", "t", "g", "inf", ...
+// the second group is a generation number for that series, e.g. 3, 4, ...
+// the third group is optional, and is a set of single-character capabilities
 //   "g" represents ARM (graviton), "a" for AMD, "n" for fast networking, "d" for fast storage, etc.
 // the fourth and final group (after the dot) is the instance size, e.g. "large"
-var _armInstanceCapabilityRegex = regexp.MustCompile(`^\w+[0-9]+\w*g\w*\.\w+$`)
+func ParseInstanceType(instanceType string) (ParsedInstanceType, error) {
+	if err := CheckValidInstanceType(instanceType); err != nil {
+		return ParsedInstanceType{}, err
+	}
 
-// instanceType is assumed to be a valid instance type that exists in AWS, e.g. g4dn.xlarge
-func IsARMInstance(instanceType string) bool {
-	if strings.HasPrefix(instanceType, "a") {
-		return true
+	parts := strings.Split(instanceType, ".")
+	if len(parts) != 2 {
+		return ParsedInstanceType{}, errors.ErrorUnexpected("unexpected invalid instance type: " + instanceType)
+	}
+
+	prefix := parts[0]
+	size := parts[1]
+
+	digitSets := _digitsRegex.FindAllString(prefix, -1)
+	if len(digitSets) == 0 {
+		return ParsedInstanceType{}, errors.ErrorUnexpected("unexpected invalid instance type: " + instanceType)
 	}
-	return _armInstanceCapabilityRegex.MatchString(instanceType)
+
+	prefixParts := _digitsRegex.Split(prefix, -1)
+	capabilitiesStr := prefixParts[len(prefixParts)-1]
+	capabilities := strset.FromSlice(strings.Split(capabilitiesStr, ""))
+
+	generationStr := digitSets[len(digitSets)-1]
+	generation, ok := s.ParseInt(generationStr)
+	if !ok {
+		return ParsedInstanceType{}, errors.ErrorUnexpected("unexpected invalid instance type: " + instanceType)
+	}
+
+	generationIndex := strings.LastIndex(prefix, generationStr)
+	if generationIndex == -1 {
+		return ParsedInstanceType{}, errors.ErrorUnexpected("unexpected invalid instance type: " + instanceType)
+	}
+	family := prefix[:generationIndex]
+
+	return ParsedInstanceType{
+		Family:       family,
+		Generation:   generation,
+		Capabilities: capabilities,
+		Size:         size,
+	}, nil
+}
+
+func IsARMInstance(instanceType string) (bool, error) {
+	parsedType, err := ParseInstanceType(instanceType)
+	if err != nil {
+		return false, err
+	}
+
+	if parsedType.Family == "a" {
+		return true, nil
+	}
+
+	if parsedType.Capabilities.Has("g") {
+		return true, nil
+	}
+
+	return false, nil
 }
 
 func (c *Client) SpotInstancePrice(instanceType string) (float64, error) {

pkg/lib/aws/ec2_test.go

@@ -0,0 +1,63 @@
+/*
+Copyright 2021 Cortex Labs, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package aws
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/cortexlabs/cortex/pkg/lib/sets/strset"
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseInstanceType(t *testing.T) {
+	var testcases = []struct {
+		instanceType string
+		expected     ParsedInstanceType
+	}{
+		{"t3.small", ParsedInstanceType{"t", 3, strset.New(), "small"}},
+		{"g4dn.xlarge", ParsedInstanceType{"g", 4, strset.New("d", "n"), "xlarge"}},
+		{"inf1.24xlarge", ParsedInstanceType{"inf", 1, strset.New(), "24xlarge"}},
+		{"u-9tb1.metal", ParsedInstanceType{"u-9tb", 1, strset.New(), "metal"}},
+	}
+
+	invalidTypes := []string{
+		"badtype",
+		"badtype.large",
+		"badtype1.large",
+		"badtype2ad.large",
+	}
+
+	for _, testcase := range testcases {
+		parsed, err := ParseInstanceType(testcase.instanceType)
+		require.NoError(t, err)
+		require.Equal(t, testcase.expected.Family, parsed.Family, fmt.Sprintf("unexpected family for input: %s", testcase.instanceType))
+		require.Equal(t, testcase.expected.Generation, parsed.Generation, fmt.Sprintf("unexpected generation for input: %s", testcase.instanceType))
+		require.ElementsMatch(t, testcase.expected.Capabilities.Slice(), parsed.Capabilities.Slice(), fmt.Sprintf("unexpected capabilities for input: %s", testcase.instanceType))
+		require.Equal(t, testcase.expected.Size, parsed.Size, fmt.Sprintf("unexpected size for input: %s", testcase.instanceType))
+	}
+
+	for _, instanceType := range invalidTypes {
+		_, err := ParseInstanceType(instanceType)
+		require.Error(t, err)
+	}
+
+	for instanceType := range AllInstanceTypes {
+		_, err := ParseInstanceType(instanceType)
+		require.NoError(t, err)
+	}
+}

pkg/lib/aws/elb.go

@@ -28,10 +28,18 @@ import (
 // https://docs.aws.amazon.com/elasticloadbalancing/latest/network/load-balancer-target-groups.html
 var _nlbUnsupportedInstancePrefixes = strset.New("c1", "cc1", "cc2", "cg1", "cg2", "cr1", "g1", "g2", "hi1", "hs1", "m1", "m2", "m3", "t1")
 
-// instanceType must be a valid instance type that exists in AWS, e.g. g4dn.xlarge
-func IsInstanceSupportedByNLB(instanceType string) bool {
-	instancePrefix := strings.Split(instanceType, ".")[0]
-	return !_nlbUnsupportedInstancePrefixes.Has(instancePrefix)
+func IsInstanceSupportedByNLB(instanceType string) (bool, error) {
+	if err := CheckValidInstanceType(instanceType); err != nil {
+		return false, err
+	}
+
+	for prefix := range _nlbUnsupportedInstancePrefixes {
+		if strings.HasPrefix(instanceType, prefix) {
+			return false, nil
+		}
+	}
+
+	return true, nil
 }
 
 // returns the the first load balancer which has all of the specified tags, or nil if no load balancers match

pkg/lib/aws/errors.go

@@ -28,6 +28,7 @@ import (
 )
 
 const (
+	ErrInvalidInstanceType          = "aws.invalid_instance_type"
 	ErrInvalidAWSCredentials        = "aws.invalid_aws_credentials"
 	ErrInvalidS3aPath               = "aws.invalid_s3a_path"
 	ErrInvalidS3Path                = "aws.invalid_s3_path"
@@ -90,6 +91,13 @@ func IsErrCode(err error, errorCode string) bool {
 	return false
 }
 
+func ErrorInvalidInstanceType(instanceType string) error {
+	return errors.WithStack(&errors.Error{
+		Kind:    ErrInvalidInstanceType,
+		Message: fmt.Sprintf("%s is not an AWS instance type (e.g. m5.large is a valid instance type)", s.UserStr(instanceType)),
+	})
+}
+
 func ErrorInvalidAWSCredentials(awsErr error) error {
 	awsErrMsg := errors.Message(awsErr)
 	return errors.WithStack(&errors.Error{