ClearML/go-nvlib
1
0
Fork 0
mirror of https://github.com/clearml/go-nvlib synced 2025-04-20 06:04:53 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #39 from elezar/make-nvmllib-requried

Browse Source 
Make nvmllib a requried argument to devicelib
This commit is contained in:
Evan Lezar 2024-05-27 18:12:10 +02:00 committed by GitHub
commit 466340618f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 12 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -38,7 +38,7 @@ type Interface interface {
} }
type devicelib struct { type devicelib struct {
nvml nvml.Interface nvmllib nvml.Interface
skippedDevices map[string]struct{} skippedDevices map[string]struct{}
verifySymbols *bool verifySymbols *bool
migProfiles []MigProfile migProfiles []MigProfile
@ -47,14 +47,13 @@ type devicelib struct {
var _ Interface = &devicelib{} var _ Interface = &devicelib{}
// New creates a new instance of the 'device' interface. // New creates a new instance of the 'device' interface.
func New(opts ...Option) Interface { func New(nvmllib nvml.Interface, opts ...Option) Interface {
d := &devicelib{} d := &devicelib{
nvmllib: nvmllib,
}
for _, opt := range opts { for _, opt := range opts {
opt(d) opt(d)
} }
if d.nvml == nil {
d.nvml = nvml.New()
}
if d.verifySymbols == nil { if d.verifySymbols == nil {
verify := true verify := true
d.verifySymbols = &verify d.verifySymbols = &verify
@ -68,13 +67,6 @@ func New(opts ...Option) Interface {
return d 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. // WithVerifySymbols provides an option to toggle whether to verify select symbols exist in dynamic libraries before calling them.
func WithVerifySymbols(verify bool) Option { func WithVerifySymbols(verify bool) Option {
return func(d *devicelib) { return func(d *devicelib) {

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

@ -51,7 +51,7 @@ func (d *devicelib) NewDevice(dev nvml.Device) (Device, error) {
// NewDeviceByUUID builds a new Device from a UUID. // NewDeviceByUUID builds a new Device from a UUID.
func (d *devicelib) NewDeviceByUUID(uuid string) (Device, error) { func (d *devicelib) NewDeviceByUUID(uuid string) (Device, error) {
dev, ret := d.nvml.DeviceGetHandleByUUID(uuid) dev, ret := d.nvmllib.DeviceGetHandleByUUID(uuid)
if ret != nvml.SUCCESS { if ret != nvml.SUCCESS {
return nil, fmt.Errorf("error getting device handle for uuid '%v': %v", uuid, ret) 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. // VisitDevices visits each top-level device and invokes a callback function for it.
func (d *devicelib) VisitDevices(visit func(int, Device) error) error { func (d *devicelib) VisitDevices(visit func(int, Device) error) error {
count, ret := d.nvml.DeviceGetCount() count, ret := d.nvmllib.DeviceGetCount()
if ret != nvml.SUCCESS { if ret != nvml.SUCCESS {
return fmt.Errorf("error getting device count: %v", ret) return fmt.Errorf("error getting device count: %v", ret)
} }
for i := 0; i < count; i++ { for i := 0; i < count; i++ {
device, ret := d.nvml.DeviceGetHandleByIndex(i) device, ret := d.nvmllib.DeviceGetHandleByIndex(i)
if ret != nvml.SUCCESS { if ret != nvml.SUCCESS {
return fmt.Errorf("error getting device handle for index '%v': %v", i, ret) 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 true
} }
return d.nvml.Extensions().LookupSymbol(symbol) == nil return d.nvmllib.Extensions().LookupSymbol(symbol) == nil
} }

2
pkg/nvlib/device/mig_device.go
View File

@ -50,7 +50,7 @@ func (d *devicelib) NewMigDevice(handle nvml.Device) (MigDevice, error) {
// NewMigDeviceByUUID builds a new MigDevice from a UUID. // NewMigDeviceByUUID builds a new MigDevice from a UUID.
func (d *devicelib) NewMigDeviceByUUID(uuid string) (MigDevice, error) { func (d *devicelib) NewMigDeviceByUUID(uuid string) (MigDevice, error) {
dev, ret := d.nvml.DeviceGetHandleByUUID(uuid) dev, ret := d.nvmllib.DeviceGetHandleByUUID(uuid)
if ret != nvml.SUCCESS { if ret != nvml.SUCCESS {
return nil, fmt.Errorf("error getting device handle for uuid '%v': %v", uuid, ret) return nil, fmt.Errorf("error getting device handle for uuid '%v': %v", uuid, ret)
} }

2
pkg/nvlib/device/mig_profile_test.go
View File

@ -79,7 +79,7 @@ func newMockDeviceLib() Interface {
}, },
} }
return New(WithNvml(mockNvml), WithVerifySymbols(false)) return New(mockNvml, WithVerifySymbols(false))
} }
func TestParseMigProfile(t *testing.T) { func TestParseMigProfile(t *testing.T) {

2
pkg/nvlib/info/builder.go
View File

@ -55,7 +55,7 @@ func New(opts ...Option) Interface {
) )
} }
if o.devicelib == nil { if o.devicelib == nil {
o.devicelib = device.New(device.WithNvml(o.nvmllib)) o.devicelib = device.New(o.nvmllib)
} }
if o.platform == "" { if o.platform == "" {
o.platform = PlatformAuto o.platform = PlatformAuto