feat(xunix): add SameDirSymlinks, add symlinks pointing to mounted libs

Author: johnstcn

xunix/gpu.go

@@ -6,6 +6,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"slices"
 	"sort"
 	"strings"
 
@@ -39,6 +40,7 @@ func GPUEnvs(ctx context.Context) []string {
 
 func GPUs(ctx context.Context, log slog.Logger, usrLibDir string) ([]Device, []mount.MountPoint, error) {
 	var (
+		afs     = GetFS(ctx)
 		mounter = Mounter(ctx)
 		devices = []Device{}
 		binds   = []mount.MountPoint{}
@@ -63,7 +65,19 @@ 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)
+			if err != nil {
+				log.Error(ctx, "find symlinks", slog.F("path", m.Path), slog.Error(err))
+			} else {
+				for _, link := range links {
+					log.Debug(ctx, "found symlink", slog.F("link", link), slog.F("target", m.Path))
+					binds = append(binds, mount.MountPoint{
+						Path: link,
+						Opts: []string{"ro"},
+					})
+				}
+			}
 		}
 	}
 
@@ -176,6 +190,68 @@ func recursiveSymlinks(afs FS, mountpoint string, path string) ([]string, error)
 	return paths, nil
 }
 
+// SameDirSymlinks returns all links in the same directory as `target` that
+// point to target, either indirectly or directly. Only symlinks in the same
+// directory as `target` are considered.
+func SameDirSymlinks(afs FS, target string) ([]string, error) {
+	var (
+		found         = make([]string, 0)
+		maxIterations = 10 // arbitrary upper limit to prevent infinite loops
+	)
+	for range maxIterations {
+		foundThisTime := false
+		fis, err := afero.ReadDir(afs, filepath.Dir(target))
+		if err != nil {
+			return nil, xerrors.Errorf("read dir %q: %w", filepath.Dir(target), err)
+		}
+		for _, fi := range fis {
+			// Ignore the target itself.
+			if fi.Name() == filepath.Base(target) {
+				continue
+			}
+			// Ignore non-symlinks.
+			if fi.Mode()&os.ModeSymlink == 0 {
+				continue
+			}
+			// Get the target of the symlink.
+			link, err := afs.Readlink(filepath.Join(filepath.Dir(target), fi.Name()))
+			if err != nil {
+				return nil, xerrors.Errorf("readlink %q: %w", fi.Name(), err)
+			}
+			// Make the link absolute.
+			if !filepath.IsAbs(link) {
+				link = filepath.Join(filepath.Dir(target), link)
+			}
+			// Ignore symlinks that point outside of target's directory.
+			if filepath.Dir(link) != filepath.Dir(target) {
+				continue
+			}
+
+			// Check if the symlink points to to the target, or if it points
+			// to one of the symlinks we've already found.
+			if link != target {
+				if !slices.Contains(found, link) {
+					continue
+				}
+			}
+
+			// Have we already seen this target?
+			fullPath := filepath.Join(filepath.Dir(target), fi.Name())
+			if slices.Contains(found, fullPath) {
+				continue
+			}
+
+			found = append(found, filepath.Join(filepath.Dir(target), fi.Name()))
+			foundThisTime = true
+		}
+		// If we didn't find any symlinks this time, we're done.
+		if !foundThisTime {
+			break
+		}
+	}
+	return found, nil
+}
+
 // TryUnmountProcGPUDrivers unmounts any GPU-related mounts under /proc as it causes
 // issues when creating any container in some cases. Errors encountered while
 // unmounting are treated as non-fatal.

xunix/gpu_test.go

@@ -2,10 +2,13 @@ package xunix_test
 
 import (
 	"context"
+	"os"
 	"path/filepath"
+	"sort"
 	"testing"
 
 	"github.com/spf13/afero"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"k8s.io/mount-utils"
 
@@ -112,3 +115,103 @@ func TestGPUs(t *testing.T) {
 		}
 	})
 }
