From 020f598d9a1dd94fcb2582688df1f5b4a44cecc5 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Wed, 3 Apr 2024 16:20:43 +0200
Subject: [PATCH] Generate hook to create .so symlinks in CDI

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 cmd/nvidia-ctk/cdi/generate/generate.go |  8 ++++++++
 pkg/nvcdi/common-nvml.go                |  2 +-
 pkg/nvcdi/driver-nvml.go                | 25 ++++++++++++++++++-------
 pkg/nvcdi/lib.go                        |  2 ++
 pkg/nvcdi/options.go                    |  7 +++++++
 5 files changed, 36 insertions(+), 8 deletions(-)

diff --git a/cmd/nvidia-ctk/cdi/generate/generate.go b/cmd/nvidia-ctk/cdi/generate/generate.go
index 5653b342..71b23da0 100644
--- a/cmd/nvidia-ctk/cdi/generate/generate.go
+++ b/cmd/nvidia-ctk/cdi/generate/generate.go
@@ -60,6 +60,8 @@ type options struct {
 		files          cli.StringSlice
 		ignorePatterns cli.StringSlice
 	}
+
+	noDotSoSymlinks bool
 }
 
 // NewCommand constructs a generate-cdi command with the specified logger
@@ -166,6 +168,11 @@ func (m command) build() *cli.Command {
 			Usage:       "Specify a pattern the CSV mount specifications.",
 			Destination: &opts.csv.ignorePatterns,
 		},
+		&cli.BoolFlag{
+			Name:        "no-dot-so-symlinks",
+			Usage:       "Skip the generation of a hook for creating .so symlinks to driver files in the container",
+			Destination: &opts.noDotSoSymlinks,
+		},
 	}
 
 	return &c
@@ -270,6 +277,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		nvcdi.WithLibrarySearchPaths(opts.librarySearchPaths.Value()),
 		nvcdi.WithCSVFiles(opts.csv.files.Value()),
 		nvcdi.WithCSVIgnorePatterns(opts.csv.ignorePatterns.Value()),
+		nvcdi.WithNoDotSoSymlinks(opts.noDotSoSymlinks),
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create CDI library: %v", err)
diff --git a/pkg/nvcdi/common-nvml.go b/pkg/nvcdi/common-nvml.go
index f4bfe30a..ce7708c1 100644
--- a/pkg/nvcdi/common-nvml.go
+++ b/pkg/nvcdi/common-nvml.go
@@ -41,7 +41,7 @@ func (l *nvmllib) newCommonNVMLDiscoverer() (discover.Discover, error) {
 		l.logger.Warningf("failed to create discoverer for graphics mounts: %v", err)
 	}
 
-	driverFiles, err := NewDriverDiscoverer(l.logger, l.driver, l.nvidiaCTKPath, l.ldconfigPath, l.nvmllib)
+	driverFiles, err := l.newDriverDiscoverer()
 	if err != nil {
 		return nil, fmt.Errorf("failed to create discoverer for driver files: %v", err)
 	}
diff --git a/pkg/nvcdi/driver-nvml.go b/pkg/nvcdi/driver-nvml.go
index 10d154d6..ffee6282 100644
--- a/pkg/nvcdi/driver-nvml.go
+++ b/pkg/nvcdi/driver-nvml.go
@@ -32,24 +32,35 @@ import (
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/lookup/root"
 )
 
-// NewDriverDiscoverer creates a discoverer for the libraries and binaries associated with a driver installation.
+// newDriverDiscoverer creates a discoverer for the libraries and binaries associated with a driver installation.
 // The supplied NVML Library is used to query the expected driver version.
-func NewDriverDiscoverer(logger logger.Interface, driver *root.Driver, nvidiaCTKPath string, ldconfigPath string, nvmllib nvml.Interface) (discover.Discover, error) {
-	if r := nvmllib.Init(); r != nvml.SUCCESS {
+func (l *nvmllib) newDriverDiscoverer() (discover.Discover, error) {
+	if r := l.nvmllib.Init(); r != nvml.SUCCESS {
 		return nil, fmt.Errorf("failed to initialize NVML: %v", r)
 	}
 	defer func() {
-		if r := nvmllib.Shutdown(); r != nvml.SUCCESS {
-			logger.Warningf("failed to shutdown NVML: %v", r)
+		if r := l.nvmllib.Shutdown(); r != nvml.SUCCESS {
+			l.logger.Warningf("failed to shutdown NVML: %v", r)
 		}
 	}()
 
-	version, r := nvmllib.SystemGetDriverVersion()
+	version, r := l.nvmllib.SystemGetDriverVersion()
 	if r != nvml.SUCCESS {
 		return nil, fmt.Errorf("failed to determine driver version: %v", r)
 	}
 
-	return newDriverVersionDiscoverer(logger, driver, nvidiaCTKPath, ldconfigPath, version)
+	driver, err := newDriverVersionDiscoverer(l.logger, l.driver, l.nvidiaCTKPath, l.ldconfigPath, version)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create discoverer: %w", err)
+	}
+	discoverers := []discover.Discover{driver}
+
+	if !l.noDotSoSymlinks {
+		createDotSoSymlinksHook := discover.NewDotSoSymlinksDiscoverer(l.nvidiaCTKPath, version)
+		discoverers = append(discoverers, createDotSoSymlinksHook)
+	}
+
+	return discover.Merge(discoverers...), nil
 }
 
 func newDriverVersionDiscoverer(logger logger.Interface, driver *root.Driver, nvidiaCTKPath, ldconfigPath, version string) (discover.Discover, error) {
diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go
index ad75a239..fcb834ff 100644
--- a/pkg/nvcdi/lib.go
+++ b/pkg/nvcdi/lib.go
@@ -63,6 +63,8 @@ type nvcdilib struct {
 	infolib info.Interface
 
 	mergedDeviceOptions []transform.MergedDeviceOption
+
+	noDotSoSymlinks bool
 }
 
 // New creates a new nvcdi library
diff --git a/pkg/nvcdi/options.go b/pkg/nvcdi/options.go
index a4e49b60..fde5a9dc 100644
--- a/pkg/nvcdi/options.go
+++ b/pkg/nvcdi/options.go
@@ -140,3 +140,10 @@ func WithLibrarySearchPaths(paths []string) Option {
 		o.librarySearchPaths = paths
 	}
 }
+
+// WithNoDotSoSymlinks sets the no-dot-so-symlinks feature.
+func WithNoDotSoSymlinks(noDotSoSymlinks bool) Option {
+	return func(o *nvcdilib) {
+		o.noDotSoSymlinks = noDotSoSymlinks
+	}
+}