Use GetLibrary().Lookup() in nvml package

This change uses the GetLibrary().Lookup() function in the nvml package
to check whether a particular function is available. This avoids
the need to explicitly open a library, for example.

Signed-off-by: Evan Lezar <elezar@nvidia.com>

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
 }

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()

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 {

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)