diff --git a/pkg/nvlib/device/device.go b/pkg/nvlib/device/device.go index 62e6c3e..4183095 100644 --- a/pkg/nvlib/device/device.go +++ b/pkg/nvlib/device/device.go @@ -19,7 +19,6 @@ package device import ( "fmt" - "github.com/NVIDIA/go-nvml/pkg/dl" "gitlab.com/nvidia/cloud-native/go-nvlib/pkg/nvml" ) @@ -152,8 +151,7 @@ func (d *device) GetCudaComputeCapabilityAsString() (string, error) { // IsMigCapable checks if a device is capable of having MIG paprtitions created on it func (d *device) IsMigCapable() (bool, error) { - err := d.lib.nvmlLookupSymbol("nvmlDeviceGetMigMode") - if err != nil { + if !d.lib.hasSymbol("nvmlDeviceGetMigMode") { return false, nil } @@ -170,8 +168,7 @@ func (d *device) IsMigCapable() (bool, error) { // IsMigEnabled checks if a device has MIG mode currently enabled on it func (d *device) IsMigEnabled() (bool, error) { - err := d.lib.nvmlLookupSymbol("nvmlDeviceGetMigMode") - if err != nil { + if !d.lib.hasSymbol("nvmlDeviceGetMigMode") { return false, nil } @@ -465,22 +462,12 @@ func (d *devicelib) GetMigProfiles() ([]MigProfile, error) { return profiles, nil } -// nvmlLookupSymbol checks to see if the given symbol is present in the NVML library -func (d *devicelib) nvmlLookupSymbol(symbol string) error { - // If devicelib is configured to not verify symbols, then we short-circuit here +// hasSymbol checks to see if the given symbol is present in the NVML library. +// If devicelib is configured to not verify symbols, then all symbols are assumed to exist. +func (d *devicelib) hasSymbol(symbol string) bool { if !*d.verifySymbols { - return nil + return true } - // Otherwise we lookup the provided symbol and verify it is available - lib := dl.New("libnvidia-ml.so.1", dl.RTLD_LAZY|dl.RTLD_GLOBAL) - if lib == nil { - return fmt.Errorf("error instantiating DynamicLibrary for NVML") - } - err := lib.Open() - if err != nil { - return fmt.Errorf("error opening DynamicLibrary for NVML: %v", err) - } - defer lib.Close() - return lib.Lookup(symbol) + return d.nvml.Lookup(symbol) == nil } diff --git a/pkg/nvml/nvml.go b/pkg/nvml/nvml.go index 4e3eaa6..f5edb5d 100644 --- a/pkg/nvml/nvml.go +++ b/pkg/nvml/nvml.go @@ -34,6 +34,12 @@ func New() Interface { return &nvmlLib{} } +// Lookup checks whether the specified symbol exists in the configured NVML library. +func (n *nvmlLib) Lookup(name string) error { + // TODO: For now we rely on the default NVML library and perform the lookups against this. + return nvml.GetLibrary().Lookup(name) +} + // Init initializes an NVML Interface func (n *nvmlLib) Init() Return { ret := nvml.Init() diff --git a/pkg/nvml/nvml_mock.go b/pkg/nvml/nvml_mock.go index 5243a1f..00957f8 100644 --- a/pkg/nvml/nvml_mock.go +++ b/pkg/nvml/nvml_mock.go @@ -35,6 +35,9 @@ var _ Interface = &InterfaceMock{} // InitFunc: func() Return { // panic("mock out the Init method") // }, +// LookupFunc: func(s string) error { +// panic("mock out the Lookup method") +// }, // ShutdownFunc: func() Return { // panic("mock out the Shutdown method") // }, @@ -69,6 +72,9 @@ type InterfaceMock struct { // InitFunc mocks the Init method. InitFunc func() Return + // LookupFunc mocks the Lookup method. + LookupFunc func(s string) error + // ShutdownFunc mocks the Shutdown method. ShutdownFunc func() Return @@ -104,6 +110,11 @@ type InterfaceMock struct { // Init holds details about calls to the Init method. Init []struct { } + // Lookup holds details about calls to the Lookup method. + Lookup []struct { + // S is the s argument value. + S string + } // Shutdown holds details about calls to the Shutdown method. Shutdown []struct { } @@ -120,6 +131,7 @@ type InterfaceMock struct { lockErrorString sync.RWMutex lockEventSetCreate sync.RWMutex lockInit sync.RWMutex + lockLookup sync.RWMutex lockShutdown sync.RWMutex lockSystemGetCudaDriverVersion sync.RWMutex lockSystemGetDriverVersion sync.RWMutex @@ -302,6 +314,38 @@ func (mock *InterfaceMock) InitCalls() []struct { return calls } +// Lookup calls LookupFunc. +func (mock *InterfaceMock) Lookup(s string) error { + if mock.LookupFunc == nil { + panic("InterfaceMock.LookupFunc: method is nil but Interface.Lookup was just called") + } + callInfo := struct { + S string + }{ + S: s, + } + mock.lockLookup.Lock() + mock.calls.Lookup = append(mock.calls.Lookup, callInfo) + mock.lockLookup.Unlock() + return mock.LookupFunc(s) +} + +// LookupCalls gets all the calls that were made to Lookup. +// Check the length with: +// +// len(mockedInterface.LookupCalls()) +func (mock *InterfaceMock) LookupCalls() []struct { + S string +} { + var calls []struct { + S string + } + mock.lockLookup.RLock() + calls = mock.calls.Lookup + mock.lockLookup.RUnlock() + return calls +} + // Shutdown calls ShutdownFunc. func (mock *InterfaceMock) Shutdown() Return { if mock.ShutdownFunc == nil { diff --git a/pkg/nvml/types.go b/pkg/nvml/types.go index 39d005f..2aaaf08 100644 --- a/pkg/nvml/types.go +++ b/pkg/nvml/types.go @@ -30,6 +30,7 @@ type Interface interface { ErrorString(r Return) string EventSetCreate() (EventSet, Return) Init() Return + Lookup(string) error Shutdown() Return SystemGetCudaDriverVersion() (int, Return) SystemGetDriverVersion() (string, Return)