Skip to content

Commit bc41666

Browse files
mistercodedormistercodedor
committed
fix: address hardware profile PR comments
1 parent 1cc4b80 commit bc41666

5 files changed

Lines changed: 210 additions & 17 deletions

File tree

‎pkg/cmd/register/gpu_nvml.go‎

Lines changed: 34 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -8,18 +8,38 @@ import (
88
"github.com/NVIDIA/go-nvml/pkg/nvml"
99
)
1010

11-
// archNames maps NVML compute capability (major version) to GPU architecture name.
12-
var archNames = map[int]string{
13-
1: "Tesla",
14-
2: "Fermi",
15-
3: "Kepler",
16-
5: "Maxwell",
17-
6: "Pascal",
18-
7: "Volta/Turing",
19-
8: "Ampere",
20-
9: "Hopper/Ada Lovelace",
21-
10: "Blackwell",
22-
12: "Vera Rubin",
11+
// archName returns the GPU architecture name for the given CUDA compute capability.
12+
func archName(major, minor int) string {
13+
switch major {
14+
case 1:
15+
return "Tesla"
16+
case 2:
17+
return "Fermi"
18+
case 3:
19+
return "Kepler"
20+
case 5:
21+
return "Maxwell"
22+
case 6:
23+
return "Pascal"
24+
case 7:
25+
if minor >= 5 {
26+
return "Turing"
27+
}
28+
return "Volta"
29+
case 8:
30+
if minor >= 9 {
31+
return "Ada Lovelace"
32+
}
33+
return "Ampere"
34+
case 9:
35+
return "Hopper"
36+
case 10:
37+
return "Blackwell"
38+
case 12:
39+
return "Vera Rubin"
40+
default:
41+
return ""
42+
}
2343
}
2444

2545
// probeGPUsNVML uses NVML to detect GPUs and interconnects.
@@ -65,8 +85,8 @@ func probeGPUsNVML() ([]GPU, []Interconnect) {
6585
arch := ""
6686
major, minor, ret := device.GetCudaComputeCapability()
6787
if ret == nvml.SUCCESS {
68-
if archName, ok := archNames[major]; ok {
69-
arch = archName
88+
if name := archName(major, minor); name != "" {
89+
arch = name
7090
} else {
7191
arch = fmt.Sprintf("sm_%d%d", major, minor)
7292
}

‎pkg/cmd/register/hardware.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ type Interconnect struct {
3535
type StorageDevice struct {
3636
Name string `json:"name,omitempty"`
3737
StorageBytes int64 `json:"storage_bytes"`
38-
StorageType string `json:"storage_type,omitempty"` // "SSD" or "HDD"
38+
StorageType string `json:"storage_type,omitempty"` // "SSD", "HDD", or "NVMe"
3939
}
4040

4141
// HardwareProfile is the full hardware snapshot collected by a HardwareProfiler.

‎pkg/cmd/register/hardware_darwin.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,16 @@ package register
44

55
import (
66
"os/exec"
7+
"regexp"
78
"runtime"
9+
"strconv"
810
"strings"
911

1012
"golang.org/x/sys/unix"
1113
)
1214

15+
var diskutilBytesRe = regexp.MustCompile(`\((\d[\d,]*)\s+Bytes\)`)
16+
1317
// SystemHardwareProfiler probes hardware on macOS using sysctl and system_profiler.
1418
type SystemHardwareProfiler struct{}
1519

@@ -129,6 +133,14 @@ func isWholeDisk(name string) bool {
129133
}
130134

131135
func parseDiskutilSize(val string, dev *StorageDevice) {
136+
if m := diskutilBytesRe.FindStringSubmatch(val); len(m) == 2 {
137+
cleaned := strings.ReplaceAll(m[1], ",", "")
138+
if size, err := strconv.ParseInt(cleaned, 10, 64); err == nil {
139+
dev.StorageBytes = size
140+
return
141+
}
142+
}
143+
// Fallback: extract digits from first token.
132144
parts := strings.Fields(val)
133145
if len(parts) < 1 {
134146
return
@@ -143,6 +155,10 @@ func parseDiskutilSize(val string, dev *StorageDevice) {
143155
}
144156

145157
func parseDiskutilSolidState(val string, dev *StorageDevice) {
158+
// Don't overwrite a more specific type (e.g. NVMe from Protocol line).
159+
if dev.StorageType != "" {
160+
return
161+
}
146162
if strings.EqualFold(val, "yes") {
147163
dev.StorageType = "SSD"
148164
} else {

‎pkg/cmd/register/hardware_darwin_test.go‎

Lines changed: 155 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,155 @@
1+
//go:build darwin
2+
3+
package register
4+
5+
import (
6+
"testing"
7+
)
8+
9+
func Test_parseDiskutilSize(t *testing.T) {
10+
tests := []struct {
11+
name string
12+
val string
13+
want int64
14+
}{
15+
{
16+
"BytesFirst",
17+
"500107862016 Bytes (exactly 976773168 512-Byte-Units)",
18+
500107862016,
19+
},
20+
{
21+
"HumanReadableFirst",
22+
"1.0 TB (1,000,000,000,000 Bytes)",
23+
1000000000000,
24+
},
25+
{
26+
"GBWithParenBytes",
27+
"500.1 GB (500,107,862,016 Bytes)",
28+
500107862016,
29+
},
30+
{
31+
"NoParensFallback",
32+
"500107862016",
33+
500107862016,
34+
},
35+
{
36+
"Empty",
37+
"",
38+
0,
39+
},
40+
}
41+
for _, tt := range tests {
42+
t.Run(tt.name, func(t *testing.T) {
43+
dev := &StorageDevice{}
44+
parseDiskutilSize(tt.val, dev)
45+
if dev.StorageBytes != tt.want {
46+
t.Errorf("parseDiskutilSize(%q) = %d, want %d", tt.val, dev.StorageBytes, tt.want)
47+
}
48+
})
49+
}
50+
}
51+
52+
func Test_isWholeDisk(t *testing.T) {
53+
tests := []struct {
54+
name string
55+
want bool
56+
}{
57+
{"disk0", true},
58+
{"disk1", true},
59+
{"disk0s1", false},
60+
{"disk2s3", false},
61+
{"notadisk", false},
62+
}
63+
for _, tt := range tests {
64+
t.Run(tt.name, func(t *testing.T) {
65+
if got := isWholeDisk(tt.name); got != tt.want {
66+
t.Errorf("isWholeDisk(%q) = %v, want %v", tt.name, got, tt.want)
67+
}
68+
})
69+
}
70+
}
71+
72+
func Test_parseDiskutilSolidState(t *testing.T) {
73+
tests := []struct {
74+
val string
75+
want string
76+
}{
77+
{"Yes", "SSD"},
78+
{"yes", "SSD"},
79+
{"No", "HDD"},
80+
{"no", "HDD"},
81+
}
82+
for _, tt := range tests {
83+
t.Run(tt.val, func(t *testing.T) {
84+
dev := &StorageDevice{}
85+
parseDiskutilSolidState(tt.val, dev)
86+
if dev.StorageType != tt.want {
87+
t.Errorf("parseDiskutilSolidState(%q) → %q, want %q", tt.val, dev.StorageType, tt.want)
88+
}
89+
})
90+
}
91+
}
92+
93+
func Test_parseDiskutilInfoOutput(t *testing.T) {
94+
output := `**********
95+
96+
Device Identifier: disk0
97+
Device Node: /dev/disk0
98+
Whole: Yes
99+
Part of Whole: disk0
100+
Disk Size: 500.1 GB (500,107,862,016 Bytes)
101+
Protocol: NVMe
102+
Solid State: Yes
103+
Device / Media Name: APPLE SSD AP0512Q
104+
105+
**********
106+
107+
Device Identifier: disk0s1
108+
Device Node: /dev/disk0s1
109+
Whole: No
110+
Part of Whole: disk0
111+
Disk Size: 524.3 MB (524,288,000 Bytes)
112+
113+
**********
114+
115+
Device Identifier: disk1
116+
Device Node: /dev/disk1
117+
Whole: Yes
118+
Part of Whole: disk1
119+
Disk Size: 1.0 TB (1,000,000,000,000 Bytes)
120+
Solid State: No
121+
122+
`
123+
124+
devices := parseDiskutilInfoOutput(output)
125+
if len(devices) != 2 {
126+
t.Fatalf("expected 2 devices, got %d", len(devices))
127+
}
128+
129+
if devices[0].Name != "disk0" {
130+
t.Errorf("device 0 name = %q, want disk0", devices[0].Name)
131+
}
132+
if devices[0].StorageBytes != 500107862016 {
133+
t.Errorf("device 0 bytes = %d, want 500107862016", devices[0].StorageBytes)
134+
}
135+
if devices[0].StorageType != "NVMe" {
136+
t.Errorf("device 0 type = %q, want NVMe", devices[0].StorageType)
137+
}
138+
139+
if devices[1].Name != "disk1" {
140+
t.Errorf("device 1 name = %q, want disk1", devices[1].Name)
141+
}
142+
if devices[1].StorageBytes != 1000000000000 {
143+
t.Errorf("device 1 bytes = %d, want 1000000000000", devices[1].StorageBytes)
144+
}
145+
if devices[1].StorageType != "HDD" {
146+
t.Errorf("device 1 type = %q, want HDD", devices[1].StorageType)
147+
}
148+
}
149+
150+
func Test_parseDiskutilInfoOutput_Empty(t *testing.T) {
151+
devices := parseDiskutilInfoOutput("")
152+
if len(devices) != 0 {
153+
t.Errorf("expected 0 devices, got %d", len(devices))
154+
}
155+
}

‎pkg/cmd/register/hardware_stub.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
1-
//go:build !linux && !darwin && !windows
1+
//go:build !linux && !darwin
22

33
package register
44

5+
import "runtime"
6+
57
// SystemHardwareProfiler is a no-op adapter for unsupported platforms.
68
type SystemHardwareProfiler struct{}
79

810
func (p *SystemHardwareProfiler) Profile() (*HardwareProfile, error) {
9-
return &HardwareProfile{}, nil
11+
return &HardwareProfile{Architecture: runtime.GOARCH}, nil
1012
}

0 commit comments

Comments
 (0)