Create GPU index in actual driver

datarhei · Oct 31, 2024 · abc821f · abc821f
1 parent d591a23
commit abc821f
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 3 deletions.
diff --git a/resources/psutil/gpu/gpu.go b/resources/psutil/gpu/gpu.go
@@ -12,6 +12,7 @@ type Process struct {
 }
 
 type Stats struct {
+	Index        int
 	ID           string
 	Name         string
 	Architecture string

diff --git a/resources/psutil/gpu/nvidia/nvidia.go b/resources/psutil/gpu/nvidia/nvidia.go
@@ -444,8 +444,9 @@ func (n *nvidia) Stats() ([]gpu.Stats, error) {
 		return stats, n.err
 	}
 
-	for _, nv := range n.stats.GPU {
+	for i, nv := range n.stats.GPU {
 		s := gpu.Stats{
+			Index:        i,
 			ID:           nv.ID,
 			Name:         nv.Name,
 			Architecture: nv.Architecture,

diff --git a/resources/psutil/gpu/nvidia/nvidia_test.go b/resources/psutil/gpu/nvidia/nvidia_test.go
@@ -324,6 +324,7 @@ func TestNvidiaGPUStats(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, []gpu.Stats{
 		{
+			Index:        0,
 			ID:           "00000000:01:00.0",
 			Name:         "NVIDIA L4",
 			Architecture: "Ada Lovelace",
@@ -360,6 +361,7 @@ func TestNvidiaGPUStats(t *testing.T) {
 			},
 		},
 		{
+			Index:        1,
 			ID:           "00000000:C1:00.0",
 			Name:         "NVIDIA L4",
 			Architecture: "Ada Lovelace",

diff --git a/resources/psutil/psutil.go b/resources/psutil/psutil.go
@@ -639,9 +639,9 @@ func (u *util) GPU() ([]GPUInfo, error) {
 
 	stats := []GPUInfo{}
 
-	for i, nv := range nvstats {
+	for _, nv := range nvstats {
 		stats = append(stats, GPUInfo{
-			Index:       i,
+			Index:       nv.Index,
 			ID:          nv.ID,
 			Name:        nv.Name,
 			MemoryTotal: nv.MemoryTotal,