pre-Generate EC2 instance type details out of band

Signed-off-by: Davanum Srinivas <[email protected]>

cmd/ec2geninfo/main.go

@@ -0,0 +1,187 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"text/template"
+
+	"github.com/aws/aws-sdk-go-v2/aws"
+	"github.com/aws/aws-sdk-go-v2/config"
+	"github.com/aws/aws-sdk-go-v2/service/ec2"
+	"github.com/aws/aws-sdk-go-v2/service/ec2/types"
+)
+
+type InstanceInfo struct {
+	InstanceType             string
+	InstanceStorageSupported bool
+	EFASupported             bool
+	NvidiaGPUSupported       bool
+	NvidiaGPUType            string
+	NeuronSupported          bool
+	NeuronDeviceType         string
+	CBRSupported             bool
+	CPUArch                  string
+}
+
+const ec2InstancesTemplate = `// / Generated by ` + "`" + `ec2geninfo` + "`" + `
+
+package instance
+
+var InstanceTypesMap = map[string]InstanceInfo{}
+
+func init() {
+	for _, instance := range InstanceTypes {
+		InstanceTypesMap[instance.InstanceType] = instance
+	}
+}
+
+type InstanceInfo struct {
+	InstanceType             string
+	InstanceStorageSupported bool
+	EFASupported             bool
+	NvidiaGPUSupported       bool
+	NvidiaGPUType            string
+	NeuronSupported          bool
+	NeuronDeviceType         string
+	CBRSupported             bool
+	CPUArch                  string
+}
+
+var InstanceTypes = []InstanceInfo{
+{{- range . }}
+	{
+		InstanceType:             "{{ .InstanceType }}",
+		InstanceStorageSupported: {{ .InstanceStorageSupported }},
+		EFASupported:             {{ .EFASupported }},
+		NvidiaGPUSupported:       {{ .NvidiaGPUSupported }},
+		NvidiaGPUType:            "{{ .NvidiaGPUType }}",
+		NeuronSupported:          {{ .NeuronSupported }},
+		NeuronDeviceType:         "{{ .NeuronDeviceType }}",
+		CBRSupported:             {{ .CBRSupported }},
+		CPUArch:                  "{{ .CPUArch }}",
+	},
+{{- end }}
+}
+`
+
+func main() {
+	err := updateEC2Instances()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error updating EC2 instances: %v\n", err)
+		os.Exit(1)
+	}
+}
+
+func updateEC2Instances() error {
+	regions := []string{"us-east-1", "us-east-2", "us-west-2"}
+	instances := make(map[string]InstanceInfo)
+
+	for _, region := range regions {
+		var err error
+		instances, err = getEC2Instances(region, instances)
+		if err != nil {
+			return err
+		}
+	}
+
+	tmpl, err := template.New("ec2InstancesTemplate").Parse(ec2InstancesTemplate)
+	if err != nil {
+		return err
+	}
+
+	file, err := os.Create("../../pkg/utils/instance/instance_types.go")
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	err = tmpl.Execute(file, instances)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func getEC2Instances(region string, instances map[string]InstanceInfo) (map[string]InstanceInfo, error) {
+	cfg, err := config.LoadDefaultConfig(context.TODO(), config.WithRegion(region))
+	if err != nil {
+		return nil, err
+	}
+
+	client := ec2.NewFromConfig(cfg)
+
+	input := &ec2.DescribeInstanceTypesInput{
+		Filters: []types.Filter{
+			{Name: aws.String("current-generation"), Values: []string{"true"}},
+			{Name: aws.String("bare-metal"), Values: []string{"false"}},
+		},
+	}
+
+	paginator := ec2.NewDescribeInstanceTypesPaginator(client, input)
+
+	for paginator.HasMorePages() {
+		page, err := paginator.NextPage(context.TODO())
+		if err != nil {
+			return nil, err
+		}
+
+		for _, inst := range page.InstanceTypes {
+			itype := string(inst.InstanceType)
+
+			efaSupported := inst.NetworkInfo != nil && inst.NetworkInfo.EfaSupported != nil && *inst.NetworkInfo.EfaSupported
+
+			nvidiaGPUSupported := false
+			nvidiaGPUType := ""
+			if inst.GpuInfo != nil && len(inst.GpuInfo.Gpus) > 0 {
+				nvidiaGPUSupported = *inst.GpuInfo.Gpus[0].Manufacturer == "NVIDIA"
+				nvidiaGPUType = *inst.GpuInfo.Gpus[0].Name
+			}
+
+			neuronSupported := inst.NeuronInfo != nil
+			neuronDeviceType := ""
+			if neuronSupported {
+				for _, acc := range inst.NeuronInfo.NeuronDevices {
+					neuronDeviceType = *acc.Name
+					break
+				}
+			}
+
+			cbrSupported := false
+			if inst.SupportedUsageClasses != nil {
+				for _, usageClass := range inst.SupportedUsageClasses {
+					if usageClass == types.UsageClassTypeCapacityBlock {
+						cbrSupported = true
+						break
+					}
+				}
+			}
+
+			cpuArch := "unknown"
+			if inst.ProcessorInfo != nil && inst.ProcessorInfo.SupportedArchitectures != nil {
+				for _, arch := range inst.ProcessorInfo.SupportedArchitectures {
+					if arch == types.ArchitectureTypeArm64 || arch == types.ArchitectureTypeArm64Mac {
+						cpuArch = "arm64"
+					} else if arch == types.ArchitectureTypeX8664 || arch == types.ArchitectureTypeX8664Mac {
+						cpuArch = "x86-64"
+					}
+				}
+			}
+
+			instances[itype] = InstanceInfo{
+				InstanceType:             itype,
+				InstanceStorageSupported: inst.InstanceStorageSupported != nil && *inst.InstanceStorageSupported,
+				EFASupported:             efaSupported,
+				NvidiaGPUSupported:       nvidiaGPUSupported,
+				NvidiaGPUType:            nvidiaGPUType,
+				NeuronSupported:          neuronSupported,
+				NeuronDeviceType:         neuronDeviceType,
+				CBRSupported:             cbrSupported,
+				CPUArch:                  cpuArch,
+			}
+		}
+	}
+
+	return instances, nil
+}

pkg/utils/instance/instance.go

@@ -1,72 +1,52 @@
 package instance
 
 import (
-	"strings"
-
 	"github.com/aws/aws-sdk-go-v2/aws"
 	ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
 )
 
 // IsARMInstanceType returns true if the instance type is ARM architecture
 func IsARMInstanceType(instanceType string) bool {
-	return strings.HasPrefix(instanceType, "a1") ||
-		strings.HasPrefix(instanceType, "t4g") ||
-		strings.HasPrefix(instanceType, "m6g") ||
-		strings.HasPrefix(instanceType, "m7g") ||
-		strings.HasPrefix(instanceType, "m8g") ||
-		strings.HasPrefix(instanceType, "r8g") ||
-		strings.HasPrefix(instanceType, "c8g") ||
-		strings.HasPrefix(instanceType, "c6g") ||
-		strings.HasPrefix(instanceType, "c7g") ||
-		strings.HasPrefix(instanceType, "r6g") ||
-		strings.HasPrefix(instanceType, "r7g") ||
-		strings.HasPrefix(instanceType, "im4g") ||
-		strings.HasPrefix(instanceType, "is4g") ||
-		strings.HasPrefix(instanceType, "g5g") ||
-		strings.HasPrefix(instanceType, "hpc7g") ||
-		strings.HasPrefix(instanceType, "x2g")
+	if InstanceTypesMap[instanceType].CPUArch == "arm64" {
+		return true
+	}
+	return false
 }
 
 // IsGPUInstanceType returns true if the instance type is GPU optimised
 func IsGPUInstanceType(instanceType string) bool {
-	return IsNvidiaInstanceType(instanceType) ||
-		IsInferentiaInstanceType(instanceType) ||
-		IsTrainiumInstanceType(instanceType)
+	itype := InstanceTypesMap[instanceType]
+	return itype.NvidiaGPUSupported || itype.NeuronSupported
 }
 
 // IsNeuronInstanceType returns true if the instance type requires AWS Neuron
 func IsNeuronInstanceType(instanceType string) bool {
-	return IsInferentiaInstanceType(instanceType) ||
-		IsTrainiumInstanceType(instanceType)
+	return InstanceTypesMap[instanceType].NeuronSupported
 }
 
 // IsARMGPUInstanceType returns true if the instance type is ARM-GPU architecture
 func IsARMGPUInstanceType(instanceType string) bool {
-	return strings.HasPrefix(instanceType, "g5g")
+	itype := InstanceTypesMap[instanceType]
+	return itype.CPUArch == "arm64" && itype.NvidiaGPUSupported
 }
 
 // IsNvidiaInstanceType returns true if the instance type has NVIDIA accelerated hardware
 func IsNvidiaInstanceType(instanceType string) bool {
-	return strings.HasPrefix(instanceType, "p2") ||
-		strings.HasPrefix(instanceType, "p3") ||
-		strings.HasPrefix(instanceType, "p4") ||
-		strings.HasPrefix(instanceType, "p5") ||
-		strings.HasPrefix(instanceType, "g3") ||
-		strings.HasPrefix(instanceType, "g4") ||
-		strings.HasPrefix(instanceType, "g5") ||
-		strings.HasPrefix(instanceType, "g6")
+	return InstanceTypesMap[instanceType].NvidiaGPUSupported
 }
 
-// IsInferentiaInstanceType returns true if the instance type requires AWS Neuron
+// IsInferentiaInstanceType returns true if the instance type requires AWS Neuron Inferentia/Inferentia2
 func IsInferentiaInstanceType(instanceType string) bool {
-	return strings.HasPrefix(instanceType, "inf1") ||
-		strings.HasPrefix(instanceType, "inf2")
+	itype := InstanceTypesMap[instanceType]
+	return itype.NeuronSupported &&
+		(itype.NeuronDeviceType == "Inferentia" || itype.NeuronDeviceType == "Inferentia2")
 }
 
-// IsTrainiumnstanceType returns true if the instance type requires AWS Neuron
+// IsTrainiumnstanceType returns true if the instance type requires AWS Neuron Trainium/Trainium2
 func IsTrainiumInstanceType(instanceType string) bool {
-	return strings.HasPrefix(instanceType, "trn1") ||
-		strings.HasPrefix(instanceType, "trn2")
+	itype := InstanceTypesMap[instanceType]
+	return itype.NeuronSupported &&
+		(itype.NeuronDeviceType == "Trainium" || itype.NeuronDeviceType == "Trainium2")
 }
 
 // GetSmallestInstanceType returns the smallest instance type in instanceTypes.