add inference acceleartor filter support

Author: bwagner5

README.md

@@ -194,6 +194,9 @@ Filter Flags:
       --gpus-min int                                   Minimum Total Number of GPUs (Example: 4) If --gpus-max is not specified, the upper bound will be infinity
       --hibernation-support                            Hibernation supported
       --hypervisor string                              Hypervisor: [xen or nitro]
+      --inference-accelerators int                     Total Number of inference accelerators (Example: 4) (sets --inference-accelerators-min and -max to the same value)
+      --inference-accelerators-max int                 Maximum Total Number of inference accelerators (Example: 4) If --inference-accelerators-min is not specified, the lower bound will be 0
+      --inference-accelerators-min int                 Minimum Total Number of inference accelerators (Example: 4) If --inference-accelerators-max is not specified, the upper bound will be infinity
       --instance-storage string                        Amount of local instance storage (Example: 4 GiB) (sets --instance-storage-min and -max to the same value)
       --instance-storage-max string                    Maximum Amount of local instance storage (Example: 4 GiB) If --instance-storage-min is not specified, the lower bound will be 0
       --instance-storage-min string                    Minimum Amount of local instance storage (Example: 4 GiB) If --instance-storage-max is not specified, the upper bound will be infinity

cmd/main.go

@@ -60,6 +60,7 @@ const (
 	cpuArchitecture                = "cpu-architecture"
 	gpus                           = "gpus"
 	gpuMemoryTotal                 = "gpu-memory-total"
+	inferenceAccelerators          = "inference-accelerators"
 	placementGroupStrategy         = "placement-group-strategy"
 	usageClass                     = "usage-class"
 	rootDeviceType                 = "root-device-type"
@@ -147,6 +148,7 @@ Full docs can be found at github.com/aws/amazon-` + binName
 	cli.StringOptionsFlag(cpuArchitecture, cli.StringMe("a"), nil, "CPU architecture [x86_64/amd64, x86_64_mac, i386, or arm64]", []string{"x86_64", "x86_64_mac", "amd64", "i386", "arm64"})
 	cli.IntMinMaxRangeFlags(gpus, cli.StringMe("g"), nil, "Total Number of GPUs (Example: 4)")
 	cli.ByteQuantityMinMaxRangeFlags(gpuMemoryTotal, nil, nil, "Number of GPUs' total memory (Example: 4 GiB)")
+	cli.IntMinMaxRangeFlags(inferenceAccelerators, nil, nil, "Total Number of inference accelerators (Example: 4)")
 	cli.StringOptionsFlag(placementGroupStrategy, nil, nil, "Placement group strategy: [cluster, partition, spread]", []string{"cluster", "partition", "spread"})
 	cli.StringOptionsFlag(usageClass, cli.StringMe("u"), nil, "Usage class: [spot or on-demand]", []string{"spot", "on-demand"})
 	cli.StringOptionsFlag(rootDeviceType, nil, nil, "Supported root device types: [ebs or instance-store]", []string{"ebs", "instance-store"})
@@ -256,6 +258,7 @@ Full docs can be found at github.com/aws/amazon-` + binName
 		CPUArchitecture:                cli.StringMe(flags[cpuArchitecture]),
 		GpusRange:                      cli.IntRangeMe(flags[gpus]),
 		GpuMemoryRange:                 cli.ByteQuantityRangeMe(flags[gpuMemoryTotal]),
+		InferenceAcceleratorsRange:     cli.IntRangeMe(flags[inferenceAccelerators]),
 		PlacementGroupStrategy:         cli.StringMe(flags[placementGroupStrategy]),
 		UsageClass:                     cli.StringMe(flags[usageClass]),
 		RootDeviceType:                 cli.StringMe(flags[rootDeviceType]),

pkg/selector/comparators.go

@@ -111,6 +111,17 @@ func isSupportedWithBool(instanceTypeValue *bool, target *bool) bool {
 
 // Helper functions for aggregating data parsed from AWS API calls
 
+func getTotalAcceleratorsCount(acceleratorInfo *ec2.InferenceAcceleratorInfo) *int64 {
+	if acceleratorInfo == nil {
+		return nil
+	}
+	total := aws.Int64(0)
+	for _, accel := range acceleratorInfo.Accelerators {
+		total = aws.Int64(*total + *accel.Count)
+	}
+	return total
+}
+
 func getTotalGpusCount(gpusInfo *ec2.GpuInfo) *int64 {
 	if gpusInfo == nil {
 		return nil

pkg/selector/selector.go

@@ -55,6 +55,7 @@ const (
 	memoryRange                    = "memoryRange"
 	gpuMemoryRange                 = "gpuMemoryRange"
 	gpusRange                      = "gpusRange"
+	inferenceAcceleratorsRange     = "inferenceAcceleratorsRange"
 	placementGroupStrategy         = "placementGroupStrategy"
 	hypervisor                     = "hypervisor"
 	baremetal                      = "baremetal"
@@ -255,6 +256,7 @@ func (itf Selector) rawFilter(filters Filters) ([]*instancetypes.Details, error)
 			memoryRange:                    {filters.MemoryRange, instanceTypeInfo.MemoryInfo.SizeInMiB},
 			gpuMemoryRange:                 {filters.GpuMemoryRange, getTotalGpuMemory(instanceTypeInfo.GpuInfo)},
 			gpusRange:                      {filters.GpusRange, getTotalGpusCount(instanceTypeInfo.GpuInfo)},
+			inferenceAcceleratorsRange:     {filters.InferenceAcceleratorsRange, getTotalAcceleratorsCount(instanceTypeInfo.InferenceAcceleratorInfo)},
 			placementGroupStrategy:         {filters.PlacementGroupStrategy, instanceTypeInfo.PlacementGroupInfo.SupportedStrategies},
 			hypervisor:                     {filters.Hypervisor, instanceTypeInfo.Hypervisor},
 			baremetal:                      {filters.BareMetal, instanceTypeInfo.BareMetal},

pkg/selector/types.go

@@ -137,6 +137,9 @@ type Filters struct {
 	// GpuMemoryRange filter is a range of acceptable GPU memory in Gibibytes (GiB) available to an EC2 instance type in aggreagte across all GPUs.
 	GpuMemoryRange *ByteQuantityRangeFilter
 
+	// InferenceAcceleratorsRange filters inference acclerators available to the instance type
+	InferenceAcceleratorsRange *IntRangeFilter
+
 	// HibernationSupported denotes whether EC2 hibernate is supported
 	// Possible values are: true or false
 	HibernationSupported *bool