+
+func Test_SameDirSymlinks(t *testing.T) {
+	t.Parallel()
+
+	var (
+		ctx = context.Background()
+		// We need to test with a real filesystem as the fake one doesn't
+		// support creating symlinks.
+		tmpDir = t.TempDir()
+		// We do test with the interface though!
+		afs = xunix.GetFS(ctx)
+	)
+
+	// Create some files in the temporary directory.
+	_, err := os.Create(filepath.Join(tmpDir, "file1.real"))
+	require.NoError(t, err, "create file")
+	_, err = os.Create(filepath.Join(tmpDir, "file2.real"))
+	require.NoError(t, err, "create file2")
+	_, err = os.Create(filepath.Join(tmpDir, "file3.real"))
+	require.NoError(t, err, "create file3")
+	_, err = os.Create(filepath.Join(tmpDir, "file4.real"))
+	require.NoError(t, err, "create file4")
+
+	// Create a symlink to the file in the temporary directory.
+	// This needs to be done by the real os package.
+	err = os.Symlink(filepath.Join(tmpDir, "file1.real"), filepath.Join(tmpDir, "file1.link1"))
+	require.NoError(t, err, "create first symlink")
+
+	// Create another symlink to the previous symlink.
+	err = os.Symlink(filepath.Join(tmpDir, "file1.link1"), filepath.Join(tmpDir, "file1.link2"))
+	require.NoError(t, err, "create second symlink")
+
+	// Create a symlink to a file outside of the temporary directory.
+	err = os.MkdirAll(filepath.Join(tmpDir, "dir"), 0o755)
+	require.NoError(t, err, "create dir")
+	// Create a symlink from file2 to inside the dir.
+	err = os.Symlink(filepath.Join(tmpDir, "file2.real"), filepath.Join(tmpDir, "dir", "file2.link1"))
+	require.NoError(t, err, "create dir symlink")
+
+	// Create a symlink with a relative path. To do this, we need to
+	// change the working directory to the temporary directory.
+	oldWorkingDir, err := os.Getwd()
+	require.NoError(t, err, "get working dir")
+	// Change the working directory to the temporary directory.
+	require.NoError(t, os.Chdir(tmpDir), "change working dir")
+	err = os.Symlink(filepath.Join(tmpDir, "file4.real"), "file4.link1")
+	require.NoError(t, err, "create relative symlink")
+	// Change the working directory back to the original.
+	require.NoError(t, os.Chdir(oldWorkingDir), "change working dir back")
+
+	for _, tt := range []struct {
+		name     string
+		expected []string
+	}{
+		{
+			// Two symlinks to the same file.
+			name: "file1.real",
+			expected: []string{
+				filepath.Join(tmpDir, "file1.link1"),
+				filepath.Join(tmpDir, "file1.link2"),
+			},
+		},
+		{
+			// Mid-way in the symlink chain.
+			name: "file1.link1",
+			expected: []string{
+				filepath.Join(tmpDir, "file1.link2"),
+			},
+		},
+		{
+			// End of the symlink chain.
+			name:     "file1.link2",
+			expected: []string{},
+		},
+		{
+			// Symlink to a file outside of the temporary directory.
+			name:     "file2.real",
+			expected: []string{},
+		},
+		{
+			// No symlinks to this file.
+			name:     "file3.real",
+			expected: []string{},
+		},
+		{
+			// One relative symlink.
+			name:     "file4.real",
+			expected: []string{filepath.Join(tmpDir, "file4.link1")},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+			fullPath := filepath.Join(tmpDir, tt.name)
+			actual, err := xunix.SameDirSymlinks(afs, fullPath)
+			require.NoError(t, err, "find symlink")
+			sort.Strings(actual)
+			assert.Equal(t, tt.expected, actual, "find symlinks %q", tt.name)
+		})
+	}
+}