From aeb802448bd25e04b478f19ca48b23f5ad0909ed Mon Sep 17 00:00:00 2001 From: Julian Sommer <115516918+juliansommer@users.noreply.github.com> Date: Tue, 17 Feb 2026 10:39:59 +1000 Subject: [PATCH] feat: more instance types --- hive/daemon/daemon.go | 281 +++++++++++++++++++++++--------- hive/daemon/daemon_test.go | 326 ++++++++++++++++++++++--------------- 2 files changed, 403 insertions(+), 204 deletions(-) diff --git a/hive/daemon/daemon.go b/hive/daemon/daemon.go index c1d9a153d..4090829e1 100644 --- a/hive/daemon/daemon.go +++ b/hive/daemon/daemon.go @@ -16,6 +16,7 @@ import ( "path/filepath" "regexp" "runtime" + "slices" "strconv" "strings" "sync" @@ -125,28 +126,33 @@ type Daemon struct { mu sync.Mutex } -// cpuToInstanceFamily maps CPU model patterns to AWS instance family prefixes -var cpuToInstanceFamily = map[string]string{ - "EPYC": "m8a", // AMD EPYC processors - "Xeon": "m7i", // Intel Xeon processors - "ARM": "m8g", // ARM-based processors - "Apple": "m8g", // Apple Silicon (ARM-based) - "Graviton": "m8g", // AWS Graviton +// cpuToVendor maps CPU model patterns to CPU vendors. +// Patterns are checked in order; more specific patterns should come first. +var cpuToVendor = []struct { + pattern string + vendor cpuVendor +}{ + {"EPYC", vendorAMD}, + {"Xeon", vendorIntel}, + {"Graviton", vendorARM}, + {"Apple", vendorARM}, + {"ARM", vendorARM}, + {"AMD", vendorAMD}, // catches Ryzen, Threadripper, etc. } -// getInstanceFamilyFromCPU returns the AWS instance family based on CPU model -func getInstanceFamilyFromCPU(cpuModel string) string { +// vendorFromCPU returns the CPU vendor based on the CPU model string and architecture. +func vendorFromCPU(cpuModel, arch string) cpuVendor { + if arch == "arm64" { + return vendorARM + } cpuUpper := strings.ToUpper(cpuModel) - for pattern, family := range cpuToInstanceFamily { - if strings.Contains(cpuUpper, strings.ToUpper(pattern)) { - return family + for _, entry := range cpuToVendor { + if strings.Contains(cpuUpper, strings.ToUpper(entry.pattern)) { + return entry.vendor } } - // Default fallback based on architecture - if runtime.GOARCH == "arm64" { - return "t4g" - } - return "t3" // fallback for unknown x86_64 + slog.Warn("Unrecognized CPU model, defaulting to Intel vendor", "cpuModel", cpuModel) + return vendorIntel } // getSystemMemory returns the total system memory in GB @@ -229,49 +235,175 @@ func getCPUModel() (string, error) { } } -// generateInstanceTypes creates the instance type map for the detected CPU family -// It includes both the CPU-specific family and the burstable family (t3 for x86, t4g for ARM) -func generateInstanceTypes(family, arch string) map[string]*ec2.InstanceTypeInfo { - sizes := []struct { - suffix string - vcpus int - memoryGB float64 - }{ - {"nano", 2, 0.5}, - {"micro", 2, 1.0}, - {"small", 2, 2.0}, - {"medium", 2, 4.0}, - {"large", 2, 8.0}, - {"xlarge", 4, 16.0}, - {"2xlarge", 8, 32.0}, - } +type instanceSize struct { + suffix string + vcpus int + memoryGB float64 +} - // Determine the burstable family based on architecture and CPU vendor - burstableFamily := "t3" - if arch == "arm64" { - burstableFamily = "t4g" - } else if strings.HasSuffix(family, "a") { - burstableFamily = "t3a" - } +type cpuVendor int - // Build list of families to generate: CPU-specific + burstable (if different) - families := []struct { - name string - burstable bool - }{ - {family, false}, - } - if family != burstableFamily { - families = append(families, struct { - name string - burstable bool - }{burstableFamily, true}) - } +const ( + vendorIntel cpuVendor = iota + vendorAMD + vendorARM +) +func (v cpuVendor) String() string { + switch v { + case vendorIntel: + return "Intel" + case vendorAMD: + return "AMD" + case vendorARM: + return "ARM" + default: + return "Unknown" + } +} + +type instanceFamilyDef struct { + name string + vendor cpuVendor + sizes []instanceSize + currentGen bool +} + +// Size tables for each instance category + +var burstableSizes = []instanceSize{ + {"nano", 2, 0.5}, + {"micro", 2, 1}, + {"small", 2, 2}, + {"medium", 2, 4}, + {"large", 2, 8}, + {"xlarge", 4, 16}, + {"2xlarge", 8, 32}, +} + +var gpSizes = []instanceSize{ + {"large", 2, 8}, + {"xlarge", 4, 16}, + {"2xlarge", 8, 32}, + {"4xlarge", 16, 64}, + {"8xlarge", 32, 128}, + {"12xlarge", 48, 192}, + {"16xlarge", 64, 256}, + {"24xlarge", 96, 384}, +} + +// gpSizesSmall is gpSizes without 12xlarge and 24xlarge (older/ARM families). +var gpSizesSmall = slices.Clone(gpSizes[:6]) + +var computeSizes = []instanceSize{ + {"large", 2, 4}, + {"xlarge", 4, 8}, + {"2xlarge", 8, 16}, + {"4xlarge", 16, 32}, + {"8xlarge", 32, 64}, + {"12xlarge", 48, 96}, + {"16xlarge", 64, 128}, + {"24xlarge", 96, 192}, +} + +// computeSizesSmall is computeSizes without 12xlarge and 24xlarge (older/ARM families). +var computeSizesSmall = slices.Clone(computeSizes[:6]) + +var memorySizes = []instanceSize{ + {"large", 2, 16}, + {"xlarge", 4, 32}, + {"2xlarge", 8, 64}, + {"4xlarge", 16, 128}, + {"8xlarge", 32, 256}, + {"12xlarge", 48, 384}, + {"16xlarge", 64, 512}, + {"24xlarge", 96, 768}, +} + +// memorySizesSmall is memorySizes without 12xlarge and 24xlarge (older/ARM families). +var memorySizesSmall = slices.Clone(memorySizes[:6]) + +// instanceFamilyDefs defines all supported instance families with their vendor and sizes. +// +// We support the core families across burstable, general purpose, compute optimized, +// and memory optimized categories. The following AWS family categories are intentionally +// excluded because they require specialized hardware not available on standard bare-metal hosts: +// +// - Local disk variants (d/n suffixes): c5d, c5ad, c5n, m5d, m5ad, m5n, m5dn, m5zn, r5d, r5ad, +// r5n, r5dn, r5b, c6gd, c6gn, c6id, c6in, m6gd, m6id, m6idn, m6in, r6gd, r6id, r6idn, r6in, +// c7gd, c7gn, c7i-flex, m7gd, m7i-flex, r7gd, r7iz, c8gd, c8gn, c8i-flex, m8gd, m8i-flex, +// r8gd, r8gn, r8gb, r8i-flex — require NVMe instance storage or enhanced networking +// - GPU/accelerator: g2-g6, g6e, g6f, gr6, gr6f, p2-p6, inf1-inf2, trn1-trn2, dl1, dl2q — require +// GPU, Inferentia, Trainium, or other accelerator hardware +// - Storage optimized: d2, d3, d3en, h1, i2-i8g, i7ie, i8ge, im4gn, is4gen — require dense HDD/NVMe +// - FPGA: f1, f2 — require FPGA hardware +// - High memory: u-*, u7i-*, x1, x1e, x2gd, x2idn, x2iedn, x2iezn, x8g — require TB-scale memory +// - High frequency: z1d — specialized high clock-speed instances +// - Dedicated host: mac*, hpc* — require macOS/Apple hardware or HPC interconnects +// - Video: vt1 — requires video transcoding hardware +// - Legacy (pre-gen4): a1, c1, c3, cc1, cc2, cg1, cr1, hi1, hs1, m1, m2, m3, r3, t1 +var instanceFamilyDefs = []instanceFamilyDef{ + // Burstable + {name: "t2", vendor: vendorIntel, sizes: burstableSizes, currentGen: false}, + {name: "t3", vendor: vendorIntel, sizes: burstableSizes, currentGen: true}, + {name: "t3a", vendor: vendorAMD, sizes: burstableSizes, currentGen: true}, + {name: "t4g", vendor: vendorARM, sizes: burstableSizes, currentGen: true}, + + // General Purpose (1:4 vCPU:memory) + {name: "m4", vendor: vendorIntel, sizes: gpSizesSmall, currentGen: false}, + {name: "m5", vendor: vendorIntel, sizes: gpSizes, currentGen: true}, + {name: "m5a", vendor: vendorAMD, sizes: gpSizes, currentGen: true}, + {name: "m6i", vendor: vendorIntel, sizes: gpSizes, currentGen: true}, + {name: "m6a", vendor: vendorAMD, sizes: gpSizes, currentGen: true}, + {name: "m6g", vendor: vendorARM, sizes: gpSizesSmall, currentGen: true}, + {name: "m7i", vendor: vendorIntel, sizes: gpSizes, currentGen: true}, + {name: "m7a", vendor: vendorAMD, sizes: gpSizes, currentGen: true}, + {name: "m7g", vendor: vendorARM, sizes: gpSizesSmall, currentGen: true}, + {name: "m8i", vendor: vendorIntel, sizes: gpSizes, currentGen: true}, + {name: "m8a", vendor: vendorAMD, sizes: gpSizes, currentGen: true}, + {name: "m8g", vendor: vendorARM, sizes: gpSizesSmall, currentGen: true}, + + // Compute Optimized (1:2 vCPU:memory) + {name: "c4", vendor: vendorIntel, sizes: computeSizesSmall, currentGen: false}, + {name: "c5", vendor: vendorIntel, sizes: computeSizes, currentGen: true}, + {name: "c5a", vendor: vendorAMD, sizes: computeSizes, currentGen: true}, + {name: "c6i", vendor: vendorIntel, sizes: computeSizes, currentGen: true}, + {name: "c6a", vendor: vendorAMD, sizes: computeSizes, currentGen: true}, + {name: "c6g", vendor: vendorARM, sizes: computeSizesSmall, currentGen: true}, + {name: "c7i", vendor: vendorIntel, sizes: computeSizes, currentGen: true}, + {name: "c7a", vendor: vendorAMD, sizes: computeSizes, currentGen: true}, + {name: "c7g", vendor: vendorARM, sizes: computeSizesSmall, currentGen: true}, + {name: "c8i", vendor: vendorIntel, sizes: computeSizes, currentGen: true}, + {name: "c8a", vendor: vendorAMD, sizes: computeSizes, currentGen: true}, + {name: "c8g", vendor: vendorARM, sizes: computeSizesSmall, currentGen: true}, + + // Memory Optimized (1:8 vCPU:memory) + {name: "r4", vendor: vendorIntel, sizes: memorySizesSmall, currentGen: false}, + {name: "r5", vendor: vendorIntel, sizes: memorySizes, currentGen: true}, + {name: "r5a", vendor: vendorAMD, sizes: memorySizes, currentGen: true}, + {name: "r6i", vendor: vendorIntel, sizes: memorySizes, currentGen: true}, + {name: "r6a", vendor: vendorAMD, sizes: memorySizes, currentGen: true}, + {name: "r6g", vendor: vendorARM, sizes: memorySizesSmall, currentGen: true}, + {name: "r7i", vendor: vendorIntel, sizes: memorySizes, currentGen: true}, + {name: "r7a", vendor: vendorAMD, sizes: memorySizes, currentGen: true}, + {name: "r7g", vendor: vendorARM, sizes: memorySizesSmall, currentGen: true}, + {name: "r8i", vendor: vendorIntel, sizes: memorySizes, currentGen: true}, + {name: "r8a", vendor: vendorAMD, sizes: memorySizes, currentGen: true}, + {name: "r8g", vendor: vendorARM, sizes: memorySizesSmall, currentGen: true}, +} + +// generateInstanceTypes creates the instance type map for the given CPU vendor. +// It generates all instance families matching the vendor across burstable, +// general purpose, compute optimized, and memory optimized categories. +func generateInstanceTypes(v cpuVendor, arch string) map[string]*ec2.InstanceTypeInfo { instanceTypes := make(map[string]*ec2.InstanceTypeInfo) - for _, fam := range families { - for _, size := range sizes { - name := fmt.Sprintf("%s.%s", fam.name, size.suffix) + for _, def := range instanceFamilyDefs { + if def.vendor != v { + continue + } + burstable := strings.HasPrefix(def.name, "t") + for _, size := range def.sizes { + name := fmt.Sprintf("%s.%s", def.name, size.suffix) instanceTypes[name] = &ec2.InstanceTypeInfo{ InstanceType: aws.String(name), VCpuInfo: &ec2.VCpuInfo{ @@ -283,12 +415,11 @@ func generateInstanceTypes(family, arch string) map[string]*ec2.InstanceTypeInfo ProcessorInfo: &ec2.ProcessorInfo{ SupportedArchitectures: []*string{aws.String(arch)}, }, - CurrentGeneration: aws.Bool(true), - BurstablePerformanceSupported: aws.Bool(false), - // BurstablePerformanceSupported: aws.Bool(fam.burstable), - Hypervisor: aws.String("kvm"), - SupportedVirtualizationTypes: []*string{aws.String("hvm")}, - SupportedRootDeviceTypes: []*string{aws.String("ebs")}, + CurrentGeneration: aws.Bool(def.currentGen), + BurstablePerformanceSupported: aws.Bool(burstable), + Hypervisor: aws.String("kvm"), + SupportedVirtualizationTypes: []*string{aws.String("hvm")}, + SupportedRootDeviceTypes: []*string{aws.String("ebs")}, } } } @@ -314,28 +445,24 @@ func NewResourceManager() *ResourceManager { cpuModel = "Unknown" } - // Determine instance family from CPU model - instanceFamily := getInstanceFamilyFromCPU(cpuModel) - // Determine architecture arch := "x86_64" if runtime.GOARCH == "arm64" { arch = "arm64" } - // Generate instance types based on detected CPU family - instanceTypes := generateInstanceTypes(instanceFamily, arch) + // Detect CPU vendor and generate instance types + vendor := vendorFromCPU(cpuModel, arch) + instanceTypes := generateInstanceTypes(vendor, arch) - // Determine burstable family for logging - burstableFamily := "t3" - if runtime.GOARCH == "arm64" { - burstableFamily = "t4g" - } else if strings.HasSuffix(instanceFamily, "a") { - burstableFamily = "t3a" + if len(instanceTypes) == 0 { + slog.Error("No instance types generated, daemon will be unable to run VMs", + "vendor", vendor, "arch", arch, "cpuModel", cpuModel) } + slog.Info("System resources detected", "vCPUs", numCPU, "memGB", totalMemGB, "cpu", cpuModel, - "family", instanceFamily, "burstableFamily", burstableFamily, "os", runtime.GOOS) + "vendor", vendor, "instanceTypes", len(instanceTypes), "os", runtime.GOOS) return &ResourceManager{ availableVCPU: numCPU, @@ -526,6 +653,7 @@ func (d *Daemon) subscribeAll() error { {"ec2.start", d.handleEC2StartStoppedInstance, "hive-workers"}, {"ec2.terminate", d.handleEC2TerminateStoppedInstance, "hive-workers"}, {"ec2.DescribeStoppedInstances", d.handleEC2DescribeStoppedInstances, "hive-workers"}, + // these 2 fan out to all nodes and gateway aggregates the results {"ec2.DescribeInstances", d.handleEC2DescribeInstances, ""}, {"ec2.DescribeInstanceTypes", d.handleEC2DescribeInstanceTypes, ""}, {"ec2.EnableEbsEncryptionByDefault", d.handleEC2EnableEbsEncryptionByDefault, "hive-workers"}, @@ -829,7 +957,10 @@ func (d *Daemon) restoreInstances() { } instanceType, ok := d.resourceMgr.instanceTypes[instance.InstanceType] - if ok { + if !ok { + slog.Error("Instance type not found in current type map, cannot re-allocate resources", + "instanceId", instance.ID, "instanceType", instance.InstanceType) + } else { slog.Info("Re-allocating resources for instance", "instanceId", instance.ID, "type", instance.InstanceType) if err := d.resourceMgr.allocate(instanceType); err != nil { slog.Error("Failed to re-allocate resources for instance on startup", "instanceId", instance.ID, "err", err) diff --git a/hive/daemon/daemon_test.go b/hive/daemon/daemon_test.go index 2a5d81118..d6917f917 100644 --- a/hive/daemon/daemon_test.go +++ b/hive/daemon/daemon_test.go @@ -450,10 +450,9 @@ func TestGetInstanceTypeInfos(t *testing.T) { infos := rm.GetInstanceTypeInfos() require.NotEmpty(t, infos, "Should return at least one instance type") - // Number of instance types depends on whether CPU family differs from burstable family - // 7 sizes per family, 1-2 families depending on detected CPU - assert.True(t, len(infos) == 7 || len(infos) == 14, - "Should have 7 or 14 instance types (7 sizes x 1-2 families), got %d", len(infos)) + // With comprehensive families, expect 49+ types (ARM) up to 128 (Intel) + assert.True(t, len(infos) >= 49, + "Should have at least 49 instance types, got %d", len(infos)) // Verify structure of returned instance type info for _, info := range infos { @@ -465,7 +464,6 @@ func TestGetInstanceTypeInfos(t *testing.T) { assert.NotNil(t, info.ProcessorInfo, "ProcessorInfo should not be nil") assert.NotEmpty(t, info.ProcessorInfo.SupportedArchitectures, "SupportedArchitectures should not be empty") assert.NotNil(t, info.CurrentGeneration, "CurrentGeneration should not be nil") - assert.True(t, *info.CurrentGeneration, "CurrentGeneration should be true") t.Logf("Instance type: %s, vCPUs: %d, Memory: %d MiB", *info.InstanceType, *info.VCpuInfo.DefaultVCpus, *info.MemoryInfo.SizeInMiB) @@ -480,22 +478,25 @@ func TestCPUDetection(t *testing.T) { t.Logf("Detected CPU model: %s", cpuModel) - // Test instance family mapping + // Test vendor mapping with correct architectures testCases := []struct { cpuModel string + arch string + expected cpuVendor }{ - {"AMD EPYC 7551"}, - {"Intel Xeon E5-2686"}, - {"Apple M1"}, - {"AWS Graviton2"}, - {"Unknown CPU"}, + {"AMD EPYC 7551", "x86_64", vendorAMD}, + {"AMD Ryzen 9 7950X", "x86_64", vendorAMD}, + {"Intel Xeon E5-2686", "x86_64", vendorIntel}, + {"Apple M1", "arm64", vendorARM}, + {"AWS Graviton2", "arm64", vendorARM}, + {"Unknown CPU", "x86_64", vendorIntel}, } for _, tt := range testCases { t.Run(tt.cpuModel, func(t *testing.T) { - family := getInstanceFamilyFromCPU(tt.cpuModel) - assert.NotEmpty(t, family, "Should return a family for %s", tt.cpuModel) - t.Logf("CPU: %s -> Family: %s", tt.cpuModel, family) + v := vendorFromCPU(tt.cpuModel, tt.arch) + assert.Equal(t, tt.expected, v, "Wrong vendor for %s on %s", tt.cpuModel, tt.arch) + t.Logf("CPU: %s (%s) -> Vendor: %s", tt.cpuModel, tt.arch, v) }) } } @@ -787,34 +788,12 @@ func TestHandleEC2DescribeInstanceTypes(t *testing.T) { // Test 4: Verify "capacity" filter returns duplicates t.Run("VerifyCapacityFilter_Duplicates", func(t *testing.T) { - // Determine instance family dynamically from ResourceManager - var family string - for name := range daemon.resourceMgr.instanceTypes { - parts := strings.Split(name, ".") - if len(parts) > 0 { - family = parts[0] - break - } - } - require.NotEmpty(t, family, "Should have a detected instance family") - - // Count unique families in resourceMgr - familySet := make(map[string]bool) - for name := range daemon.resourceMgr.instanceTypes { - parts := strings.Split(name, ".") - if len(parts) > 0 { - familySet[parts[0]] = true - } - } - numFamilies := len(familySet) - t.Logf("Detected %d unique instance families", numFamilies) - - // Force resources to a predictable state for the first part of the test (2 vCPUs) + // Force resources to a predictable state daemon.resourceMgr.mu.Lock() oldAvailableVCPU := daemon.resourceMgr.availableVCPU oldAvailableMem := daemon.resourceMgr.availableMem daemon.resourceMgr.availableVCPU = 2 - daemon.resourceMgr.availableMem = 16.0 // Plenty of memory + daemon.resourceMgr.availableMem = 16.0 daemon.resourceMgr.mu.Unlock() defer func() { @@ -824,7 +803,6 @@ func TestHandleEC2DescribeInstanceTypes(t *testing.T) { daemon.resourceMgr.mu.Unlock() }() - // Create input with capacity=true filter, no specific instance types input := &ec2.DescribeInstanceTypesInput{ Filters: []*ec2.Filter{ { @@ -842,15 +820,23 @@ func TestHandleEC2DescribeInstanceTypes(t *testing.T) { err = json.Unmarshal(reply.Data, &output) require.NoError(t, err) - // With 2 vCPUs, we expect 1 slot for each 2-vCPU type (nano, micro, small, medium, large) - // and 0 slots for 4+ vCPU types. 5 sizes per family. - expectedSlots := 5 * numFamilies - assert.Equal(t, expectedSlots, len(output.InstanceTypes), "Should have %d total slots available with 2 vCPUs (5 sizes x %d families)", expectedSlots, numFamilies) + // With 2 vCPUs and 16GB, every type with 2 vCPUs and <=16GB memory should have 1 slot. + // Calculate expected by counting fitting types directly. + expectedSlots := 0 + for _, it := range daemon.resourceMgr.instanceTypes { + vcpus := *it.VCpuInfo.DefaultVCpus + memGB := float64(*it.MemoryInfo.SizeInMiB) / 1024.0 + if vcpus <= 2 && memGB <= 16.0 { + expectedSlots++ + } + } + assert.Equal(t, expectedSlots, len(output.InstanceTypes), + "Should have %d slots for types fitting 2 vCPU / 16GB", expectedSlots) - // Now let's "fake" more capacity to test multiple slots (duplicates) per type + // Now increase capacity to test duplicate slots daemon.resourceMgr.mu.Lock() - daemon.resourceMgr.availableVCPU = 4 // Give us 4 cores - daemon.resourceMgr.availableMem = 15.0 // Limited memory to trigger large/xlarge filtering + daemon.resourceMgr.availableVCPU = 4 + daemon.resourceMgr.availableMem = 15.0 daemon.resourceMgr.mu.Unlock() reply, err = daemon.natsConn.Request("ec2.DescribeInstanceTypes", msgData, 5*time.Second) @@ -858,27 +844,23 @@ func TestHandleEC2DescribeInstanceTypes(t *testing.T) { err = json.Unmarshal(reply.Data, &output) require.NoError(t, err) - // With 4 cores and 15GB memory, per family: - // nano (2 cores, 0.5GB): 2 slots - // micro (2 cores, 1.0GB): 2 slots - // small (2 cores, 2.0GB): 2 slots - // medium (2 cores, 4.0GB): 2 slots - // large (2 cores, 8.0GB): 1 slot (15GB < 16GB) - // xlarge (4 cores, 16.0GB): 0 slots (15GB < 16GB) - // 2xlarge (8 cores, 32.0GB): 0 slots - // Total per family: 2+2+2+2+1 = 9 slots - expectedSlots = 9 * numFamilies - assert.Equal(t, expectedSlots, len(output.InstanceTypes), "Should have %d total slots with 4 cores and 15GB memory (9 slots x %d families)", expectedSlots, numFamilies) - - // Verify duplicates exist + // Verify duplicate slots exist — find a nano type and confirm it has 2 slots typeCounts := make(map[string]int) for _, info := range output.InstanceTypes { if info.InstanceType != nil { typeCounts[*info.InstanceType]++ } } - nanoType := fmt.Sprintf("%s.nano", family) - assert.Equal(t, 2, typeCounts[nanoType], "Should have 2 slots for %s", nanoType) + // Find any nano type in the generated types + var nanoType string + for name := range daemon.resourceMgr.instanceTypes { + if strings.HasSuffix(name, ".nano") { + nanoType = name + break + } + } + require.NotEmpty(t, nanoType, "Should have at least one nano type") + assert.Equal(t, 2, typeCounts[nanoType], "Should have 2 slots for %s with 4 vCPUs", nanoType) }) } @@ -3230,78 +3212,79 @@ func TestSaveClusterConfig_ErrorOnInvalidPath(t *testing.T) { // --- generateInstanceTypes --- -func TestGenerateInstanceTypes_DualFamily(t *testing.T) { - types := generateInstanceTypes("m7i", "x86_64") - assert.Len(t, types, 14) // 7 m7i + 7 t3 - - m7iCount, t3Count := 0, 0 +func hasFamily(types map[string]*ec2.InstanceTypeInfo, prefix string) bool { for name := range types { - if strings.HasPrefix(name, "m7i.") { - m7iCount++ - } - if strings.HasPrefix(name, "t3.") { - t3Count++ + if strings.HasPrefix(name, prefix) { + return true } } - assert.Equal(t, 7, m7iCount) - assert.Equal(t, 7, t3Count) + return false } -func TestGenerateInstanceTypes_SingleFamily(t *testing.T) { - // When family matches burstable, should not duplicate - types := generateInstanceTypes("t3", "x86_64") - assert.Len(t, types, 7) +func countFamily(types map[string]*ec2.InstanceTypeInfo, prefix string) int { + count := 0 for name := range types { - assert.True(t, strings.HasPrefix(name, "t3.")) + if strings.HasPrefix(name, prefix) { + count++ + } + } + return count +} + +func TestGenerateInstanceTypes_Intel(t *testing.T) { + types := generateInstanceTypes(vendorIntel, "x86_64") + // Intel: t2(7)+t3(7) + m4(6)+m5(8)+m6i(8)+m7i(8)+m8i(8) + c4(6)+c5(8)+c6i(8)+c7i(8)+c8i(8) + r4(6)+r5(8)+r6i(8)+r7i(8)+r8i(8) = 128 + assert.Len(t, types, 128) + + for _, prefix := range []string{"t2.", "t3.", "m7i.", "m8i.", "c5.", "c7i.", "r5.", "r7i."} { + assert.True(t, hasFamily(types, prefix), "expected Intel types to include %s family", prefix) + } + + // Verify AMD/ARM families are NOT present + for name := range types { + assert.False(t, strings.HasPrefix(name, "t3a."), "Intel should not have t3a: %s", name) + assert.False(t, strings.HasPrefix(name, "t4g."), "Intel should not have t4g: %s", name) + assert.False(t, strings.HasPrefix(name, "m7a."), "Intel should not have m7a: %s", name) + assert.False(t, strings.HasPrefix(name, "m7g."), "Intel should not have m7g: %s", name) } } -func TestGenerateInstanceTypes_AMDBurstable(t *testing.T) { - types := generateInstanceTypes("m8a", "x86_64") - assert.Len(t, types, 14) // 7 m8a + 7 t3a +func TestGenerateInstanceTypes_AMD(t *testing.T) { + types := generateInstanceTypes(vendorAMD, "x86_64") + // AMD: t3a(7) + m5a(8)+m6a(8)+m7a(8)+m8a(8) + c5a(8)+c6a(8)+c7a(8)+c8a(8) + r5a(8)+r6a(8)+r7a(8)+r8a(8) = 103 + assert.Len(t, types, 103) - m8aCount, t3aCount := 0, 0 - for name := range types { - if strings.HasPrefix(name, "m8a.") { - m8aCount++ - } - if strings.HasPrefix(name, "t3a.") { - t3aCount++ - } + for _, prefix := range []string{"t3a.", "m5a.", "m8a.", "c5a.", "c8a.", "r5a.", "r8a."} { + assert.True(t, hasFamily(types, prefix), "expected AMD types to include %s family", prefix) } - assert.Equal(t, 7, m8aCount) - assert.Equal(t, 7, t3aCount) - // Verify t3 (Intel burstable) is NOT present + // Verify Intel/ARM families are NOT present for name := range types { - assert.False(t, strings.HasPrefix(name, "t3."), "AMD should get t3a, not t3: %s", name) + assert.False(t, strings.HasPrefix(name, "t3."), "AMD should not have t3: %s", name) + assert.False(t, strings.HasPrefix(name, "t4g."), "AMD should not have t4g: %s", name) + assert.False(t, strings.HasPrefix(name, "m7i."), "AMD should not have m7i: %s", name) } } -func TestGenerateInstanceTypes_ARM64(t *testing.T) { - types := generateInstanceTypes("m8g", "arm64") - // m8g is not t4g, so we should get both families - assert.Len(t, types, 14) +func TestGenerateInstanceTypes_ARM(t *testing.T) { + types := generateInstanceTypes(vendorARM, "arm64") + // ARM: t4g(7) + m6g(6)+m7g(6)+m8g(6) + c6g(6)+c7g(6)+c8g(6) + r6g(6)+r7g(6)+r8g(6) = 61 + assert.Len(t, types, 61) - t4gCount := 0 - for name := range types { - if strings.HasPrefix(name, "t4g.") { - t4gCount++ - } + for _, prefix := range []string{"t4g.", "m7g.", "m8g.", "c7g.", "c8g.", "r7g.", "r8g."} { + assert.True(t, hasFamily(types, prefix), "expected ARM types to include %s family", prefix) } - assert.Equal(t, 7, t4gCount) -} -func TestGenerateInstanceTypes_ARM64_SingleFamily(t *testing.T) { - types := generateInstanceTypes("t4g", "arm64") - assert.Len(t, types, 7) + // Verify Intel/AMD families are NOT present for name := range types { - assert.True(t, strings.HasPrefix(name, "t4g.")) + assert.False(t, strings.HasPrefix(name, "t3."), "ARM should not have t3: %s", name) + assert.False(t, strings.HasPrefix(name, "t3a."), "ARM should not have t3a: %s", name) + assert.False(t, strings.HasPrefix(name, "m7i."), "ARM should not have m7i: %s", name) } } -func TestGenerateInstanceTypes_VerifySizes(t *testing.T) { - types := generateInstanceTypes("t3", "x86_64") +func TestGenerateInstanceTypes_VerifyBurstableSizes(t *testing.T) { + types := generateInstanceTypes(vendorIntel, "x86_64") expected := map[string]struct { vcpus int64 @@ -3324,27 +3307,112 @@ func TestGenerateInstanceTypes_VerifySizes(t *testing.T) { } } -// --- getInstanceFamilyFromCPU --- +func TestGenerateInstanceTypes_ComputeRatio(t *testing.T) { + types := generateInstanceTypes(vendorIntel, "x86_64") -func TestGetInstanceFamilyFromCPU_ExactMappings(t *testing.T) { - assert.Equal(t, "m8a", getInstanceFamilyFromCPU("AMD EPYC 7551")) - assert.Equal(t, "m7i", getInstanceFamilyFromCPU("Intel Xeon E5-2686")) - assert.Equal(t, "m8g", getInstanceFamilyFromCPU("Apple M1")) - assert.Equal(t, "m8g", getInstanceFamilyFromCPU("AWS Graviton2")) - assert.Equal(t, "m8g", getInstanceFamilyFromCPU("ARM Cortex-A72")) -} - -func TestGetInstanceFamilyFromCPU_CaseInsensitive(t *testing.T) { - assert.Equal(t, "m8a", getInstanceFamilyFromCPU("amd epyc 9654")) -} - -func TestGetInstanceFamilyFromCPU_DefaultFallback(t *testing.T) { - family := getInstanceFamilyFromCPU("Unknown CPU Model XYZ") - if runtime.GOARCH == "arm64" { - assert.Equal(t, "t4g", family) - } else { - assert.Equal(t, "t3", family) + // Compute optimized: 1:2 vCPU:memory ratio + expected := map[string]struct { + vcpus int64 + memMB int64 + }{ + "c5.large": {2, 4096}, + "c5.xlarge": {4, 8192}, + "c5.2xlarge": {8, 16384}, + "c7i.4xlarge": {16, 32768}, } + + for name, exp := range expected { + it, ok := types[name] + require.True(t, ok, "missing instance type %s", name) + assert.Equal(t, exp.vcpus, *it.VCpuInfo.DefaultVCpus, "%s vCPUs", name) + assert.Equal(t, exp.memMB, *it.MemoryInfo.SizeInMiB, "%s memory", name) + } +} + +func TestGenerateInstanceTypes_MemoryRatio(t *testing.T) { + types := generateInstanceTypes(vendorIntel, "x86_64") + + // Memory optimized: 1:8 vCPU:memory ratio + expected := map[string]struct { + vcpus int64 + memMB int64 + }{ + "r5.large": {2, 16384}, + "r5.xlarge": {4, 32768}, + "r5.2xlarge": {8, 65536}, + "r7i.4xlarge": {16, 131072}, + } + + for name, exp := range expected { + it, ok := types[name] + require.True(t, ok, "missing instance type %s", name) + assert.Equal(t, exp.vcpus, *it.VCpuInfo.DefaultVCpus, "%s vCPUs", name) + assert.Equal(t, exp.memMB, *it.MemoryInfo.SizeInMiB, "%s memory", name) + } +} + +func TestGenerateInstanceTypes_NoSmallSizesForNonBurstable(t *testing.T) { + types := generateInstanceTypes(vendorIntel, "x86_64") + + // Non-burstable families should not have nano/micro/small/medium sizes + for name := range types { + if strings.HasPrefix(name, "t") { + continue // skip all burstable families + } + for _, small := range []string{".nano", ".micro", ".small", ".medium"} { + assert.False(t, strings.HasSuffix(name, small), + "non-burstable type %s should not have %s size", name, small) + } + } +} + +func TestGenerateInstanceTypes_OlderFamiliesHaveSmallerSizeRange(t *testing.T) { + types := generateInstanceTypes(vendorIntel, "x86_64") + + // m4 (older Intel GP) should have 6 sizes (large → 16xlarge, no 24xlarge) + assert.Equal(t, 6, countFamily(types, "m4."), "m4 should have 6 sizes (large → 16xlarge)") + // m5 (newer Intel GP) should have 8 sizes (large → 24xlarge) + assert.Equal(t, 8, countFamily(types, "m5."), "m5 should have 8 sizes (large → 24xlarge)") +} + +func TestGenerateInstanceTypes_BurstableFlag(t *testing.T) { + types := generateInstanceTypes(vendorIntel, "x86_64") + + // Previous-gen families: t2, m4, c4, r4 + prevGen := map[string]bool{"t2": true, "m4": true, "c4": true, "r4": true} + + for name, info := range types { + isBurstable := strings.HasPrefix(name, "t") + family := strings.SplitN(name, ".", 2)[0] + assert.Equal(t, isBurstable, *info.BurstablePerformanceSupported, + "%s burstable flag mismatch", name) + assert.Equal(t, !prevGen[family], *info.CurrentGeneration, + "%s current generation flag mismatch", name) + } +} + +// --- vendorFromCPU --- + +func TestVendorFromCPU_ExactMappings(t *testing.T) { + assert.Equal(t, vendorAMD, vendorFromCPU("AMD EPYC 7551", "x86_64")) + assert.Equal(t, vendorAMD, vendorFromCPU("AMD Ryzen 9 7950X", "x86_64")) + assert.Equal(t, vendorAMD, vendorFromCPU("AMD Ryzen Threadripper PRO 5995WX", "x86_64")) + assert.Equal(t, vendorIntel, vendorFromCPU("Intel Xeon E5-2686", "x86_64")) + assert.Equal(t, vendorARM, vendorFromCPU("Apple M1", "arm64")) + assert.Equal(t, vendorARM, vendorFromCPU("AWS Graviton2", "arm64")) + assert.Equal(t, vendorARM, vendorFromCPU("ARM Cortex-A72", "arm64")) +} + +func TestVendorFromCPU_CaseInsensitive(t *testing.T) { + assert.Equal(t, vendorAMD, vendorFromCPU("amd epyc 9654", "x86_64")) +} + +func TestVendorFromCPU_DefaultFallback(t *testing.T) { + // Unknown CPU on x86_64 defaults to Intel + assert.Equal(t, vendorIntel, vendorFromCPU("Unknown CPU Model XYZ", "x86_64")) + // Any CPU on arm64 always returns ARM + assert.Equal(t, vendorARM, vendorFromCPU("Unknown CPU Model XYZ", "arm64")) + assert.Equal(t, vendorARM, vendorFromCPU("Intel Xeon E5-2686", "arm64")) } // --- instanceTypeVCPUs / instanceTypeMemoryMiB nil safety ---