diff --git a/pkg/telemetry/collector.go b/pkg/telemetry/collector.go index 5a130c89c4..659e847255 100644 --- a/pkg/telemetry/collector.go +++ b/pkg/telemetry/collector.go @@ -18,6 +18,7 @@ import ( "context" "encoding/json" "errors" + "fmt" "hpc-toolkit/pkg/config" "hpc-toolkit/pkg/shell" "net" @@ -66,6 +67,7 @@ func (c *Collector) CollectMetrics(errorCode int, err error) { c.metadata[IS_SLURM] = getIsSlurm(bpModulesList) c.metadata[IS_VM_INSTANCE] = getIsVmInstance(bpModulesList) c.metadata[MACHINE_TYPE] = getMachineType(c.blueprint) + c.metadata[MACHINE_CATEGORY] = getMachineCategory(c.blueprint) c.metadata[STORAGE_TYPE] = getStorageType(c.blueprint) c.metadata[REGION] = getRegion(c.blueprint) c.metadata[ZONE] = getZone(c.blueprint) @@ -231,6 +233,21 @@ func getMachineType(bp config.Blueprint) string { return strings.Join(machineTypes, ",") } +// getMachineCategory maps each machine type to its inferred hardware category (CPU/GPU/TPU/Other) +func getMachineCategory(bp config.Blueprint) string { + machineTypes := getMachineType(bp) + if machineTypes == "" { + return "" + } + + var categories []string + for _, mt := range strings.Split(machineTypes, ",") { + categories = append(categories, fmt.Sprintf("%s:%s", mt, detectMachineCategory(mt))) + } + + return strings.Join(categories, ",") +} + func getStorageType(bp config.Blueprint) string { var storageTypes []string seen := make(map[string]bool) diff --git a/pkg/telemetry/collector_test.go b/pkg/telemetry/collector_test.go index 9898de22db..4bbec0300f 100644 --- a/pkg/telemetry/collector_test.go +++ b/pkg/telemetry/collector_test.go @@ -59,6 +59,7 @@ func TestCollectMetrics_Extensible(t *testing.T) { expectedKeys := []string{ COMMAND_FLAGS, MACHINE_TYPE, + MACHINE_CATEGORY, REGION, ZONE, STATIC_NODE_COUNTS, @@ -124,6 +125,7 @@ func TestCollectMetrics_Extensible(t *testing.T) { REGION: "us-central1", ZONE: "us-central1-a", MACHINE_TYPE: "c2-standard-8", + MACHINE_CATEGORY: "c2-standard-8:CPU", STATIC_NODE_COUNTS: "c2-standard-8:1", OS_NAME: getOSName(), // Dynamically expect the current OS name OS_VERSION: getOSVersion(), // Dynamically expect the current OS version @@ -156,6 +158,7 @@ func TestCollectMetrics_Extensible(t *testing.T) { OS_VERSION: getOSVersion(), // Verify OS info is still collected on failure TERRAFORM_VERSION: getTerraformVersion(), // Verify Terraform version is still collected on failure MACHINE_TYPE: "", // Verify empty machine type when no matching modules exist + MACHINE_CATEGORY: "", STATIC_NODE_COUNTS: "", INSTALLATION_MODE: BINARY, }, @@ -3010,3 +3013,163 @@ func TestGetIsAIAssisted(t *testing.T) { }) } } + +func TestGetMachineCategory(t *testing.T) { + tests := []struct { + name string + bp config.Blueprint + want string + }{ + { + name: "GPU, CPU and TPU mapping", + bp: config.Blueprint{ + Groups: []config.Group{ + { + Name: config.GroupName("primary"), + Modules: []config.Module{ + { + ID: config.ModuleID("compute_node_1"), + Settings: config.NewDict(map[string]cty.Value{ + "machine_type": cty.StringVal("g4-standard-48"), + }), + }, + { + ID: config.ModuleID("compute_node_2"), + Settings: config.NewDict(map[string]cty.Value{ + "machine_type": cty.StringVal("tpu7x-standard-4t"), + }), + }, + { + ID: config.ModuleID("compute_node_3"), + Settings: config.NewDict(map[string]cty.Value{ + "machine_type": cty.StringVal("c2-standard-4"), + }), + }, + }, + }, + }, + }, + want: "g4-standard-48:GPU,tpu7x-standard-4t:TPU,c2-standard-4:CPU", + }, + { + name: "Handles shorthand mapping for TPU and GPU", + bp: config.Blueprint{ + Groups: []config.Group{ + { + Name: config.GroupName("primary"), + Modules: []config.Module{ + { + ID: config.ModuleID("compute_node_4"), + Settings: config.NewDict(map[string]cty.Value{ + "machine_type": cty.StringVal("a100-40gb-1"), // g2 short hand actually a2 + }), + }, + { + ID: config.ModuleID("compute_node_5"), + Settings: config.NewDict(map[string]cty.Value{ + "machine_type": cty.StringVal("v6e-4"), // v6e TPU shorthand + }), + }, + }, + }, + }, + }, + want: "a100-40gb-1:GPU,v6e-4:TPU", + }, + { + name: "Handles unknown machine types mapped to Other", + bp: config.Blueprint{ + Groups: []config.Group{ + { + Name: config.GroupName("primary"), + Modules: []config.Module{ + { + ID: config.ModuleID("compute_unknown"), + Settings: config.NewDict(map[string]cty.Value{ + "machine_type": cty.StringVal("bizarre-type-1"), + }), + }, + }, + }, + }, + }, + want: "bizarre-type-1:Other", + }, + { + name: "Handles trailing spaces and upper casing in machine types to normalize appropriately", + bp: config.Blueprint{ + Groups: []config.Group{ + { + Name: config.GroupName("primary"), + Modules: []config.Module{ + { + ID: config.ModuleID("compute_noisy"), + Settings: config.NewDict(map[string]cty.Value{ + "machine_type": cty.StringVal(" N2-Standard-4 "), + }), + }, + }, + }, + }, + }, + want: "N2-Standard-4:CPU", + }, + { + name: "Handles deduping duplicate machine types", + bp: config.Blueprint{ + Groups: []config.Group{ + { + Name: config.GroupName("primary"), + Modules: []config.Module{ + { + ID: config.ModuleID("compute_1"), + Settings: config.NewDict(map[string]cty.Value{ + "machine_type": cty.StringVal("e2-micro"), + }), + }, + { + ID: config.ModuleID("compute_2"), + Settings: config.NewDict(map[string]cty.Value{ + "machine_type": cty.StringVal("e2-micro"), + }), + }, + }, + }, + }, + }, + want: "e2-micro:CPU", + }, + { + name: "Returns empty string for blueprint with no compute instances", + bp: config.Blueprint{ + Groups: []config.Group{ + { + Name: config.GroupName("primary"), + Modules: []config.Module{ + { + ID: config.ModuleID("storage"), + Settings: config.NewDict(map[string]cty.Value{ + "disk_type": cty.StringVal("pd-standard"), + }), + }, + }, + }, + }, + }, + want: "", + }, + { + name: "Returns empty string for empty blueprint", + bp: config.Blueprint{}, + want: "", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if got := getMachineCategory(tt.bp); got != tt.want { + t.Errorf("getMachineCategory() = %v, want %v", got, tt.want) + } + }) + } +} diff --git a/pkg/telemetry/collector_util.go b/pkg/telemetry/collector_util.go index 918d38c28b..3f35ac5f6e 100644 --- a/pkg/telemetry/collector_util.go +++ b/pkg/telemetry/collector_util.go @@ -1075,3 +1075,52 @@ func getModuleDynamicNodeCounts(m config.Module, bp config.Blueprint, targetKeys } return counts } + +// detectMachineCategory attempts to identify the hardware nature of an instantiated machine type by interpreting available cluster toolkit mappings (e.g. associating mapped nvidia labels with GPUs). +// Returns "CPU", "GPU", "TPU", or "Other". +func detectMachineCategory(mType string) string { + mType = strings.ToLower(strings.TrimSpace(mType)) + mType = config.ResolveMachineType(mType) + + if config.IsTPU(mType) { + return "TPU" + } + + mappings := config.GetMachineMappings() + + for family, label := range mappings.MachineFamilyToLabelMap { + if strings.HasPrefix(mType, family) && strings.Contains(label, "nvidia") { + return "GPU" + } + } + + if strings.Contains(mType, "gpu") { + return "GPU" + } + + for _, p := range mappings.CpuMachineFamilies { + if strings.HasPrefix(mType, p) { + return "CPU" + } + } + + if isCPUFallback(mType) { + return "CPU" + } + + return "Other" +} + +// 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-", "-micro", "-small", "-medium", "-metal", + } + for _, kw := range cpuKeywords { + if strings.Contains(mType, kw) { + return true + } + } + return false +} diff --git a/pkg/telemetry/types.go b/pkg/telemetry/types.go index c8d55c6452..d9bc7c095d 100644 --- a/pkg/telemetry/types.go +++ b/pkg/telemetry/types.go @@ -89,6 +89,7 @@ const ( IS_SLURM = "CLUSTER_TOOLKIT_IS_SLURM" IS_VM_INSTANCE = "CLUSTER_TOOLKIT_IS_VM_INSTANCE" MACHINE_TYPE = "CLUSTER_TOOLKIT_MACHINE_TYPE" + MACHINE_CATEGORY = "CLUSTER_TOOLKIT_MACHINE_CATEGORY" STORAGE_TYPE = "CLUSTER_TOOLKIT_STORAGE_TYPE" REGION = "CLUSTER_TOOLKIT_REGION" ZONE = "CLUSTER_TOOLKIT_ZONE"