Skip to content

[Telemetry] Add machine_category metric for hardware classification (CPU/GPU/TPU) #6010

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
kadupoornima wants to merge 6 commits into
base: develop
Choose a base branch
from
+230 −0
Open

[Telemetry] Add machine_category metric for hardware classification (CPU/GPU/TPU) #6010

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
180ccf6
feat: Add machine_category telemetry metric
kadupoornima Jul 27, 2026
2cedfed
feat: Add machine_category telemetry metric
kadupoornima Jul 27, 2026
f662697
.
kadupoornima Jul 27, 2026
2fa16c6
Update pkg/telemetry/collector_util.go
kadupoornima Jul 27, 2026
3c31c86
.
kadupoornima Jul 27, 2026
4d3c7a8
comment-1
kadupoornima Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions pkg/telemetry/collector.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import (
"context"
"encoding/json"
"errors"
"fmt"
"hpc-toolkit/pkg/config"
"hpc-toolkit/pkg/shell"
"net"
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 {

@SwarnaBharathiMantena SwarnaBharathiMantena Jul 28, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a way we can reuse the getMachineType function here? The current getMachineCategory function will probably get simpler if we use getMachineType.

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)
Expand Down
163 changes: 163 additions & 0 deletions pkg/telemetry/collector_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ func TestCollectMetrics_Extensible(t *testing.T) {
expectedKeys := []string{
COMMAND_FLAGS,
MACHINE_TYPE,
MACHINE_CATEGORY,
REGION,
ZONE,
STATIC_NODE_COUNTS,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
},
Expand Down Expand Up @@ -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: "",
},
{
Comment thread
kadupoornima marked this conversation as resolved.
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)
}
})
}
}
49 changes: 49 additions & 0 deletions pkg/telemetry/collector_util.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
1 change: 1 addition & 0 deletions pkg/telemetry/types.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
Loading