diff --git a/pkg/monitor/nvidia/v0/spec.go b/pkg/monitor/nvidia/v0/spec.go index 9163e76279..444bddf827 100644 --- a/pkg/monitor/nvidia/v0/spec.go +++ b/pkg/monitor/nvidia/v0/spec.go @@ -81,9 +81,17 @@ func (s Spec) DeviceNum() int { return int(s.sr.num) } +// activeProcs returns the process slots currently in use. procnum is read from +// the shared-memory region and may be corrupt (negative or larger than the +// backing array); clamp it to a valid range so slicing can never panic. +func (s Spec) activeProcs() []shrregProcSlotT { + n := min(max(int(s.sr.procnum), 0), len(s.sr.procs)) + return s.sr.procs[:n] +} + func (s Spec) DeviceMemoryContextSize(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs { + for _, p := range s.activeProcs() { v += p.used[idx].contextSize } return v @@ -91,7 +99,7 @@ func (s Spec) DeviceMemoryContextSize(idx int) uint64 { func (s Spec) DeviceMemoryModuleSize(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs { + for _, p := range s.activeProcs() { v += p.used[idx].moduleSize } return v @@ -99,7 +107,7 @@ func (s Spec) DeviceMemoryModuleSize(idx int) uint64 { func (s Spec) DeviceMemoryBufferSize(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs { + for _, p := range s.activeProcs() { v += p.used[idx].bufferSize } return v @@ -107,7 +115,7 @@ func (s Spec) DeviceMemoryBufferSize(idx int) uint64 { func (s Spec) DeviceMemoryOffset(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs { + for _, p := range s.activeProcs() { v += p.used[idx].offset } return v @@ -115,7 +123,7 @@ func (s Spec) DeviceMemoryOffset(idx int) uint64 { func (s Spec) DeviceMemoryTotal(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs { + for _, p := range s.activeProcs() { v += p.used[idx].total } return v @@ -123,7 +131,7 @@ func (s Spec) DeviceMemoryTotal(idx int) uint64 { func (s Spec) DeviceSmUtil(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs { + for _, p := range s.activeProcs() { v += p.deviceUtil[idx].smUtil } return v diff --git a/pkg/monitor/nvidia/v0/spec_test.go b/pkg/monitor/nvidia/v0/spec_test.go index 5234a8fe77..0e68e78250 100644 --- a/pkg/monitor/nvidia/v0/spec_test.go +++ b/pkg/monitor/nvidia/v0/spec_test.go @@ -65,7 +65,8 @@ func TestSpec_DeviceMemoryContextSize(t *testing.T) { { name: "device memory context size for index 1", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{contextSize: 100}, {contextSize: 200}}}, {used: [16]deviceMemory{{contextSize: 300}, {contextSize: 400}}}, @@ -77,7 +78,8 @@ func TestSpec_DeviceMemoryContextSize(t *testing.T) { { name: "device memory context size for index 0", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{contextSize: 100}, {contextSize: 200}}}, {used: [16]deviceMemory{{contextSize: 300}, {contextSize: 400}}}, @@ -86,6 +88,19 @@ func TestSpec_DeviceMemoryContextSize(t *testing.T) { input: 0, expected: uint64(400), }, + { + name: "stale proc slot beyond procnum is ignored", + spec: &Spec{sr: &sharedRegionT{ + num: 2, + procnum: 1, + procs: [1024]shrregProcSlotT{ + {used: [16]deviceMemory{{contextSize: 100}, {contextSize: 200}}}, + {used: [16]deviceMemory{{contextSize: 300}, {contextSize: 400}}}, + }, + }}, + input: 1, + expected: uint64(200), + }, } for _, tt := range tests { @@ -103,7 +118,8 @@ func TestSpec_DeviceMemoryModuleSize(t *testing.T) { { name: "device memory module size for index 1", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{moduleSize: 100}, {moduleSize: 200}}}, {used: [16]deviceMemory{{moduleSize: 300}, {moduleSize: 400}}}, @@ -115,7 +131,8 @@ func TestSpec_DeviceMemoryModuleSize(t *testing.T) { { name: "device memory module size for index 0", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{moduleSize: 100}, {moduleSize: 200}}}, {used: [16]deviceMemory{{moduleSize: 300}, {moduleSize: 400}}}, @@ -124,6 +141,19 @@ func TestSpec_DeviceMemoryModuleSize(t *testing.T) { input: 0, expected: uint64(400), }, + { + name: "stale proc slot beyond procnum is ignored", + spec: &Spec{sr: &sharedRegionT{ + num: 2, + procnum: 1, + procs: [1024]shrregProcSlotT{ + {used: [16]deviceMemory{{moduleSize: 100}, {moduleSize: 200}}}, + {used: [16]deviceMemory{{moduleSize: 300}, {moduleSize: 400}}}, + }, + }}, + input: 1, + expected: uint64(200), + }, } for _, tt := range tests { @@ -141,7 +171,8 @@ func TestSpec_DeviceMemoryBufferSize(t *testing.T) { { name: "device memory buffer size for index 1", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{bufferSize: 100}, {bufferSize: 200}}}, {used: [16]deviceMemory{{bufferSize: 300}, {bufferSize: 400}}}, @@ -153,7 +184,8 @@ func TestSpec_DeviceMemoryBufferSize(t *testing.T) { { name: "device memory buffer size for index 0", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{bufferSize: 100}, {bufferSize: 200}}}, {used: [16]deviceMemory{{bufferSize: 300}, {bufferSize: 400}}}, @@ -162,6 +194,19 @@ func TestSpec_DeviceMemoryBufferSize(t *testing.T) { input: 0, expected: uint64(400), }, + { + name: "stale proc slot beyond procnum is ignored", + spec: &Spec{sr: &sharedRegionT{ + num: 2, + procnum: 1, + procs: [1024]shrregProcSlotT{ + {used: [16]deviceMemory{{bufferSize: 100}, {bufferSize: 200}}}, + {used: [16]deviceMemory{{bufferSize: 300}, {bufferSize: 400}}}, + }, + }}, + input: 1, + expected: uint64(200), + }, } for _, tt := range tests { @@ -179,7 +224,8 @@ func TestSpec_DeviceMemoryOffset(t *testing.T) { { name: "device memory offset for index 1", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{offset: 100}, {offset: 200}}}, {used: [16]deviceMemory{{offset: 300}, {offset: 400}}}, @@ -191,7 +237,8 @@ func TestSpec_DeviceMemoryOffset(t *testing.T) { { name: "device memory offset for index 0", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{offset: 100}, {offset: 200}}}, {used: [16]deviceMemory{{offset: 300}, {offset: 400}}}, @@ -200,6 +247,19 @@ func TestSpec_DeviceMemoryOffset(t *testing.T) { input: 0, expected: uint64(400), }, + { + name: "stale proc slot beyond procnum is ignored", + spec: &Spec{sr: &sharedRegionT{ + num: 2, + procnum: 1, + procs: [1024]shrregProcSlotT{ + {used: [16]deviceMemory{{offset: 100}, {offset: 200}}}, + {used: [16]deviceMemory{{offset: 300}, {offset: 400}}}, + }, + }}, + input: 1, + expected: uint64(200), + }, } for _, tt := range tests { @@ -217,7 +277,8 @@ func TestSpec_DeviceMemoryTotal(t *testing.T) { { name: "device memory total for index 1", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{total: 100}, {total: 200}}}, {used: [16]deviceMemory{{total: 300}, {total: 400}}}, @@ -229,7 +290,8 @@ func TestSpec_DeviceMemoryTotal(t *testing.T) { { name: "device memory total for index 0", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {used: [16]deviceMemory{{total: 100}, {total: 200}}}, {used: [16]deviceMemory{{total: 300}, {total: 400}}}, @@ -238,6 +300,19 @@ func TestSpec_DeviceMemoryTotal(t *testing.T) { input: 0, expected: uint64(400), }, + { + name: "stale proc slot beyond procnum is ignored", + spec: &Spec{sr: &sharedRegionT{ + num: 2, + procnum: 1, + procs: [1024]shrregProcSlotT{ + {used: [16]deviceMemory{{total: 100}, {total: 200}}}, + {used: [16]deviceMemory{{total: 300}, {total: 400}}}, + }, + }}, + input: 1, + expected: uint64(200), + }, } for _, tt := range tests { @@ -255,7 +330,8 @@ func TestSpec_DeviceSmUtil(t *testing.T) { { name: "device sm util for index 1", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {deviceUtil: [16]deviceUtilization{{smUtil: 100}, {smUtil: 200}}}, {deviceUtil: [16]deviceUtilization{{smUtil: 300}, {smUtil: 400}}}, @@ -267,7 +343,8 @@ func TestSpec_DeviceSmUtil(t *testing.T) { { name: "device sm util for index 0", spec: &Spec{sr: &sharedRegionT{ - num: 2, + num: 2, + procnum: 2, procs: [1024]shrregProcSlotT{ {deviceUtil: [16]deviceUtilization{{smUtil: 100}, {smUtil: 200}}}, {deviceUtil: [16]deviceUtilization{{smUtil: 300}, {smUtil: 400}}}, @@ -276,6 +353,19 @@ func TestSpec_DeviceSmUtil(t *testing.T) { input: 0, expected: uint64(400), }, + { + name: "stale proc slot beyond procnum is ignored", + spec: &Spec{sr: &sharedRegionT{ + num: 2, + procnum: 1, + procs: [1024]shrregProcSlotT{ + {deviceUtil: [16]deviceUtilization{{smUtil: 100}, {smUtil: 200}}}, + {deviceUtil: [16]deviceUtilization{{smUtil: 300}, {smUtil: 400}}}, + }, + }}, + input: 1, + expected: uint64(200), + }, } for _, tt := range tests { @@ -288,6 +378,36 @@ func TestSpec_DeviceSmUtil(t *testing.T) { } } +func TestSpec_CorruptProcnumIsClamped(t *testing.T) { + tests := []struct { + name string + procnum int32 + expected uint64 + }{ + // Negative procnum clamps to 0 active slots, so nothing is summed. + {name: "negative procnum", procnum: -5, expected: 0}, + // procnum larger than the backing array clamps to its length; only the + // single populated slot contributes. + {name: "procnum over backing array", procnum: 2000, expected: 100}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sr := &sharedRegionT{num: 2, procnum: tt.procnum} + sr.procs[0].used[0].total = 100 + sr.procs[0].deviceUtil[0].smUtil = 100 + s := Spec{sr: sr} + // A corrupt procnum must never panic the slice bound. + if got := s.DeviceMemoryTotal(0); got != tt.expected { + t.Errorf("DeviceMemoryTotal(0) = %d, want %d", got, tt.expected) + } + if got := s.DeviceSmUtil(0); got != tt.expected { + t.Errorf("DeviceSmUtil(0) = %d, want %d", got, tt.expected) + } + }) + } +} + func TestDeviceMemoryLimit(t *testing.T) { testCases := []struct { name string diff --git a/pkg/monitor/nvidia/v1/spec.go b/pkg/monitor/nvidia/v1/spec.go index bf01ca4113..f297703aaa 100644 --- a/pkg/monitor/nvidia/v1/spec.go +++ b/pkg/monitor/nvidia/v1/spec.go @@ -88,9 +88,17 @@ func (s Spec) DeviceNum() int { return int(s.sr.num) } +// activeProcs returns the process slots currently in use. procnum is read from +// the shared-memory region and may be corrupt (negative or larger than the +// backing array); clamp it to a valid range so slicing can never panic. +func (s Spec) activeProcs() []shrregProcSlotT { + n := min(max(int(s.sr.procnum), 0), len(s.sr.procs)) + return s.sr.procs[:n] +} + func (s Spec) DeviceMemoryContextSize(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs[:int(s.sr.procnum)] { + for _, p := range s.activeProcs() { v += p.used[idx].contextSize } return v @@ -98,7 +106,7 @@ func (s Spec) DeviceMemoryContextSize(idx int) uint64 { func (s Spec) DeviceMemoryModuleSize(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs[:int(s.sr.procnum)] { + for _, p := range s.activeProcs() { v += p.used[idx].moduleSize } return v @@ -106,7 +114,7 @@ func (s Spec) DeviceMemoryModuleSize(idx int) uint64 { func (s Spec) DeviceMemoryBufferSize(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs[:int(s.sr.procnum)] { + for _, p := range s.activeProcs() { v += p.used[idx].bufferSize } return v @@ -114,7 +122,7 @@ func (s Spec) DeviceMemoryBufferSize(idx int) uint64 { func (s Spec) DeviceMemoryOffset(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs[:int(s.sr.procnum)] { + for _, p := range s.activeProcs() { v += p.used[idx].offset } return v @@ -122,7 +130,7 @@ func (s Spec) DeviceMemoryOffset(idx int) uint64 { func (s Spec) DeviceMemoryTotal(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs[:int(s.sr.procnum)] { + for _, p := range s.activeProcs() { v += p.used[idx].total } return v @@ -130,7 +138,7 @@ func (s Spec) DeviceMemoryTotal(idx int) uint64 { func (s Spec) DeviceSmUtil(idx int) uint64 { v := uint64(0) - for _, p := range s.sr.procs[:int(s.sr.procnum)] { + for _, p := range s.activeProcs() { v += p.deviceUtil[idx].smUtil } return v diff --git a/pkg/monitor/nvidia/v1/spec_test.go b/pkg/monitor/nvidia/v1/spec_test.go index 175ca3a226..f5a3ee71d7 100644 --- a/pkg/monitor/nvidia/v1/spec_test.go +++ b/pkg/monitor/nvidia/v1/spec_test.go @@ -1083,3 +1083,29 @@ func Test_SetUtilizationSwitch(t *testing.T) { }) } } + +func TestSpec_CorruptProcnumIsClamped(t *testing.T) { + tests := []struct { + name string + procnum int32 + expected uint64 + }{ + // Negative procnum clamps to 0 active slots, so nothing is summed. + {name: "negative procnum", procnum: -5, expected: 0}, + // procnum larger than the backing array clamps to its length; only the + // single populated slot contributes. + {name: "procnum over backing array", procnum: 2000, expected: 100}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + sr := &sharedRegionT{num: 2, procnum: tt.procnum} + sr.procs[0].used[0].total = 100 + sr.procs[0].deviceUtil[0].smUtil = 100 + s := Spec{sr: sr} + // A corrupt procnum must never panic the slice bound. + assert.Equal(t, tt.expected, s.DeviceMemoryTotal(0)) + assert.Equal(t, tt.expected, s.DeviceSmUtil(0)) + }) + } +}