fix(xunix): tighten up gpuExtraRegex to avoid extraneous matches

Author: johnstcn

xunix/gpu.go

@@ -18,10 +18,9 @@ import (
 )
 
 var (
-	gpuMountRegex     = regexp.MustCompile("(?i)(nvidia|vulkan|cuda)")
-	gpuExtraRegex     = regexp.MustCompile("(?i)(libgl|nvidia|vulkan|cuda)")
-	gpuEnvRegex       = regexp.MustCompile("(?i)nvidia")
-	sharedObjectRegex = regexp.MustCompile(`\.so(\.[0-9\.]+)?$`)
+	gpuMountRegex = regexp.MustCompile(`(?i)(nvidia|vulkan|cuda)`)
+	gpuExtraRegex = regexp.MustCompile(`(?i)(libgl(e|sx|\.)|nvidia|vulkan|cuda)`)
+	gpuEnvRegex   = regexp.MustCompile(`(?i)nvidia`)
 )
 
 func GPUEnvs(ctx context.Context) []string {
@@ -65,8 +64,12 @@ func GPUs(ctx context.Context, log slog.Logger, usrLibDir string) ([]Device, []m
 			}
 
 			// If it's not in /dev treat it as a bind mount.
-			links, err := SameDirSymlinks(afs, m.Path)
 			binds = append(binds, m)
+			// We also want to find any symlinks that point to the target.
+			// This is important for the nvidia driver as it mounts the driver
+			// files with the driver version appended to the end, and creates
+			// symlinks that point to the actual files.
+			links, err := SameDirSymlinks(afs, m.Path)
 			if err != nil {
 				log.Error(ctx, "find symlinks", slog.F("path", m.Path), slog.Error(err))
 			} else {
@@ -118,7 +121,7 @@ func usrLibGPUs(ctx context.Context, log slog.Logger, usrLibDir string) ([]mount
 				return nil
 			}
 
-			if !sharedObjectRegex.MatchString(path) || !gpuExtraRegex.MatchString(path) {
+			if !gpuExtraRegex.MatchString(path) {
 				return nil
 			}
 

xunix/gpu_test.go

@@ -62,11 +62,17 @@ func TestGPUs(t *testing.T) {
 				filepath.Join(usrLibMountpoint, "nvidia", "libglxserver_nvidia.so.1"),
 			}
 
-			// fakeUsrLibFiles are files that should be written to the "mounted"
-			// /usr/lib directory. It includes files that shouldn't be returned.
-			fakeUsrLibFiles = append([]string{
+			// fakeUsrLibFiles are files that we do not expect to be returned
+			// bind mounts for.
+			fakeUsrLibFiles = []string{
 				filepath.Join(usrLibMountpoint, "libcurl-gnutls.so"),
-			}, expectedUsrLibFiles...)
+				filepath.Join(usrLibMountpoint, "libglib.so"),
+			}
+
+			// allUsrLibFiles are all the files that should be written to the
+			// "mounted" /usr/lib directory. It includes files that shouldn't
+			// be returned.
+			allUsrLibFiles = append(expectedUsrLibFiles, fakeUsrLibFiles...)
 		)
 
 		ctx := xunix.WithFS(context.Background(), fs)
@@ -93,15 +99,19 @@ func TestGPUs(t *testing.T) {
 		err := fs.MkdirAll(filepath.Join(usrLibMountpoint, "nvidia"), 0o755)
 		require.NoError(t, err)
 
-		for _, file := range fakeUsrLibFiles {
+		for _, file := range allUsrLibFiles {
 			_, err = fs.Create(file)
 			require.NoError(t, err)
 		}
+		for _, mp := range mounter.MountPoints {
+			_, err = fs.Create(mp.Path)
+			require.NoError(t, err)
+		}
 
 		devices, binds, err := xunix.GPUs(ctx, log, usrLibMountpoint)
 		require.NoError(t, err)
 		require.Len(t, devices, 2, "unexpected 2 nvidia devices")
-		require.Len(t, binds, 4, "expected 4 nvidia binds")
+		require.Len(t, binds, 5, "expected 5 nvidia binds")
 		require.Contains(t, binds, mount.MountPoint{
 			Device: "/dev/sda1",
 			Path:   "/usr/local/nvidia",
@@ -113,6 +123,12 @@ func TestGPUs(t *testing.T) {
 				Opts: []string{"ro"},
 			})
 		}
+		for _, file := range fakeUsrLibFiles {
+			require.NotContains(t, binds, mount.MountPoint{
+				Path: file,
+				Opts: []string{"ro"},
+			})
+		}
 	})
 }