[no-relnote] Add LibCudaPath to Driver type

This change moves the logic to locate libcuda.so files
to the Driver type.

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

internal/discover/graphics.go

@@ -27,7 +27,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"
 )
 
@@ -221,14 +220,11 @@ func (d graphicsDriverLibraries) isDriverLibrary(filename string, libraryName st
 // 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(".*.*")
+	libcudaPath, err := driver.LibCudaPath()
 	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.")

internal/lookup/cuda/cuda.go

@@ -17,12 +17,14 @@
 package root
 
 import (
+	"fmt"
 	"os"
 	"path/filepath"
 	"strings"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/logger"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup"
+	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/symlinks"
 )
 
 // Driver represents a filesystem in which a set of drivers or devices is defined.
@@ -79,6 +81,34 @@ func (r *Driver) Libraries() lookup.Locator {
 	)
 }
 
+// LibCudaPath returns the path to libcuda.so.RM_VERSION in the driver installation.
+func (r *Driver) LibCudaPath() (string, error) {
+	candidates, err := r.Libraries().Locate("libcuda.so.1")
+	if err != nil {
+		return "", err
+	}
+
+	pattern := "*.*"
+	var matched []string
+	for _, candidate := range candidates {
+		// It is not guaranteed that candidate is a resolved symlink.
+		resolved, err := symlinks.Resolve(candidate)
+		if err != nil {
+			continue
+		}
+		matches, err := filepath.Match("libcuda.so"+pattern, filepath.Base(resolved))
+		if err != nil || !matches {
+			continue
+		}
+		matched = append(matched, candidate)
+	}
+
+	if len(matched) == 0 {
+		return "", fmt.Errorf("%v %w", "libcuda.so"+pattern, lookup.ErrNotFound)
+	}
+	return matched[0], nil
+}
+
 // Configs returns a locator for driver configs.
 // If configSearchPaths is specified, these paths are used as absolute paths,
 // otherwise, /etc and /usr/share are searched.

internal/lookup/root/root.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,20 @@ 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(".*")
+			libCudaPath, err := l.LibCudaPath()
 			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))
-			}
+			libCudaPath = strings.TrimPrefix(libCudaPath, "/private")
 
-			require.EqualValues(t, expectedWithRoot, strippedCandidates)
+			if tc.expected == "" {
+				require.EqualValues(t, "", libCudaPath)
+			} else {
+				require.EqualValues(t, filepath.Join(driverRoot, tc.expected), libCudaPath)
+			}
 		})
 	}
 }
@@ -101,5 +95,8 @@ func setupDriverRoot(t *testing.T, libCudaPath string) (string, error) {
 	}
 	defer libCuda.Close()
 
+	if err := os.Symlink(filepath.Base(libCudaPath), filepath.Join(driverRoot, filepath.Dir(libCudaPath), "libcuda.so.1")); err != nil {
+		return "", err
+	}
 	return driverRoot, nil
 }

internal/lookup/cuda/cuda_test.go renamed to internal/lookup/root/root_test.go

@@ -28,7 +28,6 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
 	"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"
 )
 
@@ -194,13 +193,16 @@ func NewDriverBinariesDiscoverer(logger logger.Interface, driverRoot string) dis
 func getVersionLibs(logger logger.Interface, driver *root.Driver, version string) ([]string, error) {
 	logger.Infof("Using driver version %v", version)
 
-	libCudaPaths, err := cuda.New(
-		driver.Libraries(),
-	).Locate("." + version)
+	libCudaPath, err := driver.LibCudaPath()
 	if err != nil {
-		return nil, fmt.Errorf("failed to locate libcuda.so.%v: %v", version, err)
+		return nil, fmt.Errorf("failed to locate libcuda.so: %w", err)
 	}
-	libRoot := filepath.Dir(libCudaPaths[0])
+	matched, err := filepath.Match("libcuda.so."+version, libCudaPath)
+	if err != nil || !matched {
+		return nil, fmt.Errorf("failed to locate libcuda.so.%v: %w", version, err)
+	}
+
+	libRoot := filepath.Dir(libCudaPath)
 
 	libraries := lookup.NewFileLocator(
 		lookup.WithLogger(logger),

pkg/nvcdi/driver-nvml.go

@@ -27,7 +27,6 @@ import (
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
-	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/cuda"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/nvsandboxutils"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/spec"
 )
@@ -100,15 +99,11 @@ func (m *managementlib) getCudaVersion() (string, error) {
 		return version, nil
 	}
 
-	libCudaPaths, err := cuda.New(
-		m.driver.Libraries(),
-	).Locate(".*.*")
+	libCudaPath, err := m.driver.LibCudaPath()
 	if err != nil {
 		return "", fmt.Errorf("failed to locate libcuda.so: %v", err)
 	}
 
-	libCudaPath := libCudaPaths[0]
-
 	version = strings.TrimPrefix(filepath.Base(libCudaPath), "libcuda.so.")
 
 	return version, nil