comment-1

Author: kadupoornima

pkg/telemetry/collector.go

@@ -235,14 +235,14 @@ func getMachineType(bp config.Blueprint) string {
 
 // getMachineCategory maps each machine type to its inferred hardware category (CPU/GPU/TPU/Other)
 func getMachineCategory(bp config.Blueprint) string {
-	var categories []string
-	seen := make(map[string]bool)
+	machineTypes := getMachineType(bp)
+	if machineTypes == "" {
+		return ""
+	}
 
-	for _, m := range config.GetAllBpModules(&bp) {
-		if mt := getMachineTypeFromModule(m, bp); mt != "" && !seen[mt] {
-			categories = append(categories, fmt.Sprintf("%s:%s", mt, detectMachineCategory(mt)))
-			seen[mt] = true
-		}
+	var categories []string
+	for _, mt := range strings.Split(machineTypes, ",") {
+		categories = append(categories, fmt.Sprintf("%s:%s", mt, detectMachineCategory(mt)))
 	}
 
 	return strings.Join(categories, ",")

pkg/telemetry/collector_test.go

@@ -3076,31 +3076,6 @@ func TestGetMachineCategory(t *testing.T) {
 			},
 			want: "a100-40gb-1:GPU,v6e-4:TPU",
 		},
-		{
-			name: "Handles memory-optimized machine types correctly as CPU",
-			bp: config.Blueprint{
-				Groups: []config.Group{
-					{
-						Name: config.GroupName("primary"),
-						Modules: []config.Module{
-							{
-								ID: config.ModuleID("compute_mem_1"),
-								Settings: config.NewDict(map[string]cty.Value{
-									"machine_type": cty.StringVal("m1-megamem-96"),
-								}),
-							},
-							{
-								ID: config.ModuleID("compute_mem_2"),
-								Settings: config.NewDict(map[string]cty.Value{
-									"machine_type": cty.StringVal("m2-ultramem-208"),
-								}),
-							},
-						},
-					},
-				},
-			},
-			want: "m1-megamem-96:CPU,m2-ultramem-208:CPU",
-		},
 		{
 			name: "Handles unknown machine types mapped to Other",
 			bp: config.Blueprint{

pkg/telemetry/collector_util.go

@@ -1114,8 +1114,8 @@ func detectMachineCategory(mType string) string {
 // isCPUFallback checks if the machine type contains any of the standard CPU identifiers.
 func isCPUFallback(mType string) bool {
 	cpuKeywords := []string{
-		"-standard-", "-highmem-", "-highcpu-",
-		"-megamem-", "-ultramem-", "custom-",
+		"-standard-", "-highmem-", "-highcpu-", "-megamem-",
+		"-ultramem-", "custom-", "-micro", "-small", "-medium", "-metal",
 	}
 	for _, kw := range cpuKeywords {
 		if strings.Contains(mType, kw) {