Consolidate logic to determine driver version

internal/discover/compat_libs.go

@@ -4,16 +4,14 @@ import (
 	"strings"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/root"
 )
 
 // NewCUDACompatHookDiscoverer creates a discoverer for a enable-cuda-compat hook.
 // This hook is responsible for setting up CUDA compatibility in the container and depends on the host driver version.
-func NewCUDACompatHookDiscoverer(logger logger.Interface, hookCreator HookCreator, driver *root.Driver) Discover {
-	_, cudaVersionPattern := getCUDALibRootAndVersionPattern(logger, driver)
+func NewCUDACompatHookDiscoverer(logger logger.Interface, hookCreator HookCreator, version string) Discover {
 	var args []string
-	if !strings.Contains(cudaVersionPattern, "*") {
-		args = append(args, "--host-driver-version="+cudaVersionPattern)
+	if version != "" && !strings.Contains(version, "*") {
+		args = append(args, "--host-driver-version="+version)
 	}
 
 	return hookCreator.Create("enable-cuda-compat", args...)

internal/discover/graphics.go

@@ -28,7 +28,6 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/info/proc"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/cuda"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/root"
 )
 
@@ -51,7 +50,10 @@ func NewDRMNodesDiscoverer(logger logger.Interface, devices image.VisibleDevices
 
 // NewGraphicsMountsDiscoverer creates a discoverer for the mounts required by graphics tools such as vulkan.
 func NewGraphicsMountsDiscoverer(logger logger.Interface, driver *root.Driver, hookCreator HookCreator) (Discover, error) {
-	libraries := newGraphicsLibrariesDiscoverer(logger, driver, hookCreator)
+	libraries, err := newGraphicsLibrariesDiscoverer(logger, driver, hookCreator)
+	if err != nil {
+		return nil, fmt.Errorf("failed to construct discoverer for graphics libraries: %w", err)
+	}
 
 	configs := NewMounts(
 		logger,
@@ -113,8 +115,15 @@ type graphicsDriverLibraries struct {
 
 var _ Discover = (*graphicsDriverLibraries)(nil)
 
-func newGraphicsLibrariesDiscoverer(logger logger.Interface, driver *root.Driver, hookCreator HookCreator) Discover {
-	cudaLibRoot, cudaVersionPattern := getCUDALibRootAndVersionPattern(logger, driver)
+func newGraphicsLibrariesDiscoverer(logger logger.Interface, driver *root.Driver, hookCreator HookCreator) (Discover, error) {
+	cudaVersionPattern, err := driver.Version()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get driver version: %w", err)
+	}
+	cudaLibRoot, err := driver.GetLibcudaParentDir()
+	if err != nil {
+		return nil, fmt.Errorf("failed to get libcuda.so parent directory: %w", err)
+	}
 
 	libraries := NewMounts(
 		logger,
@@ -155,7 +164,7 @@ func newGraphicsLibrariesDiscoverer(logger logger.Interface, driver *root.Driver
 		Discover:    Merge(libraries, xorgLibraries),
 		logger:      logger,
 		hookCreator: hookCreator,
-	}
+	}, nil
 }
 
 // Mounts discovers the required libraries and filters out libnvidia-allocator.so.
@@ -228,30 +237,6 @@ func (d graphicsDriverLibraries) isDriverLibrary(filename string, libraryName st
 	return match
 }
 
-// getCUDALibRootAndVersionPattern returns the parent directory and the version
-// suffix of the libcuda.so.*.* library at the driver root.
-// If the library cannot be located an empty root is returned.
-// If the version string cannot be extracted, the generic *.* pattern is returned.
-func getCUDALibRootAndVersionPattern(logger logger.Interface, driver *root.Driver) (string, string) {
-	libCudaPaths, err := cuda.New(
-		driver.Libraries(),
-	).Locate(".*.*")
-	if err != nil {
-		logger.Warningf("failed to locate libcuda.so: %v; using *.*", err)
-		return "", "*.*"
-	}
-	libcudaPath := libCudaPaths[0]
-
-	libRoot := filepath.Dir(libcudaPath)
-	version := strings.TrimPrefix(filepath.Base(libcudaPath), "libcuda.so.")
-	if version == "" {
-		logger.Warningf("failed to extract version from %v; using *.*", libcudaPath)
-		version = "*.*"
-	}
-
-	return driver.RelativeToRoot(libRoot), version
-}
-
 // buildXOrgSearchPaths returns the ordered list of search paths for XOrg files.
 func buildXOrgSearchPaths(libRoot string) []string {
 	var paths []string

internal/lookup/cuda/cuda_test.go → internal/lookup/root/cuda_test.go

@@ -14,7 +14,7 @@
 # limitations under the License.
 **/
 
-package cuda
+package root
 
 import (
 	"fmt"
@@ -35,19 +35,19 @@ func TestLocate(t *testing.T) {
 	testCases := []struct {
 		description   string
 		libcudaPath   string
-		expected      []string
+		expected      string
 		expectedError error
 	}{
 		{
 			description:   "no libcuda does not resolve library",
 			libcudaPath:   "",
-			expected:      []string{},
+			expected:      "",
 			expectedError: lookup.ErrNotFound,
 		},
 		{
 			description:   "no-ldcache searches /usr/lib64",
 			libcudaPath:   "/usr/lib64/libcuda.so.123.34",
-			expected:      []string{"/usr/lib64/libcuda.so.123.34"},
+			expected:      "/usr/lib64/libcuda.so.123.34",
 			expectedError: nil,
 		},
 	}
@@ -58,26 +58,17 @@ func TestLocate(t *testing.T) {
 			require.NoError(t, err)
 
 			l := New(
-				lookup.NewLibraryLocator(
-					lookup.WithLogger(logger),
-					lookup.WithRoot(driverRoot),
-				),
+				WithLogger(logger),
+				WithDriverRoot(driverRoot),
 			)
 
-			candidates, err := l.Locate(".*")
+			libcudasoPath, err := l.GetLibcudasoPath()
 			require.ErrorIs(t, err, tc.expectedError)
 
-			var strippedCandidates []string
-			for _, c := range candidates {
-				// NOTE: We need to strip `/private` on MacOs due to symlink resolution
-				strippedCandidates = append(strippedCandidates, strings.TrimPrefix(c, "/private"))
-			}
-			var expectedWithRoot []string
-			for _, e := range tc.expected {
-				expectedWithRoot = append(expectedWithRoot, filepath.Join(driverRoot, e))
-			}
-
-			require.EqualValues(t, expectedWithRoot, strippedCandidates)
+			// NOTE: We need to strip `/private` on MacOs due to symlink resolution
+			stripped := strings.TrimPrefix(libcudasoPath, "/private")
+
+			require.Equal(t, tc.expected, stripped)
 		})
 	}
 }
@@ -101,5 +92,5 @@ func setupDriverRoot(t *testing.T, libCudaPath string) (string, error) {
 	}
 	defer libCuda.Close()
 
-	return driverRoot, nil
+	return filepath.EvalSymlinks(driverRoot)
 }

internal/lookup/root/options.go

@@ -18,28 +18,45 @@ package root
 
 import "github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 
-type Option func(*Driver)
+type options struct {
+	logger logger.Interface
+	// Root represents the root from the perspective of the driver libraries and binaries.
+	Root string
+	// librarySearchPaths specifies explicit search paths for discovering libraries.
+	librarySearchPaths []string
+	// configSearchPaths specified explicit search paths for discovering driver config files.
+	configSearchPaths []string
+	versioner         Versioner
+}
+
+type Option func(*options)
 
 func WithLogger(logger logger.Interface) Option {
-	return func(d *Driver) {
-		d.logger = logger
+	return func(o *options) {
+		o.logger = logger
 	}
 }
 
 func WithDriverRoot(root string) Option {
-	return func(d *Driver) {
-		d.Root = root
+	return func(o *options) {
+		o.Root = root
 	}
 }
 
 func WithLibrarySearchPaths(paths ...string) Option {
-	return func(d *Driver) {
-		d.librarySearchPaths = paths
+	return func(o *options) {
+		o.librarySearchPaths = paths
 	}
 }
 
 func WithConfigSearchPaths(paths ...string) Option {
-	return func(d *Driver) {
-		d.configSearchPaths = paths
+	return func(o *options) {
+		o.configSearchPaths = paths
+	}
+}
+
+func WithVersioner(versioner Versioner) Option {
+	return func(o *options) {
+		o.versioner = versioner
 	}
 }