Skip to content

Commit ef65112

Browse files
kadupoornimakadupoornima
authored
[Telemetry] Add support to collect machine type information (#5494)
1 parent 587e4d1 commit ef65112

6 files changed

Lines changed: 303 additions & 11 deletions

File tree

‎pkg/config/config.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -961,3 +961,12 @@ func (bp *Blueprint) evalVars() (Dict, error) {
961961
}
962962
return NewDict(res), nil
963963
}
964+
965+
// GetAllModules returns a slice of all modules defined in the blueprint.
966+
func GetAllModules(bp *Blueprint) []Module {
967+
var modules []Module
968+
bp.WalkModulesSafe(func(_ ModulePath, m *Module) {
969+
modules = append(modules, *m)
970+
})
971+
return modules
972+
}

‎pkg/config/config_test.go‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@ import (
2424
"testing"
2525

2626
"hpc-toolkit/pkg/modulereader"
27+
"reflect"
2728

2829
"github.com/pkg/errors"
2930
"github.com/zclconf/go-cty/cty"
@@ -924,3 +925,88 @@ func (s *zeroSuite) TestValidateSlurmClusterName(c *C) {
924925
c.Check(err, ErrorMatches, ".*lowercase letter.*")
925926
}
926927
}
928+
929+
// TestGetAllModules verifies that the method correctly flattens and extracts
930+
// all modules defined across multiple deployment groups in a blueprint.
931+
func TestGetAllModules(t *testing.T) {
932+
tests := []struct {
933+
name string
934+
setupBp func() *Blueprint
935+
expected []ModuleID
936+
}{
937+
{
938+
name: "No modules or groups",
939+
setupBp: func() *Blueprint {
940+
return &Blueprint{
941+
Groups: []Group{},
942+
}
943+
},
944+
expected: []ModuleID{},
945+
},
946+
{
947+
name: "Single module in a single group",
948+
setupBp: func() *Blueprint {
949+
return &Blueprint{
950+
Groups: []Group{
951+
{
952+
Name: GroupName("group-1"),
953+
Modules: []Module{
954+
{ID: ModuleID("module-1"), Source: "source/1"},
955+
},
956+
},
957+
},
958+
}
959+
},
960+
expected: []ModuleID{"module-1"},
961+
},
962+
{
963+
name: "Multiple modules across multiple groups",
964+
setupBp: func() *Blueprint {
965+
return &Blueprint{
966+
Groups: []Group{
967+
{
968+
Name: GroupName("group-1"),
969+
Modules: []Module{
970+
{ID: ModuleID("module-1"), Source: "source/1"},
971+
{ID: ModuleID("module-2"), Source: "source/2"},
972+
},
973+
},
974+
{
975+
Name: GroupName("group-2"),
976+
Modules: []Module{
977+
{ID: ModuleID("module-3"), Source: "source/3"},
978+
},
979+
},
980+
},
981+
}
982+
},
983+
// The expected order should match the sequential walk across groups
984+
expected: []ModuleID{"module-1", "module-2", "module-3"},
985+
},
986+
}
987+
988+
for _, tt := range tests {
989+
t.Run(tt.name, func(t *testing.T) {
990+
bp := tt.setupBp()
991+
992+
// Call the method
993+
modules := GetAllModules(bp)
994+
995+
// Extract just the ModuleIDs to make the comparison simpler and more robust
996+
var actualIDs []ModuleID
997+
for _, m := range modules {
998+
actualIDs = append(actualIDs, m.ID)
999+
}
1000+
1001+
// Handle nil vs empty slice comparisons safely
1002+
if len(tt.expected) == 0 && len(actualIDs) == 0 {
1003+
return
1004+
}
1005+
1006+
// Verify that the exact sequence of ModuleIDs matches our expectation
1007+
if !reflect.DeepEqual(actualIDs, tt.expected) {
1008+
t.Errorf("GetAllModules() returned IDs %v, want %v", actualIDs, tt.expected)
1009+
}
1010+
})
1011+
}
1012+
}

‎pkg/telemetry/collector.go‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,16 @@ import (
2020
"strings"
2121
"time"
2222

23+
"github.com/zclconf/go-cty/cty"
24+
2325
"github.com/spf13/cobra"
2426
"github.com/spf13/pflag"
2527
)
2628

