diff --git a/pkg/monitor/nvidia/v0/spec.go b/pkg/monitor/nvidia/v0/spec.go index fde2c9dcd8..240ea429f8 100644 --- a/pkg/monitor/nvidia/v0/spec.go +++ b/pkg/monitor/nvidia/v0/spec.go @@ -78,7 +78,7 @@ func (s Spec) DeviceMax() int { } func (s Spec) DeviceNum() int { - return int(s.sr.num) + return int(min(s.sr.num, uint64(maxDevices))) } // activeProcs returns the process slots currently in use. procnum is read from diff --git a/pkg/monitor/nvidia/v0/spec_test.go b/pkg/monitor/nvidia/v0/spec_test.go index 812e8d5c5c..5963e4bf25 100644 --- a/pkg/monitor/nvidia/v0/spec_test.go +++ b/pkg/monitor/nvidia/v0/spec_test.go @@ -48,6 +48,8 @@ func TestSpec_DeviceNum(t *testing.T) { tests := []specTest{ {name: "device num is 4", spec: &Spec{sr: &sharedRegionT{num: 4}}, expected: 4}, {name: "device num is 8", spec: &Spec{sr: &sharedRegionT{num: 8}}, expected: 8}, + {name: "num larger than maxDevices is clamped", spec: &Spec{sr: &sharedRegionT{num: 9999}}, expected: maxDevices}, + {name: "high-bit uint64 num is clamped not negative", spec: &Spec{sr: &sharedRegionT{num: 0x8000000000000001}}, expected: maxDevices}, } for _, tt := range tests { diff --git a/pkg/monitor/nvidia/v1/spec.go b/pkg/monitor/nvidia/v1/spec.go index a0e214bbdc..e8d71325f8 100644 --- a/pkg/monitor/nvidia/v1/spec.go +++ b/pkg/monitor/nvidia/v1/spec.go @@ -92,7 +92,7 @@ func (s Spec) DeviceMax() int { } func (s Spec) DeviceNum() int { - return int(s.sr.num) + return int(min(s.sr.num, uint64(maxDevices))) } // activeProcs returns the process slots currently in use. procnum is read from diff --git a/pkg/monitor/nvidia/v1/spec_test.go b/pkg/monitor/nvidia/v1/spec_test.go index 9761a840c7..525b02616e 100644 --- a/pkg/monitor/nvidia/v1/spec_test.go +++ b/pkg/monitor/nvidia/v1/spec_test.go @@ -134,6 +134,24 @@ func Test_DeviceNum(t *testing.T) { }, want: int(4), }, + { + name: "num larger than maxDevices is clamped", + args: &Spec{ + sr: &sharedRegionT{ + num: 9999, + }, + }, + want: maxDevices, + }, + { + name: "high-bit uint64 num is clamped not negative", + args: &Spec{ + sr: &sharedRegionT{ + num: 0x8000000000000001, + }, + }, + want: maxDevices, + }, } for _, test := range tests { t.Run(test.name, func(t *testing.T) {