ClearML/go-nvlib
1
0
Fork 0
mirror of https://github.com/clearml/go-nvlib synced 2025-01-31 02:47:02 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Skip display devices based on device names

Browse Source 
This allows devices to be skipped based on device names and
skips "NVIDIA DGX Display" devices by default.

Signed-off-by: Evan Lezar <elezar@nvidia.com>
This commit is contained in:
Evan Lezar 2022-11-16 16:07:58 +01:00
parent 0e10f084d1
commit 655eb9795c
2 changed files with 27 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
pkg/nvlib/device/api.go
View File

@ -36,7 +36,7 @@ type Interface interface {
type devicelib struct {
nvml nvml.Interface
selectedDeviceClasses map[Class]struct{}
skippedDevices map[string]struct{}
}
var _ Interface = &devicelib{}
@ -50,10 +50,8 @@ func New(opts ...Option) Interface {
if d.nvml == nil {
d.nvml = nvml.New()
}
if d.selectedDeviceClasses == nil {
d.selectedDeviceClasses = map[Class]struct{}{
ClassCompute: {},
}
if d.skippedDevices == nil {
WithSkippedDevices("NVIDIA DGX Display")(d)
}
return d
}
@ -65,14 +63,14 @@ func WithNvml(nvml nvml.Interface) Option {
}
}
// WithSelectedDeviceClasses selects the specified device classes when filtering devices
func WithSelectedDeviceClasses(classes ...Class) Option {
// WithSkippedDevices provides an Option to set devices to be skipped by model name
func WithSkippedDevices(names ...string) Option {
return func(d *devicelib) {
if d.selectedDeviceClasses == nil {
d.selectedDeviceClasses = make(map[Class]struct{})
if d.skippedDevices == nil {
d.skippedDevices = make(map[string]struct{})
}
for _, c := range classes {
d.selectedDeviceClasses[c] = struct{}{}
for _, name := range names {
d.skippedDevices[name] = struct{}{}
}
}
}

20
pkg/nvlib/device/device.go
View File

@ -238,6 +238,20 @@ func (d *device) getClass() (Class, error) {
return Class(device.Class), nil
}
// isSkipped checks whether the device should be skipped.
func (d *device) isSkipped() (bool, error) {
name, ret := d.GetName()
if ret != nvml.SUCCESS {
return false, fmt.Errorf("error getting device name: %v", ret)
}
if _, exists := d.lib.skippedDevices[name]; exists {
return true, nil
}
return false, nil
}
// VisitDevices visits each top-level device and invokes a callback function for it
func (d *devicelib) VisitDevices(visit func(int, Device) error) error {
count, ret := d.nvml.DeviceGetCount()
@ -255,11 +269,11 @@ func (d *devicelib) VisitDevices(visit func(int, Device) error) error {
return fmt.Errorf("error creating new device wrapper: %v", err)
}
class, err := dev.getClass()
isSkipped, err := dev.isSkipped()
if err != nil {
return fmt.Errorf("error getting PCI device class for device: %v", err)
return fmt.Errorf("error checking whether device is skipped: %v", err)
}
if !d.classIsSelected(class) {
if isSkipped {
continue
}