29+
var (
30+
machineTypeModulePattern = ".*modules.compute.*"
31+
)
32+
2733
// NewCollector creates and initializes a new Telemetry Collector.
2834
func NewCollector(cmd *cobra.Command, args []string) *Collector {
2935
return &Collector{
@@ -41,6 +47,7 @@ func (c *Collector) CollectMetrics(errorCode int) {
4147
defer c.mu.Unlock()
4248

4349
c.metadata[COMMAND_FLAGS] = getCmdFlags(c.eventCmd)
50+
c.metadata[MACHINE_TYPE] = getMachineType(c.blueprint)
4451
c.metadata[REGION] = getRegion(c.blueprint)
4552
c.metadata[ZONE] = getZone(c.blueprint)
4653
c.metadata[IS_TEST_DATA] = getIsTestData()
@@ -95,6 +102,38 @@ func getCmdFlags(cmd *cobra.Command) string {
95102
return strings.Join(flags, ",")
96103
}
97104

105+
func getMachineType(bp config.Blueprint) string {
106+
var machineTypes []string
107+
seen := make(map[string]bool) // To keep track of added machine types to avoid duplication
108+
modules := getModulesWithPattern(machineTypeModulePattern, bp)
109+
110+
evalAndAdd := func(key string, m config.Module) {
111+
if m.Settings.Has(key) {
112+
keyValue := m.Settings.Get(key)
113+
// Evaluate the value to resolve expressions like $(vars.key)
114+
evaluatedKey, err := bp.Eval(keyValue)
115+
if err != nil {
116+
return
117+
}
118+
// Some module outputs or references carry cty marks, so we unmark them safely before use.
119+
unmarkedKey, _ := evaluatedKey.Unmark()
120+
if !unmarkedKey.IsNull() && unmarkedKey.Type() == cty.String {
121+
mType := unmarkedKey.AsString()
122+
if !seen[mType] {
123+
machineTypes = append(machineTypes, mType)
124+
seen[mType] = true
125+
}
126+
}
127+
}
128+
}
129+
130+
for _, m := range modules {
131+
evalAndAdd("machine_type", m)
132+
evalAndAdd("node_type", m) // For schedmd-slurm-gcp-v6-nodeset-tpu module. It uses node_type setting instead of machine_type.
133+
}
134+
return strings.Join(machineTypes, ",")
135+
}
136+
98137
func getRegion(bp config.Blueprint) string {
99138
return getKeyFromBlueprint("region", bp)
100139
}

‎pkg/telemetry/collector_test.go‎

Lines changed: 153 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ func TestCollectMetrics_Extensible(t *testing.T) {
4747
// Define all expected metric keys from types.go
4848
expectedKeys := []string{
4949
COMMAND_FLAGS,
50+
MACHINE_TYPE,
5051
REGION,
5152
ZONE,
5253
IS_TEST_DATA,
@@ -73,47 +74,66 @@ func TestCollectMetrics_Extensible(t *testing.T) {
7374
_ = cmd.Flags().Set("project", "test-project")
7475
},
7576
setupCollector: func(c *Collector) {
76-
// Mock the blueprint variables to include region and zone
77+
// Mock the blueprint variables and modules
7778
c.blueprint = config.Blueprint{
7879
Vars: config.NewDict(map[string]cty.Value{
7980
"region": cty.StringVal("us-central1"),
8081
"zone": cty.StringVal("us-central1-a"),
8182
}),
83+
Groups: []config.Group{
84+
{
85+
Name: config.GroupName("primary"),
86+
Modules: []config.Module{
87+
{
88+
ID: config.ModuleID("compute_pool"),
89+
// Ensure the source matches the machineTypeModulePattern ".*modules.compute.*"
90+
Source: "modules/compute/vm-instance",
91+
Settings: config.NewDict(map[string]cty.Value{
92+
"machine_type": cty.StringVal("c2-standard-8"),
93+
}),
94+
},
95+
},
96+
},
97+
},
8298
}
8399
},
84100
expectedValues: map[string]string{
85-
COMMAND_FLAGS: "force,project",
86101
IS_TEST_DATA: "true",
87102
EXIT_CODE: "0",
103+
COMMAND_FLAGS: "force,project",
88104
REGION: "us-central1",
89105
ZONE: "us-central1-a",
106+
MACHINE_TYPE: "c2-standard-8",
90107
},
91108
},
92109
{
93-
name: "Failure exit code with missing region and zone",
94-
errorCode: 1, // Simulating a failure
110+
name: "Failure exit code with missing region, zone, and machine type",
111+
errorCode: 1,
95112
setupCmd: func(cmd *cobra.Command) {
96-
// No flags set for this test case
113+
// No flags set
97114
},
98115
setupCollector: func(c *Collector) {
99-
// Blueprint with empty vars to simulate missing region and zone
116+
// Blueprint with empty vars and no compute modules
100117
c.blueprint = config.Blueprint{
101-
Vars: config.NewDict(map[string]cty.Value{}),
118+
Vars: config.NewDict(map[string]cty.Value{}),
119+
Groups: []config.Group{},
102120
}
103121
},
104122
expectedValues: map[string]string{
105123
IS_TEST_DATA: "true",
106-
EXIT_CODE: "1", // Verify failure code is captured
107-
COMMAND_FLAGS: "", // Verify empty flags
108-
REGION: "", // Verify missing region
109-
ZONE: "", // Verify missing zone
124+
EXIT_CODE: "1",
125+
COMMAND_FLAGS: "",
126+
REGION: "",
127+
ZONE: "",
128+
MACHINE_TYPE: "", // Verify empty machine type when no matching modules exist
110129
},
111130
},
112131
}
113132

114133
for _, tt := range tests {
115134
t.Run(tt.name, func(t *testing.T) {
116135
cmd := &cobra.Command{Use: "mock"}
136+
117137
// Execute the setup function to apply flags to the command
118138
if tt.setupCmd != nil {
119139
tt.setupCmd(cmd)
@@ -423,3 +443,125 @@ func TestGetKeyFromBlueprint(t *testing.T) {
423443
})
424444
}
425445
}
446+
447+
// TestGetMachineType verifies that machine types are correctly extracted from the blueprint.
448+
func TestGetMachineType(t *testing.T) {
449+
tests := []struct {
450+
name string
451+
setupBp func() config.Blueprint
452+
expected string
453+
}{
454+
{
455+
name: "Single machine type in module",
456+
setupBp: func() config.Blueprint {
457+
return config.Blueprint{
458+
Groups: []config.Group{
459+
{
460+
Name: config.GroupName("primary"),
461+
Modules: []config.Module{
462+
{
463+
ID: config.ModuleID("compute_pool"),
464+
Source: "modules/compute/vm-instance",
465+
Settings: config.NewDict(map[string]cty.Value{
466+
"machine_type": cty.StringVal("c2-standard-8"),
467+
}),
468+
},
469+
},
470+
},
471+
},
472+
}
473+
},
474+
expected: "c2-standard-8",
475+
},
476+
{
477+
name: "Multiple different machine types",
478+
setupBp: func() config.Blueprint {
479+
return config.Blueprint{
480+
Groups: []config.Group{
481+
{
482+
Name: config.GroupName("primary"),
483+
Modules: []config.Module{
484+
{
485+
ID: config.ModuleID("login_node"),
486+
Source: "modules/compute/vm-instance",
487+
Settings: config.NewDict(map[string]cty.Value{
488+
"machine_type": cty.StringVal("n2-standard-2"),
489+
}),
490+
},
491+
{
492+
ID: config.ModuleID("lcontroller_node"),
493+
Source: "modules/compute/vm-instance",
494+
Settings: config.NewDict(map[string]cty.Value{
495+
"machine_type": cty.StringVal("n2-standard-2"),
496+
}),
497+
},
498+
{
499+
ID: config.ModuleID("compute_pool"),
500+
Source: "modules/compute/gke-node-pool",
501+
Settings: config.NewDict(map[string]cty.Value{
502+
"machine_type": cty.StringVal("c2-standard-8"),
503+
}),
504+
},
505+
},
506+
},
507+
},
508+
}
509+
},
510+
expected: "n2-standard-2,c2-standard-8",
511+
},
512+
{
513+
name: "TPU node type module (schedmd-slurm-gcp-v6-nodeset-tpu)",
514+
setupBp: func() config.Blueprint {
515+
return config.Blueprint{
516+
Groups: []config.Group{
517+
{
518+
Name: config.GroupName("primary"),
519+
Modules: []config.Module{
520+
{
521+
ID: config.ModuleID("tpu_nodeset"),
522+
Source: "community/modules/compute/schedmd-slurm-gcp-v6-nodeset-tpu",
523+
Settings: config.NewDict(map[string]cty.Value{
524+
"node_type": cty.StringVal("v4-8"),
525+
}),
526+
},
527+
},
528+
},
529+
},
530+
}
531+
},
532+
expected: "v4-8",
533+
},
534+
{
535+
name: "No machine types in modules",
536+
setupBp: func() config.Blueprint {
537+
return config.Blueprint{
538+
Groups: []config.Group{
539+
{
540+
Name: config.GroupName("primary"),
541+
Modules: []config.Module{
542+
{
543+
ID: config.ModuleID("vpc_network"),
544+
Source: "modules/network/vpc",
545+
Settings: config.NewDict(map[string]cty.Value{}),
546+
},
547+
},
548+
},
549+
},
550+
}
551+
},
552+
expected: "",
553+
},
554+
}
555+
556+
for _, tt := range tests {
557+
t.Run(tt.name, func(t *testing.T) {
558+
bp := tt.setupBp()
559+
560+
actual := getMachineType(bp)
561+
562+
if actual != tt.expected {
563+
t.Errorf("getMachineType() = %q, want %q", actual, tt.expected)
564+
}
565+
})
566+
}
567+
}

0 commit comments

Comments
 (0)