diff --git a/pkg/nvlib/device/api.go b/pkg/nvlib/device/api.go index 11aa139..c2a6517 100644 --- a/pkg/nvlib/device/api.go +++ b/pkg/nvlib/device/api.go @@ -38,7 +38,7 @@ type Interface interface { } type devicelib struct { - nvml nvml.Interface + nvmllib nvml.Interface skippedDevices map[string]struct{} verifySymbols *bool migProfiles []MigProfile @@ -47,14 +47,13 @@ type devicelib struct { var _ Interface = &devicelib{} // New creates a new instance of the 'device' interface. -func New(opts ...Option) Interface { - d := &devicelib{} +func New(nvmllib nvml.Interface, opts ...Option) Interface { + d := &devicelib{ + nvmllib: nvmllib, + } for _, opt := range opts { opt(d) } - if d.nvml == nil { - d.nvml = nvml.New() - } if d.verifySymbols == nil { verify := true d.verifySymbols = &verify @@ -68,13 +67,6 @@ func New(opts ...Option) Interface { return d } -// WithNvml provides an Option to set the NVML library used by the 'device' interface. -func WithNvml(nvml nvml.Interface) Option { - return func(d *devicelib) { - d.nvml = nvml - } -} - // WithVerifySymbols provides an option to toggle whether to verify select symbols exist in dynamic libraries before calling them. func WithVerifySymbols(verify bool) Option { return func(d *devicelib) { diff --git a/pkg/nvlib/device/device.go b/pkg/nvlib/device/device.go index 1051459..5e1510c 100644 --- a/pkg/nvlib/device/device.go +++ b/pkg/nvlib/device/device.go @@ -51,7 +51,7 @@ func (d *devicelib) NewDevice(dev nvml.Device) (Device, error) { // NewDeviceByUUID builds a new Device from a UUID. func (d *devicelib) NewDeviceByUUID(uuid string) (Device, error) { - dev, ret := d.nvml.DeviceGetHandleByUUID(uuid) + dev, ret := d.nvmllib.DeviceGetHandleByUUID(uuid) if ret != nvml.SUCCESS { return nil, fmt.Errorf("error getting device handle for uuid '%v': %v", uuid, ret) } @@ -334,13 +334,13 @@ func (d *device) isSkipped() (bool, error) { // 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() + count, ret := d.nvmllib.DeviceGetCount() if ret != nvml.SUCCESS { return fmt.Errorf("error getting device count: %v", ret) } for i := 0; i < count; i++ { - device, ret := d.nvml.DeviceGetHandleByIndex(i) + device, ret := d.nvmllib.DeviceGetHandleByIndex(i) if ret != nvml.SUCCESS { return fmt.Errorf("error getting device handle for index '%v': %v", i, ret) } @@ -469,5 +469,5 @@ func (d *devicelib) hasSymbol(symbol string) bool { return true } - return d.nvml.Extensions().LookupSymbol(symbol) == nil + return d.nvmllib.Extensions().LookupSymbol(symbol) == nil } diff --git a/pkg/nvlib/device/mig_device.go b/pkg/nvlib/device/mig_device.go index b02d417..7145a06 100644 --- a/pkg/nvlib/device/mig_device.go +++ b/pkg/nvlib/device/mig_device.go @@ -50,7 +50,7 @@ func (d *devicelib) NewMigDevice(handle nvml.Device) (MigDevice, error) { // NewMigDeviceByUUID builds a new MigDevice from a UUID. func (d *devicelib) NewMigDeviceByUUID(uuid string) (MigDevice, error) { - dev, ret := d.nvml.DeviceGetHandleByUUID(uuid) + dev, ret := d.nvmllib.DeviceGetHandleByUUID(uuid) if ret != nvml.SUCCESS { return nil, fmt.Errorf("error getting device handle for uuid '%v': %v", uuid, ret) } diff --git a/pkg/nvlib/device/mig_profile_test.go b/pkg/nvlib/device/mig_profile_test.go index df30502..7c17573 100644 --- a/pkg/nvlib/device/mig_profile_test.go +++ b/pkg/nvlib/device/mig_profile_test.go @@ -79,7 +79,7 @@ func newMockDeviceLib() Interface { }, } - return New(WithNvml(mockNvml), WithVerifySymbols(false)) + return New(mockNvml, WithVerifySymbols(false)) } func TestParseMigProfile(t *testing.T) { diff --git a/pkg/nvlib/info/builder.go b/pkg/nvlib/info/builder.go index 87f20f0..6168440 100644 --- a/pkg/nvlib/info/builder.go +++ b/pkg/nvlib/info/builder.go @@ -55,7 +55,7 @@ func New(opts ...Option) Interface { ) } if o.devicelib == nil { - o.devicelib = device.New(device.WithNvml(o.nvmllib)) + o.devicelib = device.New(o.nvmllib) } if o.platform == "" { o.platform = PlatformAuto